roms-tools 0.0.6__py3-none-any.whl → 0.20__py3-none-any.whl

roms_tools/setup/topography.py

@@ -0,0 +1,257 @@
+import xarray as xr
+import numpy as np
+import gcm_filters
+from scipy.interpolate import RegularGridInterpolator
+from scipy.ndimage import label
+from roms_tools.setup.datasets import fetch_topo
+from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
+import warnings
+from itertools import count
+
+
+def _add_topography_and_mask(
+    ds, topography_source, smooth_factor, hmin, rmax
+) -> xr.Dataset:
+    lon = ds.lon_rho.values
+    lat = ds.lat_rho.values
+
+    # interpolate topography onto desired grid
+    hraw = _make_raw_topography(lon, lat, topography_source)
+    hraw = xr.DataArray(data=hraw, dims=["eta_rho", "xi_rho"])
+
+    # Mask is obtained by finding locations where ocean depth is positive
+    mask = xr.where(hraw > 0, 1.0, 0.0)
+
+    # smooth topography domain-wide with Gaussian kernel to avoid grid scale instabilities
+    ds["hraw"] = _smooth_topography_globally(hraw, mask, smooth_factor)
+    ds["hraw"].attrs = {
+        "long_name": "Working bathymetry at rho-points",
+        "source": f"Raw bathymetry from {topography_source} (smoothing diameter {smooth_factor})",
+        "units": "meter",
+    }
+
+    # fill enclosed basins with land
+    mask = _fill_enclosed_basins(mask.values)
+    ds["mask_rho"] = xr.DataArray(mask, dims=("eta_rho", "xi_rho"))
+    ds["mask_rho"].attrs = {
+        "long_name": "Mask at rho-points",
+        "units": "land/water (0/1)",
+    }
+
+    ds = _add_velocity_masks(ds)
+
+    # smooth topography locally to satisfy r < rmax
+    ds["h"] = _smooth_topography_locally(ds["hraw"] * ds["mask_rho"], hmin, rmax)
+    ds["h"].attrs = {
+        "long_name": "Final bathymetry at rho-points",
+        "units": "meter",
+    }
+
+    ds = _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax)
+
+    return ds
+
+
+def _make_raw_topography(lon, lat, topography_source) -> np.ndarray:
+    """
+    Given a grid of (lon, lat) points, fetch the topography file and interpolate height values onto the desired grid.
+    """
+
+    topo_ds = fetch_topo(topography_source)
+
+    # the following will depend on the topography source
+    if topography_source == "ETOPO5":
+        topo_lon = topo_ds["topo_lon"].copy()
+        # Modify longitude values where necessary
+        topo_lon = xr.where(topo_lon < 0, topo_lon + 360, topo_lon)
+        topo_lon_minus360 = topo_lon - 360
+        topo_lon_plus360 = topo_lon + 360
+        # Concatenate along the longitude axis
+        topo_lon_concatenated = xr.concat(
+            [topo_lon_minus360, topo_lon, topo_lon_plus360], dim="lon"
+        )
+        topo_concatenated = xr.concat(
+            [-topo_ds["topo"], -topo_ds["topo"], -topo_ds["topo"]], dim="lon"
+        )
+
+        interp = RegularGridInterpolator(
+            (topo_ds["topo_lat"].values, topo_lon_concatenated.values),
+            topo_concatenated.values,
+            method="linear",
+        )
+
+    # Interpolate onto desired domain grid points
+    hraw = interp((lat, lon))
+
+    return hraw
+
+
+def _smooth_topography_globally(hraw, wet_mask, factor) -> xr.DataArray:
+    # since GCM-Filters assumes periodic domain, we extend the domain by one grid cell in each dimension
+    # and set that margin to land
+    margin_mask = xr.concat([wet_mask, 0 * wet_mask.isel(eta_rho=-1)], dim="eta_rho")
+    margin_mask = xr.concat(
+        [margin_mask, 0 * margin_mask.isel(xi_rho=-1)], dim="xi_rho"
+    )
+
+    # we choose a Gaussian filter kernel corresponding to a Gaussian with standard deviation factor/sqrt(12);
+    # this standard deviation matches the standard deviation of a boxcar kernel with total width equal to factor.
+    filter = gcm_filters.Filter(
+        filter_scale=factor,
+        dx_min=1,
+        filter_shape=gcm_filters.FilterShape.GAUSSIAN,
+        grid_type=gcm_filters.GridType.REGULAR_WITH_LAND,
+        grid_vars={"wet_mask": margin_mask},
+    )
+    hraw_extended = xr.concat([hraw, hraw.isel(eta_rho=-1)], dim="eta_rho")
+    hraw_extended = xr.concat(
+        [hraw_extended, hraw_extended.isel(xi_rho=-1)], dim="xi_rho"
+    )
+
+    hsmooth = filter.apply(hraw_extended, dims=["eta_rho", "xi_rho"])
+    hsmooth = hsmooth.isel(eta_rho=slice(None, -1), xi_rho=slice(None, -1))
+
+    return hsmooth
+
+
+def _fill_enclosed_basins(mask) -> np.ndarray:
+    """
+    Fills in enclosed basins with land
+    """
+
+    # Label connected regions in the mask
+    reg, nreg = label(mask)
+    # Find the largest region
+    lint = 0
+    lreg = 0
+    for ireg in range(nreg):
+        int_ = np.sum(reg == ireg)
+        if int_ > lint and mask[reg == ireg].sum() > 0:
+            lreg = ireg
+            lint = int_
+
+    # Remove regions other than the largest one
+    for ireg in range(nreg):
+        if ireg != lreg:
+            mask[reg == ireg] = 0
+
+    return mask
+
+
+def _smooth_topography_locally(h, hmin=5, rmax=0.2):
+    """
+    Smoothes topography locally to satisfy r < rmax
+    """
+    # Compute rmax_log
+    if rmax > 0.0:
+        rmax_log = np.log((1.0 + rmax * 0.9) / (1.0 - rmax * 0.9))
+    else:
+        rmax_log = 0.0
+
+    # Apply hmin threshold
+    h = xr.where(h < hmin, hmin, h)
+
+    # We will smooth logarithmically
+    h_log = np.log(h / hmin)
+
+    cf1 = 1.0 / 6
+    cf2 = 0.25
+
+    for iter in count():
+        # Compute gradients in domain interior
+
+        # in eta-direction
+        cff = h_log.diff("eta_rho").isel(xi_rho=slice(1, -1))
+        cr = np.abs(cff)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")  # Ignore division by zero warning
+            Op1 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
+
+        # in xi-direction
+        cff = h_log.diff("xi_rho").isel(eta_rho=slice(1, -1))
+        cr = np.abs(cff)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")  # Ignore division by zero warning
+            Op2 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
+
+        # in diagonal direction
+        cff = (h_log - h_log.shift(eta_rho=1, xi_rho=1)).isel(
+            eta_rho=slice(1, None), xi_rho=slice(1, None)
+        )
+        cr = np.abs(cff)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")  # Ignore division by zero warning
+            Op3 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
+
+        # in the other diagonal direction
+        cff = (h_log.shift(eta_rho=1) - h_log.shift(xi_rho=1)).isel(
+            eta_rho=slice(1, None), xi_rho=slice(1, None)
+        )
+        cr = np.abs(cff)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")  # Ignore division by zero warning
+            Op4 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
+
+        # Update h_log in domain interior
+        h_log[1:-1, 1:-1] += cf1 * (
+            Op1[1:, :]
+            - Op1[:-1, :]
+            + Op2[:, 1:]
+            - Op2[:, :-1]
+            + cf2 * (Op3[1:, 1:] - Op3[:-1, :-1] + Op4[:-1, 1:] - Op4[1:, :-1])
+        )
+
+        # No gradient at the domain boundaries
+        h_log[0, :] = h_log[1, :]
+        h_log[-1, :] = h_log[-2, :]
+        h_log[:, 0] = h_log[:, 1]
+        h_log[:, -1] = h_log[:, -2]
+
+        # Update h
+        h = hmin * np.exp(h_log)
+        # Apply hmin threshold again
+        h = xr.where(h < hmin, hmin, h)
+
+        # compute maximum slope parameter r
+        r_eta, r_xi = _compute_rfactor(h)
+        rmax0 = np.max([r_eta.max(), r_xi.max()])
+        if rmax0 < rmax:
+            break
+
+    return h
+
+
+def _compute_rfactor(h):
+    """
+    Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid directions.
+    """
+    # compute r_{i-1/2} = |h_i - h_{i-1}| / (h_i + h_{i+1})
+    r_eta = np.abs(h.diff("eta_rho")) / (h + h.shift(eta_rho=1)).isel(
+        eta_rho=slice(1, None)
+    )
+    r_xi = np.abs(h.diff("xi_rho")) / (h + h.shift(xi_rho=1)).isel(
+        xi_rho=slice(1, None)
+    )
+
+    return r_eta, r_xi
+
+
+def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
+    ds.attrs["topography_source"] = topography_source
+    ds.attrs["smooth_factor"] = smooth_factor
+    ds.attrs["hmin"] = hmin
+    ds.attrs["rmax"] = rmax
+
+    return ds
+
+
+def _add_velocity_masks(ds):
+
+    # add u- and v-masks
+    ds["mask_u"] = interpolate_from_rho_to_u(ds["mask_rho"], method="multiplicative")
+    ds["mask_v"] = interpolate_from_rho_to_v(ds["mask_rho"], method="multiplicative")
+
+    ds["mask_u"].attrs = {"long_name": "Mask at u-points", "units": "land/water (0/1)"}
+    ds["mask_v"].attrs = {"long_name": "Mask at v-points", "units": "land/water (0/1)"}
+
+    return ds

roms_tools/setup/utils.py

@@ -0,0 +1,162 @@
+import xarray as xr
+
+
+def nan_check(field, mask) -> 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.
+    If any NaN values are found at these wet points, a ValueError is raised.
+
+    Parameters
+    ----------
+    field : array-like
+        The data array to be checked for NaN values. This is typically an xarray.DataArray or numpy array.
+
+    mask : array-like
+        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`.
+
+    Raises
+    ------
+    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
+    da = xr.where(mask == 1, field, 0)
+
+    # Check if any NaN values exist in the modified field
+    if da.isnull().any().values:
+        raise ValueError(
+            "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."
+        )
+
+
+def interpolate_from_rho_to_u(field, method="additive"):
+
+    """
+    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
+    (additive) or multiplies (multiplicative) the field values between adjacent rho points
+    along the xi dimension. It also handles the removal of unnecessary coordinate variables
+    and updates the dimensions accordingly.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
+        named "xi_rho".
+
+    method : str, optional, default='additive'
+        The method to use for interpolation. Options are:
+        - 'additive': Average the field values between adjacent rho points.
+        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
+          binary masks.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the u grid with the dimension "xi_u".
+    """
+
+    if method == "additive":
+        field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
+            xi_rho=slice(1, None)
+        )
+    elif method == "multiplicative":
+        field_interpolated = (field * field.shift(xi_rho=1)).isel(xi_rho=slice(1, None))
+    else:
+        raise NotImplementedError(f"Unsupported method '{method}' specified.")
+
+    if "lat_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lat_rho"])
+    if "lon_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lon_rho"])
+
+    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
+
+    return field_interpolated
+
+
+def interpolate_from_rho_to_v(field, method="additive"):
+
+    """
+    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
+    (additive) or multiplies (multiplicative) the field values between adjacent rho points
+    along the eta dimension. It also handles the removal of unnecessary coordinate variables
+    and updates the dimensions accordingly.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
+        named "eta_rho".
+
+    method : str, optional, default='additive'
+        The method to use for interpolation. Options are:
+        - 'additive': Average the field values between adjacent rho points.
+        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
+          binary masks.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the v grid with the dimension "eta_v".
+    """
+
+    if method == "additive":
+        field_interpolated = 0.5 * (field + field.shift(eta_rho=1)).isel(
+            eta_rho=slice(1, None)
+        )
+    elif method == "multiplicative":
+        field_interpolated = (field * field.shift(eta_rho=1)).isel(
+            eta_rho=slice(1, None)
+        )
+    else:
+        raise NotImplementedError(f"Unsupported method '{method}' specified.")
+
+    if "lat_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lat_rho"])
+    if "lon_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lon_rho"])
+
+    field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
+
+    return field_interpolated
+
+
+def extrapolate_deepest_to_bottom(field: xr.DataArray, dim: str) -> xr.DataArray:
+    """
+    Extrapolate the deepest non-NaN values to the bottom along a specified dimension.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array containing NaN values that need to be filled. This array
+        should have at least one dimension named by `dim`.
+    dim : str
+        The name of the dimension along which to perform the interpolation and extrapolation.
+        Typically, this would be a vertical dimension such as 'depth' or 's_rho'.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        A new data array with NaN values along the specified dimension filled by nearest
+        neighbor interpolation and extrapolation to the bottom. The original data array is not modified.
+
+    """
+    field_interpolated = field.interpolate_na(
+        dim=dim, method="nearest", fill_value="extrapolate"
+    )
+
+    return field_interpolated