feat: add Symbolics.fast_substitute for affects

Author: AayushSabharwal

src/systems/callbacks.jl

@@ -25,6 +25,12 @@ function SymbolicAffect(affect::SymbolicAffect; kwargs...)
 end
 SymbolicAffect(affect; kwargs...) = make_affect(affect; kwargs...)
 
+function Symbolics.fast_substitute(aff::SymbolicAffect, rules)
+    substituter = Base.Fix2(fast_substitute, rules)
+    SymbolicAffect(map(substituter, aff.affect), map(substituter, aff.alg_eqs),
+        map(substituter, aff.discrete_parameters))
+end
+
 struct AffectSystem
     """The internal implicit discrete system whose equations are solved to obtain values after the affect."""
     system::AbstractSystem
@@ -36,6 +42,19 @@ struct AffectSystem
     discretes::Vector
 end
 
+function Symbolics.fast_substitute(aff::AffectSystem, rules)
+    substituter = Base.Fix2(fast_substitute, rules)
+    sys = aff.system
+    @set! sys.eqs = map(substituter, get_eqs(sys))
+    @set! sys.parameter_dependencies = map(substituter, get_parameter_dependencies(sys))
+    @set! sys.defaults = Dict([k => substituter(v) for (k, v) in defaults(sys)])
+    @set! sys.guesses = Dict([k => substituter(v) for (k, v) in guesses(sys)])
+    @set! sys.unknowns = map(substituter, get_unknowns(sys))
+    @set! sys.ps = map(substituter, get_ps(sys))
+    AffectSystem(sys, map(substituter, aff.unknowns),
+        map(substituter, aff.parameters), map(substituter, aff.discretes))
+end
+
 function AffectSystem(spec::SymbolicAffect; iv = nothing, alg_eqs = Equation[], kwargs...)
     AffectSystem(spec.affect; alg_eqs = vcat(spec.alg_eqs, alg_eqs), iv,
         discrete_parameters = spec.discrete_parameters, kwargs...)

src/systems/imperative_affect.jl

@@ -67,6 +67,12 @@ function ImperativeAffect(; f, kwargs...)
     ImperativeAffect(f; kwargs...)
 end
 
+function Symbolics.fast_substitute(aff::ImperativeAffect, rules)
+    substituter = Base.Fix2(fast_substitute, rules)
+    ImperativeAffect(aff.f, map(substituter, aff.obs), aff.obs_syms,
+        map(substituter, aff.modified), aff.mod_syms, aff.ctx, aff.skip_checks)
+end
+
 function Base.show(io::IO, mfa::ImperativeAffect)
     obs_vals = join(map((ob, nm) -> "$ob => $nm", mfa.obs, mfa.obs_syms), ", ")
     mod_vals = join(map((md, nm) -> "$md => $nm", mfa.modified, mfa.mod_syms), ", ")