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

roms_tools/utils.py

@@ -1,9 +1,12 @@
-import xarray as xr
-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
 
 
 def _load_data(
@@ -13,6 +16,7 @@ def _load_data(
     time_chunking=True,
     decode_times=True,
     force_combine_nested=False,
+    read_zarr: bool = False,
 ):
     """Load dataset from the specified file.
 
@@ -32,11 +36,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
     -------
@@ -53,6 +60,20 @@ 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"[\*\?\[\]]")
 
@@ -139,27 +160,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:
-            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":
@@ -172,6 +209,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
 
 
@@ -371,3 +411,261 @@ def get_dask_chunks(location, chunk_size):
         "v": {"eta_v": chunk_size, "xi_rho": chunk_size},
     }
     return chunk_mapping.get(location, {})
+
+
+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_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_val : float
+        The minimum value (in degrees) of the coordinate range (inclusive).
+
+    max_val : float
+        The maximum value (in degrees) of the coordinate range (inclusive).
+
+    resolution : float
+        The spacing (in degrees) between each coordinate in the array. The resolution will
+        be rounded to the nearest value of the form `1/n` or an integer, depending on the size
+        of the resolution value.
+
+    Returns
+    -------
+    numpy.ndarray
+        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
+
+    # Search for the best fraction or integer approximation to the resolution
+    for n in range(1, 1000):  # Try fractions 1/n, where n ranges from 1 to 999
+        if resolution <= 1:
+            fraction = (
+                1 / n
+            )  # For small resolutions (<= 1), consider fractions of the form 1/n
+        else:
+            fraction = n  # For larger resolutions (>1), consider integers (n)
+
+        diff = abs(
+            fraction - resolution
+        )  # Calculate the difference between the fraction and the resolution
+
+        if diff < min_diff:  # If the current fraction is a better approximation
+            min_diff = diff
+            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_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)
+
+    # Truncate any values that exceed max (including small floating point errors)
+    target = target[target <= end_int + 1e-10]
+
+    return target.astype(np.float32)
+
+
+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.
+
+    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
+        Input array to trim. May contain NaNs at the edges along `xdim`.
+    xdim : str
+        The dimension along which to remove NaN-only slices.
+    layer_depth : xr.DataArray, optional
+        Optional array to evaluate NaN positions. If not provided, `field` is used instead.
+
+    Returns
+    -------
+    field : xr.DataArray
+        Trimmed `field`, with leading and trailing NaN-only slices removed along `xdim`.
+    layer_depth : xr.DataArray or None
+        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:
+        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:
+            raise ValueError(
+                f"Cannot trim along dimension '{xdim}': expected at most one other dimension, "
+                f"but got {len(other_dims)} ({other_dims})."
+            )
+
+        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(
+                    {xdim: slice(first_valid, last_valid + 1)}
+                )
+
+    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

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))

roms_tools-3.0.0.dist-info/METADATA

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: roms-tools
+Version: 3.0.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)