refactor: use initialization_data instead of initializeprob, etc.

Author: AayushSabharwal

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, UIProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
+    SYS, ID} <: AbstractODEFunction{iip}
     f::F
     mass_matrix::TMM
     analytic::Ta
@@ -419,10 +419,7 @@ struct ODEFunction{iip, specialize, F, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TW
     observed::O
     colorvec::TCV
     sys::SYS
-    initializeprob::IProb
-    update_initializeprob!::UIProb
-    initializeprobmap::IProbMap
-    initializeprobpmap::IProbPmap
+    initialization_data::ID
 end
 
 @doc doc"""
@@ -526,7 +523,7 @@ information on generating the SplitFunction from this symbolic engine.
 struct SplitFunction{
     iip, specialize, F1, F2, TMM, C, Ta, Tt, TJ, JVP, VJP, JP, WP, SP, TW, TWt,
     TPJ, O,
-    TCV, SYS, IProb, UIProb, IProbMap, IProbPmap} <: AbstractODEFunction{iip}
+    TCV, SYS, ID} <: AbstractODEFunction{iip}
     f1::F1
     f2::F2
     mass_matrix::TMM
@@ -545,10 +542,7 @@ struct SplitFunction{
     observed::O
     colorvec::TCV
     sys::SYS
-    initializeprob::IProb
-    update_initializeprob!::UIProb
-    initializeprobmap::IProbMap
-    initializeprobpmap::IProbPmap
+    initialization_data::ID
 end
 
 @doc doc"""
@@ -1523,7 +1517,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, UIProb, IProbMap, IProbPmap} <:
+    SYS, ID} <:
        AbstractDAEFunction{iip}
     f::F
     analytic::Ta
@@ -1539,10 +1533,7 @@ struct DAEFunction{iip, specialize, F, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TP
     observed::O
     colorvec::TCV
     sys::SYS
-    initializeprob::IProb
-    update_initializeprob!::UIProb
-    initializeprobmap::IProbMap
-    initializeprobpmap::IProbPmap
+    initialization_data::ID
 end
 
 """
@@ -2432,7 +2423,9 @@ function ODEFunction{iip, specialize}(f;
         update_initializeprob! = __has_update_initializeprob!(f) ?
                                  f.update_initializeprob! : nothing,
         initializeprobmap = __has_initializeprobmap(f) ? f.initializeprobmap : nothing,
-        initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing
+        initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing,
+        initialization_data = __has_initialization_data(f) ? f.initialization_data :
+                              nothing
 ) where {iip,
         specialize
 }
@@ -2479,8 +2472,11 @@ function ODEFunction{iip, specialize}(f;
     _f = prepare_function(f)
 
     sys = sys_or_symbolcache(sys, syms, paramsyms, indepsym)
+    initdata = reconstruct_initialization_data(
+        initialization_data, initializeprob, update_initializeprob!,
+        initializeprobmap, initializeprobpmap)
 
-    @assert typeof(initializeprob) <:
+    @assert typeof(initdata.initializeprob) <:
             Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
 
     if specialize === NoSpecialize
@@ -2490,11 +2486,10 @@ function ODEFunction{iip, specialize}(f;
             typeof(sparsity), Any, Any, typeof(W_prototype), Any,
             Any,
             typeof(_colorvec),
-            typeof(sys), Any, Any, Any, Any}(_f, mass_matrix, analytic, tgrad, jac,
+            typeof(sys), Any}(_f, mass_matrix, analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
-            initializeprobpmap)
+            observed, _colorvec, sys, initdata)
     elseif specialize === false
         ODEFunction{iip, FunctionWrapperSpecialize,
             typeof(_f), typeof(mass_matrix), typeof(analytic), typeof(tgrad),
@@ -2503,13 +2498,11 @@ function ODEFunction{iip, specialize}(f;
             typeof(paramjac),
             typeof(observed),
             typeof(_colorvec),
-            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
-            typeof(initializeprobmap), typeof(initializeprobpmap)}(_f, mass_matrix,
+            typeof(sys), typeof(initdata)}(_f, mass_matrix,
             analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
-            initializeprobpmap)
+            observed, _colorvec, sys, initdata)
     else
         ODEFunction{iip, specialize,
             typeof(_f), typeof(mass_matrix), typeof(analytic), typeof(tgrad),
@@ -2518,13 +2511,10 @@ function ODEFunction{iip, specialize}(f;
             typeof(paramjac),
             typeof(observed),
             typeof(_colorvec),
-            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
-            typeof(initializeprobmap),
-            typeof(initializeprobpmap)}(_f, mass_matrix, analytic, tgrad, jac,
+            typeof(sys), typeof(initdata)}(_f, mass_matrix, analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity, Wfact,
             Wfact_t, W_prototype, paramjac,
-            observed, _colorvec, sys, initializeprob, update_initializeprob!, initializeprobmap,
-            initializeprobpmap)
+            observed, _colorvec, sys, initdata)
     end
 end
 
@@ -2541,27 +2531,23 @@ 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, Any}(
+            typeof(f.sys), 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.update_initializeprob!, f.initializeprobmap,
-            f.initializeprobpmap)
+            f.observed, f.colorvec, f.sys, f.initialization_data)
     else
         ODEFunction{isinplace(f), specialization(f), typeof(newf), typeof(f.mass_matrix),
             typeof(f.analytic), typeof(f.tgrad),
             typeof(f.jac), typeof(f.jvp), typeof(f.vjp), typeof(f.jac_prototype),
             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.update_initializeprob!),
-            typeof(f.initializeprobmap),
-            typeof(f.initializeprobpmap)}(newf, f.mass_matrix, f.analytic, f.tgrad, f.jac,
+            typeof(f.sys), typeof(f.initialization_data)}(
+            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.update_initializeprob!,
-            f.initializeprobmap, f.initializeprobpmap)
+            f.observed, f.colorvec, f.sys, f.initialization_data)
     end
 end
 
@@ -2692,8 +2678,8 @@ end
 
 @add_kwonly function SplitFunction(f1, f2, mass_matrix, cache, analytic, tgrad, jac, jvp,
         vjp, jac_prototype, W_prototype, sparsity, Wfact, Wfact_t, paramjac,
-        observed, colorvec, sys, initializeprob, update_initializeprob!,
-        initializeprobmap, initializeprobpmap)
+        observed, colorvec, sys, initializeprob = nothing, update_initializeprob! = nothing,
+        initializeprobmap = nothing, initializeprobpmap = nothing, initialization_data = nothing)
     f1 = ODEFunction(f1)
     f2 = ODEFunction(f2)
 
@@ -2702,17 +2688,20 @@ end
         throw(NonconformingFunctionsError(["f2"]))
     end
 
+    initdata = reconstruct_initialization_data(
+        initialization_data, initializeprob, update_initializeprob!,
+        initializeprobmap, initializeprobpmap)
+
     SplitFunction{isinplace(f2), FullSpecialize, typeof(f1), typeof(f2),
         typeof(mass_matrix),
         typeof(cache), typeof(analytic), typeof(tgrad), typeof(jac), typeof(jvp),
         typeof(vjp), typeof(jac_prototype), typeof(W_prototype), typeof(sparsity),
         typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(observed), typeof(colorvec),
-        typeof(sys), typeof(initializeprob), typeof(update_initializeprob!), typeof(initializeprobmap),
-        typeof(initializeprobpmap)}(
+        typeof(sys), typeof(initdata)}(
         f1, f2, mass_matrix,
         cache, analytic, tgrad, jac, jvp, vjp,
         jac_prototype, W_prototype, sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys,
-        initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
+        initdata)
 end
 function SplitFunction{iip, specialize}(f1, f2;
         mass_matrix = __has_mass_matrix(f1) ?
@@ -2748,37 +2737,39 @@ function SplitFunction{iip, specialize}(f1, f2;
         update_initializeprob! = __has_update_initializeprob!(f1) ?
                                  f1.update_initializeprob! : nothing,
         initializeprobmap = __has_initializeprobmap(f1) ? f1.initializeprobmap : nothing,
-        initializeprobpmap = __has_initializeprobpmap(f1) ? f1.initializeprobpmap : nothing
+        initializeprobpmap = __has_initializeprobpmap(f1) ? f1.initializeprobpmap : nothing,
+        initialization_data = __has_initialization_data(f1) ? f1.initialization_data :
+                              nothing
 ) where {iip,
         specialize
 }
     sys = sys_or_symbolcache(sys, syms, paramsyms, indepsym)
-    @assert typeof(initializeprob) <:
+    initdata = reconstruct_initialization_data(
+        initialization_data, initializeprob, update_initializeprob!,
+        initializeprobmap, initializeprobpmap)
+    @assert typeof(initdata.initializeprob) <:
             Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
 
     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}(f1, f2, mass_matrix, _func_cache,
+            Any, Any, Any}(f1, f2, mass_matrix, _func_cache,
             analytic,
             tgrad, jac, jvp, vjp, jac_prototype, W_prototype,
             sparsity, Wfact, Wfact_t, paramjac,
-            observed, colorvec, sys, initializeprob.update_initializeprob!, initializeprobmap,
-            initializeprobpmap, initializeprobpmap)
+            observed, colorvec, sys, initdata)
     else
         SplitFunction{iip, specialize, typeof(f1), typeof(f2), typeof(mass_matrix),
             typeof(_func_cache), typeof(analytic),
             typeof(tgrad), typeof(jac), typeof(jvp), typeof(vjp),
             typeof(jac_prototype), typeof(W_prototype), typeof(sparsity),
             typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(observed),
             typeof(colorvec),
-            typeof(sys), typeof(initializeprob), typeof(update_initializeprob!),
-            typeof(initializeprobmap),
-            typeof(initializeprobpmap)}(f1, f2,
+            typeof(sys), typeof(initdata)}(f1, f2,
             mass_matrix, _func_cache, analytic, tgrad, jac,
             jvp, vjp, jac_prototype, W_prototype,
             sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys,
-            initializeprob, update_initializeprob!, initializeprobmap, initializeprobpmap)
+            initdata)
     end
 end
 
@@ -3407,7 +3398,9 @@ function DAEFunction{iip, specialize}(f;
         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 {
+        initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing,
+        initialization_data = __has_initialization_data(f) ? f.initialization_data :
+                              nothing) where {
         iip,
         specialize
 }
@@ -3439,33 +3432,32 @@ function DAEFunction{iip, specialize}(f;
 
     _f = prepare_function(f)
     sys = sys_or_symbolcache(sys, syms, paramsyms, indepsym)
+    initdata = reconstruct_initialization_data(
+        initialization_data, initializeprob, update_initializeprob!,
+        initializeprobmap, initializeprobpmap)
 
-    @assert typeof(initializeprob) <:
+    @assert typeof(initdata.initializeprob) <:
             Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
 
     if specialize === NoSpecialize
         DAEFunction{iip, specialize, Any, Any, Any,
             Any, Any, Any, Any, Any,
             Any, Any, Any,
-            Any, typeof(_colorvec), Any, Any, Any, Any, Any}(_f, analytic, tgrad, jac, jvp,
+            Any, typeof(_colorvec), Any, Any}(_f, analytic, tgrad, jac, jvp,
             vjp, jac_prototype, sparsity,
             Wfact, Wfact_t, paramjac, observed,
-            _colorvec, sys, initializeprob, update_initializeprob!,
-            initializeprobmap, initializeprobpmap)
+            _colorvec, sys, initdata)
     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(update_initializeprob!),
-            typeof(initializeprobmap),
-            typeof(initializeprobpmap)}(
+            typeof(sys), typeof(initdata)}(
             _f, analytic, tgrad, jac, jvp, vjp,
             jac_prototype, sparsity, Wfact, Wfact_t,
             paramjac, observed,
-            _colorvec, sys, initializeprob, update_initializeprob!,
-            initializeprobmap, initializeprobpmap)
+            _colorvec, sys, initdata)
     end
 end
 
@@ -4384,6 +4376,14 @@ function sys_or_symbolcache(sys, syms, paramsyms, indepsym = nothing)
     return sys
 end
 
+function reconstruct_initialization_data(
+        initdata, initprob, update_initprob!, initprobmap, initprobpmap)
+    if initdata === nothing && initprob !== nothing
+        initdata = InitializationData(initprob, update_initprob!, initprobmap, initprobpmap)
+    end
+    return initprob
+end
+
 ########## Existence Functions
 
 # Check that field/property exists (may be nothing)
@@ -4407,10 +4407,19 @@ __has_colorvec(f) = isdefined(f, :colorvec)
 __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)
+function __has_initializeprob(f)
+    has_initialization_data(f) && isdefined(f.initialization_data, :initializeprob)
+end
+function __has_update_initializeprob!(f)
+    has_initialization_data(f) && isdefined(f.initialization_data, :update_initializeprob!)
+end
+function __has_initializeprobmap(f)
+    has_initialization_data(f) && isdefined(f.initialization_data, :initializeprobmap)
+end
+function __has_initializeprobpmap(f)
+    has_initialization_data(f) && isdefined(f.initialization_data, :initializeprobpmap)
+end
+__has_initialization_data(f) = isdefined(f, :initialization_data)
 
 # compatibility
 has_invW(f::AbstractSciMLFunction) = false
@@ -4424,16 +4433,20 @@ has_Wfact_t(f::AbstractSciMLFunction) = __has_Wfact_t(f) && f.Wfact_t !== nothin
 has_paramjac(f::AbstractSciMLFunction) = __has_paramjac(f) && f.paramjac !== nothing
 has_sys(f::AbstractSciMLFunction) = __has_sys(f) && f.sys !== nothing
 function has_initializeprob(f::AbstractSciMLFunction)
-    __has_initializeprob(f) && f.initializeprob !== nothing
+    __has_initializeprob(f) && f.initialization_data.initializeprob !== nothing
 end
 function has_update_initializeprob!(f::AbstractSciMLFunction)
-    __has_update_initializeprob!(f) && f.update_initializeprob! !== nothing
+    __has_update_initializeprob!(f) &&
+        f.initialization_data.update_initializeprob! !== nothing
 end
 function has_initializeprobmap(f::AbstractSciMLFunction)
-    __has_initializeprobmap(f) && f.initializeprobmap !== nothing
+    __has_initializeprobmap(f) && f.initialization_data.initializeprobmap !== nothing
 end
 function has_initializeprobpmap(f::AbstractSciMLFunction)
-    __has_initializeprobpmap(f) && f.initializeprobpmap !== nothing
+    __has_initializeprobpmap(f) && f.initialization_data.initializeprobpmap !== nothing
+end
+function has_initialization_data(f::AbstractSciMLFunction)
+    __has_initialization_data(f) && f.initialization_data !== nothing
 end
 
 function has_syms(f::AbstractSciMLFunction)