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

roms_tools/plot.py

@@ -96,12 +96,6 @@ def plot_2d_horizontal_field(
     lon_deg = field.lon
     lat_deg = field.lat
 
-    # check if North or South pole are in domain
-    if lat_deg.max().values > 89 or lat_deg.min().values < -89:
-        raise NotImplementedError(
-            "Plotting is not implemented for the case that the domain contains the North or South pole."
-        )
-
     trans = get_projection(lon_deg, lat_deg)
 
     if ax is None:
@@ -186,25 +180,8 @@ def plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=
         with_dim_names=with_dim_names,
     )
 
-    vmax = 3
-    vmin = 0
-    cmap = plt.colormaps.get_cmap("Blues")
-    cmap.set_bad(color="gray")
-    kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-    field = parent_grid_ds.mask_rho
-    field = field.where(field)
-
-    _add_field_to_ax(
-        ax,
-        parent_lon_deg,
-        parent_lat_deg,
-        field,
-        add_colorbar=False,
-        kwargs=kwargs,
-    )
-
     _add_gridlines(ax)
+    ax.coastlines()
 
     ax.legend(loc="best")
 
@@ -659,7 +636,7 @@ def _add_field_to_ax(
     proj = ccrs.PlateCarree()
 
     p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-    if hasattr(field, "long_name"):
+    if hasattr(field, "long_name") and hasattr(field, "units"):
         label = f"{field.long_name} [{field.units}]"
     elif hasattr(field, "Long_name"):
         # this is the case for matlab generated grids
@@ -674,10 +651,36 @@ def _add_field_to_ax(
         ax.clabel(cs, inline=True, fontsize=FONT_SZ)
 
 
-def get_projection(lon, lat):
-    return ccrs.NearsidePerspective(
-        central_longitude=lon.mean().values, central_latitude=lat.mean().values
-    )
+def get_projection(lon: xr.DataArray, lat: xr.DataArray):
+    """
+    Return a Cartopy projection appropriate for the given lon/lat domain.
+
+    - Raises an error if the domain includes either pole.
+    - Uses PlateCarree if the longitudinal span > 90°.
+    - Otherwise uses NearsidePerspective centered on the domain.
+
+    Parameters
+    ----------
+    lon, lat : xr.DataArray
+        Longitude and latitude values in degrees.
+
+    Returns
+    -------
+    cartopy.crs.Projection
+        The chosen Cartopy projection.
+    """
+    # check if North or South pole are in domain
+    if lat.max() > 89 or lat.min() < -89:
+        raise NotImplementedError(
+            "Plotting is not implemented for the case that the domain contains the North or South pole."
+        )
+
+    if lon.max() - lon.min() > 90:
+        return ccrs.PlateCarree(central_longitude=lon.mean().values)
+    else:
+        return ccrs.NearsidePerspective(
+            central_longitude=lon.mean().values, central_latitude=lat.mean().values
+        )
 
 
 def _validate_plot_inputs(
@@ -816,6 +819,7 @@ def plot(
     yincrease: bool | None = None,
     use_coarse_grid: bool = False,
     with_dim_names: bool = False,
+    apply_mask: bool = True,
     ax: Axes | None = None,
     save_path: str | None = None,
     cmap_name: str = "YlOrRd",
@@ -889,6 +893,9 @@ def plot(
     with_dim_names : bool, optional
         Add grid dimension names (`xi`, `eta`) to the outer plot edges. Only for 2D plots. Default is False.
 
+    apply_mask: bool, optional
+        Whether to apply the land mask to the field. Default is True.
+
     ax : Axes, optional
         Matplotlib axes object. If None, a new figure is created. Default is None.
 
@@ -971,7 +978,8 @@ def plot(
     field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
     # Mask the field
-    field = field.where(mask)
+    if apply_mask:
+        field = field.where(mask)
 
     # Assign eta and xi as coordinates
     coords_to_assign = {dim: field[dim] for dim in horizontal_dims.values()}
@@ -1268,7 +1276,7 @@ def plot_uptake_efficiency(ds: xr.Dataset) -> None:
     ----------
     ds : xarray.Dataset
         Dataset containing the following variables:
-        - "abs_time": array of timestamps (datetime-like)
+        - "time": array of timestamps (datetime-like)
         - "cdr_efficiency": uptake efficiency from flux differences
         - "cdr_efficiency_from_delta_diff": uptake efficiency from DIC differences
 
@@ -1281,16 +1289,16 @@ def plot_uptake_efficiency(ds: xr.Dataset) -> None:
     -------
     None
     """
-    required_vars = ["abs_time", "cdr_efficiency", "cdr_efficiency_from_delta_diff"]
+    required_vars = ["time", "cdr_efficiency", "cdr_efficiency_from_delta_diff"]
     for var in required_vars:
         if var not in ds or ds[var].size == 0:
             raise ValueError(f"Dataset must contain non-empty variable '{var}'.")
 
-    times = ds["abs_time"]
+    times = ds["time"]
 
     # Check for monotonically increasing times
     if not np.all(times[1:] >= times[:-1]):
-        raise ValueError("abs_time must be strictly increasing.")
+        raise ValueError("time must be strictly increasing.")
 
     fig, ax = plt.subplots(figsize=(10, 4))
 

roms_tools/setup/boundary_forcing.py

@@ -11,17 +11,19 @@ import xarray as xr
 from scipy.ndimage import label
 
 from roms_tools import Grid
-from roms_tools.plot import line_plot, section_plot
-from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.setup.datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
-    RawDataSource,
     UnifiedBGCDataset,
 )
+from roms_tools.fill import one_dim_fill
+from roms_tools.plot import line_plot, section_plot
+from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.setup.utils import (
+    RawDataSource,
     add_time_info_to_ds,
+    check_and_set_boundaries,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
     from_yaml,
@@ -30,7 +32,6 @@ from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
     nan_check,
-    one_dim_fill,
     rotate_velocities,
     substitute_nans_by_fillvalue,
     to_dict,
@@ -61,8 +62,9 @@ class BoundaryForcing:
         The end time of the desired surface forcing data. This time is used to filter the dataset
         to include only records on or before this time, with a single record at or after this time.
         If no time filtering is desired, set it to None. Default is None.
-    boundaries : Dict[str, bool], optional
-        Dictionary specifying which boundaries are forced (south, east, north, west). Default is all True.
+    boundaries : dict[str, bool], optional
+        Specifies which grid boundaries ('south', 'east', 'north', 'west') are active and to be processed.
+        if not provided, valid (non-land) boundaries are enabled automatically.
     source : RawDataSource
         Dictionary specifying the source of the boundary forcing data. Keys include:
 
@@ -86,10 +88,6 @@ class BoundaryForcing:
         Indicates whether to perform a two-dimensional horizontal fill on the source data prior to regridding to boundaries.
         If `False`, a one-dimensional horizontal fill is performed separately on each of the four regridded boundaries.
         Defaults to `False`.
-    adjust_depth_for_sea_surface_height : bool, optional
-        Whether to account for sea surface height (`zeta`) variations when computing depth coordinates.
-        This adjustment is only applicable for `type="physics"`, as for biogeochemical fields usually `zeta` is not available.
-        Defaults to `False`.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
     use_dask: bool, optional
@@ -117,14 +115,7 @@ class BoundaryForcing:
     """The start time of the desired surface forcing data."""
     end_time: datetime | None = None
     """The end time of the desired surface forcing data."""
-    boundaries: dict[str, bool] = field(
-        default_factory=lambda: {
-            "south": True,
-            "east": True,
-            "north": True,
-            "west": True,
-        }
-    )
+    boundaries: dict[str, bool] | None = None
     """Dictionary specifying which boundaries are forced (south, east, north, west)."""
     source: RawDataSource
     """Dictionary specifying the source of the boundary forcing data."""
@@ -133,9 +124,6 @@ class BoundaryForcing:
     apply_2d_horizontal_fill: bool = False
     """Whether to perform a two-dimensional horizontal fill on the source data prior to
     regridding to boundaries."""
-    adjust_depth_for_sea_surface_height: bool = False
-    """Whether to account for sea surface height (`zeta`) variations when computing
-    depth coordinates."""
     model_reference_date: datetime = datetime(2000, 1, 1)
     """Reference date for the model."""
     use_dask: bool = False
@@ -146,12 +134,18 @@ class BoundaryForcing:
     ds: xr.Dataset = field(init=False, repr=False)
     """An xarray Dataset containing post-processed variables ready for input into
     ROMS."""
+    adjust_depth_for_sea_surface_height: bool = field(init=False)
+    """Whether to account for sea surface height when computing depth coordinates."""
+    ds_depth_coords: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing the depth coordinates."""
 
     def __post_init__(self):
-        self._input_checks()
-        # Dataset for depth coordinates
+        # Initialize depth coordinates
+        self.adjust_depth_for_sea_surface_height = False
         self.ds_depth_coords = xr.Dataset()
 
+        self._input_checks()
+
         target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
@@ -388,51 +382,58 @@ class BoundaryForcing:
 
         self.ds = ds
 
-    def _input_checks(self):
-        # Check that start_time and end_time are both None or none of them is
+    def _input_checks(self) -> None:
+        """Validate and normalize user-provided input parameters."""
+        # -------------------------------------------------------
+        # Time range checks
+        # -------------------------------------------------------
         if (self.start_time is None) != (self.end_time is None):
             raise ValueError(
                 "Both `start_time` and `end_time` must be provided together as datetime objects or both should be None."
             )
 
-        # Trigger a warning if both are None
         if self.start_time is None and self.end_time is None:
             logging.warning(
                 "Both `start_time` and `end_time` are None. No time filtering will be applied to the source data."
             )
 
-        # Validate the 'type' parameter
-        if self.type not in ["physics", "bgc"]:
+        # -------------------------------------------------------
+        # Type check
+        # -------------------------------------------------------
+        if self.type not in {"physics", "bgc"}:
             raise ValueError("`type` must be either 'physics' or 'bgc'.")
 
-        # Ensure 'source' dictionary contains required keys
+        # -------------------------------------------------------
+        # Source configuration checks
+        # -------------------------------------------------------
         if "name" not in self.source:
             raise ValueError("`source` must include a 'name'.")
+
         if "path" not in self.source:
             if self.source["name"] != "GLORYS":
                 raise ValueError("`source` must include a 'path'.")
-
             self.source["path"] = GLORYSDefaultDataset.dataset_name
 
-        # Set 'climatology' to False if not provided in 'source'
-        self.source = {
-            **self.source,
-            "climatology": self.source.get("climatology", False),
-        }
+        # Assign default value
+        self.source["climatology"] = self.source.get("climatology", False)
+
+        # -------------------------------------------------------
+        # Boundary selection defaults and validation
+        # -------------------------------------------------------
+
+        self.boundaries = check_and_set_boundaries(
+            self.boundaries, self.grid.ds.mask_rho
+        )
 
-        # Ensure adjust_depth_for_sea_surface_height is only used with type="physics"
+        # -------------------------------------------------------
+        # Depth adjustment checks
+        # -------------------------------------------------------
         if self.type == "bgc" and self.adjust_depth_for_sea_surface_height:
             logging.warning(
                 "adjust_depth_for_sea_surface_height is not applicable for BGC fields. "
                 "Setting it to False."
             )
             self.adjust_depth_for_sea_surface_height = False
-        elif self.adjust_depth_for_sea_surface_height:
-            logging.info("Sea surface height will be used to adjust depth coordinates.")
-        else:
-            logging.info(
-                "Sea surface height will NOT be used to adjust depth coordinates."
-            )
 
     def _get_data(
         self,
@@ -1061,7 +1062,14 @@ class BoundaryForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        forcing_dict = to_dict(self, exclude=["use_dask"])
+        forcing_dict = to_dict(
+            self,
+            exclude=[
+                "ds_depth_coords",
+                "adjust_depth_for_sea_surface_height",
+                "use_dask",
+            ],
+        )
         write_to_yaml(forcing_dict, filepath)
 
     @classmethod

roms_tools/setup/cdr_release.py

@@ -573,7 +573,7 @@ class TracerPerturbation(Release):
 
     tracer_fluxes : dict, optional
 
-        Dictionary of tracer names and their non-negative flux values. The flux values can be either
+        Dictionary of tracer names and their flux values. The flux values can be either
         a float (constant in time) or a list of float (time-varying).
 
         - Constant: applies uniformly across the entire simulation period.
@@ -589,9 +589,7 @@ class TracerPerturbation(Release):
     """
 
     times: list[datetime] = Field([])
-    tracer_fluxes: dict[str, Flux | NonNegativeFloat | list[NonNegativeFloat]] = Field(
-        {}
-    )
+    tracer_fluxes: dict[str, Flux | float | list[float]] = Field({})
     """Dictionary of tracer names and their non-negative flux values."""
 
     release_type: Literal[ReleaseType.tracer_perturbation] = (

roms_tools/setup/grid.py

@@ -201,46 +201,71 @@ class Grid:
             self.mask_shapefile = mask_shapefile
 
     def update_topography(
-        self, topography_source=None, hmin=None, verbose=False
+        self,
+        topography_source: dict | None = None,
+        hmin: float | None = None,
+        verbose: bool = False,
     ) -> None:
-        """Update the grid dataset with processed topography.
+        """
+        Update the grid dataset with processed topography.
 
-        This method performs several key operations, including regridding the topography, smoothing the
-        topography over the entire domain and locally.
-        The processed topography is then added to the grid's dataset.
+        This method performs the following operations:
+
+        1. Regrids the topography from the specified source.
+        2. Applies domain-wide and local smoothing.
+        3. Enforces the minimum depth constraint ``hmin``.
+        4. Updates the internal dataset (``self.ds``) with the processed bathymetry.
 
         Parameters
         ----------
         topography_source : dict, optional
-            A dictionary specifying the source of the topography data. The dictionary should
-            contain the following keys:
-            - "name" (str): The name of the topography data source (e.g., "SRTM15").
-            - "path" (Union[str, Path): The path to the raw data file.
+            Dictionary specifying the topography data source with keys:
+
+            - ``"name"`` (str): Name of the topography dataset (e.g. ``"SRTM15"``)
+            - ``"path"`` (str or pathlib.Path): Path to the raw data file
 
             If not provided, `topography_source` will remain unchanged (i.e., the existing value will not be overwritten).
 
         hmin : float, optional
-            The minimum ocean depth (in meters).
+            Minimum ocean depth in meters.
             If not provided, `hmin` will remain unchanged (i.e., the existing value will not be overwritten).
 
         verbose : bool, optional
-            If True, the method will print detailed information about the grid generation process,
-            including the timing of each step. Defaults to False.
+            If True, print detailed information about the processing steps and timing.
+            Defaults to False.
 
         Returns
         -------
         None
-            This method updates the internal dataset (`self.ds`) in place by adding or overwriting the
-            topography variable. It does not return any value.
+            Updates ``self.ds`` in place.
         """
-        topography_source = topography_source or self.topography_source
-        hmin = hmin or self.hmin
+        if topography_source is None:
+            topography_source = self.topography_source
+
+        if hmin is None:
+            hmin = self.hmin
+
+        # This can only happen for externally generated grids read via Grid.from_file()
+        if topography_source is None:
+            raise ValueError(
+                "Topography source information is not available. "
+                "Please provide `topography_source` explicitly when calling "
+                "`update_topography()`."
+            )
+        # This can only happen for externally generated grids read via Grid.from_file()
+        if hmin is None:
+            raise ValueError(
+                "Minimal ocean depth is not available. "
+                "Please provide `hmin` explicitly when calling `update_topography()`."
+            )
+
+        name = topography_source["name"]  # type: ignore[index]
 
         # Extract target coordinates for processing
         target_coords = get_target_coords(self)
 
         with Timed(
-            f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ===",
+            f"=== Generating the topography using {name} data and hmin = {hmin} meters ===",
             verbose=verbose,
         ):
             ds = add_topography(

roms_tools/setup/initial_conditions.py

@@ -1,5 +1,4 @@
 import importlib.metadata
-import logging
 from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -11,17 +10,17 @@ import xarray as xr
 from matplotlib.axes import Axes
 
 from roms_tools import Grid
-from roms_tools.plot import plot
-from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.setup.datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
-    Dataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
-    RawDataSource,
+    LatLonDataset,
     UnifiedBGCDataset,
 )
+from roms_tools.plot import plot
+from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.setup.utils import (
+    RawDataSource,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
     from_yaml,
@@ -81,9 +80,6 @@ class InitialConditions:
             - A list of strings or Path objects containing multiple files.
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
-    adjust_depth_for_sea_surface_height : bool, optional
-        Whether to account for sea surface height variations when computing depth coordinates.
-        Defaults to `False`.
     model_reference_date : datetime, optional
         The reference date for the model. Defaults to January 1, 2000.
     use_dask: bool, optional
@@ -129,9 +125,6 @@ class InitialConditions:
     data."""
     model_reference_date: datetime = datetime(2000, 1, 1)
     """The reference date for the model."""
-    adjust_depth_for_sea_surface_height: bool = False
-    """Whether to account for sea surface height variations when computing depth
-    coordinates."""
     allow_flex_time: bool = False
     """Whether to handle ini_time flexibly."""
     use_dask: bool = False
@@ -145,12 +138,18 @@ class InitialConditions:
     ds: xr.Dataset = field(init=False, repr=False)
     """An xarray Dataset containing post-processed variables ready for input into
     ROMS."""
+    adjust_depth_for_sea_surface_height: bool = field(init=False)
+    """Whether to account for sea surface height when computing depth coordinates."""
+    ds_depth_coords: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing the depth coordinates."""
 
     def __post_init__(self):
-        self._input_checks()
-        # Dataset for depth coordinates
+        # Initialize depth coordinates
+        self.adjust_depth_for_sea_surface_height = False
         self.ds_depth_coords = xr.Dataset()
 
+        self._input_checks()
+
         processed_fields = {}
         processed_fields = self._process_data(processed_fields, type="physics")
 
@@ -319,14 +318,8 @@ class InitialConditions:
                 **self.bgc_source,
                 "climatology": self.bgc_source.get("climatology", False),
             }
-        if self.adjust_depth_for_sea_surface_height:
-            logging.info("Sea surface height will be used to adjust depth coordinates.")
-        else:
-            logging.info(
-                "Sea surface height will NOT be used to adjust depth coordinates."
-            )
 
-    def _get_data(self, forcing_type=Literal["physics", "bgc"]) -> Dataset:
+    def _get_data(self, forcing_type=Literal["physics", "bgc"]) -> LatLonDataset:
         """Determine the correct `Dataset` type and return an instance.
 
         forcing_type : str
@@ -339,7 +332,7 @@ class InitialConditions:
         Dataset
             The `Dataset` instance
         """
-        dataset_map: dict[str, dict[str, dict[str, type[Dataset]]]] = {
+        dataset_map: dict[str, dict[str, dict[str, type[LatLonDataset]]]] = {
             "physics": {
                 "GLORYS": {
                     "external": GLORYSDataset,
@@ -826,7 +819,14 @@ class InitialConditions:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        forcing_dict = to_dict(self, exclude=["use_dask"])
+        forcing_dict = to_dict(
+            self,
+            exclude=[
+                "ds_depth_coords",
+                "adjust_depth_for_sea_surface_height",
+                "use_dask",
+            ],
+        )
         write_to_yaml(forcing_dict, filepath)
 
     @classmethod