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

roms_tools/tests/test_setup/test_fill.py

@@ -0,0 +1,177 @@
+import pytest
+from roms_tools.setup.datasets import (
+    GLORYSDataset,
+    ERA5Dataset,
+    CESMBGCDataset,
+    CESMBGCSurfaceForcingDataset,
+    TPXODataset,
+)
+from roms_tools.setup.download import download_test_data
+from roms_tools.setup.fill import LateralFill
+from roms_tools.setup.utils import extrapolate_deepest_to_bottom
+from datetime import datetime
+import numpy as np
+import xarray as xr
+
+
+@pytest.fixture()
+def era5_data(request, use_dask):
+    fname = download_test_data("ERA5_regional_test_data.nc")
+    data = ERA5Dataset(
+        filename=fname,
+        start_time=datetime(2020, 1, 31),
+        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")
+
+    data = GLORYSDataset(
+        filename=fname,
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        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:
+        if var != "zeta":
+            data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
+                data.ds[data.var_names[var]], data.dim_names["depth"]
+            )
+
+    return data
+
+
+@pytest.fixture()
+def tpxo_data(request, use_dask):
+    fname = download_test_data("TPXO_regional_test_data.nc")
+
+    data = TPXODataset(
+        filename=fname,
+        use_dask=use_dask,
+    )
+    data.post_process()
+
+    return data
+
+
+@pytest.fixture()
+def cesm_bgc_data(request, use_dask):
+    fname = download_test_data("CESM_BGC_2012.nc")
+
+    data = CESMBGCDataset(
+        filename=fname,
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        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:
+        data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
+            data.ds[data.var_names[var]], data.dim_names["depth"]
+        )
+
+    return data
+
+
+@pytest.fixture()
+def cesm_surface_bgc_data(request, use_dask):
+    fname = download_test_data("CESM_BGC_SURFACE_2012.nc")
+
+    data = CESMBGCSurfaceForcingDataset(
+        filename=fname,
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        climatology=False,
+        use_dask=use_dask,
+    )
+    data.post_process()
+
+    return data
+
+
+@pytest.mark.parametrize(
+    "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):
+    data = request.getfixturevalue(data_fixture)
+    lateral_fill = LateralFill(
+        data.ds["mask"],
+        [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))
+        assert not filled.isnull().any()
+
+
+def test_lateral_fill_correct_order_of_magnitude(cesm_bgc_data):
+
+    lateral_fill = LateralFill(
+        cesm_bgc_data.ds["mask"],
+        [cesm_bgc_data.dim_names["latitude"], cesm_bgc_data.dim_names["longitude"]],
+    )
+
+    ALK = cesm_bgc_data.ds["ALK"]
+
+    # zero out alkalinity field in all depth levels but the uppermost
+    ALK = xr.where(cesm_bgc_data.ds.ALK.depth > 25, 0, cesm_bgc_data.ds.ALK)
+    ALK = ALK.where(cesm_bgc_data.ds.mask)
+
+    filled = lateral_fill.apply(ALK.astype(np.float64))
+
+    # check that alkalinity values in the uppermost values are of the correct order of magnitude
+    # and that no new minima and maxima are introduced
+    assert filled.isel(depth=0).min() == ALK.isel(depth=0).min()
+    assert filled.isel(depth=0).max() == ALK.isel(depth=0).max()
+
+    # check that the filled alkalinity values are zero in all deeper layers
+    assert filled.isel(depth=slice(1, None)).equals(
+        xr.zeros_like(filled.isel(depth=slice(1, None)))
+    )
+
+
+@pytest.mark.parametrize(
+    "data_fixture",
+    [
+        "era5_data",
+        "glorys_data",
+    ],
+)
+def test_lateral_fill_reproducibility(data_fixture, request, use_dask):
+
+    data = request.getfixturevalue(data_fixture)
+    lateral_fill0 = LateralFill(
+        data.ds["mask"],
+        [data.dim_names["latitude"], data.dim_names["longitude"]],
+    )
+    lateral_fill1 = LateralFill(
+        data.ds["mask"],
+        [data.dim_names["latitude"], data.dim_names["longitude"]],
+    )
+
+    ds0 = data.ds.copy()
+    ds1 = data.ds.copy()
+
+    for var in data.var_names:
+
+        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

@@ -8,6 +8,26 @@ from conftest import calculate_file_hash
 from pathlib import Path
 
 
+@pytest.fixture()
+def counter_clockwise_rotated_grid():
+
+    grid = Grid(
+        nx=1, ny=1, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=20
+    )
+
+    return grid
+
+
+@pytest.fixture()
+def clockwise_rotated_grid():
+
+    grid = Grid(
+        nx=1, ny=1, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=-20
+    )
+
+    return grid
+
+
 def test_grid_creation(grid):
 
     assert grid.nx == 1
@@ -20,6 +40,42 @@ def test_grid_creation(grid):
     assert isinstance(grid.ds, xr.Dataset)
 
 
+@pytest.mark.parametrize(
+    "grid_fixture",
+    ["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.
+    """
+    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()
+
+    # Assertion with tolerance is necessary for non-rotated grids
+    def assert_larger_equal_than_with_tolerance(value1, value2, tolerance=1e-5):
+        assert value1 >= value2 - tolerance
+
+    def assert_smaller_equal_than_with_tolerance(value1, value2, tolerance=1e-5):
+        assert value1 <= value2 + tolerance
+
+    # u versus rho
+    assert_larger_equal_than_with_tolerance(grid.ds.lon_u.min(), grid.ds.lon_rho.min())
+    assert_larger_equal_than_with_tolerance(grid.ds.lat_u.min(), grid.ds.lat_rho.min())
+    assert_smaller_equal_than_with_tolerance(grid.ds.lon_u.max(), grid.ds.lon_rho.max())
+    assert_smaller_equal_than_with_tolerance(grid.ds.lon_u.max(), grid.ds.lon_rho.max())
+
+    # v versus rho
+    assert_larger_equal_than_with_tolerance(grid.ds.lon_v.min(), grid.ds.lon_rho.min())
+    assert_larger_equal_than_with_tolerance(grid.ds.lat_v.min(), grid.ds.lat_rho.min())
+    assert_smaller_equal_than_with_tolerance(grid.ds.lon_v.max(), grid.ds.lon_rho.max())
+    assert_smaller_equal_than_with_tolerance(grid.ds.lon_v.max(), grid.ds.lon_rho.max())
+
+
 def test_plot_save_methods(tmp_path):
 
     grid = Grid(

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -6,8 +6,8 @@ import numpy as np
 import textwrap
 from roms_tools.setup.download import download_test_data
 from roms_tools.setup.datasets import CESMBGCDataset
-from conftest import calculate_file_hash
 from pathlib import Path
+from conftest import calculate_file_hash
 
 
 @pytest.mark.parametrize(
@@ -143,7 +143,7 @@ def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
         use_dask=use_dask,
     )
 
-    assert initial_conditions.bgc_source["climatology"] is True
+    assert initial_conditions.bgc_source["climatology"] is False
 
 
 def test_interpolation_from_climatology(
@@ -159,6 +159,7 @@ def test_interpolation_from_climatology(
         start_time=datetime(2012, 1, 15),
         climatology=True,
         use_dask=use_dask,
+        apply_post_processing=False,
     )
     assert np.allclose(ds["ALK"].sel(month=1), bgc_data.ds["ALK"], equal_nan=True)
 
@@ -168,6 +169,7 @@ def test_interpolation_from_climatology(
         start_time=datetime(2012, 1, 30),
         climatology=True,
         use_dask=use_dask,
+        apply_post_processing=False,
     )
     assert np.allclose(
         0.5 * (ds["ALK"].sel(month=1) + ds["ALK"].sel(month=2)),
@@ -272,7 +274,7 @@ def test_roundtrip_yaml(initial_conditions, tmp_path, use_dask):
 
 def test_files_have_same_hash(initial_conditions, tmp_path, use_dask):
 
-    yaml_filepath = tmp_path / "test_yaml"
+    yaml_filepath = tmp_path / "test_yaml.yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
 

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -3,8 +3,8 @@ from datetime import datetime
 from roms_tools import Grid, SurfaceForcing
 from roms_tools.setup.download import download_test_data
 import textwrap
-from conftest import calculate_file_hash
 from pathlib import Path
+from conftest import calculate_file_hash
 
 
 @pytest.fixture
@@ -13,8 +13,8 @@ def grid_that_straddles_dateline():
     Fixture for creating a domain that straddles the dateline and lies within the bounds of the regional ERA5 data.
     """
     grid = Grid(
-        nx=5,
-        ny=5,
+        nx=20,
+        ny=20,
         size_x=1800,
         size_y=2400,
         center_lon=-10,
@@ -643,7 +643,7 @@ def test_files_have_same_hash(sfc_forcing_fixture, request, tmp_path, use_dask):
 
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
-    yaml_filepath = tmp_path / "test_yaml"
+    yaml_filepath = tmp_path / "test_yaml.yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
 

roms_tools/tests/test_setup/test_tides.py

@@ -3,8 +3,8 @@ from roms_tools import Grid, TidalForcing
 import xarray as xr
 from roms_tools.setup.download import download_test_data
 import textwrap
-from conftest import calculate_file_hash
 from pathlib import Path
+from conftest import calculate_file_hash
 
 
 @pytest.fixture
@@ -232,10 +232,13 @@ def test_roundtrip_yaml(tidal_forcing, tmp_path, use_dask):
 
 def test_files_have_same_hash(tidal_forcing, tmp_path, use_dask):
 
-    yaml_filepath = tmp_path / "test_yaml"
+    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_topography.py

@@ -56,3 +56,43 @@ def test_hmin_criterion():
 
     assert grid.hmin == 10.0
     assert np.less_equal(grid.hmin, grid.ds.h.min())
+
+
+def test_mask_topography_boundary():
+    """
+    Test that the mask and topography along the grid boundaries (north, south, east, west)
+    are identical to the adjacent inland cells.
+    """
+
+    # Create a grid with some land along the northern boundary
+    grid = Grid(
+        nx=10, ny=10, size_x=1000, size_y=1000, center_lon=-20, center_lat=60, rot=0
+    )
+
+    # Toopography
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(eta_rho=0).data, grid.ds.h.isel(eta_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(eta_rho=-1).data, grid.ds.h.isel(eta_rho=-2).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(xi_rho=0).data, grid.ds.h.isel(xi_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(xi_rho=-1).data, grid.ds.h.isel(xi_rho=-2).data
+    )
+
+    # Mask
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(eta_rho=0).data, grid.ds.mask_rho.isel(eta_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(eta_rho=-1).data, grid.ds.mask_rho.isel(eta_rho=-2).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(xi_rho=0).data, grid.ds.mask_rho.isel(xi_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(xi_rho=-1).data, grid.ds.mask_rho.isel(xi_rho=-2).data
+    )

roms_tools/tests/test_setup/test_validation.py

@@ -69,4 +69,4 @@ def test_check_results(forcing_fixture, request):
 
     forcing = request.getfixturevalue(forcing_fixture)
 
-    xr.testing.assert_allclose(forcing.ds, expected_forcing_ds)
+    xr.testing.assert_allclose(forcing.ds, expected_forcing_ds, rtol=1.0e-5)

roms_tools/tests/test_utils.py

@@ -23,6 +23,8 @@ class TestPartitionGrid:
             "xi_rho": 32,
             "xi_u": 31,
             "eta_v": 10,
+            "eta_psi": 11,
+            "xi_psi": 33,
             "eta_coarse": 6,
             "xi_coarse": 17,
             "s_rho": 100,
@@ -33,6 +35,8 @@ class TestPartitionGrid:
             "xi_rho": 32,
             "xi_u": 31,
             "eta_v": 10,
+            "eta_psi": 10,
+            "xi_psi": 33,
             "eta_coarse": 5,
             "xi_coarse": 17,
             "s_rho": 100,
@@ -43,6 +47,8 @@ class TestPartitionGrid:
             "xi_rho": 32,
             "xi_u": 31,
             "eta_v": 11,
+            "eta_psi": 12,
+            "xi_psi": 33,
             "eta_coarse": 6,
             "xi_coarse": 17,
             "s_rho": 100,
@@ -57,6 +63,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 10,
             "eta_v": 31,
+            "eta_psi": 33,
+            "xi_psi": 11,
             "eta_coarse": 17,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -67,6 +75,8 @@ class TestPartitionGrid:
             "xi_rho": 10,
             "xi_u": 10,
             "eta_v": 31,
+            "eta_psi": 33,
+            "xi_psi": 10,
             "eta_coarse": 17,
             "xi_coarse": 5,
             "s_rho": 100,
@@ -77,6 +87,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 11,
             "eta_v": 31,
+            "eta_psi": 33,
+            "xi_psi": 12,
             "eta_coarse": 17,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -96,6 +108,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 10,
             "eta_v": 10,
+            "eta_psi": 11,
+            "xi_psi": 11,
             "eta_coarse": 6,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -107,6 +121,8 @@ class TestPartitionGrid:
             "xi_u": 10,
             "eta_v": 10,
             "eta_coarse": 5,
+            "eta_psi": 10,
+            "xi_psi": 11,
             "xi_coarse": 6,
             "s_rho": 100,
             "s_w": 101,
@@ -116,6 +132,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 10,
             "eta_v": 11,
+            "eta_psi": 12,
+            "xi_psi": 11,
             "eta_coarse": 6,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -126,6 +144,8 @@ class TestPartitionGrid:
             "xi_rho": 10,
             "xi_u": 10,
             "eta_v": 10,
+            "eta_psi": 11,
+            "xi_psi": 10,
             "eta_coarse": 6,
             "xi_coarse": 5,
             "s_rho": 100,
@@ -136,6 +156,8 @@ class TestPartitionGrid:
             "xi_rho": 10,
             "xi_u": 10,
             "eta_v": 10,
+            "eta_psi": 10,
+            "xi_psi": 10,
             "eta_coarse": 5,
             "xi_coarse": 5,
             "s_rho": 100,
@@ -146,6 +168,8 @@ class TestPartitionGrid:
             "xi_rho": 10,
             "xi_u": 10,
             "eta_v": 11,
+            "eta_psi": 12,
+            "xi_psi": 10,
             "eta_coarse": 6,
             "xi_coarse": 5,
             "s_rho": 100,
@@ -156,6 +180,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 11,
             "eta_v": 10,
+            "eta_psi": 11,
+            "xi_psi": 12,
             "eta_coarse": 6,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -166,6 +192,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 11,
             "eta_v": 10,
+            "eta_psi": 10,
+            "xi_psi": 12,
             "eta_coarse": 5,
             "xi_coarse": 6,
             "s_rho": 100,
@@ -176,6 +204,8 @@ class TestPartitionGrid:
             "xi_rho": 11,
             "xi_u": 11,
             "eta_v": 11,
+            "eta_psi": 12,
+            "xi_psi": 12,
             "eta_coarse": 6,
             "xi_coarse": 6,
             "s_rho": 100,

roms_tools/utils.py

@@ -14,7 +14,7 @@ def partition(
 
     This function divides the input dataset into `np_eta` by `np_xi` tiles, where each tile
     represents a subdomain of the original dataset. The partitioning is performed along
-    the spatial dimensions `eta_rho`, `xi_rho`, `eta_v`, `xi_u`, `eta_coarse`, and `xi_coarse`,
+    the spatial dimensions `eta_rho`, `xi_rho`, `eta_v`, `xi_u`, `eta_psi`, `xi_psi`, `eta_coarse`, and `xi_coarse`,
     depending on which dimensions are present in the dataset.
 
     Parameters
@@ -69,6 +69,8 @@ def partition(
         "xi_rho",
         "eta_v",
         "xi_u",
+        "eta_psi",
+        "xi_psi",
         "eta_coarse",
         "xi_coarse",
     ]
@@ -123,6 +125,7 @@ def partition(
         eta_rho_domain_size = integer_division_or_raise(
             ds.sizes["eta_rho"] - 2 * n_eta_ghost_cells, np_eta, "eta_rho"
         )
+
     if "xi_rho" in dims_to_partition:
         xi_rho_domain_size = integer_division_or_raise(
             ds.sizes["xi_rho"] - 2 * n_xi_ghost_cells, np_xi, "xi_rho"
@@ -132,11 +135,22 @@ def partition(
         eta_v_domain_size = integer_division_or_raise(
             ds.sizes["eta_v"] - 1 * n_eta_ghost_cells, np_eta, "eta_v"
         )
+
     if "xi_u" in dims_to_partition:
         xi_u_domain_size = integer_division_or_raise(
             ds.sizes["xi_u"] - 1 * n_xi_ghost_cells, np_xi, "xi_u"
         )
 
+    if "eta_psi" in dims_to_partition:
+        eta_psi_domain_size = integer_division_or_raise(
+            ds.sizes["eta_psi"] - 3 * n_eta_ghost_cells, np_eta, "eta_psi"
+        )
+
+    if "xi_psi" in dims_to_partition:
+        xi_psi_domain_size = integer_division_or_raise(
+            ds.sizes["xi_psi"] - 3 * n_xi_ghost_cells, np_xi, "xi_psi"
+        )
+
     if "eta_coarse" in dims_to_partition:
         eta_coarse_domain_size = integer_division_or_raise(
             ds.sizes["eta_coarse"] - 2 * n_eta_ghost_cells, np_eta, "eta_coarse"
@@ -167,7 +181,12 @@ def partition(
             + [eta_rho_domain_size] * (np_eta - 2)
             + [eta_rho_domain_size + n_eta_ghost_cells]
         )
-
+        if "eta_psi" in dims_to_partition:
+            partitioned_sizes["eta_psi"] = (
+                [n_eta_ghost_cells + eta_psi_domain_size]
+                + [eta_psi_domain_size] * (np_eta - 2)
+                + [eta_psi_domain_size + 2 * n_eta_ghost_cells]
+            )
         if "eta_coarse" in dims_to_partition:
             partitioned_sizes["eta_coarse"] = (
                 [eta_coarse_domain_size + n_eta_ghost_cells]
@@ -181,7 +200,12 @@ def partition(
             + [xi_rho_domain_size] * (np_xi - 2)
             + [xi_rho_domain_size + n_xi_ghost_cells]
         )
-
+        if "xi_psi" in dims_to_partition:
+            partitioned_sizes["xi_psi"] = (
+                [n_xi_ghost_cells + xi_psi_domain_size]
+                + [xi_psi_domain_size] * (np_xi - 2)
+                + [xi_psi_domain_size + 2 * n_xi_ghost_cells]
+            )
         if "xi_coarse" in dims_to_partition:
             partitioned_sizes["xi_coarse"] = (
                 [xi_coarse_domain_size + n_xi_ghost_cells]
@@ -229,6 +253,18 @@ def partition(
                 indexers["xi_u"] = slice(
                     int(xi_u_partition_indices[j]), int(xi_u_partition_indices[j + 1])
                 )
+            if "eta_psi" in dims_to_partition:
+                eta_psi_partition_indices = cumsum(partitioned_sizes["eta_psi"])
+                indexers["eta_psi"] = slice(
+                    int(eta_psi_partition_indices[i]),
+                    int(eta_psi_partition_indices[i + 1]),
+                )
+            if "xi_psi" in dims_to_partition:
+                xi_psi_partition_indices = cumsum(partitioned_sizes["xi_psi"])
+                indexers["xi_psi"] = slice(
+                    int(xi_psi_partition_indices[j]),
+                    int(xi_psi_partition_indices[j + 1]),
+                )
 
             if "eta_coarse" in dims_to_partition:
                 eta_coarse_partition_indices = cumsum(partitioned_sizes["eta_coarse"])