roms-tools 1.4.0__py3-none-any.whl → 1.4.2__py3-none-any.whl

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -9,12 +9,11 @@ from conftest import calculate_file_hash
 
 @pytest.fixture
 def grid_that_straddles_dateline():
-    """
-    Fixture for creating a domain that straddles the dateline and lies within the bounds of the regional ERA5 data.
-    """
+    """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,
@@ -27,8 +26,9 @@ def grid_that_straddles_dateline():
 
 @pytest.fixture
 def grid_that_straddles_dateline_but_is_too_big_for_regional_test_data():
-    """
-    Fixture for creating a domain that straddles the dateline but exceeds the bounds of the regional ERA5 data.
+    """Fixture for creating a domain that straddles the dateline but exceeds the bounds
+    of the regional ERA5 data.
+
     Centered east of dateline.
     """
     grid = Grid(
@@ -46,9 +46,11 @@ def grid_that_straddles_dateline_but_is_too_big_for_regional_test_data():
 
 @pytest.fixture
 def another_grid_that_straddles_dateline_but_is_too_big_for_regional_test_data():
-    """
-    Fixture for creating a domain that straddles the dateline but exceeds the bounds of the regional ERA5 data.
-    Centered west of dateline. This one was hard to catch for the nan_check for a long time, but should work now.
+    """Fixture for creating a domain that straddles the dateline but exceeds the bounds
+    of the regional ERA5 data.
+
+    Centered west of dateline. This one was hard to catch for the nan_check for a long
+    time, but should work now.
     """
     grid = Grid(
         nx=5,
@@ -65,9 +67,11 @@ def another_grid_that_straddles_dateline_but_is_too_big_for_regional_test_data()
 
 @pytest.fixture
 def grid_that_lies_east_of_dateline_less_than_five_degrees_away():
-    """
-    Fixture for creating a domain that lies east of Greenwich meridian, but less than 5 degrees away.
-    We care about the 5 degree mark because it decides whether the code handles the longitudes as straddling the dateline or not.
+    """Fixture for creating a domain that lies east of Greenwich meridian, but less than
+    5 degrees away.
+
+    We care about the 5 degree mark because it decides whether the code handles the
+    longitudes as straddling the dateline or not.
     """
 
     grid = Grid(
@@ -85,9 +89,11 @@ def grid_that_lies_east_of_dateline_less_than_five_degrees_away():
 
 @pytest.fixture
 def grid_that_lies_east_of_dateline_more_than_five_degrees_away():
-    """
-    Fixture for creating a domain that lies east of Greenwich meridian, more than 5 degrees away.
-    We care about the 5 degree mark because it decides whether the code handles the longitudes as straddling the dateline or not.
+    """Fixture for creating a domain that lies east of Greenwich meridian, more than 5
+    degrees away.
+
+    We care about the 5 degree mark because it decides whether the code handles the
+    longitudes as straddling the dateline or not.
     """
     grid = Grid(
         nx=5,
@@ -104,9 +110,11 @@ def grid_that_lies_east_of_dateline_more_than_five_degrees_away():
 
 @pytest.fixture
 def grid_that_lies_west_of_dateline_less_than_five_degrees_away():
-    """
-    Fixture for creating a domain that lies west of Greenwich meridian, less than 5 degrees away.
-    We care about the 5 degree mark because it decides whether the code handles the longitudes as straddling the dateline or not.
+    """Fixture for creating a domain that lies west of Greenwich meridian, less than 5
+    degrees away.
+
+    We care about the 5 degree mark because it decides whether the code handles the
+    longitudes as straddling the dateline or not.
     """
 
     grid = Grid(
@@ -124,9 +132,11 @@ def grid_that_lies_west_of_dateline_less_than_five_degrees_away():
 
 @pytest.fixture
 def grid_that_lies_west_of_dateline_more_than_five_degrees_away():
-    """
-    Fixture for creating a domain that lies west of Greenwich meridian, more than 5 degrees away.
-    We care about the 5 degree mark because it decides whether the code handles the longitudes as straddling the dateline or not.
+    """Fixture for creating a domain that lies west of Greenwich meridian, more than 5
+    degrees away.
+
+    We care about the 5 degree mark because it decides whether the code handles the
+    longitudes as straddling the dateline or not.
     """
 
     grid = Grid(
@@ -144,8 +154,10 @@ def grid_that_lies_west_of_dateline_more_than_five_degrees_away():
 
 @pytest.fixture
 def grid_that_straddles_180_degree_meridian():
-    """
-    Fixture for creating a domain that straddles 180 degree meridian. This is a good test grid for the global ERA5 data, which comes on an [-180, 180] longitude grid.
+    """Fixture for creating a domain that straddles 180 degree meridian.
+
+    This is a good test grid for the global ERA5 data, which comes on an [-180, 180]
+    longitude grid.
     """
 
     grid = Grid(
@@ -172,8 +184,7 @@ def grid_that_straddles_180_degree_meridian():
     ],
 )
 def test_successful_initialization_with_regional_data(grid_fixture, request, use_dask):
-    """
-    Test the initialization of SurfaceForcing with regional ERA5 data.
+    """Test the initialization of SurfaceForcing with regional ERA5 data.
 
     This test checks the following:
     1. SurfaceForcing object initializes successfully with provided regional data.
@@ -242,8 +253,7 @@ def test_successful_initialization_with_regional_data(grid_fixture, request, use
 def test_nan_detection_initialization_with_regional_data(
     grid_fixture, request, use_dask
 ):
-    """
-    Test handling of NaN values during initialization with regional data.
+    """Test handling of NaN values during initialization with regional data.
 
     Ensures ValueError is raised if NaN values are detected in the dataset.
     """
@@ -270,10 +280,11 @@ def test_nan_detection_initialization_with_regional_data(
 def test_no_longitude_intersection_initialization_with_regional_data(
     grid_that_straddles_180_degree_meridian, use_dask
 ):
-    """
-    Test initialization of SurfaceForcing with a grid that straddles the 180° meridian.
+    """Test initialization of SurfaceForcing with a grid that straddles the 180°
+    meridian.
 
-    Ensures ValueError is raised when the longitude range does not intersect with the dataset.
+    Ensures ValueError is raised when the longitude range does not intersect with the
+    dataset.
     """
     start_time = datetime(2020, 1, 31)
     end_time = datetime(2020, 2, 2)
@@ -309,11 +320,11 @@ def test_no_longitude_intersection_initialization_with_regional_data(
     ],
 )
 def test_successful_initialization_with_global_data(grid_fixture, request, use_dask):
-    """
-    Test initialization of SurfaceForcing with global data.
+    """Test initialization of SurfaceForcing with global data.
 
     Verifies that the SurfaceForcing object is correctly initialized with global data,
-    including the correct handling of the grid and physics data. Checks both coarse and fine grid initialization.
+    including the correct handling of the grid and physics data. Checks both coarse and
+    fine grid initialization.
     """
     start_time = datetime(2020, 1, 31)
     end_time = datetime(2020, 2, 2)
@@ -357,8 +368,7 @@ def test_successful_initialization_with_global_data(grid_fixture, request, use_d
 
 
 def test_nans_filled_in(grid_that_straddles_dateline, use_dask):
-    """
-    Test that the surface forcing fields contain no NaNs.
+    """Test that the surface forcing fields contain no NaNs.
 
     The test is performed twice:
     - First with the default fine grid.
@@ -405,10 +415,8 @@ def test_nans_filled_in(grid_that_straddles_dateline, use_dask):
 
 
 def test_time_attr_climatology(bgc_surface_forcing_from_climatology):
-    """
-    Test that the 'cycle_length' attribute is present in the time coordinate of the BGC dataset
-    when using climatology data.
-    """
+    """Test that the 'cycle_length' attribute is present in the time coordinate of the
+    BGC dataset when using climatology data."""
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert hasattr(
             bgc_surface_forcing_from_climatology.ds[time_coord],
@@ -418,10 +426,8 @@ def test_time_attr_climatology(bgc_surface_forcing_from_climatology):
 
 
 def test_time_attr(bgc_surface_forcing):
-    """
-    Test that the 'cycle_length' attribute is not present in the time coordinate of the BGC dataset
-    when not using climatology data.
-    """
+    """Test that the 'cycle_length' attribute is not present in the time coordinate of
+    the BGC dataset when not using climatology data."""
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert not hasattr(
             bgc_surface_forcing.ds[time_coord],
@@ -448,12 +454,11 @@ def test_time_attr(bgc_surface_forcing):
 def test_surface_forcing_creation(
     sfc_forcing_fixture, expected_climatology, expected_fname, request
 ):
-    """
-    Test the creation and initialization of the SurfaceForcing object with BGC.
+    """Test the creation and initialization of the SurfaceForcing object with BGC.
 
     Verifies that the SurfaceForcing object is properly created with correct attributes.
-    Ensures that expected variables are present in the dataset
-    and that attributes match the given configurations.
+    Ensures that expected variables are present in the dataset and that attributes match
+    the given configurations.
     """
 
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
@@ -491,9 +496,7 @@ def test_surface_forcing_creation(
     ],
 )
 def test_surface_forcing_plot_save(sfc_forcing_fixture, request, tmp_path):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
     sfc_forcing.plot(varname="uwnd", time=0)
 
@@ -525,9 +528,7 @@ def test_surface_forcing_plot_save(sfc_forcing_fixture, request, tmp_path):
 
 
 def test_surface_forcing_bgc_plot_save(bgc_surface_forcing, tmp_path):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
 
     # Check the values in the dataset
     bgc_surface_forcing.plot(varname="pco2_air", time=0)
@@ -562,9 +563,7 @@ def test_surface_forcing_bgc_plot_save(bgc_surface_forcing, tmp_path):
 def test_surface_forcing_bgc_from_clim_plot_save(
     bgc_surface_forcing_from_climatology, tmp_path
 ):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
 
     # Check the values in the dataset
     bgc_surface_forcing_from_climatology.plot(varname="pco2_air", time=0)
@@ -609,7 +608,8 @@ def test_surface_forcing_bgc_from_clim_plot_save(
     ],
 )
 def test_roundtrip_yaml(sfc_forcing_fixture, request, tmp_path, use_dask):
-    """Test that creating an SurfaceForcing object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+    """Test that creating an SurfaceForcing object, saving its parameters to yaml file,
+    and re-opening yaml file creates the same object."""
 
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 

roms_tools/tests/test_setup/test_tides.py

@@ -25,9 +25,7 @@ def grid_that_is_out_of_bounds_of_regional_tpxo_data():
 
 @pytest.fixture
 def grid_that_straddles_dateline():
-    """
-    Fixture for creating a domain that straddles the dateline.
-    """
+    """Fixture for creating a domain that straddles the dateline."""
     grid = Grid(
         nx=5,
         ny=5,
@@ -43,9 +41,7 @@ def grid_that_straddles_dateline():
 
 @pytest.fixture
 def grid_that_straddles_180_degree_meridian():
-    """
-    Fixture for creating a domain that straddles 180 degree meridian.
-    """
+    """Fixture for creating a domain that straddles 180 degree meridian."""
 
     grid = Grid(
         nx=5,
@@ -174,9 +170,8 @@ def test_insufficient_number_of_consituents(grid_that_straddles_dateline, use_da
 
 
 def test_tidal_forcing_plot_save(tidal_forcing, tmp_path):
-    """
-    Test plot and save methods in the same test since we dask arrays are already computed.
-    """
+    """Test plot and save methods in the same test since we dask arrays are already
+    computed."""
     tidal_forcing.ds.load()
 
     tidal_forcing.plot(varname="ssh_Re", ntides=0)
@@ -211,7 +206,8 @@ def test_tidal_forcing_plot_save(tidal_forcing, tmp_path):
 
 
 def test_roundtrip_yaml(tidal_forcing, tmp_path, use_dask):
-    """Test that creating a TidalForcing object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+    """Test that creating a TidalForcing 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_topography.py

@@ -56,3 +56,41 @@ 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_vertical_coordinate.py

@@ -3,9 +3,7 @@ from roms_tools import Grid
 
 
 def test_invalid_theta_s_value():
-    """
-    Test the validation of the theta_s value.
-    """
+    """Test the validation of the theta_s value."""
     with pytest.raises(ValueError):
 
         Grid(
@@ -24,9 +22,7 @@ def test_invalid_theta_s_value():
 
 
 def test_invalid_theta_b_value():
-    """
-    Test the validation of the theta_b value.
-    """
+    """Test the validation of the theta_b value."""
     with pytest.raises(ValueError):
         Grid(
             nx=2,

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

@@ -9,12 +9,12 @@ from pathlib import Path
 def partition(
     ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1
 ) -> tuple[list[int], list[xr.Dataset]]:
-    """
-    Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial tiles.
+    """Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial
+    tiles.
 
     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
@@ -32,6 +32,7 @@ def partition(
     -------
     tuple[list[int], list[xr.Dataset]]
         A tuple containing two elements:
+
         - A list of integers representing the file numbers associated with each partition.
         - A list of `xarray.Dataset` objects, each representing a partitioned subdomain of the original dataset.
 
@@ -69,6 +70,8 @@ def partition(
         "xi_rho",
         "eta_v",
         "xi_u",
+        "eta_psi",
+        "xi_psi",
         "eta_coarse",
         "xi_coarse",
     ]
@@ -88,8 +91,7 @@ def partition(
         n_xi_ghost_cells = 1
 
     def integer_division_or_raise(a: int, b: int, dimension: str) -> int:
-        """
-        Perform integer division and ensure that the division is exact.
+        """Perform integer division and ensure that the division is exact.
 
         Parameters
         ----------
@@ -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]
@@ -190,7 +214,7 @@ def partition(
             )
 
     def cumsum(pmf):
-        """Implementation of cumsum which ensures the result starts with zero"""
+        """Implementation of cumsum which ensures the result starts with zero."""
         cdf = np.empty(len(pmf) + 1, dtype=int)
         cdf[0] = 0
         np.cumsum(pmf, out=cdf[1:])
@@ -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"])
@@ -254,8 +290,7 @@ def partition(
 def partition_netcdf(
     filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
 ) -> None:
-    """
-    Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
+    """Partition a ROMS NetCDF file 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.
     Each tile is saved as a separate NetCDF file.