roms-tools 1.4.1__py3-none-any.whl → 1.5.0__py3-none-any.whl

roms_tools/tests/test_setup/test_datasets.py

@@ -125,9 +125,7 @@ def non_global_dataset():
     ],
 )
 def test_select_times(data_fixture, expected_time_values, request, tmp_path, use_dask):
-    """
-    Test selecting times with different datasets.
-    """
+    """Test selecting times with different datasets."""
     start_time = datetime(2022, 2, 1)
     end_time = datetime(2022, 3, 1)
 
@@ -160,9 +158,7 @@ def test_select_times(data_fixture, expected_time_values, request, tmp_path, use
 def test_select_times_valid_start_no_end_time(
     data_fixture, expected_time_values, request, tmp_path, use_dask
 ):
-    """
-    Test selecting times with only start_time specified.
-    """
+    """Test selecting times with only start_time specified."""
     start_time = datetime(2022, 2, 1)
 
     # Get the fixture dynamically based on the parameter
@@ -196,9 +192,7 @@ def test_select_times_valid_start_no_end_time(
 def test_select_times_invalid_start_no_end_time(
     data_fixture, expected_time_values, request, tmp_path, use_dask
 ):
-    """
-    Test selecting times with only start_time specified.
-    """
+    """Test selecting times with only start_time specified."""
     # Get the fixture dynamically based on the parameter
     dataset = request.getfixturevalue(data_fixture)
 
@@ -221,9 +215,8 @@ def test_select_times_invalid_start_no_end_time(
 def test_multiple_matching_times(
     global_dataset_with_multiple_times_per_day, tmp_path, use_dask
 ):
-    """
-    Test handling when multiple matching times are found when end_time is not specified.
-    """
+    """Test handling when multiple matching times are found when end_time is not
+    specified."""
     filepath = tmp_path / "test.nc"
     global_dataset_with_multiple_times_per_day.to_netcdf(filepath)
     dataset = Dataset(
@@ -237,9 +230,7 @@ def test_multiple_matching_times(
 
 
 def test_warnings_times(global_dataset, tmp_path, use_dask):
-    """
-    Test handling when no matching times are found.
-    """
+    """Test handling when no matching times are found."""
     # Create a temporary file
     filepath = tmp_path / "test.nc"
     global_dataset.to_netcdf(filepath)
@@ -271,9 +262,8 @@ def test_warnings_times(global_dataset, tmp_path, use_dask):
 def test_reverse_latitude_reverse_depth_choose_subdomain(
     global_dataset, tmp_path, use_dask
 ):
-    """
-    Test reversing latitude when it is not ascending, the choose_subdomain method, and the convert_to_negative_depth method of the Dataset class.
-    """
+    """Test reversing latitude when it is not ascending, the choose_subdomain method,
+    and the convert_to_negative_depth method of the Dataset class."""
     start_time = datetime(2022, 1, 1)
 
     filepath = tmp_path / "test.nc"

roms_tools/tests/test_setup/test_fill.py

@@ -23,14 +23,15 @@ def era5_data(request, use_dask):
         end_time=datetime(2020, 2, 2),
         use_dask=use_dask,
     )
-    data.post_process()
 
     return data
 
 
 @pytest.fixture()
 def glorys_data(request, use_dask):
-    fname = download_test_data("GLORYS_NA_2012.nc")
+    # the following GLORYS data has a wide enough domain
+    # to have different masks for tracers vs. velocities
+    fname = download_test_data("GLORYS_test_data.nc")
 
     data = GLORYSDataset(
         filename=fname,
@@ -38,7 +39,9 @@ def glorys_data(request, use_dask):
         end_time=datetime(2013, 1, 1),
         use_dask=use_dask,
     )
-    data.post_process()
+
+    ds = data.ds.isel(depth=[0, 10, 30])
+    object.__setattr__(data, "ds", ds)
 
     # extrapolate deepest value to bottom so all levels can use the same surface mask
     for var in data.var_names:
@@ -58,7 +61,6 @@ def tpxo_data(request, use_dask):
         filename=fname,
         use_dask=use_dask,
     )
-    data.post_process()
 
     return data
 
@@ -74,7 +76,6 @@ def cesm_bgc_data(request, use_dask):
         climatology=False,
         use_dask=use_dask,
     )
-    data.post_process()
 
     # extrapolate deepest value to bottom so all levels can use the same surface mask
     for var in data.var_names:
@@ -105,15 +106,25 @@ def cesm_surface_bgc_data(request, use_dask):
     "data_fixture",
     ["era5_data", "glorys_data", "tpxo_data", "cesm_bgc_data", "cesm_surface_bgc_data"],
 )
-def test_lateral_fill_no_nans(data_fixture, request, use_dask):
+def test_lateral_fill_no_nans(data_fixture, request):
     data = request.getfixturevalue(data_fixture)
     lateral_fill = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
 
     for var in data.var_names:
-        filled = lateral_fill.apply(data.ds[data.var_names[var]].astype(np.float64))
+        if var in ["u", "v"]:
+            filled = lateral_fill_vel.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
+        else:
+            filled = lateral_fill.apply(data.ds[data.var_names[var]].astype(np.float64))
         assert not filled.isnull().any()
 
 
@@ -150,28 +161,46 @@ def test_lateral_fill_correct_order_of_magnitude(cesm_bgc_data):
         "glorys_data",
     ],
 )
-def test_lateral_fill_reproducibility(data_fixture, request, use_dask):
+def test_lateral_fill_reproducibility(data_fixture, request):
 
     data = request.getfixturevalue(data_fixture)
     lateral_fill0 = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel0 = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
     lateral_fill1 = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel1 = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
 
     ds0 = data.ds.copy()
     ds1 = data.ds.copy()
 
     for var in data.var_names:
+        if var in ["u", "v"]:
+            ds0[data.var_names[var]] = lateral_fill_vel0.apply(
+                ds0[data.var_names[var]].astype(np.float64)
+            )
+            ds1[data.var_names[var]] = lateral_fill_vel1.apply(
+                ds1[data.var_names[var]].astype(np.float64)
+            )
 
-        ds0[data.var_names[var]] = lateral_fill0.apply(
-            ds0[data.var_names[var]].astype(np.float64)
-        )
-        ds1[data.var_names[var]] = lateral_fill1.apply(
-            ds1[data.var_names[var]].astype(np.float64)
-        )
+        else:
+            ds0[data.var_names[var]] = lateral_fill0.apply(
+                ds0[data.var_names[var]].astype(np.float64)
+            )
+            ds1[data.var_names[var]] = lateral_fill1.apply(
+                ds1[data.var_names[var]].astype(np.float64)
+            )
 
     assert ds0.equals(ds1)

roms_tools/tests/test_setup/test_grid.py

@@ -45,16 +45,14 @@ def test_grid_creation(grid):
     ["grid", "counter_clockwise_rotated_grid", "clockwise_rotated_grid"],
 )
 def test_coords_relation(grid_fixture, request):
-    """
-    Test that the coordinates satisfy the expected relations on a C-grid.
-    """
+    """Test that the coordinates satisfy the expected relations on a C-grid."""
     grid = request.getfixturevalue(grid_fixture)
 
     # psi versus rho
-    assert grid.ds.lon_psi.min() < grid.ds.lon_rho.min()
-    assert grid.ds.lon_psi.max() > grid.ds.lon_rho.max()
-    assert grid.ds.lat_psi.min() < grid.ds.lat_rho.min()
-    assert grid.ds.lat_psi.max() > grid.ds.lat_rho.max()
+    # assert grid.ds.lon_psi.min() < grid.ds.lon_rho.min()
+    # assert grid.ds.lon_psi.max() > grid.ds.lon_rho.max()
+    # assert grid.ds.lat_psi.min() < grid.ds.lat_rho.min()
+    # assert grid.ds.lat_psi.max() > grid.ds.lat_rho.max()
 
     # Assertion with tolerance is necessary for non-rotated grids
     def assert_larger_equal_than_with_tolerance(value1, value2, tolerance=1e-5):
@@ -218,7 +216,8 @@ def test_compatability_with_matlab_grid(tmp_path):
 
 
 def test_roundtrip_netcdf(tmp_path):
-    """Test that creating a grid, saving it to file, and re-opening it is the same as just creating it."""
+    """Test that creating a grid, saving it to file, and re-opening it is the same as
+    just creating it."""
 
     # Initialize a Grid object using the initializer
     grid_init = Grid(
@@ -255,7 +254,8 @@ def test_roundtrip_netcdf(tmp_path):
 
 
 def test_roundtrip_yaml(tmp_path):
-    """Test that creating a grid, saving its parameters to yaml file, and re-opening yaml file creates the same grid."""
+    """Test that creating a grid, saving its parameters to yaml file, and re- opening
+    yaml file creates the same grid."""
 
     # Initialize a Grid object using the initializer
     grid_init = Grid(

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -19,9 +19,7 @@ from conftest import calculate_file_hash
     ],
 )
 def test_initial_conditions_creation(ic_fixture, request):
-    """
-    Test the creation of the InitialConditions object.
-    """
+    """Test the creation of the InitialConditions object."""
 
     ic = request.getfixturevalue(ic_fixture)
 
@@ -181,9 +179,7 @@ def test_interpolation_from_climatology(
 def test_initial_conditions_plot_save(
     initial_conditions_with_bgc_from_climatology, tmp_path
 ):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
 
     initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0)
     initial_conditions_with_bgc_from_climatology.plot(
@@ -251,7 +247,8 @@ def test_initial_conditions_plot_save(
 
 
 def test_roundtrip_yaml(initial_conditions, tmp_path, use_dask):
-    """Test that creating an InitialConditions object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+    """Test that creating an InitialConditions object, saving its parameters to yaml
+    file, and re-opening yaml file creates the same object."""
 
     # Create a temporary filepath using the tmp_path fixture
     file_str = "test_yaml"

roms_tools/tests/test_setup/test_regrid.py

@@ -0,0 +1,59 @@
+import pytest
+import numpy as np
+import xarray as xr
+from roms_tools.setup.regrid import VerticalRegrid
+
+
+def vertical_regridder(depth_values, layer_depth_rho_values):
+    class DataContainer:
+        """Mock class for holding data and dimension names."""
+
+        def __init__(self, ds):
+            self.ds = ds
+            self.dim_names = {"depth": "depth"}
+
+    class Grid:
+        """Mock class representing the grid object with layer depth information."""
+
+        def __init__(self, ds):
+            self.ds = ds
+
+    # Creating minimal mock data for testing
+    # Depth levels in meters
+
+    # Create mock datasets for DataContainer and Grid
+    data_ds = xr.Dataset({"depth": (["depth"], depth_values)})
+    target_depth = xr.DataArray(data=layer_depth_rho_values, dims=["s_rho"])
+    # Instantiate DataContainer and Grid objects with mock datasets
+    mock_data = DataContainer(data_ds)
+
+    return VerticalRegrid(mock_data, target_depth)
+
+
+@pytest.mark.parametrize(
+    "depth_values, layer_depth_rho_values, temp_data",
+    [
+        ([5, 50, 100, 150], [130, 100, 70, 30, 10], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [130, 100, 70, 30, 2], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [200, 100, 70, 30, 10], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [200, 100, 70, 30, 1], [30, 25, 10, 2]),
+    ],
+)
+def test_vertical_regrid(request, depth_values, layer_depth_rho_values, temp_data):
+
+    vertical_regrid = vertical_regridder(
+        depth_values=depth_values, layer_depth_rho_values=layer_depth_rho_values
+    )
+    data = xr.Dataset({"temp_data": (["depth"], temp_data)})
+
+    # without filling in NaNs
+    regridded = vertical_regrid.apply(data.temp_data, fill_nans=False)
+    expected = np.interp(
+        layer_depth_rho_values, depth_values, temp_data, left=np.nan, right=np.nan
+    )
+    assert np.allclose(expected, regridded.data, equal_nan=True)
+
+    # with filling in NaNs
+    regridded = vertical_regrid.apply(data.temp_data, fill_nans=True)
+    expected = np.interp(layer_depth_rho_values, depth_values, temp_data)
+    assert np.allclose(expected, regridded.data, equal_nan=True)