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

roms_tools/tests/test_tiling/test_partition.py

@@ -1,9 +1,10 @@
-import pytest
 from pathlib import Path
+
+import pytest
 import xarray.testing as xrt
 
-from roms_tools.tiling.partition import partition, partition_netcdf
 from roms_tools import Grid
+from roms_tools.tiling.partition import partition, partition_netcdf
 
 
 @pytest.fixture
@@ -231,6 +232,19 @@ class TestPartitionGrid:
         with pytest.raises(ValueError, match="cannot be evenly divided"):
             partition(grid.ds, np_eta=4, np_xi=1)
 
+        with pytest.raises(ValueError, match="cannot be evenly divided"):
+            partition(grid.ds, np_eta=10, np_xi=10)
+
+    def test_skip_coarse_dims(self, grid):
+        """Test that coarse dimensions remain unchanged when excluded from
+        partitioning."""
+        _, partitioned_datasets = partition(
+            grid.ds, np_eta=10, np_xi=10, include_coarse_dims=False
+        )
+        for ds in partitioned_datasets:
+            assert ds.sizes["eta_coarse"] == grid.ds.sizes["eta_coarse"]
+            assert ds.sizes["xi_coarse"] == grid.ds.sizes["xi_coarse"]
+
 
 class TestPartitionMissingDims:
     def test_partition_missing_dims(self, grid):
@@ -271,11 +285,11 @@ class TestPartitionNetcdf:
         filepath = tmp_path / "test_grid.nc"
         grid.save(filepath)
 
-        saved_filenames = partition_netcdf(filepath, np_eta=3, np_xi=3)
+        saved_filenames = partition_netcdf(filepath, np_eta=3, np_xi=5)
 
         filepath_str = str(filepath.with_suffix(""))
         expected_filepath_list = [
-            Path(filepath_str + f".{index}.nc") for index in range(9)
+            Path(filepath_str + f".{index:02d}.nc") for index in range(15)
         ]
 
         assert saved_filenames == expected_filepath_list

roms_tools/tests/test_utils.py

@@ -0,0 +1,21 @@
+import numpy as np
+import pytest
+
+from roms_tools.utils import _generate_focused_coordinate_range
+
+
+@pytest.mark.parametrize(
+    "min_val, max_val, center, sc, N",
+    [
+        (-20.0, 5.5, -3.1, 1.0, 100),
+        (100.0, 200.0, 150.0, 30.0, 100),
+        (0.0, 2000.0, 150.0, 0.0, 100),
+        (0.0, 2000.0, 150.0, 30.0, 100),
+    ],
+)
+def test_coordinate_range_monotonicity(min_val, max_val, center, sc, N):
+    centers, faces = _generate_focused_coordinate_range(
+        min_val=min_val, max_val=max_val, center=center, sc=sc, N=N
+    )
+    assert np.all(np.diff(faces) > 0), "faces is not strictly increasing"
+    assert np.all(np.diff(centers) > 0), "centers is not strictly increasing"

roms_tools/tests/test_vertical_coordinate.py

@@ -1,12 +1,13 @@
-import pytest
 import numpy as np
+import pytest
 import xarray as xr
 
+from roms_tools.constants import UPPER_BOUND_THETA_B, UPPER_BOUND_THETA_S
 from roms_tools.vertical_coordinate import (
     compute_cs,
-    sigma_stretch,
     compute_depth,
     compute_depth_coordinates,
+    sigma_stretch,
 )
 
 
@@ -17,11 +18,15 @@ def test_compute_cs():
     assert cs.shape == sigma.shape
     assert np.all(cs <= 0) and np.all(cs >= -1)
 
-    with pytest.raises(ValueError, match="theta_s must be between 0 and 10"):
-        compute_cs(sigma, 15, 2)
+    with pytest.raises(
+        ValueError, match=f"theta_s must be between 0 and {UPPER_BOUND_THETA_S}"
+    ):
+        compute_cs(sigma, UPPER_BOUND_THETA_S + 1, theta_b)
 
-    with pytest.raises(ValueError, match="theta_b must be between 0 and 4"):
-        compute_cs(sigma, 5, 5)
+    with pytest.raises(
+        ValueError, match=f"theta_b must be between 0 and {UPPER_BOUND_THETA_B}"
+    ):
+        compute_cs(sigma, theta_s, UPPER_BOUND_THETA_B + 1)
 
 
 def test_sigma_stretch():

roms_tools/tiling/partition.py

@@ -1,14 +1,15 @@
 from numbers import Integral
+from pathlib import Path
+from typing import Union
 
 import numpy as np
 import xarray as xr
-from typing import Union
-from pathlib import Path
+
 from roms_tools.utils import save_datasets
 
 
 def partition(
-    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1
+    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1, include_coarse_dims: bool = True
 ) -> tuple[list[int], list[xr.Dataset]]:
     """Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial
     tiles.
@@ -29,6 +30,11 @@ def partition(
     np_xi : int, optional
         The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
 
+    include_coarse_dims : bool, optional
+        Whether to include coarse grid dimensions (`eta_coarse`, `xi_coarse`) in the partitioning.
+        If False, these dimensions will not be split. Relevant if none of the coarse resolution variables are actually used by ROMS.
+        Default is True.
+
     Returns
     -------
     tuple[list[int], list[xr.Dataset]]
@@ -66,17 +72,18 @@ def partition(
     ):
         raise ValueError("np_eta and np_xi must be positive integers")
 
-    partitionable_dims_maybe_present = [
+    base_dims = [
         "eta_rho",
         "xi_rho",
         "eta_v",
         "xi_u",
         "eta_psi",
         "xi_psi",
-        "eta_coarse",
-        "xi_coarse",
     ]
-    dims_to_partition = [d for d in partitionable_dims_maybe_present if d in ds.dims]
+    coarse_dims = ["eta_coarse", "xi_coarse"]
+    partitionable_dims = base_dims + coarse_dims if include_coarse_dims else base_dims
+
+    dims_to_partition = [d for d in partitionable_dims if d in ds.dims]
 
     # if eta is periodic there are no ghost cells along those dimensions
     if "eta_v" in ds.sizes and ds.sizes["eta_rho"] == ds.sizes["eta_v"]:
@@ -291,7 +298,10 @@ def partition(
 
 
 def partition_netcdf(
-    filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
+    filepath: Union[str, Path],
+    np_eta: int = 1,
+    np_xi: int = 1,
+    include_coarse_dims: bool = True,
 ) -> None:
     """Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
 
@@ -309,6 +319,11 @@ def partition_netcdf(
     np_xi : int, optional
         The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
 
+    include_coarse_dims : bool, optional
+        Whether to include coarse grid dimensions (`eta_coarse`, `xi_coarse`) in the partitioning.
+        If False, these dimensions will not be split. Relevant if none of the coarse resolution variables are actually used by ROMS.
+        Default is True.
+
     Returns
     -------
     List[Path]
@@ -319,15 +334,19 @@ def partition_netcdf(
     filepath = Path(filepath)
 
     # Open the dataset
-    ds = xr.open_dataset(filepath.with_suffix(".nc"))
+    ds = xr.open_dataset(filepath.with_suffix(".nc"), decode_timedelta=False)
 
     # Partition the dataset
-    file_numbers, partitioned_datasets = partition(ds, np_eta=np_eta, np_xi=np_xi)
+    file_numbers, partitioned_datasets = partition(
+        ds, np_eta=np_eta, np_xi=np_xi, include_coarse_dims=include_coarse_dims
+    )
 
     # Generate paths to the partitioned files
     base_filepath = filepath.with_suffix("")
+    ndigits = len(str(max(np.array(file_numbers))))
     paths_to_partitioned_files = [
-        Path(f"{base_filepath}.{file_number}") for file_number in file_numbers
+        Path(f"{base_filepath}.{file_number:0{ndigits}d}")
+        for file_number in file_numbers
     ]
 
     # Save the partitioned datasets to files

roms_tools/utils.py

@@ -1,10 +1,12 @@
-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
 
 
 def _load_data(
@@ -14,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.
 
@@ -33,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
     -------
@@ -54,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"[\*\?\[\]]")
 
@@ -140,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 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 +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
 
 
@@ -374,21 +413,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
@@ -425,8 +464,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 +476,161 @@ 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.
 
-    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`.
+    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
-        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 +638,34 @@ 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

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