Enable symbolic remake for optimization problems

Project.toml

@@ -71,7 +71,7 @@ NaNMath = "0.3, 1"
 RecursiveArrayTools = "2.3"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.4.3, 0.5"
-SciMLBase = "1.75.0"
+SciMLBase = "1.76.1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"

src/variables.jl

@@ -119,19 +119,16 @@ 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 `process_p_u0_symbolic` and process symbolic maps of `p` and/or `u0` if the 
 user has `ModelingToolkit` loaded.
 """
-function SciMLBase.process_p_u0_symbolic(prob::ODEProblem, p, u0)
+function SciMLBase.process_p_u0_symbolic(prob::Union{SciMLBase.AbstractDEProblem,
+                                                     NonlinearProblem, OptimizationProblem},
+                                         p,
+                                         u0)
     # check if a symbolic remake is possible
     if eltype(p) <: Pair
         hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||

test/nonlinearsystem.jl

@@ -189,8 +189,7 @@ eq = [v1 ~ sin(2pi * t * h)
 @named sys = ODESystem(eq)
 @test length(equations(structural_simplify(sys))) == 0
 
-#1504
-let
+@testset "Issue: 1504" begin
     @variables u[1:4]
 
     eqs = [u[1] ~ 1,
@@ -213,3 +212,23 @@ eqs = [0 ~ a * x]
 testdict = Dict([:test => 1])
 @named sys = NonlinearSystem(eqs, [x], [a], metadata = testdict)
 @test get_metadata(sys) == testdict
+
+@testset "Remake" begin
+    @parameters a=1.0 b=1.0 c=1.0
+    @constants h = 1
+    @variables x y z
+
+    eqs = [0 ~ a * (y - x) * h,
+        0 ~ x * (b - z) - y,
+        0 ~ x * y - c * z]
+    @named sys = NonlinearSystem(eqs, [x, y, z], [a, b, c], defaults = Dict(x => 2.0))
+    prob = NonlinearProblem(sys, ones(length(states(sys))))
+
+    prob_ = remake(prob, u0 = [1.0, 2.0, 3.0], p = [1.1, 1.2, 1.3])
+    @test prob_.u0 == [1.0, 2.0, 3.0]
+    @test prob_.p == [1.1, 1.2, 1.3]
+
+    prob_ = remake(prob, u0 = Dict(y => 2.0), p = Dict(a => 2.0))
+    @test prob_.u0 == [1.0, 2.0, 1.0]
+    @test prob_.p == [2.0, 1.0, 1.0]
+end

test/odesystem.jl

@@ -277,8 +277,8 @@ sol_dpmap = solve(prob_dpmap, Rodas5())
     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]
+    @test prob_new.p == [4.0, 3.0, 1.0]
+    @test prob_new.u0 == [1.0, 0.0]
 end
 
 # test kwargs

test/optimizationsystem.jl

@@ -137,23 +137,31 @@ end
 # nested constraints
 @testset "nested systems" begin
     @variables x y
-    o1 = (x - 1)^2
+    @parameters a = 1
+    o1 = (x - a)^2
     o2 = (y - 1 / 2)^2
     c1 = [
         x ~ 1,
     ]
     c2 = [
         y ~ 1,
     ]
-    sys1 = OptimizationSystem(o1, [x], [], name = :sys1, constraints = c1)
+    sys1 = OptimizationSystem(o1, [x], [a], name = :sys1, constraints = c1)
     sys2 = OptimizationSystem(o2, [y], [], name = :sys2, constraints = c2)
     sys = OptimizationSystem(0, [], []; name = :sys, systems = [sys1, sys2],
                              constraints = [sys1.x + sys2.y ~ 2], checks = false)
     prob = OptimizationProblem(sys, [0.0, 0.0])
-
     @test isequal(constraints(sys), vcat(sys1.x + sys2.y ~ 2, sys1.x ~ 1, sys2.y ~ 1))
-    @test isequal(equations(sys), (sys1.x - 1)^2 + (sys2.y - 1 / 2)^2)
+    @test isequal(equations(sys), (sys1.x - sys1.a)^2 + (sys2.y - 1 / 2)^2)
     @test isequal(states(sys), [sys1.x, sys2.y])
+
+    prob_ = remake(prob, u0 = [1.0, 0.0], p = [2.0])
+    @test isequal(prob_.u0, [1.0, 0.0])
+    @test isequal(prob_.p, [2.0])
+
+    prob_ = remake(prob, u0 = Dict(sys1.x => 1.0), p = Dict(sys1.a => 2.0))
+    @test isequal(prob_.u0, [1.0, 0.0])
+    @test isequal(prob_.p, [2.0])
 end
 
 @testset "time dependent var" begin