feat!: add MTKParameters to more systems/functions

AayushSabharwal · AayushSabharwal · commit 4ac6301826f9 · 2024-02-05T17:25:48.000+05:30
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -84,21 +84,19 @@ function generate_tgrad(sys::AbstractODESystem, dvs = unknowns(sys), ps = parame
         simplify = false, kwargs...)
     tgrad = calculate_tgrad(sys, simplify = simplify)
     pre = get_preprocess_constants(tgrad)
-    if ps isa Tuple
-        return build_function(tgrad,
-            dvs,
-            ps...,
-            get_iv(sys);
-            postprocess_fbody = pre,
-            kwargs...)
+    p = if has_index_cache(sys)
+        reorder_parameters(get_index_cache(sys), ps)
+    elseif ps isa Tuple
+        ps
     else
-        return build_function(tgrad,
-            dvs,
-            ps,
-            get_iv(sys);
-            postprocess_fbody = pre,
-            kwargs...)
+        (ps,)
     end
+    return build_function(tgrad,
+        dvs,
+        p...,
+        get_iv(sys);
+        postprocess_fbody = pre,
+        kwargs...)
 end
 
 function generate_jacobian(sys::AbstractODESystem, dvs = unknowns(sys), ps = parameters(sys);
@@ -135,7 +133,12 @@ function generate_dae_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
     @variables ˍ₋gamma
     jac = ˍ₋gamma * jac_du + jac_u
     pre = get_preprocess_constants(jac)
-    return build_function(jac, derivatives, dvs, ps, ˍ₋gamma, get_iv(sys);
+    p = if has_index_cache(sys)
+        reorder_parameters(get_index_cache(sys), ps)
+    else
+        (ps,)
+    end
+    return build_function(jac, derivatives, dvs, p..., ˍ₋gamma, get_iv(sys);
         postprocess_fbody = pre, kwargs...)
 end
 
@@ -399,15 +402,25 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = u
                     build_explicit_observed_function(sys, obsvar)
                 end
                 if args === ()
-                    let obs = obs
-                        (u, p, t = Inf) -> if ps isa Tuple
+                    let obs = obs, ps_T = typeof(ps)
+                        (u, p, t = Inf) -> if p isa MTKParameters
+                            obs(u, raw_vectors(p)..., t)
+                        elseif ps_T <: Tuple
                             obs(u, p..., t)
                         else
                             obs(u, p, t)
                         end
                     end
                 else
-                    if ps isa Tuple
+                    if args[2] isa MTKParameters
+                        if length(args) == 2
+                            u, p = args
+                            obs(u, raw_vectors(p)..., Inf)
+                        else
+                            u, p, t = args
+                            obs(u, raw_vectors(p)..., t)
+                        end
+                    elseif ps isa Tuple
                         if length(args) == 2
                             u, p = args
                             obs(u, p..., Inf)
@@ -437,16 +450,21 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = u
                         ps)
                 end
                 if args === ()
-                    let obs = obs
-                        (u, p, t) -> if ps isa Tuple
+                    let obs = obs, ps_T = typeof(ps)
+                        (u, p, t) -> if p isa MTKParameters
+                            obs(u, raw_vectors(p)..., t)
+                        elseif ps_T <: Tuple
                             obs(u, p..., t)
                         else
                             obs(u, p, t)
                         end
                     end
                 else
-                    if ps isa Tuple # split parameters
+                    u, p, t = args
+                    if p isa MTKParameters
                         u, p, t = args
+                        obs(u, raw_vectors(p)..., t)
+                    elseif ps isa Tuple # split parameters
                         obs(u, p..., t)
                     else
                         obs(args...)
@@ -518,7 +536,9 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
                    (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
                    f_gen
     f(du, u, p, t) = f_oop(du, u, p, t)
+    f(du, u, p::MTKParameters, t) = f_oop(du, u, raw_vectors(p)..., t)
     f(out, du, u, p, t) = f_iip(out, du, u, p, t)
+    f(out, du, u, p::MTKParameters, t) = f_iip(out, du, u, raw_vectors(p)..., t)
 
     if jac
         jac_gen = generate_dae_jacobian(sys, dvs, ps;
@@ -530,8 +550,10 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
                            (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in jac_gen) :
                            jac_gen
         _jac(du, u, p, ˍ₋gamma, t) = jac_oop(du, u, p, ˍ₋gamma, t)
+        _jac(du, u, p::MTKParameters, ˍ₋gamma, t) = jac_oop(du, u, raw_vectors(p)..., ˍ₋gamma, t)
 
         _jac(J, du, u, p, ˍ₋gamma, t) = jac_iip(J, du, u, p, ˍ₋gamma, t)
+        _jac(J, du, u, p::MTKParameters, ˍ₋gamma, t) = jac_iip(J, du, u, raw_vectors(p)..., ˍ₋gamma, t)
     else
         _jac = nothing
     end
@@ -544,10 +566,13 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
             end
             if args === ()
                 let obs = obs
-                    (u, p, t) -> obs(u, p, t)
+                    fun(u, p, t) = obs(u, p, t)
+                    fun(u, p::MTKParameters, t) = obs(u, raw_vectors(p)..., t)
+                    fun
                 end
             else
-                obs(args...)
+                u, p, t = args
+                p isa MTKParameters ? obs(u, raw_vectors(p)..., t) : obs(u, p, t)
             end
         end
     end
@@ -591,7 +616,9 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         kwargs...)
     f_oop, f_iip = (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen)
     f(u, h, p, t) = f_oop(u, h, p, t)
+    f(u, h, p::MTKParameters, t) = f_oop(u, h, raw_vectors(p)..., t)
     f(du, u, h, p, t) = f_iip(du, u, h, p, t)
+    f(du, u, h, p::MTKParameters, t) = f_iip(du, u, h, raw_vectors(p)..., t)
 
     DDEFunction{iip}(f, sys = sys)
 end
@@ -617,9 +644,13 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
         isdde = true, kwargs...)
     g_oop, g_iip = (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in g_gen)
     f(u, h, p, t) = f_oop(u, h, p, t)
+    f(u, h, p::MTKParameters, t) = f_oop(u, h, raw_vectors(p)..., t)
     f(du, u, h, p, t) = f_iip(du, u, h, p, t)
+    f(du, u, h, p::MTKParameters, t) = f_iip(du, u, h, raw_vectors(p)..., t)
     g(u, h, p, t) = g_oop(u, h, p, t)
+    g(u, h, p::MTKParameters, t) = g_oop(u, h, raw_vectors(p)..., t)
     g(du, u, h, p, t) = g_iip(du, u, h, p, t)
+    g(du, u, h, p::MTKParameters, t) = g_iip(du, u, h, raw_vectors(p)..., t)
 
     SDDEFunction{iip}(f, g, sys = sys)
 end
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -401,7 +401,9 @@ function build_explicit_observed_function(sys, ts;
     if inputs !== nothing
         ps = setdiff(ps, inputs) # Inputs have been converted to parameters by io_preprocessing, remove those from the parameter list
     end
-    if ps isa Tuple
+    if has_index_cache(sys)
+        ps = DestructuredArgs.(reorder_parameters(get_index_cache(sys), ps))
+    elseif ps isa Tuple
         ps = DestructuredArgs.(ps, inbounds = !checkbounds)
     else
         ps = (DestructuredArgs(ps, inbounds = !checkbounds),)
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
@@ -116,6 +116,10 @@ struct SDESystem <: AbstractODESystem
     """
     complete::Bool
     """
+    Cached data for fast symbolic indexing.
+    """
+    index_cache::Union{Nothing, IndexCache}
+    """
     The hierarchical parent system before simplification.
     """
     parent::Any
@@ -125,7 +129,7 @@ struct SDESystem <: AbstractODESystem
             jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
             cevents, devents, metadata = nothing, gui_metadata = nothing,
-            complete = false, parent = nothing;
+            complete = false, index_cache = nothing, parent = nothing;
             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
@@ -140,7 +144,7 @@ struct SDESystem <: AbstractODESystem
         new(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac,
             Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents,
-            metadata, gui_metadata, complete, parent)
+            metadata, gui_metadata, complete, index_cache, parent)
     end
 end
 
@@ -221,11 +225,15 @@ function generate_diffusion_function(sys::SDESystem, dvs = unknowns(sys),
         eqs = delay_to_function(sys, eqs)
     end
     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
     if isdde
-        return build_function(eqs, u, DDE_HISTORY_FUN, p, get_iv(sys); kwargs...)
+        return build_function(eqs, u, DDE_HISTORY_FUN, p..., get_iv(sys); kwargs...)
     else
-        return build_function(eqs, u, p, get_iv(sys); kwargs...)
+        return build_function(eqs, u, p..., get_iv(sys); kwargs...)
     end
 end
 
@@ -408,9 +416,13 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
                    (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in g_gen) : g_gen
 
     f(u, p, t) = f_oop(u, p, t)
+    f(u, p::MTKParameters, t) = f_oop(u, raw_vectors(p)..., t)
     f(du, u, p, t) = f_iip(du, u, p, t)
+    f(du, u, p::MTKParameters, t) = f_iip(du, u, raw_vectors(p)..., t)
     g(u, p, t) = g_oop(u, p, t)
+    g(u, p::MTKParameters, t) = g_oop(u, raw_vectors(p)..., t)
     g(du, u, p, t) = g_iip(du, u, p, t)
+    g(du, u, p::MTKParameters, t) = g_iip(du, u, raw_vectors(p)..., t)
 
     if tgrad
         tgrad_gen = generate_tgrad(sys, dvs, ps; expression = Val{eval_expression},
@@ -419,7 +431,9 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
                                (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in tgrad_gen) :
                                tgrad_gen
         _tgrad(u, p, t) = tgrad_oop(u, p, t)
+        _tgrad(u, p::MTKParameters, t) = tgrad_oop(u, raw_vectors(p)..., t)
         _tgrad(J, u, p, t) = tgrad_iip(J, u, p, t)
+        _tgrad(J, u, p::MTKParameters, t) = tgrad_iip(J, u, raw_vectors(p)..., t)
     else
         _tgrad = nothing
     end
@@ -431,7 +445,9 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
                            (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in jac_gen) :
                            jac_gen
         _jac(u, p, t) = jac_oop(u, p, t)
+        _jac(u, p::MTKParameters, t) = jac_oop(u, raw_vectors(p)..., t)
         _jac(J, u, p, t) = jac_iip(J, u, p, t)
+        _jac(J, u, p::MTKParameters, t) = jac_iip(J, u, raw_vectors(p)..., t)
     else
         _jac = nothing
     end
@@ -446,9 +462,13 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
                                    (drop_expr(@RuntimeGeneratedFunction(ex)) for ex in tmp_Wfact_t) :
                                    tmp_Wfact_t
         _Wfact(u, p, dtgamma, t) = Wfact_oop(u, p, dtgamma, t)
+        _Wfact(u, p::MTKParameters, dtgamma, t) = Wfact_oop(u, raw_vectors(p)..., dtgamma, t)
         _Wfact(W, u, p, dtgamma, t) = Wfact_iip(W, u, p, dtgamma, t)
+        _Wfact(W, u, p::MTKParameters, dtgamma, t) = Wfact_iip(W, u, raw_vectors(p)..., dtgamma, t)
         _Wfact_t(u, p, dtgamma, t) = Wfact_oop_t(u, p, dtgamma, t)
+        _Wfact_t(u, p::MTKParameters, dtgamma, t) = Wfact_oop_t(u, raw_vectors(p)..., dtgamma, t)
         _Wfact_t(W, u, p, dtgamma, t) = Wfact_iip_t(W, u, p, dtgamma, t)
+        _Wfact_t(W, u, p::MTKParameters, dtgamma, t) = Wfact_iip_t(W, u, raw_vectors(p)..., dtgamma, t)
     else
         _Wfact, _Wfact_t = nothing, nothing
     end
@@ -462,11 +482,14 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = unknowns(sys),
             obs = get!(dict, value(obsvar)) do
                 build_explicit_observed_function(sys, obsvar; checkbounds = checkbounds)
             end
-            obs(u, p, t)
+            if p isa MTKParameters
+                obs(u, raw_vectors(p)..., t)
+            else
+                obs(u, p, t)
+            end
         end
     end
 
-    sts = unknowns(sys)
     SDEFunction{iip}(f, g,
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
@@ -101,11 +101,15 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
     If a model `sys` is complete, then `sys.x` no longer performs namespacing.
     """
     complete::Bool
+    """
+    Cached data for fast symbolic indexing.
+    """
+    index_cache::Union{Nothing, IndexCache}
 
     function JumpSystem{U}(tag, ap::U, iv, unknowns, ps, var_to_name, observed, name, systems,
             defaults, connector_type, devents,
             metadata = nothing, gui_metadata = nothing,
-            complete = false;
+            complete = false, index_cache = nothing;
             checks::Union{Bool, Int} = true) where {U <: ArrayPartition}
         if checks == true || (checks & CheckComponents) > 0
             check_variables(unknowns, iv)
@@ -116,7 +120,7 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
             check_units(u, ap, iv)
         end
         new{U}(tag, ap, iv, unknowns, ps, var_to_name, observed, name, systems, defaults,
-            connector_type, devents, metadata, gui_metadata, complete)
+            connector_type, devents, metadata, gui_metadata, complete, index_cache)
     end
 end
 JumpSystem(tag, ap, iv, states, ps, var_to_name, args...; kwargs...) = JumpSystem{typeof(ap)}(tag, ap, iv, states, ps, var_to_name, args...; kwargs...)
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
