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

roms_tools/setup/tides.py

@@ -51,6 +51,10 @@ class TidalForcing:
         The reference date for the ROMS simulation. Default is datetime(2000, 1, 1).
     use_dask: bool, optional
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
+    bypass_validation: bool, optional
+        Indicates whether to skip validation checks in the processed data. When set to True,
+        the validation process that ensures no NaN values exist at wet points
+        in the processed dataset is bypassed. Defaults to False.
 
     Examples
     --------
@@ -65,6 +69,7 @@ class TidalForcing:
     allan_factor: float = 2.0
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
+    bypass_validation: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -129,7 +134,8 @@ class TidalForcing:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        if not self.bypass_validation:
+            self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -420,7 +426,10 @@ class TidalForcing:
 
     @classmethod
     def from_yaml(
-        cls, filepath: Union[str, Path], use_dask: bool = False
+        cls,
+        filepath: Union[str, Path],
+        use_dask: bool = False,
+        bypass_validation: bool = False,
     ) -> "TidalForcing":
         """Create an instance of the TidalForcing class from a YAML file.
 
@@ -430,6 +439,10 @@ class TidalForcing:
             The path to the YAML file from which the parameters will be read.
         use_dask: bool, optional
             Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
+        bypass_validation: bool, optional
+            Indicates whether to skip validation checks in the processed data. When set to True,
+            the validation process that ensures no NaN values exist at wet points
+            in the processed dataset is bypassed. Defaults to False.
 
         Returns
         -------
@@ -440,7 +453,12 @@ class TidalForcing:
 
         grid = Grid.from_yaml(filepath)
         tidal_forcing_params = _from_yaml(cls, filepath)
-        return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
+        return cls(
+            grid=grid,
+            **tidal_forcing_params,
+            use_dask=use_dask,
+            bypass_validation=bypass_validation
+        )
 
     def _correct_tides(self, data):
         """Apply tidal corrections to the dataset. This method corrects the dataset for

roms_tools/setup/topography.py

@@ -94,7 +94,18 @@ def _add_topography(
 
 
 def _get_topography_data(source):
+    """Load topography data based on the specified source.
 
+    Parameters
+    ----------
+    source : dict
+        A dictionary containing the source details (e.g., "name" and "path").
+
+    Returns
+    -------
+    data : object
+        The loaded topography dataset (ETOPO5 or SRTM15).
+    """
     kwargs = {"use_dask": False}
 
     if source["name"] == "ETOPO5":
@@ -115,7 +126,24 @@ def _get_topography_data(source):
 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:
@@ -134,9 +162,22 @@ def _make_raw_topography(
 
 
 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(
@@ -164,7 +205,27 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
 
 
 def _smooth_topography_locally(h, hmin=5, rmax=0.2):
-    """Smoothes topography locally to satisfy r < rmax."""
+    """Smooths topography locally to ensure the slope (r-factor) is below the specified
+    threshold.
+
+    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.
+
+    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))
@@ -174,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)
 
-        # 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:
@@ -242,8 +328,23 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
 
 
 def _compute_rfactor(h):
-    """Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid
-    directions."""
+    """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
+    ----------
+    h : xarray.DataArray
+        The topography data.
+
+    Returns
+    -------
+    tuple of xarray.DataArray
+        r_eta : r-factor in the eta direction.
+        r_xi : r-factor in the xi direction.
+    """
     # 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)
@@ -256,6 +357,26 @@ def _compute_rfactor(h):
 
 
 def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
+    """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
 
@@ -263,6 +384,18 @@ def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
 
 
 def nan_check(hraw):
+    """Checks for NaN values in the topography data.
+
+    Parameters
+    ----------
+    hraw : xarray.DataArray
+        Input topography data to check for NaN values.
+
+    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 "

roms_tools/setup/utils.py

@@ -927,37 +927,40 @@ def get_vector_pairs(variable_info):
     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).
+def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
+    """Calculate the great circle distance between two points on the Earth's surface
+    using the Haversine formula.
 
-    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.
+    Latitude and longitude are assumed to be in degrees by default. If `input_in_degrees` is set to `False`,
+    the input is assumed to already be in radians.
 
     Parameters
     ----------
     lon1, lat1 : float
-        Longitude and latitude of the first point in degrees.
+        Longitude and latitude of the first point.
     lon2, lat2 : float
-        Longitude and latitude of the second point in degrees.
+        Longitude and latitude of the second point.
+    input_in_degrees : bool, optional
+        If True (default), the input coordinates are assumed to be in degrees and will be converted to radians.
+        If False, the input is assumed to be in radians and no conversion is applied.
 
     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
+    if input_in_degrees:
+        d2r = np.pi / 180
+        lon1 = lon1 * d2r
+        lat1 = lat1 * d2r
+        lon2 = lon2 * d2r
+        lat2 = lat2 * d2r
 
     # Difference in latitudes and longitudes
     dlat = lat2 - lat1
@@ -1058,12 +1061,17 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     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}
+    # Check if the forcing_object has a grid attribute
+    if hasattr(forcing_object, "grid") and forcing_object.grid is not None:
+        grid_data = asdict(forcing_object.grid)
+        grid_yaml_data = {"Grid": _pop_grid_data(grid_data)}
+    else:
+        parent_grid_data = asdict(forcing_object.parent_grid)
+        parent_grid_yaml_data = {"ParentGrid": _pop_grid_data(parent_grid_data)}
+        child_grid_data = asdict(forcing_object.child_grid)
+        child_grid_yaml_data = {"ChildGrid": _pop_grid_data(child_grid_data)}
+
+        grid_yaml_data = {**parent_grid_yaml_data, **child_grid_yaml_data}
 
     # Step 2: Get ROMS Tools version
     # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
@@ -1082,7 +1090,16 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     filtered_field_names = [
         param
         for param in field_names
-        if param not in ("grid", "ds", "use_dask", "climatology")
+        if param
+        not in (
+            "grid",
+            "parent_grid",
+            "child_grid",
+            "ds",
+            "use_dask",
+            "bypass_validation",
+            "climatology",
+        )
     ]
 
     for field_name in filtered_field_names:
@@ -1111,6 +1128,14 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
 
+def _pop_grid_data(grid_data):
+    grid_data.pop("ds", None)  # Remove 'ds' attribute (non-serializable)
+    grid_data.pop("straddle", None)
+    grid_data.pop("verbose", None)
+
+    return grid_data
+
+
 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.
 
@@ -1203,3 +1228,84 @@ def handle_boundaries(field):
     field[:, -1] = field[:, -2]
 
     return field
+
+
+def get_boundary_coords():
+    """This function determines the boundary points for the grid variables by specifying
+    the indices for the south, east, north, and west boundaries.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the boundary coordinates for different variable types.
+        The dictionary has the following structure:
+        - Keys: Variable types ("rho", "u", "v", "vector").
+        - Values: Nested dictionaries that map each direction ("south", "east", "north", "west")
+          to another dictionary specifying the boundary coordinates, represented by grid indices
+          for the respective variable types. For example:
+          - "rho" variables (e.g., `eta_rho`, `xi_rho`)
+          - "u" variables (e.g., `xi_u`)
+          - "v" variables (e.g., `eta_v`)
+          - "vector" variables with lists of indices for multiple grid points (e.g., `eta_rho`, `xi_rho`).
+    """
+
+    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},
+        },
+        "vector": {
+            "south": {"eta_rho": [0, 1]},
+            "east": {"xi_rho": [-2, -1]},
+            "north": {"eta_rho": [-2, -1]},
+            "west": {"xi_rho": [0, 1]},
+        },
+    }
+
+    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

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata

@@ -120,8 +120,7 @@
                 "nriver"
             ],
             "coordinates": "abs_time river_name tracer_name",
-            "long_name": "River tracer data",
-            "units": "degrees C [temperature]; psu [salinity]"
+            "long_name": "River tracer data"
         },
         "river_volume/.zarray": {
             "chunks": [
@@ -165,7 +164,7 @@
                 "id": "blosc",
                 "shuffle": 1
             },
-            "dtype": "<U11",
+            "dtype": "<U4",
             "fill_value": null,
             "filters": null,
             "order": "C",

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zattrs

@@ -5,6 +5,5 @@
         "nriver"
     ],
     "coordinates": "abs_time river_name tracer_name",
-    "long_name": "River tracer data",
-    "units": "degrees C [temperature]; psu [salinity]"
+    "long_name": "River tracer data"
 }

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zarray

@@ -9,7 +9,7 @@
         "id": "blosc",
         "shuffle": 1
     },
-    "dtype": "<U11",
+    "dtype": "<U4",
     "fill_value": null,
     "filters": null,
     "order": "C",

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zmetadata

@@ -121,7 +121,7 @@
         "river_tracer/.zarray": {
             "chunks": [
                 12,
-                2,
+                34,
                 6
             ],
             "compressor": {
@@ -137,7 +137,7 @@
             "order": "C",
             "shape": [
                 12,
-                2,
+                34,
                 6
             ],
             "zarr_format": 2
@@ -149,8 +149,7 @@
                 "nriver"
             ],
             "coordinates": "abs_time month river_name tracer_name",
-            "long_name": "River tracer data",
-            "units": "degrees C [temperature]; psu [salinity]"
+            "long_name": "River tracer data"
         },
         "river_volume/.zarray": {
             "chunks": [
@@ -185,7 +184,7 @@
         },
         "tracer_name/.zarray": {
             "chunks": [
-                2
+                34
             ],
             "compressor": {
                 "blocksize": 0,
@@ -199,7 +198,7 @@
             "filters": null,
             "order": "C",
             "shape": [
-                2
+                34
             ],
             "zarr_format": 2
         },

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zarray

@@ -1,7 +1,7 @@
 {
     "chunks": [
         12,
-        2,
+        34,
         6
     ],
     "compressor": {
@@ -17,7 +17,7 @@
     "order": "C",
     "shape": [
         12,
-        2,
+        34,
         6
     ],
     "zarr_format": 2

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zattrs

@@ -5,6 +5,5 @@
         "nriver"
     ],
     "coordinates": "abs_time month river_name tracer_name",
-    "long_name": "River tracer data",
-    "units": "degrees C [temperature]; psu [salinity]"
+    "long_name": "River tracer data"
 }