feat: propagate p_constructor to InitializationProblem

Author: AayushSabharwal

src/systems/problem_utils.jl

@@ -869,15 +869,17 @@ All other keyword arguments are forwarded to `InitializationProblem`.
 function maybe_build_initialization_problem(
         sys::AbstractSystem, op::AbstractDict, u0map, pmap, t, defs,
         guesses, missing_unknowns; implicit_dae = false, u0_constructor = identity,
-        floatT = Float64, initialization_eqs = [], use_scc = true, kwargs...)
+        p_constructor = identity, floatT = Float64, initialization_eqs = [],
+        use_scc = true, kwargs...)
     guesses = merge(ModelingToolkit.guesses(sys), todict(guesses))
 
     if t === nothing && is_time_dependent(sys)
         t = zero(floatT)
     end
 
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
-        sys, t, u0map, pmap; guesses, initialization_eqs, use_scc, kwargs...)
+        sys, t, u0map, pmap; guesses, initialization_eqs,
+        use_scc, u0_constructor, p_constructor, kwargs...)
     if state_values(initializeprob) !== nothing
         initializeprob = remake(initializeprob; u0 = floatT.(state_values(initializeprob)))
     end
@@ -1109,7 +1111,7 @@ function process_SciMLProblem(
             warn_cyclic_dependency, check_units = check_initialization_units,
             circular_dependency_max_cycle_length, circular_dependency_max_cycles, use_scc,
             force_time_independent = force_initialization_time_independent, algebraic_only, allow_incomplete,
-            u0_constructor, floatT)
+            u0_constructor, p_constructor, floatT)
 
         kwargs = merge(kwargs, kws)
     end

test/initial_values.jl

@@ -3,7 +3,7 @@ using ModelingToolkit: t_nounits as t, D_nounits as D, get_u0
 using OrdinaryDiffEq
 using DataInterpolations
 using StaticArrays
-using SymbolicIndexingInterface: getu
+using SymbolicIndexingInterface
 
 @variables x(t)[1:3]=[1.0, 2.0, 3.0] y(t) z(t)[1:2]
 
@@ -337,29 +337,28 @@ end
 end
 
 @testset "`p_constructor` keyword argument" begin
-    @parameters σ ρ β
-    @variables x(t) y(t) z(t)
-
-    eqs = [D(D(x)) ~ σ * (y - x),
-        D(y) ~ x * (ρ - z) - y,
-        D(z) ~ x * y - β * z]
-
-    @mtkbuild sys = ODESystem(eqs, t)
-
-    u0 = [D(x) => 2.0f0,
-        x => 1.0f0,
-        y => 0.0f0,
-        z => 0.0f0]
+    @parameters g = 1.0
+    @variables x(t) y(t) [state_priority = 10, guess = 1.0] λ(t) [guess = 1.0]
+    eqs = [D(D(x)) ~ λ * x
+           D(D(y)) ~ λ * y - g
+           x^2 + y^2 ~ 1]
+    @mtkbuild pend = ODESystem(eqs, t)
 
-    p = [σ => 28.0f0,
-        ρ => 10.0f0,
-        β => 8.0f0 / 3]
+    u0 = [x => 1.0, D(x) => 0.0]
     u0_constructor = p_constructor = vals -> SVector{length(vals)}(vals...)
-    prob = ODEProblem(sys, u0, tspan, p; u0_constructor, p_constructor)
+    tspan = (0.0, 5.0)
+    prob = ODEProblem(pend, u0, tspan; u0_constructor, p_constructor)
+    @test prob.u0 isa SVector
     @test prob.p.tunable isa SVector
     @test prob.p.initials isa SVector
+    initdata = prob.f.initialization_data
+    @test state_values(initdata.initializeprob) isa SVector
+    @test parameter_values(initdata.initializeprob).tunable isa SVector
 
-    @mtkbuild sys=ODESystem(eqs, t) split=false
-    prob = ODEProblem(sys, u0, tspan, p; u0_constructor, p_constructor)
+    @mtkbuild pend=ODESystem(eqs, t) split=false
+    prob = ODEProblem(pend, u0, tspan; u0_constructor, p_constructor)
     @test prob.p isa SVector
+    initdata = prob.f.initialization_data
+    @test state_values(initdata.initializeprob) isa SVector
+    @test parameter_values(initdata.initializeprob) isa SVector
 end