roms-tools 1.6.2__py3-none-any.whl → 2.0.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,81 +42,93 @@ 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):
 
-    topo_ds = fetch_topo(topography_source)
+    kwargs = {"use_dask": False}
 
-    # 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"
+    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:
+
+    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
 
@@ -148,28 +163,6 @@ 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
-
-    return mask
-
-
 def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     """Smoothes topography locally to satisfy r < rmax."""
     # Compute rmax_log
@@ -232,7 +225,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         )
 
         # No gradient at the domain boundaries
-        h_log = _handle_boundaries(h_log)
+        h_log = handle_boundaries(h_log)
 
         # Update h
         h = hmin * np.exp(h_log)
@@ -248,29 +241,6 @@ 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.
-
-    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.
-
-    Returns
-    -------
-    field : numpy.ndarray or xarray.DataArray
-        The input field with adjusted boundary values.
-    """
-
-    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."""
@@ -286,23 +256,17 @@ def _compute_rfactor(h):
 
 
 def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
-    ds.attrs["topography_source"] = topography_source
+    ds.attrs["topography_source"] = topography_source["name"]
     ds.attrs["hmin"] = hmin
 
     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"
-    ).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)"}
-
-    return ds
+def nan_check(hraw):
+    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)

roms_tools/setup/utils.py

@@ -1,10 +1,14 @@
 import xarray as xr
 import numpy as np
-from typing import Union
+from typing import Union, Any, Dict, Type
 import pandas as pd
 import cftime
 from roms_tools.utils import partition
 from pathlib import Path
+from datetime import datetime
+from dataclasses import fields, asdict
+import importlib.metadata
+import yaml
 
 
 def nan_check(field, mask, error_message=None) -> None:
@@ -37,7 +41,7 @@ def nan_check(field, mask, error_message=None) -> None:
     da = xr.where(mask == 1, field, 0)
     if error_message is None:
         error_message = (
-            "NaN values found in interpolated field. This likely occurs because the ROMS grid, including "
+            "NaN values found in regridded 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."
         )
@@ -103,10 +107,10 @@ def interpolate_from_rho_to_u(field, method="additive"):
     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"])
+    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
+    for var in vars_to_drop:
+        if var in field_interpolated.coords:
+            field_interpolated = field_interpolated.drop_vars(var)
 
     field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
 
@@ -151,10 +155,10 @@ def interpolate_from_rho_to_v(field, method="additive"):
     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"])
+    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
+    for var in vars_to_drop:
+        if var in field_interpolated.coords:
+            field_interpolated = field_interpolated.drop_vars(var)
 
     field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
 
@@ -729,27 +733,27 @@ def get_target_coords(grid, use_coarse_grid=False):
     """
     # Select grid variables based on whether the coarse grid is used
     if use_coarse_grid:
-        lat, lon, angle, mask = (
-            grid.ds.lat_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
-            grid.ds.lon_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
-            grid.ds.angle_coarse.rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            ),
-            grid.ds.mask_coarse.rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            ),
+        lat = grid.ds.lat_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
         )
+        lon = grid.ds.lon_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+        )
+        angle = grid.ds.angle_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+        )
+        mask = grid.ds.get("mask_coarse")
+        if mask is not None:
+            mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
 
     else:
-        lat, lon, angle, mask = (
-            grid.ds.lat_rho,
-            grid.ds.lon_rho,
-            grid.ds.angle,
-            grid.ds.mask_rho,
-        )
+        lat = grid.ds.lat_rho
+        lon = grid.ds.lon_rho
+        angle = grid.ds.angle
+        mask = grid.ds.get("mask_rho")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -921,3 +925,281 @@ def get_vector_pairs(variable_info):
                 processed.update([var_name, vector_pair])
 
     return vector_pairs
+
+
+def gc_dist(lon1, lat1, lon2, lat2):
+    """Calculate the great circle distance between two points on the Earth's surface.
+    Latitude and longitude must be provided in degrees (they will be converted to
+    radians).
+
+    The function uses the Haversine formula to compute the shortest distance
+    along the surface of a sphere (Earth), assuming the Earth is a perfect sphere.
+
+    Parameters
+    ----------
+    lon1, lat1 : float
+        Longitude and latitude of the first point in degrees.
+    lon2, lat2 : float
+        Longitude and latitude of the second point in degrees.
+
+    Returns
+    -------
+    dis : float
+        The great circle distance between the two points in meters.
+        This is the shortest distance along the surface of a sphere (Earth).
+
+    Notes
+    -----
+    The radius of the Earth is taken to be 6371315 meters.
+    """
+    # Convert degrees to radians
+    d2r = np.pi / 180
+    lon1 = lon1 * d2r
+    lat1 = lat1 * d2r
+    lon2 = lon2 * d2r
+    lat2 = lat2 * d2r
+
+    # Difference in latitudes and longitudes
+    dlat = lat2 - lat1
+    dlon = lon2 - lon1
+
+    # Haversine formula
+    dang = 2 * np.arcsin(
+        np.sqrt(
+            np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
+        )
+    )
+
+    # Radius of the Earth in meters
+    r_earth = 6371315.0
+
+    # Distance in meters
+    dis = r_earth * dang
+
+    return dis
+
+
+def convert_to_roms_time(ds, model_reference_date, climatology, time_name="time"):
+
+    if climatology:
+        ds.attrs["climatology"] = str(True)
+        month = xr.DataArray(range(1, 13), dims=time_name)
+        month.attrs["long_name"] = "Month index (1-12)"
+        ds = ds.assign_coords({"month": month})
+        # Preserve absolute time coordinate for readability
+        abs_time = np.datetime64(model_reference_date) + ds[time_name]
+        # Convert to pandas TimedeltaIndex
+        timedelta_index = pd.to_timedelta(ds[time_name].values)
+
+        # Determine the start of the year for the base_datetime
+        start_of_year = datetime(model_reference_date.year, 1, 1)
+
+        # Calculate the offset from midnight of the new year
+        offset = model_reference_date - start_of_year
+
+        # Convert the timedelta to nanoseconds first, then to days
+        time = xr.DataArray(
+            (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
+            dims=time_name,
+        )
+        time.attrs["cycle_length"] = 365.25
+
+    else:
+        # Preserve absolute time coordinate for readability
+        abs_time = ds[time_name]
+
+        time = (
+            (ds[time_name] - np.datetime64(model_reference_date)).astype("float64")
+            / 3600
+            / 24
+            * 1e-9
+        )
+
+    attrs = [key for key in abs_time.attrs]
+    for attr in attrs:
+        del abs_time.attrs[attr]
+    abs_time.attrs["long_name"] = "absolute time"
+    ds = ds.assign_coords({"abs_time": abs_time})
+
+    time.attrs["long_name"] = f"relative time: days since {str(model_reference_date)}"
+    time.encoding["units"] = "days"
+    time.attrs["units"] = "days"
+    ds.encoding["unlimited_dims"] = "time"
+
+    return ds, time
+
+
+def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
+    """Serialize a forcing object (including its grid) into a YAML file.
+
+    This function serializes a dataclass object (forcing_object) and its associated
+    `grid` attribute into a YAML file. It includes additional metadata, such as
+    the version of the `roms-tools` package, and omits fields like `grid` and `ds`
+    that are not serializable or meant to be excluded.
+
+    The function also converts datetime fields to ISO format strings for proper
+    serialization.
+
+    Parameters
+    ----------
+    forcing_object : object
+        The object that contains the forcing data, typically a dataclass with attributes
+        such as `grid`, `start_time`, `end_time`, etc.
+    filepath : Union[str, Path]
+        The path where the serialized YAML file will be saved.
+
+    Returns
+    -------
+    None
+        The function writes the serialized data directly to a YAML file at the specified path.
+    """
+
+    # Convert the filepath to a Path object
+    filepath = Path(filepath)
+
+    # Step 1: Serialize Grid data
+    # Convert the grid attribute to a dictionary and remove non-serializable fields
+    grid_data = asdict(forcing_object.grid)
+    grid_data.pop("ds", None)  # Remove 'ds' attribute (non-serializable)
+    grid_data.pop("straddle", None)
+    grid_data.pop("verbose", None)
+    grid_yaml_data = {"Grid": grid_data}
+
+    # Step 2: Get ROMS Tools version
+    # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
+    try:
+        roms_tools_version = importlib.metadata.version("roms-tools")
+    except importlib.metadata.PackageNotFoundError:
+        roms_tools_version = "unknown"
+
+    # Create YAML header with version information
+    header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+    # Step 3: Prepare Forcing Data
+    # Prepare the forcing object fields, excluding 'grid' and 'ds'
+    forcing_data = {}
+    field_names = [field.name for field in fields(forcing_object)]
+    filtered_field_names = [
+        param
+        for param in field_names
+        if param not in ("grid", "ds", "use_dask", "climatology")
+    ]
+
+    for field_name in filtered_field_names:
+        # Retrieve the value of each field using getattr
+        value = getattr(forcing_object, field_name)
+
+        # If the field is a datetime object, convert it to ISO format
+        if isinstance(value, datetime):
+            value = value.isoformat()
+
+        # Add the field and its value to the forcing_data dictionary
+        forcing_data[field_name] = value
+
+    # Step 4: Combine Grid and Forcing Data
+    # Combine grid and forcing data into a single dictionary for the final YAML content
+    yaml_data = {
+        **grid_yaml_data,  # Add the grid data to the final YAML structure
+        forcing_object.__class__.__name__: forcing_data,  # Include the serialized forcing object data
+    }
+
+    # Step 5: Write to YAML file
+    with filepath.open("w") as file:
+        # Write the header first
+        file.write(header)
+        # Write the serialized YAML data
+        yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+
+
+def _from_yaml(forcing_object: Type, filepath: Union[str, Path]) -> Dict[str, Any]:
+    """Extract the configuration data for a given forcing object from a YAML file.
+
+    This function reads a YAML file, searches for the configuration data associated
+    with the class name of the forcing object, and returns the configuration data
+    as a dictionary. The dictionary contains the forcing parameters extracted from
+    the YAML file, with any date fields converted from ISO format.
+
+    Parameters
+    ----------
+    filepath : Union[str, Path]
+        The path to the YAML file from which the parameters will be read.
+    forcing_object : Type
+        The class type (e.g., TidalForcing) whose configuration data is to be loaded
+        from the YAML file. The class name is used to locate the relevant data in
+        the YAML structure.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the forcing parameters extracted from the YAML file.
+        This dictionary contains key-value pairs where the keys are the parameter
+        names, and the values are the corresponding values from the YAML file.
+        Any date fields are converted from ISO format if necessary.
+
+    Raises
+    ------
+    ValueError
+        If no configuration for the specified class name is found in the YAML file.
+    """
+
+    # Read the entire file content
+    with filepath.open("r") as file:
+        file_content = file.read()
+
+    # Split the content into YAML documents
+    documents = list(yaml.safe_load_all(file_content))
+
+    forcing_data = None
+    forcing_object_name = forcing_object.__name__
+
+    # Process the YAML documents to find the forcing data for the given object
+    for doc in documents:
+        if doc is None:
+            continue
+        if forcing_object_name in doc:
+            forcing_data = doc[forcing_object_name]
+            break
+
+    if forcing_data is None:
+        raise ValueError(
+            f"No {forcing_object_name} configuration found in the YAML file."
+        )
+
+    # Convert any date fields from ISO format if necessary
+    for key, value in forcing_data.items():
+        forcing_data[key] = _convert_from_iso_format(value)
+
+    # Return the forcing data as a dictionary
+    return forcing_data
+
+
+def _convert_from_iso_format(value):
+    try:
+        # Return the parsed datetime object if successful
+        return datetime.fromisoformat(str(value))
+    except ValueError:
+        # Return None or raise an exception if parsing fails
+        return value
+
+
+def handle_boundaries(field):
+    """Adjust the boundaries of a 2D field by copying values from adjacent cells.
+
+    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.
+
+    Returns
+    -------
+    field : numpy.ndarray or xarray.DataArray
+        The input field with adjusted boundary values.
+    """
+
+    field[0, :] = field[1, :]
+    field[-1, :] = field[-2, :]
+    field[:, 0] = field[:, 1]
+    field[:, -1] = field[:, -2]
+
+    return field