Merge pull request #1772 from OceanParcels/CROCO_fix_sigma_calculation

Implementing correct depth-to-sigma calculation

parcels/compilation/codegenerator.py

@@ -425,10 +425,13 @@ def __init__(self, fieldset=None, ptype=JITParticle):
         self.fieldset = fieldset
         self.ptype = ptype
         self.field_args = collections.OrderedDict()
-        if isinstance(fieldset.U, Field) and fieldset.U.gridindexingtype == "croco" and hasattr(fieldset, "H"):
-            self.field_args["H"] = fieldset.H  # CROCO requires H field
         self.vector_field_args = collections.OrderedDict()
         self.const_args = collections.OrderedDict()
+        if isinstance(fieldset.U, Field) and fieldset.U.gridindexingtype == "croco" and hasattr(fieldset, "H"):
+            self.field_args["H"] = fieldset.H  # CROCO requires H field
+            self.field_args["Zeta"] = fieldset.Zeta  # CROCO requires Zeta field
+            self.field_args["Cs_w"] = fieldset.Cs_w  # CROCO requires CS_w field
+            self.const_args["hc"] = fieldset.hc  # CROCO requires hc constant
 
     def generate(self, py_ast, funcvars: list[str]):
         # Replace occurrences of intrinsic objects in Python AST
@@ -825,16 +828,18 @@ def visit_FieldEvalNode(self, node):
         self.visit(node.field)
         self.visit(node.args)
         args = self._check_FieldSamplingArguments(node.args.ccode)
-        statements_croco = []
-        if "croco" in node.field.obj.gridindexingtype and node.field.obj.name != "H":
-            statements_croco.append(
-                c.Assign(
-                    "parcels_interp_state",
-                    f"temporal_interpolation({args[3]}, {args[2]}, 0, time, H, &particles->xi[pnum*ngrid], &particles->yi[pnum*ngrid], &particles->zi[pnum*ngrid], &particles->ti[pnum*ngrid], &{node.var}, LINEAR, {node.field.obj.gridindexingtype.upper()})",
-                )
-            )
-            statements_croco.append(c.Statement(f"{node.var} = {args[1]}/{node.var}"))
+        if "croco" in node.field.obj.gridindexingtype and node.field.obj.name != "H" and node.field.obj.name != "Zeta":
+            # Get Cs_w values directly from fieldset (since they are 1D in vertical only)
+            Cs_w = [float(self.fieldset.Cs_w.data[0][zi][0][0]) for zi in range(self.fieldset.Cs_w.data.shape[1])]
+            statements_croco = [
+                c.Statement(f"float cs_w[] = {*Cs_w, }".replace("(", "{").replace(")", "}")),
+                c.Statement(
+                    f"{node.var} = croco_from_z_to_sigma(U, H, Zeta, {args[3]}, {args[2]}, {args[1]}, time, &particles->xi[pnum*ngrid], &particles->yi[pnum*ngrid], &particles->zi[pnum*ngrid], &particles->ti[pnum*ngrid], hc, &cs_w)"
+                ),
+            ]
             args = (args[0], node.var, args[2], args[3])
+        else:
+            statements_croco = []
         ccode_eval = node.field.obj._ccode_eval(node.var, *args)
         stmts = [
             c.Assign("parcels_interp_state", ccode_eval),
@@ -852,16 +857,18 @@ def visit_VectorFieldEvalNode(self, node):
         self.visit(node.field)
         self.visit(node.args)
         args = self._check_FieldSamplingArguments(node.args.ccode)
-        statements_croco = []
         if "3DSigma" in node.field.obj.vector_type:
-            statements_croco.append(
-                c.Assign(
-                    "parcels_interp_state",
-                    f"temporal_interpolation({args[3]}, {args[2]}, 0, time, H, &particles->xi[pnum*ngrid], &particles->yi[pnum*ngrid], &particles->zi[pnum*ngrid], &particles->ti[pnum*ngrid], &{node.var}, LINEAR, {node.field.obj.U.gridindexingtype.upper()})",
-                )
-            )
-            statements_croco.append(c.Statement(f"{node.var4} = {args[1]}/{node.var}"))
+            # Get Cs_w values directly from fieldset (since they are 1D in vertical only)
+            Cs_w = [float(self.fieldset.Cs_w.data[0][zi][0][0]) for zi in range(self.fieldset.Cs_w.data.shape[1])]
+            statements_croco = [
+                c.Statement(f"float cs_w[] = {*Cs_w, }".replace("(", "{").replace(")", "}")),
+                c.Statement(
+                    f"{node.var4} = croco_from_z_to_sigma(U, H, Zeta, {args[3]}, {args[2]}, {args[1]}, time, &particles->xi[pnum*ngrid], &particles->yi[pnum*ngrid], &particles->zi[pnum*ngrid], &particles->ti[pnum*ngrid], hc, &cs_w)"
+                ),
+            ]
             args = (args[0], node.var4, args[2], args[3])
+        else:
+            statements_croco = []
         ccode_eval = node.field.obj._ccode_eval(
             node.var, node.var2, node.var3, node.field.obj.U, node.field.obj.V, node.field.obj.W, *args
         )

parcels/field.py

@@ -76,6 +76,26 @@ def _deal_with_errors(error, key, vector_type: VectorType):
         return 0
 
 
+def _croco_from_z_to_sigma_scipy(fieldset, time, z, y, x, particle):
+    """Calculate local sigma level of the particle, by linearly interpolating the
+    scaling function that maps sigma to depth (using local ocean depth H,
+    sea-surface Zeta and stretching parameters Cs_w and hc).
+    See also https://croco-ocean.gitlabpages.inria.fr/croco_doc/model/model.grid.html#vertical-grid-parameters
+    """
+    h = fieldset.H.eval(time, 0, y, x, particle=particle, applyConversion=False)
+    zeta = fieldset.Zeta.eval(time, 0, y, x, particle=particle, applyConversion=False)
+    sigma_levels = fieldset.U.grid.depth
+    z0 = fieldset.hc * sigma_levels + (h - fieldset.hc) * fieldset.Cs_w.data[0, :, 0, 0]
+    zvec = z0 + zeta * (1 + (z0 / h))
+    zinds = zvec <= z
+    if z >= zvec[-1]:
+        zi = len(zvec) - 2
+    else:
+        zi = zinds.argmin() - 1 if z >= zvec[0] else 0
+
+    return sigma_levels[zi] + (z - zvec[zi]) * (sigma_levels[zi + 1] - sigma_levels[zi]) / (zvec[zi + 1] - zvec[zi])
+
+
 class Field:
     """Class that encapsulates access to field data.
 
@@ -617,18 +637,23 @@ def from_netcdf(
 
         _grid_fb_class = NetcdfFileBuffer
 
-        with _grid_fb_class(
-            lonlat_filename,
-            dimensions,
-            indices,
-            netcdf_engine,
-            gridindexingtype=gridindexingtype,
-        ) as filebuffer:
-            lon, lat = filebuffer.lonlat
-            indices = filebuffer.indices
-            # Check if parcels_mesh has been explicitly set in file
-            if "parcels_mesh" in filebuffer.dataset.attrs:
-                mesh = filebuffer.dataset.attrs["parcels_mesh"]
+        if "lon" in dimensions and "lat" in dimensions:
+            with _grid_fb_class(
+                lonlat_filename,
+                dimensions,
+                indices,
+                netcdf_engine,
+                gridindexingtype=gridindexingtype,
+            ) as filebuffer:
+                lon, lat = filebuffer.lonlat
+                indices = filebuffer.indices
+                # Check if parcels_mesh has been explicitly set in file
+                if "parcels_mesh" in filebuffer.dataset.attrs:
+                    mesh = filebuffer.dataset.attrs["parcels_mesh"]
+        else:
+            lon = 0
+            lat = 0
+            mesh = "flat"
 
         if "depth" in dimensions:
             with _grid_fb_class(
@@ -1537,8 +1562,8 @@ def eval(self, time, z, y, x, particle=None, applyConversion=True):
         """
         (ti, periods) = self._time_index(time)
         time -= periods * (self.grid.time_full[-1] - self.grid.time_full[0])
-        if self.gridindexingtype == "croco" and self is not self.fieldset.H:
-            z = z / self.fieldset.H.eval(time, 0, y, x, particle=particle, applyConversion=False)
+        if self.gridindexingtype == "croco" and self not in [self.fieldset.H, self.fieldset.Zeta]:
+            z = _croco_from_z_to_sigma_scipy(self.fieldset, time, z, y, x, particle=particle)
         if ti < self.grid.tdim - 1 and time > self.grid.time[ti]:
             f0 = self._spatial_interpolation(ti, z, y, x, time, particle=particle)
             f1 = self._spatial_interpolation(ti + 1, z, y, x, time, particle=particle)
@@ -2250,7 +2275,7 @@ def spatial_c_grid_interpolation3D(self, ti, z, y, x, time, particle=None, apply
             (u, v, w) = self.spatial_c_grid_interpolation3D_full(ti, z, y, x, time, particle=particle)
         else:
             if self.gridindexingtype == "croco":
-                z = z / self.fieldset.H.eval(time, 0, y, x, particle=particle, applyConversion=False)
+                z = _croco_from_z_to_sigma_scipy(self.fieldset, time, z, y, x, particle=particle)
             (u, v) = self.spatial_c_grid_interpolation2D(ti, z, y, x, time, particle=particle)
             w = self.W.eval(time, z, y, x, particle=particle, applyConversion=False)
             if applyConversion:

parcels/fieldfilebuffer.py

@@ -188,7 +188,10 @@ def _check_extend_depth(self, data, di):
         )
 
     def _apply_indices(self, data, ti):
-        if len(data.shape) == 2:
+        if len(data.shape) == 1:
+            if self.indices["depth"] is not None:
+                data = data[self.indices["depth"]]
+        elif len(data.shape) == 2:
             if self.nolonlatindices:
                 pass
             else:

parcels/fieldset.py

@@ -713,6 +713,7 @@ def from_croco(
         filenames,
         variables,
         dimensions,
+        hc: float | None = None,
         indices=None,
         mesh="spherical",
         allow_time_extrapolation=None,
@@ -723,11 +724,14 @@ def from_croco(
     ):
         """Initialises FieldSet object from NetCDF files of CROCO fields.
         All parameters and keywords are exactly the same as for FieldSet.from_nemo(), except that
-        the vertical coordinate is scaled by the bathymetry (``h``) field from CROCO, in order to
-        account for the sigma-grid. The horizontal interpolation uses the MITgcm grid indexing
-        as described in FieldSet.from_mitgcm().
+        in order to scale the vertical coordinate in CROCO, the following fields are required:
+        the bathymetry (``h``), the sea-surface height (``zeta``), the S-coordinate stretching curves
+        at W-points (``Cs_w``), and the stretching parameter (``hc``).
+        The horizontal interpolation uses the MITgcm grid indexing as described in FieldSet.from_mitgcm().
 
-        The sigma grid scaling means that FieldSet.from_croco() requires a variable ``H: h`` to work.
+        In 3D, when there is a ``depth`` dimension, the sigma grid scaling means that FieldSet.from_croco()
+        requires variables ``H: h`` and ``Zeta: zeta``, ``Cs_w: Cs_w``, as well as the stretching parameter ``hc``
+        (as an extra input) parameter to work.
 
         See `the CROCO 3D tutorial <../examples/tutorial_croco_3D.ipynb>`__ for more infomation.
         """
@@ -739,14 +743,23 @@ def from_croco(
             )
 
         dimsU = dimensions["U"] if "U" in dimensions else dimensions
-        if "depth" in dimsU:
-            warnings.warn(
-                "Note that it is unclear which vertical velocity ('w' or 'omega') to use in 3D CROCO fields.\nSee https://docs.oceanparcels.org/en/latest/examples/tutorial_croco_3D.html for more information",
-                FieldSetWarning,
-                stacklevel=2,
-            )
+        croco3D = True if "depth" in dimsU else False
+
+        if croco3D:
+            if "W" in variables and variables["W"] == "omega":
+                warnings.warn(
+                    "Note that Parcels expects 'w' for vertical velicites in 3D CROCO fields.\nSee https://docs.oceanparcels.org/en/latest/examples/tutorial_croco_3D.html for more information",
+                    FieldSetWarning,
+                    stacklevel=2,
+                )
             if "H" not in variables:
-                raise ValueError("FieldSet.from_croco() requires a field 'H' for the bathymetry")
+                raise ValueError("FieldSet.from_croco() requires a bathymetry field 'H' for 3D CROCO fields")
+            if "Zeta" not in variables:
+                raise ValueError("FieldSet.from_croco() requires a free-surface field 'Zeta' for 3D CROCO fields")
+            if "Cs_w" not in variables:
+                raise ValueError(
+                    "FieldSet.from_croco() requires the S-coordinate stretching curves at W-points 'Cs_w' for 3D CROCO fields"
+                )
 
         interp_method = {}
         for v in variables:
@@ -776,6 +789,10 @@ def from_croco(
             gridindexingtype="croco",
             **kwargs,
         )
+        if croco3D:
+            if hc is None:
+                raise ValueError("FieldSet.from_croco() requires the hc parameter for 3D CROCO fields")
+            fieldset.add_constant("hc", hc)
         return fieldset
 
     @classmethod

parcels/include/parcels.h

@@ -1242,6 +1242,33 @@ static inline StatusCode temporal_interpolationUVW(type_coord x, type_coord y, t
   return SUCCESS;
 }
 
+
+static inline double croco_from_z_to_sigma(CField *U, CField *H, CField *Zeta,
+                                           type_coord x, type_coord y, type_coord z, double time,
+                                           int *xi, int *yi, int *zi, int *ti, double hc, float *cs_w)
+{
+  float local_h, local_zeta, z0;
+  int status, zii;
+  CStructuredGrid *grid = U->grid->grid;
+  float *sigma_levels = grid->depth;
+  int zdim = grid->zdim;
+  float zvec[zdim];
+  status = temporal_interpolation(x, y, 0, time, H, xi, yi, zi, ti, &local_h, LINEAR, CROCO); CHECKSTATUS(status);
+  status = temporal_interpolation(x, y, 0, time, Zeta, xi, yi, zi, ti, &local_zeta, LINEAR, CROCO); CHECKSTATUS(status);
+  for (zii = 0; zii < zdim; zii++)  {
+    z0 = hc*sigma_levels[zii] + (local_h - hc) *cs_w[zii];
+    zvec[zii] = z0 + local_zeta * (1 + z0 / local_h);
+  }
+  if (z >= zvec[zdim-1])
+    zii = zdim - 2;
+  else
+    for (zii = 0; zii < zdim-1; zii++)
+      if ((z >= zvec[zii])  && (z < zvec[zii+1]))
+        break;
+
+  return sigma_levels[zii] + (z - zvec[zii]) * (sigma_levels[zii + 1] - sigma_levels[zii]) / (zvec[zii + 1] - zvec[zii]);
+}
+
 #ifdef __cplusplus
 }
 #endif

tests/test_advection.py

@@ -195,6 +195,47 @@ def test_advection_RK45(lon, lat, mode, rk45_tol):
     print(fieldset.RK45_tol)
 
 
+def test_conversion_3DCROCO():
+    """Test of the (SciPy) version of the conversion from depth to sigma in CROCO
+
+    Values below are retrieved using xroms and hardcoded in the method (to avoid dependency on xroms):
+    ```py
+    x, y = 10, 20
+    s_xroms = ds.s_w.values
+    z_xroms = ds.z_w.isel(time=0).isel(eta_rho=y).isel(xi_rho=x).values
+    lat, lon = ds.y_rho.values[y, x], ds.x_rho.values[y, x]
+    ```
+    """
+    fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py")
+
+    lat, lon = 78000.0, 38000.0
+    s_xroms = np.array([-1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0.0], dtype=np.float32)
+    z_xroms = np.array(
+        [
+            -1.26000000e02,
+            -1.10585846e02,
+            -9.60985413e01,
+            -8.24131317e01,
+            -6.94126511e01,
+            -5.69870148e01,
+            -4.50318756e01,
+            -3.34476166e01,
+            -2.21383114e01,
+            -1.10107975e01,
+            2.62768921e-02,
+        ],
+        dtype=np.float32,
+    )
+
+    sigma = np.zeros_like(z_xroms)
+    from parcels.field import _croco_from_z_to_sigma_scipy
+
+    for zi, z in enumerate(z_xroms):
+        sigma[zi] = _croco_from_z_to_sigma_scipy(fieldset, 0, z, lat, lon, None)
+
+    assert np.allclose(sigma, s_xroms, atol=1e-3)
+
+
 @pytest.mark.parametrize("mode", ["scipy", "jit"])
 def test_advection_3DCROCO(mode):
     fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py")

tests/test_data/fieldset_CROCO3D.py

@@ -1,18 +1,22 @@
 import os
 
+import xarray as xr
+
 import parcels
 
 
 def create_fieldset(indices=None):
     example_dataset_folder = parcels.download_example_dataset("CROCOidealized_data")
     file = os.path.join(example_dataset_folder, "CROCO_idealized.nc")
 
-    variables = {"U": "u", "V": "v", "W": "w", "H": "h"}
+    variables = {"U": "u", "V": "v", "W": "w", "H": "h", "Zeta": "zeta", "Cs_w": "Cs_w"}
     dimensions = {
         "U": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"},
         "V": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"},
         "W": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"},
         "H": {"lon": "x_rho", "lat": "y_rho"},
+        "Zeta": {"lon": "x_rho", "lat": "y_rho", "time": "time"},
+        "Cs_w": {"depth": "s_w"},
     }
     fieldset = parcels.FieldSet.from_croco(
         file,
@@ -21,6 +25,7 @@ def create_fieldset(indices=None):
         allow_time_extrapolation=True,
         mesh="flat",
         indices=indices,
+        hc=xr.open_dataset(file).hc.values,
     )
 
     return fieldset