roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/tests/test_utils.py

@@ -1,16 +1,255 @@
-from roms_tools.setup.utils import interpolate_from_climatology
-from roms_tools.setup.datasets import ERA5Correction
-from roms_tools.setup.download import download_test_data
-import xarray as xr
+import pytest
+from pathlib import Path
+import xarray.testing as xrt
 
+from roms_tools.utils import partition, partition_netcdf
+from roms_tools import Grid
 
-def test_interpolate_from_climatology():
 
-    fname = download_test_data("ERA5_regional_test_data.nc")
-    era5_times = xr.open_dataset(fname).time
+@pytest.fixture
+def grid():
 
-    climatology = ERA5Correction()
-    field = climatology.ds["ssr_corr"]
+    grid = Grid(nx=30, ny=30, size_x=80, size_y=80, center_lon=-20, center_lat=0, rot=0)
 
-    interpolated_field = interpolate_from_climatology(field, "time", era5_times)
-    assert len(interpolated_field.time) == len(era5_times)
+    return grid
+
+
+class TestPartitionGrid:
+    def test_partition_grid_along_x(self, grid):
+        _, [ds1, ds2, ds3] = partition(grid.ds, np_eta=3, np_xi=1)
+
+        assert ds1.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 32,
+            "xi_u": 31,
+            "eta_v": 10,
+            "eta_coarse": 6,
+            "xi_coarse": 17,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds2.sizes == {
+            "eta_rho": 10,
+            "xi_rho": 32,
+            "xi_u": 31,
+            "eta_v": 10,
+            "eta_coarse": 5,
+            "xi_coarse": 17,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds3.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 32,
+            "xi_u": 31,
+            "eta_v": 11,
+            "eta_coarse": 6,
+            "xi_coarse": 17,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+
+    def test_partition_grid_along_y(self, grid):
+        _, [ds1, ds2, ds3] = partition(grid.ds, np_eta=1, np_xi=3)
+
+        assert ds1.sizes == {
+            "eta_rho": 32,
+            "xi_rho": 11,
+            "xi_u": 10,
+            "eta_v": 31,
+            "eta_coarse": 17,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds2.sizes == {
+            "eta_rho": 32,
+            "xi_rho": 10,
+            "xi_u": 10,
+            "eta_v": 31,
+            "eta_coarse": 17,
+            "xi_coarse": 5,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds3.sizes == {
+            "eta_rho": 32,
+            "xi_rho": 11,
+            "xi_u": 11,
+            "eta_v": 31,
+            "eta_coarse": 17,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+
+    def test_partition_grid_along_xy(self, grid):
+        # decomposition is increasing eta to the right, increasing xi down
+        # fmt: off
+        _, [ds1, ds2, ds3,
+         ds4, ds5, ds6,
+         ds7, ds8, ds9] = partition(grid.ds, np_eta=3, np_xi=3)
+        # fmt: on
+
+        assert ds1.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 11,
+            "xi_u": 10,
+            "eta_v": 10,
+            "eta_coarse": 6,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds4.sizes == {
+            "eta_rho": 10,
+            "xi_rho": 11,
+            "xi_u": 10,
+            "eta_v": 10,
+            "eta_coarse": 5,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds7.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 11,
+            "xi_u": 10,
+            "eta_v": 11,
+            "eta_coarse": 6,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds2.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 10,
+            "xi_u": 10,
+            "eta_v": 10,
+            "eta_coarse": 6,
+            "xi_coarse": 5,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds5.sizes == {
+            "eta_rho": 10,
+            "xi_rho": 10,
+            "xi_u": 10,
+            "eta_v": 10,
+            "eta_coarse": 5,
+            "xi_coarse": 5,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds8.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 10,
+            "xi_u": 10,
+            "eta_v": 11,
+            "eta_coarse": 6,
+            "xi_coarse": 5,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds3.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 11,
+            "xi_u": 11,
+            "eta_v": 10,
+            "eta_coarse": 6,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds6.sizes == {
+            "eta_rho": 10,
+            "xi_rho": 11,
+            "xi_u": 11,
+            "eta_v": 10,
+            "eta_coarse": 5,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+        assert ds9.sizes == {
+            "eta_rho": 11,
+            "xi_rho": 11,
+            "xi_u": 11,
+            "eta_v": 11,
+            "eta_coarse": 6,
+            "xi_coarse": 6,
+            "s_rho": 100,
+            "s_w": 101,
+        }
+
+    def test_partition_grid_no_op(self, grid):
+        _, partitioned_datasets = partition(grid.ds, np_eta=1, np_xi=1)
+
+        xrt.assert_identical(partitioned_datasets[0], grid.ds)
+
+    def test_invalid_partitioning(self, grid):
+        with pytest.raises(
+            ValueError, match="np_eta and np_xi must be positive integers"
+        ):
+            partition(grid.ds, np_eta=3.0, np_xi=1)
+
+        with pytest.raises(
+            ValueError, match="np_eta and np_xi must be positive integers"
+        ):
+            partition(grid.ds, np_eta=-3, np_xi=1)
+
+        with pytest.raises(ValueError, match="cannot be evenly divided"):
+            partition(grid.ds, np_eta=4, np_xi=1)
+
+
+class TestPartitionMissingDims:
+    def test_partition_missing_dims(self, grid):
+        dims_to_drop = ["xi_u", "eta_v", "eta_coarse", "xi_coarse"]
+
+        ds_missing_dims = grid.ds.drop_dims(dims_to_drop)
+
+        _, partitioned_datasets = partition(ds_missing_dims, np_eta=1, np_xi=1)
+
+        xrt.assert_identical(partitioned_datasets[0], ds_missing_dims)
+
+    def test_partition_missing_all_dims(self, grid):
+        # this is all the partitionable dims, so in this case the file will just be copied np_eta * np_xi times
+        dims_to_drop = ["eta_rho", "xi_rho", "xi_u", "eta_v", "eta_coarse", "xi_coarse"]
+
+        ds_missing_dims = grid.ds.drop_dims(dims_to_drop)
+
+        _, partitioned_datasets = partition(ds_missing_dims, np_eta=1, np_xi=1)
+
+        xrt.assert_identical(partitioned_datasets[0], ds_missing_dims)
+
+
+class TestFileNumbers:
+    def test_partition_file_numbers(self, grid):
+        np_eta = 3
+        np_xi = 5
+        file_numbers, _ = partition(grid.ds, np_eta=np_eta, np_xi=np_xi)
+
+        # Generate the expected file numbers
+        expected_file_numbers = list(range(np_eta * np_xi))
+
+        # Check if file_numbers is a continuous range without gaps
+        assert set(file_numbers) == set(expected_file_numbers)
+
+
+class TestPartitionNetcdf:
+    def test_partition_netcdf(self, grid, tmp_path):
+        filepath = tmp_path / "test_grid.nc"
+        grid.save(filepath)
+
+        saved_filenames = partition_netcdf(filepath, np_eta=3, np_xi=3)
+
+        filepath_str = str(filepath.with_suffix(""))
+        expected_filepath_list = [
+            Path(filepath_str + f".{index}.nc") for index in range(9)
+        ]
+
+        assert saved_filenames == expected_filepath_list
+
+        for expected_filepath in expected_filepath_list:
+            assert expected_filepath.exists()
+            expected_filepath.unlink()

roms_tools/utils.py

@@ -0,0 +1,298 @@
+from numbers import Integral
+
+import numpy as np
+import xarray as xr
+from typing import Union
+from pathlib import Path
+
+
+def partition(
+    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1
+) -> tuple[list[int], list[xr.Dataset]]:
+    """
+    Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial tiles.
+
+    This function divides the input dataset into `np_eta` by `np_xi` tiles, where each tile
+    represents a subdomain of the original dataset. The partitioning is performed along
+    the spatial dimensions `eta_rho`, `xi_rho`, `eta_v`, `xi_u`, `eta_coarse`, and `xi_coarse`,
+    depending on which dimensions are present in the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The input ROMS dataset that is to be partitioned.
+
+    np_eta : int, optional
+        The number of partitions along the `eta` direction. Must be a positive integer. Default is 1.
+
+    np_xi : int, optional
+        The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
+
+    Returns
+    -------
+    tuple[list[int], list[xr.Dataset]]
+        A tuple containing two elements:
+        - A list of integers representing the file numbers associated with each partition.
+        - A list of `xarray.Dataset` objects, each representing a partitioned subdomain of the original dataset.
+
+    Raises
+    ------
+    ValueError
+        If `np_eta` or `np_xi` is not a positive integer, or if the dataset cannot be evenly partitioned
+        into the specified number of tiles.
+
+
+    Example
+    -------
+    >>> partitioned_file_numbers, partitioned_datasets = partition(
+    ...     ds, np_eta=2, np_xi=3
+    ... )
+    >>> print(partitioned_file_numbers)
+    [0, 1, 2, 3, 4, 5]
+    >>> print([ds.sizes for ds in partitioned_datasets])
+    [{'eta_rho': 50, 'xi_rho': 50}, {'eta_rho': 50, 'xi_rho': 50}, ...]
+
+    This example partitions the dataset into 2 tiles along the `eta` direction and 3 tiles
+    along the `xi` direction, resulting in a total of 6 partitions.
+    """
+
+    if (
+        not isinstance(np_eta, Integral)
+        or np_eta < 1
+        or not isinstance(np_xi, Integral)
+        or np_xi < 1
+    ):
+        raise ValueError("np_eta and np_xi must be positive integers")
+
+    partitionable_dims_maybe_present = [
+        "eta_rho",
+        "xi_rho",
+        "eta_v",
+        "xi_u",
+        "eta_coarse",
+        "xi_coarse",
+    ]
+    dims_to_partition = [d for d in partitionable_dims_maybe_present 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"]:
+        # TODO how are we supposed to know if eta is periodic if eta_v doesn't appear? partit.F doesn't say...
+        n_eta_ghost_cells = 0
+    else:
+        n_eta_ghost_cells = 1
+
+    # if xi is periodic there are no ghost cells along those dimensions
+    if "xi_u" in ds.sizes and ds.sizes["xi_rho"] == ds.sizes["xi_u"]:
+        n_xi_ghost_cells = 0
+    else:
+        n_xi_ghost_cells = 1
+
+    def integer_division_or_raise(a: int, b: int, dimension: str) -> int:
+        """
+        Perform integer division and ensure that the division is exact.
+
+        Parameters
+        ----------
+        a : int
+            The numerator for the division.
+        b : int
+            The denominator for the division.
+        dimension : str
+            The name of the dimension being partitioned, used for error reporting.
+
+        Returns
+        -------
+        int
+            The result of the integer division.
+
+        Raises
+        ------
+        ValueError
+            If the division is not exact, indicating that the domain cannot be evenly divided
+            along the specified dimension.
+        """
+        remainder = a % b
+        if remainder == 0:
+            return a // b
+        else:
+            raise ValueError(
+                f"Dimension '{dimension}' of size {a} cannot be evenly divided into {b} partitions."
+            )
+
+    if "eta_rho" in dims_to_partition:
+        eta_rho_domain_size = integer_division_or_raise(
+            ds.sizes["eta_rho"] - 2 * n_eta_ghost_cells, np_eta, "eta_rho"
+        )
+    if "xi_rho" in dims_to_partition:
+        xi_rho_domain_size = integer_division_or_raise(
+            ds.sizes["xi_rho"] - 2 * n_xi_ghost_cells, np_xi, "xi_rho"
+        )
+
+    if "eta_v" in dims_to_partition:
+        eta_v_domain_size = integer_division_or_raise(
+            ds.sizes["eta_v"] - 1 * n_eta_ghost_cells, np_eta, "eta_v"
+        )
+    if "xi_u" in dims_to_partition:
+        xi_u_domain_size = integer_division_or_raise(
+            ds.sizes["xi_u"] - 1 * n_xi_ghost_cells, np_xi, "xi_u"
+        )
+
+    if "eta_coarse" in dims_to_partition:
+        eta_coarse_domain_size = integer_division_or_raise(
+            ds.sizes["eta_coarse"] - 2 * n_eta_ghost_cells, np_eta, "eta_coarse"
+        )
+    if "xi_coarse" in dims_to_partition:
+        xi_coarse_domain_size = integer_division_or_raise(
+            ds.sizes["xi_coarse"] - 2 * n_xi_ghost_cells, np_xi, "xi_coarse"
+        )
+
+    # unpartitioned dimensions should have sizes unchanged
+    partitioned_sizes = {
+        dim: [size] for dim, size in ds.sizes.items() if dim in dims_to_partition
+    }
+
+    # TODO refactor to use two functions for odd- and even-length dimensions
+    if "eta_v" in dims_to_partition:
+        partitioned_sizes["eta_v"] = [eta_v_domain_size] * (np_eta - 1) + [
+            eta_v_domain_size + n_eta_ghost_cells
+        ]
+    if "xi_u" in dims_to_partition:
+        partitioned_sizes["xi_u"] = [xi_u_domain_size] * (np_xi - 1) + [
+            xi_u_domain_size + n_xi_ghost_cells
+        ]
+
+    if np_eta > 1:
+        partitioned_sizes["eta_rho"] = (
+            [eta_rho_domain_size + n_eta_ghost_cells]
+            + [eta_rho_domain_size] * (np_eta - 2)
+            + [eta_rho_domain_size + n_eta_ghost_cells]
+        )
+
+        if "eta_coarse" in dims_to_partition:
+            partitioned_sizes["eta_coarse"] = (
+                [eta_coarse_domain_size + n_eta_ghost_cells]
+                + [eta_coarse_domain_size] * (np_eta - 2)
+                + [eta_coarse_domain_size + n_eta_ghost_cells]
+            )
+
+    if np_xi > 1:
+        partitioned_sizes["xi_rho"] = (
+            [xi_rho_domain_size + n_xi_ghost_cells]
+            + [xi_rho_domain_size] * (np_xi - 2)
+            + [xi_rho_domain_size + n_xi_ghost_cells]
+        )
+
+        if "xi_coarse" in dims_to_partition:
+            partitioned_sizes["xi_coarse"] = (
+                [xi_coarse_domain_size + n_xi_ghost_cells]
+                + [xi_coarse_domain_size] * (np_xi - 2)
+                + [xi_coarse_domain_size + n_xi_ghost_cells]
+            )
+
+    def cumsum(pmf):
+        """Implementation of cumsum which ensures the result starts with zero"""
+        cdf = np.empty(len(pmf) + 1, dtype=int)
+        cdf[0] = 0
+        np.cumsum(pmf, out=cdf[1:])
+        return cdf
+
+    file_numbers = []
+    partitioned_datasets = []
+    for i in range(np_eta):
+        for j in range(np_xi):
+            file_number = j + (i * np_xi)
+            file_numbers.append(file_number)
+
+            indexers = {}
+
+            if "eta_rho" in dims_to_partition:
+                eta_rho_partition_indices = cumsum(partitioned_sizes["eta_rho"])
+                indexers["eta_rho"] = slice(
+                    int(eta_rho_partition_indices[i]),
+                    int(eta_rho_partition_indices[i + 1]),
+                )
+            if "xi_rho" in dims_to_partition:
+                xi_rho_partition_indices = cumsum(partitioned_sizes["xi_rho"])
+                indexers["xi_rho"] = slice(
+                    int(xi_rho_partition_indices[j]),
+                    int(xi_rho_partition_indices[j + 1]),
+                )
+
+            if "eta_v" in dims_to_partition:
+                eta_v_partition_indices = cumsum(partitioned_sizes["eta_v"])
+                indexers["eta_v"] = slice(
+                    int(eta_v_partition_indices[i]),
+                    int(eta_v_partition_indices[i + 1]),
+                )
+            if "xi_u" in dims_to_partition:
+                xi_u_partition_indices = cumsum(partitioned_sizes["xi_u"])
+                indexers["xi_u"] = slice(
+                    int(xi_u_partition_indices[j]), int(xi_u_partition_indices[j + 1])
+                )
+
+            if "eta_coarse" in dims_to_partition:
+                eta_coarse_partition_indices = cumsum(partitioned_sizes["eta_coarse"])
+                indexers["eta_coarse"] = slice(
+                    int(eta_coarse_partition_indices[i]),
+                    int(eta_coarse_partition_indices[i + 1]),
+                )
+
+            if "xi_coarse" in dims_to_partition:
+                xi_coarse_partition_indices = cumsum(partitioned_sizes["xi_coarse"])
+                indexers["xi_coarse"] = slice(
+                    int(xi_coarse_partition_indices[j]),
+                    int(xi_coarse_partition_indices[j + 1]),
+                )
+
+            partitioned_ds = ds.isel(**indexers)
+
+            partitioned_datasets.append(partitioned_ds)
+
+    return file_numbers, partitioned_datasets
+
+
+def partition_netcdf(
+    filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
+) -> None:
+    """
+    Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
+
+    This function divides the dataset in the specified NetCDF file into `np_eta` by `np_xi` tiles.
+    Each tile is saved as a separate NetCDF file.
+
+    Parameters
+    ----------
+    filepath : Union[str, Path]
+        The path to the input NetCDF file.
+
+    np_eta : int, optional
+        The number of partitions along the `eta` direction. Must be a positive integer. Default is 1.
+
+    np_xi : int, optional
+        The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
+
+    Returns
+    -------
+    List[Path]
+        A list of Path objects for the filenames that were saved.
+    """
+
+    # Ensure filepath is a Path object
+    filepath = Path(filepath)
+
+    # Open the dataset
+    ds = xr.open_dataset(filepath.with_suffix(".nc"))
+
+    # Partition the dataset
+    file_numbers, partitioned_datasets = partition(ds, np_eta=np_eta, np_xi=np_xi)
+
+    # Generate paths to the partitioned files
+    base_filepath = filepath.with_suffix("")
+    paths_to_partitioned_files = [
+        Path(f"{base_filepath}.{file_number}.nc") for file_number in file_numbers
+    ]
+
+    # Save the partitioned datasets to files
+    xr.save_mfdataset(partitioned_datasets, paths_to_partitioned_files)
+
+    return paths_to_partitioned_files

{roms_tools-1.2.0.dist-info → roms_tools-1.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: roms-tools
-Version: 1.2.0
+Version: 1.4.0
 Summary: Tools for running and analysing UCLA-ROMS simulations
 Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>
 License: Apache-2
@@ -19,16 +19,19 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: xarray>=2022.6.0
-Requires-Dist: xarray-datatree
 Requires-Dist: numpy
-Requires-Dist: netcdf4
 Requires-Dist: pooch
 Requires-Dist: matplotlib
 Requires-Dist: cartopy
 Requires-Dist: packaging
 Requires-Dist: scipy
 Requires-Dist: gcm-filters
-Requires-Dist: numba
+Requires-Dist: netcdf4
+Requires-Dist: PyYAML
+Requires-Dist: pyamg
+Requires-Dist: bottleneck
+Provides-Extra: dask
+Requires-Dist: dask; extra == "dask"
 
 # ROMS-Tools
 [![PyPI version](https://img.shields.io/pypi/v/roms-tools.svg)](https://pypi.org/project/roms-tools/)
@@ -50,29 +53,54 @@ A suite of python tools for setting up a [ROMS](https://github.com/CESR-lab/ucla
 pip install roms-tools
 ```
 
+If you want to use `ROMS-Tools` together with dask (which we recommend), you can
+install `ROMS-Tools` along with the additional dependency via:
+
+```bash
+pip install roms-tools[dask]
+```
+
 ### Installation from GitHub
 
-To obtain the latest development version, clone the source repository and install it:
+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.yml
+conda activate romstools
+```
+
+Finally, install `ROMS-Tools` in the same environment:
+
+```bash
 pip install -e .
 ```
 
+If you want to use `ROMS-Tools` together with dask (which we recommend), you can
+install `ROMS-Tools` along with the additional dependency via:
+
+```bash
+pip install -e .[dask]
+```
+
 
 ### Run the tests
 
-Before running the tests, you can install and activate the following conda environment:
+Before running the tests, you can activate the conda environment created in the previous section:
 
 ```bash
-cd roms-tools
-conda env create -f ci/environment.yml
 conda activate romstools
 ```
 
 Check the installation of `ROMS-Tools` has worked by running the test suite
 ```bash
+cd roms-tools
 pytest
 ```