Pushing the time interpolation also to Interpolators

Author: erikvansebille

parcels/_interpolation.py

@@ -14,6 +14,8 @@ class InterpolationContext2D:
     ----------
     data: np.ndarray
         field data of shape (time, y, x)
+    tau: float
+        time interpolation coordinate in unit length
     eta: float
         y-direction interpolation coordinate in unit cube (between 0 and 1)
     xsi: float
@@ -24,15 +26,19 @@ class InterpolationContext2D:
         y index of cell containing particle
     xi: int
         x index of cell containing particle
+    interptime: bool = True
+        whether to interpolate in time
 
     """
 
     data: np.ndarray
+    tau: float
     eta: float
     xsi: float
     ti: int
     yi: int
     xi: int
+    interptime: bool = True
 
 
 @dataclass
@@ -45,6 +51,8 @@ class InterpolationContext3D:
         field data of shape (time, z, y, x). This needs to be complete in the vertical
         direction as some interpolation methods need to know whether they are at the
         surface or bottom.
+    tau: float
+        time interpolation coordinate in unit length
     zeta: float
         vertical interpolation coordinate in unit cube
     eta: float
@@ -61,10 +69,13 @@ class InterpolationContext3D:
         x index of cell containing particle
     gridindexingtype: GridIndexingType
         grid indexing type
+    interptime: bool = True
+        whether to interpolate in time
 
     """
 
     data: np.ndarray
+    tau: float
     zeta: float
     eta: float
     xsi: float
@@ -73,6 +84,7 @@ class InterpolationContext3D:
     yi: int
     xi: int
     gridindexingtype: GridIndexingType  # included in 3D as z-face is indexed differently with MOM5 and POP
+    interptime: bool = True
 
 
 _interpolator_registry_2d: dict[str, Callable[[InterpolationContext2D], float]] = {}
@@ -110,7 +122,11 @@ def decorator(interpolator: Callable[[InterpolationContext3D], float]):
 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
-    return ctx.data[ctx.ti, yii, xii]
+    ft0 = ctx.data[ctx.ti, yii, xii]
+    if ctx.interptime:
+        ft1 = ctx.data[ctx.ti + 1, yii, xii]
+        return (1 - ctx.tau) * ft0 + ctx.tau * ft1
+    return ft0
 
 
 def _interp_on_unit_square(*, eta: float, xsi: float, data: np.ndarray, yi: int, xi: int) -> float:
@@ -128,11 +144,15 @@ def _interp_on_unit_square(*, eta: float, xsi: float, data: np.ndarray, yi: int,
 @register_2d_interpolator("partialslip")
 @register_2d_interpolator("freeslip")
 def _linear_2d(ctx: InterpolationContext2D) -> float:
-    return _interp_on_unit_square(eta=ctx.eta, xsi=ctx.xsi, data=ctx.data[ctx.ti, :, :], yi=ctx.yi, xi=ctx.xi)
+    ft0 = _interp_on_unit_square(eta=ctx.eta, xsi=ctx.xsi, data=ctx.data[ctx.ti, :, :], yi=ctx.yi, xi=ctx.xi)
+    if ctx.interptime:
+        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
+    return ft0
 
 
 @register_2d_interpolator("linear_invdist_land_tracer")
-def _linear_invdist_land_tracer_2d(ctx: InterpolationContext2D) -> float:
+def _linear_invdist_land_tracer_2d(ctx: InterpolationContext2D) -> float:  # TODO make time-varying
     xsi = ctx.xsi
     eta = ctx.eta
     data = ctx.data
@@ -166,19 +186,27 @@ def _linear_invdist_land_tracer_2d(ctx: InterpolationContext2D) -> float:
 @register_2d_interpolator("cgrid_tracer")
 @register_2d_interpolator("bgrid_tracer")
 def _tracer_2d(ctx: InterpolationContext2D) -> float:
-    return ctx.data[ctx.ti, ctx.yi + 1, ctx.xi + 1]
+    ft0 = ctx.data[ctx.ti, ctx.yi + 1, ctx.xi + 1]
+    if ctx.interptime:
+        ft1 = ctx.data[ctx.ti + 1, ctx.yi + 1, ctx.xi + 1]
+        return (1 - ctx.tau) * ft0 + ctx.tau * ft1
+    return ft0
 
 
 @register_3d_interpolator("nearest")
 def _nearest_3d(ctx: InterpolationContext3D) -> 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
     zii = ctx.zi if ctx.zeta <= 0.5 else ctx.zi + 1
-    return ctx.data[ctx.ti, zii, yii, xii]
+    ft0 = ctx.data[ctx.ti, zii, yii, xii]
+    if ctx.interptime:
+        ft1 = ctx.data[ctx.ti + 1, zii, yii, xii]
+        return (1 - ctx.tau) * ft0 + ctx.tau * ft1
+    return ft0
 
 
 @register_3d_interpolator("cgrid_velocity")
-def _cgrid_velocity_3d(ctx: InterpolationContext3D) -> float:
+def _cgrid_velocity_3d(ctx: InterpolationContext3D) -> float:  # TODO make time-varying
     # evaluating W velocity in c_grid
     if ctx.gridindexingtype == "nemo":
         f0 = ctx.data[ctx.ti, ctx.zi, ctx.yi + 1, ctx.xi + 1]
@@ -190,7 +218,7 @@ def _cgrid_velocity_3d(ctx: InterpolationContext3D) -> float:
 
 
 @register_3d_interpolator("linear_invdist_land_tracer")
-def _linear_invdist_land_tracer_3d(ctx: InterpolationContext3D) -> float:
+def _linear_invdist_land_tracer_3d(ctx: InterpolationContext3D) -> float:  # TODO make time-varying
     land = np.isclose(ctx.data[ctx.ti, ctx.zi : ctx.zi + 2, ctx.yi : ctx.yi + 2, ctx.xi : ctx.xi + 2], 0.0)
     nb_land = np.sum(land)
     if nb_land == 8:
@@ -250,7 +278,7 @@ def _z_layer_interp(
 @register_3d_interpolator("linear")
 @register_3d_interpolator("partialslip")
 @register_3d_interpolator("freeslip")
-def _linear_3d(ctx: InterpolationContext3D) -> float:
+def _linear_3d(ctx: InterpolationContext3D) -> float:  # TODO make time-varying
     zdim = ctx.data.shape[1]
     data_3d = ctx.data[ctx.ti, :, :, :]
     f0, f1 = _get_3d_f0_f1(eta=ctx.eta, xsi=ctx.xsi, data=data_3d, zi=ctx.zi, yi=ctx.yi, xi=ctx.xi)
@@ -277,4 +305,8 @@ def _linear_3d_bgrid_w_velocity(ctx: InterpolationContext3D) -> float:
 @register_3d_interpolator("bgrid_tracer")
 @register_3d_interpolator("cgrid_tracer")
 def _tracer_3d(ctx: InterpolationContext3D) -> float:
-    return ctx.data[ctx.ti, ctx.zi, ctx.yi + 1, ctx.xi + 1]
+    ft0 = ctx.data[ctx.ti, ctx.zi, ctx.yi + 1, ctx.xi + 1]
+    if ctx.interptime:
+        ft1 = ctx.data[ctx.ti + 1, ctx.zi, ctx.yi + 1, ctx.xi + 1]
+        return (1 - ctx.tau) * ft0 + ctx.tau * ft1
+    return ft0

parcels/field.py

@@ -893,15 +893,9 @@ def _interpolator2D(self, time, z, y, x, particle=None):
 
         (tau, _, eta, xsi, ti, _, yi, xi) = self._search_indices(time, z, y, x, particle=particle)
 
-        ctx = InterpolationContext2D(self.data, eta, xsi, ti, yi, xi)
-        f0 = f(ctx)
-
-        if ti < self.grid.tdim - 1 and time > self.grid.time[ti]:
-            ctx = InterpolationContext2D(self.data, eta, xsi, ti + 1, yi, xi)
-            f1 = f(ctx)
-            return f0 * (1 - tau) + f1 * tau
-        else:
-            return f0
+        interptime = True if (ti < self.grid.tdim - 1 and tau > 0) else False
+        ctx = InterpolationContext2D(self.data, tau, eta, xsi, ti, yi, xi, interptime=interptime)
+        return f(ctx)
 
     def _interpolator3D(self, time, z, y, x, particle=None):
         """Impelement 3D interpolation with coordinate transformations as seen in Delandmeter and Van Sebille (2019), 10.5194/gmd-12-3571-2019.."""
@@ -912,15 +906,11 @@ def _interpolator3D(self, time, z, y, x, particle=None):
 
         (tau, zeta, eta, xsi, ti, zi, yi, xi) = self._search_indices(time, z, y, x, particle=particle)
 
-        ctx = InterpolationContext3D(self.data, zeta, eta, xsi, ti, zi, yi, xi, self.gridindexingtype)
-        f0 = f(ctx)
-
-        if ti < self.grid.tdim - 1 and time > self.grid.time[ti]:
-            ctx = InterpolationContext3D(self.data, zeta, eta, xsi, ti + 1, zi, yi, xi, self.gridindexingtype)
-            f1 = f(ctx)
-            return f0 * (1 - tau) + f1 * tau
-        else:
-            return f0
+        interptime = True if (ti < self.grid.tdim - 1 and tau > 0) else False
+        ctx = InterpolationContext3D(
+            self.data, tau, zeta, eta, xsi, ti, zi, yi, xi, self.gridindexingtype, interptime=interptime
+        )
+        return f(ctx)
 
     def _spatial_interpolation(self, time, z, y, x, particle=None):
         """Interpolate spatial field values."""