Merge pull request #50 from NOC-MSM/47-add-option-to-retain-ghost-points-in-nested-nemodatatrees

47 add option to retain ghost points in nested nemodatatrees

Author: oj-tooth

docs/docs/howto.md

@@ -26,7 +26,9 @@ NEMODataTree.from_paths(paths, iperio=True, nftype="T")
 
 In the example above, we consider only a global parent domain, which is zonally periodic (`iperio=True`) and north-folding on **T** grid points (`nftype="T"`). Note, that we are only required to specify paths for one or more NEMO model grid types (e.g., `*_gridT.nc`).
 
-For NEMO models using a linear free surface approximation (i.e., vertical scale factors are time-independent), we should also specify `key_linssh=True` to read these directly from the domain_cfg file included in the `paths` dictionary.
+For NEMO models using a linear free surface approximation (i.e., vertical scale factors are time-independent), we should also specify `linssh=True` to read these directly from the domain_cfg file included in the `paths` dictionary.
+
+Additionally, for NEMO models configured with more complex vertical coordinates (e.g., MEs or sigma-coordinates), such that vertical reference variables (e.g., `depth`) vary spatially, we should also specify `vco="3d"` to include all vertical reference variables as 3-dimensional arrays analogously to using `key_vco_3d` within NEMO itself. By default, a `NEMODataTree` is constructed using 1-dimensional vertical reference variables.
 
 ### Create a NEMODataTree from `xarray.Datasets`
 

docs/docs/user_guide.md

@@ -178,9 +178,9 @@ This is because domain variables are assigned to their respective grid nodes dur
 
     Typically, NEMO model simulations use a quasi-eulerian vertical coordinate which absorbs the divergence of horizontal barotropic velocities (e.g., $z^{*}$ or $s^{*}$), meaning that vertical grid scale factors evolve through time (i.e., a time-varying free surface translates into variations in grid cell thickness).
 
-    `NEMODataTree` considers the case of time-evolving vertical grid scale factors to be the default as the `key_linssh` argument to the `.from_paths()` and `.from_datasets()` constructors is set to be `False` by default. This means that vertical grid scale factors must be provided in the netCDF files or `xarray.Datasets` used to define each NEMO model grid node in the `NEMODataTree`.
+    `NEMODataTree` considers the case of time-evolving vertical grid scale factors to be the default as the `linssh` argument to the `.from_paths()` and `.from_datasets()` constructors is set to be `False` by default. This means that vertical grid scale factors must be provided in the netCDF files or `xarray.Datasets` used to define each NEMO model grid node in the `NEMODataTree`.
 
-    For NEMO model simulations using a linear free surface approximation (i.e., variations in the free surface are neglected compared to the depth of the ocean), we should use `key_linssh=True` to indicate that vertical grid scale factors remain fixed through time and should be read directly from the reference variables contained within the domain_cfg netCDF file or `xarray.Dataset` (e.g., `e3t_0`, `e3u_0` etc.).
+    For NEMO model simulations using a linear free surface approximation (i.e., variations in the free surface are neglected compared to the depth of the ocean), we should use `linssh=True` to indicate that vertical grid scale factors remain fixed through time and should be read directly from the reference variables contained within the domain_cfg netCDF file or `xarray.Dataset` (e.g., `e3t_0`, `e3u_0` etc.).
 
 **Dimensions & Coordinates**
 
@@ -268,11 +268,11 @@ In summary, defining a `NEMODataTree` for a nested configuration includes two im
 
     4. Redefine the `dims` and `coords` of each grid dataset to use `i{dom}`, `j{dom}`, `k{dom}` as used to define the semi-discrete equations in NEMO, where *dom* is the unique domain number.
 
-    5. **Clip nested child domains to remove ghost points along the boundaries & add a mapping from the parent grid indices to the child grid indices to the `coords`.**
+    5. **Clip nested child domains to remove ghost points along the boundaries & add a mapping from the parent grid indices to the child grid indices to the `coords`.** Alternatively, the entire child domain, including ghost points, can be retained by passing `nbghost_child=None` to the `.from_paths()` and `.from_datasets()` constructors.
 
     6. **Assemble dictionaries of processed NEMO model grid datasets for each of the parent, child and grandchild domains.**
 
-    7. Assemble the `xarray.DataTree` using a nested dictionary of NEMO model domains.
+    7. Assemble the underlying `xarray.DataTree` using a nested dictionary of NEMO model domains.
 
 ### NEMODataArray :ocean: :simple-databricks:
 ---

nemo_cookbook/nemodatatree.py

@@ -74,7 +74,7 @@ def from_paths(
         linssh: bool = False,
         vco: str = "1d",
         vco_ref: bool = False,
-        nbghost_child: int = 4,
+        nbghost_child: int | None = 4,
         **open_kwargs: dict[str, any],
     ) -> Self:
         """
@@ -150,8 +150,9 @@ def from_paths(
         vco_ref: bool = False
             If True, add reference vertical scale factors and compute reference water column heights from domain files. Default is False.
 
-        nbghost_child : int = 4
+        nbghost_child : int | None = 4
             Number of ghost cells to remove from the western/southern boundaries of the (grand)child domains. Default is 4.
+            If None, no ghost cells are removed and the full (grand)child domain is used.
 
         **open_kwargs : dict, optional
             Additional keyword arguments to pass to `xarray.open_dataset` or `xr.open_mfdataset` when opening NEMO model output files.
@@ -210,9 +211,9 @@ def from_paths(
             )
         if not isinstance(vco_ref, bool):
             raise TypeError("reference vertical coordinates (`vco_ref`) must be a boolean.")
-        if not isinstance(nbghost_child, int):
+        if not isinstance(nbghost_child, (int, type(None))):
             raise TypeError(
-                "number of ghost cells along the western/southern boundaries (`nbghost_child`) must be an integer."
+                "number of ghost cells along the western/southern boundaries (`nbghost_child`) must be an integer or None."
             )
         if not isinstance(open_kwargs, dict):
             raise TypeError("`open_kwargs` must be a dictionary.")
@@ -274,7 +275,7 @@ def from_datasets(
         vco: str = "1d",
         vco_ref: bool = False,
         maskcs: bool = False,
-        nbghost_child: int = 4,
+        nbghost_child: int | None = 4,
     ) -> Self:
         """
         Create a NEMODataTree from a dictionary of `xarray.Dataset` objects created from NEMO model output files,
@@ -333,8 +334,9 @@ def from_datasets(
         vco_ref: bool = False
             If True, add reference vertical scale factors and compute reference water column heights from domain files. Default is False.
 
-        nbghost_child : int = 4
+        nbghost_child : int | None = 4
             Number of ghost cells to remove from the western/southern boundaries of the (grand)child domains. Default is 4.
+            If None, no ghost cells are removed and the full (grand)child domain is used.
 
         Returns
         -------
@@ -387,9 +389,9 @@ def from_datasets(
             )
         if not isinstance(vco_ref, bool):
             raise TypeError("reference vertical coordinates (`vco_ref`) must be a boolean.")
-        if not isinstance(nbghost_child, int):
+        if not isinstance(nbghost_child, (int, type(None))):
             raise TypeError(
-                "number of ghost cells along the western/southern boundaries (`nbghost_child`) must be an integer."
+                "number of ghost cells along the western/southern boundaries (`nbghost_child`) must be an integer or None."
             )
 
         # -- Define parent, child, grandchild dataset collections -- #

nemo_cookbook/processing.py

@@ -58,9 +58,16 @@ def _add_parent_indices(
     i_ic = xr.DataArray(
         np.repeat(i_child, repeats=ds.attrs["rx"]),
         dims=[f"i{label}"],
-        coords={f"i{label}": ds[f"i{label}"]},
     )
+    if i_ic.size < ds[f"i{label}"].size:
+        # Pad i-mapping with NaNs if ghost cells remain in child domain:
+        width = int((ds[f"i{label}"].size - i_ic.size) / 2)
+        i_ic = i_ic.pad({f"i{label}": (width, width)})
+      
     ds[f"i{plabel}_i{label}"] = i_ic
+    ds[f"i{plabel}_i{label}"] = ds[f"i{plabel}_i{label}"].assign_coords(
+        {f"i{label}": ds[f"i{label}"]}
+        )
     ds[f"i{plabel}_i{label}"] = ds[f"i{plabel}_i{label}"].assign_attrs(
         name=f"i{plabel}_i{label}",
         long_name=f"i{plabel} indices of child domain i{label} indices",
@@ -73,9 +80,16 @@ def _add_parent_indices(
     j_jc = xr.DataArray(
         np.repeat(j_child, repeats=ds.attrs["ry"]),
         dims=[f"j{label}"],
-        coords={f"j{label}": ds[f"j{label}"]},
     )
+    if j_jc.size < ds[f"j{label}"].size:
+        # Pad j-mapping with NaNs if ghost cells remain in child domain:
+        width = int((ds[f"j{label}"].size - j_jc.size) / 2)
+        j_jc = j_jc.pad({f"j{label}": (width, width)})
+      
     ds[f"j{plabel}_j{label}"] = j_jc
+    ds[f"j{plabel}_j{label}"] =  ds[f"j{plabel}_j{label}"].assign_coords(
+        {f"j{label}": ds[f"j{label}"]}
+        )
     ds[f"j{plabel}_j{label}"] = ds[f"j{plabel}_j{label}"].assign_attrs(
         name=f"j{plabel}_j{label}",
         long_name=f"j{plabel} indices of child domain j{label} indices",
@@ -906,9 +920,10 @@ def _process_child(
     maskcs : bool = False
         If True, all closed seas are masked using mask_opensea variables from domain files. Default is False.
 
-    nbghost_child : int = _DEFAULT_NBGHOST_CHILD
+    nbghost_child : int | None = _DEFAULT_NBGHOST_CHILD
         Number of ghost cells to remove from the western/southern boundaries of the (grand)child domain.
-        Default is 4 (`_DEFAULT_NBGHOST_CHILD`).
+        Default is 4 (`_DEFAULT_NBGHOST_CHILD`). If None, no ghost cells are removed and the full
+        (grand)child domain is used.
 
     linssh: bool = False
         Linear free-surface approximation. If True, vertical coordinates are time-independent and given by
@@ -971,18 +986,23 @@ def _process_child(
         read_mask=read_mask, maskcs=maskcs, vco=vco, vco_ref=vco_ref
     )
 
-    # Get child domain indices excluding ghost cells:
-    ind_child = _get_child_indices(
-        rx=d_nests.get("rx"),
-        ry=d_nests.get("ry"),
-        imin=d_nests.get("imin"),
-        imax=d_nests.get("imax"),
-        jmin=d_nests.get("jmin"),
-        jmax=d_nests.get("jmax"),
-        nbghost_child=nbghost_child,
-    )
-    i_slice = slice(ind_child[0], ind_child[1] + 1)
-    j_slice = slice(ind_child[2], ind_child[3] + 1)
+    if nbghost_child is not None:
+        # Get child domain indices excluding ghost cells:
+        ind_child = _get_child_indices(
+            rx=d_nests.get("rx"),
+            ry=d_nests.get("ry"),
+            imin=d_nests.get("imin"),
+            imax=d_nests.get("imax"),
+            jmin=d_nests.get("jmin"),
+            jmax=d_nests.get("jmax"),
+            nbghost_child=nbghost_child,
+        )
+        i_slice = slice(ind_child[0], ind_child[1] + 1)
+        j_slice = slice(ind_child[2], ind_child[3] + 1)
+    else:
+        # Use full child domain including ghost cells:
+        i_slice = slice(None)
+        j_slice = slice(None)
 
     # Process T / U / V / W / F grids:
     d_proc_grids = {}
@@ -1077,8 +1097,9 @@ def create_datatree_dict(
         domain variables. Default is False.
     maskcs: bool = False
         If True, all closed seas are masked using mask_opensea variables from domain files. Default is False.
-    nbghost_child : int = _DEFAULT_NBGHOST_CHILD
+    nbghost_child : int | None = _DEFAULT_NBGHOST_CHILD
         Number of ghost cells to remove from the western/southern boundaries of the (grand)child domain. Default is 4 (`_DEFAULT_NBGHOST_CHILD`).
+        If None, no ghost cells are removed and the full (grand)child domain is used.
     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). If False, vertical
         coordinates are time-dependent and must be included in grid datasets. Default is False.