Merge pull request #2846 from hersle/extend_initialization_equations

Extend initialization equations

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -2517,6 +2517,7 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nam
         end
     end
 
+    # collect fields common to all system types
     eqs = union(get_eqs(basesys), get_eqs(sys))
     sts = union(get_unknowns(basesys), get_unknowns(sys))
     ps = union(get_ps(basesys), get_ps(sys))
@@ -2529,16 +2530,19 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nam
     devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
     syss = union(get_systems(basesys), get_systems(sys))
+    args = length(ivs) == 0 ? (eqs, sts, ps) : (eqs, ivs[1], sts, ps)
+    kwargs = (parameter_dependencies = dep_ps, observed = obs, continuous_events = cevs,
+        discrete_events = devs, defaults = defs, systems = syss,
+        name = name, gui_metadata = gui_metadata)
 
-    if length(ivs) == 0
-        T(eqs, sts, ps, observed = obs, defaults = defs, name = name, systems = syss,
-            continuous_events = cevs, discrete_events = devs, gui_metadata = gui_metadata,
-            parameter_dependencies = dep_ps)
-    elseif length(ivs) == 1
-        T(eqs, ivs[1], sts, ps, observed = obs, defaults = defs, name = name,
-            systems = syss, continuous_events = cevs, discrete_events = devs,
-            gui_metadata = gui_metadata, parameter_dependencies = dep_ps)
+    # collect fields specific to some system types
+    if basesys isa ODESystem
+        ieqs = union(get_initialization_eqs(basesys), get_initialization_eqs(sys))
+        guesses = merge(get_guesses(basesys), get_guesses(sys)) # prefer `sys`
+        kwargs = merge(kwargs, (initialization_eqs = ieqs, guesses = guesses))
     end
+
+    return T(args...; kwargs...)
 end
 
 function Base.:(&)(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nameof(sys))

test/initializationsystem.jl

@@ -447,3 +447,12 @@ prob = ODEProblem(pend, [x => 1], (0.0, 1.5), [g => 1],
 unsimp = generate_initializesystem(pend; u0map = [x => 1], initialization_eqs = [y ~ 1])
 sys = structural_simplify(unsimp; fully_determined = false)
 @test length(equations(sys)) == 3
+
+# Extend two systems with initialization equations and guesses
+# https://github.com/SciML/ModelingToolkit.jl/issues/2845
+@variables x(t) y(t)
+@named sysx = ODESystem([D(x) ~ 0], t; initialization_eqs = [x ~ 1])
+@named sysy = ODESystem([D(y) ~ 0], t; initialization_eqs = [y^2 ~ 2], guesses = [y => 1])
+sys = extend(sysx, sysy)
+@test length(equations(generate_initializesystem(sys))) == 2
+@test length(ModelingToolkit.guesses(sys)) == 1