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

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

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