Merge pull request #49 from NOC-MSM/46-add-support-for-3-dimensional-depth-coordinates

46 add support for 3 dimensional depth coordinates

Author: oj-tooth

README.md

@@ -39,7 +39,7 @@ If you already use `xarray`, NEMO Cookbook should feel immediately natural:
 * `NEMODataArray` behaves like `xarray.DataArray`.
 * All standard `xarray` operations are still available!
 
-What’s new is that these objects understand the NEMO grid, meaning  you no longer need to manually track:
+What’s new is that these objects understand the NEMO grid, meaning you no longer need to manually track:
 
 * which NEMO model grid a variable belongs to (e.g., T, U, V, F, W).
 * how variables relate across NEMO model grids.

nemo_cookbook/extract.py

@@ -68,9 +68,9 @@ def create_section_polygon(
             [lat_sec, np.array([lat_sec[-1] - dlat, lat_sec[-1] - dlat, lat_sec[0]])]
         )
     elif all(lat == lat_sec[0] for lat in lat_sec):
-        raise ValueError("extracting zonal hydrographic sections is not supported.")
+        raise ValueError("extracting zonal hydrographic sections is not supported. Use NEMODataTree.extract_zonal_section() instead.")
     else:
-        raise ValueError("hydrographic section endpoints must not have same latitude.")
+        raise ValueError("hydrographic section endpoints must not have same latitude to form a closed polygon.")
 
     return lon_poly, lat_poly
 
@@ -154,8 +154,19 @@ def create_boundary_dataset(
     ds_bdy[f"{dom_prefix}gphib"][vbdy_mask] = nemo[gridV][f"{dom_prefix}gphiv"].sel(
         i=ds_bdy["i_bdy"][vbdy_mask], j=ds_bdy["j_bdy"][vbdy_mask]
     )
-    ds_bdy[f"{dom_prefix}depthb"][:, ubdy_mask] = nemo[gridU][f"{dom_prefix}depthu"]
-    ds_bdy[f"{dom_prefix}depthb"][:, vbdy_mask] = nemo[gridV][f"{dom_prefix}depthv"]
+
+    if (nemo[gridU][f"{dom_prefix}depthu"].ndim == 1) and (nemo[gridV][f"{dom_prefix}depthv"].ndim == 1):
+        # Using 1-dimensional vertical reference (depth) coordinate:
+        ds_bdy[f"{dom_prefix}depthb"][:, ubdy_mask] = nemo[gridU][f"{dom_prefix}depthu"]
+        ds_bdy[f"{dom_prefix}depthb"][:, vbdy_mask] = nemo[gridV][f"{dom_prefix}depthv"]
+    else:
+        # Using 3-dimensional vertical reference (depth) coordinate:
+        ds_bdy[f"{dom_prefix}depthb"][:, ubdy_mask] = nemo[gridU][f"{dom_prefix}depthu"].sel(
+            i=ds_bdy["i_bdy"][ubdy_mask], j=ds_bdy["j_bdy"][ubdy_mask]
+        )
+        ds_bdy[f"{dom_prefix}depthb"][:, vbdy_mask] = nemo[gridV][f"{dom_prefix}depthv"].sel(
+            i=ds_bdy["i_bdy"][vbdy_mask], j=ds_bdy["j_bdy"][vbdy_mask]
+        )
 
     return ds_bdy
 

nemo_cookbook/nemodataarray.py

@@ -882,8 +882,17 @@ def depth_integral(
             )
 
         # -- Define input variables -- #
-        var_in = self.masked.data
-        e3_in = self.metrics["e3"].masked.data
+        var_in = (self
+                 .masked
+                 .data
+                 .chunk({self.k_name: -1})
+                 )
+ 
+        e3_in = (self.metrics["e3"]
+                 .masked
+                 .data
+                 .chunk({self.k_name: -1})
+                 )
 
         # -- Vertically integrate w.r.t depth -- #
         result = xr.apply_ufunc(
@@ -1150,14 +1159,23 @@ def transform_vertical_grid(
             )
 
         # -- Define input variables -- #
-        var_in = self.masked.data
-        e3_in = self.metrics["e3"].masked.data
+        var_in = (self
+                 .masked
+                 .data
+                 .chunk({self.k_name: -1})
+                 )
+ 
+        e3_in = (self.metrics["e3"]
+                 .masked
+                 .data
+                 .chunk({self.k_name: -1})
+                 )
 
         # Ensure total depth of new vertical grid >= total depth of NEMO model vertical grid:
-        depth_max = self._tree[self._grid][f"{self._dom_prefix}depth{self._grid_suffix}"].max(self.k_name)
+        depth_max = self._tree[self._grid][f"{self._dom_prefix}depth{self._grid_suffix}"].max()
         if e3_new.sum(dim="k_new") < depth_max:
             raise ValueError(
-                f"e3_new must sum to at least the maximum depth ({depth_max.item()} m) of the original vertical grid."
+                f"e3_new must sum to at least the maximum depth ({depth_max.values.item()} m) of the original vertical grid."
             )
 
         # -- Transform variable to target vertical grid -- #
@@ -1183,7 +1201,8 @@ def transform_vertical_grid(
         result = xr.Dataset(
             data_vars={self.name: var_out, f"e3{self._grid_suffix}_new": e3_out},
             coords={
-                f"depth{self._grid_suffix}_new": ("k_new", e3_new.cumsum(dim="k_new").data)
+                # Vertical reference (depth of cell center) coordinates of new vertical grid:
+                f"depth{self._grid_suffix}_new": ("k_new", (e3_new.cumsum(dim="k_new") - (e3_new / 2)).data),
             },
         )
 

nemo_cookbook/nemodatatree.py

@@ -71,7 +71,8 @@ def from_paths(
         nftype: str | None = None,
         read_mask: bool = False,
         maskcs: bool = False,
-        key_linssh: bool = False,
+        linssh: bool = False,
+        vco: str = "1d",
         vco_ref: bool = False,
         nbghost_child: int = 4,
         **open_kwargs: dict[str, any],
@@ -139,10 +140,13 @@ def from_paths(
         maskcs: bool = False
             If True, all closed seas are masked using mask_opensea variables from domain files. Default is False.
 
-        key_linssh: bool = False
+        linssh: bool = False
             Linear free-surface approximation. If True, vertical coordinates are time-independent and given by (e3t_0, e3u_0, e3v_0, e3w_0) in domain_cfg.
             If False, vertical coordinates are time-dependent and must be specified in NEMO model grid datasets. Default is False.
 
+        vco : str = "1d"
+            Vertical reference variables. Options are '1d' to use 1-dimensional vertical reference coordinates or '3d' to use 3-dimensional vertical reference coordinates (deptht, depthu, depthv, depthw, depthf). Default is '1d'.
+
         vco_ref: bool = False
             If True, add reference vertical scale factors and compute reference water column heights from domain files. Default is False.
 
@@ -173,7 +177,7 @@ def from_paths(
 
         Create a regional `NEMODataTree` using a linear free-surface approximation from a dictionary of paths to remote netCDF files:
 
-        >>> nemo = NEMODataTree.from_paths(paths, name="My NEMO model", iperio=False, nftype=None, key_linssh=True)
+        >>> nemo = NEMODataTree.from_paths(paths, name="My NEMO model", iperio=False, nftype=None, linssh=True)
 
         See Also
         --------
@@ -196,8 +200,14 @@ def from_paths(
             raise TypeError("reading land-sea masks from domain_cfg (`read_mask`) must be a boolean.")
         if not isinstance(maskcs, bool):
             raise TypeError("masking of closed seas (`maskcs`) must be a boolean.")
-        if not isinstance(key_linssh, bool):
-            raise TypeError("linear free-surface approximation (`key_linssh`) must be a boolean.")
+        if not isinstance(linssh, bool):
+            raise TypeError("linear free-surface approximation (`linssh`) must be a boolean.")
+        if not isinstance(vco, str):
+            raise TypeError("vertical reference coordinates (`vco`) must be a string.")
+        if vco not in ("1d", "3d"):
+            raise ValueError(
+                "vertical reference coordinates (`vco`) must be '1d' (1-dimensional) or '3d' (3-dimensional)."
+            )
         if not isinstance(vco_ref, bool):
             raise TypeError("reference vertical coordinates (`vco_ref`) must be a boolean.")
         if not isinstance(nbghost_child, int):
@@ -240,7 +250,8 @@ def from_paths(
             read_mask=read_mask,
             maskcs=maskcs,
             nbghost_child=nbghost_child,
-            key_linssh=key_linssh,
+            linssh=linssh,
+            vco=vco,
             vco_ref=vco_ref,
             open_kwargs=dict(**open_kwargs),
         )
@@ -259,7 +270,8 @@ def from_datasets(
         iperio: bool = False,
         nftype: str | None = None,
         read_mask: bool = False,
-        key_linssh: bool = False,
+        linssh: bool = False,
+        vco: str = "1d",
         vco_ref: bool = False,
         maskcs: bool = False,
         nbghost_child: int = 4,
@@ -311,10 +323,13 @@ def from_datasets(
         maskcs: bool = False
             If True, all closed seas are masked using mask_opensea variables from domain files. Default is False.
 
-        key_linssh: bool = False
+        linssh: bool = False
             Linear free-surface approximation. If True, vertical coordinates are time-independent and given by (e3t_0, e3u_0, e3v_0, e3w_0) in domain_cfg.
             If False, vertical coordinates are time-dependent and must be specified in NEMO model grid datasets. Default is False.
 
+        vco : str = "1d"
+            Vertical reference variables. Options are '1d' to use 1-dimensional vertical reference coordinates or '3d' to use 3-dimensional vertical reference coordinates (deptht, depthu, depthv, depthw, depthf). Default is '1d'.        
+    
         vco_ref: bool = False
             If True, add reference vertical scale factors and compute reference water column heights from domain files. Default is False.
 
@@ -362,8 +377,14 @@ def from_datasets(
             raise TypeError("reading land-sea masks from domain_cfg (`read_mask`) must be a boolean.")
         if not isinstance(maskcs, bool):
             raise TypeError("masking of closed seas (`maskcs`) must be a boolean.")
-        if not isinstance(key_linssh, bool):
-            raise TypeError("linear free-surface approximation (`key_linssh`) must be a boolean.")
+        if not isinstance(linssh, bool):
+            raise TypeError("linear free-surface approximation (`linssh`) must be a boolean.")
+        if not isinstance(vco, str):
+            raise TypeError("vertical reference coordinates (`vco`) must be a string.")
+        if vco not in ("1d", "3d"):
+            raise ValueError(
+                "vertical reference coordinates (`vco`) must be '1d' (1-dimensional) or '3d' (3-dimensional)."
+            )
         if not isinstance(vco_ref, bool):
             raise TypeError("reference vertical coordinates (`vco_ref`) must be a boolean.")
         if not isinstance(nbghost_child, int):
@@ -403,7 +424,8 @@ def from_datasets(
             nftype=nftype,
             read_mask=read_mask,
             maskcs=maskcs,
-            key_linssh=key_linssh,
+            linssh=linssh,
+            vco=vco,
             vco_ref=vco_ref,
             nbghost_child=nbghost_child,
         )