feat!: do not scalarize parameters, fix some tests

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -1397,7 +1397,6 @@ macro mtkbuild(exprs...)
     else
         (;)
     end
-        @show kwargs
     esc(quote
         $named_expr
         $name = $structural_simplify($name; $(kwargs)...)
@@ -1518,13 +1517,12 @@ function linearization_function(sys::AbstractSystem, inputs,
     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)
+    ps = parameters(sys)
     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

src/systems/callbacks.jl

@@ -390,11 +390,8 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
                 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
+                        pind = parameter_index(sys, sym)
+                        discrete_linear_index(ic, pind)
                     end
                 else
                     psind = Dict(reverse(en) for en in enumerate(ps))

src/systems/clock_inference.jl

@@ -146,6 +146,7 @@ function generate_discrete_affect(
     @static if VERSION < v"1.7"
         error("The `generate_discrete_affect` function requires at least Julia 1.7")
     end
+    has_index_cache(osys) && get_index_cache(osys) !== nothing || error("System must have index_cache for clock support")
     out = Sym{Any}(:out)
     appended_parameters = parameters(syss[continuous_id])
     param_to_idx = Dict{Any, ParameterIndex}(p => parameter_index(osys, p)
@@ -169,14 +170,14 @@ function generate_discrete_affect(
             push!(fullvars, s)
         end
         needed_disc_to_cont_obs = []
-        disc_to_cont_idxs = Int[]
+        disc_to_cont_idxs = ParameterIndex[]
         for v in inputs[continuous_id]
             vv = arguments(v)[1]
             if vv in fullvars
                 push!(needed_disc_to_cont_obs, vv)
                 # @show param_to_idx[v] v
                 # @assert param_to_idx[v].portion isa SciMLStructures.Discrete # TOOD: remove
-                push!(disc_to_cont_idxs, param_to_idx[v].idx)
+                push!(disc_to_cont_idxs, param_to_idx[v])
             end
         end
         append!(appended_parameters, input, unknowns(sys))
@@ -201,39 +202,36 @@ function generate_discrete_affect(
             ],
             [],
             let_block)
-        cont_to_disc_idxs = [parameter_index(osys, sym).idx for sym in input]
-        disc_range = [parameter_index(osys, sym).idx for sym in unknowns(sys)]
+        cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
+        disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
         save_vec = Expr(:ref, :Float64)
         for unk in unknowns(sys)
-            idx = parameter_index(osys, unk).idx
-            push!(save_vec.args, :(discretes[$idx]))
+            idx = parameter_index(osys, unk)
+            push!(save_vec.args, :($(parameter_values)(p, $idx)))
         end
         empty_disc = isempty(disc_range)
         disc_init = :(function (p, t)
             d2c_obs = $disc_to_cont_obs
+            disc_state = Tuple($(parameter_values)(p, i) for i in $disc_range)
+            result = d2c_obs(disc_state, p..., t)
+            for (val, i) in zip(result, $disc_to_cont_idxs)
+                # prevent multiple updates to dependents
+                _set_parameter_unchecked!(p, val, i; update_dependent = false)
+            end
             discretes, repack, _ = $(SciMLStructures.canonicalize)(
                 $(SciMLStructures.Discrete()), p)
-            d2c_view = view(discretes, $disc_to_cont_idxs)
-            disc_state = view(discretes, $disc_range)
-            copyto!(d2c_view, d2c_obs(disc_state, p..., t))
-            repack(discretes)
+            repack(discretes) # to force recalculation of dependents
         end)
 
         # @show disc_to_cont_idxs
         # @show cont_to_disc_idxs
         # @show disc_range
-
         affect! = :(function (integrator, saved_values)
             @unpack u, p, t = integrator
             c2d_obs = $cont_to_disc_obs
             d2c_obs = $disc_to_cont_obs
-            # Like Sample
-            discretes, repack, _ = $(SciMLStructures.canonicalize)(
-                $(SciMLStructures.Discrete()), p)
-            c2d_view = view(discretes, $cont_to_disc_idxs)
-            # Like Hold
-            d2c_view = view(discretes, $disc_to_cont_idxs)
-            disc_unknowns = view(discretes, $disc_range)
+            # TODO: find a way to do this without allocating
+            disc_unknowns = [$(parameter_values)(p, i) for i in $disc_range]
             disc = $disc
 
             push!(saved_values.t, t)
@@ -248,12 +246,25 @@ function generate_discrete_affect(
             # d2c comes last
             # @show t
             # @show "incoming", p
-            copyto!(c2d_view, c2d_obs(integrator.u, p..., t))
+            result = c2d_obs(integrator.u, p..., t)
+            for (val, i) in zip(result, $cont_to_disc_idxs)
+                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+            end
             # @show "after c2d", p
-            $empty_disc || disc(disc_unknowns, integrator.u, p..., t)
+            if !$empty_disc
+                disc(disc_unknowns, integrator.u, p..., t)
+                for (val, i) in zip(disc_unknowns, $disc_range)
+                    $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+                end
+            end
             # @show "after state update", p
-            copyto!(d2c_view, d2c_obs(disc_unknowns, p..., t))
+            result = d2c_obs(disc_unknowns, p..., t)
+            for (val, i) in zip(result, $disc_to_cont_idxs)
+                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+            end
             # @show "after d2c", p
+            discretes, repack, _ = $(SciMLStructures.canonicalize)(
+                $(SciMLStructures.Discrete()), p)
             repack(discretes)
         end)
         sv = SavedValues(Float64, Vector{Float64})

src/systems/diffeqs/abstractodesystem.jl

@@ -167,11 +167,7 @@ function generate_function(sys::AbstractODESystem, dvs = unknowns(sys),
 
     # TODO: add an optional check on the ordering of observed equations
     u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        reorder_parameters(get_index_cache(sys), ps isa Tuple ? reduce(vcat, ps) : ps)
-    else
-        (map(x -> time_varying_as_func(value(x), sys), ps),)
-    end
+    p = map.(x -> time_varying_as_func(value(x), sys), reorder_parameters(sys, ps))
     t = get_iv(sys)
 
     if isdde
@@ -775,6 +771,20 @@ function get_u0_p(sys,
     u0, p, defs
 end
 
+function get_u0(sys, u0map; symbolic_u0 = false)
+    dvs = unknowns(sys)
+
+    defs = defaults(sys)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    if symbolic_u0
+        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false, use_union = false)
+    else
+        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
+    end
+    return u0, defs
+end
+
 function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         implicit_dae = false, du0map = nothing,
         version = nothing, tgrad = false,
@@ -793,20 +803,24 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    u0, _p, defs = get_u0_p(sys,
-        u0map,
-        parammap;
-        tofloat,
-        use_union,
-        symbolic_u0)
-    if u0 !== nothing
-        u0 = u0_constructor(u0)
-    end
-
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        u0, defs = get_u0(sys, u0map; symbolic_u0)
         p = MTKParameters(sys, parammap)
     else
-        p = _p
+        u0, p, defs = get_u0_p(sys,
+            u0map,
+            parammap;
+            tofloat,
+            use_union,
+            symbolic_u0)
+        p, split_idxs = split_parameters_by_type(p)
+        if p isa Tuple
+            ps = Base.Fix1(getindex, parameters(sys)).(split_idxs)
+            ps = (ps...,) #if p is Tuple, ps should be Tuple
+        end
+    end
+    if u0 !== nothing
+        u0 = u0_constructor(u0)
     end
 
     if implicit_dae && du0map !== nothing

src/systems/diffeqs/odesystem.jl

@@ -197,21 +197,18 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = scalarize(deqs)
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
-
     iv′ = value(scalarize(iv))
-    ps′ = value.(scalarize(ps))
+    ps′ = ps
     ctrl′ = value.(scalarize(controls))
     dvs′ = value.(scalarize(dvs))
     dvs′ = filter(x -> !isdelay(x, iv), dvs′)
-
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn(
             "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
             :ODESystem, force = true)
     end
     defaults = todict(defaults)
     defaults = Dict{Any, Any}(value(k) => value(v) for (k, v) in pairs(defaults))
-
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
@@ -279,7 +276,7 @@ function ODESystem(eqs, iv; kwargs...)
     algevars = setdiff(allunknowns, diffvars)
     # the orders here are very important!
     return ODESystem(Equation[diffeq; algeeq; compressed_eqs], iv,
-        collect(Iterators.flatten((diffvars, algevars))), ps; kwargs...)
+        collect(Iterators.flatten((diffvars, algevars))), collect(ps); kwargs...)
 end
 
 # NOTE: equality does not check cached Jacobian

src/systems/diffeqs/sdesystem.jl

@@ -407,7 +407,6 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
         error("A completed `SDESystem` is required. Call `complete` or `structural_simplify` on the system before creating an `SDEFunction`")
     end
     dvs = scalarize.(dvs)
-    ps = scalarize.(ps)
 
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression}, kwargs...)
     f_oop, f_iip = eval_expression ?