fix: linearization with MTKParameters

Author: AayushSabharwal

src/ModelingToolkit.jl

@@ -129,12 +129,12 @@ include("constants.jl")
 include("utils.jl")
 include("domains.jl")
 
+include("systems/index_cache.jl")
+include("systems/parameter_buffer.jl")
 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/systems/abstractsystem.jl

@@ -1504,7 +1504,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     x0 = merge(defaults(sys), Dict(missing_variable_defaults(sys)), op)
     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)...,)
+        p = MTKParameters(sys, p)
         ps = reorder_parameters(sys, parameters(sys))
     else
         p = _p
@@ -1515,7 +1515,6 @@ function linearization_function(sys::AbstractSystem, inputs,
             ps = (ps...,) #if p is Tuple, ps should be Tuple
         end
     end
-
     lin_fun = let diff_idxs = diff_idxs,
         alge_idxs = alge_idxs,
         input_idxs = input_idxs,
@@ -1550,7 +1549,9 @@ function linearization_function(sys::AbstractSystem, inputs,
                 h_xz = fg_u = zeros(0, length(inputs))
             end
             hp = let u = u, t = t
-                p -> h(u, p, t)
+                _hp(p) = h(u, p, t)
+                _hp(p::MTKParameters) = h(u, p..., t)
+                _hp
             end
             h_u = jacobian_wrt_vars(hp, p, input_idxs, chunk)
             (f_x = fg_xz[diff_idxs, diff_idxs],
@@ -1592,13 +1593,14 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
         kwargs...)
     sts = unknowns(sys)
     t = get_iv(sys)
-    p = parameters(sys)
+    ps = parameters(sys)
+    p = reorder_parameters(sys, ps)
 
-    fun = generate_function(sys, sts, p; expression = Val{false})[1]
-    dx = fun(sts, p, t)
+    fun = generate_function(sys, sts, ps; expression = Val{false})[1]
+    dx = fun(sts, p..., t)
 
     h = build_explicit_observed_function(sys, outputs)
-    y = h(sts, p, t)
+    y = h(sts, p..., t)
 
     fg_xz = Symbolics.jacobian(dx, sts)
     fg_u = Symbolics.jacobian(dx, inputs)
@@ -1794,6 +1796,15 @@ function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives =
     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
+        if p isa SciMLBase.NullParameters
+            p = op
+        elseif p isa Dict
+            p = merge(p, op)
+        elseif p isa Vector && eltype(p) <: Pair
+            p = merge(Dict(p), op)
+        elseif p isa Vector
+            p = merge(Dict(parameters(sys) .=> p), op)
+        end
         p2 = MTKParameters(sys, p)
     end
     linres = lin_fun(u0, p2, t)

src/systems/parameter_buffer.jl

@@ -1,9 +1,10 @@
-struct MTKParameters{T, D, C, E, F}
+struct MTKParameters{T, D, C, E, F, G}
     tunable::T
     discrete::D
     constant::C
     dependent::E
-    dependent_update::F
+    dependent_update_iip::F
+    dependent_update_oop::G
 end
 
 function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = false)
@@ -19,12 +20,12 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
         length(p) == length(ps) || error("Invalid parameters")
         p = ps .=> p
     end
-    defs = Dict(default_toterm(unwrap(k)) => v for (k, v) in defaults(sys) if unwrap(k) in all_ps)
+    defs = Dict(default_toterm(unwrap(k)) => v for (k, v) in defaults(sys) if unwrap(k) in all_ps || default_toterm(unwrap(k)) in all_ps)
     if p isa SciMLBase.NullParameters
         p = defs
     else
-        extra_params = Dict(unwrap(k) => v for (k, v) in p if !in(unwrap(k), all_ps))
-        p = merge(defs, Dict(default_toterm(unwrap(k)) => v for (k, v) in p if unwrap(k) in all_ps))
+        extra_params = Dict(unwrap(k) => v for (k, v) in p if !in(unwrap(k), all_ps) && !in(default_toterm(unwrap(k)), all_ps))
+        p = merge(defs, Dict(default_toterm(unwrap(k)) => v for (k, v) in p if unwrap(k) in all_ps || default_toterm(unwrap(k)) in all_ps))
         p = Dict(k => fixpoint_sub(v, extra_params) for (k, v) in p if !haskey(extra_params, unwrap(k)))
     end
 
@@ -74,39 +75,41 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
         dep_exprs[idx] = wrap(fixpoint_sub(val, dependencies))
     end
     p = reorder_parameters(ic, parameters(sys))[begin:end-length(dep_buffer.x)]
-    update_function = if isempty(dep_exprs.x)
-        (_...) -> ()
+    update_function_iip, update_function_oop = if isempty(dep_exprs.x)
+        nothing, nothing
     else
-        RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(build_function(dep_exprs, p...)[2])
+        oop, iip = build_function(dep_exprs, p...)
+        RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(iip), RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(oop)
     end
     # everything is an ArrayPartition so it's easy to figure out how many
     # distinct vectors we have for each portion as `ArrayPartition.x`
     if isempty(tunable_buffer.x)
-        tunable_buffer = ArrayPartition(Float64[])
+        tunable_buffer = Float64[]
     end
     if isempty(disc_buffer.x)
-        disc_buffer = ArrayPartition(Float64[])
+        disc_buffer = Float64[]
     end
     if isempty(const_buffer.x)
-        const_buffer = ArrayPartition(Float64[])
+        const_buffer = Float64[]
     end
     if isempty(dep_buffer.x)
-        dep_buffer = ArrayPartition(Float64[])
+        dep_buffer = Float64[]
     end
     if use_union
-        tunable_buffer = ArrayPartition(restrict_array_to_union(tunable_buffer))
-        disc_buffer = ArrayPartition(restrict_array_to_union(disc_buffer))
-        const_buffer = ArrayPartition(restrict_array_to_union(const_buffer))
-        dep_buffer = ArrayPartition(restrict_array_to_union(dep_buffer))
+        tunable_buffer = restrict_array_to_union(tunable_buffer)
+        disc_buffer = restrict_array_to_union(disc_buffer)
+        const_buffer = restrict_array_to_union(const_buffer)
+        dep_buffer = restrict_array_to_union(dep_buffer)
     elseif tofloat
-        tunable_buffer = ArrayPartition(Float64.(tunable_buffer))
-        disc_buffer = ArrayPartition(Float64.(disc_buffer))
-        const_buffer = ArrayPartition(Float64.(const_buffer))
-        dep_buffer = ArrayPartition(Float64.(dep_buffer))
+        tunable_buffer = Float64.(tunable_buffer)
+        disc_buffer = Float64.(disc_buffer)
+        const_buffer = Float64.(const_buffer)
+        dep_buffer = Float64.(dep_buffer)
     end
     return MTKParameters{typeof(tunable_buffer), typeof(disc_buffer), typeof(const_buffer),
-                         typeof(dep_buffer), typeof(update_function)}(tunable_buffer,
-                         disc_buffer, const_buffer, dep_buffer, update_function)
+                         typeof(dep_buffer), typeof(update_function_iip), typeof(update_function_oop)}(
+                         tunable_buffer, disc_buffer, const_buffer, dep_buffer, update_function_iip,
+                         update_function_oop)
 end
 
 SciMLStructures.isscimlstructure(::MTKParameters) = true
@@ -121,19 +124,27 @@ for (Portion, field) in [
     @eval function SciMLStructures.canonicalize(::$Portion, p::MTKParameters)
         function repack(values)
             p.$field .= values
-            p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
+            if p.dependent_update_iip !== nothing
+                p.dependent_update_iip(p.dependent, p...)
+            end
+            p
         end
         return p.$field, repack, true
     end
 
     @eval function SciMLStructures.replace(::$Portion, p::MTKParameters, newvals)
         @set! p.$field = newvals
-        p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
+        if p.dependent_update_oop !== nothing
+            @set! p.dependent = ArrayPartition(p.dependent_update_oop(p...))
+        end
+        p
     end
 
     @eval function SciMLStructures.replace!(::$Portion, p::MTKParameters, newvals)
         p.$field .= newvals
-        p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
+        if p.dependent_update_iip !== nothing
+            p.dependent_update_iip(p.dependent, p...)
+        end
         nothing
     end
 end
@@ -166,26 +177,32 @@ function SymbolicIndexingInterface.set_parameter!(p::MTKParameters, val, idx::Pa
     else
         error("Unhandled portion $portion")
     end
-    p.dependent_update(p.dependent, p.tunable.x..., p.discrete.x..., p.constant.x...)
+    if p.dependent_update_iip !== nothing
+        p.dependent_update_iip(p.dependent, p...)
+    end
 end
 
+_subarrays(v::AbstractVector) = isempty(v) ? () : (v,)
+_subarrays(v::ArrayPartition) = v.x
+_num_subarrays(v::AbstractVector) = 1
+_num_subarrays(v::ArrayPartition) = length(v.x)
 # for compiling callbacks
 # getindex indexes the vectors, setindex! linearly indexes values
 # it's inconsistent, but we need it to be this way
 function Base.getindex(buf::MTKParameters, i)
     if !isempty(buf.tunable)
-        i <= length(buf.tunable.x) && return buf.tunable.x[i]
-        i -= length(buf.tunable.x)
+        i <= _num_subarrays(buf.tunable) && return _subarrays(buf.tunable)[i]
+        i -= _num_subarrays(buf.tunable)
     end
     if !isempty(buf.discrete)
-        i <= length(buf.discrete.x) && return buf.discrete.x[i]
-        i -= length(buf.discrete.x)
+        i <= _num_subarrays(buf.discrete) && return _subarrays(buf.discrete)[i]
+        i -= _num_subarrays(buf.discrete)
     end
     if !isempty(buf.constant)
-        i <= length(buf.constant.x) && return buf.constant.x[i]
-        i -= length(buf.constant.x)
+        i <= _num_subarrays(buf.constant) && return _subarrays(buf.constant)[i]
+        i -= _num_subarrays(buf.constant)
     end
-    isempty(buf.dependent) || return buf.dependent.x[i]
+    isempty(buf.dependent) || return _subarrays(buf.dependent)[i]
     throw(BoundsError(buf, i))
 end
 function Base.setindex!(buf::MTKParameters, val, i)
@@ -196,15 +213,17 @@ function Base.setindex!(buf::MTKParameters, val, i)
     else
         buf.constant[i - length(buf.tunable) - length(buf.discrete)] = val
     end
-    buf.dependent_update(buf.dependent, buf.tunable.x..., buf.discrete.x..., buf.constant.x...)
+    if buf.dependent_update_iip !== nothing
+        buf.dependent_update_iip(buf.dependent, buf...)
+    end
 end
 
 function Base.iterate(buf::MTKParameters, state = 1)
     total_len = 0
-    isempty(buf.tunable) || (total_len += length(buf.tunable.x))
-    isempty(buf.discrete) || (total_len += length(buf.discrete.x))
-    isempty(buf.constant) || (total_len += length(buf.constant.x))
-    isempty(buf.dependent) || (total_len += length(buf.dependent.x))
+    isempty(buf.tunable) || (total_len += _num_subarrays(buf.tunable))
+    isempty(buf.discrete) || (total_len += _num_subarrays(buf.discrete))
+    isempty(buf.constant) || (total_len += _num_subarrays(buf.constant))
+    isempty(buf.dependent) || (total_len += _num_subarrays(buf.dependent))
     if state <= total_len
         return (buf[state], state + 1)
     else
@@ -229,14 +248,24 @@ function jacobian_wrt_vars(pf::F, p::MTKParameters, input_idxs, chunk::C) where
         p_big = p_big
 
         function (p_small_inner)
-            tunable, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p_big)
-            tunable[input_idxs] .= p_small_inner
-            p_big = repack(tunable)
-            pf(p_big)
+            for (i, val) in zip(input_idxs, p_small_inner)
+                set_parameter!(p_big, val, i)
+            end
+            # tunable, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p_big)
+            # tunable[input_idxs] .= p_small_inner
+            # p_big = repack(tunable)
+            return if pf isa SciMLBase.ParamJacobianWrapper
+                buffer = similar(p_big.tunable, size(pf.u))
+                pf(buffer, p_big)
+                buffer
+            else
+                pf(p_big)
+            end
         end
     end
-    tunable, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
-    p_small = tunable[input_idxs]
+    # tunable, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
+    # p_small = tunable[input_idxs]
+    p_small = parameter_values.((p,), input_idxs)
     cfg = ForwardDiff.JacobianConfig(p_closure, p_small, chunk, tag)
     ForwardDiff.jacobian(p_closure, p_small, cfg, Val(false))
 end

src/systems/systems.jl

@@ -28,7 +28,8 @@ function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false
     @set! newsys.parent = complete(sys)
     newsys = complete(newsys)
     if newsys′ isa Tuple
-        return newsys, newsys′[2]
+        idxs = [parameter_index(newsys, i) for i in io[1]]
+        return newsys, idxs
     else
         return newsys
     end

test/inversemodel.jl

@@ -157,7 +157,8 @@ nsys = get_named_sensitivity(model, :y; op) # Test that we get the same result w
 # Test the same thing for comp sensitivities
 
 Sf, simplified_sys = Blocks.get_comp_sensitivity_function(model, :y) # This should work without providing an operating opint containing a dummy derivative
-x, p = ModelingToolkit.get_u0_p(simplified_sys, op)
+x, _ = ModelingToolkit.get_u0_p(simplified_sys, op)
+p = ModelingToolkit.MTKParameters(simplified_sys, op)
 matrices1 = Sf(x, p, 0)
 matrices2, _ = Blocks.get_comp_sensitivity(model, :y; op) # Test that we get the same result when calling the higher-level API
 @test matrices1.f_x ≈ matrices2.A[1:7, 1:7]

test/linearize.jl

@@ -170,7 +170,6 @@ function saturation(; y_max, y_min = y_max > 0 ? -y_max : -Inf, name)
     ]
     ODESystem(eqs, t, name = name)
 end
-
 @named sat = saturation(; y_max = 1)
 # inside the linear region, the function is identity
 @unpack u, y = sat

test/symbolic_events.jl

@@ -498,7 +498,6 @@ let rng = rng
     function testsol(jsys, u0, p, tspan; tstops = Float64[], paramtotest = nothing,
             N = 40000, kwargs...)
         jsys = complete(jsys)
-        @show ModelingToolkit.get_index_cache(jsys)
         dprob = DiscreteProblem(jsys, u0, tspan, p)
         jprob = JumpProblem(jsys, dprob, Direct(); kwargs...)
         sol = solve(jprob, SSAStepper(); tstops = tstops)