fix: make most tests pass

Author: AayushSabharwal

src/ModelingToolkit.jl

@@ -63,7 +63,7 @@ using PrecompileTools, Reexport
         ParallelForm, SerialForm, MultithreadedForm, build_function,
         rhss, lhss, prettify_expr, gradient,
         jacobian, hessian, derivative, sparsejacobian, sparsehessian,
-        substituter, scalarize, getparent
+        substituter, scalarize, getparent, hasderiv, hasdiff
 
     import DiffEqBase: @add_kwonly
     import OrdinaryDiffEq

src/systems/abstractsystem.jl

@@ -187,8 +187,14 @@ function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym)
     if unwrap(sym) isa Int    # [x, 1] coerces 1 to a Num
         return unwrap(sym) in 1:length(variable_symbols(sys))
     end
-    return any(isequal(sym), variable_symbols(sys)) ||
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        ic = get_index_cache(sys)
+        h = getsymbolhash(sym)
+        return haskey(ic.unknown_idx, h) || haskey(ic.unknown_idx, getsymbolhash(default_toterm(sym))) || hasname(sym) && is_variable(sys, getname(sym))
+    else
+        return any(isequal(sym), variable_symbols(sys)) ||
            hasname(sym) && is_variable(sys, getname(sym))
+    end
 end
 
 function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym::Symbol)
@@ -202,6 +208,22 @@ function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym)
     if unwrap(sym) isa Int
         return unwrap(sym)
     end
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        ic = get_index_cache(sys)
+        h = getsymbolhash(sym)
+        return if haskey(ic.unknown_idx, h)
+            ic.unknown_idx[h]
+        else
+            h = getsymbolhash(default_toterm(sym))
+            if haskey(ic.unknown_idx, h)
+                ic.unknown_idx[h]
+            elseif hasname(sym)
+                variable_index(sys, getname(sym))
+            else
+                nothing
+            end
+        end
+    end
     idx = findfirst(isequal(sym), variable_symbols(sys))
     if idx === nothing && hasname(sym)
         idx = variable_index(sys, getname(sym))
@@ -230,7 +252,16 @@ function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym)
     if unwrap(sym) isa Int
         return unwrap(sym) in 1:length(parameter_symbols(sys))
     end
-
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        ic = get_index_cache(sys)
+        h = getsymbolhash(sym)
+        return if haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h)
+            true
+        else
+            h = getsymbolhash(default_toterm(sym))
+            haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) || hasname(sym) && is_parameter(sys, getname(sym))
+        end
+    end
     return any(isequal(sym), parameter_symbols(sys)) ||
            hasname(sym) && is_parameter(sys, getname(sym))
 end
@@ -246,6 +277,25 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym)
     if unwrap(sym) isa Int
         return unwrap(sym)
     end
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        ic = get_index_cache(sys)
+        h = getsymbolhash(sym)
+        return if haskey(ic.param_idx, h)
+            ParameterIndex(SciMLStructures.Tunable(), ic.param_idx[h])
+        elseif haskey(ic.discrete_idx, h)
+            ParameterIndex(SciMLStructures.Discrete(), ic.discrete_idx[h])
+        else
+            h = getsymbolhash(default_toterm(sym))
+            if haskey(ic.param_idx, h)
+                ParameterIndex(SciMLStructures.Tunable(), ic.param_idx[h])
+            elseif haskey(ic.discrete_idx, h)
+                ParameterIndex(SciMLStructures.Discrete(), ic.discrete_idx[h])
+            else
+                nothing
+            end
+        end
+    end
+    
     idx = findfirst(isequal(sym), parameter_symbols(sys))
     if idx === nothing && hasname(sym)
         idx = parameter_index(sys, getname(sym))
@@ -1441,12 +1491,18 @@ function linearization_function(sys::AbstractSystem, inputs,
     end
     sys = ssys
     x0 = merge(defaults(sys), Dict(missing_variable_defaults(sys)), op)
-    u0, p, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
-    p, split_idxs = split_parameters_by_type(p)
-    ps = parameters(sys)
-    if p isa Tuple
-        ps = Base.Fix1(getindex, ps).(split_idxs)
-        ps = (ps...,) #if p is Tuple, ps should be Tuple
+    u0, _p, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p = (MTKParameters(sys, p)...,)
+        ps = reorder_parameters(sys, parameters(sys))
+    else
+        p = _p
+        p, split_idxs = split_parameters_by_type(p)
+        ps = parameters(sys)
+        if p isa Tuple
+            ps = Base.Fix1(getindex, ps).(split_idxs)
+            ps = (ps...,) #if p is Tuple, ps should be Tuple
+        end
     end
 
     lin_fun = let diff_idxs = diff_idxs,
@@ -1472,7 +1528,7 @@ function linearization_function(sys::AbstractSystem, inputs,
                 uf = SciMLBase.UJacobianWrapper(fun, t, p)
                 fg_xz = ForwardDiff.jacobian(uf, u)
                 h_xz = ForwardDiff.jacobian(let p = p, t = t
-                        xz -> h(xz, p, t)
+                        xz ->  p isa MTKParameters ? h(xz, p..., t) : h(xz, p, t)
                     end, u)
                 pf = SciMLBase.ParamJacobianWrapper(fun, t, u)
                 fg_u = jacobian_wrt_vars(pf, p, input_idxs, chunk)
@@ -1726,7 +1782,9 @@ function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives =
         p = DiffEqBase.NullParameters())
     x0 = merge(defaults(sys), op)
     u0, p2, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
-
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p2 = MTKParameters(sys, p)
+    end
     linres = lin_fun(u0, p2, t)
     f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u = linres
 

src/systems/callbacks.jl

@@ -1,5 +1,5 @@
 #################################### system operations #####################################
-get_continuous_events(sys::AbstractSystem) = Equation[]
+get_continuous_events(sys::AbstractSystem) = SymbolicContinuousCallback[]
 get_continuous_events(sys::AbstractODESystem) = getfield(sys, :continuous_events)
 has_continuous_events(sys::AbstractSystem) = isdefined(sys, :continuous_events)
 
@@ -15,10 +15,11 @@ struct FunctionalAffect
     sts_syms::Vector{Symbol}
     pars::Vector
     pars_syms::Vector{Symbol}
+    discretes::Vector
     ctx::Any
 end
 
-function FunctionalAffect(f, sts, pars, ctx = nothing)
+function FunctionalAffect(f, sts, pars, discretes, ctx = nothing)
     # sts & pars contain either pairs: resistor.R => R, or Syms: R
     vs = [x isa Pair ? x.first : x for x in sts]
     vs_syms = Symbol[x isa Pair ? Symbol(x.second) : getname(x) for x in sts]
@@ -28,17 +29,18 @@ function FunctionalAffect(f, sts, pars, ctx = nothing)
     ps_syms = Symbol[x isa Pair ? Symbol(x.second) : getname(x) for x in pars]
     length(ps_syms) == length(unique(ps_syms)) || error("Parameters are not unique")
 
-    FunctionalAffect(f, vs, vs_syms, ps, ps_syms, ctx)
+    FunctionalAffect(f, vs, vs_syms, ps, ps_syms, discretes, ctx)
 end
 
-FunctionalAffect(; f, sts, pars, ctx = nothing) = FunctionalAffect(f, sts, pars, ctx)
+FunctionalAffect(; f, sts, pars, discretes, ctx = nothing) = FunctionalAffect(f, sts, pars, discretes, ctx)
 
 func(f::FunctionalAffect) = f.f
 context(a::FunctionalAffect) = a.ctx
 parameters(a::FunctionalAffect) = a.pars
 parameters_syms(a::FunctionalAffect) = a.pars_syms
 unknowns(a::FunctionalAffect) = a.sts
 unknowns_syms(a::FunctionalAffect) = a.sts_syms
+discretes(a::FunctionalAffect) = a.discretes
 
 function Base.:(==)(a1::FunctionalAffect, a2::FunctionalAffect)
     isequal(a1.f, a2.f) && isequal(a1.sts, a2.sts) && isequal(a1.pars, a2.pars) &&
@@ -52,6 +54,7 @@ function Base.hash(a::FunctionalAffect, s::UInt)
     s = hash(a.sts_syms, s)
     s = hash(a.pars, s)
     s = hash(a.pars_syms, s)
+    s = hash(a.discretes, s)
     hash(a.ctx, s)
 end
 
@@ -64,6 +67,7 @@ function namespace_affect(affect::FunctionalAffect, s)
         unknowns_syms(affect),
         renamespace.((s,), parameters(affect)),
         parameters_syms(affect),
+        renamespace.((s,), discretes(affect)),
         context(affect))
 end
 
@@ -121,7 +125,7 @@ end
 
 affects(cb::SymbolicContinuousCallback) = cb.affect
 function affects(cbs::Vector{SymbolicContinuousCallback})
-    reduce(vcat, [affects(cb) for cb in cbs])
+    reduce(vcat, [affects(cb) for cb in cbs], init = [])
 end
 
 namespace_affects(af::Vector, s) = Equation[namespace_affect(a, s) for a in af]
@@ -213,7 +217,7 @@ end
 affects(cb::SymbolicDiscreteCallback) = cb.affects
 
 function affects(cbs::Vector{SymbolicDiscreteCallback})
-    reduce(vcat, affects(cb) for cb in cbs)
+    reduce(vcat, affects(cb) for cb in cbs; init = [])
 end
 
 function namespace_callback(cb::SymbolicDiscreteCallback, s)::SymbolicDiscreteCallback
@@ -241,16 +245,54 @@ end
 ################################# compilation functions ####################################
 
 # handles ensuring that affect! functions work with integrator arguments
-function add_integrator_header(integrator = gensym(:MTKIntegrator), out = :u)
-    expr -> Func([DestructuredArgs(expr.args, integrator, inds = [:u, :p, :t])], [],
-        expr.body),
-    expr -> Func([DestructuredArgs(expr.args, integrator, inds = [out, :u, :p, :t])], [],
-        expr.body)
+function add_integrator_header(sys::AbstractSystem, integrator = gensym(:MTKIntegrator), out = :u)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        function (expr)
+            p = gensym(:p)
+            Func([
+                DestructuredArgs([expr.args[1], p, expr.args[end]],
+                                 integrator, inds = [:u, :p, :t]),
+            ], [], Let([DestructuredArgs([arg.name for arg in expr.args[2:end-1]], p),
+                       expr.args[2:end-1]...], expr.body, false)
+            )
+        end,
+        function (expr)
+            p = gensym(:p)
+            Func([
+                DestructuredArgs([expr.args[1], expr.args[2], p, expr.args[end]],
+                                 integrator, inds = [out, :u, :p, :t]),
+            ], [], Let([DestructuredArgs([arg.name for arg in expr.args[3:end-1]], p),
+                       expr.args[3:end-1]...], expr.body, false)
+            )
+        end
+    else
+        expr -> Func([DestructuredArgs(expr.args, integrator, inds = [:u, :p, :t])], [],
+                expr.body),
+        expr -> Func([DestructuredArgs(expr.args, integrator, inds = [out, :u, :p, :t])], [],
+            expr.body)
+    end
 end
 
-function condition_header(integrator = gensym(:MTKIntegrator))
-    expr -> Func([expr.args[1], expr.args[2],
-            DestructuredArgs(expr.args[3:end], integrator, inds = [:p])], [], expr.body)
+function condition_header(sys::AbstractSystem, integrator = gensym(:MTKIntegrator))
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        function (expr)
+            p = gensym(:p)
+            res = Func(
+                [expr.args[1], expr.args[2], DestructuredArgs([p], integrator, inds = [:p])],
+                [],
+                Let(
+                    [
+                        DestructuredArgs([arg.name for arg in expr.args[3:end]], p),
+                        expr.args[3:end]...
+                    ], expr.body, false
+                )
+            )
+            return res
+        end
+    else
+        expr -> Func([expr.args[1], expr.args[2],
+                    DestructuredArgs(expr.args[3:end], integrator, inds = [:p])], [], expr.body)
+    end
 end
 
 """
@@ -267,15 +309,15 @@ Notes
 function compile_condition(cb::SymbolicDiscreteCallback, sys, dvs, ps;
         expression = Val{true}, kwargs...)
     u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = map(x -> time_varying_as_func(value(x), sys), ps)
+    p = map.(x -> time_varying_as_func(value(x), sys), reorder_parameters(sys, ps))
     t = get_iv(sys)
     condit = condition(cb)
     cs = collect_constants(condit)
     if !isempty(cs)
         cmap = map(x -> x => getdefault(x), cs)
         condit = substitute(condit, cmap)
     end
-    build_function(condit, u, t, p; expression, wrap_code = condition_header(),
+    build_function(condit, u, t, p...; expression, wrap_code = condition_header(sys),
         kwargs...)
 end
 
@@ -325,8 +367,19 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
                 unknownind = Dict(reverse(en) for en in enumerate(dvs))
                 update_inds = map(sym -> unknownind[sym], update_vars)
             elseif isparameter(first(lhss)) && alleq
-                psind = Dict(reverse(en) for en in enumerate(ps))
-                update_inds = map(sym -> psind[sym], update_vars)
+                if has_index_cache(sys) && get_index_cache(sys) !== nothing
+                    ic = get_index_cache(sys)
+                    update_inds = map(update_vars) do sym
+                        @unpack portion, idx = parameter_index(sys, sym)
+                        if portion == SciMLStructures.Discrete()
+                            idx += length(ic.param_idx)
+                        end
+                        idx
+                    end
+                else
+                    psind = Dict(reverse(en) for en in enumerate(ps))
+                    update_inds = map(sym -> psind[sym], update_vars)
+                end
                 outvar = :p
             else
                 error("Error, building an affect function for a callback that wants to modify both parameters and unknowns. This is not currently allowed in one individual callback.")
@@ -335,9 +388,10 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
             update_inds = outputidxs
         end
 
+        ps = reorder_parameters(sys, ps)
         if checkvars
             u = map(x -> time_varying_as_func(value(x), sys), dvs)
-            p = map(x -> time_varying_as_func(value(x), sys), ps)
+            p = map.(x -> time_varying_as_func(value(x), sys), ps)
         else
             u = dvs
             p = ps
@@ -346,8 +400,8 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
         integ = gensym(:MTKIntegrator)
         getexpr = (postprocess_affect_expr! === nothing) ? expression : Val{true}
         pre = get_preprocess_constants(rhss)
-        rf_oop, rf_ip = build_function(rhss, u, p, t; expression = getexpr,
-            wrap_code = add_integrator_header(integ, outvar),
+        rf_oop, rf_ip = build_function(rhss, u, p..., t; expression = getexpr,
+            wrap_code = add_integrator_header(sys, integ, outvar),
             outputidxs = update_inds,
             postprocess_fbody = pre,
             kwargs...)
@@ -385,10 +439,10 @@ function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknow
     root_eq_vars = unique(collect(Iterators.flatten(map(ModelingToolkit.vars, rhss))))
 
     u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = map(x -> time_varying_as_func(value(x), sys), ps)
+    p = map.(x -> time_varying_as_func(value(x), sys), reorder_parameters(sys, ps))
     t = get_iv(sys)
     pre = get_preprocess_constants(rhss)
-    rf_oop, rf_ip = build_function(rhss, u, p, t; expression = Val{false},
+    rf_oop, rf_ip = build_function(rhss, u, p..., t; expression = Val{false},
         postprocess_fbody = pre, kwargs...)
 
     affect_functions = map(cbs) do cb # Keep affect function separate
@@ -400,16 +454,16 @@ function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknow
         cond = function (u, t, integ)
             if DiffEqBase.isinplace(integ.sol.prob)
                 tmp, = DiffEqBase.get_tmp_cache(integ)
-                rf_ip(tmp, u, integ.p, t)
+                rf_ip(tmp, u, parameter_values(integ)..., t)
                 tmp[1]
             else
-                rf_oop(u, integ.p, t)
+                rf_oop(u, parameter_values(integ)..., t)
             end
         end
         ContinuousCallback(cond, affect_functions[])
     else
         cond = function (out, u, t, integ)
-            rf_ip(out, u, integ.p, t)
+            rf_ip(out, u, parameter_values(integ)..., t)
         end
 
         # since there may be different number of conditions and affects,
@@ -432,9 +486,13 @@ function compile_user_affect(affect::FunctionalAffect, sys, dvs, ps; kwargs...)
     dvs_ind = Dict(reverse(en) for en in enumerate(dvs))
     v_inds = map(sym -> dvs_ind[sym], unknowns(affect))
 
-    ps_ind = Dict(reverse(en) for en in enumerate(ps))
-    p_inds = map(sym -> ps_ind[sym], parameters(affect))
-
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p_inds = [parameter_index(sys, sym) for sym in parameters(affect)]
+    else
+        ps_ind = Dict(reverse(en) for en in enumerate(ps))
+        p_inds = map(sym -> ps_ind[sym], parameters(affect))
+    end
+    
     # HACK: filter out eliminated symbols. Not clear this is the right thing to do
     # (MTK should keep these symbols)
     u = filter(x -> !isnothing(x[2]), collect(zip(unknowns_syms(affect), v_inds))) |>