fix: improve performance of implicit affect

Author: vyudu

src/systems/callbacks.jl

@@ -942,32 +942,27 @@ function compile_equational_affect(
             dvs_to_access = unknowns(affsys)
             ps_to_access = parameters(affsys)
 
-            u_getters = [getsym(sys, aff_map[u]) for u in dvs_to_access]
-            p_getters = [getsym(sys, unPre(p)) for p in ps_to_access]
-            u_setters = [setsym(sys, u) for u in dvs_to_update]
-            p_setters = [setsym(sys, p) for p in ps_to_update]
-            affu_getters = [getsym(affsys, sys_map[u]) for u in dvs_to_update]
-            affp_getters = [getsym(affsys, sys_map[p]) for p in ps_to_update]
+            u_getter = getsym(sys, [aff_map[u] for u in dvs_to_access])
+            p_getter = getsym(sys, [unPre(p) for p in ps_to_access])
+            u_setter! = setsym(sys, dvs_to_update)
+            p_setter! = setsym(sys, ps_to_update)
+            affu_getter = getsym(affsys, [sys_map[u] for u in dvs_to_update])
+            affp_getter = getsym(affsys, [sys_map[p] for p in ps_to_update])
 
             affprob = ImplicitDiscreteProblem(affsys, [dv => 0 for dv in dvs_to_access],
                 (0, 0), [p => 0 for p in ps_to_access];
                 build_initializeprob = false, check_length = false)
 
             function implicit_affect!(integ)
-                pmap = Pair[p => getp(integ) for (p, getp) in zip(ps_to_access, p_getters)]
-                u0map = Pair[u => getu(integ)
-                             for (u, getu) in zip(dvs_to_access, u_getters)]
-                affprob = remake(affprob, u0 = u0map, p = pmap, tspan = (integ.t, integ.t))
+                new_us = u_getter(integ)
+                new_ps = p_getter(integ)
+                affprob = remake(affprob, u0 = new_us, p = new_ps, tspan = (integ.t, integ.t))
                 affsol = init(affprob, IDSolve())
                 (check_error(affsol) === ReturnCode.InitialFailure) &&
                     throw(UnsolvableCallbackError(all_equations(aff)))
 
-                for (setu!, getu) in zip(u_setters, affu_getters)
-                    setu!(integ, getu(affsol))
-                end
-                for (setp!, getp) in zip(p_setters, affp_getters)
-                    setp!(integ, getp(affsol))
-                end
+                u_setter!(integ, affu_getter(affsol))
+                p_setter!(integ, affp_getter(affsol))
             end
         end
     end
@@ -1080,3 +1075,16 @@ function continuous_events_toplevel(sys::AbstractSystem)
     end
     return get_continuous_events(sys)
 end
+
+"""
+Process the symbolic events of a system.
+"""
+function create_symbolic_events(cont_events, disc_events, sys_eqs, iv)
+    alg_eqs = filter(eq -> eq.lhs isa Union{Symbolic, Number} && !is_diff_equation(eq),
+        sys_eqs)
+    cont_callbacks = to_cb_vector(cont_events; CB_TYPE = SymbolicContinuousCallback,
+        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
+    disc_callbacks = to_cb_vector(disc_events; CB_TYPE = SymbolicDiscreteCallback,
+        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
+    cont_callbacks, disc_callbacks
+end

src/systems/diffeqs/odesystem.jl

@@ -318,13 +318,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         throw(ArgumentError("System names must be unique."))
     end
 
-    alg_eqs = filter(eq -> eq.lhs isa Union{Symbolic, Number} && !is_diff_equation(eq),
-        deqs)
-    cont_callbacks = to_cb_vector(continuous_events; CB_TYPE = SymbolicContinuousCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
-    disc_callbacks = to_cb_vector(discrete_events; CB_TYPE = SymbolicDiscreteCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
-
+    cont_callbacks, disc_callbacks = create_symbolic_events(continuous_events, discrete_events, deqs, iv)
     if is_dde === nothing
         is_dde = _check_if_dde(deqs, iv′, systems)
     end

src/systems/diffeqs/sdesystem.jl

@@ -270,12 +270,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
     Wfact = RefValue(EMPTY_JAC)
     Wfact_t = RefValue(EMPTY_JAC)
 
-    alg_eqs = filter(eq -> eq.lhs isa Union{Symbolic, Number} && !is_diff_equation(eq),
-        deqs)
-    cont_callbacks = to_cb_vector(continuous_events; CB_TYPE = SymbolicContinuousCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
-    disc_callbacks = to_cb_vector(discrete_events; CB_TYPE = SymbolicDiscreteCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
+    cont_callbacks, disc_callbacks = create_symbolic_events(continuous_events, discrete_events, deqs, iv)
 
     if is_dde === nothing
         is_dde = _check_if_dde(deqs, iv′, systems)

src/systems/discrete_system/implicit_discrete_system.jl

@@ -298,9 +298,6 @@ function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
     for k in collect(keys(u0map))
         v = u0map[k]
         if !((op = operation(k)) isa Shift)
-            isnothing(getunshifted(k)) &&
-                @warn "Initial condition given in term of current state of the unknown. If `build_initializeprob = false`, this may be overridden by the implicit discrete solver."
-
             updated[k] = v
         elseif op.steps > 0
             error("Initial conditions must be for the current or past state of the unknowns. Instead of providing the condition for $k, provide the condition for $(Shift(iv, -1)(only(arguments(k)))).")