Implementing correct depth-to-sigma calculation in scipy mode

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)
@@ -2252,7 +2277,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,
         indices=None,
         mesh="spherical",
         allow_time_extrapolation=None,
@@ -723,11 +724,12 @@ 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().
+        the vertical coordinate is scaled by the bathymetry (``h``) and sea-surface height (``zeta``)
+        fields from CROCO, in order to account for the sigma-grid.
+        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.
+        The sigma grid scaling means that FieldSet.from_croco() requires variables
+        ``H: h`` and ``Zeta: zeta``, as well as he ``hc`` parameter to work.
 
         See `the CROCO 3D tutorial <../examples/tutorial_croco_3D.ipynb>`__ for more infomation.
         """
@@ -740,13 +742,16 @@ 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,
-            )
+            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")
+            if "Zeta" not in variables:
+                raise ValueError("FieldSet.from_croco() requires a field 'Zeta' for the free_surface")
 
         interp_method = {}
         for v in variables:
@@ -776,6 +781,7 @@ def from_croco(
             gridindexingtype="croco",
             **kwargs,
         )
+        fieldset.add_constant("hc", hc)
         return fieldset
 
     @classmethod

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