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

roms_tools/tests/test_setup/test_utils.py

@@ -2,13 +2,103 @@ import logging
 from datetime import datetime
 from pathlib import Path
 
+import numpy as np
 import pytest
 import xarray as xr
 
 from roms_tools import BoundaryForcing, Grid
 from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import ERA5Correction
-from roms_tools.setup.utils import interpolate_from_climatology, validate_names
+from roms_tools.setup.utils import (
+    get_target_coords,
+    interpolate_from_climatology,
+    validate_names,
+)
+
+
+class DummyGrid:
+    """Lightweight grid wrapper mimicking the real Grid API for testing."""
+
+    def __init__(self, ds: xr.Dataset, straddle: bool):
+        """Initialize grid wrapper."""
+        self.ds = ds
+        self.straddle = straddle
+
+
+class TestGetTargetCoords:
+    def make_rho_grid(self, lons, lats, with_mask=False):
+        """Helper to create a minimal rho grid dataset."""
+        eta, xi = len(lats), len(lons)
+        lon_rho, lat_rho = np.meshgrid(lons, lats)
+        ds = xr.Dataset(
+            {
+                "lon_rho": (("eta_rho", "xi_rho"), lon_rho),
+                "lat_rho": (("eta_rho", "xi_rho"), lat_rho),
+                "angle": (("eta_rho", "xi_rho"), np.zeros_like(lon_rho)),
+            },
+            coords={"eta_rho": np.arange(eta), "xi_rho": np.arange(xi)},
+        )
+        if with_mask:
+            ds["mask_rho"] = (("eta_rho", "xi_rho"), np.ones_like(lon_rho))
+        return ds
+
+    def test_basic_rho_grid(self):
+        ds = self.make_rho_grid(lons=[-10, -5, 0, 5, 10], lats=[50, 55])
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        assert "lat" in result and "lon" in result
+        assert np.allclose(result["lon"], ds.lon_rho)
+
+    def test_wrap_longitudes_to_minus180_180(self):
+        ds = self.make_rho_grid(lons=[190, 200], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        # longitudes >180 should wrap to -170, -160
+        expected = np.array([[-170, -160], [-170, -160]])
+        assert np.allclose(result["lon"].values, expected)
+
+    def test_convert_to_0_360_if_far_from_greenwich(self):
+        ds = self.make_rho_grid(lons=[-170, -160], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=False)
+        result = get_target_coords(grid)
+        # Should convert to 190, 200 since domain is far from Greenwich
+        expected = np.array([[190, 200], [190, 200]])
+        assert np.allclose(result["lon"].values, expected)
+        assert result["straddle"] is False
+
+    def test_close_to_greenwich_stays_minus180_180(self):
+        ds = self.make_rho_grid(lons=[-2, -1], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=False)
+        result = get_target_coords(grid)
+        # Should remain unchanged (-2, -1), not converted to 358, 359
+        expected = np.array([[-2, -1], [-2, -1]])
+        assert np.allclose(result["lon"].values, expected)
+        assert result["straddle"] is True
+
+    def test_includes_optional_fields(self):
+        ds = self.make_rho_grid(lons=[-10, -5], lats=[0, 1], with_mask=True)
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        assert result["mask"] is not None
+
+    def test_coarse_grid_selection(self):
+        lon = np.array([[190, 200]])
+        lat = np.array([[10, 10]])
+        ds = xr.Dataset(
+            {
+                "lon_coarse": (("eta_coarse", "xi_coarse"), lon),
+                "lat_coarse": (("eta_coarse", "xi_coarse"), lat),
+                "angle_coarse": (("eta_coarse", "xi_coarse"), np.zeros_like(lon)),
+                "mask_coarse": (("eta_coarse", "xi_coarse"), np.ones_like(lon)),
+            },
+            coords={"eta_coarse": [0], "xi_coarse": [0, 1]},
+        )
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid, use_coarse_grid=True)
+        # Should wrap longitudes to -170, -160
+        expected = np.array([[-170, -160]])
+        assert np.allclose(result["lon"].values, expected)
+        assert "mask" in result
 
 
 def test_interpolate_from_climatology(use_dask):

roms_tools/tests/test_setup/test_validation.py

@@ -1,15 +1,11 @@
-import os
 import shutil
+from collections.abc import Callable
+from pathlib import Path
 
 import pytest
 import xarray as xr
 
 
-def _get_fname(name):
-    dirname = os.path.dirname(__file__)
-    return os.path.join(dirname, "test_data", f"{name}.zarr")
-
-
 @pytest.mark.parametrize(
     "forcing_fixture",
     [
@@ -34,7 +30,11 @@ def _get_fname(name):
 # this test will not be run by default
 # to run it and overwrite the test data, invoke pytest as follows
 # pytest --overwrite=tidal_forcing --overwrite=boundary_forcing
-def test_save_results(forcing_fixture, request):
+def test_save_results(
+    forcing_fixture,
+    request: pytest.FixtureRequest,
+    get_test_data_path: Callable[[str], Path],
+) -> None:
     overwrite = request.config.getoption("--overwrite")
 
     # Skip the test if the fixture isn't marked for overwriting, unless 'all' is specified
@@ -42,10 +42,10 @@ def test_save_results(forcing_fixture, request):
         pytest.skip(f"Skipping overwrite for {forcing_fixture}")
 
     forcing = request.getfixturevalue(forcing_fixture)
-    fname = _get_fname(forcing_fixture)
+    fname = get_test_data_path(forcing_fixture)
 
     # Check if the Zarr directory exists and delete it if it does
-    if os.path.exists(fname):
+    if fname.exists():
         shutil.rmtree(fname)
 
     forcing.ds.to_zarr(fname)
@@ -72,8 +72,12 @@ def test_save_results(forcing_fixture, request):
         "river_forcing_no_climatology",
     ],
 )
-def test_check_results(forcing_fixture, request):
-    fname = _get_fname(forcing_fixture)
+def test_check_results(
+    forcing_fixture,
+    request: pytest.FixtureRequest,
+    get_test_data_path: Callable[[str], Path],
+) -> None:
+    fname = get_test_data_path(forcing_fixture)
     expected_forcing_ds = xr.open_zarr(fname, decode_timedelta=False)
     forcing = request.getfixturevalue(forcing_fixture)
 
@@ -83,6 +87,7 @@ def test_check_results(forcing_fixture, request):
     )
 
 
+@pytest.mark.use_dask
 @pytest.mark.parametrize(
     "forcing_fixture",
     [
@@ -97,11 +102,12 @@ def test_check_results(forcing_fixture, request):
         "bgc_boundary_forcing_from_climatology",
     ],
 )
-def test_dask_vs_no_dask(forcing_fixture, request, tmp_path, use_dask):
+def test_dask_vs_no_dask(
+    forcing_fixture: str,
+    request: pytest.FixtureRequest,
+    tmp_path: Path,
+) -> None:
     """Test comparing the forcing created with and without Dask on same platform."""
-    if not use_dask:
-        pytest.skip("Test only runs when --use_dask is specified")
-
     # Get the forcing with Dask
     forcing_with_dask = request.getfixturevalue(forcing_fixture)
 

roms_tools/tests/test_setup/utils.py

@@ -0,0 +1,71 @@
+from datetime import datetime
+from pathlib import Path
+
+from roms_tools import Grid, get_glorys_bounds
+
+try:
+    import copernicusmarine  # type: ignore
+except ImportError:
+    copernicusmarine = None
+
+
+def download_regional_and_bigger(
+    tmp_path: Path,
+    grid: Grid,
+    start_time: datetime,
+    variables: list[str] = ["thetao", "so", "uo", "vo", "zos"],
+) -> tuple[Path, Path]:
+    """
+    Helper: download minimal and slightly bigger GLORYS subsets.
+
+    Parameters
+    ----------
+    tmp_path : Path
+        Directory to store the downloaded NetCDF files.
+    grid : Grid
+        ROMS-Tools Grid object defining the target domain.
+    start_time : datetime
+        Start time of the requested subset.
+    variables : list[str]
+        What variables to download.
+
+    Returns
+    -------
+    Tuple[Path, Path]
+        Paths to the minimal and slightly bigger GLORYS subset files.
+    """
+    bounds = get_glorys_bounds(grid)
+
+    # minimal dataset
+    regional_file = tmp_path / "regional_GLORYS.nc"
+    copernicusmarine.subset(
+        dataset_id="cmems_mod_glo_phy_my_0.083deg_P1D-m",
+        variables=variables,
+        **bounds,
+        start_datetime=start_time,
+        end_datetime=start_time,
+        coordinates_selection_method="outside",
+        output_filename=str(regional_file),
+    )
+
+    # slightly bigger dataset
+    for key, delta in {
+        "minimum_latitude": -1,
+        "minimum_longitude": -1,
+        "maximum_latitude": +1,
+        "maximum_longitude": +1,
+    }.items():
+        bounds[key] += delta
+
+    bigger_regional_file = tmp_path / "bigger_regional_GLORYS.nc"
+    copernicusmarine.subset(
+        dataset_id="cmems_mod_glo_phy_my_0.083deg_P1D-m",
+        variables=variables,
+        **bounds,
+        start_datetime=start_time,
+        end_datetime=start_time,
+        coordinates_selection_method="outside",
+        output_filename=str(bigger_regional_file),
+    )
+
+    return regional_file, bigger_regional_file

roms_tools/tests/test_tiling/test_join.py

@@ -0,0 +1,241 @@
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import xarray as xr
+
+from roms_tools import Grid
+from roms_tools.tiling.join import (
+    _find_common_dims,
+    _find_transitions,
+    _infer_partition_layout_from_datasets,
+    join_netcdf,
+    open_partitions,
+)
+from roms_tools.tiling.partition import partition_netcdf
+
+
+@pytest.fixture
+def partitioned_grid_factory(
+    tmp_path, large_grid
+) -> Callable[[int, int], tuple[Grid, list[Path]]]:
+    """
+    A fixture factory that returns a function to generate partitioned files
+    with a configurable layout.
+    """
+
+    def _partitioned_files(np_xi: int, np_eta: int) -> tuple[Grid, list[Path]]:
+        partable_grid = large_grid
+        partable_grid.save(tmp_path / "test_grid.nc")
+        parted_files = partition_netcdf(
+            tmp_path / "test_grid.nc", np_xi=np_xi, np_eta=np_eta
+        )
+        return partable_grid, parted_files
+
+    return _partitioned_files
+
+
+@pytest.fixture
+def partitioned_ic_factory(
+    tmp_path, initial_conditions_on_large_grid
+) -> Callable[[int, int], tuple[Path, list[Path]]]:
+    def _partitioned_files(np_xi: int, np_eta: int) -> tuple[Path, list[Path]]:
+        whole_ics = initial_conditions_on_large_grid
+        whole_path = whole_ics.save(tmp_path / "test_ic.nc")[0]
+        parted_paths = partition_netcdf(whole_path, np_xi=np_xi, np_eta=np_eta)
+        return whole_path, parted_paths
+
+    return _partitioned_files
+
+
+class TestHelperFunctions:
+    def test_find_common_dims(self):
+        """Test _find_common_dims with different datasets and directions.
+
+        Tests for a valid common dimension, a valid dimension that is not common,
+        an invalid direction, and a ValueError when no common dimension is found.
+        """
+        # Create mock xarray.Dataset objects for testing
+        ds1 = xr.Dataset(coords={"xi_rho": [0], "xi_u": [0]})
+        ds2 = xr.Dataset(coords={"xi_rho": [0], "xi_u": [0]})
+        ds3 = xr.Dataset(coords={"xi_rho": [0], "xi_v": [0]})
+        datasets_common = [ds1, ds2]
+        datasets_not_common = [ds1, ds3]
+
+        # Test case with a common dimension ("xi_rho" and "xi_u")
+        # The function should find 'xi_rho' and 'xi_u' and return them in a list.
+        assert _find_common_dims("xi", datasets_common) == ["xi_rho", "xi_u"]
+
+        # Test case where a dimension is not common to all datasets
+        # The function should find only "xi_rho" as "xi_u" is not in ds3.
+        assert _find_common_dims("xi", datasets_not_common) == ["xi_rho"]
+
+        # Test case for an invalid direction, should raise a ValueError
+        with pytest.raises(ValueError, match="'direction' must be 'xi' or 'eta'"):
+            _find_common_dims("zeta", datasets_common)
+
+        # Test case where no common dimensions exist
+        ds_no_common1 = xr.Dataset(coords={"xi_rho": [0]})
+        ds_no_common2 = xr.Dataset(coords={"xi_v": [0]})
+        with pytest.raises(
+            ValueError, match="No common point found along direction xi"
+        ):
+            _find_common_dims("xi", [ds_no_common1, ds_no_common2])
+
+    def test_find_transitions(self):
+        """Test _find_transitions with various input lists.
+
+        Test cases include lists with no transitions, a single transition,
+        multiple transitions, and edge cases like empty and single-element lists.
+        """
+        # Test case with no transitions
+        assert _find_transitions([10, 10, 10, 10]) == []
+
+        # Test case with a single transition
+        assert _find_transitions([10, 10, 12, 12]) == [2]
+
+        # Test case with multiple transitions
+        assert _find_transitions([10, 12, 12, 14, 14, 14]) == [1, 3]
+
+        # Test case with transitions on every element
+        assert _find_transitions([10, 12, 14, 16]) == [1, 2, 3]
+
+        # Edge case: empty list
+        assert _find_transitions([]) == []
+
+        # Edge case: single-element list
+        assert _find_transitions([10]) == []
+
+    def test_infer_partition_layout_from_datasets(self):
+        """Test _infer_partition_layout_from_datasets with various layouts.
+
+        Tests include a single dataset, a 2x2 grid, a 4x1 grid (single row),
+        and a 1x4 grid (single column).
+        """
+        # Test case 1: Single dataset (1x1 partition)
+        ds1 = xr.Dataset(coords={"eta_rho": [0], "xi_rho": [0]})
+        assert _infer_partition_layout_from_datasets([ds1]) == (1, 1)
+
+        # Test case 2: 2x2 grid partition.
+        # The eta dimension will transition after the second dataset (np_xi=2).
+        ds_2x2_1 = xr.Dataset(coords={"eta_rho": [0] * 20, "xi_rho": [0] * 10})
+        ds_2x2_2 = xr.Dataset(coords={"eta_rho": [0] * 20, "xi_rho": [0] * 10})
+        ds_2x2_3 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 10})
+        ds_2x2_4 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 10})
+        datasets_2x2 = [ds_2x2_1, ds_2x2_2, ds_2x2_3, ds_2x2_4]
+        assert _infer_partition_layout_from_datasets(datasets_2x2) == (2, 2)
+
+        # Test case 3: 4x1 grid partition (single row).
+        # The eta dimension sizes are all the same, so no transition is detected.
+        # The function falls back to returning nd, 1.
+        ds_4x1_1 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 5})
+        ds_4x1_2 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 5})
+        ds_4x1_3 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 5})
+        ds_4x1_4 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 5})
+        datasets_4x1 = [ds_4x1_1, ds_4x1_2, ds_4x1_3, ds_4x1_4]
+        assert _infer_partition_layout_from_datasets(datasets_4x1) == (4, 1)
+
+        # Test case 4: 1x4 grid partition (single column).
+        # The xi dimension is partitioned, so the eta dimensions must change at every step.
+        ds_1x4_1 = xr.Dataset(coords={"eta_rho": [0] * 10, "xi_rho": [0] * 20})
+        ds_1x4_2 = xr.Dataset(coords={"eta_rho": [0] * 12, "xi_rho": [0] * 20})
+        ds_1x4_3 = xr.Dataset(coords={"eta_rho": [0] * 14, "xi_rho": [0] * 20})
+        ds_1x4_4 = xr.Dataset(coords={"eta_rho": [0] * 16, "xi_rho": [0] * 20})
+        datasets_1x4 = [ds_1x4_1, ds_1x4_2, ds_1x4_3, ds_1x4_4]
+        # In this case, `_find_transitions` for eta will find a transition at index 1, so np_xi=1.
+        # This will correctly return (1, 4).
+        assert _infer_partition_layout_from_datasets(datasets_1x4) == (1, 4)
+
+
+class TestJoinROMSData:
+    @pytest.mark.parametrize(
+        "np_xi, np_eta",
+        [
+            (1, 1),
+            (1, 6),
+            (6, 1),
+            (2, 2),
+            (3, 3),
+            (3, 4),
+            (4, 3),  # (12,24)
+            # # All possible:
+            # # Single-partition grid
+            # (1,1)
+            # # Single-row grids
+            # (2, 1), (3, 1), (4, 1), (6, 1), (12, 1),
+            # # Single-column grids
+            # (1, 2), (1, 3), (1, 4), (1, 6), (1, 8), (1, 12), (1, 24),
+            # # Multi-row, multi-column grids
+            # (2, 2), (2, 3), (2, 4), (2, 6), (2, 8), (2, 12), (2, 24),
+            # (3, 2), (3, 3), (3, 4), (3, 6), (3, 8), (3, 12), (3, 24),
+            # (4, 2), (4, 3), (4, 4), (4, 6), (4, 8), (4, 12), (4, 24),
+            # (6, 2), (6, 3), (6, 4), (6, 6), (6, 8), (6, 12), (6, 24),
+            # (12, 2), (12, 3), (12, 4), (12, 6), (12, 8), (12, 12), (12, 24)
+        ],
+    )
+    def test_open_grid_partitions(self, partitioned_grid_factory, np_xi, np_eta):
+        grid, partitions = partitioned_grid_factory(np_xi=np_xi, np_eta=np_eta)
+        joined_grid = open_partitions(partitions)
+
+        for v in grid.ds.variables:
+            assert (grid.ds[v].values == joined_grid[v].values).all(), (
+                f"{v} does not match in joined dataset"
+            )
+        assert grid.ds.attrs == joined_grid.attrs
+
+    def test_join_grid_netcdf(self, partitioned_grid_factory):
+        grid, partitions = partitioned_grid_factory(np_xi=3, np_eta=4)
+        joined_netcdf = join_netcdf(
+            partitions, output_path=partitions[0].parent / "joined_grid.nc"
+        )
+        assert joined_netcdf.exists()
+        joined_grid = xr.open_dataset(joined_netcdf)
+
+        for v in grid.ds.variables:
+            assert (grid.ds[v].values == joined_grid[v].values).all(), (
+                f"{v} does not match in joined dataset"
+            )
+        assert grid.ds.attrs == joined_grid.attrs
+
+    @pytest.mark.parametrize(
+        "np_xi, np_eta",
+        [
+            (1, 1),
+            (1, 6),
+            (6, 1),
+            (2, 2),
+            (3, 3),
+            (3, 4),
+            (4, 3),  # (12,24)
+        ],
+    )
+    def test_open_initial_condition_partitions(
+        self, partitioned_ic_factory, np_xi, np_eta
+    ):
+        whole_file, partitioned_files = partitioned_ic_factory(
+            np_xi=np_xi, np_eta=np_eta
+        )
+        joined_ics = open_partitions(partitioned_files)
+        whole_ics = xr.open_dataset(whole_file, decode_timedelta=True)
+
+        for v in whole_ics.variables:
+            assert (whole_ics[v].values == joined_ics[v].values).all(), (
+                f"{v} does not match in joined dataset: {joined_ics[v].values} vs {whole_ics[v].values}"
+            )
+        assert whole_ics.attrs == joined_ics.attrs
+
+    def test_join_initial_condition_netcdf(self, tmp_path, partitioned_ic_factory):
+        whole_file, partitioned_files = partitioned_ic_factory(np_xi=3, np_eta=4)
+        whole_ics = xr.open_dataset(whole_file, decode_timedelta=True)
+
+        joined_netcdf = join_netcdf(
+            partitioned_files, output_path=partitioned_files[0].parent / "joined_ics.nc"
+        )
+        assert joined_netcdf.exists()
+        joined_ics = xr.open_dataset(joined_netcdf, decode_timedelta=True)
+
+        for v in whole_ics.variables:
+            assert (whole_ics[v].values == joined_ics[v].values).all(), (
+                f"{v} does not match in joined dataset: {joined_ics[v].values} vs {whole_ics[v].values}"
+            )
+        assert whole_ics.attrs == joined_ics.attrs

roms_tools/tests/test_tiling/test_partition.py

@@ -297,3 +297,48 @@ class TestPartitionNetcdf:
         for expected_filepath in expected_filepath_list:
             assert expected_filepath.exists()
             expected_filepath.unlink()
+
+    def test_partition_netcdf_with_output_dir(self, grid, tmp_path):
+        # Save the input file
+        input_file = tmp_path / "input_grid.nc"
+        grid.save(input_file)
+
+        # Create a custom output directory
+        output_dir = tmp_path / "custom_output"
+        output_dir.mkdir()
+
+        saved_filenames = partition_netcdf(
+            input_file, np_eta=3, np_xi=5, output_dir=output_dir
+        )
+
+        base_name = input_file.stem  # "input_grid"
+        expected_filenames = [output_dir / f"{base_name}.{i:02d}.nc" for i in range(15)]
+
+        assert saved_filenames == expected_filenames
+
+        for f in expected_filenames:
+            assert f.exists()
+            f.unlink()
+
+    def test_partition_netcdf_multiple_files(self, grid, tmp_path):
+        # Create two test input files
+        file1 = tmp_path / "grid1.nc"
+        file2 = tmp_path / "grid2.nc"
+        grid.save(file1)
+        grid.save(file2)
+
+        # Run partitioning with 2x2 tiles on both files
+        saved_filenames = partition_netcdf([file1, file2], np_eta=3, np_xi=5)
+
+        # Expect 4 tiles per file → 8 total output files
+        expected_filepaths = []
+        for file in [file1, file2]:
+            base = file.with_suffix("")
+            expected_filepaths += [Path(f"{base}.{i:02d}.nc") for i in range(15)]
+
+        assert len(saved_filenames) == 30
+        assert saved_filenames == expected_filepaths
+
+        for path in expected_filepaths:
+            assert path.exists()
+            path.unlink()