Merge pull request #3901 from AayushSabharwal/as/fix-disc-save

fix: use improved discrete saving API

Author: AayushSabharwal

Project.toml

@@ -100,8 +100,8 @@ ConstructionBase = "1"
 DataInterpolations = "7, 8"
 DataStructures = "0.17, 0.18"
 DeepDiffs = "1"
-DelayDiffEq = "5.50"
-DiffEqBase = "6.170.1"
+DelayDiffEq = "5.61"
+DiffEqBase = "6.189.1"
 DiffEqCallbacks = "2.16, 3, 4"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
@@ -123,7 +123,7 @@ ImplicitDiscreteSolve = "0.1.2, 1"
 InfiniteOpt = "0.5"
 InteractiveUtils = "1"
 JuliaFormatter = "1.0.47, 2"
-JumpProcesses = "9.13.1"
+JumpProcesses = "9.19"
 LabelledArrays = "1.3"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
 Libdl = "1"
@@ -138,7 +138,7 @@ NonlinearSolve = "4.3"
 OffsetArrays = "1"
 OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
-OrdinaryDiffEqCore = "1.15.0"
+OrdinaryDiffEqCore = "1.34.0"
 OrdinaryDiffEqDefault = "1.2"
 OrdinaryDiffEqNonlinearSolve = "1.5.0"
 PrecompileTools = "1"
@@ -148,7 +148,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.4.0"
-SciMLBase = "2.108.0"
+SciMLBase = "2.115.0"
 SciMLPublic = "1.0.0"
 SciMLStructures = "1.7"
 Serialization = "1"
@@ -157,8 +157,8 @@ SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "1, 2"
 StaticArrays = "1.9.14"
-StochasticDelayDiffEq = "1.10"
-StochasticDiffEq = "6.72.1"
+StochasticDelayDiffEq = "1.11"
+StochasticDiffEq = "6.82.0"
 SymbolicIndexingInterface = "0.3.39"
 SymbolicUtils = "3.30.0"
 Symbolics = "6.40"

src/systems/callbacks.jl

@@ -716,12 +716,18 @@ function generate_callback(cbs::Vector{SymbolicContinuousCallback}, sys; kwargs.
         return generate_callback(cbs[cb_ind], sys; kwargs...)
     end
 
+    if is_split(sys)
+        ic = get_index_cache(sys)
+    else
+        ic = nothing
+    end
     trigger = compile_condition(
         cbs, sys, unknowns(sys), parameters(sys; initial_parameters = true); kwargs...)
     affects = []
     affect_negs = []
     inits = []
     finals = []
+    saved_clock_partitions = Vector{Int}[]
     for cb in cbs
         affect = compile_affect(cb.affect, cb, sys; default = EMPTY_AFFECT, kwargs...)
         push!(affects, affect)
@@ -731,8 +737,15 @@ function generate_callback(cbs::Vector{SymbolicContinuousCallback}, sys; kwargs.
         push!(affect_negs, affect_neg)
         push!(inits,
             compile_affect(
-                cb.initialize, cb, sys; default = nothing, is_init = true, kwargs...))
+                cb.initialize, cb, sys; default = nothing, kwargs...))
         push!(finals, compile_affect(cb.finalize, cb, sys; default = nothing, kwargs...))
+
+        if ic !== nothing
+            save_idxs = get(ic.callback_to_clocks, cb, Int[])
+            for _ in conditions(cb)
+                push!(saved_clock_partitions, save_idxs)
+            end
+        end
     end
 
     # Since there may be different number of conditions and affects,
@@ -758,7 +771,8 @@ function generate_callback(cbs::Vector{SymbolicContinuousCallback}, sys; kwargs.
 
     return VectorContinuousCallback(
         trigger, affect, affect_neg, length(eqs); initialize, finalize,
-        rootfind = cbs[1].rootfind, initializealg = cbs[1].reinitializealg)
+        rootfind = cbs[1].rootfind, initializealg = cbs[1].reinitializealg,
+        saved_clock_partitions)
 end
 
 function generate_callback(cb, sys; kwargs...)
@@ -775,27 +789,33 @@ function generate_callback(cb, sys; kwargs...)
         compile_affect(cb.affect_neg, cb, sys; default = EMPTY_AFFECT, kwargs...)
     end
     init = compile_affect(cb.initialize, cb, sys; default = SciMLBase.INITIALIZE_DEFAULT,
-        is_init = true, kwargs...)
+        kwargs...)
     final = compile_affect(
         cb.finalize, cb, sys; default = SciMLBase.FINALIZE_DEFAULT, kwargs...)
 
     initialize = isnothing(cb.initialize) ? init : ((c, u, t, i) -> init(i))
     finalize = isnothing(cb.finalize) ? final : ((c, u, t, i) -> final(i))
 
+    saved_clock_partitions = if is_split(sys)
+        get(get_index_cache(sys).callback_to_clocks, cb, ())
+    else
+        ()
+    end
     if is_discrete(cb)
         if is_timed && conditions(cb) isa AbstractVector
             return PresetTimeCallback(trigger, affect; initialize,
-                finalize, initializealg = cb.reinitializealg)
+                finalize, initializealg = cb.reinitializealg, saved_clock_partitions)
         elseif is_timed
             return PeriodicCallback(
-                affect, trigger; initialize, finalize, initializealg = cb.reinitializealg)
+                affect, trigger; initialize, finalize, initializealg = cb.reinitializealg,
+                saved_clock_partitions)
         else
             return DiscreteCallback(trigger, affect; initialize,
-                finalize, initializealg = cb.reinitializealg)
+                finalize, initializealg = cb.reinitializealg, saved_clock_partitions)
         end
     else
         return ContinuousCallback(trigger, affect, affect_neg; initialize, finalize,
-            rootfind = cb.rootfind, initializealg = cb.reinitializealg)
+            rootfind = cb.rootfind, initializealg = cb.reinitializealg, saved_clock_partitions)
     end
 end
 
@@ -810,41 +830,13 @@ Notes
 """
 function compile_affect(
         aff::Union{Nothing, Affect}, cb::AbstractCallback, sys::AbstractSystem;
-        default = nothing, is_init = false, kwargs...)
-    save_idxs = if !(has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing)
-        Int[]
-    else
-        get(ic.callback_to_clocks, cb, Int[])
-    end
-
+        default = nothing, kwargs...)
     if isnothing(aff)
-        is_init ? wrap_save_discretes(default, save_idxs) : default
+        default
     elseif aff isa AffectSystem
-        f = compile_equational_affect(aff, sys; kwargs...)
-        wrap_save_discretes(f, save_idxs)
+        compile_equational_affect(aff, sys; kwargs...)
     elseif aff isa ImperativeAffect
-        f = compile_functional_affect(aff, sys; kwargs...)
-        wrap_save_discretes(f, save_idxs)
-    end
-end
-
-function wrap_save_discretes(f, save_idxs)
-    let save_idxs = save_idxs, f = f
-        if f === SciMLBase.INITIALIZE_DEFAULT
-            (c, u, t, i) -> begin
-                f(c, u, t, i)
-                for idx in save_idxs
-                    SciMLBase.save_discretes!(i, idx)
-                end
-            end
-        else
-            (i) -> begin
-                isnothing(f) || f(i)
-                for idx in save_idxs
-                    SciMLBase.save_discretes!(i, idx)
-                end
-            end
-        end
+        compile_functional_affect(aff, sys; kwargs...)
     end
 end
 

test/jacobiansparsity.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, SparseArrays, OrdinaryDiffEq
+using ModelingToolkit, SparseArrays, OrdinaryDiffEq, DiffEqBase
 
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)

test/odesystem.jl

@@ -865,7 +865,7 @@ prob = ODEProblem(sys, [x => 1.0], (0.0, 10.0))
     prob2 = @test_nowarn ODEProblem(sys, [x => ones(3), p => ones(3, 3)], (0.0, 10.0))
     sol2 = @test_nowarn solve(prob2, Tsit5())
 
-    @test sol1.u ≈ sol2.u[2:end]
+    @test sol1.u ≈ sol2.u
 end
 
 # Requires fix in symbolics for `linear_expansion(p * x, D(y))`