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

ci/environment-with-xesmf.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   - xesmf
   # testing
   - zarr

ci/environment.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   # testing
   - zarr
   - pytest

roms_tools/__init__.py

@@ -1,5 +1,5 @@
+import logging
 from importlib.metadata import version as _version
-import logging  # noqa: F811
 
 try:
     __version__ = _version("roms_tools")
@@ -7,17 +7,18 @@ except ImportError:  # pragma: no cover
     # Local copy or not installed with setuptools
     __version__ = "9999"
 
-
-from roms_tools.setup.grid import Grid  # noqa: F401
-from roms_tools.setup.tides import TidalForcing  # noqa: F401
-from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
-from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
+# grid must be imported first
+from roms_tools.setup.grid import Grid  # noqa: I001, F401
+from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
-from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
-from roms_tools.setup.cdr_forcing import CDRVolumePointSource  # noqa: F401
+from roms_tools.setup.cdr_forcing import CDRForcing  # noqa: F401
+from roms_tools.setup.cdr_release import TracerPerturbation, VolumeRelease  # noqa: F401
+from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.nesting import ChildGrid  # noqa: F401
+from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
+from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
+from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
-from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 
 # Configure logging when the package is imported
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")

roms_tools/analysis/roms_output.py

@@ -1,27 +1,27 @@
-import xarray as xr
-import numpy as np
-import matplotlib.pyplot as plt
-from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
-from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
-from dataclasses import dataclass, field
-from typing import Union, Optional
-from pathlib import Path
-import re
 import logging
+import re
 import warnings
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+from matplotlib.axes import Axes
+
 from roms_tools import Grid
-from roms_tools.vertical_coordinate import (
-    compute_depth_coordinates,
-)
+from roms_tools.plot import plot
+from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from roms_tools.utils import (
+    _generate_coordinate_range,
     _load_data,
+    infer_nominal_horizontal_resolution,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    _generate_coordinate_range,
-    _remove_edge_nans,
 )
-from roms_tools.analysis.utils import _validate_plot_inputs
+from roms_tools.vertical_coordinate import (
+    compute_depth_coordinates,
+)
 
 
 @dataclass(kw_only=True)
@@ -35,6 +35,7 @@ class ROMSOutput:
     path : Union[str, Path, List[Union[str, Path]]]
         Filename, or list of filenames with model output.
     model_reference_date : datetime, optional
+        Reference date of ROMS simulation.
         If not specified, this is inferred from metadata of the model output
         If specified and does not coincide with metadata, a warning is raised.
     adjust_depth_for_sea_surface_height : bool, optional
@@ -45,14 +46,21 @@ class ROMSOutput:
     """
 
     grid: Grid
-    path: Union[str, Path]
+    """Object representing the grid information."""
+    path: str | Path
+    """Filename, or list of filenames with model output."""
     use_dask: bool = False
-    model_reference_date: Optional[datetime] = None
-    adjust_depth_for_sea_surface_height: Optional[bool] = False
+    """Whether to use dask for processing."""
+    model_reference_date: datetime | None = None
+    """Reference date of ROMS simulation."""
+    adjust_depth_for_sea_surface_height: bool | None = False
+    """Whether to account for sea surface height variations when computing depth
+    coordinates."""
+
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing the ROMS output."""
 
     def __post_init__(self):
-
         ds = self._load_model_output()
         self._infer_model_reference_date_from_metadata(ds)
         self._check_vertical_coordinate(ds)
@@ -65,18 +73,18 @@ class ROMSOutput:
 
     def plot(
         self,
-        var_name,
-        time=0,
-        s=None,
-        eta=None,
-        xi=None,
-        depth=None,
-        lat=None,
-        lon=None,
-        include_boundary=False,
-        depth_contours=False,
-        ax=None,
-        save_path=None,
+        var_name: str,
+        time: int = 0,
+        s: int | None = None,
+        eta: int | None = None,
+        xi: int | None = None,
+        depth: float | None = None,
+        lat: float | None = None,
+        lon: float | None = None,
+        include_boundary: bool = False,
+        depth_contours: bool = False,
+        ax: Axes | None = None,
+        save_path: str | None = None,
     ) -> None:
         """Generate a plot of a ROMS output field for a specified vertical or horizontal
         slice.
@@ -136,7 +144,7 @@ class ROMSOutput:
             By default, depth contours are not shown (False).
 
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots. Default is None.
+            The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots. Default is None.
 
         save_path : str, optional
             Path to save the generated plot. If None, the plot is shown interactively.
@@ -181,56 +189,6 @@ class ROMSOutput:
                     f"Invalid input: The field does not have a 'time' dimension, "
                     f"but a time index ({time}) greater than 0 was provided."
                 )
-
-        # Input checks
-        _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary)
-
-        # Get horizontal dimensions and grid location
-        horizontal_dims_dict = {
-            "rho": {"eta": "eta_rho", "xi": "xi_rho"},
-            "u": {"eta": "eta_rho", "xi": "xi_u"},
-            "v": {"eta": "eta_v", "xi": "xi_rho"},
-        }
-        for loc, horizontal_dims in horizontal_dims_dict.items():
-            if all(dim in field.dims for dim in horizontal_dims.values()):
-                break
-
-        # Convert relative to absolute indices
-        def _get_absolute_index(idx, field, dim_name):
-            index = field[dim_name].isel(**{dim_name: idx}).item()
-            return index
-
-        if eta is not None and eta < 0:
-            eta = _get_absolute_index(eta, field, horizontal_dims["eta"])
-        if xi is not None and xi < 0:
-            xi = _get_absolute_index(xi, field, horizontal_dims["xi"])
-        if s is not None and s < 0:
-            s = _get_absolute_index(s, field, "s_rho")
-
-        # Set spatial coordinates
-        lat_deg = self.grid.ds[f"lat_{loc}"]
-        lon_deg = self.grid.ds[f"lon_{loc}"]
-        if self.grid.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-        # Mask the field
-        mask = self.grid.ds[f"mask_{loc}"]
-        field = field.where(mask)
-
-        # Assign eta and xi as coordinates
-        coords_to_assign = {dim: field[dim] for dim in horizontal_dims.values()}
-        field = field.assign_coords(**coords_to_assign)
-
-        # Remove horizontal boundary if desired
-        slice_dict = {
-            "rho": {"eta_rho": slice(1, -1), "xi_rho": slice(1, -1)},
-            "u": {"eta_rho": slice(1, -1), "xi_u": slice(1, -1)},
-            "v": {"eta_v": slice(1, -1), "xi_rho": slice(1, -1)},
-        }
-        if not include_boundary:
-            field = field.isel(**slice_dict[loc])
-
         # Load the data
         if self.use_dask:
             from dask.diagnostics import ProgressBar
@@ -238,162 +196,36 @@ class ROMSOutput:
             with ProgressBar():
                 field.load()
 
-        # Compute layer depth for 3D fields when depth contours are requested or no vertical layer is specified.
-        compute_layer_depth = len(field.dims) > 2 and (depth_contours or s is None)
-        if compute_layer_depth:
-            if eta is not None or xi is not None:
-                # Computing depth coordinates directly for the slice in question is more efficient
-                # than using .ds_depth_coords, which computes depth coordinates for full field
-                if self.adjust_depth_for_sea_surface_height:
-                    zeta = self.ds.zeta.isel(time=time)
-                else:
-                    zeta = 0
-                if compute_layer_depth:
-                    layer_depth = compute_depth_coordinates(
-                        self.grid.ds,
-                        zeta,
-                        depth_type="layer",
-                        location=loc,
-                        eta=eta,
-                        xi=xi,
-                    )
-            else:
-                self._get_depth_coordinates(depth_type="layer", locations=[loc])
-                layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"]
-                if self.adjust_depth_for_sea_surface_height:
-                    layer_depth = layer_depth.isel(time=time)
-
-            if not include_boundary:
-                # Apply valid slices only for dimensions that exist in layer_depth.dims
-                layer_depth = layer_depth.isel(
-                    **{
-                        dim: s
-                        for dim, s in slice_dict.get(loc, {}).items()
-                        if dim in layer_depth.dims
-                    }
-                )
-            layer_depth.load()
-
-        # Prepare figure title
-        formatted_time = np.datetime_as_string(field.abs_time.values, unit="m")
-        title = f"time: {formatted_time}"
-
-        # Slice the field horizontally as desired
-        def _slice_along_dimension(field, title, dim_name, idx):
-            field = field.sel(**{dim_name: idx})
-            title = title + f", {dim_name} = {idx}"
-            return field, title
-
-        if eta is not None:
-            field, title = _slice_along_dimension(
-                field, title, horizontal_dims["eta"], eta
-            )
-        if xi is not None:
-            field, title = _slice_along_dimension(
-                field, title, horizontal_dims["xi"], xi
-            )
-        if s is not None:
-            field, title = _slice_along_dimension(field, title, "s_rho", s)
-            if compute_layer_depth:
-                layer_depth = layer_depth.isel(s_rho=s)
+        if self.adjust_depth_for_sea_surface_height:
+            zeta = self.ds.zeta.isel(time=time)
         else:
-            depth_contours = False
-
-        # Regrid laterally
-        if lat is not None or lon is not None:
-
-            if lat is not None:
-                lats = [lat]
-                title = title + f", lat = {lat}°N"
-            else:
-                resolution = self.grid._infer_nominal_horizontal_resolution()
-                lats = _generate_coordinate_range(
-                    field.lat.min().values, field.lat.max().values, resolution
-                )
-            lats = xr.DataArray(lats, dims=["lat"], attrs={"units": "°N"})
-
-            if lon is not None:
-                lons = [lon]
-                title = title + f", lon = {lon}°E"
-            else:
-                resolution = self.grid._infer_nominal_horizontal_resolution(lat)
-                lons = _generate_coordinate_range(
-                    field.lon.min().values, field.lon.max().values, resolution
-                )
-            lons = xr.DataArray(lons, dims=["lon"], attrs={"units": "°E"})
-
-            target_coords = {"lat": lats, "lon": lons}
-            lateral_regrid = LateralRegridFromROMS(field, target_coords)
-            field = lateral_regrid.apply(field).squeeze()
-            if compute_layer_depth:
-                layer_depth = lateral_regrid.apply(layer_depth).squeeze()
-
-        # Assign depth as coordinate
-        if compute_layer_depth:
-            field = field.assign_coords({"layer_depth": layer_depth})
-
-        if lat is not None:
-            field, layer_depth = _remove_edge_nans(
-                field, "lon", layer_depth if "layer_depth" in locals() else None
-            )
-        if lon is not None:
-            field, layer_depth = _remove_edge_nans(
-                field, "lat", layer_depth if "layer_depth" in locals() else None
-            )
-
-        # Regrid vertically
-        if depth is not None:
-            vertical_regrid = VerticalRegridFromROMS(self.ds)
-            # Save attributes before vertical regridding
-            attrs = field.attrs
-            field = vertical_regrid.apply(
-                field, layer_depth, np.array([depth])
-            ).squeeze()
-            # Reset attributes
-            field.attrs = attrs
-            title = title + f", depth = {depth}m"
+            zeta = 0
 
         # Choose colorbar
         if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
-            vmax = max(field.max().values, -field.min().values)
-            vmin = -vmax
-            cmap = plt.colormaps.get_cmap("RdBu_r")
-        else:
-            vmax = field.max().values
-            vmin = field.min().values
-            if var_name in ["temp", "salt"]:
-                cmap = plt.colormaps.get_cmap("YlOrRd")
-            else:
-                cmap = plt.colormaps.get_cmap("YlGn")
-        cmap.set_bad(color="gray")
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-        # Plotting
-        if (eta is None and xi is None) and (lat is None and lon is None):
-            fig = _plot(
-                field=field,
-                depth_contours=depth_contours,
-                title=title,
-                kwargs=kwargs,
-                c=None,
-            )
+            cmap_name = "RdBu_r"
+        elif var_name in ["temp", "salt"]:
+            cmap_name = "YlOrRd"
         else:
-            if len(field.dims) == 2:
-                fig = _section_plot(
-                    field,
-                    interface_depth=None,
-                    title=title,
-                    kwargs=kwargs,
-                    ax=ax,
-                )
-            else:
-                if "s_rho" in field.dims:
-                    fig = _profile_plot(field, title=title, ax=ax)
-                else:
-                    fig = _line_plot(field, title=title, ax=ax)
-
-        if save_path:
-            plt.savefig(save_path, dpi=300, bbox_inches="tight")
+            cmap_name = "YlGn"
+
+        plot(
+            field=field,
+            grid_ds=self.grid.ds,
+            zeta=zeta,
+            s=s,
+            eta=eta,
+            xi=xi,
+            depth=depth,
+            lat=lat,
+            lon=lon,
+            include_boundary=include_boundary,
+            depth_contours=depth_contours,
+            layer_contours=False,
+            ax=ax,
+            save_path=save_path,
+            cmap_name=cmap_name,
+        )
 
     def regrid(self, var_names=None, horizontal_resolution=None, depth_levels=None):
         """Regrid the dataset both horizontally and vertically.
@@ -416,7 +248,6 @@ class ROMSOutput:
         xarray.Dataset
             The regridded dataset.
         """
-
         if var_names is None:
             var_names = list(self.ds.data_vars)
 
@@ -437,7 +268,7 @@ class ROMSOutput:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
 
         if horizontal_resolution is None:
-            horizontal_resolution = self.grid._infer_nominal_horizontal_resolution()
+            horizontal_resolution = infer_nominal_horizontal_resolution(self.grid.ds)
         lons = _generate_coordinate_range(
             lon_deg.min().values, lon_deg.max().values, horizontal_resolution
         )
@@ -450,7 +281,7 @@ class ROMSOutput:
 
         # Prepare vertical regrid
         if depth_levels is None:
-            depth_levels, _ = self._compute_exponential_depth_levels()
+            depth_levels, _ = self.grid._compute_exponential_depth_levels()
 
         # Ensure depth_levels is an xarray.DataArray
         if not isinstance(depth_levels, xr.DataArray):
@@ -568,7 +399,6 @@ class ROMSOutput:
         - If `adjust_depth_for_sea_surface_height` is `True`, the method accounts for variations
           in sea surface height (`zeta`).
         """
-
         if self.adjust_depth_for_sea_surface_height:
             if "zeta" not in self.ds:
                 raise ValueError(
@@ -587,7 +417,6 @@ class ROMSOutput:
 
     def _load_model_output(self) -> xr.Dataset:
         """Load the model output."""
-
         # Load the dataset
         ds = _load_data(
             self.path,
@@ -688,7 +517,6 @@ class ROMSOutput:
         - `theta_s`, `theta_b`, and `hc` are checked for exact equality using `np.array_equal`.
         - `Cs_r` and `Cs_w` are checked for numerical closeness using `np.allclose`.
         """
-
         # Check exact equality for theta_s, theta_b, and hc
         if not np.array_equal(self.grid.theta_s, ds.attrs["theta_s"]):
             raise ValueError(
@@ -786,52 +614,3 @@ class ROMSOutput:
         # Add all necessary coordinates in one go
         ds = ds.assign_coords(coords_to_add)
         return ds
-
-    def _compute_exponential_depth_levels(self, Nz=None, depth=None, h=None):
-        """Compute vertical grid center and face depths using an exponential profile.
-
-        Parameters
-        ----------
-        Nz : int, optional
-            Number of vertical levels. Defaults to `len(self.ds.s_rho)`.
-
-        depth : float, optional
-            Total depth of the domain. Defaults to `grid.ds.h.max().values`.
-
-        h : float, optional
-            Scaling parameter for the exponential profile. Defaults to `Nz / 4.5`.
-
-        Returns
-        -------
-        tuple of numpy.ndarray
-            Depth values at the vertical grid centers (`z_centers`) and grid faces (`z_faces`),
-            both rounded to two decimal places.
-        """
-        if Nz is None:
-            Nz = len(self.ds.s_rho)
-        if depth is None:
-            depth = self.grid.ds.h.max().values
-        if h is None:
-            h = Nz / 4.5
-
-        k = np.arange(1, Nz + 2)
-
-        # Define the exponential profile function
-        def exponential_profile(k, Nz, h):
-            return np.exp(k / h)
-
-        z_faces = np.vectorize(exponential_profile)(k, Nz, h)
-
-        # Normalize
-        z_faces -= z_faces[0]
-        z_faces *= depth / z_faces[-1]
-        z_faces[0] = 0.0
-
-        # Calculate center depths (average between adjacent face depths)
-        z_centers = (z_faces[:-1] + z_faces[1:]) / 2
-
-        # Round both z_faces and z_centers to two decimal places
-        z_faces = np.round(z_faces, 2)
-        z_centers = np.round(z_centers, 2)
-
-        return z_centers, z_faces

roms_tools/constants.py

@@ -1,2 +1,5 @@
 R_EARTH = 6371315.0  # Earth radius in meters
+MAXIMUM_GRID_SIZE = 25000  # in km
+UPPER_BOUND_THETA_S = 10  # upper bound for surface vertical stretching parameter
+UPPER_BOUND_THETA_B = 10  # upper bound for bottom vertical stretching parameter
 NUM_TRACERS = 34  # Number of tracers (temperature, salinity, BGC tracers)

roms_tools/download.py

@@ -79,7 +79,7 @@ pup_test_data = pooch.create(
         "CESM_BGC_SURFACE_2012.nc": "3c4d156adca97909d0fac36bf50b99583ab37d8020d7a3e8511e92abf2331b38",
         "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
         "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
-        "coarsened_UNIFIED_bgc_dataset.nc": "58b510dfd9125e8bfd92342af8cbf7c2639ef9a42c21f3ee74fc31b361e3750c",
+        "coarsened_UNIFIED_bgc_dataset.nc": "c5eaeca6f7640f2358afe5c46ace99ff4879948df8a778a2f21967de034f8d0a",
         "grid_created_with_matlab.nc": "fd537ef8159fabb18e38495ec8d44e2fa1b7fb615fcb1417dd4c0e1bb5f4e41d",
         "etopo5_coarsened_and_shifted.nc": "9a5cb4b38c779d22ddb0ad069b298b9722db34ca85a89273eccca691e89e6f96",
         "srtm15_coarsened.nc": "48bc8f4beecfdca9c192b13f4cbeef1455f49d8261a82563aaec5757e100dff9",
@@ -123,7 +123,6 @@ def download_river_data(filename: str) -> str:
     str
         The path to the downloaded test data file.
     """
-
     # Fetch the file using Pooch, downloading if necessary
     fname = river_data.fetch(filename)