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

roms_tools/tests/test_setup/test_river_forcing.py

@@ -9,6 +9,7 @@ import xarray as xr
 
 from conftest import calculate_file_hash
 from roms_tools import Grid, RiverForcing
+from roms_tools.constants import MAX_DISTINCT_COLORS
 
 
 @pytest.fixture
@@ -57,6 +58,29 @@ def river_forcing_for_grid_that_straddles_dateline():
     )
 
 
+@pytest.fixture
+def river_forcing_for_gulf_of_mexico():
+    """Fixture for creating a RiverForcing object for the Gulf of Mexico with 45 rivers."""
+    grid = Grid(
+        nx=20,
+        ny=15,
+        size_x=2000,
+        size_y=1500,
+        center_lon=-89,
+        center_lat=24,
+        rot=0,
+        N=3,
+    )
+    start_time = datetime(2012, 1, 1)
+    end_time = datetime(2012, 1, 31)
+
+    return RiverForcing(
+        grid=grid,
+        start_time=start_time,
+        end_time=end_time,
+    )
+
+
 @pytest.fixture
 def single_cell_indices():
     # These are the indices that the `river_forcing` fixture generates automatically.
@@ -247,13 +271,46 @@ class TestRiverForcingGeneral:
             )
 
     def test_river_forcing_plot(self, river_forcing_with_bgc):
-        """Test plot method."""
+        """Test plot methods with and without specifying river_names."""
+        river_names = list(river_forcing_with_bgc.indices.keys())[0:2]
+
+        # Test plot_locations
         river_forcing_with_bgc.plot_locations()
-        river_forcing_with_bgc.plot("river_volume")
-        river_forcing_with_bgc.plot("river_temp")
-        river_forcing_with_bgc.plot("river_salt")
-        river_forcing_with_bgc.plot("river_ALK")
-        river_forcing_with_bgc.plot("river_PO4")
+        river_forcing_with_bgc.plot_locations(river_names=river_names)
+
+        # Fields to test
+        variables = [
+            "river_volume",
+            "river_temp",
+            "river_salt",
+            "river_ALK",
+            "river_PO4",
+        ]
+
+        for var in variables:
+            river_forcing_with_bgc.plot(var)
+            river_forcing_with_bgc.plot(var, river_names=river_names)
+
+    def test_plot_max_releases(self, caplog, river_forcing_for_gulf_of_mexico):
+        river_names = list(river_forcing_for_gulf_of_mexico.indices.keys())
+
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            river_forcing_for_gulf_of_mexico.plot_locations()
+        assert any(
+            f"Only the first {MAX_DISTINCT_COLORS} rivers will be plotted" in message
+            for message in caplog.messages
+        )
+
+        with caplog.at_level("WARNING"):
+            river_forcing_for_gulf_of_mexico.plot(
+                "river_volume", river_names=river_names
+            )
+
+        assert any(
+            f"Only the first {MAX_DISTINCT_COLORS} rivers will be plotted" in message
+            for message in caplog.messages
+        )
 
     @pytest.mark.parametrize(
         "river_forcing_fixture",

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_utils.py

@@ -1,3 +1,4 @@
+import logging
 from datetime import datetime
 from pathlib import Path
 
@@ -7,9 +8,7 @@ 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,
-)
+from roms_tools.setup.utils import interpolate_from_climatology, validate_names
 
 
 def test_interpolate_from_climatology(use_dask):
@@ -71,3 +70,53 @@ def test_roundtrip_yaml(
 
         filepath = Path(filepath)
         filepath.unlink()
+
+
+# test validate_names function
+
+VALID_NAMES = ["a", "b", "c", "d"]
+SENTINEL = "ALL"
+MAX_TO_PLOT = 3
+
+
+def test_valid_names_no_truncation():
+    names = ["a", "b"]
+    result = validate_names(names, VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="test")
+    assert result == names
+
+
+def test_valid_names_with_truncation(caplog):
+    names = ["a", "b", "c", "d"]
+    with caplog.at_level(logging.WARNING):
+        result = validate_names(
+            names, VALID_NAMES, SENTINEL, max_to_plot=2, label="test"
+        )
+        assert result == ["a", "b"]
+        assert "Only the first 2 tests will be plotted" in caplog.text
+
+
+def test_include_all_sentinel():
+    result = validate_names(SENTINEL, VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="test")
+    assert result == VALID_NAMES[:MAX_TO_PLOT]
+
+
+def test_invalid_name_raises():
+    with pytest.raises(ValueError, match="Invalid tests: z"):
+        validate_names(["a", "z"], VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="test")
+
+
+def test_non_list_input_raises():
+    with pytest.raises(ValueError, match="`test_names` should be a list of strings."):
+        validate_names("a", VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="test")
+
+
+def test_non_string_elements_in_list_raises():
+    with pytest.raises(
+        ValueError, match="All elements in `test_names` must be strings."
+    ):
+        validate_names(["a", 2], VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="test")
+
+
+def test_custom_label_in_errors():
+    with pytest.raises(ValueError, match="Invalid foozs: z"):
+        validate_names(["z"], VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="fooz")

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