Merge pull request #3135 from AayushSabharwal/as/initialization-nonnum-pdeps

fix: retain nonnumeric parameter dependencies in initialization system

Author: ChrisRackauckas

src/systems/nonlinear/initializesystem.jl

@@ -138,8 +138,13 @@ function generate_initializesystem(sys::ODESystem;
     end
 
     # 5) parameter dependencies become equations, their LHS become unknowns
+    # non-numeric dependent parameters stay as parameter dependencies
+    new_parameter_deps = Equation[]
     for eq in parameter_dependencies(sys)
-        is_variable_floatingpoint(eq.lhs) || continue
+        if !is_variable_floatingpoint(eq.lhs)
+            push!(new_parameter_deps, eq)
+            continue
+        end
         varp = tovar(eq.lhs)
         paramsubs[eq.lhs] = varp
         push!(eqs_ics, eq)
@@ -171,6 +176,7 @@ function generate_initializesystem(sys::ODESystem;
         pars;
         defaults = defs,
         checks = check_units,
+        parameter_dependencies = new_parameter_deps,
         name,
         metadata = meta,
         kwargs...)

test/downstream/linearize.jl

@@ -95,10 +95,10 @@ lsys = ModelingToolkit.reorder_unknowns(lsys0, unknowns(ssys), desired_order)
 ## Symbolic linearization
 lsyss, _ = ModelingToolkit.linearize_symbolic(cl, [f.u], [p.x])
 
-@test substitute(lsyss.A, ModelingToolkit.defaults(cl)) == lsys.A
-@test substitute(lsyss.B, ModelingToolkit.defaults(cl)) == lsys.B
-@test substitute(lsyss.C, ModelingToolkit.defaults(cl)) == lsys.C
-@test substitute(lsyss.D, ModelingToolkit.defaults(cl)) == lsys.D
+@test ModelingToolkit.fixpoint_sub(lsyss.A, ModelingToolkit.defaults(cl)) == lsys.A
+@test ModelingToolkit.fixpoint_sub(lsyss.B, ModelingToolkit.defaults(cl)) == lsys.B
+@test ModelingToolkit.fixpoint_sub(lsyss.C, ModelingToolkit.defaults(cl)) == lsys.C
+@test ModelingToolkit.fixpoint_sub(lsyss.D, ModelingToolkit.defaults(cl)) == lsys.D
 ##
 using ModelingToolkitStandardLibrary.Blocks: LimPID
 k = 400

test/initializationsystem.jl

@@ -815,3 +815,24 @@ end
     prob2 = @test_nowarn remake(prob; u0 = [y => 0.5])
     @test is_variable(prob.f.initializeprob, p)
 end
+
+struct Multiplier{T}
+    a::T
+    b::T
+end
+
+function (m::Multiplier)(x, y)
+    m.a * x + m.b * y
+end
+
+@register_symbolic Multiplier(x::Real, y::Real)
+
+@testset "Nonnumeric parameter dependencies are retained" begin
+    @variables x(t) y(t)
+    @parameters foo(::Real, ::Real) p
+    @mtkbuild sys = ODESystem([D(x) ~ t, 0 ~ foo(x, y)], t;
+        parameter_dependencies = [foo ~ Multiplier(p, 2p)], guesses = [y => -1.0])
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
+    integ = init(prob, Rosenbrock23())
+    @test integ[y] ≈ -0.5
+end