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 diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 8708e7358..5d2167494 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -37,6 +37,29 @@ } +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. + + Returns + ------- + clipped_dt : + The clipped timestep array. + dt_to_restore : + List of (indices, values) pairs with the original dt of each newly clipped + particle, so the caller can restore them after the step. + """ + 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) + dt_to_restore = [(indices, dt[indices].copy())] if indices.size else [] + + return clipped_dt, dt_to_restore + + class Kernel: """Kernel object that encapsulates auto-generated code. @@ -94,16 +117,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 +183,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 +192,26 @@ 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 # 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, dt_to_restore = _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: @@ -216,18 +230,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, newly_clipped = _clip_dt_to_endtime(pset.dt, time_to_endtime, sign_dt) + dt_to_restore.extend(newly_clipped) 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): diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 192f483ff..9743f6263 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -446,7 +446,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: diff --git a/tests/test_kernel.py b/tests/test_kernel.py index bc5c4bea9..4acf908f8 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -5,7 +5,9 @@ FieldSet, KernelWarning, Particle, + ParticleFile, ParticleSet, + StatusCode, Variable, ) from parcels._core.kernel import Kernel @@ -166,17 +168,57 @@ 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) + + +@pytest.mark.parametrize("changed_dt", [1, 2, 3]) +def test_dt_modify_age_in_kernel(fieldset, changed_dt): 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 + particles.dt = changed_dt runtime = 10 - expected_age = 1 + 2 * (runtime - 2) # 1 for the first step; 2 for the remaining steps (except last) + expected_age = 1 + runtime - changed_dt # 1 for the first step; changed_dt 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)