feat: add update_initializeprob! to relevant SciMLFunctions, remake

Author: AayushSabharwal

src/remake.jl

@@ -123,7 +123,7 @@ function remake(prob::ODEProblem; f = missing,
     iip = isinplace(prob)
 
     if f === missing
-        initializeprob, initializeprobmap, initializeprobpmap = remake_initializeprob(
+        initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap = remake_initializeprob(
             prob.f.sys, prob.f, u0, tspan[1], p)
         if specialization(prob.f) === FunctionWrapperSpecialize
             ptspan = promote_tspan(tspan)
@@ -133,14 +133,14 @@ function remake(prob::ODEProblem; f = missing,
                         unwrapped_f(prob.f.f),
                         (newu0, newu0, newp,
                             ptspan[1]));
-                    initializeprob, initializeprobmap, initializeprobpmap)
+                    initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
             else
                 _f = ODEFunction{iip, FunctionWrapperSpecialize}(
                     wrapfun_oop(
                         unwrapped_f(prob.f.f),
                         (newu0, newp,
                             ptspan[1]));
-                    initializeprob, initializeprobmap, initializeprobpmap)
+                    initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
             end
         else
             _f = prob.f
@@ -151,6 +151,13 @@ function remake(prob::ODEProblem; f = missing,
                 _f = parameterless_type(_f){
                     iip, specialization(_f), map(typeof, props)...}(props...)
             end
+            if __has_update_initializeprob!(_f)
+                props = getproperties(_f)
+                @reset props.update_initializeprob! = update_initializeprob!
+                props = values(props)
+                _f = parameterless_type(_f){
+                    iip, specialization(_f), map(typeof, props)...}(props...)
+            end
             if __has_initializeprobmap(_f)
                 props = getproperties(_f)
                 @reset props.initializeprobmap = initializeprobmap
@@ -196,18 +203,19 @@ end
 
 Re-create the initialization problem present in the function `scimlfn`, using the
 associated system `sys`, and the user-provided new values of `u0`, initial time `t0` and
-`p`. By default, returns `nothing, nothing, nothing` if `scimlfn` does not have an
+`p`. By default, returns `nothing, nothing, nothing, nothing` if `scimlfn` does not have an
 initialization problem, and
-`scimlfn.initializeprob, scimlfn.initializeprobmap, scimlfn.initializeprobpmap` if it
-does.
+`scimlfn.initializeprob, scimlfn.update_initializeprob!, scimlfn.initializeprobmap, scimlfn.initializeprobpmap`
+if it does.
 
 Note that `u0` or `p` may be `missing` if the user does not provide a value for them.
 """
 function remake_initializeprob(sys, scimlfn, u0, t0, p)
     if !has_initializeprob(scimlfn)
         return nothing, nothing, nothing
     end
-    return scimlfn.initializeprob, scimlfn.initializeprobmap, scimlfn.initializeprobpmap
+    return scimlfn.initializeprob,
+    scimlfn.update_initializeprob!, scimlfn.initializeprobmap, scimlfn.initializeprobpmap
 end
 
 """
@@ -689,7 +697,8 @@ function _updated_u0_p_symmap(prob, u0, ::Val{true}, p, ::Val{true}, t0)
     remake_buffer(prob, parameter_values(prob), keys(p), values(p))
 end
 
-function updated_u0_p(prob, u0, p, t0 = nothing; interpret_symbolicmap = true, use_defaults = false)
+function updated_u0_p(
+        prob, u0, p, t0 = nothing; interpret_symbolicmap = true, use_defaults = false)
     if u0 === missing && p === missing
         return state_values(prob), parameter_values(prob)
     end

src/scimlfunctions.jl

@@ -402,7 +402,7 @@ numerically-defined functions.
 """
 struct ODEFunction{iip, specialize, F, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, WP, TPJ,
     O, TCV,
-    SYS, IProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
+    SYS, IProb, UIProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
     f::F
     mass_matrix::TMM
     analytic::Ta
@@ -420,6 +420,7 @@ struct ODEFunction{iip, specialize, F, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TW
     colorvec::TCV
     sys::SYS
     initializeprob::IProb
+    update_initializeprob!::UIProb
     initializeprobmap::IProbMap
     initializeprobpmap::IProbPmap
 end
@@ -519,7 +520,7 @@ information on generating the SplitFunction from this symbolic engine.
 struct SplitFunction{
     iip, specialize, F1, F2, TMM, C, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt,
     TPJ, O,
-    TCV, SYS, IProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
+    TCV, SYS, IProb, UIProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
     f1::F1
     f2::F2
     mass_matrix::TMM
@@ -538,6 +539,7 @@ struct SplitFunction{
     colorvec::TCV
     sys::SYS
     initializeprob::IProb
+    update_initializeprob!::UIProb
     initializeprobmap::IProbMap
     initializeprobpmap::IProbPmap
 end
@@ -1508,7 +1510,7 @@ automatically symbolically generating the Jacobian and more from the
 numerically-defined functions.
 """
 struct DAEFunction{iip, specialize, F, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TPJ, O, TCV,
-    SYS, IProb, IProbMap, IProbPmap} <:
+    SYS, IProb, UIProb, IProbMap, IProbPmap} <:
        AbstractDAEFunction{iip}
     f::F
     analytic::Ta
@@ -1525,6 +1527,7 @@ struct DAEFunction{iip, specialize, F, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TP
     colorvec::TCV
     sys::SYS
     initializeprob::IProb
+    update_initializeprob!::UIProb
     initializeprobmap::IProbMap
     initializeprobpmap::IProbPmap
 end
@@ -2418,6 +2421,8 @@ function ODEFunction{iip, specialize}(f;
         colorvec = __has_colorvec(f) ? f.colorvec : nothing,
         sys = __has_sys(f) ? f.sys : nothing,
         initializeprob = __has_initializeprob(f) ? f.initializeprob : nothing,
+        update_initializeprob! = __has_update_initializeprob!(f) ?
+                                 f.update_initializeprob! : nothing,
         initializeprobmap = __has_initializeprobmap(f) ? f.initializeprobmap : nothing,
         initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing
 ) where {iip,
@@ -2477,10 +2482,10 @@ function ODEFunction{iip, specialize}(f;
             typeof(sparsity), Any, Any, typeof(W_prototype), Any,
             Any,
             typeof(_colorvec),
-            typeof(sys), Any, Any, Any}(_f, mass_matrix, analytic, tgrad, jac,
+            typeof(sys), Any, Any, Any, Any}(_f, mass_matrix, analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, initializeprobmap,
+            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
             initializeprobpmap)
     elseif specialize === false
         ODEFunction{iip, FunctionWrapperSpecialize,
@@ -2490,12 +2495,12 @@ function ODEFunction{iip, specialize}(f;
             typeof(paramjac),
             typeof(observed),
             typeof(_colorvec),
-            typeof(sys), typeof(initializeprob),
+            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
             typeof(initializeprobmap), typeof(initializeprobpmap)}(_f, mass_matrix,
             analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, initializeprobmap,
+            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
             initializeprobpmap)
     else
         ODEFunction{iip, specialize,
@@ -2505,11 +2510,12 @@ function ODEFunction{iip, specialize}(f;
             typeof(paramjac),
             typeof(observed),
             typeof(_colorvec),
-            typeof(sys), typeof(initializeprob), typeof(initializeprobmap),
+            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
+            typeof(initializeprobmap),
             typeof(initializeprobpmap)}(_f, mass_matrix, analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, initializeprobmap,
+            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
             initializeprobpmap)
     end
 end
@@ -2527,10 +2533,12 @@ function unwrapped_f(f::ODEFunction, newf = unwrapped_f(f.f))
             Any, Any, Any, Any, typeof(f.jac_prototype),
             typeof(f.sparsity), Any, Any, Any,
             Any, typeof(f.colorvec),
-            typeof(f.sys), Any, Any, Any}(newf, f.mass_matrix, f.analytic, f.tgrad, f.jac,
+            typeof(f.sys), Any, Any, Any, Any}(
+            newf, f.mass_matrix, f.analytic, f.tgrad, f.jac,
             f.jvp, f.vjp, f.jac_prototype, f.sparsity, f.Wfact,
             f.Wfact_t, f.W_prototype, f.paramjac,
-            f.observed, f.colorvec, f.sys, f.initializeprob, f.initializeprobmap,
+            f.observed, f.colorvec, f.sys, f.initializeprob,
+            f.update_initializeprob!, f.initializeprobmap,
             f.initializeprobpmap)
     else
         ODEFunction{isinplace(f), specialization(f), typeof(newf), typeof(f.mass_matrix),
@@ -2539,11 +2547,12 @@ function unwrapped_f(f::ODEFunction, newf = unwrapped_f(f.f))
             typeof(f.sparsity), typeof(f.Wfact), typeof(f.Wfact_t), typeof(f.W_prototype),
             typeof(f.paramjac),
             typeof(f.observed), typeof(f.colorvec),
-            typeof(f.sys), typeof(f.initializeprob), typeof(f.initializeprobmap),
+            typeof(f.sys), typeof(f.initializeprob), typeof(f.update_initializeprob!),
+            typeof(f.initializeprobmap),
             typeof(f.initializeprobpmap)}(newf, f.mass_matrix, f.analytic, f.tgrad, f.jac,
             f.jvp, f.vjp, f.jac_prototype, f.sparsity, f.Wfact,
             f.Wfact_t, f.W_prototype, f.paramjac,
-            f.observed, f.colorvec, f.sys, f.initializeprob,
+            f.observed, f.colorvec, f.sys, f.initializeprob, f.update_initializeprob!,
             f.initializeprobmap, f.initializeprobpmap)
     end
 end
@@ -2646,7 +2655,8 @@ end
 
 @add_kwonly function SplitFunction(f1, f2, mass_matrix, cache, analytic, tgrad, jac, jvp,
         vjp, jac_prototype, sparsity, Wfact, Wfact_t, paramjac,
-        observed, colorvec, sys, initializeprob, initializeprobmap, initializeprobpmap)
+        observed, colorvec, sys, initializeprob, update_initializeprob!,
+        initializeprobmap, initializeprobpmap)
     f1 = ODEFunction(f1)
     f2 = ODEFunction(f2)
 
@@ -2660,12 +2670,12 @@ end
         typeof(cache), typeof(analytic), typeof(tgrad), typeof(jac), typeof(jvp),
         typeof(vjp), typeof(jac_prototype), typeof(sparsity),
         typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(observed), typeof(colorvec),
-        typeof(sys), typeof(initializeprob), typeof(initializeprobmap),
+        typeof(sys), typeof(initializeprob), typeof(update_initializeprob!), typeof(initializeprobmap),
         typeof(initializeprobpmap)}(
         f1, f2, mass_matrix,
         cache, analytic, tgrad, jac, jvp, vjp,
         jac_prototype, sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys,
-        initializeprob, initializeprobmap, initializeprobpmap)
+        initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
 end
 function SplitFunction{iip, specialize}(f1, f2;
         mass_matrix = __has_mass_matrix(f1) ?
@@ -2695,6 +2705,8 @@ function SplitFunction{iip, specialize}(f1, f2;
                    nothing,
         sys = __has_sys(f1) ? f1.sys : nothing,
         initializeprob = __has_initializeprob(f1) ? f1.initializeprob : nothing,
+        update_initializeprob! = __has_update_initializeprob!(f1) ?
+                                 f1.update_initializeprob! : nothing,
         initializeprobmap = __has_initializeprobmap(f1) ? f1.initializeprobmap : nothing,
         initializeprobpmap = __has_initializeprobpmap(f1) ? f1.initializeprobpmap : nothing
 ) where {iip,
@@ -2706,12 +2718,12 @@ function SplitFunction{iip, specialize}(f1, f2;
 
     if specialize === NoSpecialize
         SplitFunction{iip, specialize, Any, Any, Any, Any, Any, Any, Any, Any, Any,
-            Any, Any, Any, Any, Any,
+            Any, Any, Any, Any, Any, Any,
             Any, Any, Any, Any, Any, Any}(f1, f2, mass_matrix, _func_cache,
             analytic,
             tgrad, jac, jvp, vjp, jac_prototype,
             sparsity, Wfact, Wfact_t, paramjac,
-            observed, colorvec, sys, initializeprob, initializeprobmap,
+            observed, colorvec, sys, initializeprob.update_initializeprob!, initializeprobmap,
             initializeprobpmap, initializeprobpmap)
     else
         SplitFunction{iip, specialize, typeof(f1), typeof(f2), typeof(mass_matrix),
@@ -2720,12 +2732,13 @@ function SplitFunction{iip, specialize}(f1, f2;
             typeof(jac_prototype), typeof(sparsity),
             typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(observed),
             typeof(colorvec),
-            typeof(sys), typeof(initializeprob), typeof(initializeprobmap),
+            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
+            typeof(initializeprobmap),
             typeof(initializeprobpmap)}(f1, f2,
             mass_matrix, _func_cache, analytic, tgrad, jac,
             jvp, vjp, jac_prototype,
             sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys,
-            initializeprob, initializeprobmap, initializeprobpmap)
+            initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
     end
 end
 
@@ -3351,6 +3364,8 @@ function DAEFunction{iip, specialize}(f;
         colorvec = __has_colorvec(f) ? f.colorvec : nothing,
         sys = __has_sys(f) ? f.sys : nothing,
         initializeprob = __has_initializeprob(f) ? f.initializeprob : nothing,
+        update_initializeprob! = __has_update_initializeprob!(f) ?
+                                 f.update_initializeprob! : nothing,
         initializeprobmap = __has_initializeprobmap(f) ? f.initializeprobmap : nothing,
         initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing) where {
         iip,
@@ -3392,22 +3407,25 @@ function DAEFunction{iip, specialize}(f;
         DAEFunction{iip, specialize, Any, Any, Any,
             Any, Any, Any, Any, Any,
             Any, Any, Any,
-            Any, typeof(_colorvec), Any, Any, Any, Any}(_f, analytic, tgrad, jac, jvp,
+            Any, typeof(_colorvec), Any, Any, Any, Any, Any}(_f, analytic, tgrad, jac, jvp,
             vjp, jac_prototype, sparsity,
             Wfact, Wfact_t, paramjac, observed,
-            _colorvec, sys, initializeprob, initializeprobmap, initializeprobpmap)
+            _colorvec, sys, initializeprob, update_initializeprob!,
+            initializeprobmap, initializeprobpmap)
     else
         DAEFunction{iip, specialize, typeof(_f), typeof(analytic), typeof(tgrad),
             typeof(jac), typeof(jvp), typeof(vjp), typeof(jac_prototype),
             typeof(sparsity), typeof(Wfact), typeof(Wfact_t),
             typeof(paramjac),
             typeof(observed), typeof(_colorvec),
-            typeof(sys), typeof(initializeprob), typeof(initializeprobmap),
+            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
+            typeof(initializeprobmap),
             typeof(initializeprobpmap)}(
             _f, analytic, tgrad, jac, jvp, vjp,
             jac_prototype, sparsity, Wfact, Wfact_t,
             paramjac, observed,
-            _colorvec, sys, initializeprob, initializeprobmap, initializeprobpmap)
+            _colorvec, sys, initializeprob, update_initializeprob!,
+            initializeprobmap, initializeprobpmap)
     end
 end
 
@@ -4350,6 +4368,7 @@ __has_sys(f) = isdefined(f, :sys)
 __has_analytic_full(f) = isdefined(f, :analytic_full)
 __has_resid_prototype(f) = isdefined(f, :resid_prototype)
 __has_initializeprob(f) = isdefined(f, :initializeprob)
+__has_update_initializeprob!(f) = isdefined(f, :update_initializeprob!)
 __has_initializeprobmap(f) = isdefined(f, :initializeprobmap)
 __has_initializeprobpmap(f) = isdefined(f, :initializeprobpmap)
 
@@ -4367,6 +4386,9 @@ has_sys(f::AbstractSciMLFunction) = __has_sys(f) && f.sys !== nothing
 function has_initializeprob(f::AbstractSciMLFunction)
     __has_initializeprob(f) && f.initializeprob !== nothing
 end
+function has_update_initializeprob!(f::AbstractSciMLFunction)
+    __has_update_initializeprob!(f) && f.update_initializeprob! !== nothing
+end
 function has_initializeprobmap(f::AbstractSciMLFunction)
     __has_initializeprobmap(f) && f.initializeprobmap !== nothing
 end