roms-tools 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -7,12 +7,14 @@ import matplotlib.pyplot as plt
 import yaml
 import importlib.metadata
 
+from typing import Union
 from roms_tools.setup.topography import _add_topography_and_mask, _add_velocity_masks
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
 from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
-
+from roms_tools.setup.utils import extract_single_value, save_datasets
 import warnings
+from pathlib import Path
 
 RADIUS_OF_EARTH = 6371315.0  # in m
 
@@ -236,10 +238,17 @@ class Grid:
                 object.__setattr__(self, "ds", ds)
             else:
                 self.ds[coarse_var] = coarse_field
+
         self.ds["mask_coarse"] = xr.where(self.ds["mask_coarse"] > 0.5, 1, 0).astype(
             np.int32
         )
 
+        for fine_var, coarse_var in d.items():
+            self.ds[coarse_var].attrs[
+                "long_name"
+            ] = f"{self.ds[fine_var].attrs['long_name']} on coarsened grid"
+            self.ds[coarse_var].attrs["units"] = self.ds[fine_var].attrs["units"]
+
     def update_vertical_coordinate(self, N, theta_s, theta_b, hc) -> None:
         """
         Create vertical coordinate variables for the ROMS grid.
@@ -275,7 +284,7 @@ class Grid:
             "interface_depth_rho",
             "interface_depth_u",
             "interface_depth_v",
-            "sc_r",
+            "Cs_w",
             "Cs_r",
         ]
 
@@ -290,14 +299,14 @@ class Grid:
         cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
         zw = compute_depth(h * 0, h, hc, cs_w, sigma_w)
 
-        ds["sc_r"] = sigma_r.astype(np.float32)
-        ds["sc_r"].attrs["long_name"] = "S-coordinate at rho-points"
-        ds["sc_r"].attrs["units"] = "nondimensional"
-
         ds["Cs_r"] = cs_r.astype(np.float32)
         ds["Cs_r"].attrs["long_name"] = "S-coordinate stretching curves at rho-points"
         ds["Cs_r"].attrs["units"] = "nondimensional"
 
+        ds["Cs_w"] = cs_w.astype(np.float32)
+        ds["Cs_w"].attrs["long_name"] = "S-coordinate stretching curves at w-points"
+        ds["Cs_w"].attrs["units"] = "nondimensional"
+
         ds.attrs["theta_s"] = np.float32(theta_s)
         ds.attrs["theta_b"] = np.float32(theta_b)
         ds.attrs["hc"] = np.float32(hc)
@@ -361,11 +370,20 @@ class Grid:
         """
 
         if bathymetry:
-            kwargs = {"cmap": "YlGnBu"}
+            field = self.ds.h.where(self.ds.mask_rho)
+            field = field.assign_coords(
+                {"lon": self.ds.lon_rho, "lat": self.ds.lat_rho}
+            )
+
+            vmax = field.max().values
+            vmin = field.min().values
+            cmap = plt.colormaps.get_cmap("YlGnBu")
+            cmap.set_bad(color="gray")
+            kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
             _plot(
                 self.ds,
-                field=self.ds.h.where(self.ds.mask_rho),
+                field=field,
                 straddle=self.straddle,
                 kwargs=kwargs,
             )
@@ -419,10 +437,18 @@ class Grid:
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
             interface_depth = self.ds.interface_depth_rho
+            field = field.where(self.ds.mask_rho)
+            field = field.assign_coords(
+                {"lon": self.ds.lon_rho, "lat": self.ds.lat_rho}
+            )
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
             interface_depth = self.ds.interface_depth_u
+            field = field.where(self.ds.mask_u)
+            field = field.assign_coords({"lon": self.ds.lon_u, "lat": self.ds.lat_u})
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
             interface_depth = self.ds.interface_depth_v
+            field = field.where(self.ds.mask_v)
+            field = field.assign_coords({"lon": self.ds.lon_v, "lat": self.ds.lat_v})
 
         # slice the field as desired
         title = field.long_name
@@ -476,10 +502,9 @@ class Grid:
                 self.ds,
                 field=field,
                 straddle=self.straddle,
-                depth_contours=True,
+                depth_contours=False,
                 title=title,
                 kwargs=kwargs,
-                c="g",
             )
         else:
             if len(field.dims) == 2:
@@ -499,24 +524,60 @@ class Grid:
                 else:
                     _line_plot(field, title=title)
 
-    def save(self, filepath: str) -> None:
+    def save(
+        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+    ) -> None:
         """
         Save the grid information to a netCDF4 file.
 
+        This method supports saving the dataset in two modes:
+
+        1. **Single File Mode (default)**:
+           - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single file at the specified `filepath.nc`.
+
+        2. **Partitioned Mode**:
+           - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
+           - The files are saved as `filepath.0.nc`, `filepath.1.nc`, ..., where the numbering corresponds to the partition index.
+
         Parameters
         ----------
-        filepath
+        filepath : Union[str, Path]
+            The base path or filename where the dataset should be saved.
+        np_eta : int, optional
+            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
+        np_xi : int, optional
+            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
+
+        Returns
+        -------
+        List[Path]
+            A list of Path objects for the filenames that were saved.
         """
-        self.ds.to_netcdf(filepath)
+
+        # Ensure filepath is a Path object
+        filepath = Path(filepath)
+
+        # Remove ".nc" suffix if present
+        if filepath.suffix == ".nc":
+            filepath = filepath.with_suffix("")
+
+        dataset_list = [self.ds.load()]
+        output_filenames = [str(filepath)]
+
+        saved_filenames = save_datasets(
+            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+        )
+
+        return saved_filenames
 
     @classmethod
-    def from_file(cls, filepath: str) -> "Grid":
+    def from_file(cls, filepath: Union[str, Path]) -> "Grid":
         """
         Create a Grid instance from an existing file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             Path to the file containing the grid information.
 
         Returns
@@ -555,8 +616,14 @@ class Grid:
         ):
             grid._coarsen()
 
+        # Move variables to coordinates if necessary
+        for var in ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]:
+            if var not in ds.coords:
+                ds = grid.ds.set_coords(var)
+                object.__setattr__(grid, "ds", ds)
+
         # Update vertical coordinate if necessary
-        if not all(var in grid.ds for var in ["sc_r", "Cs_r"]):
+        if not all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
             N = 100
             theta_s = 5.0
             theta_b = 2.0
@@ -565,13 +632,45 @@ class Grid:
             grid.update_vertical_coordinate(
                 N=N, theta_s=theta_s, theta_b=theta_b, hc=hc
             )
+        else:
+            object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
+            object.__setattr__(grid, "theta_b", ds.attrs["theta_b"].item())
+            object.__setattr__(grid, "hc", ds.attrs["hc"].item())
+            object.__setattr__(grid, "N", len(ds.s_rho))
 
         # Manually set the remaining attributes by extracting parameters from dataset
         object.__setattr__(grid, "nx", ds.sizes["xi_rho"] - 2)
         object.__setattr__(grid, "ny", ds.sizes["eta_rho"] - 2)
-        object.__setattr__(grid, "center_lon", ds.attrs["center_lon"])
-        object.__setattr__(grid, "center_lat", ds.attrs["center_lat"])
-        object.__setattr__(grid, "rot", ds.attrs["rot"])
+        if "center_lon" in ds.attrs:
+            center_lon = ds.attrs["center_lon"]
+        elif "tra_lon" in ds:
+            center_lon = extract_single_value(ds["tra_lon"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'center_lon' attribute or 'tra_lon' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "center_lon", center_lon)
+        if "center_lat" in ds.attrs:
+            center_lat = ds.attrs["center_lat"]
+        elif "tra_lat" in ds:
+            center_lat = extract_single_value(ds["tra_lat"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'center_lat' attribute or 'tra_lat' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "center_lat", center_lat)
+        if "rot" in ds.attrs:
+            rot = ds.attrs["rot"]
+        elif "rotate" in ds:
+            rot = extract_single_value(ds["rotate"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'rot' attribute or 'rotate' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "rot", rot)
 
         for attr in [
             "size_x",
@@ -580,19 +679,25 @@ class Grid:
             "hmin",
         ]:
             if attr in ds.attrs:
-                object.__setattr__(grid, attr, ds.attrs[attr])
+                a = ds.attrs[attr]
+            else:
+                a = None
+            object.__setattr__(grid, attr, a)
 
         return grid
 
-    def to_yaml(self, filepath: str) -> None:
+    def to_yaml(self, filepath: Union[str, Path]) -> None:
         """
         Export the parameters of the class to a YAML file, including the version of roms-tools.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
+
+        filepath = Path(filepath)
+
         data = asdict(self)
         data.pop("ds", None)
         data.pop("straddle", None)
@@ -609,20 +714,20 @@ class Grid:
         # Use the class name as the top-level key
         yaml_data = {self.__class__.__name__: data}
 
-        with open(filepath, "w") as file:
+        with filepath.open("w") as file:
             # Write header
             file.write(header)
             # Write YAML data
             yaml.dump(yaml_data, file, default_flow_style=False)
 
     @classmethod
-    def from_yaml(cls, filepath: str) -> "Grid":
+    def from_yaml(cls, filepath: Union[str, Path]) -> "Grid":
         """
         Create an instance of the class from a YAML file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
 
         Returns
@@ -630,8 +735,10 @@ class Grid:
         Grid
             An instance of the Grid class.
         """
+
+        filepath = Path(filepath)
         # Read the entire file content
-        with open(filepath, "r") as file:
+        with filepath.open("r") as file:
             file_content = file.read()
 
         # Split the content into YAML documents