Merge pull request #956 from AayushSabharwal/as/steadystate-remake

feat: add dedicated `remake` method for `SteadyStateProblem`

Author: ChrisRackauckas

src/problems/problem_utils.jl

@@ -194,3 +194,4 @@ Base.copy(p::SciMLBase.NullParameters) = p
 
 SymbolicIndexingInterface.is_time_dependent(::AbstractDEProblem) = true
 SymbolicIndexingInterface.is_time_dependent(::AbstractNonlinearProblem) = false
+SymbolicIndexingInterface.is_time_dependent(::AbstractSteadyStateProblem) = true

src/problems/steady_state_problems.jl

@@ -113,6 +113,17 @@ function SteadyStateProblem(f, u0, p = NullParameters(); kwargs...)
     SteadyStateProblem(ODEFunction(f), u0, p; kwargs...)
 end
 
+function ConstructionBase.constructorof(::Type{P}) where {P <: SteadyStateProblem}
+    function ctor(f, u0, p, kw)
+        if f isa AbstractODEFunction
+            iip = isinplace(f)
+        else
+            iip = isinplace(f, 4)
+        end
+        return SteadyStateProblem{iip}(f, u0, p; kw...)
+    end
+end
+
 """
 $(SIGNATURES)
 

src/remake.jl

@@ -712,6 +712,47 @@ function remake(prob::NonlinearProblem;
     return prob
 end
 
+function remake(prob::SteadyStateProblem;
+        f = missing,
+        u0 = missing,
+        p = missing,
+        kwargs = missing,
+        interpret_symbolicmap = true,
+        use_defaults = false,
+        lazy_initialization = nothing,
+        build_initializeprob = true,
+        _kwargs...)
+    newu0, newp = updated_u0_p(prob, u0, p; interpret_symbolicmap, use_defaults)
+
+    if build_initializeprob
+        if f !== missing && has_initialization_data(f)
+            initialization_data = remake_initialization_data(
+                prob.f.sys, f, u0, nothing, p, newu0, newp)
+        else
+            initialization_data = remake_initialization_data(
+                prob.f.sys, prob.f, u0, nothing, p, newu0, newp)
+        end
+    else
+        initialization_data = nothing
+    end
+
+    f = coalesce(f, prob.f)
+    f = remake(prob.f; f, initialization_data)
+
+    prob = if kwargs === missing
+        SteadyStateProblem{isinplace(prob)}(f = f, u0 = newu0, p = newp; prob.kwargs...,
+            _kwargs...)
+    else
+        SteadyStateProblem{isinplace(prob)}(f = f, u0 = newu0, p = newp; kwargs...)
+    end
+
+    u0, p = maybe_eager_initialize_problem(prob, initialization_data, lazy_initialization)
+    @reset prob.u0 = u0
+    @reset prob.p = p
+
+    return prob
+end
+
 """
     remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p = missing,
         kwargs = missing, _kwargs...)

test/downstream/modelingtoolkit_remake.jl

@@ -33,6 +33,9 @@ tspan = (0.0, 100.0)
 push!(syss, sys)
 push!(probs, ODEProblem(sys, u0, tspan, p, jac = true))
 
+push!(syss, sys)
+push!(probs, SteadyStateProblem(ODEProblem(sys, u0, tspan, p, jac = true)))
+
 noise_eqs = [0.1x, 0.1y, 0.1z]
 @named sdesys = SDESystem(sys, noise_eqs)
 sdesys = complete(sdesys)

test/remake_tests.jl

@@ -20,6 +20,9 @@ fn = ODEFunction(lorenz!; sys)
 for T in containerTypes
     push!(probs, ODEProblem(fn, u0, tspan, T(p)))
 end
+for T in containerTypes
+    push!(probs, SteadyStateProblem(fn, u0, T(p)))
+end
 
 function ddelorenz!(du, u, h, p, t)
     du[1] = p[1] * (u[2] - u[1])