Fix changing dt in kernel

docs/user_guide/examples/explanation_kernelloop.md

@@ -38,6 +38,8 @@ Below is a structured overview of how the Kernel loop is implemented. Note that
 
    5. If `outputdt` is a multiple of `particles.t`, write `particles.x` and `particles.t` to parquet output file
 
+Changing `particles.dt` in a Kernel affects the current particle time update and the changed value is retained for subsequent steps. A shortened timestep used to land exactly on an output time or the simulation end time is temporary and does not replace the particle's retained timestep.
+
 Besides having commutable Kernels, the main advantage of this implementation is that, when using Field Sampling with e.g. `particles.var = fieldset.Var[particles.t, particles.z, particles.y, particles.x]`, the particle location stays the same throughout the entire Kernel loop. Additionally, this implementation ensures that the particle location is the same as the location of the sampled field in the output file.
 
 ## Example with currents and winds

src/parcels/_core/kernel.py

@@ -37,6 +37,21 @@
 }
 
 
+def _clip_dt_to_endtime(dt, time_to_endtime, compute_time_direction, dt_to_restore):
+    """Clip dt so that particles do not step beyond the endtime, in the direction of integration."""
+    if compute_time_direction == 1:
+        clipped_dt = np.maximum(np.minimum(dt, time_to_endtime), 0)
+    else:
+        clipped_dt = np.minimum(np.maximum(dt, -time_to_endtime), 0)
+
+    newly_clipped = clipped_dt != dt
+    indices = np.flatnonzero(newly_clipped)
+    if indices.size:
+        dt_to_restore.append((indices, dt[indices].copy()))
+
+    return clipped_dt
+
+
 class Kernel:
     """Kernel object that encapsulates auto-generated code.
 
@@ -94,16 +109,6 @@ def pclass(self):
     def fieldset(self):
         return self._fieldset
 
-    def remove_deleted(self, pset):
-        """Utility to remove all particles that signalled deletion."""
-        bool_indices = pset._data["state"] == StatusCode.Delete
-        indices = np.where(bool_indices)[0]
-        # TODO v4: need to implement ParticleFile writing of deleted particles
-        # if len(indices) > 0 and self.fieldset.particlefile is not None:
-        #     self.fieldset.particlefile.write(pset, None, indices=indices)
-        if len(indices) > 0:
-            pset.remove_indices(indices)
-
     def _position_update(self, particles, fieldset):
         particles.x += particles.dx
         particles.y += particles.dy
@@ -170,7 +175,7 @@ def merge(self, kernel):
             self.pclass,
         )
 
-    def execute(self, pset, endtime, dt):
+    def execute(self, pset, endtime):
         """Execute this Kernel over a ParticleSet for several timesteps.
 
         Parameters
@@ -179,25 +184,28 @@ def execute(self, pset, endtime, dt):
             object of (sub-)type ParticleSet
         endtime :
             endtime of this overall kernel evaluation step
-        dt :
-            computational integration timestep from pset.execute
         """
-        compute_time_direction = 1 if dt > 0 else -1
+        if len(pset) == 0:
+            return StatusCode.StopAllExecution
+
+        sign_dt = np.sign(pset.dt[0])
+        if np.any(np.sign(pset.dt) != sign_dt):
+            raise ValueError("All particle timesteps must have the same sign dt")
 
         pset._data["state"][:] = StatusCode.Success
+        preserve_particle_dt = not hasattr(self.fieldset, "RK45_tol")
 
         while (len(pset) > 0) and np.any(np.isin(pset.state, [StatusCode.Success, StatusCode.Repeat])):
-            time_to_endtime = compute_time_direction * (endtime - pset.t)
+            time_to_endtime = sign_dt * (endtime - pset.t)
 
             evaluate_particles = (pset.state == StatusCode.Success) & (time_to_endtime >= 0)
             if not np.any(evaluate_particles):
                 return StatusCode.Success
 
+            dt_to_restore = []
+
             # adapt dt to end exactly on endtime
-            if compute_time_direction == 1:
-                pset.dt = np.maximum(np.minimum(pset.dt, time_to_endtime), 0)
-            else:
-                pset.dt = np.minimum(np.maximum(pset.dt, -time_to_endtime), 0)
+            pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt, dt_to_restore)
 
             # run kernels for all particles that need to be evaluated
             for f in self._kernels:
@@ -216,18 +224,21 @@ def execute(self, pset, endtime, dt):
             # (particles that signalled Stop*/Delete/errors should not have time/position advanced)
             update_particles = evaluate_particles & (pset.state == StatusCode.Success)
             if np.any(update_particles):
+                pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt, dt_to_restore)
                 self._position_update(pset[update_particles], self._fieldset)
 
-            # revert to original dt (unless in RK45 mode)
-            if not hasattr(self.fieldset, "RK45_tol"):
-                pset._data["dt"][:] = dt
-
             # Set particle state for particles that reached endtime
             particles_endofloop = (pset.state == StatusCode.Success) & (pset.t == endtime)
             pset[particles_endofloop].state = StatusCode.EndofLoop
 
+            if preserve_particle_dt:
+                for indices, values in dt_to_restore:
+                    pset.dt[indices] = values
+
             # delete particles that signalled deletion
-            self.remove_deleted(pset)
+            indices = np.flatnonzero(pset.state == StatusCode.Delete)
+            if len(indices) > 0:
+                pset.remove_indices(indices)
 
             # check and throw errors
             if np.any(pset.state == StatusCode.StopAllExecution):

src/parcels/_core/particleset.py

@@ -445,7 +445,7 @@ def execute(
                 else:
                     next_time = end_time
 
-                self._kernel.execute(self, endtime=next_time, dt=dt)
+                self._kernel.execute(self, endtime=next_time)
 
                 if next_output is not None:
                     if np.abs(next_time - next_output) < 0.001:

tests/test_kernel.py

@@ -5,7 +5,9 @@
     FieldSet,
     KernelWarning,
     Particle,
+    ParticleFile,
     ParticleSet,
+    StatusCode,
     Variable,
 )
 from parcels._core.kernel import Kernel
@@ -166,8 +168,47 @@ def SampleP(particles, fieldset):  # pragma: no cover
         assert np.allclose(lons[0], 0.2)
 
 
-@pytest.mark.xfail(reason="Modifying dt in a kernel doesn't work GH2765")
-def test_dt_modify_in_kernel(fieldset):
+@pytest.mark.parametrize(("with_output", "expected_num_kernel_loops"), [(True, 7), (False, 5)])
+def test_dt_modify_in_kernel(fieldset, tmp_parquet, with_output, expected_num_kernel_loops):
+    npart = 10
+    TestParticle = Particle.add_variable(Variable("num_kernel_loops", dtype=np.float32, initial=0))
+    pset = ParticleSet(fieldset, pclass=TestParticle, x=np.linspace(0, 1, npart), y=[0] * npart)
+
+    def ModifyDt(particles, fieldset):  # pragma: no cover
+        particles.num_kernel_loops += 1
+        particles.dt = 4
+        particles[particles.x > 0.8].state = StatusCode.Delete
+
+    runtime = 17
+    output = ParticleFile(tmp_parquet, outputdt=5.0) if with_output else None
+    pset.execute(ModifyDt, runtime=runtime, dt=1.0, output_file=output, verbose_progress=False)
+
+    np.testing.assert_allclose(pset.t, runtime)
+    np.testing.assert_allclose(pset.num_kernel_loops, expected_num_kernel_loops)
+
+
+def test_dt_modify_in_kernel_vector(fieldset):
+    TestParticle = Particle.add_variable(Variable("num_kernel_loops", dtype=np.float32, initial=0))
+    pset = ParticleSet(fieldset, pclass=TestParticle, x=[0.5], y=[0])
+
+    dt_in_kernel = np.array([1, 4, 3, 4, 5, 5, 10, 2, 2])
+    fieldset.add_context("dt_in_kernel", dt_in_kernel)
+
+    def ModifyDt_FromArray(particles, fieldset):  # pragma: no cover
+        i = np.asarray(particles.num_kernel_loops).astype(int)
+        particles.dt = fieldset.dt_in_kernel[i]
+        assert particles.t == np.cumsum(np.r_[0, fieldset.dt_in_kernel])[i]
+        particles.num_kernel_loops += 1
+
+    runtime = sum(dt_in_kernel)
+    expected_num_kernel_loops = len(dt_in_kernel)
+    pset.execute(ModifyDt_FromArray, runtime=runtime, dt=1.0)
+
+    np.testing.assert_allclose(pset.t[0], runtime)
+    np.testing.assert_allclose(pset.num_kernel_loops[0], expected_num_kernel_loops)
+
+
+def test_dt_modify_age_in_kernel(fieldset):
     TestParticle = Particle.add_variable(Variable("age", dtype=np.float32, initial=0))
     pset = ParticleSet(fieldset, pclass=TestParticle, x=[0.5], y=[0])
 
@@ -176,7 +217,7 @@ def ModifyDt(particles, fieldset):  # pragma: no cover
         particles.dt = 2
 
     runtime = 10
-    expected_age = 1 + 2 * (runtime - 2)  # 1 for the first step; 2 for the remaining steps (except last)
+    expected_age = 1 + 2 * (runtime / 2 - 1)  # 1 for the first step; 2 for the remaining steps (except last)
     pset.execute(ModifyDt, runtime=runtime, dt=1.0)
     np.testing.assert_allclose(pset.t[0], runtime)
     np.testing.assert_allclose(pset.age[0], expected_age)