From 219aee396ed4d29606fcb3fa6922bb140091a0e6 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Dec 2024 13:07:15 +0530
Subject: [PATCH 01/18] refactor: add guesses to `SDESystem`,
 `NonlinearSystem`, `JumpSystem`

---
 src/systems/diffeqs/odesystem.jl              | 27 +++------
 src/systems/diffeqs/sdesystem.jl              | 52 ++++++++++++-----
 .../discrete_system/discrete_system.jl        | 47 +++++++++++----
 src/systems/jumps/jumpsystem.jl               | 43 ++++++++++----
 src/systems/nonlinear/nonlinearsystem.jl      | 57 ++++++++++++++-----
 .../optimization/constraints_system.jl        |  4 +-
 .../optimization/optimizationsystem.jl        |  4 +-
 src/utils.jl                                  | 27 +++++++++
 8 files changed, 187 insertions(+), 74 deletions(-)

diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 34003f40c2..61f16fd926 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -256,29 +256,16 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
             :ODESystem, force = true)
     end
     defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, dvs′)
-    process_variables!(var_to_name, defaults, ps′)
-    process_variables!(var_to_name, defaults, [eq.lhs for eq in parameter_dependencies])
-    process_variables!(var_to_name, defaults, [eq.rhs for eq in parameter_dependencies])
+    process_variables!(var_to_name, defaults, guesses, dvs′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
     defaults = Dict{Any, Any}(value(k) => value(v)
     for (k, v) in pairs(defaults) if v !== nothing)
-
-    sysdvsguesses = [ModelingToolkit.getguess(st) for st in dvs′]
-    hasaguess = findall(!isnothing, sysdvsguesses)
-    var_guesses = dvs′[hasaguess] .=> sysdvsguesses[hasaguess]
-    sysdvsguesses = isempty(var_guesses) ? Dict() : todict(var_guesses)
-    syspsguesses = [ModelingToolkit.getguess(st) for st in ps′]
-    hasaguess = findall(!isnothing, syspsguesses)
-    ps_guesses = ps′[hasaguess] .=> syspsguesses[hasaguess]
-    syspsguesses = isempty(ps_guesses) ? Dict() : todict(ps_guesses)
-    syspdepguesses = [ModelingToolkit.getguess(eq.lhs) for eq in parameter_dependencies]
-    hasaguess = findall(!isnothing, syspdepguesses)
-    pdep_guesses = [eq.lhs for eq in parameter_dependencies][hasaguess] .=>
-        syspdepguesses[hasaguess]
-    syspdepguesses = isempty(pdep_guesses) ? Dict() : todict(pdep_guesses)
-
-    guesses = merge(sysdvsguesses, syspsguesses, syspdepguesses, todict(guesses))
     guesses = Dict{Any, Any}(value(k) => value(v)
     for (k, v) in pairs(guesses) if v !== nothing)
 
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index ac47f4c45c..d604863024 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -93,6 +93,19 @@ struct SDESystem <: AbstractODESystem
     """
     defaults::Dict
     """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
     Type of the system.
     """
     connector_type::Any
@@ -144,9 +157,8 @@ struct SDESystem <: AbstractODESystem
     isscheduled::Bool
 
     function SDESystem(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed,
-            tgrad,
-            jac,
-            ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, connector_type,
+            tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults,
+            guesses, initializesystem, initialization_eqs, connector_type,
             cevents, devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
             complete = false, index_cache = nothing, parent = nothing, is_scalar_noise = false,
             is_dde = false,
@@ -171,9 +183,9 @@ struct SDESystem <: AbstractODESystem
             check_units(u, deqs, neqs)
         end
         new(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
-            ctrl_jac,
-            Wfact, Wfact_t, name, description, systems,
-            defaults, connector_type, cevents, devents,
+            ctrl_jac, Wfact, Wfact_t, name, description, systems,
+            defaults, guesses, initializesystem, initialization_eqs, connector_type, cevents,
+            devents,
             parameter_dependencies, metadata, gui_metadata, complete, index_cache, parent, is_scalar_noise,
             is_dde, isscheduled)
     end
@@ -187,6 +199,9 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
         default_u0 = Dict(),
         default_p = Dict(),
         defaults = _merge(Dict(default_u0), Dict(default_p)),
+        guesses = Dict(),
+        initializesystem = nothing,
+        initialization_eqs = Equation[],
         name = nothing,
         description = "",
         connector_type = nothing,
@@ -207,6 +222,8 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
     dvs′ = value.(dvs)
     ps′ = value.(ps)
     ctrl′ = value.(controls)
+    parameter_dependencies, ps′ = process_parameter_dependencies(
+        parameter_dependencies, ps′)
 
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
@@ -217,13 +234,21 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
             "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
             :SDESystem, force = true)
     end
-    defaults = todict(defaults)
-    defaults = Dict(value(k) => value(v)
-    for (k, v) in pairs(defaults) if value(v) !== nothing)
 
+    defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, dvs′)
-    process_variables!(var_to_name, defaults, ps′)
+    process_variables!(var_to_name, defaults, guesses, dvs′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
+    defaults = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(defaults) if v !== nothing)
+    guesses = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(guesses) if v !== nothing)
+
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     tgrad = RefValue(EMPTY_TGRAD)
@@ -233,14 +258,13 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
     Wfact_t = RefValue(EMPTY_JAC)
     cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
     disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
-    parameter_dependencies, ps′ = process_parameter_dependencies(
-        parameter_dependencies, ps′)
     if is_dde === nothing
         is_dde = _check_if_dde(deqs, iv′, systems)
     end
     SDESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, neqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
-        ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, connector_type,
+        ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, guesses,
+        initializesystem, initialization_eqs, connector_type,
         cont_callbacks, disc_callbacks, parameter_dependencies, metadata, gui_metadata,
         complete, index_cache, parent, is_scalar_noise, is_dde; checks = checks)
 end
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 01fca30235..7458237333 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -55,6 +55,19 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     defaults::Dict
     """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
     Inject assignment statements before the evaluation of the RHS function.
     """
     preface::Any
@@ -98,9 +111,8 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     isscheduled::Bool
 
     function DiscreteSystem(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name,
-            observed,
-            name, description,
-            systems, defaults, preface, connector_type, parameter_dependencies = Equation[],
+            observed, name, description, systems, defaults, guesses, initializesystem,
+            initialization_eqs, preface, connector_type, parameter_dependencies = Equation[],
             metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing, substitutions = nothing,
             complete = false, index_cache = nothing, parent = nothing,
@@ -116,8 +128,7 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
             check_units(u, discreteEqs)
         end
         new(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name, observed, name, description,
-            systems,
-            defaults,
+            systems, defaults, guesses, initializesystem, initialization_eqs,
             preface, connector_type, parameter_dependencies, metadata, gui_metadata,
             tearing_state, substitutions, complete, index_cache, parent, isscheduled)
     end
@@ -135,6 +146,9 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
         description = "",
         default_u0 = Dict(),
         default_p = Dict(),
+        guesses = Dict(),
+        initializesystem = nothing,
+        initialization_eqs = Equation[],
         defaults = _merge(Dict(default_u0), Dict(default_p)),
         preface = nothing,
         connector_type = nothing,
@@ -155,13 +169,21 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
             "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
             :DiscreteSystem, force = true)
     end
-    defaults = todict(defaults)
-    defaults = Dict(value(k) => value(v)
-    for (k, v) in pairs(defaults) if value(v) !== nothing)
 
+    defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, dvs′)
-    process_variables!(var_to_name, defaults, ps′)
+    process_variables!(var_to_name, defaults, guesses, dvs′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
+    defaults = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(defaults) if v !== nothing)
+    guesses = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(guesses) if v !== nothing)
+
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     sysnames = nameof.(systems)
@@ -170,7 +192,8 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
     end
     DiscreteSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         eqs, iv′, dvs′, ps′, tspan, var_to_name, observed, name, description, systems,
-        defaults, preface, connector_type, parameter_dependencies, metadata, gui_metadata, kwargs...)
+        defaults, guesses, initializesystem, initialization_eqs, preface, connector_type,
+        parameter_dependencies, metadata, gui_metadata, kwargs...)
 end
 
 function DiscreteSystem(eqs, iv; kwargs...)
@@ -225,6 +248,8 @@ function flatten(sys::DiscreteSystem, noeqs = false)
             parameters(sys),
             observed = observed(sys),
             defaults = defaults(sys),
+            guesses = guesses(sys),
+            initialization_eqs = initialization_equations(sys),
             name = nameof(sys),
             description = description(sys),
             metadata = get_metadata(sys),
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
index e5e17fb5f9..efc5a9be7d 100644
--- a/src/systems/jumps/jumpsystem.jl
+++ b/src/systems/jumps/jumpsystem.jl
@@ -84,6 +84,19 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
     """
     defaults::Dict
     """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
     Type of the system.
     """
     connector_type::Any
@@ -125,8 +138,9 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
 
     function JumpSystem{U}(
             tag, ap::U, iv, unknowns, ps, var_to_name, observed, name, description,
-            systems, defaults, connector_type, cevents, devents, parameter_dependencies,
-            metadata = nothing, gui_metadata = nothing,
+            systems, defaults, guesses, initializesystem, initialization_eqs, connector_type,
+            cevents, devents,
+            parameter_dependencies, metadata = nothing, gui_metadata = nothing,
             complete = false, index_cache = nothing, isscheduled = false;
             checks::Union{Bool, Int} = true) where {U <: ArrayPartition}
         if checks == true || (checks & CheckComponents) > 0
@@ -139,7 +153,8 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
             check_units(u, ap, iv)
         end
         new{U}(tag, ap, iv, unknowns, ps, var_to_name,
-            observed, name, description, systems, defaults,
+            observed, name, description, systems, defaults, guesses, initializesystem,
+            initialization_eqs,
             connector_type, cevents, devents, parameter_dependencies, metadata,
             gui_metadata, complete, index_cache, isscheduled)
     end
@@ -154,6 +169,9 @@ function JumpSystem(eqs, iv, unknowns, ps;
         default_u0 = Dict(),
         default_p = Dict(),
         defaults = _merge(Dict(default_u0), Dict(default_p)),
+        guesses = Dict(),
+        initializesystem = nothing,
+        initialization_eqs = Equation[],
         name = nothing,
         description = "",
         connector_type = nothing,
@@ -179,13 +197,17 @@ function JumpSystem(eqs, iv, unknowns, ps;
             :JumpSystem, force = true)
     end
     defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, us′)
-    process_variables!(var_to_name, defaults, ps′)
-    process_variables!(var_to_name, defaults, [eq.lhs for eq in parameter_dependencies])
-    process_variables!(var_to_name, defaults, [eq.rhs for eq in parameter_dependencies])
+    process_variables!(var_to_name, defaults, guesses, us′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
     #! format: off
     defaults = Dict{Any, Any}(value(k) => value(v) for (k, v) in pairs(defaults) if value(v) !== nothing)
+    guesses = Dict{Any, Any}(value(k) => value(v) for (k, v) in pairs(guesses) if v !== nothing)
     #! format: on
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
@@ -219,8 +241,9 @@ function JumpSystem(eqs, iv, unknowns, ps;
 
     JumpSystem{typeof(ap)}(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         ap, iv′, us′, ps′, var_to_name, observed, name, description, systems,
-        defaults, connector_type, cont_callbacks, disc_callbacks, parameter_dependencies,
-        metadata, gui_metadata, checks = checks)
+        defaults, guesses, initializesystem, initialization_eqs, connector_type,
+        cont_callbacks, disc_callbacks,
+        parameter_dependencies, metadata, gui_metadata, checks = checks)
 end
 
 ##### MTK dispatches for JumpSystems #####
@@ -494,7 +517,7 @@ function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothi
     if has_equations(sys)
         osys = ODESystem(equations(sys).x[4], get_iv(sys), unknowns(sys), parameters(sys);
             observed = observed(sys), name = nameof(sys), description = description(sys),
-            systems = get_systems(sys), defaults = defaults(sys),
+            systems = get_systems(sys), defaults = defaults(sys), guesses = guesses(sys),
             parameter_dependencies = parameter_dependencies(sys),
             metadata = get_metadata(sys), gui_metadata = get_gui_metadata(sys))
         osys = complete(osys)
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
index b2abac5184..46d39822bf 100644
--- a/src/systems/nonlinear/nonlinearsystem.jl
+++ b/src/systems/nonlinear/nonlinearsystem.jl
@@ -57,6 +57,19 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
     """
     defaults::Dict
     """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
     Type of the system.
     """
     connector_type::Any
@@ -97,9 +110,8 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
 
     function NonlinearSystem(
             tag, eqs, unknowns, ps, var_to_name, observed, jac, name, description,
-            systems,
-            defaults, connector_type, parameter_dependencies = Equation[], metadata = nothing,
-            gui_metadata = nothing,
+            systems, defaults, guesses, initializesystem, initialization_eqs, connector_type,
+            parameter_dependencies = Equation[], metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing, substitutions = nothing,
             complete = false, index_cache = nothing, parent = nothing,
             isscheduled = false; checks::Union{Bool, Int} = true)
@@ -107,8 +119,8 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
             u = __get_unit_type(unknowns, ps)
             check_units(u, eqs)
         end
-        new(tag, eqs, unknowns, ps, var_to_name, observed,
-            jac, name, description, systems, defaults,
+        new(tag, eqs, unknowns, ps, var_to_name, observed, jac, name, description,
+            systems, defaults, guesses, initializesystem, initialization_eqs,
             connector_type, parameter_dependencies, metadata, gui_metadata, tearing_state,
             substitutions, complete, index_cache, parent, isscheduled)
     end
@@ -121,6 +133,9 @@ function NonlinearSystem(eqs, unknowns, ps;
         default_u0 = Dict(),
         default_p = Dict(),
         defaults = _merge(Dict(default_u0), Dict(default_p)),
+        guesses = Dict(),
+        initializesystem = nothing,
+        initialization_eqs = Equation[],
         systems = NonlinearSystem[],
         connector_type = nothing,
         continuous_events = nothing, # this argument is only required for ODESystems, but is added here for the constructor to accept it without error
@@ -151,21 +166,32 @@ function NonlinearSystem(eqs, unknowns, ps;
     eqs = [wrap(eq.lhs) isa Symbolics.Arr ? eq : 0 ~ eq.rhs - eq.lhs for eq in eqs]
 
     jac = RefValue{Any}(EMPTY_JAC)
-    defaults = todict(defaults)
-    defaults = Dict{Any, Any}(value(k) => value(v)
-    for (k, v) in pairs(defaults) if value(v) !== nothing)
 
-    unknowns, ps = value.(unknowns), value.(ps)
+    ps′ = value.(ps)
+    dvs′ = value.(unknowns)
+    parameter_dependencies, ps′ = process_parameter_dependencies(
+        parameter_dependencies, ps′)
+
+    defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, unknowns)
-    process_variables!(var_to_name, defaults, ps)
+    process_variables!(var_to_name, defaults, guesses, dvs′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
+    defaults = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(defaults) if v !== nothing)
+    guesses = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(guesses) if v !== nothing)
+
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
-    parameter_dependencies, ps = process_parameter_dependencies(
-        parameter_dependencies, ps)
     NonlinearSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
-        eqs, unknowns, ps, var_to_name, observed, jac, name, description, systems, defaults,
-        connector_type, parameter_dependencies, metadata, gui_metadata, checks = checks)
+        eqs, dvs′, ps′, var_to_name, observed, jac, name, description, systems, defaults,
+        guesses, initializesystem, initialization_eqs, connector_type, parameter_dependencies,
+        metadata, gui_metadata, checks = checks)
 end
 
 function NonlinearSystem(eqs; kwargs...)
@@ -857,6 +883,7 @@ function flatten(sys::NonlinearSystem, noeqs = false)
             parameters(sys),
             observed = observed(sys),
             defaults = defaults(sys),
+            guesses = guesses(sys),
             name = nameof(sys),
             description = description(sys),
             metadata = get_metadata(sys),
diff --git a/src/systems/optimization/constraints_system.jl b/src/systems/optimization/constraints_system.jl
index a2756994ac..03225fc900 100644
--- a/src/systems/optimization/constraints_system.jl
+++ b/src/systems/optimization/constraints_system.jl
@@ -143,8 +143,8 @@ function ConstraintsSystem(constraints, unknowns, ps;
     for (k, v) in pairs(defaults) if value(v) !== nothing)
 
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, unknowns′)
-    process_variables!(var_to_name, defaults, ps′)
+    process_variables!(var_to_name, defaults, Dict(), unknowns′)
+    process_variables!(var_to_name, defaults, Dict(), ps′)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     ConstraintsSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
index 0398c892eb..0b20fdef79 100644
--- a/src/systems/optimization/optimizationsystem.jl
+++ b/src/systems/optimization/optimizationsystem.jl
@@ -132,8 +132,8 @@ function OptimizationSystem(op, unknowns, ps;
     for (k, v) in pairs(defaults) if value(v) !== nothing)
 
     var_to_name = Dict()
-    process_variables!(var_to_name, defaults, unknowns′)
-    process_variables!(var_to_name, defaults, ps′)
+    process_variables!(var_to_name, defaults, Dict(), unknowns′)
+    process_variables!(var_to_name, defaults, Dict(), ps′)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     OptimizationSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
diff --git a/src/utils.jl b/src/utils.jl
index e9ddad3a07..c3011c2a79 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -244,6 +244,13 @@ function setdefault(v, val)
     val === nothing ? v : wrap(setdefaultval(unwrap(v), value(val)))
 end
 
+function process_variables!(var_to_name, defs, guesses, vars)
+    collect_defaults!(defs, vars)
+    collect_guesses!(guesses, vars)
+    collect_var_to_name!(var_to_name, vars)
+    return nothing
+end
+
 function process_variables!(var_to_name, defs, vars)
     collect_defaults!(defs, vars)
     collect_var_to_name!(var_to_name, vars)
@@ -261,6 +268,17 @@ function collect_defaults!(defs, vars)
     return defs
 end
 
+function collect_guesses!(guesses, vars)
+    for v in vars
+        symbolic_type(v) == NotSymbolic() && continue
+        if haskey(guesses, v) || !hasguess(unwrap(v)) || (def = getguess(v)) === nothing
+            continue
+        end
+        guesses[v] = getguess(v)
+    end
+    return guesses
+end
+
 function collect_var_to_name!(vars, xs)
     for x in xs
         symbolic_type(x) == NotSymbolic() && continue
@@ -1146,3 +1164,12 @@ function similar_variable(var::BasicSymbolic, name = :anon)
     end
     return sym
 end
+
+function guesses_from_metadata!(guesses, vars)
+    varguesses = [getguess(v) for v in vars]
+    hasaguess = findall(!isnothing, varguesses)
+    for i in hasaguess
+        haskey(guesses, vars[i]) && continue
+        guesses[vars[i]] = varguesses[i]
+    end
+end

From af8cd679571d6c2d078fbebdd9219536b270f335 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Dec 2024 13:07:33 +0530
Subject: [PATCH 02/18] feat: support arbitrary systems in
 `generate_initializesystem`

---
 src/systems/nonlinear/initializesystem.jl | 65 +++++++++++++----------
 1 file changed, 38 insertions(+), 27 deletions(-)

diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index 2344727920..229d462911 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -3,7 +3,7 @@ $(TYPEDSIGNATURES)
 
 Generate `NonlinearSystem` which initializes an ODE problem from specified initial conditions of an `ODESystem`.
 """
-function generate_initializesystem(sys::ODESystem;
+function generate_initializesystem(sys::AbstractSystem;
         u0map = Dict(),
         pmap = Dict(),
         initialization_eqs = [],
@@ -12,28 +12,36 @@ function generate_initializesystem(sys::ODESystem;
         algebraic_only = false,
         check_units = true, check_defguess = false,
         name = nameof(sys), extra_metadata = (;), kwargs...)
-    trueobs, eqs = unhack_observed(observed(sys), equations(sys))
+    eqs = equations(sys)
+    eqs = filter(x -> x isa Equation, eqs)
+    trueobs, eqs = unhack_observed(observed(sys), eqs)
     vars = unique([unknowns(sys); getfield.(trueobs, :lhs)])
     vars_set = Set(vars) # for efficient in-lookup
 
-    idxs_diff = isdiffeq.(eqs)
-    idxs_alge = .!idxs_diff
-
-    # prepare map for dummy derivative substitution
-    eqs_diff = eqs[idxs_diff]
-    D = Differential(get_iv(sys))
-    diffmap = merge(
-        Dict(eq.lhs => eq.rhs for eq in eqs_diff),
-        Dict(D(eq.lhs) => D(eq.rhs) for eq in trueobs)
-    )
-
-    # 1) process dummy derivatives and u0map into initialization system
-    eqs_ics = eqs[idxs_alge] # start equation list with algebraic equations
+    eqs_ics = Equation[]
     defs = copy(defaults(sys)) # copy so we don't modify sys.defaults
     additional_guesses = anydict(guesses)
     guesses = merge(get_guesses(sys), additional_guesses)
-    schedule = getfield(sys, :schedule)
-    if !isnothing(schedule)
+    idxs_diff = isdiffeq.(eqs)
+
+    # 1) Use algebraic equations of time-dependent systems as initialization constraints
+    if has_iv(sys)
+        idxs_alge = .!idxs_diff
+        append!(eqs_ics, eqs[idxs_alge]) # start equation list with algebraic equations
+
+        eqs_diff = eqs[idxs_diff]
+        D = Differential(get_iv(sys))
+        diffmap = merge(
+            Dict(eq.lhs => eq.rhs for eq in eqs_diff),
+            Dict(D(eq.lhs) => D(eq.rhs) for eq in trueobs)
+        )
+    else
+        diffmap = Dict()
+    end
+
+    if has_schedule(sys) && (schedule = get_schedule(sys); !isnothing(schedule))
+        # 2) process dummy derivatives and u0map into initialization system
+        # prepare map for dummy derivative substitution
         for x in filter(x -> !isnothing(x[1]), schedule.dummy_sub)
             # set dummy derivatives to default_dd_guess unless specified
             push!(defs, x[1] => get(guesses, x[1], default_dd_guess))
@@ -61,9 +69,14 @@ function generate_initializesystem(sys::ODESystem;
                 process_u0map_with_dummysubs(y, x)
             end
         end
+    else
+        # 2) System doesn't have a schedule, so dummy derivatives don't exist/aren't handled (SDESystem)
+        for (k, v) in u0map
+            defs[k] = v
+        end
     end
 
-    # 2) process other variables
+    # 3) process other variables
     for var in vars
         if var ∈ keys(defs)
             push!(eqs_ics, var ~ defs[var])
@@ -74,7 +87,7 @@ function generate_initializesystem(sys::ODESystem;
         end
     end
 
-    # 3) process explicitly provided initialization equations
+    # 4) process explicitly provided initialization equations
     if !algebraic_only
         initialization_eqs = [get_initialization_eqs(sys); initialization_eqs]
         for eq in initialization_eqs
@@ -83,7 +96,7 @@ function generate_initializesystem(sys::ODESystem;
         end
     end
 
-    # 4) process parameters as initialization unknowns
+    # 5) process parameters as initialization unknowns
     paramsubs = Dict()
     if pmap isa SciMLBase.NullParameters
         pmap = Dict()
@@ -138,7 +151,7 @@ function generate_initializesystem(sys::ODESystem;
         end
     end
 
-    # 5) parameter dependencies become equations, their LHS become unknowns
+    # 6) parameter dependencies become equations, their LHS become unknowns
     # non-numeric dependent parameters stay as parameter dependencies
     new_parameter_deps = Equation[]
     for eq in parameter_dependencies(sys)
@@ -153,7 +166,7 @@ function generate_initializesystem(sys::ODESystem;
         push!(defs, varp => guessval)
     end
 
-    # 6) handle values provided for dependent parameters similar to values for observed variables
+    # 7) handle values provided for dependent parameters similar to values for observed variables
     for (k, v) in merge(defaults(sys), pmap)
         if is_variable_floatingpoint(k) && has_parameter_dependency_with_lhs(sys, k)
             push!(eqs_ics, paramsubs[k] ~ v)
@@ -161,12 +174,10 @@ function generate_initializesystem(sys::ODESystem;
     end
 
     # parameters do not include ones that became initialization unknowns
-    pars = vcat(
-        [get_iv(sys)], # include independent variable as pseudo-parameter
-        [p for p in parameters(sys) if !haskey(paramsubs, p)]
-    )
+    pars = Vector{SymbolicParam}(filter(p -> !haskey(paramsubs, p), parameters(sys)))
+    is_time_dependent(sys) && push!(pars, get_iv(sys))
 
-    # 7) use observed equations for guesses of observed variables if not provided
+    # 8) use observed equations for guesses of observed variables if not provided
     for eq in trueobs
         haskey(defs, eq.lhs) && continue
         any(x -> isequal(default_toterm(x), eq.lhs), keys(defs)) && continue

From 39281945c17108c2b4a3434d89a4e57911905927 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Dec 2024 11:32:03 +0530
Subject: [PATCH 03/18] refactor: use `initialization_data` in SciMLFunction
 constructors

---
 src/systems/diffeqs/abstractodesystem.jl  | 28 ++++++++---------------
 src/systems/diffeqs/sdesystem.jl          |  6 ++---
 src/systems/nonlinear/initializesystem.jl |  8 +------
 src/systems/nonlinear/nonlinearsystem.jl  | 12 ++++++----
 4 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index f4e29346ff..7d9369d441 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/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
 
 """
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index d604863024..37e743218a 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/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
 
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index 229d462911..624ee1bd71 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/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
 
 """
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
index 46d39822bf..8cd8175668 100644
--- a/src/systems/nonlinear/nonlinearsystem.jl
+++ b/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),

From 4044317b17b3f239f0a4fa6a87f16f368f8f28aa Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Dec 2024 11:33:14 +0530
Subject: [PATCH 04/18] fix: don't build initializeprob for initializeprob

---
 src/systems/diffeqs/abstractodesystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 7d9369d441..b3d4903191 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1395,5 +1395,5 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
     else
         NonlinearLeastSquaresProblem
     end
-    TProb(isys, u0map, parammap; kwargs...)
+    TProb(isys, u0map, parammap; kwargs..., build_initializeprob = false)
 end

From 180b97801da7d134754c16becc9daebbfb688912 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Dec 2024 11:34:02 +0530
Subject: [PATCH 05/18] feat: build initialization system for all system types
 in `process_SciMLProblem`

---
 src/systems/problem_utils.jl | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
index 6b75fb2c6d..6ce5704daa 100644
--- a/src/systems/problem_utils.jl
+++ b/src/systems/problem_utils.jl
@@ -540,8 +540,9 @@ function maybe_build_initialization_problem(
     end
     if (((implicit_dae || has_observed_u0s || !isempty(missing_unknowns) ||
           !isempty(solvablepars) || has_dependent_unknowns) &&
-         get_tearing_state(sys) !== nothing) ||
-        !isempty(initialization_equations(sys))) && t !== nothing
+         (!has_tearing_state(sys) || get_tearing_state(sys) !== nothing)) ||
+        !isempty(initialization_equations(sys))) &&
+       (!is_time_dependent(sys) || t !== nothing)
         initializeprob = ModelingToolkit.InitializationProblem(
             sys, t, u0map, pmap; guesses, kwargs...)
         initializeprobmap = getu(initializeprob, unknowns(sys))
@@ -567,7 +568,9 @@ function maybe_build_initialization_problem(
         end
         empty!(missing_unknowns)
         return (;
-            initializeprob, initializeprobmap, initializeprobpmap, update_initializeprob!)
+            initialization_data = SciMLBase.OverrideInitData(
+                initializeprob, update_initializeprob!, initializeprobmap,
+                initializeprobpmap))
     end
     return (;)
 end
@@ -662,7 +665,7 @@ function process_SciMLProblem(
     op, missing_unknowns, missing_pars = build_operating_point(
         u0map, pmap, defs, cmap, dvs, ps)
 
-    if sys isa ODESystem && build_initializeprob
+    if build_initializeprob
         kws = maybe_build_initialization_problem(
             sys, op, u0map, pmap, t, defs, guesses, missing_unknowns;
             implicit_dae, warn_initialize_determined, initialization_eqs,

From c9c613f314cddb9ade459920cbb102a3b57b22da Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Dec 2024 17:18:07 +0530
Subject: [PATCH 06/18] fix: retain system data on `structural_simplify` of
 `SDESystem`

---
 src/systems/systems.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 47acd81a82..04c50bc766 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -164,6 +164,7 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
         return SDESystem(Vector{Equation}(full_equations(ode_sys)), noise_eqs,
             get_iv(ode_sys), unknowns(ode_sys), parameters(ode_sys);
             name = nameof(ode_sys), is_scalar_noise, observed = observed(ode_sys), defaults = defaults(sys),
-            parameter_dependencies = parameter_dependencies(sys))
+            parameter_dependencies = parameter_dependencies(sys),
+            guesses = guesses(sys), initialization_eqs = initialization_equations(sys))
     end
 end

From 4d5daa3caeaad6ee4ed6f9698194aa467d1c4ea4 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 3 Dec 2024 16:32:11 +0530
Subject: [PATCH 07/18] fix: pass `t` to `process_SciMLProblem` in `SDEProblem`

---
 src/systems/diffeqs/sdesystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 37e743218a..5e0d0e3208 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -738,7 +738,7 @@ function DiffEqBase.SDEProblem{iip, specialize}(
     end
     f, u0, p = process_SciMLProblem(
         SDEFunction{iip, specialize}, sys, u0map, parammap; check_length,
-        kwargs...)
+        t = tspan === nothing ? nothing : tspan[1], kwargs...)
     cbs = process_events(sys; callback, kwargs...)
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 

From 8d09409a011a4d91dfe6d22d81b255f9a0da753f Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 3 Dec 2024 16:33:03 +0530
Subject: [PATCH 08/18] feat: support arbitrary systems in
 `remake_initialization_data`

---
 src/systems/nonlinear/initializesystem.jl | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index 624ee1bd71..e32f967f4a 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -223,16 +223,19 @@ function is_parameter_solvable(p, pmap, defs, guesses)
              _val1 === nothing && _val2 !== nothing)) && _val3 !== nothing
 end
 
-function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p, newu0, newp)
+function SciMLBase.remake_initialization_data(
+        sys::AbstractSystem, odefn, u0, t0, p, newu0, newp)
     if u0 === missing && p === missing
         return odefn.initialization_data
     end
     if !(eltype(u0) <: Pair) && !(eltype(p) <: Pair)
-        oldinitprob = odefn.initializeprob
+        oldinitdata = odefn.initialization_data
+        oldinitdata === nothing && return nothing
+
+        oldinitprob = oldinitdata.initializeprob
         oldinitprob === nothing && return nothing
         if !SciMLBase.has_sys(oldinitprob.f) || !(oldinitprob.f.sys isa NonlinearSystem)
-            return SciMLBase.OverrideInitData(oldinitprob, odefn.update_initializeprob!,
-                odefn.initializeprobmap, odefn.initializeprobpmap)
+            return oldinitdata
         end
         pidxs = ParameterIndex[]
         pvals = []
@@ -254,7 +257,7 @@ function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p,
         if p !== missing
             for sym in parameter_symbols(oldinitprob)
                 push!(pidxs, parameter_index(oldinitprob, sym))
-                if isequal(sym, get_iv(sys))
+                if is_time_dependent(sys) && isequal(sym, get_iv(sys))
                     push!(pvals, t0)
                 else
                     push!(pvals, getp(sys, sym)(p))
@@ -283,8 +286,8 @@ function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p,
                     length(oldinitprob.f.resid_prototype), newu0, newp))
         end
         initprob = remake(oldinitprob; f = newf, u0 = newu0, p = newp)
-        return SciMLBase.OverrideInitData(initprob, odefn.update_initializeprob!,
-            odefn.initializeprobmap, odefn.initializeprobpmap)
+        return SciMLBase.OverrideInitData(initprob, oldinitdata.update_initializeprob!,
+            oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap)
     end
     dvs = unknowns(sys)
     ps = parameters(sys)
@@ -298,7 +301,7 @@ function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p,
     use_scc = true
 
     if SciMLBase.has_initializeprob(odefn)
-        oldsys = odefn.initializeprob.f.sys
+        oldsys = odefn.initialization_data.initializeprob.f.sys
         meta = get_metadata(oldsys)
         if meta isa InitializationSystemMetadata
             u0map = merge(meta.u0map, u0map)
@@ -336,7 +339,7 @@ function SciMLBase.remake_initialization_data(sys::ODESystem, odefn, u0, t0, p,
             pmap[p] = getp(sys, p)(newp)
         end
     end
-    if t0 === nothing
+    if t0 === nothing && is_time_dependent(sys)
         t0 = 0.0
     end
     filter_missing_values!(u0map)

From def207b92a83baf5363ebaded009c5b95706d645 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 3 Dec 2024 16:33:13 +0530
Subject: [PATCH 09/18] fix: fix type promotion bug in `remake_buffer`

---
 src/systems/parameter_buffer.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/systems/parameter_buffer.jl b/src/systems/parameter_buffer.jl
index 7d64054acc..bc4e62a773 100644
--- a/src/systems/parameter_buffer.jl
+++ b/src/systems/parameter_buffer.jl
@@ -574,6 +574,10 @@ function _remake_buffer(indp, oldbuf::MTKParameters, idxs, vals; validate = true
 
     @set! newbuf.tunable = narrow_buffer_type_and_fallback_undefs(
         oldbuf.tunable, newbuf.tunable)
+    if eltype(newbuf.tunable) <: Integer
+        T = promote_type(eltype(newbuf.tunable), Float64)
+        @set! newbuf.tunable = T.(newbuf.tunable)
+    end
     @set! newbuf.discrete = narrow_buffer_type_and_fallback_undefs.(
         oldbuf.discrete, newbuf.discrete)
     @set! newbuf.constant = narrow_buffer_type_and_fallback_undefs.(

From d971b185ad5d4d9f1714448131b2bb02aff58b2d Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 3 Dec 2024 16:33:51 +0530
Subject: [PATCH 10/18] test: test initialization on `SDEProblem`,
 `DDEProblem`, `SDDEProblem`

---
 test/initializationsystem.jl | 428 ++++++++++++++++++++---------------
 1 file changed, 247 insertions(+), 181 deletions(-)

diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 9c06dd2030..045b7d2d2c 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -1,4 +1,5 @@
 using ModelingToolkit, OrdinaryDiffEq, NonlinearSolve, Test
+using StochasticDiffEq, DelayDiffEq, StochasticDelayDiffEq
 using ForwardDiff
 using SymbolicIndexingInterface, SciMLStructures
 using SciMLStructures: Tunable
@@ -583,108 +584,123 @@ sol = solve(oprob_2nd_order_2, Rosenbrock23()) # retcode: Success
 end
 
 @testset "Initialization of parameters" begin
-    function test_parameter(prob, sym, val, initialval = zero(val))
-        @test prob.ps[sym] ≈ initialval
-        @test init(prob, Tsit5()).ps[sym] ≈ val
-        @test solve(prob, Tsit5()).ps[sym] ≈ val
-    end
-    function test_initializesystem(sys, u0map, pmap, p, equation)
-        isys = ModelingToolkit.generate_initializesystem(
-            sys; u0map, pmap, guesses = ModelingToolkit.guesses(sys))
-        @test is_variable(isys, p)
-        @test equation in equations(isys) || (0 ~ -equation.rhs) in equations(isys)
-    end
-    @variables x(t) y(t)
+    @variables _x(..) y(t)
     @parameters p q
-    u0map = Dict(x => 1.0, y => 1.0)
-    pmap = Dict()
-    pmap[q] = 1.0
-    # `missing` default, equation from ODEProblem
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x * q, D(y) ~ y * p], t; defaults = [p => missing], guesses = [p => 1.0])
-    pmap[p] = 2q
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), pmap)
-    test_parameter(prob, p, 2.0)
-    prob2 = remake(prob; u0 = u0map, p = pmap)
-    prob2.ps[p] = 0.0
-    test_parameter(prob2, p, 2.0)
-    # `missing` default, provided guess
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x, p ~ x + y], t; defaults = [p => missing], guesses = [p => 0.0])
-    prob = ODEProblem(sys, u0map, (0.0, 1.0))
-    test_parameter(prob, p, 2.0)
-    test_initializesystem(sys, u0map, pmap, p, 0 ~ p - x - y)
-    prob2 = remake(prob; u0 = u0map)
-    prob2.ps[p] = 0.0
-    test_parameter(prob2, p, 2.0)
-
-    # `missing` to ODEProblem, equation from default
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x * q, D(y) ~ y * p], t; defaults = [p => 2q], guesses = [p => 1.0])
-    pmap[p] = missing
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), pmap)
-    test_parameter(prob, p, 2.0)
-    test_initializesystem(sys, u0map, pmap, p, 0 ~ 2q - p)
-    prob2 = remake(prob; u0 = u0map, p = pmap)
-    prob2.ps[p] = 0.0
-    test_parameter(prob2, p, 2.0)
-    # `missing` to ODEProblem, provided guess
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x, p ~ x + y], t; guesses = [p => 0.0])
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), pmap)
-    test_parameter(prob, p, 2.0)
-    test_initializesystem(sys, u0map, pmap, p, 0 ~ x + y - p)
-    prob2 = remake(prob; u0 = u0map, p = pmap)
-    prob2.ps[p] = 0.0
-    test_parameter(prob2, p, 2.0)
-
-    # No `missing`, default and guess
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x * q, D(y) ~ y * p], t; defaults = [p => 2q], guesses = [p => 0.0])
-    delete!(pmap, p)
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), pmap)
-    test_parameter(prob, p, 2.0)
-    test_initializesystem(sys, u0map, pmap, p, 0 ~ 2q - p)
-    prob2 = remake(prob; u0 = u0map, p = pmap)
-    prob2.ps[p] = 0.0
-    test_parameter(prob2, p, 2.0)
-
-    # Default overridden by ODEProblem, guess provided
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ q * x, D(y) ~ y * p], t; defaults = [p => 2q], guesses = [p => 1.0])
-    _pmap = merge(pmap, Dict(p => q))
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), _pmap)
-    test_parameter(prob, p, _pmap[q])
-    test_initializesystem(sys, u0map, _pmap, p, 0 ~ q - p)
-
-    # ODEProblem dependent value with guess, no `missing`
-    @mtkbuild sys = ODESystem([D(x) ~ x * q, D(y) ~ y * p], t; guesses = [p => 0.0])
-    _pmap = merge(pmap, Dict(p => 3q))
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), _pmap)
-    test_parameter(prob, p, 3pmap[q])
-
-    # Should not be solved for:
-
-    # Override dependent default with direct value
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ q * x, D(y) ~ y * p], t; defaults = [p => 2q], guesses = [p => 1.0])
-    _pmap = merge(pmap, Dict(p => 1.0))
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), _pmap)
-    @test prob.ps[p] ≈ 1.0
-    @test prob.f.initializeprob === nothing
-
-    # Non-floating point
-    @parameters r::Int s::Int
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ s * x, D(y) ~ y * r], t; defaults = [s => 2r], guesses = [s => 1.0])
-    prob = ODEProblem(sys, u0map, (0.0, 1.0), [r => 1])
-    @test prob.ps[r] == 1
-    @test prob.ps[s] == 2
-    @test prob.f.initializeprob === nothing
-
-    @mtkbuild sys = ODESystem([D(x) ~ x, p ~ x + y], t; guesses = [p => 0.0])
-    @test_throws ModelingToolkit.MissingParametersError ODEProblem(
-        sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+    @brownian a b
+    x = _x(t)
+
+    # `System` constructor creates appropriate type with mtkbuild
+    # `Problem` and `alg` create the problem to test and allow calling `init` with
+    # the correct solver.
+    # `rhss` allows adding terms to the end of equations (only 2 equations allowed) to influence
+    # the system type (brownian vars to turn it into an SDE).
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
+        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+        function test_parameter(prob, sym, val, initialval = zero(val))
+            @test prob.ps[sym] ≈ initialval
+            @test init(prob, alg).ps[sym] ≈ val
+            @test solve(prob, alg).ps[sym] ≈ val
+        end
+        function test_initializesystem(sys, u0map, pmap, p, equation)
+            isys = ModelingToolkit.generate_initializesystem(
+                sys; u0map, pmap, guesses = ModelingToolkit.guesses(sys))
+            @test is_variable(isys, p)
+            @test equation in equations(isys) || (0 ~ -equation.rhs) in equations(isys)
+        end
+        u0map = Dict(x => 1.0, y => 1.0)
+        pmap = Dict()
+        pmap[q] = 1.0
+        # `missing` default, equation from Problem
+        @mtkbuild sys = System(
+            [D(x) ~ x * q + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => missing], guesses = [p => 1.0])
+        pmap[p] = 2q
+        prob = Problem(sys, u0map, (0.0, 1.0), pmap)
+        test_parameter(prob, p, 2.0)
+        prob2 = remake(prob; u0 = u0map, p = pmap)
+        prob2.ps[p] = 0.0
+        test_parameter(prob2, p, 2.0)
+        # `missing` default, provided guess
+        @mtkbuild sys = System(
+            [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; defaults = [p => missing], guesses = [p => 0.0])
+        prob = Problem(sys, u0map, (0.0, 1.0))
+        test_parameter(prob, p, 2.0)
+        test_initializesystem(sys, u0map, pmap, p, 0 ~ p - x - y)
+        prob2 = remake(prob; u0 = u0map)
+        prob2.ps[p] = 0.0
+        test_parameter(prob2, p, 2.0)
+
+        # `missing` to Problem, equation from default
+        @mtkbuild sys = System(
+            [D(x) ~ x * q + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => 2q], guesses = [p => 1.0])
+        pmap[p] = missing
+        prob = Problem(sys, u0map, (0.0, 1.0), pmap)
+        test_parameter(prob, p, 2.0)
+        test_initializesystem(sys, u0map, pmap, p, 0 ~ 2q - p)
+        prob2 = remake(prob; u0 = u0map, p = pmap)
+        prob2.ps[p] = 0.0
+        test_parameter(prob2, p, 2.0)
+        # `missing` to Problem, provided guess
+        @mtkbuild sys = System(
+            [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; guesses = [p => 0.0])
+        prob = Problem(sys, u0map, (0.0, 1.0), pmap)
+        test_parameter(prob, p, 2.0)
+        test_initializesystem(sys, u0map, pmap, p, 0 ~ x + y - p)
+        prob2 = remake(prob; u0 = u0map, p = pmap)
+        prob2.ps[p] = 0.0
+        test_parameter(prob2, p, 2.0)
+
+        # No `missing`, default and guess
+        @mtkbuild sys = System(
+            [D(x) ~ x * q + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => 2q], guesses = [p => 0.0])
+        delete!(pmap, p)
+        prob = Problem(sys, u0map, (0.0, 1.0), pmap)
+        test_parameter(prob, p, 2.0)
+        test_initializesystem(sys, u0map, pmap, p, 0 ~ 2q - p)
+        prob2 = remake(prob; u0 = u0map, p = pmap)
+        prob2.ps[p] = 0.0
+        test_parameter(prob2, p, 2.0)
+
+        # Default overridden by Problem, guess provided
+        @mtkbuild sys = System(
+            [D(x) ~ q * x + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => 2q], guesses = [p => 1.0])
+        _pmap = merge(pmap, Dict(p => q))
+        prob = Problem(sys, u0map, (0.0, 1.0), _pmap)
+        test_parameter(prob, p, _pmap[q])
+        test_initializesystem(sys, u0map, _pmap, p, 0 ~ q - p)
+
+        # Problem dependent value with guess, no `missing`
+        @mtkbuild sys = System(
+            [D(x) ~ x * q + rhss[1], D(y) ~ y * p + rhss[2]], t; guesses = [p => 0.0])
+        _pmap = merge(pmap, Dict(p => 3q))
+        prob = Problem(sys, u0map, (0.0, 1.0), _pmap)
+        test_parameter(prob, p, 3pmap[q])
+
+        # Should not be solved for:
+
+        # Override dependent default with direct value
+        @mtkbuild sys = System(
+            [D(x) ~ q * x + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => 2q], guesses = [p => 1.0])
+        _pmap = merge(pmap, Dict(p => 1.0))
+        prob = Problem(sys, u0map, (0.0, 1.0), _pmap)
+        @test prob.ps[p] ≈ 1.0
+        @test prob.f.initialization_data === nothing
+
+        # Non-floating point
+        @parameters r::Int s::Int
+        @mtkbuild sys = System(
+            [D(x) ~ s * x + rhss[1], D(y) ~ y * r + rhss[2]], t; defaults = [s => 2r], guesses = [s => 1.0])
+        prob = Problem(sys, u0map, (0.0, 1.0), [r => 1])
+        @test prob.ps[r] == 1
+        @test prob.ps[s] == 2
+        @test prob.f.initialization_data === nothing
+
+        @mtkbuild sys = System(
+            [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; guesses = [p => 0.0])
+        @test_throws ModelingToolkit.MissingParametersError Problem(
+            sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+    end
 
     @testset "Null system" begin
         @variables x(t) y(t) s(t)
@@ -718,103 +734,153 @@ end
 end
 
 @testset "Update initializeprob parameters" begin
-    @variables x(t) y(t)
+    @variables _x(..) y(t)
     @parameters p q
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x, p ~ x + y], t; guesses = [x => 0.0, p => 0.0])
-    prob = ODEProblem(sys, [y => 1.0], (0.0, 1.0), [p => 3.0])
-    @test prob.f.initializeprob.ps[p] ≈ 3.0
-    @test init(prob, Tsit5())[x] ≈ 2.0
-    prob.ps[p] = 2.0
-    @test prob.f.initializeprob.ps[p] ≈ 3.0
-    @test init(prob, Tsit5())[x] ≈ 1.0
-    ModelingToolkit.defaults(prob.f.sys)[p] = missing
-    prob2 = remake(prob; u0 = [y => 1.0], p = [p => 3x])
-    @test !is_variable(prob2.f.initializeprob, p) &&
-          !is_parameter(prob2.f.initializeprob, p)
-    @test init(prob2, Tsit5())[x] ≈ 0.5
-    @test_nowarn solve(prob2, Tsit5())
+    @brownian a b
+    x = _x(t)
+
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
+        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+        @mtkbuild sys = System(
+            [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; guesses = [x => 0.0, p => 0.0])
+        prob = Problem(sys, [y => 1.0], (0.0, 1.0), [p => 3.0])
+        @test prob.f.initialization_data.initializeprob.ps[p] ≈ 3.0
+        @test init(prob, alg)[x] ≈ 2.0
+        prob.ps[p] = 2.0
+        @test prob.f.initialization_data.initializeprob.ps[p] ≈ 3.0
+        @test init(prob, alg)[x] ≈ 1.0
+        ModelingToolkit.defaults(prob.f.sys)[p] = missing
+        prob2 = remake(prob; u0 = [y => 1.0], p = [p => 3x])
+        @test !is_variable(prob2.f.initialization_data.initializeprob, p) &&
+              !is_parameter(prob2.f.initialization_data.initializeprob, p)
+        @test init(prob2, alg)[x] ≈ 0.5
+        @test_nowarn solve(prob2, alg)
+    end
 end
 
 @testset "Equations for dependent parameters" begin
-    @variables x(t)
+    @variables _x(..)
     @parameters p q=5 r
-    @mtkbuild sys = ODESystem(
-        D(x) ~ 2x + r, t; parameter_dependencies = [r ~ p + 2q, q ~ p + 3],
-        guesses = [p => 1.0])
-    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => missing])
-    @test length(equations(ModelingToolkit.get_parent(prob.f.initializeprob.f.sys))) == 4
-    integ = init(prob, Tsit5())
-    @test integ.ps[p] ≈ 2
+    @brownian a
+    x = _x(t)
+
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
+        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+        @mtkbuild sys = System(
+            [D(x) ~ 2x + r + rhss], t; parameter_dependencies = [r ~ p + 2q, q ~ p + 3],
+            guesses = [p => 1.0])
+        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => missing])
+        @test length(equations(ModelingToolkit.get_parent(prob.f.initialization_data.initializeprob.f.sys))) ==
+              4
+        integ = init(prob, alg)
+        @test integ.ps[p] ≈ 2
+    end
 end
 
 @testset "Re-creating initialization problem on remake" begin
-    @variables x(t) y(t)
+    @variables _x(..) y(t)
     @parameters p q
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x, p ~ x + y], t; defaults = [p => missing], guesses = [x => 0.0, p => 0.0])
-    prob = ODEProblem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
-    @test init(prob, Tsit5()).ps[p] ≈ 2.0
-    # nonsensical value for y just to test that equations work
-    prob2 = remake(prob; u0 = [x => 1.0, y => 2x + exp(t)])
-    @test init(prob2, Tsit5()).ps[p] ≈ 4.0
-    # solve for `x` given `p` and `y`
-    prob3 = remake(prob; u0 = [x => nothing, y => 1.0], p = [p => 2x + exp(t)])
-    @test init(prob3, Tsit5())[x] ≈ 0.0
-    @test_logs (:warn, r"overdetermined") remake(
-        prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
-    prob4 = remake(prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
-    @test solve(prob4, Tsit5()).retcode == ReturnCode.InitialFailure
-    prob5 = remake(prob)
-    @test init(prob, Tsit5()).ps[p] ≈ 2.0
+    @brownian a b
+    x = _x(t)
+
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
+        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+        @mtkbuild sys = System(
+            [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; defaults = [p => missing], guesses = [
+                x => 0.0, p => 0.0])
+        prob = Problem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+        @test init(prob, alg).ps[p] ≈ 2.0
+        # nonsensical value for y just to test that equations work
+        prob2 = remake(prob; u0 = [x => 1.0, y => 2x + exp(t)])
+        @test init(prob2, alg).ps[p] ≈ 4.0
+        # solve for `x` given `p` and `y`
+        prob3 = remake(prob; u0 = [x => nothing, y => 1.0], p = [p => 2x + exp(t)])
+        @test init(prob3, alg)[x] ≈ 0.0
+        @test_logs (:warn, r"overdetermined") remake(
+            prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
+        prob4 = remake(prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
+        @test solve(prob4, alg).retcode == ReturnCode.InitialFailure
+        prob5 = remake(prob)
+        @test init(prob, alg).ps[p] ≈ 2.0
+    end
 end
 
 @testset "`remake` changes initialization problem types" begin
-    @variables x(t) y(t) z(t)
+    @variables _x(..) y(t) z(t)
     @parameters p q
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x * p + y * q, y^2 * q + q^2 * x ~ 0, z * p - p^2 * x * z ~ 0],
-        t; guesses = [x => 0.0, y => 0.0, z => 0.0, p => 0.0, q => 0.0])
-    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0, q => missing])
-    @test is_variable(prob.f.initializeprob, q)
-    ps = prob.p
-    newps = SciMLStructures.replace(Tunable(), ps, ForwardDiff.Dual.(ps.tunable))
-    prob2 = remake(prob; p = newps)
-    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
-    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
-    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
-
-    prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0))
-    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
-    @test eltype(prob2.f.initializeprob.p.tunable) <: Float64
-    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
-
-    prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0), p = newps)
-    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
-    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
-    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
-
-    prob2 = remake(prob; u0 = [x => ForwardDiff.Dual(1.0)],
-        p = [p => ForwardDiff.Dual(1.0), q => missing])
-    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
-    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
-    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
+    @brownian a
+    x = _x(t)
+
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
+        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+        @mtkbuild sys = System(
+            [D(x) ~ x * p + y * q + rhss, y^2 * q + q^2 * x ~ 0, z * p - p^2 * x * z ~ 0],
+            t; guesses = [x => 0.0, y => 0.0, z => 0.0, p => 0.0, q => 0.0])
+        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0, q => missing])
+        @test is_variable(prob.f.initialization_data.initializeprob, q)
+        ps = prob.p
+        newps = SciMLStructures.replace(Tunable(), ps, ForwardDiff.Dual.(ps.tunable))
+        prob2 = remake(prob; p = newps)
+        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
+              ForwardDiff.Dual
+        @test prob2.f.initialization_data.initializeprob.u0 ≈
+              prob.f.initialization_data.initializeprob.u0
+
+        prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0))
+        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <: Float64
+        @test prob2.f.initialization_data.initializeprob.u0 ≈
+              prob.f.initialization_data.initializeprob.u0
+
+        prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0), p = newps)
+        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
+              ForwardDiff.Dual
+        @test prob2.f.initialization_data.initializeprob.u0 ≈
+              prob.f.initialization_data.initializeprob.u0
+
+        prob2 = remake(prob; u0 = [x => ForwardDiff.Dual(1.0)],
+            p = [p => ForwardDiff.Dual(1.0), q => missing])
+        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
+              ForwardDiff.Dual
+        @test prob2.f.initialization_data.initializeprob.u0 ≈
+              prob.f.initialization_data.initializeprob.u0
+    end
 end
 
 @testset "`remake` preserves old u0map and pmap" begin
-    @variables x(t) y(t)
+    @variables _x(..) y(t)
     @parameters p
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ x + p * y, y^2 + 4y * p^2 ~ x], t; guesses = [y => 1.0, p => 1.0])
-    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
-    @test is_variable(prob.f.initializeprob, y)
-    prob2 = @test_nowarn remake(prob; p = [p => 3.0]) # ensure no over/under-determined warning
-    @test is_variable(prob.f.initializeprob, y)
-
-    prob = ODEProblem(sys, [y => 1.0, x => 2.0], (0.0, 1.0), [p => missing])
-    @test is_variable(prob.f.initializeprob, p)
-    prob2 = @test_nowarn remake(prob; u0 = [y => 0.5])
-    @test is_variable(prob.f.initializeprob, p)
+    @brownian a
+    x = _x(t)
+
+    @testset "$Problem" for (Problem, alg, rhss) in [
+        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
+        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+        @mtkbuild sys = System(
+            [D(x) ~ x + p * y + rhss, y^2 + 4y * p^2 ~ x], t; guesses = [
+                y => 1.0, p => 1.0])
+        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
+        @test is_variable(prob.f.initialization_data.initializeprob, y)
+        prob2 = @test_nowarn remake(prob; p = [p => 3.0]) # ensure no over/under-determined warning
+        @test is_variable(prob.f.initialization_data.initializeprob, y)
+
+        prob = Problem(sys, [y => 1.0, x => 2.0], (0.0, 1.0), [p => missing])
+        @test is_variable(prob.f.initialization_data.initializeprob, p)
+        prob2 = @test_nowarn remake(prob; u0 = [y => 0.5])
+        @test is_variable(prob.f.initialization_data.initializeprob, p)
+    end
 end
 
 struct Multiplier{T}

From 39bb59c3ccfcd8efee174829c58b2d546864cea9 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 6 Dec 2024 14:12:35 +0530
Subject: [PATCH 11/18] fix: handle integer `u0` in `DDEProblem`

---
 src/systems/diffeqs/abstractodesystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index b3d4903191..40c44b7bae 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -923,7 +923,7 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     h_oop, h_iip = eval_or_rgf.(h_gen; eval_expression, eval_module)
     h(p, t) = h_oop(p, t)
     h(p::MTKParameters, t) = h_oop(p..., t)
-    u0 = h(p, tspan[1])
+    u0 = float.(h(p, tspan[1]))
     if u0 !== nothing
         u0 = u0_constructor(u0)
     end

From 2f2e62524f42d521e69b692b2801d5ca5e90cfc1 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 6 Dec 2024 14:13:00 +0530
Subject: [PATCH 12/18] feat: enable creating `InitializationProblem` for
 non-`AbstractODESystem`s

---
 src/systems/diffeqs/abstractodesystem.jl | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 40c44b7bae..a6607e0c55 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1247,11 +1247,11 @@ Generates a NonlinearProblem or NonlinearLeastSquaresProblem from an ODESystem
 which represents the initialization, i.e. the calculation of the consistent
 initial conditions for the given DAE.
 """
-function InitializationProblem(sys::AbstractODESystem, args...; kwargs...)
+function InitializationProblem(sys::AbstractSystem, args...; kwargs...)
     InitializationProblem{true}(sys, args...; kwargs...)
 end
 
-function InitializationProblem(sys::AbstractODESystem, t,
+function InitializationProblem(sys::AbstractSystem, t,
         u0map::StaticArray,
         args...;
         kwargs...)
@@ -1259,11 +1259,11 @@ function InitializationProblem(sys::AbstractODESystem, t,
         sys, t, u0map, args...; kwargs...)
 end
 
-function InitializationProblem{true}(sys::AbstractODESystem, args...; kwargs...)
+function InitializationProblem{true}(sys::AbstractSystem, args...; kwargs...)
     InitializationProblem{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
 end
 
-function InitializationProblem{false}(sys::AbstractODESystem, args...; kwargs...)
+function InitializationProblem{false}(sys::AbstractSystem, args...; kwargs...)
     InitializationProblem{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 
@@ -1282,8 +1282,8 @@ function Base.showerror(io::IO, e::IncompleteInitializationError)
     println(io, e.uninit)
 end
 
-function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
-        t::Number, u0map = [],
+function InitializationProblem{iip, specialize}(sys::AbstractSystem,
+        t, u0map = [],
         parammap = DiffEqBase.NullParameters();
         guesses = [],
         check_length = true,
@@ -1347,13 +1347,13 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         @warn "Initialization system is underdetermined. $neqs equations for $nunknown unknowns. Initialization will default to using least squares. $(scc_message)To suppress this warning pass warn_initialize_determined = false. To make this warning into an error, pass fully_determined = true"
     end
 
-    parammap = parammap isa DiffEqBase.NullParameters || isempty(parammap) ?
-               [get_iv(sys) => t] :
-               merge(todict(parammap), Dict(get_iv(sys) => t))
-    parammap = Dict(k => v for (k, v) in parammap if v !== missing)
-    if isempty(u0map)
-        u0map = Dict()
+    parammap = recursive_unwrap(anydict(parammap))
+    if t !== nothing
+        parammap[get_iv(sys)] = t
     end
+    filter!(kvp -> kvp[2] !== missing, parammap)
+
+    u0map = to_varmap(u0map, unknowns(sys))
     if isempty(guesses)
         guesses = Dict()
     end

From 671b93feea590a1a7756fd469f9426cc6105766d Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 6 Dec 2024 14:13:10 +0530
Subject: [PATCH 13/18] fix: filter kwargs in `SDEProblem`

---
 src/systems/diffeqs/sdesystem.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 5e0d0e3208..0d600cfb5f 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -760,6 +760,8 @@ function DiffEqBase.SDEProblem{iip, specialize}(
         noise = nothing
     end
 
+    kwargs = filter_kwargs(kwargs)
+
     SDEProblem{iip}(f, u0, tspan, p; callback = cbs, noise,
         noise_rate_prototype = noise_rate_prototype, kwargs...)
 end

From 9987da0b19e5e2c41aed9af214a215e7dc9b714a Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 6 Dec 2024 14:15:05 +0530
Subject: [PATCH 14/18] test: test initialization on `NonlinearProblem` and
 `NonlinearLeastSquaresProblem`

---
 test/initializationsystem.jl | 269 ++++++++++++++++++++++++++++-------
 1 file changed, 218 insertions(+), 51 deletions(-)

diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 045b7d2d2c..2745b9d81f 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -583,6 +583,14 @@ sol = solve(oprob_2nd_order_2, Rosenbrock23()) # retcode: Success
     @test all(sol(1.0, idxs = sys.x) .≈ +exp(1)) && all(sol(1.0, idxs = sys.y) .≈ -exp(1))
 end
 
+NonlinearSystemWrapper(eqs, t; kws...) = NonlinearSystem(eqs; kws...)
+function NonlinearProblemWrapper(sys, u0, tspan, args...; kwargs...)
+    NonlinearProblem(sys, u0, args...; kwargs...)
+end
+function NLLSProblemWrapper(sys, u0, tspan, args...; kwargs...)
+    NonlinearLeastSquaresProblem(sys, u0, args...; kwargs...)
+end
+
 @testset "Initialization of parameters" begin
     @variables _x(..) y(t)
     @parameters p q
@@ -594,13 +602,37 @@ end
     # the correct solver.
     # `rhss` allows adding terms to the end of equations (only 2 equations allowed) to influence
     # the system type (brownian vars to turn it into an SDE).
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
-        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
-        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+    @testset "$Problem with $(SciMLBase.parameterless_type(alg))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), zeros(2)),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), [a, b]),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b]),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), zeros(2)),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), zeros(2)),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), zeros(2)),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), zeros(2))
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+
         function test_parameter(prob, sym, val, initialval = zero(val))
             @test prob.ps[sym] ≈ initialval
-            @test init(prob, alg).ps[sym] ≈ val
+            if !is_nlsolve || prob.u0 !== nothing
+                @test init(prob, alg).ps[sym] ≈ val
+            end
             @test solve(prob, alg).ps[sym] ≈ val
         end
         function test_initializesystem(sys, u0map, pmap, p, equation)
@@ -609,6 +641,8 @@ end
             @test is_variable(isys, p)
             @test equation in equations(isys) || (0 ~ -equation.rhs) in equations(isys)
         end
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+
         u0map = Dict(x => 1.0, y => 1.0)
         pmap = Dict()
         pmap[q] = 1.0
@@ -669,16 +703,14 @@ end
         prob = Problem(sys, u0map, (0.0, 1.0), _pmap)
         test_parameter(prob, p, _pmap[q])
         test_initializesystem(sys, u0map, _pmap, p, 0 ~ q - p)
-
         # Problem dependent value with guess, no `missing`
         @mtkbuild sys = System(
-            [D(x) ~ x * q + rhss[1], D(y) ~ y * p + rhss[2]], t; guesses = [p => 0.0])
+            [D(x) ~ y * q + p + rhss[1], D(y) ~ x * p + q + rhss[2]], t; guesses = [p => 0.0])
         _pmap = merge(pmap, Dict(p => 3q))
         prob = Problem(sys, u0map, (0.0, 1.0), _pmap)
         test_parameter(prob, p, 3pmap[q])
 
         # Should not be solved for:
-
         # Override dependent default with direct value
         @mtkbuild sys = System(
             [D(x) ~ q * x + rhss[1], D(y) ~ y * p + rhss[2]], t; defaults = [p => 2q], guesses = [p => 1.0])
@@ -700,6 +732,14 @@ end
             [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; guesses = [p => 0.0])
         @test_throws ModelingToolkit.MissingParametersError Problem(
             sys, [x => 1.0, y => 1.0], (0.0, 1.0))
+
+        # Unsatisfiable initialization
+        prob = Problem(sys, [x => 1.0, y => 1.0], (0.0, 1.0),
+            [p => 2.0]; initialization_eqs = [x^2 + y^2 ~ 3])
+        @test prob.f.initialization_data !== nothing
+        @test solve(prob, alg).retcode == ReturnCode.InitialFailure
+        cache = init(prob, alg)
+        @test solve!(cache).retcode == ReturnCode.InitialFailure
     end
 
     @testset "Null system" begin
@@ -707,7 +747,9 @@ end
         @parameters x0 y0
         @mtkbuild sys = ODESystem([x ~ x0, y ~ y0, s ~ x + y], t; guesses = [y0 => 0.0])
         prob = ODEProblem(sys, [s => 1.0], (0.0, 1.0), [x0 => 0.3, y0 => missing])
-        test_parameter(prob, y0, 0.7)
+        @test prob.ps[y0] ≈ 0.0
+        @test init(prob, Tsit5()).ps[y0] ≈ 0.7
+        @test solve(prob, Tsit5()).ps[y0] ≈ 0.7
     end
 
     using ModelingToolkitStandardLibrary.Mechanical.TranslationalModelica: Fixed, Mass,
@@ -730,7 +772,9 @@ end
         systems = [fixed, spring, mass, gravity, constant, damper],
         guesses = [spring.s_rel0 => 1.0])
     prob = ODEProblem(sys, [], (0.0, 1.0), [spring.s_rel0 => missing])
-    test_parameter(prob, spring.s_rel0, -3.905)
+    @test prob.ps[spring.s_rel0] ≈ 0.0
+    @test init(prob, Tsit5()).ps[spring.s_rel0] ≈ -3.905
+    @test solve(prob, Tsit5()).ps[spring.s_rel0] ≈ -3.905
 end
 
 @testset "Update initializeprob parameters" begin
@@ -739,10 +783,33 @@ end
     @brownian a b
     x = _x(t)
 
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
-        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
-        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+    @testset "$Problem with $(SciMLBase.parameterless_type(typeof(alg)))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), zeros(2)),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), [a, b]),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b]),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), zeros(2)),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), zeros(2)),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), zeros(2)),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), zeros(2))
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+
         @mtkbuild sys = System(
             [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; guesses = [x => 0.0, p => 0.0])
         prob = Problem(sys, [y => 1.0], (0.0, 1.0), [p => 3.0])
@@ -766,10 +833,33 @@ end
     @brownian a
     x = _x(t)
 
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
-        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
-        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+    @testset "$Problem with $(SciMLBase.parameterless_type(typeof(alg)))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), 0),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), a),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), _x(t - 0.1) + a),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), 0),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), 0),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), 0),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), 0)
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+
         @mtkbuild sys = System(
             [D(x) ~ 2x + r + rhss], t; parameter_dependencies = [r ~ p + 2q, q ~ p + 3],
             guesses = [p => 1.0])
@@ -787,21 +877,44 @@ end
     @brownian a b
     x = _x(t)
 
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), zeros(2)), (SDEProblem, ImplicitEM(), [a, b]),
-        (DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
-        (SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b])]
+    @testset "$Problem with $(SciMLBase.parameterless_type(typeof(alg)))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), zeros(2)),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), [a, b]),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), [_x(t - 0.1), 0.0]),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), [_x(t - 0.1) + a, b]),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), zeros(2)),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), zeros(2)),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), zeros(2)),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), zeros(2)),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), zeros(2)),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), zeros(2))
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+
         @mtkbuild sys = System(
             [D(x) ~ x + rhss[1], p ~ x + y + rhss[2]], t; defaults = [p => missing], guesses = [
                 x => 0.0, p => 0.0])
         prob = Problem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
         @test init(prob, alg).ps[p] ≈ 2.0
         # nonsensical value for y just to test that equations work
-        prob2 = remake(prob; u0 = [x => 1.0, y => 2x + exp(t)])
-        @test init(prob2, alg).ps[p] ≈ 4.0
+        prob2 = remake(prob; u0 = [x => 1.0, y => 2x + exp(x)])
+        @test init(prob2, alg).ps[p] ≈ 3 + exp(1)
         # solve for `x` given `p` and `y`
-        prob3 = remake(prob; u0 = [x => nothing, y => 1.0], p = [p => 2x + exp(t)])
-        @test init(prob3, alg)[x] ≈ 0.0
+        prob3 = remake(prob; u0 = [x => nothing, y => 1.0], p = [p => 2x + exp(y)])
+        @test init(prob3, alg)[x] ≈ 1 - exp(1)
         @test_logs (:warn, r"overdetermined") remake(
             prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
         prob4 = remake(prob; u0 = [x => 1.0, y => 2.0], p = [p => 4.0])
@@ -817,44 +930,73 @@ end
     @brownian a
     x = _x(t)
 
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
-        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
-        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+    @testset "$Problem with $(SciMLBase.parameterless_type(typeof(alg)))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), 0),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), a),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), _x(t - 0.1) + a),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), 0),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), 0),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), 0),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), 0)
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+        alge_eqs = [y^2 * q + q^2 * x ~ 0, z * p - p^2 * x * z ~ 0]
+
         @mtkbuild sys = System(
-            [D(x) ~ x * p + y * q + rhss, y^2 * q + q^2 * x ~ 0, z * p - p^2 * x * z ~ 0],
+            [D(x) ~ x * p + y^2 * q + rhss; alge_eqs],
             t; guesses = [x => 0.0, y => 0.0, z => 0.0, p => 0.0, q => 0.0])
-        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0, q => missing])
+        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0, q => missing];
+            initialization_eqs = is_nlsolve ? alge_eqs : [])
         @test is_variable(prob.f.initialization_data.initializeprob, q)
         ps = prob.p
         newps = SciMLStructures.replace(Tunable(), ps, ForwardDiff.Dual.(ps.tunable))
         prob2 = remake(prob; p = newps)
-        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(state_values(prob2.f.initialization_data.initializeprob)) <:
+              ForwardDiff.Dual
         @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
               ForwardDiff.Dual
-        @test prob2.f.initialization_data.initializeprob.u0 ≈
-              prob.f.initialization_data.initializeprob.u0
+        @test state_values(prob2.f.initialization_data.initializeprob) ≈
+              state_values(prob.f.initialization_data.initializeprob)
 
         prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0))
-        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(state_values(prob2.f.initialization_data.initializeprob)) <:
+              ForwardDiff.Dual
         @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <: Float64
-        @test prob2.f.initialization_data.initializeprob.u0 ≈
-              prob.f.initialization_data.initializeprob.u0
+        @test state_values(prob2.f.initialization_data.initializeprob) ≈
+              state_values(prob.f.initialization_data.initializeprob)
 
         prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0), p = newps)
-        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(state_values(prob2.f.initialization_data.initializeprob)) <:
+              ForwardDiff.Dual
         @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
               ForwardDiff.Dual
-        @test prob2.f.initialization_data.initializeprob.u0 ≈
-              prob.f.initialization_data.initializeprob.u0
+        @test state_values(prob2.f.initialization_data.initializeprob) ≈
+              state_values(prob.f.initialization_data.initializeprob)
 
         prob2 = remake(prob; u0 = [x => ForwardDiff.Dual(1.0)],
             p = [p => ForwardDiff.Dual(1.0), q => missing])
-        @test eltype(prob2.f.initialization_data.initializeprob.u0) <: ForwardDiff.Dual
+        @test eltype(state_values(prob2.f.initialization_data.initializeprob)) <:
+              ForwardDiff.Dual
         @test eltype(prob2.f.initialization_data.initializeprob.p.tunable) <:
               ForwardDiff.Dual
-        @test prob2.f.initialization_data.initializeprob.u0 ≈
-              prob.f.initialization_data.initializeprob.u0
+        @test state_values(prob2.f.initialization_data.initializeprob) ≈
+              state_values(prob.f.initialization_data.initializeprob)
     end
 end
 
@@ -864,19 +1006,44 @@ end
     @brownian a
     x = _x(t)
 
-    @testset "$Problem" for (Problem, alg, rhss) in [
-        (ODEProblem, Tsit5(), 0), (SDEProblem, ImplicitEM(), a),
-        (DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
-        (SDDEProblem, ImplicitEM(), _x(t - 0.1) + a)]
+    @testset "$Problem with $(SciMLBase.parameterless_type(typeof(alg)))" for (System, Problem, alg, rhss) in [
+        (ModelingToolkit.System, ODEProblem, Tsit5(), 0),
+        (ModelingToolkit.System, SDEProblem, ImplicitEM(), a),
+        (ModelingToolkit.System, DDEProblem, MethodOfSteps(Tsit5()), _x(t - 0.1)),
+        (ModelingToolkit.System, SDDEProblem, ImplicitEM(), _x(t - 0.1) + a),
+        # polyalg cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper,
+            FastShortcutNonlinearPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, NewtonRaphson(), 0),
+        # quasi newton cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, Klement(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, SimpleNewtonRaphson(), 0),
+        # DFSane cache
+        (NonlinearSystemWrapper, NonlinearProblemWrapper, DFSane(), 0),
+        # Least squares
+        # polyalg cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, FastShortcutNLLSPolyalg(), 0),
+        # generalized first order cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, LevenbergMarquardt(), 0),
+        # noinit cache
+        (NonlinearSystemWrapper, NLLSProblemWrapper, SimpleGaussNewton(), 0)
+    ]
+        is_nlsolve = alg isa SciMLBase.AbstractNonlinearAlgorithm
+        D = is_nlsolve ? v -> v^3 : Differential(t)
+        alge_eqs = [y^2 + 4y * p^2 ~ x^3]
         @mtkbuild sys = System(
-            [D(x) ~ x + p * y + rhss, y^2 + 4y * p^2 ~ x], t; guesses = [
+            [D(x) ~ x + p * y^2 + rhss; alge_eqs], t; guesses = [
                 y => 1.0, p => 1.0])
-        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
+        prob = Problem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0];
+            initialization_eqs = is_nlsolve ? alge_eqs : [])
         @test is_variable(prob.f.initialization_data.initializeprob, y)
         prob2 = @test_nowarn remake(prob; p = [p => 3.0]) # ensure no over/under-determined warning
         @test is_variable(prob.f.initialization_data.initializeprob, y)
 
-        prob = Problem(sys, [y => 1.0, x => 2.0], (0.0, 1.0), [p => missing])
+        prob = Problem(sys, [y => 1.0, x => 2.0], (0.0, 1.0), [p => missing];
+            initialization_eqs = is_nlsolve ? alge_eqs : [])
         @test is_variable(prob.f.initialization_data.initializeprob, p)
         prob2 = @test_nowarn remake(prob; u0 = [y => 0.5])
         @test is_variable(prob.f.initialization_data.initializeprob, p)

From 51eeeeb27ee5094424bd92eb589f8b99f843c592 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 6 Dec 2024 17:43:57 +0530
Subject: [PATCH 15/18] fix: store and propagate `initialization_eqs` provided
 to Problem

---
 src/systems/nonlinear/initializesystem.jl | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index e32f967f4a..cf6c6d4d42 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -21,6 +21,7 @@ function generate_initializesystem(sys::AbstractSystem;
     eqs_ics = Equation[]
     defs = copy(defaults(sys)) # copy so we don't modify sys.defaults
     additional_guesses = anydict(guesses)
+    additional_initialization_eqs = Vector{Equation}(initialization_eqs)
     guesses = merge(get_guesses(sys), additional_guesses)
     idxs_diff = isdiffeq.(eqs)
 
@@ -191,7 +192,7 @@ function generate_initializesystem(sys::AbstractSystem;
         defs[k] = substitute(defs[k], paramsubs)
     end
     meta = InitializationSystemMetadata(
-        anydict(u0map), anydict(pmap), additional_guesses, extra_metadata)
+        anydict(u0map), anydict(pmap), additional_guesses, additional_initialization_eqs, extra_metadata)
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
@@ -207,6 +208,7 @@ struct InitializationSystemMetadata
     u0map::Dict{Any, Any}
     pmap::Dict{Any, Any}
     additional_guesses::Dict{Any, Any}
+    additional_initialization_eqs::Vector{Equation}
     extra_metadata::NamedTuple
 end
 
@@ -299,6 +301,7 @@ function SciMLBase.remake_initialization_data(
     defs = defaults(sys)
     cmap, cs = get_cmap(sys)
     use_scc = true
+    initialization_eqs = Equation[]
 
     if SciMLBase.has_initializeprob(odefn)
         oldsys = odefn.initialization_data.initializeprob.f.sys
@@ -308,6 +311,7 @@ function SciMLBase.remake_initialization_data(
             pmap = merge(meta.pmap, pmap)
             merge!(guesses, meta.additional_guesses)
             use_scc = get(meta.extra_metadata, :use_scc, true)
+            initialization_eqs = meta.additional_initialization_eqs
         end
     else
         # there is no initializeprob, so the original problem construction
@@ -348,7 +352,7 @@ function SciMLBase.remake_initialization_data(
     op, missing_unknowns, missing_pars = build_operating_point(
         u0map, pmap, defs, cmap, dvs, ps)
     kws = maybe_build_initialization_problem(
-        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns; use_scc)
+        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns; use_scc, initialization_eqs)
     return get(kws, :initialization_data, nothing)
 end
 

From 96f8d5de87813b300b61e3d8359b7231ec661529 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sat, 14 Dec 2024 13:18:38 +0530
Subject: [PATCH 16/18] build: bump compats

---
 Project.toml | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3c6126b7c0..d672c007c2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -89,6 +89,7 @@ ConstructionBase = "1"
 DataInterpolations = "6.4"
 DataStructures = "0.17, 0.18"
 DeepDiffs = "1"
+DelayDiffEq = "5.50"
 DiffEqBase = "6.157"
 DiffEqCallbacks = "2.16, 3, 4"
 DiffEqNoiseProcess = "5"
@@ -117,7 +118,7 @@ Libdl = "1"
 LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
-NonlinearSolve = "3.14, 4"
+NonlinearSolve = "4.3"
 OffsetArrays = "1"
 OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
@@ -129,7 +130,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
-SciMLBase = "2.66"
+SciMLBase = "2.68.1"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
@@ -137,7 +138,9 @@ SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
-SymbolicIndexingInterface = "0.3.35"
+StochasticDiffEq = "6.72.1"
+StochasticDelayDiffEq = "1.8.1"
+SymbolicIndexingInterface = "0.3.36"
 SymbolicUtils = "3.7"
 Symbolics = "6.19"
 URIs = "1"

From 0a881e71d5f1b5fe77e170713e78b5a5ab3ec023 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 16 Dec 2024 16:22:15 +0530
Subject: [PATCH 17/18] fix: better handle reconstructing initializeprob with
 new types

---
 src/systems/diffeqs/abstractodesystem.jl  |  5 ++
 src/systems/nonlinear/initializesystem.jl | 69 ++++++++++-------------
 2 files changed, 34 insertions(+), 40 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index a6607e0c55..dd4165cb57 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1310,6 +1310,11 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
                 pmap = parammap, guesses, extra_metadata = (; use_scc)); fully_determined)
     end
 
+    meta = get_metadata(isys)
+    if meta isa InitializationSystemMetadata
+        @set! isys.metadata.oop_reconstruct_u0_p = ReconstructInitializeprob(sys, isys)
+    end
+
     ts = get_tearing_state(isys)
     unassigned_vars = StructuralTransformations.singular_check(ts)
     if warn_initialize_determined && !isempty(unassigned_vars)
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index cf6c6d4d42..602fc7cac9 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -192,7 +192,8 @@ function generate_initializesystem(sys::AbstractSystem;
         defs[k] = substitute(defs[k], paramsubs)
     end
     meta = InitializationSystemMetadata(
-        anydict(u0map), anydict(pmap), additional_guesses, additional_initialization_eqs, extra_metadata)
+        anydict(u0map), anydict(pmap), additional_guesses,
+        additional_initialization_eqs, extra_metadata, nothing)
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
@@ -204,12 +205,30 @@ function generate_initializesystem(sys::AbstractSystem;
         kwargs...)
 end
 
+struct ReconstructInitializeprob
+    getter::Any
+    setter::Any
+end
+
+function ReconstructInitializeprob(srcsys::AbstractSystem, dstsys::AbstractSystem)
+    syms = [unknowns(dstsys);
+            reduce(vcat, reorder_parameters(dstsys, parameters(dstsys)); init = [])]
+    getter = getu(srcsys, syms)
+    setter = setsym_oop(dstsys, syms)
+    return ReconstructInitializeprob(getter, setter)
+end
+
+function (rip::ReconstructInitializeprob)(srcvalp, dstvalp)
+    rip.setter(dstvalp, rip.getter(srcvalp))
+end
+
 struct InitializationSystemMetadata
     u0map::Dict{Any, Any}
     pmap::Dict{Any, Any}
     additional_guesses::Dict{Any, Any}
     additional_initialization_eqs::Vector{Equation}
     extra_metadata::NamedTuple
+    oop_reconstruct_u0_p::Union{Nothing, ReconstructInitializeprob}
 end
 
 function is_parameter_solvable(p, pmap, defs, guesses)
@@ -239,45 +258,15 @@ function SciMLBase.remake_initialization_data(
         if !SciMLBase.has_sys(oldinitprob.f) || !(oldinitprob.f.sys isa NonlinearSystem)
             return oldinitdata
         end
-        pidxs = ParameterIndex[]
-        pvals = []
-        u0idxs = Int[]
-        u0vals = []
-        for sym in variable_symbols(oldinitprob)
-            if is_variable(sys, sym) || has_observed_with_lhs(sys, sym)
-                u0 !== missing || continue
-                idx = variable_index(oldinitprob, sym)
-                push!(u0idxs, idx)
-                push!(u0vals, eltype(u0)(state_values(oldinitprob, idx)))
-            else
-                p !== missing || continue
-                idx = variable_index(oldinitprob, sym)
-                push!(u0idxs, idx)
-                push!(u0vals, typeof(getp(sys, sym)(p))(state_values(oldinitprob, idx)))
-            end
-        end
-        if p !== missing
-            for sym in parameter_symbols(oldinitprob)
-                push!(pidxs, parameter_index(oldinitprob, sym))
-                if is_time_dependent(sys) && isequal(sym, get_iv(sys))
-                    push!(pvals, t0)
-                else
-                    push!(pvals, getp(sys, sym)(p))
-                end
-            end
-        end
-        if isempty(u0idxs)
-            newu0 = state_values(oldinitprob)
-        else
-            newu0 = remake_buffer(
-                oldinitprob.f.sys, state_values(oldinitprob), u0idxs, u0vals)
-        end
-        if isempty(pidxs)
-            newp = parameter_values(oldinitprob)
+        oldinitsys = oldinitprob.f.sys
+        meta = get_metadata(oldinitsys)
+        if meta isa InitializationSystemMetadata && meta.oop_reconstruct_u0_p !== nothing
+            reconstruct_fn = meta.oop_reconstruct_u0_p
         else
-            newp = remake_buffer(
-                oldinitprob.f.sys, parameter_values(oldinitprob), pidxs, pvals)
+            reconstruct_fn = ReconstructInitializeprob(sys, oldinitsys)
         end
+        new_initu0, new_initp = reconstruct_fn(
+            ProblemState(; u = newu0, p = newp, t = t0), oldinitprob)
         if oldinitprob.f.resid_prototype === nothing
             newf = oldinitprob.f
         else
@@ -285,9 +274,9 @@ function SciMLBase.remake_initialization_data(
                 SciMLBase.isinplace(oldinitprob.f), SciMLBase.specialization(oldinitprob.f)}(
                 oldinitprob.f;
                 resid_prototype = calculate_resid_prototype(
-                    length(oldinitprob.f.resid_prototype), newu0, newp))
+                    length(oldinitprob.f.resid_prototype), new_initu0, new_initp))
         end
-        initprob = remake(oldinitprob; f = newf, u0 = newu0, p = newp)
+        initprob = remake(oldinitprob; f = newf, u0 = new_initu0, p = new_initp)
         return SciMLBase.OverrideInitData(initprob, oldinitdata.update_initializeprob!,
             oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap)
     end

From 2e07200e03af428b127814af85f48fdeccdaf574 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 16 Dec 2024 16:22:51 +0530
Subject: [PATCH 18/18] test: fix incorrect initial values in tests

---
 test/nonlinearsystem.jl | 2 +-
 test/reduction.jl       | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/test/nonlinearsystem.jl b/test/nonlinearsystem.jl
index f766ca0131..cbd95a50d3 100644
--- a/test/nonlinearsystem.jl
+++ b/test/nonlinearsystem.jl
@@ -293,7 +293,7 @@ sys = structural_simplify(ns; conservative = true)
     eqs = [0 ~ σ * (y - x)
            0 ~ x * (ρ - z) - y
            0 ~ x * y - β * z]
-    guesses = [x => 1.0, y => 0.0, z => 0.0]
+    guesses = [x => 1.0, z => 0.0]
     ps = [σ => 10.0, ρ => 26.0, β => 8 / 3]
     @mtkbuild ns = NonlinearSystem(eqs)
 
diff --git a/test/reduction.jl b/test/reduction.jl
index 6d7a05b99e..fa9029a652 100644
--- a/test/reduction.jl
+++ b/test/reduction.jl
@@ -158,9 +158,7 @@ eqs = [u1 ~ u2
 reducedsys = structural_simplify(sys)
 @test length(observed(reducedsys)) == 2
 
-u0 = [u1 => 1
-      u2 => 1
-      u3 => 0.3]
+u0 = [u2 => 1]
 pp = [2]
 nlprob = NonlinearProblem(reducedsys, u0, [p => pp[1]])
 reducedsol = solve(nlprob, NewtonRaphson())