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

roms_tools/tests/test_setup/test_tides.py

@@ -1,15 +1,16 @@
-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(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
     )
@@ -28,7 +29,6 @@ def tidal_forcing_9v5a(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
     )
@@ -47,7 +47,6 @@ def tidal_forcing_10v2a(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
     )
@@ -113,7 +112,6 @@ def grid_that_straddles_dateline():
 @pytest.fixture
 def grid_that_straddles_180_degree_meridian():
     """Fixture for creating a domain that straddles 180 degree meridian."""
-
     grid = Grid(
         nx=5,
         ny=5,
@@ -136,7 +134,6 @@ def grid_that_straddles_180_degree_meridian():
     ],
 )
 def test_successful_initialization_with_global_data(grid_fixture, request, use_dask):
-
     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"))
@@ -168,7 +165,6 @@ def test_successful_initialization_with_global_data(grid_fixture, request, use_d
 def test_unsuccessful_initialization_with_regional_data_due_to_nans(
     grid_that_is_out_of_bounds_of_regional_tpxo_data, use_dask
 ):
-
     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"))
@@ -190,7 +186,6 @@ def test_unsuccessful_initialization_with_regional_data_due_to_nans(
 def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
     grid_fixture, request, use_dask
 ):
-
     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"))
@@ -210,7 +205,6 @@ def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
 
 
 def test_insufficient_number_of_consituents(grid_that_straddles_dateline, use_dask):
-
     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"))
@@ -227,20 +221,17 @@ def test_insufficient_number_of_consituents(grid_that_straddles_dateline, use_da
 
 def test_tidal_forcing_plot(tidal_forcing):
     """Test plot method."""
-
     tidal_forcing.plot(var_name="ssh_Re", ntides=0)
 
 
 def test_tidal_forcing_save(tidal_forcing, tmp_path):
     """Test save method."""
-
     for file_str in ["test_tides", "test_tides.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             saved_filenames = tidal_forcing.save(filepath)
             # Check if the .nc file was created
             filepath = Path(filepath).with_suffix(".nc")
@@ -256,8 +247,8 @@ def test_tidal_forcing_save(tidal_forcing, tmp_path):
 )
 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."""
-
+    re-opening yaml file creates the same object.
+    """
     tidal_forcing = request.getfixturevalue(tidal_forcing_fixture)
 
     # Create a temporary filepath using the tmp_path fixture
@@ -266,7 +257,6 @@ def test_roundtrip_yaml(tidal_forcing_fixture, tmp_path, use_dask, request):
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         tidal_forcing.to_yaml(filepath)
 
         tidal_forcing_from_file = TidalForcing.from_yaml(filepath, use_dask=use_dask)
@@ -282,7 +272,6 @@ def test_roundtrip_yaml(tidal_forcing_fixture, tmp_path, use_dask, request):
     ["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"
@@ -331,7 +320,6 @@ def test_from_yaml_missing_tidal_forcing(tmp_path, use_dask):
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         # Write YAML content to file
         if isinstance(yaml_filepath, Path):
             yaml_filepath.write_text(yaml_content)

roms_tools/tests/test_setup/test_utils.py

@@ -1,15 +1,18 @@
-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
 
-def test_interpolate_from_climatology(use_dask):
+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):
     fname = download_test_data("ERA5_regional_test_data.nc")
     era5_times = xr.open_dataset(fname).time
 
@@ -25,7 +28,6 @@ def test_interpolate_from_climatology(use_dask):
 @pytest.fixture()
 def boundary_forcing_from_multiple_source_files(request, use_dask):
     """Fixture for creating a BoundaryForcing object."""
-
     grid = Grid(
         nx=5,
         ny=5,
@@ -53,15 +55,14 @@ def test_roundtrip_yaml(
     boundary_forcing_from_multiple_source_files, request, tmp_path, use_dask
 ):
     """Test that creating a BoundaryForcing object, saving its parameters to yaml file,
-    and re-opening yaml file creates the same object."""
-
+    and re-opening yaml file creates the same object.
+    """
     # Create a temporary filepath using the tmp_path fixture
     file_str = "test_yaml"
     for filepath in [
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         boundary_forcing_from_multiple_source_files.to_yaml(filepath)
 
         bdry_forcing_from_file = BoundaryForcing.from_yaml(filepath, use_dask=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
 
 
@@ -34,7 +35,6 @@ def _get_fname(name):
 # 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):
-
     overwrite = request.config.getoption("--overwrite")
 
     # Skip the test if the fixture isn't marked for overwriting, unless 'all' is specified
@@ -73,12 +73,14 @@ def test_save_results(forcing_fixture, request):
     ],
 )
 def test_check_results(forcing_fixture, request):
-
     fname = _get_fname(forcing_fixture)
     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(
@@ -97,7 +99,6 @@ def test_check_results(forcing_fixture, request):
 )
 def test_dask_vs_no_dask(forcing_fixture, request, tmp_path, use_dask):
     """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")
 

roms_tools/tests/test_tiling/test_partition.py

@@ -1,14 +1,14 @@
-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
 def grid():
-
     grid = Grid(nx=30, ny=30, size_x=80, size_y=80, center_lon=-20, center_lat=0, rot=0)
 
     return grid
@@ -231,6 +231,20 @@ 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,14 @@
 from numbers import Integral
+from pathlib import Path
 
 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 +29,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]]
@@ -57,7 +62,6 @@ def partition(
     This example partitions the dataset into 2 tiles along the `eta` direction and 3 tiles
     along the `xi` direction, resulting in a total of 6 partitions.
     """
-
     if (
         not isinstance(np_eta, Integral)
         or np_eta < 1
@@ -66,17 +70,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 +296,10 @@ def partition(
 
 
 def partition_netcdf(
-    filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
+    filepath: 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,25 +317,33 @@ 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]
         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"))
+    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