refactor: use initialization_data in SciMLFunction constructors

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -359,10 +359,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         sparsity = false,
         analytic = nothing,
         split_idxs = nothing,
-        initializeprob = nothing,
-        update_initializeprob! = nothing,
-        initializeprobmap = nothing,
-        initializeprobpmap = nothing,
+        initialization_data = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
@@ -463,10 +460,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,
-        initializeprob = initializeprob,
-        update_initializeprob! = update_initializeprob!,
-        initializeprobmap = initializeprobmap,
-        initializeprobpmap = initializeprobpmap)
+        initialization_data)
 end
 
 """
@@ -496,10 +490,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         sparse = false, simplify = false,
         eval_module = @__MODULE__,
         checkbounds = false,
-        initializeprob = nothing,
-        initializeprobmap = nothing,
-        initializeprobpmap = nothing,
-        update_initializeprob! = nothing,
+        initialization_data = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
@@ -547,15 +538,12 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         nothing
     end
 
-    DAEFunction{iip}(f,
+    DAEFunction{iip}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         jac_prototype = jac_prototype,
         observed = observedfun,
-        initializeprob = initializeprob,
-        initializeprobmap = initializeprobmap,
-        initializeprobpmap = initializeprobpmap,
-        update_initializeprob! = update_initializeprob!)
+        initialization_data)
 end
 
 function DiffEqBase.DDEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -567,6 +555,7 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         eval_expression = false,
         eval_module = @__MODULE__,
         checkbounds = false,
+        initialization_data = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `DDEFunction`")
@@ -579,7 +568,7 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
     f(u, h, p, t) = f_oop(u, h, p, t)
     f(du, u, h, p, t) = f_iip(du, u, h, p, t)
 
-    DDEFunction{iip}(f, sys = sys)
+    DDEFunction{iip}(f; sys = sys, initialization_data)
 end
 
 function DiffEqBase.SDDEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -591,6 +580,7 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
         eval_expression = false,
         eval_module = @__MODULE__,
         checkbounds = false,
+        initialization_data = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `SDDEFunction`")
@@ -609,7 +599,7 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
     g(u, h, p, t) = g_oop(u, h, p, t)
     g(du, u, h, p, t) = g_iip(du, u, h, p, t)
 
-    SDDEFunction{iip}(f, g, sys = sys)
+    SDDEFunction{iip}(f, g; sys = sys, initialization_data)
 end
 
 """

src/systems/diffeqs/sdesystem.jl

@@ -544,7 +544,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         version = nothing, tgrad = false, sparse = false,
         jac = false, Wfact = false, eval_expression = false,
         eval_module = @__MODULE__,
-        checkbounds = false,
+        checkbounds = false, initialization_data = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed `SDESystem` is required. Call `complete` or `structural_simplify` on the system before creating an `SDEFunction`")
@@ -615,13 +615,13 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
 
     observedfun = ObservedFunctionCache(sys; eval_expression, eval_module)
 
-    SDEFunction{iip, specialize}(f, g,
+    SDEFunction{iip, specialize}(f, g;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
         Wfact = _Wfact === nothing ? nothing : _Wfact,
         Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,
-        mass_matrix = _M,
+        mass_matrix = _M, initialization_data,
         observed = observedfun)
 end
 

src/systems/nonlinear/initializesystem.jl

@@ -346,13 +346,7 @@ function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p,
         u0map, pmap, defs, cmap, dvs, ps)
     kws = maybe_build_initialization_problem(
         sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns; use_scc)
-    initprob = get(kws, :initializeprob, nothing)
-    if initprob === nothing
-        return nothing
-    end
-    return SciMLBase.OverrideInitData(initprob, get(kws, :update_initializeprob!, nothing),
-        get(kws, :initializeprobmap, nothing),
-        get(kws, :initializeprobpmap, nothing))
+    return get(kws, :initialization_data, nothing)
 end
 
 """

src/systems/nonlinear/nonlinearsystem.jl

@@ -344,6 +344,7 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = unknowns(s
         eval_expression = false,
         eval_module = @__MODULE__,
         sparse = false, simplify = false,
+        initialization_data = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `NonlinearSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `NonlinearFunction`")
@@ -376,14 +377,14 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = unknowns(s
         resid_prototype = calculate_resid_prototype(length(equations(sys)), u0, p)
     end
 
-    NonlinearFunction{iip}(f,
+    NonlinearFunction{iip}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         resid_prototype = resid_prototype,
         jac_prototype = sparse ?
                         similar(calculate_jacobian(sys, sparse = sparse),
             Float64) : nothing,
-        observed = observedfun)
+        observed = observedfun, initialization_data)
 end
 
 """
@@ -395,7 +396,8 @@ respectively.
 """
 function SciMLBase.IntervalNonlinearFunction(
         sys::NonlinearSystem, dvs = unknowns(sys), ps = parameters(sys), u0 = nothing;
-        p = nothing, eval_expression = false, eval_module = @__MODULE__, kwargs...)
+        p = nothing, eval_expression = false, eval_module = @__MODULE__,
+        initialization_data = nothing, kwargs...)
     if !iscomplete(sys)
         error("A completed `NonlinearSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `IntervalNonlinearFunction`")
     end
@@ -411,7 +413,8 @@ function SciMLBase.IntervalNonlinearFunction(
 
     observedfun = ObservedFunctionCache(sys; eval_expression, eval_module)
 
-    IntervalNonlinearFunction{false}(f; observed = observedfun, sys = sys)
+    IntervalNonlinearFunction{false}(
+        f; observed = observedfun, sys = sys, initialization_data)
 end
 
 """
@@ -884,6 +887,7 @@ function flatten(sys::NonlinearSystem, noeqs = false)
             observed = observed(sys),
             defaults = defaults(sys),
             guesses = guesses(sys),
+            initialization_eqs = initialization_equations(sys),
             name = nameof(sys),
             description = description(sys),
             metadata = get_metadata(sys),