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

roms_tools/tests/test_setup/test_datasets.py

@@ -2,6 +2,7 @@ import logging
 from collections import OrderedDict
 from datetime import datetime
 from pathlib import Path
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -11,11 +12,14 @@ from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import (
     CESMBGCDataset,
     Dataset,
+    ERA5ARCODataset,
     ERA5Correction,
     GLORYSDataset,
+    GLORYSDefaultDataset,
     RiverDataset,
     TPXODataset,
 )
+from roms_tools.setup.surface_forcing import DEFAULT_ERA5_ARCO_PATH
 
 
 @pytest.fixture
@@ -437,6 +441,78 @@ def test_era5_correction_choose_subdomain(use_dask):
     assert (data.ds["longitude"] == lons).all()
 
 
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading_without_dask() -> None:
+    """Verify that loading the default ERA5 dataset fails if use_dask is not True."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with pytest.raises(ValueError):
+        _ = ERA5ARCODataset(
+            filename=DEFAULT_ERA5_ARCO_PATH,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=False,
+        )
+
+
+@pytest.mark.skip("Temporary skip until memory consumption issue is addressed. # TODO")
+@pytest.mark.stream
+@pytest.mark.use_dask
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading() -> None:
+    """Verify the default ERA5 dataset is loaded correctly."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    ds = ERA5ARCODataset(
+        filename=DEFAULT_ERA5_ARCO_PATH,
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"uwnd", "vwnd", "swrad", "lwrad", "Tair", "rain"}
+    assert set(ds.var_names).issuperset(expected_vars)
+
+
+@pytest.mark.use_copernicus
+def test_default_glorys_dataset_loading_dask_not_installed() -> None:
+    """Verify that loading the default GLORYS dataset fails if dask is not available."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with (
+        pytest.raises(RuntimeError),
+        mock.patch("roms_tools.utils._has_dask", return_value=False),
+    ):
+        _ = GLORYSDefaultDataset(
+            filename=GLORYSDefaultDataset.dataset_name,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=True,
+        )
+
+
+@pytest.mark.stream
+@pytest.mark.use_copernicus
+@pytest.mark.use_dask
+def test_default_glorys_dataset_loading() -> None:
+    """Verify the default GLORYS dataset is loaded correctly."""
+    start_time = datetime(2012, 1, 1)
+    end_time = datetime(2013, 1, 1)
+
+    ds = GLORYSDefaultDataset(
+        filename=GLORYSDefaultDataset.dataset_name,
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"temp", "salt", "u", "v", "zeta"}
+    assert set(ds.var_names).issuperset(expected_vars)
+
+
 def test_data_concatenation(use_dask):
     fname = download_test_data("GLORYS_NA_2012.nc")
     data = GLORYSDataset(

roms_tools/tests/test_setup/test_grid.py

@@ -21,6 +21,11 @@ from roms_tools.constants import (
 from roms_tools.download import download_test_data
 from roms_tools.setup.topography import _compute_rfactor
 
+try:
+    import xesmf  # type: ignore
+except ImportError:
+    xesmf = None
+
 
 @pytest.fixture()
 def counter_clockwise_rotated_grid():
@@ -177,13 +182,18 @@ def test_successful_initialization_with_topography(grid_fixture, request):
     assert grid is not None
 
 
-def test_plot():
-    grid = Grid(
-        nx=20, ny=20, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=0
-    )
+def test_plot(grid_that_straddles_180_degree_meridian):
+    grid_that_straddles_180_degree_meridian.plot(with_dim_names=False)
+    grid_that_straddles_180_degree_meridian.plot(with_dim_names=True)
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_plot_along_lat_lon(grid_that_straddles_180_degree_meridian):
+    grid_that_straddles_180_degree_meridian.plot(lat=61)
+    grid_that_straddles_180_degree_meridian.plot(lon=180)
 
-    grid.plot(with_dim_names=False)
-    grid.plot(with_dim_names=True)
+    with pytest.raises(ValueError, match="Specify either `lat` or `lon`, not both"):
+        grid_that_straddles_180_degree_meridian.plot(lat=61, lon=180)
 
 
 def test_save(tmp_path):

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -187,12 +187,12 @@ def _test_successful_initialization(
         if coarse_grid_mode == "always":
             assert sfc_forcing.use_coarse_grid
             assert (
-                "Data will be interpolated onto grid coarsened by factor 2."
+                "Data will be interpolated onto the grid coarsened by factor 2."
                 in caplog.text
             )
         elif coarse_grid_mode == "never":
             assert not sfc_forcing.use_coarse_grid
-            assert "Data will be interpolated onto fine grid." in caplog.text
+            assert "Data will be interpolated onto the fine grid." in caplog.text
 
         assert isinstance(sfc_forcing.ds, xr.Dataset)
         assert "uwnd" in sfc_forcing.ds
@@ -902,7 +902,9 @@ def test_from_yaml_missing_surface_forcing(tmp_path, use_dask):
         yaml_filepath.unlink()
 
 
+@pytest.mark.skip("Temporary skip until memory consumption issue is addressed. # TODO")
 @pytest.mark.stream
+@pytest.mark.use_dask
 def test_surface_forcing_arco(surface_forcing_arco, tmp_path):
     """One big integration test for cloud-based ERA5 data because the streaming takes a
     long time.
@@ -932,3 +934,25 @@ def test_surface_forcing_arco(surface_forcing_arco, tmp_path):
     yaml_filepath.unlink()
     Path(expected_filepath1).unlink()
     Path(expected_filepath2).unlink()
+
+
+@pytest.mark.skip("Temporary skip until memory consumption issue is addressed. # TODO")
+@pytest.mark.stream
+@pytest.mark.use_dask
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading(small_grid: Grid) -> None:
+    """Verify the default ERA5 dataset is loaded when a path is not provided."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    sf = SurfaceForcing(
+        grid=small_grid,
+        source={"name": "ERA5"},
+        type="physics",
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"uwnd", "vwnd", "swrad", "lwrad", "Tair", "rain"}
+    assert set(sf.ds.data_vars).issuperset(expected_vars)

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_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()

roms_tools/tests/test_utils.py

@@ -1,7 +1,18 @@
+from collections.abc import Callable
+from pathlib import Path
+from unittest import mock
+
 import numpy as np
 import pytest
+import xarray as xr
 
-from roms_tools.utils import _generate_focused_coordinate_range
+from roms_tools.utils import (
+    _generate_focused_coordinate_range,
+    _has_copernicus,
+    _has_dask,
+    _has_gcsfs,
+    _load_data,
+)
 
 
 @pytest.mark.parametrize(
@@ -19,3 +30,92 @@ def test_coordinate_range_monotonicity(min_val, max_val, center, sc, 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"
+
+
+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()
+
+
+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()
+
+
+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()
+
+
+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()
+
+
+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()
+
+
+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()
+
+
+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),
+        pytest.raises(RuntimeError),
+    ):
+        _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),
+        pytest.raises(ValueError),
+    ):
+        # read_zarr should require use_dask to be True
+        _load_data("foo.zarr", {"a": ""}, use_dask=False, read_zarr=True)
+
+
+@pytest.mark.parametrize(
+    ("dataset_name", "expected_dim"),
+    [
+        ("surface_forcing", "time"),
+        ("bgc_surface_forcing", "time"),
+        ("tidal_forcing", "eta_rho"),
+        ("coarse_surface_forcing", "eta_rho"),
+    ],
+)
+def test_load_data_open_dataset(
+    dataset_name: str,
+    expected_dim: str,
+    get_test_data_path: Callable[[str], Path],
+) -> None:
+    """Verify that a zarr file is correctly loaded when not using Dask.
+
+    This must use xr.open_dataset
+    """
+    ds_path = get_test_data_path(dataset_name)
+
+    with mock.patch(
+        "roms_tools.utils.xr.open_dataset",
+        wraps=xr.open_dataset,
+    ) as fn_od:
+        ds = _load_data(
+            ds_path,
+            {"latitude": "latitude"},
+            use_dask=False,
+        )
+        assert fn_od.called
+
+    assert expected_dim in ds.dims

roms_tools/tiling/partition.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from numbers import Integral
 from pathlib import Path
 
@@ -296,20 +297,21 @@ def partition(
 
 
 def partition_netcdf(
-    filepath: str | Path,
+    filepath: str | Path | Sequence[str | Path],
     np_eta: int = 1,
     np_xi: int = 1,
+    output_dir: str | Path | None = None,
     include_coarse_dims: bool = True,
-) -> None:
-    """Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
+) -> list[Path]:
+    """Partition one or more ROMS NetCDF files 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.
+    This function divides each dataset 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.
+    filepath : str | Path | Sequence[str | Path]
+        A path or list of paths to input NetCDF files.
 
     np_eta : int, optional
         The number of partitions along the `eta` direction. Must be a positive integer. Default is 1.
@@ -317,6 +319,10 @@ def partition_netcdf(
     np_xi : int, optional
         The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
 
+    output_dir : str | Path | None, optional
+        Directory or base path to save partitioned files.
+        If None, files are saved alongside the input file.
+
     include_coarse_dims : bool, optional
         Whether to include coarse grid dimensions (`eta_coarse`, `xi_coarse`) in the partitioning.
         If False, these dimensions will not be split. Relevant if none of the coarse resolution variables are actually used by ROMS.
@@ -324,31 +330,39 @@ def partition_netcdf(
 
     Returns
     -------
-    List[Path]
+    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"), decode_timedelta=False)
-
-    # Partition the dataset
-    file_numbers, partitioned_datasets = partition(
-        ds, np_eta=np_eta, np_xi=np_xi, include_coarse_dims=include_coarse_dims
-    )
-
-    # Generate paths to the partitioned files
-    base_filepath = filepath.with_suffix("")
-    ndigits = len(str(max(np.array(file_numbers))))
-    paths_to_partitioned_files = [
-        Path(f"{base_filepath}.{file_number:0{ndigits}d}")
-        for file_number in file_numbers
-    ]
+    if isinstance(filepath, str | Path):
+        filepaths = [Path(filepath)]
+    else:
+        filepaths = [Path(fp) for fp in filepath]
+
+    all_saved_filenames = []
+
+    for fp in filepaths:
+        input_file = fp.with_suffix(".nc")
+        ds = xr.open_dataset(input_file, decode_timedelta=False)
+
+        file_numbers, partitioned_datasets = partition(
+            ds, np_eta=np_eta, np_xi=np_xi, include_coarse_dims=include_coarse_dims
+        )
+
+        if output_dir:
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+            base_filepath = output_dir / fp.stem
+        else:
+            base_filepath = fp.with_suffix("")
 
-    # Save the partitioned datasets to files
-    saved_filenames = save_datasets(
-        partitioned_datasets, paths_to_partitioned_files, verbose=False
-    )
+        ndigits = len(str(max(file_numbers)))
+        paths_to_partitioned_files = [
+            Path(f"{base_filepath}.{num:0{ndigits}d}") for num in file_numbers
+        ]
+
+        saved = save_datasets(
+            partitioned_datasets, paths_to_partitioned_files, verbose=False
+        )
+        all_saved_filenames.extend(saved)
 
-    return saved_filenames
+    return all_saved_filenames