From 1d72831eeb4129466cf31ef3c2fbb8097b27c4dd Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 20:36:47 +0100 Subject: [PATCH 01/12] Update and add unit test for changing dt inside kernel --- tests/test_kernel.py | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index bc5c4bea9..ea9add04b 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -166,20 +166,41 @@ 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): - TestParticle = Particle.add_variable(Variable("age", dtype=np.float32, initial=0)) + TestParticle = Particle.add_variable(Variable("num_kernel_loops", dtype=np.float32, initial=0)) pset = ParticleSet(fieldset, pclass=TestParticle, x=[0.5], y=[0]) def ModifyDt(particles, fieldset): # pragma: no cover - particles.age += particles.dt - particles.dt = 2 + particles.num_kernel_loops += 1 + particles.dt = 4 - runtime = 10 - expected_age = 1 + 2 * (runtime - 2) # 1 for the first step; 2 for the remaining steps (except last) + runtime = 17 + expected_num_kernel_loops = 5 # Five kernel loops pset.execute(ModifyDt, 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_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_FromAarray(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_FromAarray, runtime=runtime, dt=1.0) + np.testing.assert_allclose(pset.t[0], runtime) - np.testing.assert_allclose(pset.age[0], expected_age) + np.testing.assert_allclose(pset.num_kernel_loops[0], expected_num_kernel_loops) @pytest.mark.parametrize("dt", [1e-2, 1e-5, 1e-6, 1e-9]) From c81e4428266755ffb8dbe2d2305a6b9a82ffea0f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 20:37:20 +0100 Subject: [PATCH 02/12] Add explanation about changing dt --- docs/user_guide/examples/explanation_kernelloop.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/user_guide/examples/explanation_kernelloop.md b/docs/user_guide/examples/explanation_kernelloop.md index 20de03651..a6a4677c2 100644 --- a/docs/user_guide/examples/explanation_kernelloop.md +++ b/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 From f3d69029cd1da89d6b20ee15254eab1c173144a0 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 20:58:56 +0100 Subject: [PATCH 03/12] Also adding test with ParticleFile (to break the while loop) --- tests/test_kernel.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index ea9add04b..e3130cf4c 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -5,6 +5,7 @@ FieldSet, KernelWarning, Particle, + ParticleFile, ParticleSet, Variable, ) @@ -166,7 +167,8 @@ def SampleP(particles, fieldset): # pragma: no cover assert np.allclose(lons[0], 0.2) -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): TestParticle = Particle.add_variable(Variable("num_kernel_loops", dtype=np.float32, initial=0)) pset = ParticleSet(fieldset, pclass=TestParticle, x=[0.5], y=[0]) @@ -175,8 +177,8 @@ def ModifyDt(particles, fieldset): # pragma: no cover particles.dt = 4 runtime = 17 - expected_num_kernel_loops = 5 # Five kernel loops - pset.execute(ModifyDt, runtime=runtime, dt=1.0) + 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[0], runtime) np.testing.assert_allclose(pset.num_kernel_loops[0], expected_num_kernel_loops) From 865cd02a83e94a9f6336f8f5c1337e6981df2b84 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:01:44 +0100 Subject: [PATCH 04/12] Adding support for changing dt in the kernel --- src/parcels/_core/kernel.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 8708e7358..73fb89681 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -37,6 +37,13 @@ } +def _clip_dt_to_endtime(dt, time_to_endtime, compute_time_direction): + """Clip dt so that particles do not step beyond the endtime, in the direction of integration.""" + if compute_time_direction == 1: + return np.maximum(np.minimum(dt, time_to_endtime), 0) + return np.minimum(np.maximum(dt, -time_to_endtime), 0) + + class Kernel: """Kernel object that encapsulates auto-generated code. @@ -185,6 +192,8 @@ def execute(self, pset, endtime, dt): compute_time_direction = 1 if dt > 0 else -1 pset._data["state"][:] = StatusCode.Success + preserve_particle_dt = not hasattr(self.fieldset, "RK45_tol") + dt_to_restore = pset.dt.copy() while (len(pset) > 0) and np.any(np.isin(pset.state, [StatusCode.Success, StatusCode.Repeat])): time_to_endtime = compute_time_direction * (endtime - pset.t) @@ -193,11 +202,11 @@ def execute(self, pset, endtime, dt): if not np.any(evaluate_particles): return StatusCode.Success + if preserve_particle_dt: + pset.dt = 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, compute_time_direction) # run kernels for all particles that need to be evaluated for f in self._kernels: @@ -216,18 +225,22 @@ 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, compute_time_direction) 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 + # revert to the retained dt (unless in RK45 mode) + if preserve_particle_dt: + pset.dt = dt_to_restore # Set particle state for particles that reached endtime particles_endofloop = (pset.state == StatusCode.Success) & (pset.t == endtime) pset[particles_endofloop].state = StatusCode.EndofLoop - # delete particles that signalled deletion + # delete particles that signalled deletion and remove their corresponding dt_to_restore + deleted_particles = pset.state == StatusCode.Delete self.remove_deleted(pset) + if preserve_particle_dt: + dt_to_restore = dt_to_restore[~deleted_particles] # check and throw errors if np.any(pset.state == StatusCode.StopAllExecution): From 3557709a07c637d74f4b6a328587be8e8e533414 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:05:58 +0100 Subject: [PATCH 05/12] Changing kernel.execute function signature to use sign_dt --- src/parcels/_core/kernel.py | 16 +++++++--------- src/parcels/_core/particleset.py | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 73fb89681..21d6d6d5d 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -177,8 +177,8 @@ def merge(self, kernel): self.pclass, ) - def execute(self, pset, endtime, dt): - """Execute this Kernel over a ParticleSet for several timesteps. + def execute(self, pset, endtime, sign_dt): + """Execute this Kernel over a ParticleSet for several timesteps in a given time direction. Parameters ---------- @@ -186,17 +186,15 @@ 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 + sign_dt : + direction of time integration, 1 for forward, -1 for backward """ - compute_time_direction = 1 if dt > 0 else -1 - pset._data["state"][:] = StatusCode.Success preserve_particle_dt = not hasattr(self.fieldset, "RK45_tol") dt_to_restore = pset.dt.copy() 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): @@ -206,7 +204,7 @@ def execute(self, pset, endtime, dt): pset.dt = dt_to_restore # adapt dt to end exactly on endtime - pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, compute_time_direction) + pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt) # run kernels for all particles that need to be evaluated for f in self._kernels: @@ -225,7 +223,7 @@ 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, compute_time_direction) + pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt) self._position_update(pset[update_particles], self._fieldset) # revert to the retained dt (unless in RK45 mode) diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 075403b9c..409c74e99 100644 --- a/src/parcels/_core/particleset.py +++ b/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, sign_dt=sign_dt) if next_output is not None: if np.abs(next_time - next_output) < 0.001: From 3233f66fe36e4e82fb3668f3d7947e11b113392f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:15:42 +0100 Subject: [PATCH 06/12] Also testing for multiple particles and deleting them --- tests/test_kernel.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index e3130cf4c..e7f631c14 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -7,6 +7,7 @@ Particle, ParticleFile, ParticleSet, + StatusCode, Variable, ) from parcels._core.kernel import Kernel @@ -169,19 +170,21 @@ def SampleP(particles, fieldset): # pragma: no cover @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=[0.5], y=[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[0], runtime) - np.testing.assert_allclose(pset.num_kernel_loops[0], expected_num_kernel_loops) + 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): From 8c3c18c2f250d7bffe8e6408c347bbe6286d92d3 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:19:54 +0100 Subject: [PATCH 07/12] Simplifying deleted particles logic --- src/parcels/_core/kernel.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 21d6d6d5d..09ea79d1e 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -101,16 +101,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 @@ -235,10 +225,11 @@ def execute(self, pset, endtime, sign_dt): pset[particles_endofloop].state = StatusCode.EndofLoop # delete particles that signalled deletion and remove their corresponding dt_to_restore - deleted_particles = pset.state == StatusCode.Delete - self.remove_deleted(pset) - if preserve_particle_dt: - dt_to_restore = dt_to_restore[~deleted_particles] + indices = np.flatnonzero(pset.state == StatusCode.Delete) + if len(indices) > 0: + pset.remove_indices(indices) + if preserve_particle_dt: + dt_to_restore = np.delete(dt_to_restore, indices) # check and throw errors if np.any(pset.state == StatusCode.StopAllExecution): From 07c0bf1ef1ea80cfc28bd92803a3d30c9847b08f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:30:42 +0100 Subject: [PATCH 08/12] Removing sign_dt/dt argument from kernel.execute --- src/parcels/_core/kernel.py | 8 +++++--- src/parcels/_core/particleset.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 09ea79d1e..33d31d810 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -167,7 +167,7 @@ def merge(self, kernel): self.pclass, ) - def execute(self, pset, endtime, sign_dt): + def execute(self, pset, endtime): """Execute this Kernel over a ParticleSet for several timesteps in a given time direction. Parameters @@ -176,9 +176,11 @@ def execute(self, pset, endtime, sign_dt): object of (sub-)type ParticleSet endtime : endtime of this overall kernel evaluation step - sign_dt : - direction of time integration, 1 for forward, -1 for backward """ + sign_dt = np.sign(pset.dt[0]) + if len(pset) > 0 and np.any(np.sign(pset.dt) != sign_dt): + raise ValueError("All particle timesteps must have the same sign") + pset._data["state"][:] = StatusCode.Success preserve_particle_dt = not hasattr(self.fieldset, "RK45_tol") dt_to_restore = pset.dt.copy() diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 409c74e99..d7ae51791 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -445,7 +445,7 @@ def execute( else: next_time = end_time - self._kernel.execute(self, endtime=next_time, sign_dt=sign_dt) + self._kernel.execute(self, endtime=next_time) if next_output is not None: if np.abs(next_time - next_output) < 0.001: From e28a80f8fef32f9840992f73391fc9e3fe9114c1 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 28 Oct 2026 21:31:26 +0100 Subject: [PATCH 09/12] Update execute docstring --- src/parcels/_core/kernel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 33d31d810..2f76af663 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -168,7 +168,7 @@ def merge(self, kernel): ) def execute(self, pset, endtime): - """Execute this Kernel over a ParticleSet for several timesteps in a given time direction. + """Execute this Kernel over a ParticleSet for several timesteps. Parameters ---------- From 6f49df56ab775348cd34b9b62cbb5725bde0d276 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 30 Oct 2026 12:37:47 +0100 Subject: [PATCH 10/12] Fix empty particlesets As raised in tutorial_Delft3D --- src/parcels/_core/kernel.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 2f76af663..9722e0d3f 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -177,8 +177,11 @@ def execute(self, pset, endtime): endtime : endtime of this overall kernel evaluation step """ + if len(pset) == 0: + return StatusCode.StopAllExecution + sign_dt = np.sign(pset.dt[0]) - if len(pset) > 0 and np.any(np.sign(pset.dt) != sign_dt): + if np.any(np.sign(pset.dt) != sign_dt): raise ValueError("All particle timesteps must have the same sign") pset._data["state"][:] = StatusCode.Success From 1d0bbe214f5aa68ab94ba405faabb406633a62df Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 30 Oct 2026 13:45:04 +0100 Subject: [PATCH 11/12] Fix typo in kernel name --- tests/test_kernel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index e7f631c14..3cb30094c 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -194,7 +194,7 @@ def test_dt_modify_in_kernel_vector(fieldset): 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_FromAarray(particles, fieldset): # pragma: no cover + 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] @@ -202,7 +202,7 @@ def ModifyDt_FromAarray(particles, fieldset): # pragma: no cover runtime = sum(dt_in_kernel) expected_num_kernel_loops = len(dt_in_kernel) - pset.execute(ModifyDt_FromAarray, runtime=runtime, dt=1.0) + 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) From 9341967a17e885216204628e72fb403c5df0831f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 2 Nov 2026 17:38:44 +0100 Subject: [PATCH 12/12] Improve dt clipping Based on feedback at https://github.com/Parcels-code/Parcels/pull/3019#discussion_r4424868826 --- src/parcels/_core/kernel.py | 36 ++++++++++++++++++++---------------- tests/test_kernel.py | 15 +++++++++++++++ 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 9722e0d3f..837751fea 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -37,11 +37,19 @@ } -def _clip_dt_to_endtime(dt, time_to_endtime, compute_time_direction): +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: - return np.maximum(np.minimum(dt, time_to_endtime), 0) - return np.minimum(np.maximum(dt, -time_to_endtime), 0) + 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: @@ -182,11 +190,10 @@ def execute(self, pset, endtime): 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") + 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") - dt_to_restore = pset.dt.copy() while (len(pset) > 0) and np.any(np.isin(pset.state, [StatusCode.Success, StatusCode.Repeat])): time_to_endtime = sign_dt * (endtime - pset.t) @@ -195,11 +202,10 @@ def execute(self, pset, endtime): if not np.any(evaluate_particles): return StatusCode.Success - if preserve_particle_dt: - pset.dt = dt_to_restore + dt_to_restore = [] # adapt dt to end exactly on endtime - pset.dt = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt) + 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: @@ -218,23 +224,21 @@ def execute(self, pset, endtime): # (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) + 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 the retained dt (unless in RK45 mode) - if preserve_particle_dt: - pset.dt = dt_to_restore - # Set particle state for particles that reached endtime particles_endofloop = (pset.state == StatusCode.Success) & (pset.t == endtime) pset[particles_endofloop].state = StatusCode.EndofLoop - # delete particles that signalled deletion and remove their corresponding dt_to_restore + if preserve_particle_dt: + for indices, values in dt_to_restore: + pset.dt[indices] = values + + # delete particles that signalled deletion indices = np.flatnonzero(pset.state == StatusCode.Delete) if len(indices) > 0: pset.remove_indices(indices) - if preserve_particle_dt: - dt_to_restore = np.delete(dt_to_restore, indices) # check and throw errors if np.any(pset.state == StatusCode.StopAllExecution): diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 3cb30094c..c73c7cef6 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -208,6 +208,21 @@ def ModifyDt_FromArray(particles, fieldset): # pragma: no cover 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]) + + def ModifyDt(particles, fieldset): # pragma: no cover + particles.age += particles.dt + particles.dt = 2 + + runtime = 10 + 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) + + @pytest.mark.parametrize("dt", [1e-2, 1e-5, 1e-6, 1e-9]) def test_small_dt(fieldset, dt): pset = ParticleSet(fieldset, x=[0], y=[0])