Adding calculate_next_ti helper function

Author: erikvansebille

parcels/_interpolation.py

@@ -4,8 +4,7 @@
 import numpy as np
 
 from parcels._typing import GridIndexingType
-
-EPS = np.finfo(float).eps
+from parcels.tools._helpers import calculate_next_ti
 
 
 @dataclass
@@ -119,7 +118,7 @@ def _nearest_2d(ctx: InterpolationContext2D) -> float:
     xii = ctx.xi if ctx.xsi <= 0.5 else ctx.xi + 1
     yii = ctx.yi if ctx.eta <= 0.5 else ctx.yi + 1
     ft0 = ctx.data[ctx.ti, yii, xii]
-    if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+    if calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
     ft1 = ctx.data[ctx.ti + 1, yii, xii]
     return (1 - ctx.tau) * ft0 + ctx.tau * ft1
@@ -141,7 +140,7 @@ def _interp_on_unit_square(*, eta: float, xsi: float, data: np.ndarray, yi: int,
 @register_2d_interpolator("freeslip")
 def _linear_2d(ctx: InterpolationContext2D) -> float:
     ft0 = _interp_on_unit_square(eta=ctx.eta, xsi=ctx.xsi, data=ctx.data[ctx.ti, :, :], yi=ctx.yi, xi=ctx.xi)
-    if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+    if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
     ft1 = _interp_on_unit_square(eta=ctx.eta, xsi=ctx.xsi, data=ctx.data[ctx.ti + 1, :, :], yi=ctx.yi, xi=ctx.xi)
     return (1 - ctx.tau) * ft0 + ctx.tau * ft1
@@ -160,7 +159,7 @@ def _linear_invdist_land_tracer_2d(ctx: InterpolationContext2D) -> float:
 
     def _get_data_temporalinterp(ti, yi, xi):
         dt0 = data[ti, yi, xi]
-        if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+        if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
             return dt0
         dt1 = data[ti + 1, yi, xi]
         return (1 - ctx.tau) * dt0 + ctx.tau * dt1
@@ -190,7 +189,7 @@ def _get_data_temporalinterp(ti, yi, xi):
 @register_2d_interpolator("bgrid_tracer")
 def _tracer_2d(ctx: InterpolationContext2D) -> float:
     ft0 = ctx.data[ctx.ti, ctx.yi + 1, ctx.xi + 1]
-    if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+    if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
     ft1 = ctx.data[ctx.ti + 1, ctx.yi + 1, ctx.xi + 1]
     return (1 - ctx.tau) * ft0 + ctx.tau * ft1
@@ -202,7 +201,7 @@ def _nearest_3d(ctx: InterpolationContext3D) -> float:
     yii = ctx.yi if ctx.eta <= 0.5 else ctx.yi + 1
     zii = ctx.zi if ctx.zeta <= 0.5 else ctx.zi + 1
     ft0 = ctx.data[ctx.ti, zii, yii, xii]
-    if ctx.tau < EPS or ctx.ti == ctx.data.shape[0] - 1:
+    if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
     ft1 = ctx.data[ctx.ti + 1, zii, yii, xii]
     return (1 - ctx.tau) * ft0 + ctx.tau * ft1
@@ -223,7 +222,7 @@ def _cgrid_W_velocity_3d(ctx: InterpolationContext3D) -> float:
         )
     elif ctx.gridindexingtype in ["mitgcm", "croco"]:
         ft0 = _get_cgrid_depth_point(zeta=ctx.zeta, data=ctx.data[ctx.ti, :, :, :], zi=ctx.zi, yi=ctx.yi, xi=ctx.xi)
-    if ctx.tau < EPS or ctx.ti == ctx.data.shape[0] - 1:
+    if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
 
     if ctx.gridindexingtype == "nemo":
@@ -242,7 +241,7 @@ def _linear_invdist_land_tracer_3d(ctx: InterpolationContext3D) -> float:
 
     def _get_data_temporalinterp(ti, zi, yi, xi):
         dt0 = ctx.data[ti, zi, yi, xi]
-        if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+        if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
             return dt0
         dt1 = data[ti + 1, zi, yi, xi]
         return (1 - ctx.tau) * dt0 + ctx.tau * dt1
@@ -308,7 +307,7 @@ def _linear_3d(ctx: InterpolationContext3D) -> float:
     zdim = ctx.data.shape[1]
     data_3d = ctx.data[ctx.ti, :, :, :]
     fz0, fz1 = _get_3d_f0_f1(eta=ctx.eta, xsi=ctx.xsi, data=data_3d, zi=ctx.zi, yi=ctx.yi, xi=ctx.xi)
-    if ctx.tau > EPS and ctx.ti < ctx.data.shape[0] - 1:
+    if calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         data_3d = ctx.data[ctx.ti + 1, :, :, :]
         fz0_t1, fz1_t1 = _get_3d_f0_f1(eta=ctx.eta, xsi=ctx.xsi, data=data_3d, zi=ctx.zi, yi=ctx.yi, xi=ctx.xi)
         fz0 = (1 - ctx.tau) * fz0 + ctx.tau * fz0_t1
@@ -338,7 +337,8 @@ def _linear_3d_bgrid_w_velocity(ctx: InterpolationContext3D) -> float:
 @register_3d_interpolator("cgrid_tracer")
 def _tracer_3d(ctx: InterpolationContext3D) -> float:
     ft0 = ctx.data[ctx.ti, ctx.zi, ctx.yi + 1, ctx.xi + 1]
-    if ctx.tau < EPS or ctx.ti >= ctx.data.shape[0] - 1:
+    if not calculate_next_ti(ctx.ti, ctx.tau, ctx.data.shape[0]):
         return ft0
-    ft1 = ctx.data[ctx.ti + 1, ctx.zi, ctx.yi + 1, ctx.xi + 1]
-    return (1 - ctx.tau) * ft0 + ctx.tau * ft1
+    else:
+        ft1 = ctx.data[ctx.ti + 1, ctx.zi, ctx.yi + 1, ctx.xi + 1]
+        return (1 - ctx.tau) * ft0 + ctx.tau * ft1

parcels/field.py

@@ -26,7 +26,7 @@
     assert_valid_gridindexingtype,
     assert_valid_interp_method,
 )
-from parcels.tools._helpers import default_repr, field_repr
+from parcels.tools._helpers import calculate_next_ti, default_repr, field_repr
 from parcels.tools.converters import (
     TimeConverter,
     UnitConverter,
@@ -1030,17 +1030,15 @@ def _calc_UV(ti, yi, xi):
             return (u, v)
 
         u, v = _calc_UV(ti, yi, xi)
-        if ti < self.U.grid.tdim - 1 and time > self.U.grid.time[ti]:
+        if calculate_next_ti(ti, tau, self.U.grid.tdim):
             ut1, vt1 = _calc_UV(ti + 1, yi, xi)
             u = (1 - tau) * u + tau * ut1
             v = (1 - tau) * v + tau * vt1
         return (u, v)
 
-    def spatial_c_grid_interpolation3D_full(self, ti, z, y, x, time, particle=None):
+    def spatial_c_grid_interpolation3D_full(self, time, z, y, x, particle=None):
         grid = self.U.grid
-        (_, zeta, eta, xsi, _, zi, yi, xi) = self.U._search_indices(
-            time, z, y, x, particle=particle
-        )  # TODO use tau here too
+        (tau, zeta, eta, xsi, ti, zi, yi, xi) = self.U._search_indices(time, z, y, x, particle=particle)
 
         if grid._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]:
             px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]])
@@ -1093,6 +1091,14 @@ def spatial_c_grid_interpolation3D_full(self, ti, z, y, x, time, particle=None):
         w0 = self.W.data[ti, zi, yi + 1, xi + 1]
         w1 = self.W.data[ti, zi + 1, yi + 1, xi + 1]
 
+        if calculate_next_ti(ti, tau, self.U.grid.tdim):
+            u0 = (1 - tau) * u0 + tau * self.U.data[ti + 1, zi, yi + 1, xi]
+            u1 = (1 - tau) * u1 + tau * self.U.data[ti + 1, zi, yi + 1, xi + 1]
+            v0 = (1 - tau) * v0 + tau * self.V.data[ti + 1, zi, yi, xi + 1]
+            v1 = (1 - tau) * v1 + tau * self.V.data[ti + 1, zi, yi + 1, xi + 1]
+            w0 = (1 - tau) * w0 + tau * self.W.data[ti + 1, zi, yi + 1, xi + 1]
+            w1 = (1 - tau) * w1 + tau * self.W.data[ti + 1, zi + 1, yi + 1, xi + 1]
+
         U0 = u0 * i_u.jacobian3D_lin_face(pz, py, px, zeta, eta, 0, "zonal", grid.mesh)
         U1 = u1 * i_u.jacobian3D_lin_face(pz, py, px, zeta, eta, 1, "zonal", grid.mesh)
         V0 = v0 * i_u.jacobian3D_lin_face(pz, py, px, zeta, 0, xsi, "meridional", grid.mesh)
@@ -1260,7 +1266,7 @@ def spatial_c_grid_interpolation3D(self, ti, z, y, x, time, particle=None, apply
         Curvilinear grids are treated properly, since the element is projected to a rectilinear parent element.
         """
         if self.U.grid._gtype in [GridType.RectilinearSGrid, GridType.CurvilinearSGrid]:
-            (u, v, w) = self.spatial_c_grid_interpolation3D_full(ti, z, y, x, time, particle=particle)
+            (u, v, w) = self.spatial_c_grid_interpolation3D_full(time, z, y, x, particle=particle)
         else:
             if self.gridindexingtype == "croco":
                 z = _croco_from_z_to_sigma_scipy(self.fieldset, time, z, y, x, particle=particle)

parcels/tools/_helpers.py

@@ -152,3 +152,8 @@ def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float:
     if isinstance(dt, np.timedelta64):
         return float(dt / np.timedelta64(1, "s"))
     return float(dt)
+
+
+def calculate_next_ti(ti, tau, tdim):
+    """Check if the time is beyond the last time in the field"""
+    return tau > np.finfo(float).eps and ti < tdim - 1