Skip to content

Commit

Permalink
fix: store and propagate initialization_eqs provided to Problem
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Dec 24, 2024
1 parent 9987da0 commit 51eeeeb
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/systems/nonlinear/initializesystem.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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;
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 51eeeeb

Please sign in to comment.