implementing time-interpolation in _linear_3d

Author: erikvansebille

parcels/_interpolation.py

@@ -274,12 +274,17 @@ def _z_layer_interp(
 @register_3d_interpolator("linear")
 @register_3d_interpolator("partialslip")
 @register_3d_interpolator("freeslip")
-def _linear_3d(ctx: InterpolationContext3D) -> float:  # TODO make time-varying
+def _linear_3d(ctx: InterpolationContext3D) -> float:
     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)
-
-    return _z_layer_interp(zeta=ctx.zeta, f0=f0, f1=f1, zi=ctx.zi, zdim=zdim, gridindexingtype=ctx.gridindexingtype)
+    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:
+        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
+        fz1 = (1 - ctx.tau) * fz1 + ctx.tau * fz1_t1
+
+    return _z_layer_interp(zeta=ctx.zeta, f0=fz0, f1=fz1, zi=ctx.zi, zdim=zdim, gridindexingtype=ctx.gridindexingtype)
 
 
 @register_3d_interpolator("bgrid_velocity")

tests/test_interpolation.py

@@ -68,9 +68,9 @@ def data_2d():
 )
 def test_raw_2d_interpolation(data_2d, func, eta, xsi, expected):
     """Test the 2D interpolation functions on the raw arrays."""
-    ti = 0
+    tau, ti = 0, 0
     yi, xi = 1, 1
-    ctx = interpolation.InterpolationContext2D(data_2d, eta, xsi, ti, yi, xi)
+    ctx = interpolation.InterpolationContext2D(data_2d, tau, eta, xsi, ti, yi, xi)
     assert func(ctx) == expected