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

roms_tools/setup/utils.py

@@ -7,7 +7,7 @@ from roms_tools.utils import partition
 from pathlib import Path
 
 
-def nan_check(field, mask) -> None:
+def nan_check(field, mask, error_message=None) -> None:
     """Checks for NaN values at wet points in the field.
 
     This function examines the interpolated input field for NaN values at positions indicated as wet points by the mask.
@@ -22,6 +22,10 @@ def nan_check(field, mask) -> None:
         A boolean mask or data array with the same shape as `field`. The wet points (usually ocean points)
         are indicated by `1` or `True`, and land points by `0` or `False`.
 
+    error_message : str, optional
+        A custom error message to be included in the ValueError if NaN values are detected. If not provided,
+        a default message will explain the potential cause and suggest ensuring the dataset's coverage.
+
     Raises
     ------
     ValueError
@@ -31,14 +35,15 @@ def nan_check(field, mask) -> None:
 
     # Replace values in field with 0 where mask is not 1
     da = xr.where(mask == 1, field, 0)
-
-    # Check if any NaN values exist in the modified field
-    if da.isnull().any().values:
-        raise ValueError(
+    if error_message is None:
+        error_message = (
             "NaN values found in interpolated field. This likely occurs because the ROMS grid, including "
             "a small safety margin for interpolation, is not fully contained within the dataset's longitude/latitude range. Please ensure that the "
             "dataset covers the entire area required by the ROMS grid."
         )
+    # Check if any NaN values exist in the modified field
+    if da.isnull().any().values:
+        raise ValueError(error_message)
 
 
 def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
@@ -156,63 +161,30 @@ def interpolate_from_rho_to_v(field, method="additive"):
     return field_interpolated
 
 
-def extrapolate_deepest_to_bottom(field: xr.DataArray, dim: str) -> xr.DataArray:
-    """Extrapolates the deepest non-NaN values to the bottom along the specified
-    dimension using forward fill.
-
-    This function assumes that the specified dimension is ordered from top to bottom (e.g., a vertical dimension like 'depth').
-    It fills `NaN` values below the deepest valid (non-NaN) entry along the given dimension by carrying forward the last valid value.
+def one_dim_fill(da: xr.DataArray, dim: str, direction="forward") -> xr.DataArray:
+    """Fill NaN values in a DataArray along a specified dimension.
 
     Parameters
     ----------
-    field : xr.DataArray
-        The input `xarray.DataArray` containing potential `NaN` values to be filled.
-        This array must have at least one dimension corresponding to `dim`, typically
-        a vertical axis such as 'depth' or 'height'.
+    da : xr.DataArray
+        The input DataArray with NaN values to be filled, which must include the specified dimension.
     dim : str
-        The name of the dimension along which to perform the forward fill operation.
-        The function assumes that this dimension is ordered from top to bottom, with
-        larger index values representing deeper or lower levels.
+        The name of the dimension along which to fill NaN values (e.g., 'depth' or 'time').
+    direction : str, optional
+        The filling direction; either "forward" to propagate non-NaN values downward or "backward" to propagate them upward.
+        Defaults to "forward".
 
     Returns
     -------
     xr.DataArray
-        A new `xarray.DataArray` with the `NaN` values along the specified dimension
-        filled by forward filling the deepest valid values down to the bottom.
-        The original input data remains unmodified.
-    """
-    if dim in field.dims:
-        return field.ffill(dim=dim)
-    else:
-        return field
-
-
-def _extrapolate_deepest_to_bottom(data_vars, data) -> dict:
-    """Extrapolate the deepest value to the bottom for variables using the dataset's
-    depth dimension.
-
-    This function fills in missing values at the bottom of each variable by
-    carrying forward the deepest available value, ensuring a complete depth profile.
-
-    Parameters
-    ----------
-    data_vars : dict
-        Existing dictionary of variables to be updated.
-    data : Dataset
-        Dataset containing variables and depth information.
-
-    Returns
-    -------
-    dict of str : xarray.DataArray
-        Dictionary of variables with the deepest value extrapolated to the bottom.
+        A new DataArray with NaN values filled in the specified direction, leaving the original data unchanged.
     """
-
-    for var in data.var_names.keys():
-        data_vars[var] = extrapolate_deepest_to_bottom(
-            data.ds[data.var_names[var]], data.dim_names["depth"]
-        )
-
-    return data_vars
+    if dim in da.dims:
+        if direction == "forward":
+            return da.ffill(dim=dim)
+        elif direction == "backward":
+            return da.bfill(dim=dim)
+    return da
 
 
 def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -7,11 +7,18 @@ from conftest import calculate_file_hash
 from pathlib import Path
 
 
-def test_boundary_forcing_creation(boundary_forcing):
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "boundary_forcing",
+        "boundary_forcing_with_2d_fill",
+    ],
+)
+def test_boundary_forcing_creation(boundary_forcing_fixture, request):
     """Test the creation of the BoundaryForcing object."""
 
     fname = download_test_data("GLORYS_coarse_test_data.nc")
-
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
     assert boundary_forcing.start_time == datetime(2021, 6, 29)
     assert boundary_forcing.end_time == datetime(2021, 6, 30)
     assert boundary_forcing.source == {
@@ -40,54 +47,61 @@ def test_boundary_forcing_creation(boundary_forcing):
     assert not hasattr(boundary_forcing.ds, "climatology")
 
 
-def test_boundary_forcing_creation_with_bgc(bgc_boundary_forcing_from_climatology):
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "bgc_boundary_forcing_from_climatology",
+        "bgc_boundary_forcing_from_climatology_with_2d_fill",
+    ],
+)
+def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
     """Test the creation of the BoundaryForcing object."""
 
     fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
-    assert bgc_boundary_forcing_from_climatology.start_time == datetime(2021, 6, 29)
-    assert bgc_boundary_forcing_from_climatology.end_time == datetime(2021, 6, 30)
-    assert bgc_boundary_forcing_from_climatology.source == {
+    assert boundary_forcing.start_time == datetime(2021, 6, 29)
+    assert boundary_forcing.end_time == datetime(2021, 6, 30)
+    assert boundary_forcing.source == {
         "path": fname_bgc,
         "name": "CESM_REGRIDDED",
         "climatology": True,
     }
-    assert bgc_boundary_forcing_from_climatology.model_reference_date == datetime(
-        2000, 1, 1
-    )
-    assert bgc_boundary_forcing_from_climatology.boundaries == {
+    assert boundary_forcing.model_reference_date == datetime(2000, 1, 1)
+    assert boundary_forcing.boundaries == {
         "south": True,
         "east": True,
         "north": True,
         "west": True,
     }
 
-    assert bgc_boundary_forcing_from_climatology.ds.source == "CESM_REGRIDDED"
+    assert boundary_forcing.ds.source == "CESM_REGRIDDED"
     for direction in ["south", "east", "north", "west"]:
         for var in ["ALK", "PO4"]:
-            assert f"{var}_{direction}" in bgc_boundary_forcing_from_climatology.ds
+            assert f"{var}_{direction}" in boundary_forcing.ds
+
+    assert len(boundary_forcing.ds.bry_time) == 12
+    assert boundary_forcing.ds.coords["bry_time"].attrs["units"] == "days"
+    assert hasattr(boundary_forcing.ds, "climatology")
 
-    assert len(bgc_boundary_forcing_from_climatology.ds.bry_time) == 12
-    assert (
-        bgc_boundary_forcing_from_climatology.ds.coords["bry_time"].attrs["units"]
-        == "days"
-    )
-    assert hasattr(bgc_boundary_forcing_from_climatology.ds, "climatology")
 
+def test_boundary_forcing_plot(boundary_forcing):
+    """Test plot."""
 
-def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
-    """Test plot and save methods."""
+    boundary_forcing.plot(var_name="temp_south", layer_contours=True)
+    boundary_forcing.plot(var_name="temp_east", layer_contours=True)
+    boundary_forcing.plot(var_name="temp_north", layer_contours=True)
+    boundary_forcing.plot(var_name="temp_west", layer_contours=True)
+    boundary_forcing.plot(var_name="zeta_south")
+    boundary_forcing.plot(var_name="zeta_east")
+    boundary_forcing.plot(var_name="zeta_north")
+    boundary_forcing.plot(var_name="zeta_west")
+    boundary_forcing.plot(var_name="vbar_north")
+    boundary_forcing.plot(var_name="ubar_west")
 
-    boundary_forcing.plot(varname="temp_south", layer_contours=True)
-    boundary_forcing.plot(varname="temp_east", layer_contours=True)
-    boundary_forcing.plot(varname="temp_north", layer_contours=True)
-    boundary_forcing.plot(varname="temp_west", layer_contours=True)
-    boundary_forcing.plot(varname="zeta_south")
-    boundary_forcing.plot(varname="zeta_east")
-    boundary_forcing.plot(varname="zeta_north")
-    boundary_forcing.plot(varname="zeta_west")
-    boundary_forcing.plot(varname="vbar_north")
-    boundary_forcing.plot(varname="ubar_west")
+
+def test_boundary_forcing_save(boundary_forcing, tmp_path):
+    """Test save method."""
 
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -96,9 +110,19 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
+            # Test saving without partitioning and grouping
             saved_filenames = boundary_forcing.save(filepath)
 
+            filepath_str = str(Path(filepath).with_suffix(""))
+            expected_filepath = Path(f"{filepath_str}.nc")
+
+            assert saved_filenames == [expected_filepath]
+            assert expected_filepath.exists()
+            expected_filepath.unlink()
+
+            # Test saving without partitioning but with grouping
+            saved_filenames = boundary_forcing.save(filepath, group=True)
+
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}_202106.nc")
 
@@ -106,8 +130,8 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
             assert expected_filepath.exists()
             expected_filepath.unlink()
 
-            # Test saving with partitioning
-            saved_filenames = boundary_forcing.save(filepath, np_eta=2)
+            # Test saving with partitioning and grouping
+            saved_filenames = boundary_forcing.save(filepath, np_eta=2, group=True)
             expected_filepath_list = [
                 Path(filepath_str + f"_202106.{index}.nc") for index in range(2)
             ]
@@ -119,15 +143,17 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
                 expected_filepath.unlink()
 
 
-def test_bgc_boundary_forcing_plot_save(
-    bgc_boundary_forcing_from_climatology, tmp_path
-):
-    """Test plot and save methods."""
+def test_bgc_boundary_forcing_plot(bgc_boundary_forcing_from_climatology):
+    """Test plot method."""
 
-    bgc_boundary_forcing_from_climatology.plot(varname="ALK_south")
-    bgc_boundary_forcing_from_climatology.plot(varname="ALK_east")
-    bgc_boundary_forcing_from_climatology.plot(varname="ALK_north")
-    bgc_boundary_forcing_from_climatology.plot(varname="ALK_west")
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_south")
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_east")
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_north")
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_west")
+
+
+def test_bgc_boundary_forcing_save(bgc_boundary_forcing_from_climatology, tmp_path):
+    """Test save method."""
 
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -136,9 +162,20 @@ def test_bgc_boundary_forcing_plot_save(
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
+            # Test saving without partitioning and grouping
             saved_filenames = bgc_boundary_forcing_from_climatology.save(filepath)
 
+            filepath_str = str(Path(filepath).with_suffix(""))
+            expected_filepath = Path(f"{filepath_str}.nc")
+            assert saved_filenames == [expected_filepath]
+            assert expected_filepath.exists()
+            expected_filepath.unlink()
+
+            # Test saving without partitioning but with grouping
+            saved_filenames = bgc_boundary_forcing_from_climatology.save(
+                filepath, group=True
+            )
+
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}_clim.nc")
             assert saved_filenames == [expected_filepath]
@@ -147,7 +184,7 @@ def test_bgc_boundary_forcing_plot_save(
 
             # Test saving with partitioning
             saved_filenames = bgc_boundary_forcing_from_climatology.save(
-                filepath, np_xi=2
+                filepath, np_xi=2, group=True
             )
 
             expected_filepath_list = [
@@ -197,9 +234,9 @@ def test_files_have_same_hash(boundary_forcing, tmp_path, use_dask):
     filepath2 = tmp_path / "test2.nc"
 
     boundary_forcing.to_yaml(yaml_filepath)
-    boundary_forcing.save(filepath1)
+    boundary_forcing.save(filepath1, group=True)
     bdry_forcing_from_file = BoundaryForcing.from_yaml(yaml_filepath, use_dask=use_dask)
-    bdry_forcing_from_file.save(filepath2)
+    bdry_forcing_from_file.save(filepath2, group=True)
 
     filepath_str1 = str(Path(filepath1).with_suffix(""))
     filepath_str2 = str(Path(filepath2).with_suffix(""))
@@ -225,9 +262,9 @@ def test_files_have_same_hash_clim(
     filepath2 = tmp_path / "test2.nc"
 
     bgc_boundary_forcing_from_climatology.to_yaml(yaml_filepath)
-    bgc_boundary_forcing_from_climatology.save(filepath1)
+    bgc_boundary_forcing_from_climatology.save(filepath1, group=True)
     bdry_forcing_from_file = BoundaryForcing.from_yaml(yaml_filepath, use_dask=use_dask)
-    bdry_forcing_from_file.save(filepath2)
+    bdry_forcing_from_file.save(filepath2, group=True)
 
     filepath_str1 = str(Path(filepath1).with_suffix(""))
     filepath_str2 = str(Path(filepath2).with_suffix(""))

roms_tools/tests/test_setup/test_data/bgc_surface_forcing.zarr/.zattrs

@@ -2,7 +2,7 @@
     "correct_radiation": "False",
     "end_time": "2020-02-01 00:00:00",
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev138+dirty",
+    "roms_tools_version": "0.1.dev157+dirty",
     "source": "CESM_REGRIDDED",
     "start_time": "2020-02-01 00:00:00",
     "title": "ROMS surface forcing file created by ROMS-Tools",

roms_tools/tests/test_setup/test_data/bgc_surface_forcing.zarr/.zmetadata

@@ -4,7 +4,7 @@
             "correct_radiation": "False",
             "end_time": "2020-02-01 00:00:00",
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev138+dirty",
+            "roms_tools_version": "0.1.dev157+dirty",
             "source": "CESM_REGRIDDED",
             "start_time": "2020-02-01 00:00:00",
             "title": "ROMS surface forcing file created by ROMS-Tools",