roms-tools 1.4.2__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,34 +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.
+        A new DataArray with NaN values filled in the specified direction, leaving the original data unchanged.
     """
-    field_interpolated = field.ffill(dim=dim)
-
-    return field_interpolated
+    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:
@@ -503,44 +504,6 @@ def get_variable_metadata():
     return d
 
 
-def get_boundary_info():
-    """This function provides information about the boundary points for the rho, u, and
-    v variables on the grid, specifying the indices for the south, east, north, and west
-    boundaries.
-
-    Returns
-    -------
-    dict
-        A dictionary where keys are variable types ("rho", "u", "v"), and values
-        are nested dictionaries mapping directions ("south", "east", "north", "west")
-        to the corresponding boundary coordinates.
-    """
-
-    # Boundary coordinates
-    bdry_coords = {
-        "rho": {
-            "south": {"eta_rho": 0},
-            "east": {"xi_rho": -1},
-            "north": {"eta_rho": -1},
-            "west": {"xi_rho": 0},
-        },
-        "u": {
-            "south": {"eta_rho": 0},
-            "east": {"xi_u": -1},
-            "north": {"eta_rho": -1},
-            "west": {"xi_u": 0},
-        },
-        "v": {
-            "south": {"eta_v": 0},
-            "east": {"xi_rho": -1},
-            "north": {"eta_v": -1},
-            "west": {"xi_rho": 0},
-        },
-    }
-
-    return bdry_coords
-
-
 def extract_single_value(data):
     """Extracts a single value from an xarray.DataArray or numpy array.
 
@@ -745,3 +708,207 @@ def save_datasets(dataset_list, output_filenames, np_eta=None, np_xi=None):
         saved_filenames.extend(Path(f) for f in partitioned_filenames)
 
     return saved_filenames
+
+
+def get_target_coords(grid, use_coarse_grid=False):
+    """Retrieves longitude and latitude coordinates from the grid, adjusting them based
+    on longitude range.
+
+    Parameters
+    ----------
+    grid : Grid
+        Object representing the grid information used for the model.
+    use_coarse_grid : bool, optional
+        Use coarse grid data if True. Defaults to False.
+
+    Returns
+    -------
+    dict
+        Dictionary containing the longitude, latitude, and angle arrays, along with a boolean indicating
+        if the grid straddles the meridian.
+    """
+    # Select grid variables based on whether the coarse grid is used
+    if use_coarse_grid:
+        lat, lon, angle = (
+            grid.ds.lat_coarse,
+            grid.ds.lon_coarse,
+            grid.ds.angle_coarse,
+        )
+        lat_psi = grid.ds.get("lat_psi_coarse")
+        lon_psi = grid.ds.get("lon_psi_coarse")
+    else:
+        lat, lon, angle = (
+            grid.ds.lat_rho,
+            grid.ds.lon_rho,
+            grid.ds.angle,
+        )
+        lat_psi = grid.ds.get("lat_psi")
+        lon_psi = grid.ds.get("lon_psi")
+
+    # Operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in lontitude away from Greenwich meridian
+    lon = xr.where(lon > 180, lon - 360, lon)
+    if lon_psi is not None:
+        lon_psi = xr.where(lon_psi > 180, lon_psi - 360, lon_psi)
+
+    straddle = True
+    if not grid.straddle and abs(lon).min() > 5:
+        lon = xr.where(lon < 0, lon + 360, lon)
+        if lon_psi is not None:
+            lon_psi = xr.where(lon_psi < 0, lon_psi + 360, lon_psi)
+        straddle = False
+
+    target_coords = {
+        "lat": lat,
+        "lon": lon,
+        "lat_psi": lat_psi,
+        "lon_psi": lon_psi,
+        "angle": angle,
+        "straddle": straddle,
+    }
+
+    return target_coords
+
+
+def rotate_velocities(
+    u: xr.DataArray, v: xr.DataArray, angle: xr.DataArray, interpolate: bool = True
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """Rotate and optionally interpolate velocity components to align with grid
+    orientation.
+
+    Parameters
+    ----------
+    u : xarray.DataArray
+        Zonal (east-west) velocity component at u-points.
+    v : xarray.DataArray
+        Meridional (north-south) velocity component at v-points.
+    angle : xarray.DataArray
+        Grid angle values for rotation.
+    interpolate : bool, optional
+        If True, interpolates rotated velocities to grid points (default is True).
+
+    Returns
+    -------
+    tuple of xarray.DataArray
+        Rotated velocity components (u_rot, v_rot).
+
+    Notes
+    -----
+    - Rotation formulas:
+      - u_rot = u * cos(angle) + v * sin(angle)
+      - v_rot = v * cos(angle) - u * sin(angle)
+    """
+
+    # Rotate velocities to grid orientation
+    u_rot = u * np.cos(angle) + v * np.sin(angle)
+    v_rot = v * np.cos(angle) - u * np.sin(angle)
+
+    # Interpolate to u- and v-points
+    if interpolate:
+        u_rot = interpolate_from_rho_to_u(u_rot)
+        v_rot = interpolate_from_rho_to_v(v_rot)
+
+    return u_rot, v_rot
+
+
+def compute_barotropic_velocity(
+    vel: xr.DataArray, interface_depth: xr.DataArray
+) -> xr.DataArray:
+    """Compute barotropic (depth-averaged) velocity from 3D velocity.
+
+    Assumes `vel` and `interface_depth` are at the same horizontal grid location.
+
+    Parameters
+    ----------
+    vel : xarray.DataArray
+        Velocity components (zonal and meridional) at u- and v-points.
+    interface_depth : xarray.DataArray
+        Depth values for computing layer thickness.
+
+    Returns
+    -------
+    xarray.DataArray
+        Depth-averaged velocity (`vel_bar`).
+
+    Notes
+    -----
+    Computed as:
+      - `vel_bar` = sum(dz * vel) / sum(dz)
+    """
+
+    # Layer thickness
+    dz = -interface_depth.diff(dim="s_w")
+    dz = dz.rename({"s_w": "s_rho"})
+
+    vel_bar = (dz * vel).sum(dim="s_rho") / dz.sum(dim="s_rho")
+
+    return vel_bar
+
+
+def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
+    """Transpose the dimensions of an xarray.DataArray to ensure that 'time', any
+    dimension starting with 's_', 'eta_', and 'xi_' are ordered first, followed by the
+    remaining dimensions in their original order.
+
+    Parameters
+    ----------
+    da : xarray.DataArray
+        The input DataArray whose dimensions are to be reordered.
+
+    Returns
+    -------
+    xarray.DataArray
+        The DataArray with dimensions reordered so that 'time', 's_*', 'eta_*',
+        and 'xi_*' are first, in that order, if they exist.
+    """
+
+    # List of preferred dimension patterns
+    preferred_order = ["time", "s_", "eta_", "xi_"]
+
+    # Get the existing dimensions in the DataArray
+    dims = list(da.dims)
+
+    # Collect dimensions that match any of the preferred patterns
+    matched_dims = []
+    for pattern in preferred_order:
+        # Find dimensions that start with the pattern
+        matched_dims += [dim for dim in dims if dim.startswith(pattern)]
+
+    # Create a new order: first the matched dimensions, then the rest
+    remaining_dims = [dim for dim in dims if dim not in matched_dims]
+    new_order = matched_dims + remaining_dims
+
+    # Transpose the DataArray to the new order
+    transposed_da = da.transpose(*new_order)
+
+    return transposed_da
+
+
+def get_vector_pairs(variable_info):
+    """Extracts all unique vector pairs from the variable_info dictionary.
+
+    Parameters
+    ----------
+    variable_info : dict
+        Dictionary containing variable information, including location,
+        whether it's a vector, and its vector pair.
+
+    Returns
+    -------
+    list of tuples
+        List of unique vector pairs, where each tuple contains the names of
+        the two vector components (e.g., ("u", "v")).
+    """
+    vector_pairs = []
+    processed = set()  # Track variables that have already been paired
+
+    for var_name, var_info in variable_info.items():
+        if var_info["is_vector"] and var_name not in processed:
+            vector_pair = var_info["vector_pair"]
+
+            # Ensure the vector_pair exists in the dictionary and has not been processed
+            if vector_pair and vector_pair in variable_info:
+                vector_pairs.append((var_name, vector_pair))
+                # Mark both the variable and its pair as processed
+                processed.update([var_name, vector_pair])
+
+    return vector_pairs

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_boundary_forcing_from_climatology.zarr/.zattrs

@@ -3,7 +3,7 @@
     "end_time": "2021-06-30 00:00:00",
     "hc": 250.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev152+dirty",
+    "roms_tools_version": "0.1.dev157+dirty",
     "source": "CESM_REGRIDDED",
     "start_time": "2021-06-29 00:00:00",
     "theta_b": 2.0,

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

@@ -5,7 +5,7 @@
             "end_time": "2021-06-30 00:00:00",
             "hc": 250.0,
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev152+dirty",
+            "roms_tools_version": "0.1.dev157+dirty",
             "source": "CESM_REGRIDDED",
             "start_time": "2021-06-29 00:00:00",
             "theta_b": 2.0,