roms-tools 1.5.0__py3-none-any.whl → 1.6.0__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,57 @@ 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(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")
+    assert len(boundary_forcing.ds.bry_time) == 12
+    assert boundary_forcing.ds.coords["bry_time"].attrs["units"] == "days"
+    assert hasattr(boundary_forcing.ds, "climatology")
 
 
 def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
     """Test plot and save methods."""
 
-    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")
+    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")
 
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -124,10 +134,10 @@ def test_bgc_boundary_forcing_plot_save(
 ):
     """Test plot and save methods."""
 
-    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")
 
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture

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",

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

@@ -3,7 +3,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_from_climatology.zarr/.zmetadata

@@ -5,7 +5,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/boundary_forcing.zarr/.zmetadata

@@ -38,7 +38,10 @@
             "_ARRAY_DIMENSIONS": [
                 "bry_time"
             ],
+            "axis": "T",
             "calendar": "gregorian",
+            "long_name": "Time",
+            "standard_name": "time",
             "units": "seconds since 1970-01-01"
         },
         "bry_time/.zarray": {
@@ -889,6 +892,7 @@
             ],
             "coordinates": "abs_time",
             "long_name": "eastern boundary sea surface height",
+            "standard_name": "sea_surface_height_above_geoid",
             "units": "m"
         },
         "zeta_north/.zarray": {
@@ -920,6 +924,7 @@
             ],
             "coordinates": "abs_time",
             "long_name": "northern boundary sea surface height",
+            "standard_name": "sea_surface_height_above_geoid",
             "units": "m"
         },
         "zeta_south/.zarray": {
@@ -951,6 +956,7 @@
             ],
             "coordinates": "abs_time",
             "long_name": "southern boundary sea surface height",
+            "standard_name": "sea_surface_height_above_geoid",
             "units": "m"
         },
         "zeta_west/.zarray": {
@@ -982,6 +988,7 @@
             ],
             "coordinates": "abs_time",
             "long_name": "western boundary sea surface height",
+            "standard_name": "sea_surface_height_above_geoid",
             "units": "m"
         }
     },