roms-tools 3.1.2__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/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_utils.py

@@ -7,11 +7,13 @@ import pytest
 import xarray as xr
 
 from roms_tools.utils import (
-    _generate_focused_coordinate_range,
-    _has_copernicus,
-    _has_dask,
-    _has_gcsfs,
-    _load_data,
+    _path_list_from_input,
+    generate_focused_coordinate_range,
+    get_dask_chunks,
+    has_copernicus,
+    has_dask,
+    has_gcsfs,
+    load_data,
 )
 
 
@@ -25,66 +27,132 @@ from roms_tools.utils import (
     ],
 )
 def test_coordinate_range_monotonicity(min_val, max_val, center, sc, N):
-    centers, faces = _generate_focused_coordinate_range(
+    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"
 
 
+class TestPathListFromInput:
+    """A collection of tests for the _path_list_from_input function."""
+
+    # Test cases that don't require I/O
+    def test_list_of_strings(self):
+        """Test with a list of file paths as strings."""
+        files_list = ["path/to/file1.txt", "path/to/file2.txt"]
+        result = _path_list_from_input(files_list)
+        assert len(result) == 2
+        assert result[0] == Path("path/to/file1.txt")
+        assert result[1] == Path("path/to/file2.txt")
+
+    def test_list_of_path_objects(self):
+        """Test with a list of pathlib.Path objects."""
+        files_list = [Path("file_a.txt"), Path("file_b.txt")]
+        result = _path_list_from_input(files_list)
+        assert len(result) == 2
+        assert result[0] == Path("file_a.txt")
+        assert result[1] == Path("file_b.txt")
+
+    def test_single_path_object(self):
+        """Test with a single pathlib.Path object."""
+        file_path = Path("a_single_file.csv")
+        result = _path_list_from_input(file_path)
+        assert len(result) == 1
+        assert result[0] == file_path
+
+    def test_invalid_input_type_raises(self):
+        """Test that an invalid input type raises a TypeError."""
+        with pytest.raises(TypeError, match="'files' should be str, Path, or List"):
+            _path_list_from_input(123)
+
+    # Test cases that require I/O and `tmp_path`
+    def test_single_file_as_str(self, tmp_path):
+        """Test with a single file given as a string, requiring a file to exist."""
+        p = tmp_path / "test_file.txt"
+        p.touch()
+        result = _path_list_from_input(str(p))
+        assert len(result) == 1
+        assert result[0] == p
+
+    def test_wildcard_pattern(self, tmp_path, monkeypatch):
+        """Test with a wildcard pattern, requiring files to exist, using monkeypatch."""
+        # Setup
+        d = tmp_path / "data"
+        d.mkdir()
+        (d / "file1.csv").touch()
+        (d / "file2.csv").touch()
+        (d / "other_file.txt").touch()
+
+        # Action: Temporarily change the current working directory
+        monkeypatch.chdir(tmp_path)
+
+        result = _path_list_from_input("data/*.csv")
+
+        # Assertion
+        assert len(result) == 2
+        assert result[0].name == "file1.csv"
+        assert result[1].name == "file2.csv"
+
+    def test_non_matching_pattern_raises(self, tmp_path):
+        """Test that a non-matching pattern raises a FileNotFoundError."""
+        with pytest.raises(FileNotFoundError, match="No files matched"):
+            _path_list_from_input(str(tmp_path / "non_existent_file_*.txt"))
+
+
 def test_has_dask() -> None:
     """Verify that dask existence is correctly reported when found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=mock.MagicMock):
-        assert _has_dask()
+        assert has_dask()
 
 
 def test_has_dask_error_when_missing() -> None:
     """Verify that dask existence is correctly reported when not found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=None):
-        assert not _has_dask()
+        assert not has_dask()
 
 
 def test_has_gcfs() -> None:
     """Verify that GCFS existence is correctly reported when found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=mock.MagicMock):
-        assert _has_gcsfs()
+        assert has_gcsfs()
 
 
 def test_has_gcfs_error_when_missing() -> None:
     """Verify that GCFS existence is correctly reported when not found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=None):
-        assert not _has_gcsfs()
+        assert not has_gcsfs()
 
 
 def test_has_copernicus() -> None:
     """Verify that copernicus existence is correctly reported when found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=mock.MagicMock):
-        assert _has_copernicus()
+        assert has_copernicus()
 
 
 def test_has_copernicus_error_when_missing() -> None:
     """Verify that copernicus existence is correctly reported when not found."""
     with mock.patch("roms_tools.utils.find_spec", return_value=None):
-        assert not _has_copernicus()
+        assert not has_copernicus()
 
 
 def test_load_data_dask_not_found() -> None:
     """Verify that load data raises an exception when dask is requested and missing."""
     with (
-        mock.patch("roms_tools.utils._has_dask", return_value=False),
+        mock.patch("roms_tools.utils.has_dask", return_value=False),
         pytest.raises(RuntimeError),
     ):
-        _load_data("foo.zarr", {"a": "a"}, use_dask=True)
+        load_data("foo.zarr", {"a": "a"}, use_dask=True)
 
 
 def test_load_data_open_zarr_without_dask() -> None:
     """Verify that load data raises an exception when zarr is requested without dask."""
     with (
-        mock.patch("roms_tools.utils._has_dask", return_value=False),
+        mock.patch("roms_tools.utils.has_dask", return_value=False),
         pytest.raises(ValueError),
     ):
         # read_zarr should require use_dask to be True
-        _load_data("foo.zarr", {"a": ""}, use_dask=False, read_zarr=True)
+        load_data("foo.zarr", {"a": ""}, use_dask=False, read_zarr=True)
 
 
 @pytest.mark.parametrize(
@@ -111,7 +179,7 @@ def test_load_data_open_dataset(
         "roms_tools.utils.xr.open_dataset",
         wraps=xr.open_dataset,
     ) as fn_od:
-        ds = _load_data(
+        ds = load_data(
             ds_path,
             {"latitude": "latitude"},
             use_dask=False,
@@ -119,3 +187,57 @@ def test_load_data_open_dataset(
         assert fn_od.called
 
     assert expected_dim in ds.dims
+
+
+# test get_dask_chunks
+
+
+def test_latlon_default_chunks():
+    dim_names = {"latitude": "lat", "longitude": "lon"}
+    expected = {"lat": -1, "lon": -1}
+    result = get_dask_chunks(dim_names)
+    assert result == expected
+
+
+def test_latlon_with_depth_and_time():
+    dim_names = {"latitude": "lat", "longitude": "lon", "depth": "z", "time": "t"}
+    expected = {"lat": -1, "lon": -1, "z": -1, "t": 1}
+    result = get_dask_chunks(dim_names)
+    assert result == expected
+
+
+def test_latlon_with_time_chunking_false():
+    dim_names = {"latitude": "lat", "longitude": "lon", "time": "t"}
+    expected = {"lat": -1, "lon": -1}
+    result = get_dask_chunks(dim_names, time_chunking=False)
+    assert result == expected
+
+
+def test_roms_default_chunks():
+    dim_names = {}
+    expected_keys = {"eta_rho", "eta_v", "xi_rho", "xi_u", "s_rho"}
+    result = get_dask_chunks(dim_names)
+    assert set(result.keys()) == expected_keys
+    assert all(v == -1 for v in result.values())
+
+
+def test_roms_with_depth_and_time():
+    dim_names = {"depth": "s_rho", "time": "ocean_time"}
+    result = get_dask_chunks(dim_names)
+    # ROMS default keys + depth + time
+    expected_keys = {"eta_rho", "eta_v", "xi_rho", "xi_u", "s_rho", "ocean_time"}
+    assert set(result.keys()) == expected_keys
+    assert result["ocean_time"] == 1
+    assert result["s_rho"] == -1
+
+
+def test_roms_with_ntides():
+    dim_names = {"ntides": "nt"}
+    result = get_dask_chunks(dim_names)
+    assert result["nt"] == 1
+
+
+def test_time_chunking_false_roms():
+    dim_names = {"time": "ocean_time"}
+    result = get_dask_chunks(dim_names, time_chunking=False)
+    assert "ocean_time" not in result