Merge pull request #3769 from AayushSabharwal/as/substitute-retain-data

fix: retain events and metadata in `substitute`

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -2683,22 +2683,27 @@ function Symbolics.substitute(sys::AbstractSystem, rules::Union{Vector{<:Pair},
     elseif sys isa System
         rules = todict(map(r -> Symbolics.unwrap(r[1]) => Symbolics.unwrap(r[2]),
             collect(rules)))
-        eqs = fast_substitute(get_eqs(sys), rules)
-        pdeps = fast_substitute(get_parameter_dependencies(sys), rules)
-        defs = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
+        newsys = @set sys.eqs = fast_substitute(get_eqs(sys), rules)
+        @set! newsys.unknowns = map(get_unknowns(sys)) do var
+            get(rules, var, var)
+        end
+        @set! newsys.ps = map(get_ps(sys)) do var
+            get(rules, var, var)
+        end
+        @set! newsys.parameter_dependencies = fast_substitute(
+            get_parameter_dependencies(sys), rules)
+        @set! newsys.defaults = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
         for (k, v) in get_defaults(sys))
-        guess = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
+        @set! newsys.guesses = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
         for (k, v) in get_guesses(sys))
-        noise_eqs = fast_substitute(get_noise_eqs(sys), rules)
-        costs = fast_substitute(get_costs(sys), rules)
-        observed = fast_substitute(get_observed(sys), rules)
-        initialization_eqs = fast_substitute(get_initialization_eqs(sys), rules)
-        cstrs = fast_substitute(get_constraints(sys), rules)
-        subsys = map(s -> substitute(s, rules), get_systems(sys))
-        newsys = System(eqs, get_iv(sys); name = nameof(sys), defaults = defs,
-            guesses = guess, systems = subsys, noise_eqs,
-            observed, initialization_eqs, constraints = cstrs)
-        @set! newsys.parameter_dependencies = pdeps
+        @set! newsys.noise_eqs = fast_substitute(get_noise_eqs(sys), rules)
+        @set! newsys.costs = Vector{Union{Real, BasicSymbolic}}(fast_substitute(
+            get_costs(sys), rules))
+        @set! newsys.observed = fast_substitute(get_observed(sys), rules)
+        @set! newsys.initialization_eqs = fast_substitute(
+            get_initialization_eqs(sys), rules)
+        @set! newsys.constraints = fast_substitute(get_constraints(sys), rules)
+        @set! newsys.systems = map(s -> substitute(s, rules), get_systems(sys))
     else
         error("substituting symbols is not supported for $(typeof(sys))")
     end

test/odesystem.jl

@@ -1574,3 +1574,28 @@ end
     prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0))
     @test prob.problem_type == "A"
 end
+
+@testset "`substitute` retains events and metadata" begin
+    @parameters p(t) = 1.0
+    @variables x(t) = 0.0
+    event = [0.5] => [p ~ Pre(t)]
+    event2 = [x ~ 0.75] => [p ~ 2 * Pre(t)]
+
+    struct TestMeta end
+
+    eq = [
+        D(x) ~ p
+    ]
+    @named sys = System(eq, t, [x], [p], discrete_events = [event],
+        continuous_events = [event2], metadata = Dict(TestMeta => "test"))
+
+    @variables x2(t) = 0.0
+    sys2 = substitute(sys, [x => x2])
+
+    @test length(ModelingToolkit.get_discrete_events(sys)) == 1
+    @test length(ModelingToolkit.get_discrete_events(sys2)) == 1
+    @test length(ModelingToolkit.get_continuous_events(sys)) == 1
+    @test length(ModelingToolkit.get_continuous_events(sys2)) == 1
+    @test getmetadata(sys, TestMeta, nothing) == "test"
+    @test getmetadata(sys2, TestMeta, nothing) == "test"
+end