feat!: add MTKParameters struct, use as ODEProblem.p

Author: AayushSabharwal

Project.toml

@@ -39,6 +39,7 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
@@ -99,6 +100,7 @@ SciMLBase = "2.0.1"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0, 1"
+SciMLStructures = "1.0"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"

src/ModelingToolkit.jl

@@ -31,6 +31,7 @@ using PrecompileTools, Reexport
     import Distributions
     import FunctionWrappersWrappers
     using URIs: URI
+    using SciMLStructures
 
     using RecursiveArrayTools
 
@@ -132,6 +133,8 @@ include("systems/abstractsystem.jl")
 include("systems/model_parsing.jl")
 include("systems/connectors.jl")
 include("systems/callbacks.jl")
+include("systems/index_cache.jl")
+include("systems/parameter_buffer.jl")
 
 include("systems/diffeqs/odesystem.jl")
 include("systems/diffeqs/sdesystem.jl")

src/clock.jl

@@ -83,7 +83,9 @@ true if `x` contains only discrete-domain signals.
 See also [`has_discrete_domain`](@ref)
 """
 function is_discrete_domain(x)
-    issym(x) && return getmetadata(x, TimeDomain, false) isa Discrete
+    if hasmetadata(x, TimeDomain) || issym(x)
+        return getmetadata(x, TimeDomain, false) isa AbstractDiscrete
+    end
     !has_discrete_domain(x) && has_continuous_domain(x)
 end
 

src/systems/diffeqs/abstractodesystem.jl

@@ -160,27 +160,26 @@ function generate_function(sys::AbstractODESystem, dvs = unknowns(sys), ps = par
 
     # TODO: add an optional check on the ordering of observed equations
     u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = map(x -> time_varying_as_func(value(x), sys), ps)
+    p = if has_index_cache(sys)
+        reorder_parameters(get_index_cache(sys), ps)
+    else
+        (map(x -> time_varying_as_func(value(x), sys), ps),)
+    end
     t = get_iv(sys)
 
     if isdde
-        build_function(rhss, u, DDE_HISTORY_FUN, p, t; kwargs...)
+        build_function(rhss, u, DDE_HISTORY_FUN, p..., t; kwargs...)
     else
         pre, sol_states = get_substitutions_and_solved_unknowns(sys)
 
         if implicit_dae
-            build_function(rhss, ddvs, u, p, t; postprocess_fbody = pre,
+            build_function(rhss, ddvs, u, p..., t; postprocess_fbody = pre,
                 states = sol_states,
                 kwargs...)
         else
-            if p isa Tuple
-                build_function(rhss, u, p..., t; postprocess_fbody = pre,
-                    states = sol_states,
-                    kwargs...)
-            else
-                build_function(rhss, u, p, t; postprocess_fbody = pre, states = sol_states,
-                    kwargs...)
-            end
+            build_function(rhss, u, p..., t; postprocess_fbody = pre,
+                states = sol_states,
+                kwargs...)
         end
     end
 end
@@ -321,6 +320,10 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = u
         g(u, p, t) = f_oop(u, p..., t)
         g(du, u, p, t) = f_iip(du, u, p..., t)
         f = g
+    elseif p isa MTKParameters
+        h(u, p, t) = f_oop(u, raw_vectors(p)..., t)
+        h(du, u, p, t) = f_iip(du, u, raw_vectors(p)..., t)
+        f = h
     else
         k(u, p, t) = f_oop(u, p, t)
         k(du, u, p, t) = f_iip(du, u, p, t)
@@ -758,7 +761,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    u0, p, defs = get_u0_p(sys,
+    u0, _, defs = get_u0_p(sys,
         u0map,
         parammap;
         tofloat,
@@ -768,11 +771,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         u0 = u0_constructor(u0)
     end
 
-    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
+    p = MTKParameters(sys, parammap; toterm = default_toterm)
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
@@ -789,7 +788,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     f = constructor(sys, dvs, ps, u0; ddvs = ddvs, tgrad = tgrad, jac = jac,
         checkbounds = checkbounds, p = p,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,
-        sparse = sparse, eval_expression = eval_expression, split_idxs,
+        sparse = sparse, eval_expression = eval_expression,
         kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end

src/systems/parameter_buffer.jl

@@ -0,0 +1,100 @@
+struct MTKParameters{T, D}
+    tunable::T
+    discrete::D
+end
+
+function MTKParameters(sys::AbstractSystem, p; toterm = default_toterm)
+    ic = if has_index_cache(sys)
+        get_index_cache(sys)
+    else
+        IndexCache(sys)
+    end
+    tunable_buffer = if length(ic.param_buffer_type_and_size) == 0
+        Float64[]
+    elseif length(ic.param_buffer_type_and_size) == 1
+        T, sz = only(ic.param_buffer_type_and_size)
+        Vector{T == Real ? Float64 : T}(undef, sz)
+    else
+        ArrayPartition((Vector{T == Real ? Float64 : T}(undef, sz) for (T, sz) in ic.param_buffer_type_and_size)...)
+    end
+
+    disc_buffer = if length(ic.discrete_buffer_type_and_size) == 0
+        Float64[]
+    elseif length(ic.discrete_buffer_type_and_size) == 1
+        T, sz = only(ic.discrete_buffer_type_and_size)
+        Vector{T == Real ? Float64 : T}(undef, sz)
+    else
+        ArrayPartition((Vector{T == Real ? Float64 : T}(undef, sz) for (T, sz) in ic.discrete_buffer_type_and_size)...)
+    end
+
+    for (sym, value) in defaults(sys)
+        sym = toterm(unwrap(sym))
+        h = hasmetadata(sym, SymbolHash) ? getmetadata(sym, SymbolHash) : hash(sym)
+        if haskey(ic.discrete_idx, h)
+            disc_buffer[ic.discrete_idx[h]] = value
+        elseif haskey(ic.param_idx, h)
+            tunable_buffer[ic.param_idx[h]] = value
+        end
+    end
+
+    if !isa(p, SciMLBase.NullParameters)
+        for (sym, value) in p
+            sym = toterm(unwrap(sym))
+            h = hasmetadata(sym, SymbolHash) ? getmetadata(sym, SymbolHash) : hash(sym)
+            if haskey(ic.discrete_idx, h)
+                disc_buffer[ic.discrete_idx[h]] = value
+            elseif haskey(ic.param_idx, h)
+                tunable_buffer[ic.param_idx[h]] = value
+            else
+                error("Invalid parameter $sym")
+            end
+        end
+    end
+
+    return MTKParameters{typeof(tunable_buffer), typeof(disc_buffer)}(tunable_buffer, disc_buffer)
+end
+
+SciMLStructures.isscimlstructure(::MTKParameters) = true
+
+SciMLStructures.ismutablescimlstructure(::MTKParameters) = true
+
+for (Portion, field) in [
+    (SciMLStructures.Tunable, :tunable)
+    (SciMLStructures.Discrete, :discrete)
+]
+    @eval function SciMLStructures.canonicalize(::$Portion, p::MTKParameters)
+        function repack(values)
+            p.$field .= values
+        end
+        return p.$field, repack, !isa(p.$field, ArrayPartition)
+    end
+
+    @eval function SciMLStructures.replace(::$Portion, p::MTKParameters, newvals)
+        new_field = similar(p.$field)
+        new_field .= newvals
+        @set p.$field = new_field
+    end
+
+    @eval function SciMLStructures.replace!(::$Portion, p::MTKParameters, newvals)
+        p.$field .= newvals
+        nothing
+    end
+end
+
+function raw_vectors(buf::MTKParameters)
+    tunable = if isempty(buf.tunable)
+        ()
+    elseif buf.tunable isa ArrayPartition
+        buf.tunable.x
+    else
+        (buf.tunable,)
+    end
+    discrete = if isempty(buf.discrete)
+        ()
+    elseif buf.discrete isa ArrayPartition
+        buf.discrete.x
+    else
+        (buf.discrete,)
+    end
+    return (tunable..., discrete...)
+end