roms-tools 3.0.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/fill.py

@@ -24,7 +24,6 @@ class LateralFill:
         NotImplementedError
             If the input mask has more than two dimensions, which is not supported by the current implementation.
         """
-
         if len(mask.shape) > 2:
             raise NotImplementedError("LateralFill currently supports only 2D masks.")
 
@@ -121,7 +120,6 @@ def _lateral_fill_np_array(x0, b, ml, tol=1.0e-4):
         The filled 2D array where NaN values have been replaced with iteratively
         computed values, and non-NaN values remain unchanged.
     """
-
     b_flat = b.flatten()
     x0_flat = x0.flatten()
     x = ml.solve(b_flat, x0_flat, tol=tol)
@@ -217,7 +215,6 @@ def stencil_grid_mod(S, grid, msk, dtype=None, format=None):
     The stencil is applied in all directions, and boundary conditions are
     respected by zeroing out connections to boundary points.
     """
-
     S = np.asarray(S, dtype=dtype)
     grid = tuple(grid)
 

roms_tools/setup/grid.py

@@ -4,25 +4,24 @@ import re
 import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Dict, List, Union
 
-import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
 import yaml
+from matplotlib.axes import Axes
 
 from roms_tools.constants import MAXIMUM_GRID_SIZE, R_EARTH
-from roms_tools.plot import _plot, _section_plot
+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.utils import (
-    _pop_grid_data,
-    _write_to_yaml,
     extract_single_value,
     gc_dist,
     get_target_coords,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    pop_grid_data,
+    write_to_yaml,
 )
 from roms_tools.utils import save_datasets
 from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
@@ -85,24 +84,42 @@ class Grid:
     """
 
     nx: int
+    """Number of grid points in the x-direction."""
     ny: int
+    """Number of grid points in the y-direction."""
     size_x: float
+    """Domain size in the x-direction (in kilometers)."""
     size_y: float
+    """Domain size in the y-direction (in kilometers)."""
     center_lon: float
+    """Longitude of grid center."""
     center_lat: float
+    """Latitude of grid center."""
     rot: float = 0
+    """Rotation of grid x-direction from lines of constant latitude, measured in
+    degrees."""
     N: int = 100
+    """The number of vertical levels."""
     theta_s: float = 5.0
+    """The surface control parameter."""
     theta_b: float = 2.0
+    """The bottom control parameter."""
     hc: float = 300.0
-    topography_source: Dict[str, Union[str, Path, List[Union[str, Path]]]] = None
+    """The critical depth (in meters)."""
+    topography_source: dict[str, str | Path | list[str | Path]] = None
+    """Dictionary specifying the source of the topography data."""
     hmin: float = 5.0
+    """The minimum ocean depth (in meters)."""
     verbose: bool = False
+    """Whether to print grid generation steps with timing."""
+
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
     straddle: bool = field(init=False, repr=False)
+    """Whether the grid straddles the dateline."""
 
     def __post_init__(self):
-
         self._input_checks()
 
         # Horizontal grid
@@ -149,7 +166,6 @@ class Grid:
                 )
 
     def _create_mask(self, verbose=False) -> None:
-
         if verbose:
             start_time = time.time()
             logging.info("=== Creating the mask ===")
@@ -196,7 +212,6 @@ class Grid:
             This method updates the internal dataset (`self.ds`) in place by adding or overwriting the
             topography variable. It does not return any value.
         """
-
         topography_source = topography_source or self.topography_source
         hmin = hmin or self.hmin
 
@@ -261,7 +276,6 @@ class Grid:
         None
             This method modifies the dataset in place by adding vertical coordinate variables.
         """
-
         N = N or self.N
         theta_s = theta_s or self.theta_s
         theta_b = theta_b or self.theta_b
@@ -296,9 +310,9 @@ class Grid:
         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["long_name"] = (
+            "Fractional vertical stretching coordinate at rho-points"
+        )
         ds["sigma_r"].attrs["units"] = "nondimensional"
 
         ds["Cs_r"] = cs_r.astype(np.float32)
@@ -306,9 +320,9 @@ class Grid:
         ds["Cs_r"].attrs["units"] = "nondimensional"
 
         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["long_name"] = (
+            "Fractional vertical stretching coordinate at w-points"
+        )
         ds["sigma_w"].attrs["units"] = "nondimensional"
 
         ds["Cs_w"] = cs_w.astype(np.float32)
@@ -341,7 +355,6 @@ class Grid:
         The check is based on whether the longitudinal differences between adjacent
         points exceed 300 degrees, indicating a potential wraparound of longitude.
         """
-
         if (
             np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
             or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
@@ -350,6 +363,8 @@ class Grid:
         else:
             self.straddle = False
 
+        self.ds.attrs["straddle"] = str(self.straddle)
+
     def _coarsen(self):
         """Update the grid by adding grid variables that are coarsened versions of the
         original fine-resoluion grid variables. The coarsening is by a factor of two.
@@ -399,57 +414,44 @@ class Grid:
         self.ds = ds
 
     def plot(
-        self, bathymetry: bool = True, title: str = None, with_dim_names: bool = False
+        self,
+        with_dim_names: bool = False,
+        save_path: str | None = None,
     ) -> None:
         """Plot the grid.
 
         Parameters
         ----------
-        bathymetry : bool, optional
-            Whether or not to plot the bathymetry. Default is True.
-        title : str, optional
-            The title of the plot. If not provided, it will be set to a default.
         with_dim_names : bool, optional
             Whether or not to plot the 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.
 
         Returns
         -------
         None
             This method does not return any value. It generates and displays a plot.
         """
+        field = self.ds["h"]
+
+        plot(
+            field=field,
+            grid_ds=self.ds,
+            with_dim_names=with_dim_names,
+            save_path=save_path,
+            cmap_name="YlGnBu",
+        )
 
-        field = self.ds.h.where(self.ds.mask_rho)
-        lat_deg = self.ds.lat_rho
-        lon_deg = self.ds.lon_rho
-        if self.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-        if bathymetry:
-            if title is None:
-                title = "ROMS grid and bathymetry"
-
-            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(
-                field=field,
-                title=title,
-                with_dim_names=with_dim_names,
-                kwargs=kwargs,
-            )
-        else:
-            if title is None:
-                title = "ROMS grid"
-            _plot(
-                field=field, title=title, with_dim_names=with_dim_names, plot_data=False
-            )
-
-    def plot_vertical_coordinate(self, s=None, eta=None, xi=None, ax=None) -> None:
+    def plot_vertical_coordinate(
+        self,
+        s: int | None = None,
+        eta: int | None = None,
+        xi: int | None = None,
+        max_nr_layer_contours: int | None = None,
+        ax: Axes | None = None,
+    ) -> None:
         """Plot the layer depth for a given eta-, xi-, or s-slice.
 
         Parameters
@@ -460,6 +462,8 @@ class Grid:
             The eta-index to plot. Default is None.
         xi : int, optional
             The xi-index to plot. Default is None.
+        max_nr_layer_contours : int, optional
+            Number of layer contours displayed. Only relevant when plotting a vertical section. Default is all.
         ax : matplotlib.axes.Axes, optional
             The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots.
 
@@ -473,62 +477,34 @@ class Grid:
         ValueError
             If not exactly one of s, eta, xi is specified.
         """
-
-        if sum(s is not None for s in [s, eta, xi]) != 1:
+        if sum(i is not None for i in [s, eta, xi]) != 1:
             raise ValueError("Exactly one of s, eta, or xi must be specified.")
 
-        depth = compute_depth_coordinates(
-            self.ds, zeta=0, depth_type="layer", location="rho", eta=eta, xi=xi
-        )
-
-        title = "Layer depth at rho-points"
-        if eta is not None:
-            title = title + f", eta_rho = {self.ds.eta_rho[eta].item()}"
-        if xi is not None:
-            title = title + f", xi_rho = {self.ds.xi_rho[xi].item()}"
         if s is not None:
-            title = title + f", s_rho = {depth.s_rho[s].item()}"
-            depth = depth.isel(s_rho=s)
-
-            vmax = depth.max().values
-            vmin = depth.min().values
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-            lat_deg = self.ds.lat_rho
-            lon_deg = self.ds.lon_rho
-            if self.straddle:
-                lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-            depth = depth.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-            _plot(
-                field=depth.where(self.ds.mask_rho),
-                depth_contours=False,
-                title=title,
-                kwargs=kwargs,
+            field = compute_depth_coordinates(
+                self.ds, zeta=0, depth_type="layer", location="rho"
             )
+            layer_contours = False
+            add_colorbar = True
         else:
-            field = xr.zeros_like(depth)
-            field = field.assign_coords({"layer_depth": depth})
-
-            interface_depth = compute_depth_coordinates(
-                self.ds, zeta=0, depth_type="interface", eta=eta, xi=xi
-            )
-
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": 0.0, "vmin": 0.0, "cmap": cmap, "add_colorbar": False}
-
-            _section_plot(
-                field=field,
-                interface_depth=interface_depth,
-                title=title,
-                kwargs=kwargs,
-                ax=ax,
-            )
+            field = xr.zeros_like(self.ds.s_rho * self.ds["h"])
+            field.attrs["long_name"] = "layer depth at rho-points"
+            layer_contours = True
+            add_colorbar = False
+
+        plot(
+            field=field,
+            grid_ds=self.ds,
+            s=s,
+            eta=eta,
+            xi=xi,
+            layer_contours=layer_contours,
+            max_nr_layer_contours=max_nr_layer_contours,
+            cmap_name="YlGnBu",
+            add_colorbar=add_colorbar,
+        )
 
-    def save(self, filepath: Union[str, Path]) -> None:
+    def save(self, filepath: str | Path) -> None:
         """Save the grid information to a netCDF4 file.
 
         Parameters
@@ -541,7 +517,6 @@ class Grid:
         List[Path]
             A list of Path objects for the filenames that were saved.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -777,7 +752,7 @@ class Grid:
 
         return grid
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -787,15 +762,15 @@ class Grid:
             The path to the YAML file where the parameters will be saved.
         """
         data = asdict(self)
-        data = _pop_grid_data(data)
+        data = pop_grid_data(data)
         forcing_dict = {self.__class__.__name__: data}
 
-        _write_to_yaml(forcing_dict, filepath)
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         section_name: str = "Grid",
         verbose: bool = False,
     ) -> "Grid":
@@ -826,7 +801,6 @@ class Grid:
         Issues a warning if the ROMS-Tools version in the YAML header does not match the
         currently installed version.
         """
-
         filepath = Path(filepath)
         # Read the entire file content
         with filepath.open("r") as file:
@@ -869,7 +843,8 @@ class Grid:
 
     def __repr__(self) -> str:
         """Return a string representation of the object with non-None attributes,
-        excluding 'ds'."""
+        excluding 'ds'.
+        """
         cls = self.__class__
         cls_name = cls.__name__
         # Filter attributes to exclude 'ds' and those with None values
@@ -999,42 +974,58 @@ class Grid:
             - lonv, latv: 2D arrays of longitudes and latitudes at v-points.
             - lonq, latq: 2D arrays of longitudes and latitudes at cell corners.
         """
-
         # initially define the domain to be longer in x-direction (dimension "length")
         # than in y-direction (dimension "width") to keep grid distortion minimal
         if self.size_y > self.size_x:
-            domain_length, domain_width = self.size_y * 1e3, self.size_x * 1e3  # in m
+            domain_length = self.size_y * 1e3  # in meters
+            domain_width = self.size_x * 1e3
             nl, nw = self.ny, self.nx
         else:
-            domain_length, domain_width = self.size_x * 1e3, self.size_y * 1e3  # in m
+            domain_length = self.size_x * 1e3
+            domain_width = self.size_y * 1e3
             nl, nw = self.nx, self.ny
 
-        domain_length_in_degrees = domain_length / R_EARTH
-        domain_width_in_degrees = domain_width / R_EARTH
+        dlon = domain_length / R_EARTH
+        dlat = domain_width / R_EARTH
 
-        # Generate 1D longitude arrays at cell centers and corners
-        lon_array_1d_in_degrees = domain_length_in_degrees * (
-            np.arange(-0.5, nl + 1.5) / nl - 0.5
-        )
-        lonq_array_1d_in_degrees_q = domain_length_in_degrees * (
-            np.arange(-1, nl + 2) / nl - 0.5
-        )
+        # 1D longitude arrays at centers and corners (cell and cell+1)
+        lon_array_1d_in_degrees = dlon * np.arange(-0.5, nl + 1.5) / nl - dlon / 2
+        lonq_array_1d_in_degrees = dlon * np.arange(-1, nl + 2) / nl - dlon / 2
 
-        # Mercator projection for latitude
-        y1 = np.log(np.tan(np.pi / 4 - domain_width_in_degrees / 4))
-        y2 = np.log(np.tan(np.pi / 4 + domain_width_in_degrees / 4))
+        mul = 1.0
+        # Central longitude spacing (in radians)
+        dlon_cen = dlon / nl
 
-        # Generate 1D latitude arrays with inverse Mercator projection
-        lat_array_1d_in_degrees = np.arctan(
-            np.sinh((y2 - y1) * (np.arange(-0.5, nw + 1.5) / nw) + y1)
-        )
+        # Iteratively adjust latitude to make grid cells square using aspect ratio correction
+        for _ in range(100):
+            # Compute Mercator y-limits
+            y1 = np.log(np.tan(np.pi / 4 - dlat / 4))
+            y2 = np.log(np.tan(np.pi / 4 + dlat / 4))
+
+            # Mercator-transformed latitude coordinates
+            y = (y2 - y1) * np.arange(-0.5, nw + 1.5) / nw + y1
+            lat_array_1d_in_degrees = np.arctan(np.sinh(y))  # Inverse Mercator
+
+            # Central latitude spacing (in radians)
+            i_mid = round(nw / 2)
+            dlat_cen = 0.5 * (
+                lat_array_1d_in_degrees[i_mid + 1] - lat_array_1d_in_degrees[i_mid - 1]
+            )
+
+            # Update multiplicative scaling factor
+            mul = dlat_cen / dlon_cen * domain_length / domain_width * nw / nl
+
+            # Update dlat for next iteration
+            dlat /= mul
+
+        # Latitudes of corners
         latq_array_1d_in_degrees = np.arctan(
-            np.sinh((y2 - y1) * (np.arange(-1, nw + 2) / nw) + y1)
+            np.sinh((y2 - y1) * np.arange(-1, nw + 2) / nw + y1)
         )
 
         # 2D grids for cell centers and corners
         lon, lat = np.meshgrid(lon_array_1d_in_degrees, lat_array_1d_in_degrees)
-        lonq, latq = np.meshgrid(lonq_array_1d_in_degrees_q, latq_array_1d_in_degrees)
+        lonq, latq = np.meshgrid(lonq_array_1d_in_degrees, latq_array_1d_in_degrees)
 
         if self.size_y > self.size_x:
             # Rotate grid by 90 degrees because until here the grid has been defined
@@ -1084,7 +1075,6 @@ class Grid:
             Dataset with variables: lon_rho, lat_rho, lon_u, lat_u, lon_v, lat_v,
             angle, f (Coriolis parameter), pm, pn.
         """
-
         ds = xr.Dataset()
 
         lon_rho = xr.Variable(
@@ -1178,45 +1168,6 @@ class Grid:
 
         return ds
 
-    def _infer_nominal_horizontal_resolution(self, lat=None):
-        """Estimate the nominal horizontal resolution of the grid in degrees at a
-        specified latitude.
-
-        This method calculates the nominal horizontal resolution of the grid by first
-        determining the average grid spacing in meters. The spacing is then converted
-        to degrees, accounting for the Earth's curvature, and the latitude where the
-        resolution is being computed.
-
-        Parameters
-        ----------
-        lat : float, optional
-            Latitude (in degrees) at which to estimate the horizontal resolution.
-            If not provided, the resolution is calculated at the average latitude of
-            the grid (`lat_rho`).
-
-        Returns
-        -------
-        float
-            The estimated horizontal resolution in degrees, adjusted for the Earth's curvature.
-        """
-        if lat is None:
-            # Center latitude in degrees
-            lat = (self.ds.lat_rho.max() + self.ds.lat_rho.min()) / 2
-
-        # Convert latitude to radians
-        lat_rad = np.deg2rad(lat)
-
-        # Mean resolution in meters
-        resolution_in_m = ((1 / self.ds.pm).mean() + (1 / self.ds.pn).mean()) / 2
-
-        # Meters per degree at the equator
-        meters_per_degree = 2 * np.pi * R_EARTH / 360
-
-        # Correct for latitude by multiplying by cos(latitude) for longitude
-        resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
-
-        return resolution_in_degrees
-
     def _compute_exponential_depth_levels(self, Nz=None, max_depth=None, h=None):
         """Compute vertical grid center and face depths using an exponential profile.
 
@@ -1265,7 +1216,6 @@ class Grid:
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""
-
     (coords["lon"], coords["lat"]) = _rot_sphere(coords["lon"], coords["lat"], rot)
     (coords["lonu"], coords["latu"]) = _rot_sphere(coords["lonu"], coords["latu"], rot)
     (coords["lonv"], coords["latv"]) = _rot_sphere(coords["lonv"], coords["latv"], rot)
@@ -1275,8 +1225,7 @@ def _rotate(coords, rot):
 
 
 def _translate(coords, tra_lat, tra_lon):
-    """Translate grid so that the centre lies at the position (tra_lat, tra_lon)"""
-
+    """Translate grid so that the centre lies at the position (tra_lat, tra_lon)."""
     (lon, lat) = _tra_sphere(coords["lon"], coords["lat"], tra_lat)
     (lonu, latu) = _tra_sphere(coords["lonu"], coords["latu"], tra_lat)
     (lonv, latv) = _tra_sphere(coords["lonv"], coords["latv"], tra_lat)
@@ -1375,7 +1324,6 @@ def _tra_sphere(lon, lat, tra):
     tuple
         Translated longitude and latitude arrays (lon, lat) in radians.
     """
-
     # Convert translation angle from degrees to radians
     tra = tra * np.pi / 180
 
@@ -1430,7 +1378,6 @@ def _compute_coordinate_metrics(coords):
     -----
     Boundary values of `pn` and `pm` are copied from adjacent interior values.
     """
-
     # pm = 1/dx
     pmu = gc_dist(
         coords["lonu"][:, :-1],
@@ -1482,7 +1429,6 @@ def _compute_angle(coords):
         An array of angles (in radians) of the local grid's positive x-axis
         relative to east for each grid point.
     """
-
     # Compute differences in latitudes and longitudes
     dellat = coords["latu"][:, 1:] - coords["latu"][:, :-1]
     dellon = coords["lonu"][:, 1:] - coords["lonu"][:, :-1]
@@ -1521,7 +1467,6 @@ def _f2c(f):
     fc : xarray.DataArray
         Output DataArray with modified dimensions and values.
     """
-
     fc = _f2c_xdir(f)
     fc = fc.transpose()
     fc = _f2c_xdir(fc)