cleaning up cgrid interpolation

Author: erikvansebille

parcels/_interpolation.py

@@ -201,16 +201,31 @@ def _nearest_3d(ctx: InterpolationContext3D) -> float:
     return (1 - ctx.tau) * ft0 + ctx.tau * ft1
 
 
+def _get_cgrid_depth_point(*, zeta: float, data: np.ndarray, zi: int, yi: int, xi: int) -> tuple[float]:
+    f0 = data[zi, yi, xi]
+    f1 = data[zi + 1, yi, xi]
+    return (1 - zeta) * f0 + zeta * f1
+
+
 @register_3d_interpolator("cgrid_velocity")
-def _cgrid_velocity_3d(ctx: InterpolationContext3D) -> float:  # TODO make time-varying
+def _cgrid_W_velocity_3d(ctx: InterpolationContext3D) -> float:
     # evaluating W velocity in c_grid
     if ctx.gridindexingtype == "nemo":
-        f0 = ctx.data[ctx.ti, ctx.zi, ctx.yi + 1, ctx.xi + 1]
-        f1 = ctx.data[ctx.ti, ctx.zi + 1, ctx.yi + 1, ctx.xi + 1]
+        ft0 = _get_cgrid_depth_point(
+            zeta=ctx.zeta, data=ctx.data[ctx.ti, :, :, :], zi=ctx.zi, yi=ctx.yi + 1, xi=ctx.xi + 1
+        )
     elif ctx.gridindexingtype in ["mitgcm", "croco"]:
-        f0 = ctx.data[ctx.ti, ctx.zi, ctx.yi, ctx.xi]
-        f1 = ctx.data[ctx.ti, ctx.zi + 1, ctx.yi, ctx.xi]
-    return (1 - ctx.zeta) * f0 + ctx.zeta * f1
+        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:
+        return ft0
+
+    if ctx.gridindexingtype == "nemo":
+        ft1 = _get_cgrid_depth_point(
+            zeta=ctx.zeta, data=ctx.data[ctx.ti + 1, :, :, :], zi=ctx.zi, yi=ctx.yi + 1, xi=ctx.xi + 1
+        )
+    elif ctx.gridindexingtype in ["mitgcm", "croco"]:
+        ft1 = _get_cgrid_depth_point(zeta=ctx.zeta, data=ctx.data[ctx.ti + 1, :, :, :], zi=ctx.zi, yi=ctx.yi, xi=ctx.xi)
+    return (1 - ctx.tau) * ft0 + ctx.tau * ft1
 
 
 @register_3d_interpolator("linear_invdist_land_tracer")

parcels/field.py

@@ -1526,66 +1526,25 @@ def eval(self, time, z, y, x, particle=None, applyConversion=True):
             if applyConversion:
                 u = self.U.units.to_target(u, z, y, x)
                 v = self.V.units.to_target(v, z, y, x)
-            if "3D" in self.vector_type:
-                w = self.W.eval(time, z, y, x, particle=particle, applyConversion=False)
-                if applyConversion:
-                    w = self.W.units.to_target(w, z, y, x)
-                return (u, v, w)
-            else:
-                return (u, v)
         else:
             interp = {
-                "cgrid_velocity": {
-                    "2D": self.spatial_c_grid_interpolation2D,
-                    "3D": self.spatial_c_grid_interpolation3D,
-                },
-                "partialslip": {"2D": self.spatial_slip_interpolation, "3D": self.spatial_slip_interpolation},
-                "freeslip": {"2D": self.spatial_slip_interpolation, "3D": self.spatial_slip_interpolation},
+                "cgrid_velocity": self.spatial_c_grid_interpolation2D,
+                "partialslip": self.spatial_slip_interpolation,
+                "freeslip": self.spatial_slip_interpolation,
             }
-            grid = self.U.grid
             tau, ti = self.U._time_index(time)
-            if ti < grid.tdim - 1 and time > grid.time[ti]:
-                t0 = grid.time[ti]
-                t1 = grid.time[ti + 1]
-                if "3D" in self.vector_type:
-                    (u0, v0, w0) = interp[self.U.interp_method]["3D"](
-                        ti,
-                        z,
-                        y,
-                        x,
-                        t0,
-                        particle=particle,
-                        applyConversion=applyConversion,  # TODO see if we can directly call time interpolation for W here
-                    )
-                    (u1, v1, w1) = interp[self.U.interp_method]["3D"](
-                        ti + 1, z, y, x, t1, particle=particle, applyConversion=applyConversion
-                    )
-                    w = w0 * (1 - tau) + w1 * tau
-                else:
-                    (u0, v0) = interp[self.U.interp_method]["2D"](
-                        ti, z, y, x, t0, particle=particle, applyConversion=applyConversion
-                    )
-                    (u1, v1) = interp[self.U.interp_method]["2D"](
-                        ti + 1, z, y, x, t1, particle=particle, applyConversion=applyConversion
-                    )
-                u = u0 * (1 - tau) + u1 * tau
-                v = v0 * (1 - tau) + v1 * tau
-                if "3D" in self.vector_type:
-                    return (u, v, w)
-                else:
-                    return (u, v)
-            else:
-                # Skip temporal interpolation if time is outside
-                # of the defined time range or if we have hit an
-                # exact value in the time array.
-                if "3D" in self.vector_type:
-                    return interp[self.U.interp_method]["3D"](
-                        ti, z, y, x, grid.time[ti], particle=particle, applyConversion=applyConversion
-                    )
-                else:
-                    return interp[self.U.interp_method]["2D"](
-                        ti, z, y, x, grid.time[ti], particle=particle, applyConversion=applyConversion
-                    )
+            (u, v) = interp[self.U.interp_method](ti, z, y, x, time, particle=particle, applyConversion=applyConversion)
+            if ti < self.U.grid.tdim - 1 and time > self.U.grid.time[ti]:
+                (u1, v1) = interp[self.U.interp_method](
+                    ti + 1, z, y, x, time, particle=particle, applyConversion=applyConversion
+                )
+                u = u * (1 - tau) + u1 * tau
+                v = v * (1 - tau) + v1 * tau
+        if "3D" in self.vector_type:
+            w = self.W.eval(time, z, y, x, particle=particle, applyConversion=applyConversion)
+            return (u, v, w)
+        else:
+            return (u, v)
 
     def __getitem__(self, key):
         try: