Extend initialization only for ODE systems

Author: hersle

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))
@@ -2528,20 +2529,20 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nam
     cevs = union(get_continuous_events(basesys), get_continuous_events(sys))
     devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
-    guesses = merge(get_guesses(basesys), get_guesses(sys)) # prefer `sys`
-    ieqs = union(get_initialization_eqs(basesys), get_initialization_eqs(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, initialization_eqs = ieqs, guesses = guesses)
-    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,
-            initialization_eqs = ieqs, guesses = guesses)
+    # 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))