fix: fix remake for symbolic indexing, add tests for remake

Project.toml

@@ -24,6 +24,7 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
@@ -78,6 +79,7 @@ RecursiveArrayTools = "3.8.0"
 Reexport = "1"
 RuntimeGeneratedFunctions = "0.5.12"
 SciMLOperators = "0.3.7"
+SciMLStructures = "1.1"
 StaticArrays = "1.7"
 StaticArraysCore = "1.4"
 Statistics = "1.10"

src/SciMLBase.jl

@@ -5,6 +5,7 @@ if isdefined(Base, :Experimental) &&
 end
 using ConstructionBase
 using RecipesBase, RecursiveArrayTools, Tables
+using SciMLStructures
 using SymbolicIndexingInterface
 using DocStringExtensions
 using LinearAlgebra

src/remake.jl

@@ -25,6 +25,10 @@ Re-construct `thing` with new field values specified by the keyword
 arguments.
 """
 function remake(thing; kwargs...)
+    _remake_internal(thing; kwargs...)
+end
+
+function _remake_internal(thing; kwargs...)
     T = remaker_of(thing)
     if :kwargs ∈ fieldnames(typeof(thing))
         if :kwargs ∉ keys(kwargs)
@@ -41,6 +45,11 @@ function isrecompile(prob::ODEProblem{iip}) where {iip}
     (prob.f isa ODEFunction) ? !isfunctionwrapper(prob.f.f) : true
 end
 
+function remake(prob::AbstractSciMLProblem; u0 = missing, p = missing, interpret_symbolicmap = true, kwargs...)
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
+    _remake_internal(prob; kwargs..., u0, p)
+end
+
 """
     remake(prob::ODEProblem; f = missing, u0 = missing, tspan = missing,
            p = missing, kwargs = missing, _kwargs...)
@@ -59,37 +68,7 @@ function remake(prob::ODEProblem; f = missing,
         tspan = prob.tspan
     end
 
-    if p === missing && u0 === missing
-        p, u0 = prob.p, prob.u0
-    else # at least one of them has a value
-        if p === missing
-            p = prob.p
-        end
-        if u0 === missing
-            u0 = prob.u0
-        end
-        isu0symbolic = eltype(u0) <: Pair && !isempty(u0)
-        ispsymbolic = eltype(p) <: Pair && !isempty(p) && interpret_symbolicmap
-        if isu0symbolic && !has_sys(prob.f)
-            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" *
-                                "parameter order."))
-        end
-        if ispsymbolic && !has_sys(prob.f)
-            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) or pass `interpret_symbolicmap = false` as a keyword argument"))
-        end
-        if isu0symbolic && ispsymbolic
-            p, u0 = process_p_u0_symbolic(prob, p, u0)
-        elseif isu0symbolic
-            _, u0 = process_p_u0_symbolic(prob, prob.p, u0)
-        elseif ispsymbolic
-            p, _ = process_p_u0_symbolic(prob, p, prob.u0)
-        end
-    end
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
 
     iip = isinplace(prob)
 
@@ -141,21 +120,12 @@ end
 Remake the given `BVProblem`.
 """
 function remake(prob::BVProblem; f = missing, bc = missing, u0 = missing, tspan = missing,
-        p = missing, kwargs = missing, problem_type = missing, _kwargs...)
+        p = missing, kwargs = missing, problem_type = missing, interpret_symbolicmap = true, _kwargs...)
     if tspan === missing
         tspan = prob.tspan
     end
 
-    if p === missing && u0 === missing
-        p, u0 = prob.p, prob.u0
-    else # at least one of them has a value
-        if p === missing
-            p = prob.p
-        end
-        if u0 === missing
-            u0 = prob.u0
-        end
-    end
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
 
     iip = isinplace(prob)
 
@@ -211,20 +181,15 @@ function remake(prob::SDEProblem;
         p = missing,
         noise = missing,
         noise_rate_prototype = missing,
+        interpret_symbolicmap = true,
         seed = missing,
         kwargs = missing,
         _kwargs...)
     if tspan === missing
         tspan = prob.tspan
     end
 
-    if p === missing
-        p = prob.p
-    end
-
-    if u0 === missing
-        u0 = prob.u0
-    end
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
 
     if noise === missing
         noise = prob.noise
@@ -280,38 +245,8 @@ function remake(prob::OptimizationProblem;
         kwargs = missing,
         interpret_symbolicmap = true,
         _kwargs...)
-    if p === missing && u0 === missing
-        p, u0 = prob.p, prob.u0
-    else # at least one of them has a value
-        if p === missing
-            p = prob.p
-        end
-        if u0 === missing
-            u0 = prob.u0
-        end
-        isu0symbolic = eltype(u0) <: Pair && !isempty(u0)
-        ispsymbolic = eltype(p) <: Pair && !isempty(p) && interpret_symbolicmap
-        if isu0symbolic && !has_sys(prob.f)
-            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" *
-                                "parameter order."))
-        end
-        if ispsymbolic && !has_sys(prob.f)
-            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) or pass `interpret_symbolicmap = false` as a keyword argument"))
-        end
-        if isu0symbolic && ispsymbolic
-            p, u0 = process_p_u0_symbolic(prob, p, u0)
-        elseif isu0symbolic
-            _, u0 = process_p_u0_symbolic(prob, prob.p, u0)
-        elseif ispsymbolic
-            p, _ = process_p_u0_symbolic(prob, p, prob.u0)
-        end
-    end
 
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
     if f === missing
         f = prob.f
     end
@@ -362,38 +297,7 @@ function remake(prob::NonlinearProblem;
         kwargs = missing,
         interpret_symbolicmap = true,
         _kwargs...)
-    if p === missing && u0 === missing
-        p, u0 = prob.p, prob.u0
-    else # at least one of them has a value
-        if p === missing
-            p = prob.p
-        end
-        if u0 === missing
-            u0 = prob.u0
-        end
-        isu0symbolic = eltype(u0) <: Pair && !isempty(u0)
-        ispsymbolic = eltype(p) <: Pair && !isempty(p) && interpret_symbolicmap
-        if isu0symbolic && !has_sys(prob.f)
-            throw(ArgumentError("This problem does not support symbolic maps with" *
-                                " remake, i.e. it does not have a symbolic origin. Please use `remke`" *
-                                "with the `u0` keyword argument as a vector of values, paying attention to" *
-                                "parameter order."))
-        end
-        if ispsymbolic && !has_sys(prob.f)
-            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) or pass `interpret_symbolicmap = false` as a keyword argument"))
-        end
-        if isu0symbolic && ispsymbolic
-            p, u0 = process_p_u0_symbolic(prob, p, u0)
-        elseif isu0symbolic
-            _, u0 = process_p_u0_symbolic(prob, prob.p, u0)
-        elseif ispsymbolic
-            p, _ = process_p_u0_symbolic(prob, p, prob.u0)
-        end
-    end
-
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
     if f === missing
         f = prob.f
     end
@@ -418,17 +322,8 @@ end
 Remake the given `NonlinearLeastSquaresProblem`.
 """
 function remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p = missing,
-        kwargs = missing, _kwargs...)
-    if p === missing && u0 === missing
-        p, u0 = prob.p, prob.u0
-    else # at least one of them has a value
-        if p === missing
-            p = prob.p
-        end
-        if u0 === missing
-            u0 = prob.u0
-        end
-    end
+        interpret_symbolicmap = true, kwargs = missing, _kwargs...)
+    u0, p = updated_u0_p(prob, u0, p; interpret_symbolicmap)
 
     if f === missing
         f = prob.f
@@ -442,6 +337,91 @@ function remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p
     end
 end
 
+function updated_u0_p(prob, u0, p; interpret_symbolicmap = true)
+    newp = updated_p(prob, p; interpret_symbolicmap)
+    newu0 = updated_u0(prob, u0, p)
+    return newu0, newp
+end
+
+function updated_u0(prob, u0, p)
+    if u0 === missing
+        return copy(state_values(prob))
+    end
+    if u0 isa Number
+        return u0
+    end
+    if eltype(u0) <: Pair
+        u0 = Dict(u0)
+    elseif u0 isa AbstractArray
+        if length(u0) == length(state_values(prob))
+            return u0
+        else
+            throw(ArgumentError("Invalid value for u0: $u0. Must be an array of appropriate length or a symbolic map"))
+        end
+    else
+        throw(ArgumentError("Invalid value for u0: $u0. Must be an array of appropriate length or a symbolic map"))
+    end
+    if !has_sys(prob.f)
+        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 the order."))
+    end
+    newu0 = copy(state_values(prob))
+    if all(==(NotSymbolic()), symbolic_type.(values(u0)))
+        setu(prob, collect(keys(u0)))(newu0, collect(values(u0)))
+    else
+        value_syms = [k for (k, v) in u0 if symbolic_type(v) === NotSymbolic()]
+        dependent_syms = [k for (k, v) in u0 if symbolic_type(v) !== NotSymbolic()]
+        setu(prob, value_syms)(newu0, getindex.((u0,), value_syms))
+        obs = SymbolicIndexingInterface.observed(prob, getindex.((u0,), dependent_syms))
+        if is_time_dependent(prob)
+            dependent_vals = obs(newu0, p, current_time(prob))
+        else
+            dependent_vals = obs(newu0, p)
+        end
+        setu(prob, dependent_syms)(newu0, dependent_vals)
+    end
+    return newu0
+end
+
+function updated_p(prob, p; interpret_symbolicmap = true)
+    if p === missing
+        return copy(parameter_values(prob))
+    end
+    if eltype(p) <: Pair
+        if interpret_symbolicmap
+            has_sys(prob.f) || 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) or pass `interpret_symbolicmap = false` as a keyword argument"))
+        else
+            return p
+        end
+        p = Dict(p)
+    elseif SciMLStructures.isscimlstructure(p) || p isa Union{AbstractArray, Tuple} || typeof(p) === typeof(parameter_values(prob))
+        return p
+    else
+        throw(ArgumentError("Invalid value for p: $p. Must be an array of appropriate length, a symbolic map, or a SciMLStructure."))
+    end
+
+    newp = copy(parameter_values(prob))
+    if all(==(NotSymbolic()), symbolic_type.(values(p)))
+        setp(prob, collect(keys(p)))(newp, collect(values(p)))
+    else
+        value_syms = [k for (k, v) in p if symbolic_type(v) === NotSymbolic()]
+        dependent_syms = [k for (k, v) in p if symbolic_type(v) !== NotSymbolic()]
+        setp(prob, value_syms)(newp, getindex.((p,), value_syms))
+        obs = SymbolicIndexingInterface.observed(prob, getindex.((p,), dependent_syms))
+        if is_time_dependent(prob)
+            dependent_vals = obs(state_values(prob), newp, current_time(prob))
+        else
+            dependent_vals = obs(state_values(prob), newp)
+        end
+        setp(prob, dependent_syms)(newp, dependent_vals)
+    end
+    return newp
+end
+
 # overloaded in MTK to intercept symbolic remake
 function process_p_u0_symbolic(prob, p, u0)
     if prob isa Union{AbstractDEProblem, OptimizationProblem, NonlinearProblem}

test/downstream/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 BoundaryValueDiffEq = "764a87c0-6b3e-53db-9096-fe964310641d"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -17,6 +18,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 BoundaryValueDiffEq = "5"
 ForwardDiff = "0.10"
+JumpProcesses = "9.10"
 ModelingToolkit = "8.37, 9"
 NonlinearSolve = "2, 3"
 Optimization = "3"