roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

roms_tools/setup/topography.py

@@ -1,34 +1,37 @@
+import time
+import logging
 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.download import fetch_topo
-from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
+from roms_tools.setup.utils import handle_boundaries
 import warnings
 from itertools import count
-
-
-def _add_topography_and_mask(
-    ds, topography_source, hmin, smooth_factor=8.0, rmax=0.2
+from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
+from roms_tools.setup.regrid import LateralRegrid
+
+
+def _add_topography(
+    ds,
+    target_coords,
+    topography_source,
+    hmin,
+    smooth_factor=8.0,
+    rmax=0.2,
+    verbose=False,
 ) -> xr.Dataset:
-    """Adds topography and a land/water mask to the dataset based on the provided
-    topography source.
-
-    This function performs the following operations:
-    1. Interpolates topography data onto the desired grid.
-    2. Applies a mask based on ocean depth.
-    3. Smooths the topography globally to reduce grid-scale instabilities.
-    4. Fills enclosed basins with land.
-    5. Smooths the topography locally to ensure the steepness ratio satisfies the rmax criterion.
-    6. Adds topography metadata.
+    """Adds topography to the dataset based on the provided topography source.
 
     Parameters
     ----------
     ds : xr.Dataset
-        The dataset to which topography and the land/water mask will be added.
-    topography_source : str
-        The source of the topography data.
+        The dataset to which topography will be added.
+    topography_source : Dict[str, Union[str, Path]], optional
+        Dictionary specifying the source of the topography data:
+
+        - "name" (str): The name of the topography data source (e.g., "SRTM15").
+        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file. Can be a string or a Path object.
+
+        The default is "ETOPO5", which does not require a path.
     hmin : float
         The minimum allowable depth for the topography.
     smooth_factor : float, optional
@@ -39,89 +42,142 @@ def _add_topography_and_mask(
         The maximum allowable steepness ratio for the topography smoothing.
         This parameter controls the local smoothing of the topography. Smaller values result in
         smoother topography, while larger values preserve more detail. The default is 0.2.
+    verbose: bool, optional
+        Indicates whether to print topography generation steps with timing. Defaults to False.
 
     Returns
     -------
     xr.Dataset
-        The dataset with added topography, mask, and metadata.
+        Updated dataset with added topography and metadata.
     """
 
-    lon = ds.lon_rho.values
-    lat = ds.lat_rho.values
+    if verbose:
+        start_time = time.time()
+    data = _get_topography_data(topography_source)
+    if verbose:
+        logging.info(
+            f"Reading the topography data: {time.time() - start_time:.3f} seconds"
+        )
 
     # 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)
+    hraw = _make_raw_topography(data, target_coords, verbose=verbose)
+    nan_check(hraw)
 
     # smooth topography domain-wide with Gaussian kernel to avoid grid scale instabilities
+    if verbose:
+        start_time = time.time()
     hraw = _smooth_topography_globally(hraw, smooth_factor)
-
-    # fill enclosed basins with land
-    mask = _fill_enclosed_basins(mask.values)
-
-    # adjust mask boundaries by copying values from adjacent cells
-    mask = _handle_boundaries(mask)
-
-    ds["mask_rho"] = xr.DataArray(mask.astype(np.int32), 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)
+    if verbose:
+        logging.info(
+            f"Smoothing the topography globally: {time.time() - start_time:.3f} seconds"
+        )
 
     # 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"].attrs = {
-        "long_name": "Final bathymetry at rho-points",
+        "long_name": "Bathymetry at rho-points",
         "units": "meter",
     }
+    if verbose:
+        logging.info(
+            f"Smoothing the topography locally: {time.time() - start_time:.3f} seconds"
+        )
 
     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."""
+def _get_topography_data(source):
+    """Load topography data based on the specified source.
 
-    topo_ds = fetch_topo(topography_source)
+    Parameters
+    ----------
+    source : dict
+        A dictionary containing the source details (e.g., "name" and "path").
 
-    # 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"
+    Returns
+    -------
+    data : object
+        The loaded topography dataset (ETOPO5 or SRTM15).
+    """
+    kwargs = {"use_dask": False}
+
+    if source["name"] == "ETOPO5":
+        if "path" in source.keys():
+            kwargs["filename"] = source["path"]
+        data = ETOPO5Dataset(**kwargs)
+    elif source["name"] == "SRTM15":
+        kwargs["filename"] = source["path"]
+        data = SRTM15Dataset(**kwargs)
+    else:
+        raise ValueError(
+            'Only "ETOPO5" and "SRTM15" are valid options for topography_source["name"].'
         )
 
-        interp = RegularGridInterpolator(
-            (topo_ds["topo_lat"].values, topo_lon_concatenated.values),
-            topo_concatenated.values,
-            method="linear",
+    return data
+
+
+def _make_raw_topography(
+    data, target_coords, method="linear", verbose=False
+) -> xr.DataArray:
+    """Regrid topography data to match target coordinates.
+
+    Parameters
+    ----------
+    data : object
+        The dataset object containing the topography data.
+    target_coords : object
+        The target coordinates to which the data will be regridded.
+    method : str, optional
+        The regridding method to use, by default "linear".
+    verbose : bool, optional
+        If True, logs the time taken for regridding, by default False.
+
+    Returns
+    -------
+    xr.DataArray
+        The regridded topography data with the sign flipped (bathymetry positive).
+    """
+    data.choose_subdomain(target_coords, buffer_points=3, verbose=verbose)
+
+    if verbose:
+        start_time = time.time()
+    lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+    hraw = lateral_regrid.apply(data.ds[data.var_names["topo"]], method=method)
+    if verbose:
+        logging.info(
+            f"Regridding the topography: {time.time() - start_time:.3f} seconds"
         )
 
-    # Interpolate onto desired domain grid points
-    hraw = interp((lat, lon))
+    # flip sign so that bathmetry is positive
+    hraw = -hraw
 
     return hraw
 
 
 def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
+    """Apply global smoothing to the topography using a Gaussian filter.
+
+    Parameters
+    ----------
+    hraw : xr.DataArray
+        The raw topography data to be smoothed.
+    factor : float
+        The smoothing factor (controls the width of the Gaussian filter).
+
+    Returns
+    -------
+    xr.DataArray
+        The smoothed topography data.
+    """
     # since GCM-Filters assumes periodic domain, we extend the domain by one grid cell in each dimension
     # and set that margin to land
-
     mask = xr.ones_like(hraw)
     margin_mask = xr.concat([mask, 0 * mask.isel(eta_rho=-1)], dim="eta_rho")
     margin_mask = xr.concat(
@@ -148,30 +204,28 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
     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
+def _smooth_topography_locally(h, hmin=5, rmax=0.2):
+    """Smooths topography locally to ensure the slope (r-factor) is below the specified
+    threshold.
 
-    return mask
+    This function applies a logarithmic transformation to the topography and iteratively smooths
+    it in four directions (eta, xi, and two diagonals) until the maximum slope parameter (r) is
+    below `rmax`. A threshold `hmin` is applied to prevent values from going below a minimum height.
 
+    Parameters
+    ----------
+    h : xarray.DataArray
+        The topography data to be smoothed.
+    hmin : float, optional
+        The minimum height threshold. Default is 5.
+    rmax : float, optional
+        The maximum allowable slope parameter (r-factor). Default is 0.2.
 
-def _smooth_topography_locally(h, hmin=5, rmax=0.2):
-    """Smoothes topography locally to satisfy r < rmax."""
+    Returns
+    -------
+    xarray.DataArray
+        The smoothed topography data.
+    """
     # Compute rmax_log
     if rmax > 0.0:
         rmax_log = np.log((1.0 + rmax * 0.9) / (1.0 - rmax * 0.9))
@@ -181,65 +235,90 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     # Apply hmin threshold
     h = xr.where(h < hmin, hmin, h)
 
-    # We will smooth logarithmically
+    # Perform logarithmic transformation of the height field
     h_log = np.log(h / hmin)
 
-    cf1 = 1.0 / 6
-    cf2 = 0.25
+    # Constants for smoothing
+    smoothing_factor_1 = 1.0 / 6
+    smoothing_factor_2 = 0.25
 
+    # Iterate until convergence
     for iter in count():
-        # Compute gradients in domain interior
+        # Compute gradients and smoothing for eta, xi, and diagonal directions
 
-        # in eta-direction
-        cff = h_log.diff("eta_rho").isel(xi_rho=slice(1, -1))
-        cr = np.abs(cff)
+        # Gradient in eta-direction
+        delta_eta = h_log.diff("eta_rho").isel(xi_rho=slice(1, -1))
+        abs_eta_gradient = np.abs(delta_eta)
         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)
+            eta_correction = xr.where(
+                abs_eta_gradient < rmax_log,
+                0,
+                delta_eta * (1 - rmax_log / abs_eta_gradient),
+            )
+
+        # Gradient in xi-direction
+        delta_xi = h_log.diff("xi_rho").isel(eta_rho=slice(1, -1))
+        abs_xi_gradient = np.abs(delta_xi)
         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(
+            xi_correction = xr.where(
+                abs_xi_gradient < rmax_log,
+                0,
+                delta_xi * (1 - rmax_log / abs_xi_gradient),
+            )
+
+        # Gradient in first diagonal direction
+        delta_diag_1 = (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)
+        abs_diag_1_gradient = np.abs(delta_diag_1)
         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(
+            diag_1_correction = xr.where(
+                abs_diag_1_gradient < rmax_log,
+                0,
+                delta_diag_1 * (1 - rmax_log / abs_diag_1_gradient),
+            )
+
+        # Gradient in second diagonal direction
+        delta_diag_2 = (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)
+        abs_diag_2_gradient = np.abs(delta_diag_2)
         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))
+            diag_2_correction = xr.where(
+                abs_diag_2_gradient < rmax_log,
+                0,
+                delta_diag_2 * (1 - rmax_log / abs_diag_2_gradient),
+            )
 
         # 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])
+        h_log[1:-1, 1:-1] += smoothing_factor_1 * (
+            eta_correction[1:, :]
+            - eta_correction[:-1, :]
+            + xi_correction[:, 1:]
+            - xi_correction[:, :-1]
+            + smoothing_factor_2
+            * (
+                diag_1_correction[1:, 1:]
+                - diag_1_correction[:-1, :-1]
+                + diag_2_correction[:-1, 1:]
+                - diag_2_correction[1:, :-1]
+            )
         )
 
         # No gradient at the domain boundaries
-        h_log = _handle_boundaries(h_log)
+        h_log = handle_boundaries(h_log)
 
-        # Update h
+        # Recompute the topography after smoothing
         h = hmin * np.exp(h_log)
+
         # Apply hmin threshold again
         h = xr.where(h < hmin, hmin, h)
 
-        # compute maximum slope parameter r
+        # Compute maximum slope parameter r
         r_eta, r_xi = _compute_rfactor(h)
         rmax0 = np.max([r_eta.max(), r_xi.max()])
         if rmax0 < rmax:
@@ -248,32 +327,24 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     return h
 
 
-def _handle_boundaries(field):
-    """Adjust the boundaries of a 2D field by copying values from adjacent cells.
+def _compute_rfactor(h):
+    """Computes the slope parameter (r-factor) in both horizontal directions.
+
+    The r-factor is calculated as |Δh| / (2h) in the eta and xi directions:
+        - r_eta = |h_i - h_{i-1}| / (h_i + h_{i+1})
+        - r_xi = |h_i - h_{i-1}| / (h_i + h_{i+1})
 
     Parameters
     ----------
-    field : numpy.ndarray or xarray.DataArray
-        A 2D array representing a field (e.g., topography or mask) whose boundary values
-        need to be adjusted.
+    h : xarray.DataArray
+        The topography data.
 
     Returns
     -------
-    field : numpy.ndarray or xarray.DataArray
-        The input field with adjusted boundary values.
+    tuple of xarray.DataArray
+        r_eta : r-factor in the eta direction.
+        r_xi : r-factor in the xi direction.
     """
-
-    field[0, :] = field[1, :]
-    field[-1, :] = field[-2, :]
-    field[:, 0] = field[:, 1]
-    field[:, -1] = field[:, -2]
-
-    return field
-
-
-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)
@@ -286,23 +357,49 @@ def _compute_rfactor(h):
 
 
 def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
-    ds.attrs["topography_source"] = topography_source
+    """Adds topography metadata to the dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset to update.
+    topography_source : dict
+        Dictionary with topography source information (requires 'name' key).
+    smooth_factor : float
+        Smoothing factor (unused in this function).
+    hmin : float
+        Minimum height threshold for smoothing.
+    rmax : float
+        Maximum slope parameter (unused in this function).
+
+    Returns
+    -------
+    xarray.Dataset
+        Updated dataset with added metadata.
+    """
+    ds.attrs["topography_source"] = topography_source["name"]
     ds.attrs["hmin"] = hmin
 
     return ds
 
 
-def _add_velocity_masks(ds):
+def nan_check(hraw):
+    """Checks for NaN values in the topography data.
 
-    # add u- and v-masks
-    ds["mask_u"] = interpolate_from_rho_to_u(
-        ds["mask_rho"], method="multiplicative"
-    ).astype(np.int32)
-    ds["mask_v"] = interpolate_from_rho_to_v(
-        ds["mask_rho"], method="multiplicative"
-    ).astype(np.int32)
-
-    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)"}
+    Parameters
+    ----------
+    hraw : xarray.DataArray
+        Input topography data to check for NaN values.
 
-    return ds
+    Raises
+    ------
+    ValueError
+        If NaN values are found in the data, raises an error with a descriptive message.
+    """
+    error_message = (
+        "NaN values found in regridded topography. This likely occurs because the ROMS grid, including "
+        "a small safety margin for interpolation, is not fully contained within the topography dataset's longitude/latitude range. Please ensure that the "
+        "dataset covers the entire area required by the ROMS grid."
+    )
+    if hraw.isnull().any().values:
+        raise ValueError(error_message)