roms-tools 3.1.2__py3-none-any.whl → 3.2.0__py3-none-any.whl

roms_tools/tiling/join.py

@@ -0,0 +1,189 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Literal, cast
+
+import xarray as xr
+
+from roms_tools.utils import FilePaths, _path_list_from_input
+
+
+def open_partitions(files: FilePaths) -> xr.Dataset:
+    """
+    Open partitioned ROMS netCDF files as a single dataset.
+
+    Parameters
+    ----------
+    files: str | List[str | Path]
+        List or wildcard pattern describing files to join,
+        e.g. "roms_rst.20121209133435.*.nc"
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset containing unified partitioned datasets
+    """
+    filepaths = _path_list_from_input(files)
+    datasets = [xr.open_dataset(p, decode_timedelta=True) for p in sorted(filepaths)]
+    joined = join_datasets(datasets)
+    return joined
+
+
+def join_netcdf(files: FilePaths, output_path: Path | None = None) -> Path:
+    """
+    Join partitioned NetCDFs into a single dataset.
+
+    Parameters
+    ----------
+    files : str | List[str | Path]
+        List or wildcard pattern describing files to join,
+        e.g. "roms_rst.20121209133435.*.nc"
+
+    output_path : Path, optional
+        If provided, the joined dataset will be saved to this path.
+        Otherwise, the common base of pattern (e.g. roms_rst.20121209133435.nc)
+        will be used.
+
+    Returns
+    -------
+    Path
+        The path of the saved file
+    """
+    filepaths = _path_list_from_input(files)
+    # Determine output path if not provided
+    if output_path is None:
+        # e.g. roms_rst.20120101120000.023.nc -> roms_rst.20120101120000.nc
+        output_path = filepaths[0].with_suffix("").with_suffix(".nc")
+
+    joined = open_partitions(cast(FilePaths, filepaths))
+    joined.to_netcdf(output_path)
+    print(f"Saved joined dataset to: {output_path}")
+
+    return output_path
+
+
+def _find_transitions(dim_sizes: list[int]) -> list[int]:
+    """Finds the indices of all transitions in a list of dimension sizes.
+
+    A transition is a point where the dimension size changes from the previous one.
+    This function is used to determine the number of partitions (e.g., np_eta or np_xi).
+
+    Parameters
+    ----------
+    dim_sizes : list[int]
+        A list of integer sizes for a given dimension across multiple datasets.
+
+    Returns
+    -------
+    List[int]
+        A list of indices where a transition was detected.
+    """
+    transitions: list[int] = []
+    if len(dim_sizes) < 2:
+        return transitions
+
+    for i in range(1, len(dim_sizes)):
+        if dim_sizes[i] != dim_sizes[i - 1]:
+            transitions.append(i)
+    return transitions
+
+
+def _find_common_dims(
+    direction: Literal["xi", "eta"], datasets: Sequence[xr.Dataset]
+) -> list[str]:
+    """Finds all common dimensions along the xi or eta direction amongst a list of Datasets.
+
+    Parameters
+    ----------
+    direction: str ("xi" or "eta")
+        The direction in which to seek a common dimension
+    datasets: Sequence[xr.Dataset]:
+        The datasets in which to look
+
+    Returns
+    -------
+    common_dim: list[str]
+        The dimensions common to all specified datasets along 'direction'
+    """
+    if direction not in ["xi", "eta"]:
+        raise ValueError("'direction' must be 'xi' or 'eta'")
+    dims = []
+    for point in ["rho", "u", "v"]:
+        if all(f"{direction}_{point}" in d.dims for d in datasets):
+            dims.append(f"{direction}_{point}")
+    if not dims:
+        raise ValueError(f"No common point found along direction {direction}")
+    return dims
+
+
+def _infer_partition_layout_from_datasets(
+    datasets: Sequence[xr.Dataset],
+) -> tuple[int, int]:
+    """Infer np_eta, np_xi from datasets."""
+    nd = len(datasets)
+    if nd == 1:
+        return 1, 1
+
+    eta_dims = _find_common_dims("eta", datasets)
+    first_eta_transition = nd
+
+    for eta_dim in eta_dims:
+        dim_sizes = [ds.sizes.get(eta_dim, 0) for ds in datasets]
+        eta_transitions = _find_transitions(dim_sizes)
+        if eta_transitions and (min(eta_transitions) < first_eta_transition):
+            first_eta_transition = min(eta_transitions)
+    if first_eta_transition < nd:
+        np_xi = first_eta_transition
+        np_eta = nd // np_xi
+        return np_xi, np_eta
+    # If we did not successfully find np_xi,np_eta using eta points
+    # then we have a single-column grid:
+
+    return nd, 1
+
+
+def join_datasets(datasets: Sequence[xr.Dataset]) -> xr.Dataset:
+    """Take a sequence of partitioned Datasets and return a joined Dataset."""
+    np_xi, np_eta = _infer_partition_layout_from_datasets(datasets)
+
+    # Arrange into grid
+    grid = [[datasets[j + i * np_xi] for j in range(np_xi)] for i in range(np_eta)]
+
+    # Join each row (along xi_*)
+    rows_joined = []
+    for row in grid:
+        all_vars = set().union(*(ds.data_vars for ds in row))
+        row_dataset = xr.Dataset()
+
+        for varname in all_vars:
+            var_slices = [ds[varname] for ds in row if varname in ds]
+            xi_dims = [dim for dim in var_slices[0].dims if dim.startswith("xi_")]
+
+            if not xi_dims:
+                row_dataset[varname] = var_slices[0]
+            else:
+                xi_dim = xi_dims[0]
+                row_dataset[varname] = xr.concat(
+                    var_slices, dim=xi_dim, combine_attrs="override"
+                )
+
+        rows_joined.append(row_dataset)
+
+    # Join all rows (along eta_*)
+    final_dataset = xr.Dataset()
+    all_vars = set().union(*(ds.data_vars for ds in rows_joined))
+
+    for varname in all_vars:
+        var_slices = [ds[varname] for ds in rows_joined if varname in ds]
+        eta_dims = [dim for dim in var_slices[0].dims if dim.startswith("eta_")]
+
+        if not eta_dims:
+            final_dataset[varname] = var_slices[0]
+        else:
+            eta_dim = eta_dims[0]
+            final_dataset[varname] = xr.concat(
+                var_slices, dim=eta_dim, combine_attrs="override"
+            )
+    # Copy attributes from first dataset
+    final_dataset.attrs = datasets[0].attrs
+
+    return final_dataset

roms_tools/utils.py

@@ -3,16 +3,51 @@ import logging
 import re
 import textwrap
 import warnings
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
+from dataclasses import dataclass
 from importlib.util import find_spec
 from pathlib import Path
+from typing import TypeAlias
 
 import numpy as np
 import xarray as xr
-from attr import dataclass
 
 from roms_tools.constants import R_EARTH
 
+FilePaths: TypeAlias = str | Path | list[Path | str]
+
+
+def _path_list_from_input(files: FilePaths) -> list[Path]:
+    """Converts a generic user input to a list of Paths.
+
+    Takes a list of strings or paths, or wildcard pattern, and
+    returns a list of pathlib.Path objects
+
+    Parameters
+    ----------
+    files: FilePaths
+        A list of files (str, Path), single path as a str or Path, or a wildcard string
+
+    Returns
+    -------
+    List[Path]
+        A list of pathlib.Paths
+    """
+    if isinstance(files, str):
+        filepaths = sorted(Path(files).parent.glob(Path(files).name))
+        if not filepaths:
+            raise FileNotFoundError(f"No files matched: {files}")
+    elif isinstance(files, Path):
+        filepaths = [
+            files,
+        ]
+    elif isinstance(files, list):
+        filepaths = [Path(f) for f in files]
+    else:
+        raise TypeError("'files' should be str, Path, or List[Path | str]")
+
+    return filepaths
+
 
 @dataclass
 class FileMatchResult:
@@ -25,56 +60,39 @@ class FileMatchResult:
 
 
 def _get_file_matches(
-    filename: str | Path | list[str | Path],
+    filename: str | Path | Sequence[str | Path],
 ) -> FileMatchResult:
     """Filter the filename using an optional wildcard search in the filename.
 
     Parameters
     ----------
-    filename : str or Path or list of str or Path
+    filename : str | Path | Sequence[str | Path]
         An item to search for matches.
     """
     # Precompile the regex for matching wildcard characters
     wildcard_regex = re.compile(r"[\*\?\[\]]")
 
-    # Convert Path objects to strings
-    if isinstance(filename, str | Path):
-        filename_str = str(filename)
-    elif isinstance(filename, list):
-        filename_str = [str(f) for f in filename]
+    # Normalize input into a list of strings
+    if isinstance(filename, (str | Path)):
+        filenames: list[str] = [str(filename)]
+    elif isinstance(filename, Sequence):
+        filenames = [str(f) for f in filename]
     else:
-        msg = "filename must be a string, Path, or a list of strings/Paths."
+        msg = "filename must be a string, Path, or a sequence of strings/Paths."
         raise ValueError(msg)
 
-    # Handle the case when filename is a string
-    contains_wildcard = False
-    matching_files = []
-
-    if isinstance(filename_str, str):
-        contains_wildcard = bool(wildcard_regex.search(filename_str))
-        if contains_wildcard:
-            matching_files = glob.glob(filename_str)
-            if not matching_files:
-                msg = f"No files found matching the pattern '{filename_str}'."
-                raise FileNotFoundError(msg)
-        else:
-            matching_files = [filename_str]
-
-    # Handle the case when filename is a list
-    elif isinstance(filename_str, list):
-        # contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
-        if contains_wildcard := any(wildcard_regex.search(f) for f in filename_str):
-            matching_files = []
-            for f in filename_str:
-                files = glob.glob(f)
-                if not files:
-                    msg = f"No files found matching the pattern '{f}'."
-                    raise FileNotFoundError(msg)
-                matching_files.extend(files)
+    contains_wildcard = any(wildcard_regex.search(f) for f in filenames)
+    matching_files: list[str] = []
+
+    for f in filenames:
+        if wildcard_regex.search(f):
+            files = glob.glob(f)
+            if not files:
+                raise FileNotFoundError(f"No files found matching the pattern '{f}'.")
+            matching_files.extend(files)
         else:
-            matching_files = filename_str
+            matching_files.append(f)
 
-    # Sort the matching files
     return FileMatchResult(
         contains_wildcard=contains_wildcard,
         matches=sorted(matching_files),
@@ -132,11 +150,60 @@ def _get_ds_combine_base_params() -> dict[str, str]:
     }
 
 
+def get_dask_chunks(
+    dim_names: dict[str, str], time_chunking: bool = True
+) -> dict[str, int]:
+    """Return the default dask chunks for ROMS datasets.
+
+    Parameters
+    ----------
+    dim_names : dict[str, str]
+        Dictionary specifying the names of dimensions in the dataset.
+        - For lat-lon datasets, provide keys "latitude" and "longitude" (and optionally "depth" and "time").
+        - For ROMS datasets, the default ROMS dimensions are assumed ("eta_rho", "xi_rho", "s_rho", etc.).
+    time_chunking : bool, optional
+        Whether to chunk along the time dimension.
+        - True: chunk time dimension with size 1 (useful for processing large time-series data with Dask).
+        - False: do not explicitly chunk time; Dask will use default auto-chunking.
+        Defaults to True.
+
+    Returns
+    -------
+    dict[str, int]
+        The default dask chunks for ROMS datasets.
+    """
+    if "latitude" in dim_names and "longitude" in dim_names:
+        # for lat-lon datasets
+        chunks = {
+            dim_names["latitude"]: -1,
+            dim_names["longitude"]: -1,
+        }
+    else:
+        # For ROMS datasets
+        chunks = {
+            "eta_rho": -1,
+            "eta_v": -1,
+            "xi_rho": -1,
+            "xi_u": -1,
+            "s_rho": -1,
+        }
+
+    if "depth" in dim_names:
+        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
+
+    return chunks
+
+
 def _load_data_dask(
     filenames: list[str],
     dim_names: dict[str, str],
     time_chunking: bool = True,
     decode_times: bool = True,
+    decode_timedelta: bool = True,
     read_zarr: bool = True,
     load_kwargs: dict[str, str] | None = None,
 ) -> xr.Dataset:
@@ -158,6 +225,9 @@ def _load_data_dask(
     decode_times: bool, optional
         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.
+    decode_timedelta: bool, optional
+        If True, decode timedeltas encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
+        Defaults to True.
     read_zarr: bool, optional
         If True, use the zarr engine to read the dataset, and don't use mfdataset.
         Defaults to False.
@@ -175,28 +245,7 @@ def _load_data_dask(
         If a list of files is provided but dim_names["time"] is not available or use_dask=False.
 
     """
-    if "latitude" in dim_names and "longitude" in dim_names:
-        # for lat-lon datasets
-        chunks = {
-            dim_names["latitude"]: -1,
-            dim_names["longitude"]: -1,
-        }
-    else:
-        # For ROMS datasets
-        chunks = {
-            "eta_rho": -1,
-            "eta_v": -1,
-            "xi_rho": -1,
-            "xi_u": -1,
-            "s_rho": -1,
-        }
-
-    if "depth" in dim_names:
-        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
+    chunks = get_dask_chunks(dim_names, time_chunking)
 
     with warnings.catch_warnings():
         warnings.filterwarnings(
@@ -209,6 +258,7 @@ def _load_data_dask(
             return xr.open_zarr(
                 filenames[0],
                 decode_times=decode_times,
+                decode_timedelta=decode_timedelta,
                 chunks=chunks,
                 consolidated=None,
                 storage_options={"token": "anon"},
@@ -218,7 +268,7 @@ def _load_data_dask(
         return xr.open_mfdataset(
             filenames,
             decode_times=decode_times,
-            decode_timedelta=decode_times,
+            decode_timedelta=decode_timedelta,
             chunks=chunks,
             **kwargs,
         )
@@ -237,7 +287,7 @@ def _check_load_data_dask(use_dask: bool) -> None:
     RuntimeError
         If dask is requested but not installed.
     """
-    if use_dask and not _has_dask():
+    if use_dask and not has_dask():
         msg = (
             "Dask is required but not installed. Install it with:\n"
             "  • `pip install roms-tools[dask]` or\n"
@@ -310,12 +360,13 @@ def _check_load_data_filename(
         raise ValueError(msg)
 
 
-def _load_data(
+def load_data(
     filename: str | Path | list[str | Path],
-    dim_names: dict[str, str],
-    use_dask: bool,
+    dim_names: dict[str, str] | None = None,
+    use_dask: bool = False,
     time_chunking: bool = True,
     decode_times: bool = True,
+    decode_timedelta: bool = True,
     force_combine_nested: bool = False,
     read_zarr: bool = False,
     ds_loader_fn: Callable[[], xr.Dataset] | None = None,
@@ -324,21 +375,24 @@ def _load_data(
 
     Parameters
     ----------
-    filename : Union[str, Path, List[Union[str, Path]]]
+    filename : str | Path | list[str | Path]
         The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
         or a list of strings or Path objects containing multiple files.
-    dim_names : Dict[str, str], optional
+    dim_names : dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset.
         Required only for lat-lon datasets to map dimension names like "latitude" and "longitude".
         For ROMS datasets, this parameter can be omitted, as default ROMS dimensions ("eta_rho", "xi_rho", "s_rho") are assumed.
-    use_dask: bool
+    use_dask: bool, optional
         Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
     time_chunking : bool, optional
         If True and `use_dask=True`, the data will be chunked along the time dimension with a chunk size of 1.
         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 and timedeltas encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
+        If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
+        Defaults to True.
+    decode_timedelta: bool, optional
+        If True, decode 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.
@@ -385,6 +439,7 @@ def _load_data(
             dim_names,
             time_chunking,
             decode_times,
+            decode_timedelta,
             read_zarr,
             load_kwargs,
         )
@@ -394,7 +449,7 @@ def _load_data(
             ds = xr.open_dataset(
                 file,
                 decode_times=decode_times,
-                decode_timedelta=decode_times,
+                decode_timedelta=decode_timedelta,
                 chunks=None,
             )
             ds_list.append(ds)
@@ -629,30 +684,7 @@ def save_datasets(dataset_list, output_filenames, use_dask=False, verbose=True):
     return saved_filenames
 
 
-def get_dask_chunks(location, chunk_size):
-    """Returns the appropriate Dask chunking dictionary based on grid location.
-
-    Parameters
-    ----------
-    location : str
-        The grid location, one of "rho", "u", or "v".
-    chunk_size : int
-        The chunk size to apply.
-
-    Returns
-    -------
-    dict
-        Dictionary specifying the chunking strategy.
-    """
-    chunk_mapping = {
-        "rho": {"eta_rho": chunk_size, "xi_rho": chunk_size},
-        "u": {"eta_rho": chunk_size, "xi_u": 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):
+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).
@@ -714,7 +746,7 @@ def _generate_coordinate_range(min_val: float, max_val: float, resolution: float
     return target.astype(np.float32)
 
 
-def _generate_focused_coordinate_range(
+def generate_focused_coordinate_range(
     center: float,
     sc: float,
     min_val: float,
@@ -800,7 +832,7 @@ def _generate_focused_coordinate_range(
     return centers, faces
 
 
-def _remove_edge_nans(
+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.
@@ -876,7 +908,7 @@ def _remove_edge_nans(
     return field, layer_depth
 
 
-def _has_dask() -> bool:
+def has_dask() -> bool:
     """Determine if the Dask package is installed.
 
     Returns
@@ -888,7 +920,7 @@ def _has_dask() -> bool:
     return find_spec("dask") is not None
 
 
-def _has_gcsfs() -> bool:
+def has_gcsfs() -> bool:
     """Determine if the GCSFS package is installed.
 
     Returns
@@ -900,7 +932,7 @@ def _has_gcsfs() -> bool:
     return find_spec("gcsfs") is not None
 
 
-def _has_copernicus() -> bool:
+def has_copernicus() -> bool:
     """Determine if the Copernicus Marine Toolkit package is installed.
 
     Returns
@@ -976,7 +1008,7 @@ def infer_nominal_horizontal_resolution(
     return float(resolution_in_degrees)
 
 
-def _get_pkg_error_msg(purpose: str, package_name: str, option_name: str) -> str:
+def get_pkg_error_msg(purpose: str, package_name: str, option_name: str) -> str:
     """Generate an error message indicating how to install an optional dependency.
 
     Parameters

{roms_tools-3.1.2.dist-info → roms_tools-3.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: roms-tools
-Version: 3.1.2
+Version: 3.2.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
@@ -42,15 +42,25 @@ Requires-Dist: dask[diagnostics]; extra == "stream"
 Requires-Dist: gcsfs; extra == "stream"
 Requires-Dist: zarr; extra == "stream"
 Requires-Dist: copernicusmarine; extra == "stream"
+Provides-Extra: notebooks
+Requires-Dist: dask[diagnostics]; extra == "notebooks"
+Requires-Dist: gcsfs; extra == "notebooks"
+Requires-Dist: zarr; extra == "notebooks"
+Requires-Dist: copernicusmarine; extra == "notebooks"
+Requires-Dist: gdown; extra == "notebooks"
+Requires-Dist: ipykernel; extra == "notebooks"
+Requires-Dist: jupyter; extra == "notebooks"
 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/)
+[![Run Tests](https://github.com/CWorthy-ocean/roms-tools/actions/workflows/tests.yaml/badge.svg)](https://github.com/CWorthy-ocean/roms-tools/actions/workflows/tests.yaml?query=branch%3Amain)
 [![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)
+[![PyPI Downloads](https://static.pepy.tech/personalized-badge/roms-tools?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/roms-tools)
 
 ## Overview