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

roms_tools/utils.py

@@ -1,10 +1,14 @@
-import xarray as xr
-import numpy as np
-from pathlib import Path
-import re
 import glob
 import logging
+import re
 import warnings
+from importlib.util import find_spec
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
+from roms_tools.constants import R_EARTH
 
 
 def _load_data(
@@ -14,6 +18,7 @@ def _load_data(
     time_chunking=True,
     decode_times=True,
     force_combine_nested=False,
+    read_zarr: bool = False,
 ):
     """Load dataset from the specified file.
 
@@ -33,11 +38,14 @@ def _load_data(
         If False, the data will not be chunked explicitly along the time dimension, but will follow the default auto chunking scheme. This option is useful for ROMS restart files.
         Defaults to True.
     decode_times: bool, optional
-        If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
+        If True, decode times and timedeltas encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
         Defaults to True.
     force_combine_nested: bool, optional
         If True, forces the use of nested combination (`combine_nested`) regardless of whether wildcards are used.
         Defaults to False.
+    read_zarr: bool, optional
+        If True, use the zarr engine to read the dataset, and don't use mfdataset.
+        Defaults to False.
 
     Returns
     -------
@@ -54,11 +62,25 @@ def _load_data(
     if dim_names is None:
         dim_names = {}
 
+    if use_dask:
+        if not _has_dask():
+            raise RuntimeError(
+                "Dask is required but not installed. Install it with:\n"
+                "  • `pip install roms-tools[dask]` or\n"
+                "  • `conda install dask`\n"
+                "Alternatively, install `roms-tools` with conda to include all dependencies."
+            )
+    if read_zarr:
+        if isinstance(filename, list):
+            raise ValueError("read_zarr requires a single path, not a list of paths")
+        if not use_dask:
+            raise ValueError("read_zarr must be used with use_dask")
+
     # Precompile the regex for matching wildcard characters
     wildcard_regex = re.compile(r"[\*\?\[\]]")
 
     # Convert Path objects to strings
-    if isinstance(filename, (str, Path)):
+    if isinstance(filename, str | Path):
         filename_str = str(filename)
     elif isinstance(filename, list):
         filename_str = [str(f) for f in filename]
@@ -119,7 +141,6 @@ def _load_data(
     }
 
     if use_dask:
-
         if "latitude" in dim_names and "longitude" in dim_names:
             # for lat-lon datasets
             chunks = {
@@ -140,27 +161,43 @@ def _load_data(
             chunks[dim_names["depth"]] = -1
         if "time" in dim_names and time_chunking:
             chunks[dim_names["time"]] = 1
+        if "ntides" in dim_names:
+            chunks[dim_names["ntides"]] = 1
 
         with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=UserWarning, module="xarray")
-
-            ds = xr.open_mfdataset(
-                matching_files,
-                decode_times=decode_times,
-                chunks=chunks,
-                **combine_kwargs,
-                **kwargs,
+            warnings.filterwarnings(
+                "ignore",
+                category=UserWarning,
+                message=r"^The specified chunks separate.*",
             )
 
-        # Rechunk the dataset along the tidal constituent dimension ("ntides") after loading
-        # because the original dataset does not have a chunk size of 1 along this dimension.
-        if "ntides" in dim_names and "ntides" in ds.dims:
-            ds = ds.chunk({dim_names["ntides"]: 1})
+            if read_zarr:
+                ds = xr.open_zarr(
+                    matching_files[0],
+                    decode_times=decode_times,
+                    chunks=chunks,
+                    consolidated=None,
+                    storage_options=dict(token="anon"),
+                )
+            else:
+                ds = xr.open_mfdataset(
+                    matching_files,
+                    decode_times=decode_times,
+                    decode_timedelta=decode_times,
+                    chunks=chunks,
+                    **combine_kwargs,
+                    **kwargs,
+                )
 
     else:
         ds_list = []
         for file in matching_files:
-            ds = xr.open_dataset(file, decode_times=decode_times, chunks=None)
+            ds = xr.open_dataset(
+                file,
+                decode_times=decode_times,
+                decode_timedelta=decode_times,
+                chunks=None,
+            )
             ds_list.append(ds)
 
         if kwargs["combine"] == "by_coords":
@@ -173,6 +210,9 @@ def _load_data(
     if "time" in dim_names and dim_names["time"] not in ds.dims:
         ds = ds.expand_dims(dim_names["time"])
 
+    if "time" in dim_names and not read_zarr:
+        ds = ds.drop_duplicates(dim=dim_names["time"])
+
     return ds
 
 
@@ -202,7 +242,6 @@ def interpolate_from_rho_to_u(field, method="additive"):
     field_interpolated : xr.DataArray
         The interpolated data array on the u grid with the dimension "xi_u".
     """
-
     if method == "additive":
         field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
             xi_rho=slice(1, None)
@@ -248,7 +287,6 @@ def interpolate_from_rho_to_v(field, method="additive"):
     field_interpolated : xr.DataArray
         The interpolated data array on the v grid with the dimension "eta_v".
     """
-
     if method == "additive":
         field_interpolated = 0.5 * (field + field.shift(eta_rho=1)).isel(
             eta_rho=slice(1, None)
@@ -286,7 +324,6 @@ def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
         The DataArray with dimensions reordered so that 'time', 's_*', 'eta_*',
         and 'xi_*' are first, in that order, if they exist.
     """
-
     # List of preferred dimension patterns
     preferred_order = ["time", "s_", "eta_", "xi_"]
 
@@ -329,7 +366,6 @@ def save_datasets(dataset_list, output_filenames, use_dask=False, verbose=True):
     List[Path]
         A list of Path objects for the filenames that were saved.
     """
-
     saved_filenames = []
 
     output_filenames = [f"{filename}.nc" for filename in output_filenames]
@@ -374,21 +410,21 @@ def get_dask_chunks(location, chunk_size):
     return chunk_mapping.get(location, {})
 
 
-def _generate_coordinate_range(min, max, resolution):
+def _generate_coordinate_range(min_val: float, max_val: float, resolution: float):
     """Generate an array of target coordinates (e.g., latitude or longitude) within a
     specified range, with a resolution that is rounded to the nearest value of the form
     `1/n` (or integer).
 
-    This method generates an array of target coordinates between the provided `min` and `max`
-    values, ensuring that both `min` and `max` are included in the resulting range. The resolution
+    This method generates an array of target coordinates between the provided `min_val` and `max_val`
+    values, ensuring that both `min_val` and `max_val` are included in the resulting range. The resolution
     is rounded to the nearest fraction of the form `1/n` or an integer, based on the input.
 
     Parameters
     ----------
-    min : float
+    min_val : float
         The minimum value (in degrees) of the coordinate range (inclusive).
 
-    max : float
+    max_val : float
         The maximum value (in degrees) of the coordinate range (inclusive).
 
     resolution : float
@@ -402,7 +438,6 @@ def _generate_coordinate_range(min, max, resolution):
         An array of target coordinates generated from the specified range, with the resolution
         rounded to a suitable fraction (e.g., `1/n`) or integer, depending on the input resolution.
     """
-
     # Find the closest fraction of the form 1/n or integer to match the resolution
     resolution_rounded = None
     min_diff = float("inf")  # Initialize the minimum difference as infinity
@@ -425,8 +460,8 @@ def _generate_coordinate_range(min, max, resolution):
             resolution_rounded = fraction  # Update the best fraction (or integer) found
 
     # Adjust the start and end of the range to include integer values
-    start_int = np.floor(min)  # Round the minimum value down to the nearest integer
-    end_int = np.ceil(max)  # Round the maximum value up to the nearest integer
+    start_int = np.floor(min_val)  # Round the minimum value down to the nearest integer
+    end_int = np.ceil(max_val)  # Round the maximum value up to the nearest integer
 
     # Generate the array of target coordinates, including both the min and max values
     target = np.arange(start_int, end_int + resolution_rounded, resolution_rounded)
@@ -437,47 +472,159 @@ def _generate_coordinate_range(min, max, resolution):
     return target.astype(np.float32)
 
 
-def _remove_edge_nans(field, xdim, layer_depth=None):
-    """Trim NaN-only edges along a specified dimension.
+def _generate_focused_coordinate_range(
+    center: float,
+    sc: float,
+    min_val: float,
+    max_val: float,
+    N: int,
+    stretch_factor: float = 3.0,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate an array of target coordinates within [min_val, max_val] with two.
+
+    resolution zones: a higher resolution region centered near `center` ± stretch_factor
+    * `sc`, and coarser resolution outside.
+
+    Parameters
+    ----------
+    min_val : float
+        Minimum coordinate value (inclusive).
+        Units: degrees or meters.
+    max_val : float
+        Maximum coordinate value (inclusive).
+        Units: degrees or meters.
+    center : float
+        Coordinate around which to increase resolution.
+        Units: degrees or meters.
+    sc : float
+        Width controlling the size of the high-resolution region.
+        If 0, treated as a small default to avoid zero-width window.
+        Units: degrees or meters.
+    N : int
+        Number of desired coordinate points.
+    stretch_factor : float, default 3.0
+        Multiplier for `sc` to define the width of the high-resolution window.
+
+    Returns
+    -------
+    tuple of np.ndarray
+        Tuple of two arrays:
+        - centers: cell centers.
+        - faces: cell faces.
+
+    Notes
+    -----
+    - About 20% of the faces are placed in the high-resolution region.
+    - Remaining faces are split evenly in the coarse resolution regions
+    - Faces and centers are returned sorted and unique.
+    """
+    # Define the bounds of the high-res region
+    if sc == 0.0:
+        sc = 1
+    low_bound = max(center - stretch_factor * sc, min_val)
+    high_bound = min(center + stretch_factor * sc, max_val)
+
+    # Decide how many points for high and low resolution
+    high_res_points = int(N * 0.2)  # 20% of the points in high-res region
+    low_res_points = N + 1 - high_res_points * 2  # rest split
+
+    # First coarse region
+    first_coarse_res_faces = (
+        np.linspace(min_val, low_bound, low_res_points // 2, endpoint=False)
+        if low_bound > 0
+        else np.array([])
+    )
+    # High-res region faces (low_bound to high_bound)
+    high_res_faces = np.linspace(low_bound, high_bound, high_res_points, endpoint=False)
+    # Second coarse region
+    second_coarse_res_faces = (
+        np.linspace(
+            high_bound,
+            max_val,
+            low_res_points - len(first_coarse_res_faces),
+            endpoint=True,
+        )
+        if high_bound < max_val
+        else np.array([])
+    )
+
+    faces = np.unique(
+        np.concatenate(
+            [first_coarse_res_faces, high_res_faces, second_coarse_res_faces]
+        )
+    )
+    centers = (faces[:-1] + faces[1:]) / 2
+
+    return centers, faces
+
+
+def _remove_edge_nans(
+    field: xr.DataArray, xdim: str, layer_depth: xr.DataArray | None = None
+) -> tuple[xr.DataArray, xr.DataArray | None]:
+    """Remove NaN-only slices at the edges of a specified dimension.
 
-    Useful when a ROMS grid has been regridded to a fixed lat/lon section,
-    leaving NaN-filled edges (e.g., over land or outside the domain).
-    Removes leading/trailing slices along `xdim` where all values are NaN,
-    based on `field` or optionally `layer_depth`.
+    This function trims leading and trailing slices along the specified `xdim` where all values
+    are NaN. It assumes that the data has only one other relevant dimension (typically depth).
+    If `layer_depth` is provided, it is used instead of `field` to determine which slices are NaN-only.
 
     Parameters
     ----------
     field : xr.DataArray
-        Data to trim.
+        Input array to trim. May contain NaNs at the edges along `xdim`.
     xdim : str
-        Dimension along which to remove NaN edges.
+        The dimension along which to remove NaN-only slices.
     layer_depth : xr.DataArray, optional
-        Optional field to determine where NaNs occur.
+        Optional array to evaluate NaN positions. If not provided, `field` is used instead.
 
     Returns
     -------
     field : xr.DataArray
-        Trimmed data.
+        Trimmed `field`, with leading and trailing NaN-only slices removed along `xdim`.
     layer_depth : xr.DataArray or None
-        Trimmed `layer_depth` if provided.
+        Trimmed `layer_depth`, if provided. Otherwise, returns None.
+
+    Raises
+    ------
+    ValueError
+        If `field` has more than one additional dimension besides `xdim`.
+
+    Notes
+    -----
+    - If `xdim` is not in `field.dims`, no trimming is performed.
+    - This is typically used for visualizing or extracting clean sections from 2D slices
+      (e.g., vertical sections) that have NaNs at the spatial boundaries.
     """
     if xdim in field.dims:
-        if layer_depth is not None:
-            nan_mask = layer_depth.isnull().sum(
-                dim=[dim for dim in layer_depth.dims if dim != xdim]
-            )
+        other_dims = [dim for dim in field.dims if dim != xdim]
+        if len(other_dims) == 0:
+            if layer_depth is not None:
+                nan_mask = layer_depth.isnull()
+            else:
+                nan_mask = field.isnull()
+
+        elif len(other_dims) == 1:
+            depth_dim = other_dims[0]
+
+            if layer_depth is not None:
+                if depth_dim in layer_depth.dims:
+                    nan_mask = layer_depth.isnull().sum(dim=depth_dim)
+                else:
+                    nan_mask = layer_depth.isnull()
+            else:
+                N = field.sizes[depth_dim]
+                nan_mask = field.isnull().sum(dim=depth_dim) == N
+
         else:
-            nan_mask = field.isnull().sum(
-                dim=[dim for dim in field.dims if dim != xdim]
+            raise ValueError(
+                f"Cannot trim along dimension '{xdim}': expected at most one other dimension, "
+                f"but got {len(other_dims)} ({other_dims})."
             )
 
-        # Find the valid indices where the sum of the nans is 0
         valid_indices = np.where(nan_mask.values == 0)[0]
 
         if len(valid_indices) > 0:
             first_valid = valid_indices[0]
             last_valid = valid_indices[-1]
-
             field = field.isel({xdim: slice(first_valid, last_valid + 1)})
             if layer_depth is not None:
                 layer_depth = layer_depth.isel(
@@ -485,3 +632,75 @@ def _remove_edge_nans(field, xdim, layer_depth=None):
                 )
 
     return field, layer_depth
+
+
+def _has_dask() -> bool:
+    return find_spec("dask") is not None
+
+
+def _has_gcsfs() -> bool:
+    return find_spec("gcsfs") is not None
+
+
+def normalize_longitude(lon: float, straddle: bool) -> float:
+    """Normalize longitude to the appropriate range depending on whether the grid
+    straddles the dateline.
+
+    Parameters
+    ----------
+    lon : float
+        Longitude in degrees (can be any value, including multiples of 360 or negative values).
+    straddle : bool
+        Whether the grid straddles the dateline. If True, output will be in (-180, 180];
+        if False, output will be in [0, 360).
+    Returns
+    -------
+    float
+        Normalized longitude.
+    """
+    lon = lon % 360
+    if straddle:
+        return lon - 360 if lon > 180 else lon
+    else:
+        return lon + 360 if lon < 0 else lon
+
+
+def infer_nominal_horizontal_resolution(
+    grid_ds: xr.Dataset, lat: float | None = None
+) -> float:
+    """Estimate the nominal horizontal resolution of the ROMS grid in degrees at a
+    specified latitude.
+
+    This function calculates the average horizontal grid spacing (based on `pm` and `pn`),
+    then converts that spacing from meters to degrees longitude, accounting for Earth's curvature
+    and the cosine of the latitude.
+
+    Parameters
+    ----------
+    grid_ds : xr.Dataset
+        ROMS grid dataset containing `pm`, `pn`, and `lat_rho` coordinates.
+    lat : float, optional
+        Latitude (in degrees) at which to estimate resolution. If None, the average latitude
+        over the grid is used.
+
+    Returns
+    -------
+    float
+        Estimated nominal horizontal resolution in degrees longitude.
+    """
+    if lat is None:
+        # Use center of the domain in latitude
+        lat = float((grid_ds.lat_rho.max() + grid_ds.lat_rho.min()) / 2)
+
+    lat_rad = np.deg2rad(lat)
+
+    # Mean grid spacing in meters (pm and pn are in 1/m)
+    resolution_in_m = ((1 / grid_ds.pm).mean() + (1 / grid_ds.pn).mean()) / 2
+
+    # Meters per degree at the equator
+    meters_per_degree = 2 * np.pi * R_EARTH / 360
+
+    # Adjust resolution to degrees longitude at the specified latitude
+    resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
+
+    return float(resolution_in_degrees)

roms_tools/vertical_coordinate.py

@@ -1,9 +1,11 @@
 import numpy as np
 import xarray as xr
+
+from roms_tools.constants import UPPER_BOUND_THETA_B, UPPER_BOUND_THETA_S
 from roms_tools.utils import (
-    transpose_dimensions,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    transpose_dimensions,
 )
 
 
@@ -30,10 +32,10 @@ def compute_cs(sigma, theta_s, theta_b):
     ValueError
         If theta_s or theta_b are not within the valid range.
     """
-    if not (0 < theta_s <= 10):
-        raise ValueError("theta_s must be between 0 and 10.")
-    if not (0 < theta_b <= 4):
-        raise ValueError("theta_b must be between 0 and 4.")
+    if not (0 < theta_s <= UPPER_BOUND_THETA_S):
+        raise ValueError(f"theta_s must be between 0 and {UPPER_BOUND_THETA_S}.")
+    if not (0 < theta_b <= UPPER_BOUND_THETA_B):
+        raise ValueError(f"theta_b must be between 0 and {UPPER_BOUND_THETA_B}.")
 
     C = (1 - np.cosh(theta_s * sigma)) / (np.cosh(theta_s) - 1)
     C = (np.exp(theta_b * C) - 1) / (1 - np.exp(-theta_b))
@@ -104,7 +106,6 @@ def compute_depth(zeta, h, hc, cs, sigma):
     z : xr.DataArray
         The depth at different sigma levels.
     """
-
     z = (hc * sigma + h * cs) / (hc + h)
     z = zeta + (zeta + h) * z
 
@@ -118,8 +119,8 @@ def compute_depth_coordinates(
     zeta: xr.DataArray | float = 0,
     depth_type: str = "layer",
     location: str = "rho",
-    eta: int = None,
-    xi: int = None,
+    eta: int | None = None,
+    xi: int | None = None,
 ) -> "xr.DataArray":
     """Compute vertical depth coordinates for a given ROMS grid location.
 

roms_tools-3.1.0.dist-info/METADATA

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: roms-tools
+Version: 3.1.0
+Summary: Tools for running and analysing UCLA-ROMS simulations
+Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>, Scott Eilerman <scott.eilerman@cworthy.org>
+License: Apache-2
+Project-URL: Home, https://github.com/CWorthy-ocean/roms-tools
+Project-URL: Documentation, https://roms-tools.readthedocs.io/en/latest/
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: xarray>=2022.6.0
+Requires-Dist: numpy<2.3,>2.0
+Requires-Dist: pooch
+Requires-Dist: matplotlib
+Requires-Dist: cartopy
+Requires-Dist: packaging
+Requires-Dist: scipy
+Requires-Dist: gcm_filters>=0.5.1
+Requires-Dist: netcdf4
+Requires-Dist: PyYAML
+Requires-Dist: pyamg
+Requires-Dist: bottleneck
+Requires-Dist: regionmask
+Requires-Dist: xgcm
+Requires-Dist: numba>=0.61.2
+Requires-Dist: pydantic<3,>2
+Provides-Extra: dask
+Requires-Dist: dask[diagnostics]; extra == "dask"
+Provides-Extra: stream
+Requires-Dist: dask[diagnostics]; extra == "stream"
+Requires-Dist: gcsfs; extra == "stream"
+Dynamic: license-file
+
+# ROMS-Tools
+[![Conda version](https://img.shields.io/conda/vn/conda-forge/roms-tools.svg)](https://anaconda.org/conda-forge/roms-tools)
+[![PyPI version](https://img.shields.io/pypi/v/roms-tools.svg)](https://pypi.org/project/roms-tools/)
+[![codecov](https://codecov.io/gh/CWorthy-ocean/roms-tools/graph/badge.svg?token=5S1oNu39xE)](https://codecov.io/gh/CWorthy-ocean/roms-tools)
+[![Documentation Status](https://readthedocs.org/projects/roms-tools/badge/?version=latest)](https://roms-tools.readthedocs.io/en/latest/?badge=latest)
+![Run Tests](https://github.com/CWorthy-ocean/roms-tools/actions/workflows/tests.yaml/badge.svg)
+![Supported Python Versions](https://img.shields.io/pypi/pyversions/roms-tools)
+
+## Overview
+
+A suite of Python tools for setting up and analyzing a [UCLA-ROMS](https://github.com/CESR-lab/ucla-roms) simulation with or without [MARBL biogeochemistry](https://marbl-ecosys.github.io/versions/latest_release/index.html).
+
+## Installation
+
+### ⚡️ **Installation from Conda-Forge**
+
+To install `ROMS-Tools` with all dependencies, including `xesmf`, `dask` and all packages required for streaming source data directly from the cloud, use:
+
+```bash
+conda install -c conda-forge roms-tools
+```
+
+> [!Note]
+>  Installation from Conda-Forge is the recommended installation method to ensure all features of `ROMS-Tools` are available.
+
+### 📦 **Installation from PyPI (pip)**
+
+You can also install `ROMS-Tools` from `pip`:
+
+```bash
+pip install roms-tools
+```
+
+If you want to use `ROMS-Tools` with `dask` (recommended for parallel and out-of-core computation), install it with the additional dependency:
+
+```bash
+pip install roms-tools[dask]
+```
+
+If you want to use `ROMS-Tools` with `dask` and all packages required for streaming source data directly from the cloud, install it with the additional dependencies:
+
+```bash
+pip install roms-tools[stream]
+```
+
+
+> [!Note]
+>  The PyPI versions of `ROMS-Tools` do not include `xesmf`, so some features will be unavailable.
+
+
+### Installation from GitHub
+
+To obtain the latest development version, first clone the source repository:
+
+```bash
+git clone https://github.com/CWorthy-ocean/roms-tools.git
+cd roms-tools
+```
+
+Next, install and activate the following conda environment:
+
+```bash
+conda env create -f ci/environment-with-xesmf.yml
+conda activate romstools-test
+```
+
+Finally, install `ROMS-Tools` in the same environment:
+
+```bash
+pip install -e .
+```
+
+If you want to use `ROMS-Tools` with `dask` (recommended for parallel and out-of-core computation), you can
+install `ROMS-Tools` along with the additional dependency via:
+
+```bash
+pip install -e ".[dask]"
+```
+
+If you want to use `ROMS-Tools` with `dask` and all packages required for streaming source data directly from the cloud, you can
+install `ROMS-Tools` along with the additional dependencies via:
+
+```bash
+pip install -e ".[stream]"
+```
+
+## Getting Started
+
+To learn how to use `ROMS-Tools`, check out the [documentation](https://roms-tools.readthedocs.io/en/latest/).
+
+## Feedback and contributions
+
+**If you find a bug, have a feature suggestion, or any other kind of feedback, please start a Discussion.**
+
+We also accept contributions in the form of Pull Requests.
+
+## See also
+
+- [ROMS source code](https://github.com/CESR-lab/ucla-roms)
+- [C-Star](https://github.com/CWorthy-ocean/C-Star)