Merge pull request #1941 from ValentinKaisermayer/vk-implement-process_p_u0_symbolic

Implements process_p_u0_symbolic from SciMLBase

Author: ChrisRackauckas

Project.toml

@@ -72,7 +72,7 @@ NonlinearSolve = "0.3.8"
 RecursiveArrayTools = "2.3"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.4.3, 0.5"
-SciMLBase = "1.70.0"
+SciMLBase = "1.72.0"
 Setfield = "0.7, 0.8, 1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"

src/variables.jl

@@ -119,15 +119,42 @@ end
     throw(ArgumentError("$vars are missing from the variable map."))
 end
 
+# FIXME: remove after: https://github.com/SciML/SciMLBase.jl/pull/311
+function SciMLBase.handle_varmap(varmap, sys::AbstractSystem; field = :states, kwargs...)
+    out = varmap_to_vars(varmap, getfield(sys, field); kwargs...)
+    return out
+end
+
 """
 $(SIGNATURES)
 
-Intercept the call to `handle_varmap` and convert it to an ordered list if the user has
-  ModelingToolkit loaded, and the problem has a symbolic origin.
+Intercept the call to `process_p_u0_symbolic` and process symbolic maps of `p` and/or `u0` if the 
+user has `ModelingToolkit` loaded.
 """
-function SciMLBase.handle_varmap(varmap, sys::AbstractSystem; field = :states, kwargs...)
-    out = varmap_to_vars(varmap, getfield(sys, field); kwargs...)
-    return out
+function SciMLBase.process_p_u0_symbolic(prob::ODEProblem, p, u0)
+    # check if a symbolic remake is possible
+    if eltype(p) <: Pair
+        hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
+            throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
+                                " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
+    end
+    if eltype(u0) <: Pair
+        hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) ||
+            throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
+                                " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order."))
+    end
+
+    # assemble defaults
+    defs = defaults(prob.f.sys)
+    defs = mergedefaults(defs, prob.p, parameters(prob.f.sys))
+    defs = mergedefaults(defs, p, parameters(prob.f.sys))
+    defs = mergedefaults(defs, prob.u0, states(prob.f.sys))
+    defs = mergedefaults(defs, u0, states(prob.f.sys))
+
+    u0 = varmap_to_vars(u0, states(prob.f.sys); defaults = defs, tofloat = true)
+    p = varmap_to_vars(p, parameters(prob.f.sys); defaults = defs)
+
+    return p, u0
 end
 
 struct IsHistory end

test/odesystem.jl

@@ -257,6 +257,30 @@ sol_dpmap = solve(prob_dpmap, Rodas5())
 
 @test sol_pmap.u ≈ sol_dpmap.u
 
+@testset "symbolic remake with nested system" begin
+    function makesys(name)
+        @parameters t a=1.0
+        @variables x(t) = 0.0
+        D = Differential(t)
+        ODESystem([D(x) ~ -a * x]; name)
+    end
+
+    function makecombinedsys()
+        sys1 = makesys(:sys1)
+        sys2 = makesys(:sys2)
+        @parameters t b=1.0
+        ODESystem(Equation[], t, [], [b]; systems = [sys1, sys2], name = :foo)
+    end
+
+    sys = makecombinedsys()
+    @unpack sys1, b = sys
+    prob = ODEProblem(sys, Pair[])
+    prob_new = SciMLBase.remake(prob, p = Dict(sys1.a => 3.0, b => 4.0),
+                                u0 = Dict(sys1.x => 1.0))
+    @test_broken prob_new.p == [4.0, 3.0, 1.0]
+    @test_broken prob_new.u0 == [1.0, 0.0]
+end
+
 # test kwargs
 prob2 = ODEProblem(sys, u0, tspan, p, jac = true)
 prob3 = ODEProblem(sys, u0, tspan, p, jac = true, sparse = true) #SparseMatrixCSC need to handle