Merge branch 'v4-dev' into kernelloop-explanation

Author: reint-fischer

docs/getting_started/tutorial_quickstart.md

@@ -99,7 +99,7 @@ TODO: link to a list of included kernels
 ```
 
 ```{code-cell}
-kernels = [parcels.kernels.AdvectionEE]
+kernels = [parcels.kernels.AdvectionRK2]
 ```
 
 ## Prepare output: `ParticleFile`

src/parcels/kernels/advection.py

@@ -18,101 +18,120 @@
 ]
 
 
+def _constrain_dt_to_within_time_interval(time_interval, time, dt):
+    """Helper function to make sure dt does not go outside time_interval.
+
+    This is especially relevant for higher-order RK methods (RK2, RK4, RK45),
+    which require interpolations at time + dt. If time is at the edges of the
+    time_interval (typically the last integration step), such an operation would
+    lead to an OutofTimeError.
+    """
+    if time_interval:
+        dt = np.where(time + dt <= time_interval.right, dt, time_interval.right - time)
+        dt = np.where(time + dt >= time_interval.left, dt, time - time_interval.left)
+    return dt.astype("timedelta64[s]")
+
+
 def AdvectionRK2(particles, fieldset):  # pragma: no cover
     """Advection of particles using second-order Runge-Kutta integration."""
-    dt = particles.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
     (u1, v1) = fieldset.UV[particles]
-    lon1, lat1 = (particles.lon + u1 * 0.5 * dt, particles.lat + v1 * 0.5 * dt)
-    (u2, v2) = fieldset.UV[particles.time + 0.5 * particles.dt, particles.z, lat1, lon1, particles]
-    particles.dlon += u2 * dt
-    particles.dlat += v2 * dt
+    lon1, lat1 = (particles.lon + u1 * 0.5 * dt_flt, particles.lat + v1 * 0.5 * dt_flt)
+    (u2, v2) = fieldset.UV[particles.time + 0.5 * dt, particles.z, lat1, lon1, particles]
+    particles.dlon += u2 * dt_flt
+    particles.dlat += v2 * dt_flt
 
 
 def AdvectionRK2_3D(particles, fieldset):  # pragma: no cover
     """Advection of particles using second-order Runge-Kutta integration including vertical velocity."""
-    dt = particles.dt / np.timedelta64(1, "s")
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")
     (u1, v1, w1) = fieldset.UVW[particles]
-    lon1 = particles.lon + u1 * 0.5 * dt
-    lat1 = particles.lat + v1 * 0.5 * dt
-    z1 = particles.z + w1 * 0.5 * dt
-    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * particles.dt, z1, lat1, lon1, particles]
-    particles.dlon += u2 * dt
-    particles.dlat += v2 * dt
-    particles.dz += w2 * dt
+    lon1 = particles.lon + u1 * 0.5 * dt_flt
+    lat1 = particles.lat + v1 * 0.5 * dt_flt
+    z1 = particles.z + w1 * 0.5 * dt_flt
+    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * dt, z1, lat1, lon1, particles]
+    particles.dlon += u2 * dt_flt
+    particles.dlat += v2 * dt_flt
+    particles.dz += w2 * dt_flt
 
 
 def AdvectionRK4(particles, fieldset):  # pragma: no cover
     """Advection of particles using fourth-order Runge-Kutta integration."""
-    dt = particles.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
     (u1, v1) = fieldset.UV[particles]
-    lon1, lat1 = (particles.lon + u1 * 0.5 * dt, particles.lat + v1 * 0.5 * dt)
-    (u2, v2) = fieldset.UV[particles.time + 0.5 * particles.dt, particles.z, lat1, lon1, particles]
-    lon2, lat2 = (particles.lon + u2 * 0.5 * dt, particles.lat + v2 * 0.5 * dt)
-    (u3, v3) = fieldset.UV[particles.time + 0.5 * particles.dt, particles.z, lat2, lon2, particles]
-    lon3, lat3 = (particles.lon + u3 * dt, particles.lat + v3 * dt)
-    (u4, v4) = fieldset.UV[particles.time + particles.dt, particles.z, lat3, lon3, particles]
-    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6.0 * dt
-    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6.0 * dt
+    lon1, lat1 = (particles.lon + u1 * 0.5 * dt_flt, particles.lat + v1 * 0.5 * dt_flt)
+    (u2, v2) = fieldset.UV[particles.time + 0.5 * dt, particles.z, lat1, lon1, particles]
+    lon2, lat2 = (particles.lon + u2 * 0.5 * dt_flt, particles.lat + v2 * 0.5 * dt_flt)
+    (u3, v3) = fieldset.UV[particles.time + 0.5 * dt, particles.z, lat2, lon2, particles]
+    lon3, lat3 = (particles.lon + u3 * dt_flt, particles.lat + v3 * dt_flt)
+    (u4, v4) = fieldset.UV[particles.time + dt, particles.z, lat3, lon3, particles]
+    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6.0 * dt_flt
+    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6.0 * dt_flt
 
 
 def AdvectionRK4_3D(particles, fieldset):  # pragma: no cover
     """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity."""
-    dt = particles.dt / np.timedelta64(1, "s")
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")
     (u1, v1, w1) = fieldset.UVW[particles]
-    lon1 = particles.lon + u1 * 0.5 * dt
-    lat1 = particles.lat + v1 * 0.5 * dt
-    z1 = particles.z + w1 * 0.5 * dt
-    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * particles.dt, z1, lat1, lon1, particles]
-    lon2 = particles.lon + u2 * 0.5 * dt
-    lat2 = particles.lat + v2 * 0.5 * dt
-    z2 = particles.z + w2 * 0.5 * dt
-    (u3, v3, w3) = fieldset.UVW[particles.time + 0.5 * particles.dt, z2, lat2, lon2, particles]
-    lon3 = particles.lon + u3 * dt
-    lat3 = particles.lat + v3 * dt
-    z3 = particles.z + w3 * dt
-    (u4, v4, w4) = fieldset.UVW[particles.time + particles.dt, z3, lat3, lon3, particles]
-    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt
-    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt
-    particles.dz += (w1 + 2 * w2 + 2 * w3 + w4) / 6 * dt
+    lon1 = particles.lon + u1 * 0.5 * dt_flt
+    lat1 = particles.lat + v1 * 0.5 * dt_flt
+    z1 = particles.z + w1 * 0.5 * dt_flt
+    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * dt, z1, lat1, lon1, particles]
+    lon2 = particles.lon + u2 * 0.5 * dt_flt
+    lat2 = particles.lat + v2 * 0.5 * dt_flt
+    z2 = particles.z + w2 * 0.5 * dt_flt
+    (u3, v3, w3) = fieldset.UVW[particles.time + 0.5 * dt, z2, lat2, lon2, particles]
+    lon3 = particles.lon + u3 * dt_flt
+    lat3 = particles.lat + v3 * dt_flt
+    z3 = particles.z + w3 * dt_flt
+    (u4, v4, w4) = fieldset.UVW[particles.time + dt, z3, lat3, lon3, particles]
+    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt_flt
+    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt_flt
+    particles.dz += (w1 + 2 * w2 + 2 * w3 + w4) / 6 * dt_flt
 
 
 def AdvectionRK4_3D_CROCO(particles, fieldset):  # pragma: no cover
     """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity.
     This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers.
     """
-    dt = particles.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
     sig_dep = particles.z / fieldset.H[particles.time, 0, particles.lat, particles.lon]
 
     (u1, v1, w1) = fieldset.UVW[particles.time, particles.z, particles.lat, particles.lon, particles]
     w1 *= sig_dep / fieldset.H[particles.time, 0, particles.lat, particles.lon]
-    lon1 = particles.lon + u1 * 0.5 * dt
-    lat1 = particles.lat + v1 * 0.5 * dt
-    sig_dep1 = sig_dep + w1 * 0.5 * dt
+    lon1 = particles.lon + u1 * 0.5 * dt_flt
+    lat1 = particles.lat + v1 * 0.5 * dt_flt
+    sig_dep1 = sig_dep + w1 * 0.5 * dt_flt
     dep1 = sig_dep1 * fieldset.H[particles.time, 0, lat1, lon1]
 
-    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * particles.dt, dep1, lat1, lon1, particles]
+    (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * dt, dep1, lat1, lon1, particles]
     w2 *= sig_dep1 / fieldset.H[particles.time, 0, lat1, lon1]
-    lon2 = particles.lon + u2 * 0.5 * dt
-    lat2 = particles.lat + v2 * 0.5 * dt
-    sig_dep2 = sig_dep + w2 * 0.5 * dt
+    lon2 = particles.lon + u2 * 0.5 * dt_flt
+    lat2 = particles.lat + v2 * 0.5 * dt_flt
+    sig_dep2 = sig_dep + w2 * 0.5 * dt_flt
     dep2 = sig_dep2 * fieldset.H[particles.time, 0, lat2, lon2]
 
-    (u3, v3, w3) = fieldset.UVW[particles.time + 0.5 * particles.dt, dep2, lat2, lon2, particles]
+    (u3, v3, w3) = fieldset.UVW[particles.time + 0.5 * dt, dep2, lat2, lon2, particles]
     w3 *= sig_dep2 / fieldset.H[particles.time, 0, lat2, lon2]
-    lon3 = particles.lon + u3 * dt
-    lat3 = particles.lat + v3 * dt
-    sig_dep3 = sig_dep + w3 * dt
+    lon3 = particles.lon + u3 * dt_flt
+    lat3 = particles.lat + v3 * dt_flt
+    sig_dep3 = sig_dep + w3 * dt_flt
     dep3 = sig_dep3 * fieldset.H[particles.time, 0, lat3, lon3]
 
-    (u4, v4, w4) = fieldset.UVW[particles.time + particles.dt, dep3, lat3, lon3, particles]
+    (u4, v4, w4) = fieldset.UVW[particles.time + dt, dep3, lat3, lon3, particles]
     w4 *= sig_dep3 / fieldset.H[particles.time, 0, lat3, lon3]
-    lon4 = particles.lon + u4 * dt
-    lat4 = particles.lat + v4 * dt
-    sig_dep4 = sig_dep + w4 * dt
+    lon4 = particles.lon + u4 * dt_flt
+    lat4 = particles.lat + v4 * dt_flt
+    sig_dep4 = sig_dep + w4 * dt_flt
     dep4 = sig_dep4 * fieldset.H[particles.time, 0, lat4, lon4]
 
-    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt
-    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt
+    particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt_flt
+    particles.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt_flt
     particles.dz += (
         (dep1 - particles.z) * 2 + 2 * (dep2 - particles.z) * 2 + 2 * (dep3 - particles.z) + dep4 - particles.z
     ) / 6
@@ -135,8 +154,9 @@ def AdvectionRK45(particles, fieldset):  # pragma: no cover
     Time-step dt is halved if error is larger than fieldset.RK45_tol,
     and doubled if error is smaller than 1/10th of tolerance.
     """
-    dt = particles.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
-    sign_dt = np.sign(dt)
+    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)
+    dt_flt = dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
+    sign_dt = np.sign(dt_flt)
 
     c = [1.0 / 4.0, 3.0 / 8.0, 12.0 / 13.0, 1.0, 1.0 / 2.0]
     A = [
@@ -150,42 +170,42 @@ def AdvectionRK45(particles, fieldset):  # pragma: no cover
     b5 = [16.0 / 135.0, 0.0, 6656.0 / 12825.0, 28561.0 / 56430.0, -9.0 / 50.0, 2.0 / 55.0]
 
     (u1, v1) = fieldset.UV[particles]
-    lon1, lat1 = (particles.lon + u1 * A[0][0] * dt, particles.lat + v1 * A[0][0] * dt)
-    (u2, v2) = fieldset.UV[particles.time + c[0] * particles.dt, particles.z, lat1, lon1, particles]
+    lon1, lat1 = (particles.lon + u1 * A[0][0] * dt_flt, particles.lat + v1 * A[0][0] * dt_flt)
+    (u2, v2) = fieldset.UV[particles.time + c[0] * dt, particles.z, lat1, lon1, particles]
     lon2, lat2 = (
-        particles.lon + (u1 * A[1][0] + u2 * A[1][1]) * dt,
-        particles.lat + (v1 * A[1][0] + v2 * A[1][1]) * dt,
+        particles.lon + (u1 * A[1][0] + u2 * A[1][1]) * dt_flt,
+        particles.lat + (v1 * A[1][0] + v2 * A[1][1]) * dt_flt,
     )
-    (u3, v3) = fieldset.UV[particles.time + c[1] * particles.dt, particles.z, lat2, lon2, particles]
+    (u3, v3) = fieldset.UV[particles.time + c[1] * dt, particles.z, lat2, lon2, particles]
     lon3, lat3 = (
-        particles.lon + (u1 * A[2][0] + u2 * A[2][1] + u3 * A[2][2]) * dt,
-        particles.lat + (v1 * A[2][0] + v2 * A[2][1] + v3 * A[2][2]) * dt,
+        particles.lon + (u1 * A[2][0] + u2 * A[2][1] + u3 * A[2][2]) * dt_flt,
+        particles.lat + (v1 * A[2][0] + v2 * A[2][1] + v3 * A[2][2]) * dt_flt,
     )
-    (u4, v4) = fieldset.UV[particles.time + c[2] * particles.dt, particles.z, lat3, lon3, particles]
+    (u4, v4) = fieldset.UV[particles.time + c[2] * dt, particles.z, lat3, lon3, particles]
     lon4, lat4 = (
-        particles.lon + (u1 * A[3][0] + u2 * A[3][1] + u3 * A[3][2] + u4 * A[3][3]) * dt,
-        particles.lat + (v1 * A[3][0] + v2 * A[3][1] + v3 * A[3][2] + v4 * A[3][3]) * dt,
+        particles.lon + (u1 * A[3][0] + u2 * A[3][1] + u3 * A[3][2] + u4 * A[3][3]) * dt_flt,
+        particles.lat + (v1 * A[3][0] + v2 * A[3][1] + v3 * A[3][2] + v4 * A[3][3]) * dt_flt,
     )
-    (u5, v5) = fieldset.UV[particles.time + c[3] * particles.dt, particles.z, lat4, lon4, particles]
+    (u5, v5) = fieldset.UV[particles.time + c[3] * dt, particles.z, lat4, lon4, particles]
     lon5, lat5 = (
-        particles.lon + (u1 * A[4][0] + u2 * A[4][1] + u3 * A[4][2] + u4 * A[4][3] + u5 * A[4][4]) * dt,
-        particles.lat + (v1 * A[4][0] + v2 * A[4][1] + v3 * A[4][2] + v4 * A[4][3] + v5 * A[4][4]) * dt,
+        particles.lon + (u1 * A[4][0] + u2 * A[4][1] + u3 * A[4][2] + u4 * A[4][3] + u5 * A[4][4]) * dt_flt,
+        particles.lat + (v1 * A[4][0] + v2 * A[4][1] + v3 * A[4][2] + v4 * A[4][3] + v5 * A[4][4]) * dt_flt,
     )
-    (u6, v6) = fieldset.UV[particles.time + c[4] * particles.dt, particles.z, lat5, lon5, particles]
+    (u6, v6) = fieldset.UV[particles.time + c[4] * dt, particles.z, lat5, lon5, particles]
 
-    lon_4th = (u1 * b4[0] + u2 * b4[1] + u3 * b4[2] + u4 * b4[3] + u5 * b4[4]) * dt
-    lat_4th = (v1 * b4[0] + v2 * b4[1] + v3 * b4[2] + v4 * b4[3] + v5 * b4[4]) * dt
-    lon_5th = (u1 * b5[0] + u2 * b5[1] + u3 * b5[2] + u4 * b5[3] + u5 * b5[4] + u6 * b5[5]) * dt
-    lat_5th = (v1 * b5[0] + v2 * b5[1] + v3 * b5[2] + v4 * b5[3] + v5 * b5[4] + v6 * b5[5]) * dt
+    lon_4th = (u1 * b4[0] + u2 * b4[1] + u3 * b4[2] + u4 * b4[3] + u5 * b4[4]) * dt_flt
+    lat_4th = (v1 * b4[0] + v2 * b4[1] + v3 * b4[2] + v4 * b4[3] + v5 * b4[4]) * dt_flt
+    lon_5th = (u1 * b5[0] + u2 * b5[1] + u3 * b5[2] + u4 * b5[3] + u5 * b5[4] + u6 * b5[5]) * dt_flt
+    lat_5th = (v1 * b5[0] + v2 * b5[1] + v3 * b5[2] + v4 * b5[3] + v5 * b5[4] + v6 * b5[5]) * dt_flt
 
     kappa = np.sqrt(np.pow(lon_5th - lon_4th, 2) + np.pow(lat_5th - lat_4th, 2))
 
-    good_particles = (kappa <= fieldset.RK45_tol) | (np.fabs(dt) <= np.fabs(fieldset.RK45_min_dt))
+    good_particles = (kappa <= fieldset.RK45_tol) | (np.fabs(dt_flt) <= np.fabs(fieldset.RK45_min_dt))
     particles.dlon += np.where(good_particles, lon_5th, 0)
     particles.dlat += np.where(good_particles, lat_5th, 0)
 
     increase_dt_particles = (
-        good_particles & (kappa <= fieldset.RK45_tol / 10) & (np.fabs(dt * 2) <= np.fabs(fieldset.RK45_max_dt))
+        good_particles & (kappa <= fieldset.RK45_tol / 10) & (np.fabs(dt_flt * 2) <= np.fabs(fieldset.RK45_max_dt))
     )
     particles.dt = np.where(increase_dt_particles, particles.dt * 2, particles.dt)
     particles.dt = np.where(

tests/test_advection.py

@@ -467,7 +467,7 @@ def test_nemo_curvilinear_fieldset():
     runtime = np.timedelta64(160, "D")
 
     pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp)
-    pset.execute(AdvectionEE, runtime=runtime, dt=np.timedelta64(6, "h"))
+    pset.execute(AdvectionEE, runtime=runtime, dt=np.timedelta64(10, "D"))
     np.testing.assert_allclose(pset.lat, latp, atol=1e-1)
 
 

tests/test_particleset_execute.py

@@ -23,7 +23,7 @@
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode, XLinear
-from parcels.kernels import AdvectionEE, AdvectionRK4, AdvectionRK4_3D
+from parcels.kernels import AdvectionEE, AdvectionRK2, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45
 from tests.common_kernels import DoNothing
 
 
@@ -162,6 +162,25 @@ def SampleU(particles, fieldset):  # pragma: no cover
     assert pset[0].time == endtime
 
 
+@pytest.mark.parametrize("kernel", [AdvectionEE, AdvectionRK2, AdvectionRK4, AdvectionRK45])
+@pytest.mark.parametrize("dt", [np.timedelta64(10, "D"), np.timedelta64(1, "D")])
+def test_particleset_run_RK_to_endtime_fwd_bwd(fieldset, kernel, dt):
+    """Test that RK kernels can be run to the endtime of a fieldset (and not throw OutsideTimeInterval)"""
+    starttime = fieldset.time_interval.left
+    endtime = fieldset.time_interval.right
+
+    # Setting zero velocities to avoid OutofBoundsErrors
+    fieldset.U.data[:] = 0.0
+    fieldset.V.data[:] = 0.0
+
+    pset = ParticleSet(fieldset, lon=[0.2], lat=[5.0], time=[starttime])
+    pset.execute(kernel, endtime=endtime, dt=dt)
+    assert pset[0].time == endtime
+
+    pset.execute(kernel, endtime=starttime, dt=-dt)
+    assert pset[0].time == starttime
+
+
 def test_particleset_interpolate_on_domainedge(zonal_flow_fieldset):
     fieldset = zonal_flow_fieldset