Updating helper function to return local variable dt

Also expanding docstring for clarity

Author: erikvansebille

src/parcels/kernels/advection.py

@@ -18,8 +18,14 @@
 ]
 
 
-def _adjust_RK_dt(time_interval, time, dt):
-    """Helper function to make sure dt does not go outside time_interval for RK methods"""
+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)
@@ -28,104 +34,104 @@ def _adjust_RK_dt(time_interval, time, dt):
 
 def AdvectionRK2(particles, fieldset):  # pragma: no cover
     """Advection of particles using second-order Runge-Kutta integration."""
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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."""
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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."""
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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."""
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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.
     """
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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
@@ -149,8 +155,8 @@ 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.
     """
-    particles.dt = _adjust_RK_dt(fieldset.time_interval, particles.time, particles.dt)
-    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
 
     c = [1.0 / 4.0, 3.0 / 8.0, 12.0 / 13.0, 1.0, 1.0 / 2.0]
     A = [
@@ -164,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(