roms-tools 3.2.0__py3-none-any.whl → 3.4.0__py3-none-any.whl

roms_tools/setup/topography.py

@@ -7,8 +7,12 @@ import gcm_filters
 import numpy as np
 import xarray as xr
 
+from roms_tools.datasets.lat_lon_datasets import (
+    EMODDataset,
+    ETOPO5Dataset,
+    SRTM15Dataset,
+)
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
 from roms_tools.setup.utils import handle_boundaries
 
 
@@ -17,7 +21,7 @@ def add_topography(
     target_coords,
     topography_source,
     hmin,
-    smooth_factor=6.0,
+    smooth_factor=5.0,
     rmax=0.2,
     verbose=False,
 ) -> xr.Dataset:
@@ -77,10 +81,8 @@ def add_topography(
     # smooth topography locally to satisfy r < rmax
     if verbose:
         start_time = time.time()
-    # inserting hraw * mask_rho into this function eliminates any inconsistencies between
-    # the land according to the topography and the land according to the mask; land points
-    # will always be set to hmin
-    ds["h"] = _smooth_topography_locally(hraw * ds["mask_rho"], hmin, rmax)
+
+    ds["h"] = _smooth_topography_locally(hraw, hmin, rmax)
     ds["h"].attrs = {
         "long_name": "Bathymetry at rho-points",
         "units": "meter",
@@ -117,9 +119,12 @@ def _get_topography_data(source):
     elif source["name"] == "SRTM15":
         kwargs["filename"] = source["path"]
         data = SRTM15Dataset(**kwargs)
+    elif source["name"] == "EMOD":
+        kwargs["filename"] = source["path"]
+        data = EMODDataset(**kwargs)
     else:
         raise ValueError(
-            'Only "ETOPO5" and "SRTM15" are valid options for topography_source["name"].'
+            'Only "ETOPO5", "SRTM15" and "EMOD" are valid options for topography_source["name"].'
         )
 
     return data
@@ -149,6 +154,7 @@ def _make_raw_topography(
     data.choose_subdomain(target_coords, buffer_points=3, verbose=verbose)
     # Enforce double precision to ensure reproducibility
     data.convert_to_float64()
+    data.apply_lateral_fill()
 
     if verbose:
         start_time = time.time()

roms_tools/setup/utils.py

@@ -7,9 +7,8 @@ from dataclasses import asdict, fields, is_dataclass
 from datetime import datetime
 from enum import StrEnum
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, TypeAlias
 
-import cftime
 import numba as nb
 import numpy as np
 import pandas as pd
@@ -31,6 +30,8 @@ yaml.SafeDumper.add_multi_representer(
 HEADER_SZ = 96
 HEADER_CHAR = "="
 
+RawDataSource: TypeAlias = dict[str, str | Path | list[str | Path] | bool]
+
 
 def log_the_separator() -> None:
     """Log a separator line using HEADER_CHAR repeated HEADER_SZ times."""
@@ -126,32 +127,6 @@ def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
     return field.fillna(fill_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
-    ----------
-    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 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 DataArray with NaN values filled in the specified direction, leaving the original data unchanged.
-    """
-    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:
     """Assigns climatology dates to the dataset's time dimension.
 
@@ -180,240 +155,6 @@ def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
     return ds
 
 
-def interpolate_cyclic_time(
-    data_array: xr.DataArray,
-    time_dim_name: str,
-    day_of_year: int | float | np.ndarray | xr.DataArray | Sequence[int | float],
-) -> xr.DataArray:
-    """Interpolates a DataArray cyclically across the start and end of the year.
-
-    This function extends the data cyclically by appending the last time step
-    (shifted back by one year) at the beginning and the first time step
-    (shifted forward by one year) at the end. It then performs linear interpolation
-    to match the specified `day_of_year` values.
-
-    Parameters
-    ----------
-    data_array : xr.DataArray
-        The input data array containing a time-like dimension.
-    time_dim_name : str
-        The name of the time dimension in the dataset.
-    day_of_year : Union[int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]]
-        The target day(s) of the year for interpolation. This can be:
-        - A single integer or float representing the day of the year.
-        - A NumPy array or xarray DataArray containing multiple days.
-        - A list or tuple of integers or floats for multiple target days.
-
-    Returns
-    -------
-    xr.DataArray
-        The interpolated DataArray, ensuring cyclic continuity across year boundaries.
-
-    Notes
-    -----
-    - This function is useful for interpolating climatological data, where the time axis
-      represents a repeating annual cycle.
-    - The `day_of_year` values should be within the range [1, 365] or [1, 366] for leap years.
-    """
-    # Concatenate across the beginning and end of the year
-    time_concat = xr.concat(
-        [
-            data_array[time_dim_name][-1] - 365.25,  # Shift last time backward
-            data_array[time_dim_name],
-            data_array[time_dim_name][0] + 365.25,  # Shift first time forward
-        ],
-        dim=time_dim_name,
-    )
-
-    data_array_concat = xr.concat(
-        [
-            data_array.isel(
-                **{time_dim_name: -1}
-            ),  # Append last value at the beginning
-            data_array,
-            data_array.isel(**{time_dim_name: 0}),  # Append first value at the end
-        ],
-        dim=time_dim_name,
-    )
-    data_array_concat[time_dim_name] = time_concat
-
-    # Interpolate to specified times
-    data_array_interpolated = data_array_concat.interp(
-        **{time_dim_name: day_of_year}, method="linear"
-    )
-
-    return data_array_interpolated
-
-
-def interpolate_from_climatology(
-    field: xr.DataArray | xr.Dataset,
-    time_dim_name: str,
-    time: xr.DataArray | pd.DatetimeIndex,
-) -> xr.DataArray | xr.Dataset:
-    """Interpolates a climatological field to specified time points.
-
-    This function interpolates the input `field` based on `day_of_year` values
-    extracted from the provided `time` points. If `field` is an `xarray.Dataset`,
-    interpolation is applied to all its data variables individually.
-
-    Parameters
-    ----------
-    field : xarray.DataArray or xarray.Dataset
-        The input field to be interpolated.
-        - If `field` is an `xarray.DataArray`, it must have a time dimension identified by `time_dim_name`.
-        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along `time_dim_name`.
-        The time dimension is assumed to represent `day_of_year` for climatological purposes.
-    time_dim_name : str
-        The name of the time dimension in `field`. This dimension is used for interpolation.
-    time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation. These are internally converted to `day_of_year`
-        before performing interpolation.
-
-    Returns
-    -------
-    xarray.DataArray or xarray.Dataset
-        The interpolated field, maintaining the same type (`xarray.DataArray` or `xarray.Dataset`)
-        but aligned to the specified `time` values.
-
-    Notes
-    -----
-    - This function assumes that `field` represents a climatological dataset, where time is expressed as `day_of_year` (1-365).
-    - The `time` input is automatically converted to `day_of_year`, so manual conversion is not required before calling this function.
-    """
-
-    def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
-        if isinstance(time, xr.DataArray):
-            # Extract day of year from xarray.DataArray
-            day_of_year = time.dt.dayofyear
-        else:
-            if np.size(time) == 1:
-                # Convert single datetime64 object to pandas.Timestamp
-                date = pd.Timestamp(time)
-                day_of_year = (
-                    date.dayofyear
-                    + (date.hour / 24)
-                    + (date.minute / 1440)
-                    + (date.second / 86400)
-                )
-            else:
-                # Convert each datetime64 object in the array to pandas.Timestamp and compute fractional day of year
-                day_of_year = np.array(
-                    [
-                        pd.Timestamp(t).dayofyear
-                        + (pd.Timestamp(t).hour / 24)
-                        + (pd.Timestamp(t).minute / 1440)
-                        + (pd.Timestamp(t).second / 86400)
-                        for t in time
-                    ]
-                )
-
-        data_array_interpolated = interpolate_cyclic_time(
-            data_array, time_dim_name, day_of_year
-        )
-
-        if np.size(time) == 1:
-            data_array_interpolated = data_array_interpolated.expand_dims(
-                {time_dim_name: 1}
-            )
-        return data_array_interpolated
-
-    if isinstance(field, xr.DataArray):
-        return interpolate_single_field(field)
-    elif isinstance(field, xr.Dataset):
-        interpolated_data_vars = {
-            var: interpolate_single_field(data_array)
-            for var, data_array in field.data_vars.items()
-        }
-        return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
-
-    else:
-        raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
-
-
-def get_time_type(data_array: xr.DataArray) -> str:
-    """Determines the type of time values in the xarray DataArray.
-
-    Parameters
-    ----------
-    data_array : xr.DataArray
-        The xarray DataArray to be checked for time data types.
-
-    Returns
-    -------
-    str
-        A string indicating the type of the time data: 'cftime', 'datetime', or 'int'.
-
-    Raises
-    ------
-    TypeError
-        If the values in the DataArray are not of type numpy.ndarray or list.
-    """
-    # List of cftime datetime types
-    cftime_types = (
-        cftime.DatetimeNoLeap,
-        cftime.DatetimeJulian,
-        cftime.DatetimeGregorian,
-        cftime.Datetime360Day,
-        cftime.DatetimeProlepticGregorian,
-    )
-
-    # Check if any of the coordinate values are of cftime, datetime, or integer type
-    if isinstance(data_array.values, np.ndarray | list):
-        # Check if the data type is numpy datetime64, indicating standard datetime objects
-        if data_array.values.dtype == "datetime64[ns]":
-            return "datetime"
-
-        # Check if any values in the array are instances of cftime types
-        if any(isinstance(value, cftime_types) for value in data_array.values):
-            return "cftime"
-
-        # Check if all values are of integer type (e.g., for indices or time steps)
-        if np.issubdtype(data_array.values.dtype, np.integer):
-            return "int"
-
-        # If none of the above conditions are met, raise a ValueError
-        raise ValueError("Unsupported data type for time values in input dataset.")
-
-    # Handle unexpected types
-    raise TypeError("DataArray values must be of type numpy.ndarray or list.")
-
-
-def convert_cftime_to_datetime(data_array: np.ndarray) -> np.ndarray:
-    """Converts cftime datetime objects to numpy datetime64 objects in a numpy ndarray.
-
-    Parameters
-    ----------
-    data_array : np.ndarray
-        The numpy ndarray containing cftime datetime objects to be converted.
-
-    Returns
-    -------
-    np.ndarray
-        The ndarray with cftime datetimes converted to numpy datetime64 objects.
-
-    Notes
-    -----
-    This function is intended to be used with numpy ndarrays. If you need to convert
-    cftime datetime objects in an xarray.DataArray, please use the appropriate function
-    to handle xarray.DataArray conversions.
-    """
-    # List of cftime datetime types
-    cftime_types = (
-        cftime.DatetimeNoLeap,
-        cftime.DatetimeJulian,
-        cftime.DatetimeGregorian,
-    )
-
-    # Define a conversion function for cftime to numpy datetime64
-    def convert_datetime(dt):
-        if isinstance(dt, cftime_types):
-            # Convert to ISO format and then to nanosecond precision
-            return np.datetime64(dt.isoformat(), "ns")
-        return np.datetime64(dt, "ns")
-
-    return np.vectorize(convert_datetime)(data_array)
-
-
 def get_variable_metadata():
     """Retrieves metadata for commonly used variables in the dataset.
 
@@ -1859,38 +1600,6 @@ def get_boundary_coords():
     return bdry_coords
 
 
-def wrap_longitudes(grid_ds, straddle):
-    """Adjusts longitude values in a dataset to handle dateline crossing.
-
-    Parameters
-    ----------
-    grid_ds : xr.Dataset
-        The dataset containing longitude variables to adjust.
-    straddle : bool
-        If True, adjusts longitudes to the range [-180, 180] for datasets
-        that straddle the dateline. If False, adjusts longitudes to the
-        range [0, 360].
-
-    Returns
-    -------
-    xr.Dataset
-        The dataset with adjusted longitude values.
-    """
-    for lon_dim in ["lon_rho", "lon_u", "lon_v"]:
-        if straddle:
-            grid_ds[lon_dim] = xr.where(
-                grid_ds[lon_dim] > 180,
-                grid_ds[lon_dim] - 360,
-                grid_ds[lon_dim],
-            )
-        else:
-            grid_ds[lon_dim] = xr.where(
-                grid_ds[lon_dim] < 0, grid_ds[lon_dim] + 360, grid_ds[lon_dim]
-            )
-
-    return grid_ds
-
-
 def to_float(val):
     """Convert a value or list of values to float.
 
@@ -1963,3 +1672,103 @@ def validate_names(
         names = names[:max_to_plot]
 
     return names
+
+
+def check_and_set_boundaries(
+    boundaries: dict[str, bool] | None,
+    mask: xr.DataArray,
+) -> dict[str, bool]:
+    """
+    Validate and finalize the `boundaries` dictionary.
+
+    Parameters
+    ----------
+    boundaries : dict[str, bool] or None
+        User-supplied dictionary controlling which boundaries are active.
+        Keys may include any subset of {"south", "east", "north", "west"}.
+        Missing keys will be filled from mask-based defaults.
+        If None, all boundaries are inferred from the land mask.
+
+    mask : xr.DataArray
+        2D land/sea mask on rho-points. Used to determine which boundaries
+        contain at least one ocean point.
+
+    Returns
+    -------
+    dict[str, bool]
+        Completed and validated boundary configuration.
+    """
+    valid_keys = {"south", "east", "north", "west"}
+
+    # --------------------------------------------
+    # Case 1: boundaries not provided → infer them
+    # --------------------------------------------
+    if boundaries is None:
+        inferred = _infer_valid_boundaries_from_mask(mask)
+        logging.info(f"No `boundaries` provided. Using mask-based defaults: {inferred}")
+        return inferred
+
+    # --------------------------------------------
+    # Case 2: boundaries provided → validate
+    # --------------------------------------------
+    if not isinstance(boundaries, dict):
+        raise TypeError(
+            "`boundaries` must be a dict mapping boundary names to booleans."
+        )
+
+    # Unknown keys?
+    unknown_keys = set(boundaries) - valid_keys
+    if unknown_keys:
+        raise ValueError(
+            f"`boundaries` contains invalid keys: {unknown_keys}. "
+            "Allowed keys are: 'south', 'east', 'north', 'west'."
+        )
+
+    # Type-check provided values
+    for key, val in boundaries.items():
+        if not isinstance(val, bool):
+            raise TypeError(f"Boundary '{key}' must be a boolean.")
+
+    # Fill missing boundaries using defaults
+    inferred_defaults = _infer_valid_boundaries_from_mask(mask)
+    completed = boundaries.copy()
+
+    for key in valid_keys:
+        if key not in completed:
+            completed[key] = inferred_defaults[key]
+            logging.info(
+                f"`boundaries[{key!r}]` not provided — defaulting to "
+                f"{inferred_defaults[key]}"
+            )
+
+    logging.info(f"Using boundary configuration: {completed}")
+    return completed
+
+
+def _infer_valid_boundaries_from_mask(mask: xr.DataArray) -> dict[str, bool]:
+    """
+    Determine which grid boundaries contain at least one ocean point.
+
+    Any boundary consisting entirely of land is considered inactive.
+
+    Parameters
+    ----------
+    mask : xr.DataArray
+        2D mask array on rho-points where 1 = ocean, 0 = land.
+
+    Returns
+    -------
+    dict[str, bool]
+        Boolean availability for {south, east, north, west}.
+    """
+    bdry_coords = get_boundary_coords()
+    boundaries = {}
+
+    for direction in ["south", "east", "north", "west"]:
+        coords = bdry_coords["rho"][direction]
+        bdry_mask = mask.isel(**coords)
+
+        # Boundary is valid if ANY ocean point exists
+        boundaries[direction] = bool(bdry_mask.values.any())
+
+    return boundaries

roms_tools/tests/test_analysis/test_cdr_ensemble.py

@@ -86,18 +86,16 @@ def test_extract_efficiency(create_member_ds: xr.Dataset) -> None:
     assert eff_rel.time.attrs.get("long_name") == "time since release start"
 
 
-def test_extract_efficiency_missing_abs_time() -> None:
-    """Test that _extract_efficiency raises an error if 'abs_time' is missing."""
+def test_extract_efficiency_missing_time() -> None:
+    """Test that _extract_efficiency raises an error if 'time' is missing."""
     times = np.array(["2000-01-01", "2000-01-02"], dtype="datetime64[ns]")
     ds = xr.Dataset(
         {"cdr_efficiency": ("time", [0.1, 0.2])},
-        coords={"time": times},  # Note: no 'abs_time' coordinate
+        coords={"abs_time": times},  # Note: no 'time' coordinate
     )
 
     ens = Ensemble.__new__(Ensemble)
-    with pytest.raises(
-        ValueError, match="Dataset must contain an 'abs_time' coordinate."
-    ):
+    with pytest.raises(ValueError, match="Dataset must contain a 'time' coordinate."):
         ens._extract_efficiency(ds)