roms-tools 3.1.1__py3-none-any.whl → 3.2.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -1,9 +1,9 @@
 import importlib.metadata
 import logging
 import re
-import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import xarray as xr
@@ -12,9 +12,10 @@ from matplotlib.axes import Axes
 
 from roms_tools.constants import MAXIMUM_GRID_SIZE, R_EARTH
 from roms_tools.plot import plot
-from roms_tools.setup.mask import _add_mask, _add_velocity_masks
-from roms_tools.setup.topography import _add_topography
+from roms_tools.setup.mask import add_mask, add_velocity_masks
+from roms_tools.setup.topography import add_topography
 from roms_tools.setup.utils import (
+    Timed,
     extract_single_value,
     gc_dist,
     get_target_coords,
@@ -64,6 +65,8 @@ class Grid:
         - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file. Can be a string or a Path object.
 
         The default is "ETOPO5", which does not require a path.
+    mask_shapefile: str | Path | None, optional
+        Path to a custom shapefile to use to determine the land mask; if None, use NaturalEarth 10m.
     hmin : float, optional
        The minimum ocean depth (in meters). The default is 5.0.
     N : int, optional
@@ -106,8 +109,10 @@ class Grid:
     """The bottom control parameter."""
     hc: float = 300.0
     """The critical depth (in meters)."""
-    topography_source: dict[str, str | Path | list[str | Path]] = None
+    topography_source: dict[str, str | Path | list[str | Path]] | None = None
     """Dictionary specifying the source of the topography data."""
+    mask_shapefile: str | Path | None = None
+    """Path to a custom shapefile to use to determine the landmask; if None, use NaturalEarth 10m."""
     hmin: float = 5.0
     """The minimum ocean depth (in meters)."""
     verbose: bool = False
@@ -129,7 +134,7 @@ class Grid:
         self._straddle()
 
         # Mask
-        self._create_mask(verbose=self.verbose)
+        self.update_mask(mask_shapefile=self.mask_shapefile, verbose=self.verbose)
 
         # Coarsen the dataset if needed
         self._coarsen()
@@ -165,19 +170,35 @@ class Grid:
                     "`topography_source` must include a 'path' key when the 'name' is not 'ETOPO5'."
                 )
 
-    def _create_mask(self, verbose=False) -> None:
-        if verbose:
-            start_time = time.time()
-            logging.info("=== Creating the mask ===")
-        ds = _add_mask(self.ds)
+    def update_mask(
+        self, mask_shapefile: str | Path | None = None, verbose: bool = False
+    ) -> None:
+        """
+        Update the land mask of the current grid dataset.
 
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
-            )
+        This method generates a land mask based on the provided coastline
+        shapefile, fills enclosed basins with lands and updates the dataset
+        stored in `self.ds`. If no shapefile is provided, a default dataset (Natural
+        Earth 10m) is used. The operation is optionally timed and logged.
 
-        self.ds = ds
+        Parameters
+        ----------
+        mask_shapefile : str or Path, optional
+            Path to a coastal shapefile to derive the land mask. If `None`,
+            the default Natural Earth 10m coastline dataset is used.
+        verbose : bool, default False
+            If True, prints timing and progress information.
+
+        Returns
+        -------
+        None
+            Updates the `self.ds` attribute in place with the new mask.
+
+        """
+        with Timed("=== Deriving the mask from coastlines ===", verbose=verbose):
+            ds = add_mask(self.ds, shapefile=mask_shapefile)
+            self.ds = ds
+            self.mask_shapefile = mask_shapefile
 
     def update_topography(
         self, topography_source=None, hmin=None, verbose=False
@@ -218,33 +239,22 @@ class Grid:
         # Extract target coordinates for processing
         target_coords = get_target_coords(self)
 
-        # If verbose is enabled, start the timer and print the start message
-        if verbose:
-            start_time = time.time()
-            logging.info(
-                f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ==="
-            )
-
-        # Add topography to the dataset
-        ds = _add_topography(
-            ds=self.ds,
-            target_coords=target_coords,
-            topography_source=topography_source,
-            hmin=hmin,
+        with Timed(
+            f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ===",
             verbose=verbose,
-        )
-
-        # If verbose is enabled, print elapsed time and a separator
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
+        ):
+            ds = add_topography(
+                ds=self.ds,
+                target_coords=target_coords,
+                topography_source=topography_source,
+                hmin=hmin,
+                verbose=verbose,
             )
 
-        # Update the grid's dataset and related attributes
-        self.ds = ds
-        self.topography_source = topography_source
-        self.hmin = hmin
+            # Update the grid's dataset and related attributes
+            self.ds = ds
+            self.topography_source = topography_source
+            self.hmin = hmin
 
     def update_vertical_coordinate(
         self, N=None, theta_s=None, theta_b=None, hc=None, verbose=False
@@ -281,69 +291,61 @@ class Grid:
         theta_b = theta_b or self.theta_b
         hc = hc or self.hc
 
-        if verbose:
-            start_time = time.time()
-            logging.info(
-                f"=== Preparing the vertical coordinate system using N = {N}, theta_s = {theta_s}, theta_b = {theta_b}, hc = {hc} ==="
-            )
+        with Timed(
+            f"=== Preparing the vertical coordinate system using N = {N}, theta_s = {theta_s}, theta_b = {theta_b}, hc = {hc} ===",
+            verbose=verbose,
+        ):
+            ds = self.ds
+            # need to drop vertical coordinates because they could cause conflict if N changed
+            vars_to_drop = [
+                "layer_depth_rho",
+                "layer_depth_u",
+                "layer_depth_v",
+                "interface_depth_rho",
+                "interface_depth_u",
+                "interface_depth_v",
+                "sigma_r",
+                "sigma_w",
+                "Cs_w",
+                "Cs_r",
+            ]
 
-        ds = self.ds
-        # need to drop vertical coordinates because they could cause conflict if N changed
-        vars_to_drop = [
-            "layer_depth_rho",
-            "layer_depth_u",
-            "layer_depth_v",
-            "interface_depth_rho",
-            "interface_depth_u",
-            "interface_depth_v",
-            "sigma_r",
-            "sigma_w",
-            "Cs_w",
-            "Cs_r",
-        ]
-
-        for var in vars_to_drop:
-            if var in ds.variables:
-                ds = ds.drop_vars(var)
-
-        cs_r, sigma_r = sigma_stretch(theta_s, theta_b, N, "r")
-        cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
-
-        ds["sigma_r"] = sigma_r.astype(np.float32)
-        ds["sigma_r"].attrs["long_name"] = (
-            "Fractional vertical stretching coordinate at rho-points"
-        )
-        ds["sigma_r"].attrs["units"] = "nondimensional"
+            for var in vars_to_drop:
+                if var in ds.variables:
+                    ds = ds.drop_vars(var)
 
-        ds["Cs_r"] = cs_r.astype(np.float32)
-        ds["Cs_r"].attrs["long_name"] = "Vertical stretching function at rho-points"
-        ds["Cs_r"].attrs["units"] = "nondimensional"
+            cs_r, sigma_r = sigma_stretch(theta_s, theta_b, N, "r")
+            cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
 
-        ds["sigma_w"] = sigma_w.astype(np.float32)
-        ds["sigma_w"].attrs["long_name"] = (
-            "Fractional vertical stretching coordinate at w-points"
-        )
-        ds["sigma_w"].attrs["units"] = "nondimensional"
-
-        ds["Cs_w"] = cs_w.astype(np.float32)
-        ds["Cs_w"].attrs["long_name"] = "Vertical stretching function at w-points"
-        ds["Cs_w"].attrs["units"] = "nondimensional"
+            ds["sigma_r"] = sigma_r.astype(np.float32)
+            ds["sigma_r"].attrs["long_name"] = (
+                "Fractional vertical stretching coordinate at rho-points"
+            )
+            ds["sigma_r"].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)
+            ds["Cs_r"] = cs_r.astype(np.float32)
+            ds["Cs_r"].attrs["long_name"] = "Vertical stretching function at rho-points"
+            ds["Cs_r"].attrs["units"] = "nondimensional"
 
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
+            ds["sigma_w"] = sigma_w.astype(np.float32)
+            ds["sigma_w"].attrs["long_name"] = (
+                "Fractional vertical stretching coordinate at w-points"
             )
+            ds["sigma_w"].attrs["units"] = "nondimensional"
 
-        self.ds = ds
-        self.theta_s = theta_s
-        self.theta_b = theta_b
-        self.hc = hc
-        self.N = N
+            ds["Cs_w"] = cs_w.astype(np.float32)
+            ds["Cs_w"].attrs["long_name"] = "Vertical stretching function 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)
+
+            self.ds = ds
+            self.theta_s = theta_s
+            self.theta_b = theta_b
+            self.hc = hc
+            self.N = N
 
     def _straddle(self) -> None:
         """Check if the Greenwich meridian goes through the domain.
@@ -415,30 +417,57 @@ class Grid:
 
     def plot(
         self,
+        lat: float | None = None,
+        lon: float | None = None,
         with_dim_names: bool = False,
         save_path: str | None = None,
     ) -> None:
-        """Plot the grid.
+        """Plot the grid with bathymetry.
+
+        Depending on the arguments, this will either:
+          * Plot the full horizontal grid (if both `lat` and `lon` are None),
+          * Plot a zonal (east-west) vertical section at a given latitude (`lat`),
+          * Plot a meridional (south-north) vertical section at a given longitude (`lon`).
 
         Parameters
         ----------
+        lat : float, optional
+            Latitude in degrees at which to plot a vertical (zonal) section. Cannot be
+            provided together with `lon`. Default is None.
+
+        lon : float, optional
+            Longitude in degrees at which to plot a vertical (meridional) section. Cannot be
+            provided together with `lat`. Default is None.
+
         with_dim_names : bool, optional
-            Whether or not to plot the dimension names. Default is False.
+            If True and no section is requested (i.e., both `lat` and `lon` are None), annotate
+            the plot with the underlying dimension names. Default is False.
 
         save_path : str, optional
             Path to save the generated plot. If None, the plot is shown interactively.
             Default is None.
 
+        Raises
+        ------
+        ValueError
+            If both `lat` and `lon` are specified simultaneously.
+
         Returns
         -------
         None
             This method does not return any value. It generates and displays a plot.
         """
+        if lat is not None and lon is not None:
+            raise ValueError("Specify either `lat` or `lon`, not both.")
+
         field = self.ds["h"]
 
         plot(
             field=field,
             grid_ds=self.ds,
+            lat=lat,
+            lon=lon,
+            yincrease=False,
             with_dim_names=with_dim_names,
             save_path=save_path,
             cmap_name="YlGnBu",
@@ -576,7 +605,7 @@ class Grid:
         ds = xr.open_dataset(filepath)
 
         if not all(mask in ds for mask in ["mask_u", "mask_v"]):
-            ds = _add_velocity_masks(ds)
+            ds = add_velocity_masks(ds)
 
         # Create a new Grid instance without calling __init__ and __post_init__
         grid = cls.__new__(cls)
@@ -731,24 +760,30 @@ class Grid:
             "hmin",
         ]:
             if attr in ds.attrs:
-                a = float(ds.attrs[attr])
+                value = float(ds.attrs[attr])
             else:
-                a = None
+                value = None
 
-            object.__setattr__(grid, attr, a)
+            object.__setattr__(grid, attr, value)
 
         if "topography_source_name" in ds.attrs:
             if "topography_source_path" in ds.attrs:
-                a = {
+                topo_source = {
                     "name": ds.attrs["topography_source_name"],
                     "path": ds.attrs["topography_source_path"],
                 }
             else:
-                a = {"name": ds.attrs["topography_source_name"]}
+                topo_source = {"name": ds.attrs["topography_source_name"]}
+        else:
+            topo_source = None
+        grid.topography_source = topo_source
+
+        if "mask_shapefile" in ds.attrs:
+            mask_shapefile = ds.attrs["mask_shapefile"]
         else:
-            a = None
+            mask_shapefile = None
 
-        object.__setattr__(grid, "topography_source", a)
+        grid.mask_shapefile = mask_shapefile
 
         return grid
 
@@ -769,10 +804,7 @@ class Grid:
 
     @classmethod
     def from_yaml(
-        cls,
-        filepath: str | Path,
-        section_name: str = "Grid",
-        verbose: bool = False,
+        cls, filepath: str | Path, verbose: bool = False, **kwargs: Any
     ) -> "Grid":
         """Create an instance of the class from a YAML file.
 
@@ -780,10 +812,13 @@ class Grid:
         ----------
         filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
-        section_name : str, optional
-            The name of the YAML section containing the grid configuration. Defaults to "Grid".
         verbose : bool, optional
             Indicates whether to print grid generation steps with timing. Defaults to False.
+        **kwargs : Any
+            Additional keyword arguments:
+
+            - section_name : str, optional (default: "Grid")
+              The name of the YAML section containing the grid configuration.
 
         Returns
         -------
@@ -801,6 +836,8 @@ class Grid:
         Issues a warning if the ROMS-Tools version in the YAML header does not match the
         currently installed version.
         """
+        section_name: str = kwargs.pop("section_name", None) or "Grid"
+
         filepath = Path(filepath)
         # Read the entire file content
         with filepath.open("r") as file:
@@ -854,7 +891,7 @@ class Grid:
         attr_str = ", ".join(f"{k}={v!r}" for k, v in attr_dict.items())
         return f"{cls_name}({attr_str})"
 
-    def _create_horizontal_grid(self) -> xr.Dataset():
+    def _create_horizontal_grid(self) -> xr.Dataset:
         """Create the horizontal grid based on a Mercator projection and store it in the
         'ds' attribute.
 
@@ -872,41 +909,32 @@ class Grid:
         - Longitude values are adjusted to fall within the range [0, 360].
         - Grid rotation and translation are applied based on the specified parameters.
         """
-        if self.verbose:
-            start_time = time.time()
-            logging.info("=== Creating the horizontal grid ===")
-
-        self._raise_if_domain_size_too_large()
+        with Timed("=== Creating the horizontal grid ===", verbose=self.verbose):
+            self._raise_if_domain_size_too_large()
 
-        coords = self._make_initial_lon_lat_ds()
+            coords = self._make_initial_lon_lat_ds()
 
-        # rotate coordinate system
-        coords = _rotate(coords, self.rot)
+            # rotate coordinate system
+            coords = _rotate(coords, self.rot)
 
-        # translate coordinate system
-        coords = _translate(coords, self.center_lat, self.center_lon)
+            # translate coordinate system
+            coords = _translate(coords, self.center_lat, self.center_lon)
 
-        # compute 1/dx and 1/dy
-        coords["pm"], coords["pn"] = _compute_coordinate_metrics(coords)
+            # compute 1/dx and 1/dy
+            coords["pm"], coords["pn"] = _compute_coordinate_metrics(coords)
 
-        # compute angle of local grid positive x-axis relative to east
-        coords["angle"] = _compute_angle(coords)
+            # compute angle of local grid positive x-axis relative to east
+            coords["angle"] = _compute_angle(coords)
 
-        # make sure lons are in [0, 360] range
-        for lon in ["lon", "lonu", "lonv", "lonq"]:
-            coords[lon][coords[lon] < 0] = coords[lon][coords[lon] < 0] + 2 * np.pi
+            # make sure lons are in [0, 360] range
+            for lon in ["lon", "lonu", "lonv", "lonq"]:
+                coords[lon][coords[lon] < 0] = coords[lon][coords[lon] < 0] + 2 * np.pi
 
-        ds = self._create_grid_ds(coords)
+            ds = self._create_grid_ds(coords)
 
-        ds = self._add_global_metadata(ds)
+            ds = self._add_global_metadata(ds)
 
-        if self.verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
-            )
-
-        self.ds = ds
+            self.ds = ds
 
     def _add_global_metadata(self, ds):
         """Add global metadata and attributes to the dataset.