feat: add discrete saving to callback structs

Author: AayushSabharwal

src/callbacks.jl

@@ -110,8 +110,13 @@ Contains a single callback whose `condition` is a continuous function. The callb
     `affect!` satisfies the constraints (or else errors). It is not recommended that `NoInit()` is
     used as that will lead to an unstable step following initialization. This warning can be
     ignored for non-DAE ODEs.
+
+# Extended help
+
+- `discrete_save_idxs`: An iterable of timeseries indexes to save after the callback triggers. MTK-only
+  API
 """
-struct ContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R} <:
+struct ContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R, DSI} <:
        AbstractContinuousCallback
     condition::F1
     affect!::F2
@@ -127,20 +132,20 @@ struct ContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R} <:
     reltol::T2
     repeat_nudge::T3
     initializealg::T4
+    discrete_save_idxs::DSI
     function ContinuousCallback(condition::F1, affect!::F2, affect_neg!::F3,
             initialize::F4, finalize::F5, idxs::I, rootfind,
             interp_points, save_positions, dtrelax::R, abstol::T,
-            reltol::T2,
-            repeat_nudge::T3,
-            initializealg::T4 = nothing) where {F1, F2, F3, F4, F5, T, T2, T3, T4, I, R
+            reltol::T2, repeat_nudge::T3, initializealg::T4 = nothing,
+            discrete_save_idxs::DSI = ()) where {F1, F2, F3, F4, F5, T, T2, T3, T4, I, R, DSI
     }
         _condition = prepare_function(condition)
-        new{typeof(_condition), F2, F3, F4, F5, T, T2, T3, T4, I, R}(_condition,
+        new{typeof(_condition), F2, F3, F4, F5, T, T2, T3, T4, I, R, DSI}(_condition,
             affect!, affect_neg!,
             initialize, finalize, idxs, rootfind,
             interp_points,
             BitArray(collect(save_positions)),
-            dtrelax, abstol, reltol, repeat_nudge, initializealg)
+            dtrelax, abstol, reltol, repeat_nudge, initializealg, discrete_save_idxs)
     end
 end
 
@@ -154,12 +159,13 @@ function ContinuousCallback(condition, affect!, affect_neg!;
         dtrelax = 1,
         abstol = 10eps(), reltol = 0,
         repeat_nudge = 1 // 100,
-        initializealg = nothing)
+        initializealg = nothing,
+        discrete_save_idxs = ())
     ContinuousCallback(condition, affect!, affect_neg!, initialize, finalize,
         idxs,
         rootfind, interp_points,
         save_positions,
-        dtrelax, abstol, reltol, repeat_nudge, initializealg)
+        dtrelax, abstol, reltol, repeat_nudge, initializealg, discrete_save_idxs)
 end
 
 function ContinuousCallback(condition, affect!;
@@ -172,11 +178,11 @@ function ContinuousCallback(condition, affect!;
         interp_points = 10,
         dtrelax = 1,
         abstol = 10eps(), reltol = 0, repeat_nudge = 1 // 100,
-        initializealg = nothing)
+        initializealg = nothing, discrete_save_idxs = ())
     ContinuousCallback(condition, affect!, affect_neg!, initialize, finalize, idxs,
         rootfind, interp_points,
         collect(save_positions),
-        dtrelax, abstol, reltol, repeat_nudge, initializealg)
+        dtrelax, abstol, reltol, repeat_nudge, initializealg, discrete_save_idxs)
 end
 
 """
@@ -219,8 +225,16 @@ multiple events.
   - `len`: Number of callbacks chained. This is compulsory to be specified.
 
 Rest of the arguments have the same meaning as in [`ContinuousCallback`](@ref).
+
+# Extended help
+
+- `discrete_save_idxs`: An iterable of `len` elements, where the `i`th element is an iterable of timeseries
+  indexes to save when the `i`th event triggers. MTK-only API.
+- `all_save_idxs`: An iterable of all unique timeseries indexes in `discrete_save_idxs`. Used to save after
+  `initialize` and `finalize`. This avoids saving twice if the same timeseries index is saved by two
+  events.
 """
-struct VectorContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R} <:
+struct VectorContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R, DSI, ASI} <:
        AbstractContinuousCallback
     condition::F1
     affect!::F2
@@ -237,21 +251,25 @@ struct VectorContinuousCallback{F1, F2, F3, F4, F5, T, T2, T3, T4, I, R} <:
     reltol::T2
     repeat_nudge::T3
     initializealg::T4
+    discrete_save_idxs::DSI
+    all_save_idxs::ASI
     function VectorContinuousCallback(
             condition::F1, affect!::F2, affect_neg!::F3, len::Int,
             initialize::F4, finalize::F5, idxs::I, rootfind,
             interp_points, save_positions, dtrelax::R,
-            abstol::T, reltol::T2,
-            repeat_nudge::T3,
-            initializealg::T4 = nothing) where {F1, F2, F3, F4, F5, T, T2,
-            T3, T4, I, R}
+            abstol::T, reltol::T2, repeat_nudge::T3,
+            initializealg::T4 = nothing, discrete_save_idxs::DSI = (),
+            all_save_idxs::ASI = ()) where {F1, F2, F3, F4, F5, T, T2,
+            T3, T4, I, R, DSI, ASI}
         _condition = prepare_function(condition)
-        new{typeof(_condition), F2, F3, F4, F5, T, T2, T3, T4, I, R}(_condition,
+        new{typeof(_condition), F2, F3, F4, F5, T, T2, T3, T4, I, R, DSI, ASI}(
+            _condition,
             affect!, affect_neg!, len,
             initialize, finalize, idxs, rootfind,
             interp_points,
             BitArray(collect(save_positions)),
-            dtrelax, abstol, reltol, repeat_nudge, initializealg)
+            dtrelax, abstol, reltol, repeat_nudge, initializealg,
+            discrete_save_idxs, all_save_idxs)
     end
 end
 
@@ -264,13 +282,15 @@ function VectorContinuousCallback(condition, affect!, affect_neg!, len;
         interp_points = 10,
         dtrelax = 1,
         abstol = 10eps(), reltol = 0, repeat_nudge = 1 // 100,
-        initializealg = nothing)
+        initializealg = nothing, discrete_save_idxs = (),
+        all_save_idxs = ())
     VectorContinuousCallback(condition, affect!, affect_neg!, len,
         initialize, finalize,
         idxs,
         rootfind, interp_points,
         save_positions, dtrelax,
-        abstol, reltol, repeat_nudge, initializealg)
+        abstol, reltol, repeat_nudge, initializealg, discrete_save_idxs,
+        all_save_idxs)
 end
 
 function VectorContinuousCallback(condition, affect!, len;
@@ -283,12 +303,14 @@ function VectorContinuousCallback(condition, affect!, len;
         interp_points = 10,
         dtrelax = 1,
         abstol = 10eps(), reltol = 0, repeat_nudge = 1 // 100,
-        initializealg = nothing)
+        initializealg = nothing, discrete_save_idxs = (),
+        all_save_idxs = ())
     VectorContinuousCallback(condition, affect!, affect_neg!, len, initialize, finalize,
         idxs,
         rootfind, interp_points,
         collect(save_positions),
-        dtrelax, abstol, reltol, repeat_nudge, initializealg)
+        dtrelax, abstol, reltol, repeat_nudge, initializealg, discrete_save_idxs,
+        all_save_idxs)
 end
 
 """
@@ -339,31 +361,39 @@ DiscreteCallback(condition, affect!;
     `affect!` satisfies the constraints (or else errors). It is not recommended that `NoInit()` is
     used as that will lead to an unstable step following initialization. This warning can be
     ignored for non-DAE ODEs.
+
+# Extended help
+
+- `discrete_save_idxs`: An iterable of timeseries indexes to save after the callback triggers. MTK-only
+  API
 """
-struct DiscreteCallback{F1, F2, F3, F4, F5} <: AbstractDiscreteCallback
+struct DiscreteCallback{F1, F2, F3, F4, F5, DSI} <: AbstractDiscreteCallback
     condition::F1
     affect!::F2
     initialize::F3
     finalize::F4
     save_positions::BitArray{1}
     initializealg::F5
+    discrete_save_idxs::DSI
     function DiscreteCallback(condition::F1, affect!::F2,
             initialize::F3, finalize::F4,
             save_positions,
-            initializealg::F5 = nothing) where {F1, F2, F3, F4, F5}
+            initializealg::F5 = nothing,
+            discrete_save_idxs::DSI = ()) where {F1, F2, F3, F4, F5, DSI}
         _condition = prepare_function(condition)
-        new{typeof(_condition), F2, F3, F4, F5}(_condition,
+        new{typeof(_condition), F2, F3, F4, F5, DSI}(_condition,
             affect!, initialize, finalize,
             BitArray(collect(save_positions)),
-            initializealg)
+            initializealg, discrete_save_idxs)
     end
 end
 function DiscreteCallback(condition, affect!;
         initialize = INITIALIZE_DEFAULT, finalize = FINALIZE_DEFAULT,
         save_positions = (true, true),
-        initializealg = nothing)
+        initializealg = nothing, discrete_save_idxs = ())
     DiscreteCallback(
-        condition, affect!, initialize, finalize, save_positions, initializealg)
+        condition, affect!, initialize, finalize, save_positions, initializealg,
+        discrete_save_idxs)
 end
 
 """
@@ -420,3 +450,21 @@ end
     split_callbacks((cs..., d.continuous_callbacks...), (ds..., d.discrete_callbacks...),
         args...)
 end
+
+function save_discretes!(integrator::DEIntegrator, cb::Union{ContinuousCallback, DiscreteCallback})
+    for idx in cb.discrete_save_idxs
+        save_discretes!(integrator, idx)
+    end
+end
+
+function save_discretes!(integrator::DEIntegrator, cb::VectorContinuousCallback)
+    for idx in cb.all_save_idxs
+        save_discretes!(integrator, idx)
+    end
+end
+
+function save_discretes!(integrator::DEIntegrator, cb::VectorContinuousCallback, i)
+    for idx in cb.discrete_save_idxs[i]
+        save_discretes!(integrator, idx)
+    end
+end