roms-tools 2.6.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

roms_tools/tests/test_setup/test_tides.py

@@ -1,18 +1,91 @@
-import pytest
-from roms_tools import Grid, TidalForcing
-import xarray as xr
-from roms_tools.download import download_test_data
 import textwrap
 from pathlib import Path
+
+import pytest
+import xarray as xr
+
 from conftest import calculate_data_hash
+from roms_tools import Grid, TidalForcing
+from roms_tools.download import download_test_data
 
 
-@pytest.fixture
-def grid_that_lies_within_bounds_of_regional_tpxo_data():
+@pytest.fixture(scope="session")
+def tidal_forcing_9v5a(use_dask):
+
     grid = Grid(
         nx=3, ny=3, size_x=1500, size_y=1500, center_lon=235, center_lat=25, rot=-20
     )
-    return grid
+    fname_grid = Path(download_test_data("regional_grid_tpxo9v5a.nc"))
+    fname_h = Path(download_test_data("regional_h_tpxo9v5a.nc"))
+    fname_u = Path(download_test_data("regional_u_tpxo9v5a.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
+
+    return TidalForcing(
+        grid=grid,
+        source={"name": "TPXO", "path": fname_dict},
+        ntides=1,
+        use_dask=use_dask,
+    )
+
+
+@pytest.fixture(scope="session")
+def tidal_forcing_10v2a(use_dask):
+
+    grid = Grid(
+        nx=3, ny=3, size_x=1500, size_y=1500, center_lon=235, center_lat=25, rot=-20
+    )
+    fname_grid = Path(download_test_data("regional_grid_tpxo10v2a.nc"))
+    fname_h = Path(download_test_data("regional_h_tpxo10v2a.nc"))
+    fname_u = Path(download_test_data("regional_u_tpxo10v2a.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
+
+    return TidalForcing(
+        grid=grid,
+        source={"name": "TPXO", "path": fname_dict},
+        ntides=1,
+        use_dask=use_dask,
+    )
+
+
+@pytest.fixture(scope="session")
+def tidal_forcing_from_global_data(use_dask):
+
+    grid = Grid(
+        nx=3, ny=3, size_x=1800, size_y=1500, center_lon=235, center_lat=25, rot=-20
+    )
+
+    fname_grid = Path(download_test_data("global_grid_tpxo10.v2.nc"))
+    fname_h = Path(download_test_data("global_h_tpxo10.v2.nc"))
+    fname_u = Path(download_test_data("global_u_tpxo10.v2.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
+
+    return TidalForcing(
+        grid=grid,
+        source={"name": "TPXO", "path": fname_dict},
+        ntides=1,
+        use_dask=use_dask,
+    )
+
+
+@pytest.mark.parametrize(
+    "tidal_forcing_fixture",
+    ["tidal_forcing", "tidal_forcing_9v5a", "tidal_forcing_10v2a"],  # 10v2
+)
+def test_successful_initialization(tidal_forcing_fixture, request, use_dask):
+    tidal_forcing = request.getfixturevalue(tidal_forcing_fixture)
+
+    assert isinstance(tidal_forcing.ds, xr.Dataset)
+    assert "omega" in tidal_forcing.ds
+    assert "ssh_Re" in tidal_forcing.ds
+    assert "ssh_Im" in tidal_forcing.ds
+    assert "pot_Re" in tidal_forcing.ds
+    assert "pot_Im" in tidal_forcing.ds
+    assert "u_Re" in tidal_forcing.ds
+    assert "u_Im" in tidal_forcing.ds
+    assert "v_Re" in tidal_forcing.ds
+    assert "v_Im" in tidal_forcing.ds
+
+    assert tidal_forcing.ntides == 1
 
 
 @pytest.fixture
@@ -59,7 +132,6 @@ def grid_that_straddles_180_degree_meridian():
 @pytest.mark.parametrize(
     "grid_fixture",
     [
-        "grid_that_lies_within_bounds_of_regional_tpxo_data",
         "grid_that_is_out_of_bounds_of_regional_tpxo_data",
         "grid_that_straddles_dateline",
         "grid_that_straddles_180_degree_meridian",
@@ -67,39 +139,17 @@ def grid_that_straddles_180_degree_meridian():
 )
 def test_successful_initialization_with_global_data(grid_fixture, request, use_dask):
 
-    fname = Path(download_test_data("TPXO_global_test_data.nc"))
+    fname_grid = Path(download_test_data("global_grid_tpxo10.v2.nc"))
+    fname_h = Path(download_test_data("global_h_tpxo10.v2.nc"))
+    fname_u = Path(download_test_data("global_u_tpxo10.v2.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
 
     grid = request.getfixturevalue(grid_fixture)
 
     tidal_forcing = TidalForcing(
-        grid=grid, source={"name": "TPXO", "path": fname}, ntides=2, use_dask=use_dask
-    )
-
-    assert isinstance(tidal_forcing.ds, xr.Dataset)
-    assert "omega" in tidal_forcing.ds
-    assert "ssh_Re" in tidal_forcing.ds
-    assert "ssh_Im" in tidal_forcing.ds
-    assert "pot_Re" in tidal_forcing.ds
-    assert "pot_Im" in tidal_forcing.ds
-    assert "u_Re" in tidal_forcing.ds
-    assert "u_Im" in tidal_forcing.ds
-    assert "v_Re" in tidal_forcing.ds
-    assert "v_Im" in tidal_forcing.ds
-
-    assert tidal_forcing.source == {"name": "TPXO", "path": fname}
-    assert tidal_forcing.ntides == 2
-
-
-def test_successful_initialization_with_regional_data(
-    grid_that_lies_within_bounds_of_regional_tpxo_data, use_dask
-):
-
-    fname = Path(download_test_data("TPXO_regional_test_data.nc"))
-
-    tidal_forcing = TidalForcing(
-        grid=grid_that_lies_within_bounds_of_regional_tpxo_data,
-        source={"name": "TPXO", "path": fname},
-        ntides=10,
+        grid=grid,
+        source={"name": "TPXO", "path": fname_dict},
+        ntides=1,
         use_dask=use_dask,
     )
 
@@ -114,20 +164,22 @@ def test_successful_initialization_with_regional_data(
     assert "v_Re" in tidal_forcing.ds
     assert "v_Im" in tidal_forcing.ds
 
-    assert tidal_forcing.source == {"name": "TPXO", "path": fname}
-    assert tidal_forcing.ntides == 10
+    assert tidal_forcing.ntides == 1
 
 
 def test_unsuccessful_initialization_with_regional_data_due_to_nans(
     grid_that_is_out_of_bounds_of_regional_tpxo_data, use_dask
 ):
 
-    fname = Path(download_test_data("TPXO_regional_test_data.nc"))
+    fname_grid = Path(download_test_data("regional_grid_tpxo10v2.nc"))
+    fname_h = Path(download_test_data("regional_h_tpxo10v2.nc"))
+    fname_u = Path(download_test_data("regional_u_tpxo10v2.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
 
     with pytest.raises(ValueError, match="NaN values found"):
         TidalForcing(
             grid=grid_that_is_out_of_bounds_of_regional_tpxo_data,
-            source={"name": "TPXO", "path": fname},
+            source={"name": "TPXO", "path": fname_dict},
             ntides=10,
             use_dask=use_dask,
         )
@@ -141,7 +193,10 @@ def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
     grid_fixture, request, use_dask
 ):
 
-    fname = Path(download_test_data("TPXO_regional_test_data.nc"))
+    fname_grid = Path(download_test_data("regional_grid_tpxo10v2.nc"))
+    fname_h = Path(download_test_data("regional_h_tpxo10v2.nc"))
+    fname_u = Path(download_test_data("regional_u_tpxo10v2.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
 
     grid = request.getfixturevalue(grid_fixture)
 
@@ -150,7 +205,7 @@ def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
     ):
         TidalForcing(
             grid=grid,
-            source={"name": "TPXO", "path": fname},
+            source={"name": "TPXO", "path": fname_dict},
             ntides=10,
             use_dask=use_dask,
         )
@@ -158,12 +213,15 @@ def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
 
 def test_insufficient_number_of_consituents(grid_that_straddles_dateline, use_dask):
 
-    fname = Path(download_test_data("TPXO_global_test_data.nc"))
+    fname_grid = Path(download_test_data("global_grid_tpxo10.v2.nc"))
+    fname_h = Path(download_test_data("global_h_tpxo10.v2.nc"))
+    fname_u = Path(download_test_data("global_u_tpxo10.v2.nc"))
+    fname_dict = {"grid": fname_grid, "h": fname_h, "u": fname_u}
 
-    with pytest.raises(ValueError, match="The dataset contains fewer"):
+    with pytest.raises(ValueError, match="The dataset contains tidal"):
         TidalForcing(
             grid=grid_that_straddles_dateline,
-            source={"name": "TPXO", "path": fname},
+            source={"name": "TPXO", "path": fname_dict},
             ntides=10,
             use_dask=use_dask,
         )
@@ -194,10 +252,16 @@ def test_tidal_forcing_save(tidal_forcing, tmp_path):
             filepath.unlink()
 
 
-def test_roundtrip_yaml(tidal_forcing, tmp_path, use_dask):
+@pytest.mark.parametrize(
+    "tidal_forcing_fixture",
+    ["tidal_forcing", "tidal_forcing_from_global_data"],
+)
+def test_roundtrip_yaml(tidal_forcing_fixture, tmp_path, use_dask, request):
     """Test that creating a TidalForcing object, saving its parameters to yaml file, and
     re-opening yaml file creates the same object."""
 
+    tidal_forcing = request.getfixturevalue(tidal_forcing_fixture)
+
     # Create a temporary filepath using the tmp_path fixture
     file_str = "test_yaml"
     for filepath in [
@@ -215,15 +279,18 @@ def test_roundtrip_yaml(tidal_forcing, tmp_path, use_dask):
         filepath.unlink()
 
 
-def test_files_have_same_hash(tidal_forcing, tmp_path, use_dask):
+@pytest.mark.parametrize(
+    "tidal_forcing_fixture",
+    ["tidal_forcing", "tidal_forcing_from_global_data"],
+)
+def test_files_have_same_hash(tidal_forcing_fixture, tmp_path, use_dask, request):
+
+    tidal_forcing = request.getfixturevalue(tidal_forcing_fixture)
 
     yaml_filepath = tmp_path / "test_yaml.yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
 
-    print(yaml_filepath)
-    print(filepath1)
-    print(filepath2)
     tidal_forcing.to_yaml(yaml_filepath)
     tidal_forcing.save(filepath1)
     tidal_forcing_from_file = TidalForcing.from_yaml(yaml_filepath, use_dask=use_dask)

roms_tools/tests/test_setup/test_utils.py

@@ -1,12 +1,16 @@
-from roms_tools import Grid, BoundaryForcing
-from roms_tools.setup.utils import interpolate_from_climatology
-from roms_tools.setup.datasets import ERA5Correction
-from roms_tools.download import download_test_data
-import xarray as xr
-import pytest
 from datetime import datetime
 from pathlib import Path
 
+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,
+)
+
 
 def test_interpolate_from_climatology(use_dask):
 

roms_tools/tests/test_setup/test_validation.py

@@ -1,6 +1,7 @@
-import pytest
 import os
 import shutil
+
+import pytest
 import xarray as xr
 
 
@@ -78,7 +79,10 @@ def test_check_results(forcing_fixture, request):
     expected_forcing_ds = xr.open_zarr(fname, decode_timedelta=False)
     forcing = request.getfixturevalue(forcing_fixture)
 
-    xr.testing.assert_allclose(forcing.ds, expected_forcing_ds, rtol=1.0e-5, atol=1e-08)
+    # Set small tolerance because some fields like NOx, NHy have values order 1e-12
+    xr.testing.assert_allclose(
+        forcing.ds, expected_forcing_ds, rtol=1.0e-12, atol=1e-13
+    )
 
 
 @pytest.mark.parametrize(

roms_tools/tests/test_tiling/test_partition.py

@@ -1,9 +1,10 @@
-import pytest
 from pathlib import Path
+
+import pytest
 import xarray.testing as xrt
 
-from roms_tools.tiling.partition import partition, partition_netcdf
 from roms_tools import Grid
+from roms_tools.tiling.partition import partition, partition_netcdf
 
 
 @pytest.fixture
@@ -231,6 +232,19 @@ class TestPartitionGrid:
         with pytest.raises(ValueError, match="cannot be evenly divided"):
             partition(grid.ds, np_eta=4, np_xi=1)
 
+        with pytest.raises(ValueError, match="cannot be evenly divided"):
+            partition(grid.ds, np_eta=10, np_xi=10)
+
+    def test_skip_coarse_dims(self, grid):
+        """Test that coarse dimensions remain unchanged when excluded from
+        partitioning."""
+        _, partitioned_datasets = partition(
+            grid.ds, np_eta=10, np_xi=10, include_coarse_dims=False
+        )
+        for ds in partitioned_datasets:
+            assert ds.sizes["eta_coarse"] == grid.ds.sizes["eta_coarse"]
+            assert ds.sizes["xi_coarse"] == grid.ds.sizes["xi_coarse"]
+
 
 class TestPartitionMissingDims:
     def test_partition_missing_dims(self, grid):
@@ -271,11 +285,11 @@ class TestPartitionNetcdf:
         filepath = tmp_path / "test_grid.nc"
         grid.save(filepath)
 
-        saved_filenames = partition_netcdf(filepath, np_eta=3, np_xi=3)
+        saved_filenames = partition_netcdf(filepath, np_eta=3, np_xi=5)
 
         filepath_str = str(filepath.with_suffix(""))
         expected_filepath_list = [
-            Path(filepath_str + f".{index}.nc") for index in range(9)
+            Path(filepath_str + f".{index:02d}.nc") for index in range(15)
         ]
 
         assert saved_filenames == expected_filepath_list

roms_tools/tests/test_utils.py

@@ -0,0 +1,21 @@
+import numpy as np
+import pytest
+
+from roms_tools.utils import _generate_focused_coordinate_range
+
+
+@pytest.mark.parametrize(
+    "min_val, max_val, center, sc, N",
+    [
+        (-20.0, 5.5, -3.1, 1.0, 100),
+        (100.0, 200.0, 150.0, 30.0, 100),
+        (0.0, 2000.0, 150.0, 0.0, 100),
+        (0.0, 2000.0, 150.0, 30.0, 100),
+    ],
+)
+def test_coordinate_range_monotonicity(min_val, max_val, center, sc, N):
+    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"

roms_tools/tests/test_vertical_coordinate.py

@@ -1,12 +1,13 @@
-import pytest
 import numpy as np
+import pytest
 import xarray as xr
 
+from roms_tools.constants import UPPER_BOUND_THETA_B, UPPER_BOUND_THETA_S
 from roms_tools.vertical_coordinate import (
     compute_cs,
-    sigma_stretch,
     compute_depth,
     compute_depth_coordinates,
+    sigma_stretch,
 )
 
 
@@ -17,11 +18,15 @@ def test_compute_cs():
     assert cs.shape == sigma.shape
     assert np.all(cs <= 0) and np.all(cs >= -1)
 
-    with pytest.raises(ValueError, match="theta_s must be between 0 and 10"):
-        compute_cs(sigma, 15, 2)
+    with pytest.raises(
+        ValueError, match=f"theta_s must be between 0 and {UPPER_BOUND_THETA_S}"
+    ):
+        compute_cs(sigma, UPPER_BOUND_THETA_S + 1, theta_b)
 
-    with pytest.raises(ValueError, match="theta_b must be between 0 and 4"):
-        compute_cs(sigma, 5, 5)
+    with pytest.raises(
+        ValueError, match=f"theta_b must be between 0 and {UPPER_BOUND_THETA_B}"
+    ):
+        compute_cs(sigma, theta_s, UPPER_BOUND_THETA_B + 1)
 
 
 def test_sigma_stretch():

roms_tools/tiling/partition.py

@@ -1,14 +1,15 @@
 from numbers import Integral
+from pathlib import Path
+from typing import Union
 
 import numpy as np
 import xarray as xr
-from typing import Union
-from pathlib import Path
+
 from roms_tools.utils import save_datasets
 
 
 def partition(
-    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1
+    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1, include_coarse_dims: bool = True
 ) -> tuple[list[int], list[xr.Dataset]]:
     """Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial
     tiles.
@@ -29,6 +30,11 @@ def partition(
     np_xi : int, optional
         The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
 
+    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.
+        Default is True.
+
     Returns
     -------
     tuple[list[int], list[xr.Dataset]]
@@ -66,17 +72,18 @@ def partition(
     ):
         raise ValueError("np_eta and np_xi must be positive integers")
 
-    partitionable_dims_maybe_present = [
+    base_dims = [
         "eta_rho",
         "xi_rho",
         "eta_v",
         "xi_u",
         "eta_psi",
         "xi_psi",
-        "eta_coarse",
-        "xi_coarse",
     ]
-    dims_to_partition = [d for d in partitionable_dims_maybe_present if d in ds.dims]
+    coarse_dims = ["eta_coarse", "xi_coarse"]
+    partitionable_dims = base_dims + coarse_dims if include_coarse_dims else base_dims
+
+    dims_to_partition = [d for d in partitionable_dims if d in ds.dims]
 
     # if eta is periodic there are no ghost cells along those dimensions
     if "eta_v" in ds.sizes and ds.sizes["eta_rho"] == ds.sizes["eta_v"]:
@@ -291,7 +298,10 @@ def partition(
 
 
 def partition_netcdf(
-    filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
+    filepath: Union[str, Path],
+    np_eta: int = 1,
+    np_xi: int = 1,
+    include_coarse_dims: bool = True,
 ) -> None:
     """Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
 
@@ -309,6 +319,11 @@ def partition_netcdf(
     np_xi : int, optional
         The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
 
+    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.
+        Default is True.
+
     Returns
     -------
     List[Path]
@@ -319,15 +334,19 @@ def partition_netcdf(
     filepath = Path(filepath)
 
     # Open the dataset
-    ds = xr.open_dataset(filepath.with_suffix(".nc"))
+    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)
+    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}") for file_number in file_numbers
+        Path(f"{base_filepath}.{file_number:0{ndigits}d}")
+        for file_number in file_numbers
     ]
 
     # Save the partitioned datasets to files