From b52a59f34ce8c888426703b8ffe723f5a48d4aa3 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Tue, 28 May 2024 07:36:45 +0200
Subject: [PATCH] add phase offset option to PeriodicCallback

---
 src/iterative_and_periodic.jl | 28 +++++++++++++++++-----------
 test/periodic_tests.jl        | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/src/iterative_and_periodic.jl b/src/iterative_and_periodic.jl
index 734dfc45..64bbabf7 100644
--- a/src/iterative_and_periodic.jl
+++ b/src/iterative_and_periodic.jl
@@ -82,8 +82,6 @@ struct PeriodicCallbackAffect{A, dT, Ref1, Ref2}
 end
 
 function (S::PeriodicCallbackAffect)(integrator)
-    @unpack affect!, Δt, t0, index = S
-
     add_next_tstop!(integrator, S)
 
     affect!(integrator)
@@ -93,7 +91,7 @@ function add_next_tstop!(integrator, S)
     @unpack Δt, t0, index = S
 
     # Schedule next call to `f` using `add_tstops!`, but be careful not to keep integrating forever
-    tnew = t0[] + (index[] + 1) * Δt
+    @show tnew = t0[] + (index[] + 1) * Δt
     tstops = integrator.opts.tstops
     #=
     Okay yeah, this is nasty
@@ -110,34 +108,42 @@ end
 
 """
 ```julia
-PeriodicCallback(f, Δt::Number; initial_affect = false,
+PeriodicCallback(f, Δt::Number; phase = 0, initial_affect = false,
     final_affect = false,
     kwargs...)
 ```
 
 `PeriodicCallback` can be used when a function should be called periodically in terms of
 integration time (as opposed to wall time), i.e. at `t = tspan[1]`, `t = tspan[1] + Δt`,
-`t = tspan[1] + 2Δt`, and so on. This callback can, for example, be used to model a
+`t = tspan[1] + 2Δt`, and so on.
+
+If a non-zero `phase` is provided, the invokations of the callback will be shifted by `phase` time units, i.e., the calls will occur at
+`t = tspan[1] + phase`, `t = tspan[1] + phase + Δt`,
+`t = tspan[1] + phase + 2Δt`, and so on.
+
+This callback can, for example, be used to model a
 discrete-time controller for a continuous-time system, running at a fixed rate.
 
 ## Arguments
 
   - `f` the `affect!(integrator)` function to be called periodically
   - `Δt` is the period
-
-## Keyword Arguments
-
+  
+  ## Keyword Arguments
+  
+  - `phase` is a phase offset
   - `initial_affect` is whether to apply the affect at `t=0`, which defaults to `false`
   - `final_affect` is whether to apply the affect at the final time, which defaults to `false`
   - `kwargs` are keyword arguments accepted by the `DiscreteCallback` constructor.
 """
 function PeriodicCallback(f, Δt::Number;
+        phase = 0,
         initial_affect = false,
         final_affect = false,
         initialize = (cb, u, t, integrator) -> u_modified!(integrator,
             initial_affect),
         kwargs...)
-
+    phase < 0 && throw(ArgumentError("phase offset must be non-negative"))
     # Value of `t` at which `f` should be called next:
     t0 = Ref(typemax(Δt))
     index = Ref(0)
@@ -154,8 +160,8 @@ function PeriodicCallback(f, Δt::Number;
     initialize_periodic = function (c, u, t, integrator)
         @assert integrator.tdir == sign(Δt)
         initialize(c, u, t, integrator)
-        t0[] = t
-        index[] = 0
+        t0[] = t + phase
+        index[] = iszero(phase) ? 0 : -1
         if initial_affect
             affect!(integrator)
         else
diff --git a/test/periodic_tests.jl b/test/periodic_tests.jl
index 567eb5ff..ae1545a4 100644
--- a/test/periodic_tests.jl
+++ b/test/periodic_tests.jl
@@ -126,3 +126,37 @@ sol2 = solve(prob, Tsit5(), callback = cb)
 @test sol2(72.0 + eps(72.0)) ≈ [10.0]
 @test sol2(96.0 + eps(96.0)) ≈ [10.0]
 @test sol2(120.0 + eps(120.0)) ≈ [10.0]
+
+# Test phase offset
+approxin(needle, haystack) = any(isapprox(needle), haystack)
+
+function integ(du, u, p, t)
+    du[1] = 1
+    du[2] = 0
+end
+
+function nullaffect!(integ)
+    integ.u[2] += 1
+end
+
+cb = PeriodicCallback(nullaffect!, 1.0; phase = 0.1)
+prob = ODEProblem(integ, [0.0, 0], (0.0, 10.0))
+sol = solve(prob, Tsit5(), callback = cb)
+@test sol.u[end][2] == 10 # Test expected number of calls to affect
+for t in 0.1:1:10
+    @test approxin(t, sol.t) # Test that the integrator stopped at all expected points
+end
+sol.t[end] ≈ 10 # test that we did not step over end
+
+# With negative phase
+@test_throws ArgumentError PeriodicCallback(nullaffect!, 1.0; phase = -0.1)
+
+# Phase offset larger than period
+cb = PeriodicCallback(nullaffect!, 1.0; phase = 1.1)
+prob = ODEProblem(integ, [0.0, 0], (0.0, 10.0))
+sol = solve(prob, Tsit5(), callback = cb)
+@test sol.u[end][2] == 9 # Test expected number of calls to affect
+for t in 1.1:1:10
+    @test approxin(t, sol.t) # Test that the integrator stopped at all expected points
+end
+sol.t[end] ≈ 10 # test that we did not step over end