Fixing W interpolation for CGrid

Author: erikvansebille

parcels/application_kernels/interpolation.py

@@ -139,7 +139,7 @@ def CGrid_Velocity(
     U = vectorfield.U.data
     V = vectorfield.V.data
     grid = vectorfield.grid
-    tdim, ydim, xdim = U.shape[0], U.shape[2], U.shape[3]
+    tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3]
 
     if grid.lon.ndim == 1:
         px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]])
@@ -171,39 +171,39 @@ def CGrid_Velocity(
     # Create arrays of corner points for xarray.isel
     # TODO C grid may not need all xi and yi cornerpoints, so could speed up here?
 
-    # Time coordinates: 8 points at ti, then 8 points at ti+1
+    # Time coordinates: 4 points at ti, then 4 points at ti+1
     if lenT == 1:
-        ti = np.repeat(ti, 4)
+        ti_full = np.repeat(ti, 4)
     else:
         ti_1 = np.clip(ti + 1, 0, tdim - 1)
-        ti = np.concatenate([np.repeat(ti, 4), np.repeat(ti_1, 4)])
+        ti_full = np.concatenate([np.repeat(ti, 4), np.repeat(ti_1, 4)])
 
     # Depth coordinates: 4 points at zi, repeated for both time levels
-    zi = np.repeat(zi, lenT * 4)
+    zi_full = np.repeat(zi, lenT * 4)
 
     # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/depth
     yi_1 = np.clip(yi + 1, 0, ydim - 1)
-    yi = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT))
+    yi_full = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT))
     # # TODO check why in some cases minus needed here!!!
     # yi_minus_1 = np.clip(yi - 1, 0, ydim - 1)
     # yi = np.tile(np.repeat(np.column_stack([yi_minus_1, yi]), 2), (lenT))
 
     # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/depth
     xi_1 = np.clip(xi + 1, 0, xdim - 1)
-    xi = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT))
+    xi_full = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT))
 
     for data in [U, V]:
         axis_dim = grid.get_axis_dim_mapping(data.dims)
 
         # Create DataArrays for indexing
         selection_dict = {
-            axis_dim["X"]: xr.DataArray(xi, dims=("points")),
-            axis_dim["Y"]: xr.DataArray(yi, dims=("points")),
+            axis_dim["X"]: xr.DataArray(xi_full, dims=("points")),
+            axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")),
         }
         if "Z" in axis_dim:
-            selection_dict[axis_dim["Z"]] = xr.DataArray(zi, dims=("points"))
+            selection_dict[axis_dim["Z"]] = xr.DataArray(zi_full, dims=("points"))
         if "time" in data.dims:
-            selection_dict["time"] = xr.DataArray(ti, dims=("points"))
+            selection_dict["time"] = xr.DataArray(ti_full, dims=("points"))
 
         corner_data = data.isel(selection_dict).data.reshape(lenT, len(xsi), 4)
 
@@ -271,7 +271,53 @@ def CGrid_Velocity(
     xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3]
     u = np.where(np.abs(xx - x) > 1e-4, np.nan, u)
 
-    return (u, v, 0)  # TODO fix and test W also
+    if vectorfield.W:
+        data = vectorfield.W.data
+        # Time coordinates: 2 points at ti, then 2 points at ti+1
+        if lenT == 1:
+            ti_full = np.repeat(ti, 2)
+        else:
+            ti_1 = np.clip(ti + 1, 0, tdim - 1)
+            ti_full = np.concatenate([np.repeat(ti, 2), np.repeat(ti_1, 2)])
+
+        # Depth coordinates: 1 points at zi, repeated for both time levels
+        zi_1 = np.clip(zi + 1, 0, zdim - 1)
+        zi_full = np.tile(np.array([zi, zi_1]).flatten(), lenT)
+
+        # Y coordinates: yi+1 for each spatial point, repeated for time/depth
+        yi_1 = np.clip(yi + 1, 0, ydim - 1)
+        yi_full = np.tile(yi_1, (lenT) * 2)
+
+        # X coordinates: xi+1 for each spatial point, repeated for time/depth
+        xi_1 = np.clip(xi + 1, 0, xdim - 1)
+        xi_full = np.tile(xi_1, (lenT) * 2)
+
+        axis_dim = grid.get_axis_dim_mapping(data.dims)
+
+        # Create DataArrays for indexing
+        selection_dict = {
+            axis_dim["X"]: xr.DataArray(xi_full, dims=("points")),
+            axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")),
+            axis_dim["Z"]: xr.DataArray(zi_full, dims=("points")),
+        }
+        if "time" in data.dims:
+            selection_dict["time"] = xr.DataArray(ti_full, dims=("points"))
+
+        corner_data = data.isel(selection_dict).data.reshape(lenT, 2, len(xsi))
+
+        if lenT == 2:
+            tau_full = tau[np.newaxis, :]
+            corner_data = corner_data[0, :, :] * (1 - tau_full) + corner_data[1, :, :] * tau_full
+        else:
+            corner_data = corner_data[0, :, :]
+
+        w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta
+        if isinstance(w, dask.Array):
+            w = w.compute()
+    else:
+        w = np.zeros_like(u)
+
+    return (u, v, w)
 
 
 def CGrid_Tracer(

tests/v4/test_advection.py

@@ -20,6 +20,7 @@
 from parcels.particleset import ParticleSet
 from parcels.tools.statuscodes import StatusCode
 from parcels.xgrid import XGrid
+from tests.utils import round_and_hash_float_array
 
 kernel = {
     "EE": AdvectionEE,
@@ -444,7 +445,8 @@ def periodicBC(particle, fieldSet, time):  # pragma: no cover
     np.testing.assert_allclose(pset.lat_nextloop, latp, atol=1e-1)
 
 
-def test_nemo_3D_curvilinear_fieldset():
+@pytest.mark.parametrize("method", ["RK4", "RK4_3D"])
+def test_nemo_3D_curvilinear_fieldset(method):
     download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data")
     ufiles = download_dir.glob("*U.nc")
     dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"])
@@ -466,10 +468,34 @@ def test_nemo_3D_curvilinear_fieldset():
 
     coord_file = f"{download_dir}/coordinates.nc"
     dscoord = xr.open_dataset(coord_file, decode_times=False).rename({"glamf": "lon", "gphif": "lat"})
+    dscoord = dscoord.isel(time=0, drop=True)
 
     ds = xr.merge([dsu, dsv, dsw, dscoord])
+    ds = ds.drop_vars(
+        [
+            "uos",
+            "vos",
+            "nav_lev",
+            "nav_lon",
+            "nav_lat",
+            "tauvo",
+            "tauuo",
+            "time_steps",
+            "gphiu",
+            "gphiv",
+            "gphit",
+            "glamu",
+            "glamv",
+            "glamt",
+            "time_centered_bounds",
+            "time_counter_bounds",
+            "time_centered",
+        ]
+    )
+    ds = ds.drop_vars(["e1f", "e1t", "e1u", "e1v", "e2f", "e2t", "e2u", "e2v"])
+    ds["time"] = [np.timedelta64(int(t), "s") + np.datetime64("1900-01-01") for t in ds["time"]]
 
-    ds.load()  # TODO remove for debug
+    ds["W"] *= -1  # Invert W velocity
 
     xgcm_grid = parcels.xgcm.Grid(
         ds,
@@ -488,14 +514,21 @@ def test_nemo_3D_curvilinear_fieldset():
     W = parcels.Field("W", ds["W"], grid, mesh_type="spherical")
     U.units = parcels.GeographicPolar()
     V.units = parcels.Geographic()
-    W.units = parcels.Geographic()
-    UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity)  # TODO remove
+    UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity)
     UVW = parcels.VectorField("UVW", U, V, W, vector_interp_method=CGrid_Velocity)
     fieldset = parcels.FieldSet([U, V, W, UV, UVW])
 
-    lons = np.linspace(1.9, 3.4, 10)
-    lats = np.linspace(52.5, 51.6, 10)
+    npart = 10
+    lons = np.linspace(1.9, 3.4, npart)
+    lats = np.linspace(52.5, 51.6, npart)
     pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, depth=np.ones_like(lons))
 
-    pset.execute(parcels.AdvectionRK4, runtime=np.timedelta64(4, "D"), dt=np.timedelta64(6, "h"))
-    print(pset.depth)
+    pset.execute(kernel[method], runtime=np.timedelta64(4, "D"), dt=np.timedelta64(6, "h"))
+
+    if method == "RK4":
+        np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546)
+    elif method == "RK4_3D":
+        # TODO check why decimals needs to be so low in RK4_3D (compare to v3)
+        np.testing.assert_equal(
+            round_and_hash_float_array([p.depth for p in pset], decimals=1), 29747210774230389239432
+        )