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

roms_tools/setup/utils.py

@@ -8,8 +8,7 @@ from pathlib import Path
 
 
 def nan_check(field, mask) -> None:
-    """
-    Checks for NaN values at wet points in the field.
+    """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.
     If any NaN values are found at these wet points, a ValueError is raised.
@@ -28,7 +27,6 @@ def nan_check(field, mask) -> None:
     ValueError
         If the field contains NaN values at any of the wet points indicated by the mask.
         The error message will explain the potential cause and suggest ensuring the dataset's coverage.
-
     """
 
     # Replace values in field with 0 where mask is not 1
@@ -44,8 +42,7 @@ def nan_check(field, mask) -> None:
 
 
 def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
-    """
-    Replace NaN values in the field with a specified fill value.
+    """Replace NaN values in the field with a specified fill value.
 
     This function replaces any NaN values in the input field with the provided fill value.
 
@@ -66,9 +63,7 @@ def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
 
 
 def interpolate_from_rho_to_u(field, method="additive"):
-
-    """
-    Interpolates the given field from rho points to u points.
+    """Interpolates the given field from rho points to u points.
 
     This function performs an interpolation from the rho grid (cell centers) to the u grid
     (cell edges in the xi direction). Depending on the chosen method, it either averages
@@ -114,9 +109,7 @@ def interpolate_from_rho_to_u(field, method="additive"):
 
 
 def interpolate_from_rho_to_v(field, method="additive"):
-
-    """
-    Interpolates the given field from rho points to v points.
+    """Interpolates the given field from rho points to v points.
 
     This function performs an interpolation from the rho grid (cell centers) to the v grid
     (cell edges in the eta direction). Depending on the chosen method, it either averages
@@ -164,8 +157,8 @@ def interpolate_from_rho_to_v(field, method="additive"):
 
 
 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.
+    """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.
@@ -187,16 +180,43 @@ def extrapolate_deepest_to_bottom(field: xr.DataArray, dim: str) -> 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.
     """
-    field_interpolated = field.ffill(dim=dim)
 
-    return field_interpolated
+    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
 
 
 def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
-    """
-    Assigns climatology dates to the dataset's time dimension.
+    """Assigns climatology dates to the dataset's time dimension.
 
     This function updates the dataset's time coordinates to reflect climatological dates.
     It defines fixed day increments for each month and assigns these to the specified time dimension.
@@ -213,7 +233,6 @@ def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
     -------
     xr.Dataset
         The updated xarray Dataset with climatological dates assigned to the specified time dimension.
-
     """
     # Define the days in each month and convert to timedelta
     increments = [15, 30, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30]
@@ -229,8 +248,7 @@ def interpolate_from_climatology(
     time_dim_name: str,
     time: Union[xr.DataArray, pd.DatetimeIndex],
 ) -> Union[xr.DataArray, xr.Dataset]:
-    """
-    Interpolates the given field temporally based on the specified time points.
+    """Interpolates the given field temporally based on the specified time points.
 
     If `field` is an xarray.Dataset, this function applies the interpolation to all data variables in the dataset.
 
@@ -305,8 +323,7 @@ def interpolate_from_climatology(
 
 
 def get_time_type(data_array: xr.DataArray) -> str:
-    """
-    Determines the type of time values in the xarray DataArray.
+    """Determines the type of time values in the xarray DataArray.
 
     Parameters
     ----------
@@ -354,8 +371,7 @@ def get_time_type(data_array: xr.DataArray) -> str:
 
 
 def convert_cftime_to_datetime(data_array: np.ndarray) -> np.ndarray:
-    """
-    Converts cftime datetime objects to numpy datetime64 objects in a numpy ndarray.
+    """Converts cftime datetime objects to numpy datetime64 objects in a numpy ndarray.
 
     Parameters
     ----------
@@ -391,8 +407,7 @@ def convert_cftime_to_datetime(data_array: np.ndarray) -> np.ndarray:
 
 
 def get_variable_metadata():
-    """
-    Retrieves metadata for commonly used variables in the dataset.
+    """Retrieves metadata for commonly used variables in the dataset.
 
     This function returns a dictionary containing the metadata for various variables, including long names
     and units for each variable.
@@ -401,7 +416,6 @@ def get_variable_metadata():
     -------
     dict of str: dict
         Dictionary where keys are variable names and values are dictionaries with "long_name" and "units" keys.
-
     """
 
     d = {
@@ -518,49 +532,8 @@ 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.
+    """Extracts a single value from an xarray.DataArray or numpy array.
 
     Parameters
     ----------
@@ -589,8 +562,8 @@ def extract_single_value(data):
 
 
 def group_dataset(ds, filepath):
-    """
-    Group the dataset into monthly or yearly subsets based on the frequency of the data.
+    """Group the dataset into monthly or yearly subsets based on the frequency of the
+    data.
 
     Parameters
     ----------
@@ -643,8 +616,7 @@ def group_dataset(ds, filepath):
 
 
 def group_by_month(ds, filepath):
-    """
-    Group the dataset by month and generate filenames with 'YYYYMM' format.
+    """Group the dataset by month and generate filenames with 'YYYYMM' format.
 
     Parameters
     ----------
@@ -681,8 +653,7 @@ def group_by_month(ds, filepath):
 
 
 def group_by_year(ds, filepath):
-    """
-    Group the dataset by year and generate filenames with 'YYYY' format.
+    """Group the dataset by year and generate filenames with 'YYYY' format.
 
     Parameters
     ----------
@@ -715,8 +686,7 @@ def group_by_year(ds, filepath):
 
 
 def save_datasets(dataset_list, output_filenames, np_eta=None, np_xi=None):
-    """
-    Save the list of datasets to netCDF4 files, with optional spatial partitioning.
+    """Save the list of datasets to netCDF4 files, with optional spatial partitioning.
 
     Parameters
     ----------
@@ -766,3 +736,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/setup/vertical_coordinate.py

@@ -3,8 +3,8 @@ import xarray as xr
 
 
 def compute_cs(sigma, theta_s, theta_b):
-    """
-    Compute the S-coordinate stretching curves according to Shchepetkin and McWilliams (2009).
+    """Compute the S-coordinate stretching curves according to Shchepetkin and
+    McWilliams (2009).
 
     Parameters
     ----------
@@ -37,8 +37,7 @@ def compute_cs(sigma, theta_s, theta_b):
 
 
 def sigma_stretch(theta_s, theta_b, N, type):
-    """
-    Compute sigma and stretching curves based on the type and parameters.
+    """Compute sigma and stretching curves based on the type and parameters.
 
     Parameters
     ----------
@@ -80,8 +79,7 @@ def sigma_stretch(theta_s, theta_b, N, type):
 
 
 def compute_depth(zeta, h, hc, cs, sigma):
-    """
-    Compute the depth at different sigma levels.
+    """Compute the depth at different sigma levels.
 
     Parameters
     ----------

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -8,9 +8,7 @@ from pathlib import Path
 
 
 def test_boundary_forcing_creation(boundary_forcing):
-    """
-    Test the creation of the BoundaryForcing object.
-    """
+    """Test the creation of the BoundaryForcing object."""
 
     fname = download_test_data("GLORYS_coarse_test_data.nc")
 
@@ -43,9 +41,7 @@ def test_boundary_forcing_creation(boundary_forcing):
 
 
 def test_boundary_forcing_creation_with_bgc(bgc_boundary_forcing_from_climatology):
-    """
-    Test the creation of the BoundaryForcing object.
-    """
+    """Test the creation of the BoundaryForcing object."""
 
     fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
 
@@ -80,9 +76,7 @@ def test_boundary_forcing_creation_with_bgc(bgc_boundary_forcing_from_climatolog
 
 
 def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
 
     boundary_forcing.plot(varname="temp_south", layer_contours=True)
     boundary_forcing.plot(varname="temp_east", layer_contours=True)
@@ -128,9 +122,7 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
 def test_bgc_boundary_forcing_plot_save(
     bgc_boundary_forcing_from_climatology, tmp_path
 ):
-    """
-    Test plot and save methods.
-    """
+    """Test plot and save methods."""
 
     bgc_boundary_forcing_from_climatology.plot(varname="ALK_south")
     bgc_boundary_forcing_from_climatology.plot(varname="ALK_east")
@@ -176,7 +168,8 @@ def test_bgc_boundary_forcing_plot_save(
     ],
 )
 def test_roundtrip_yaml(bdry_forcing_fixture, request, tmp_path, use_dask):
-    """Test that creating a BoundaryForcing object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+    """Test that creating a BoundaryForcing object, saving its parameters to yaml file,
+    and re-opening yaml file creates the same object."""
 
     bdry_forcing = request.getfixturevalue(bdry_forcing_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,