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

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.lat_lon_datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
     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

@@ -206,7 +206,8 @@ class Grid:
         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 the following operations:
 
@@ -218,29 +219,45 @@ class Grid:
         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]
 

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,15 +10,15 @@ 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.lat_lon_datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
     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,
@@ -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,12 +318,6 @@ 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"]) -> LatLonDataset:
         """Determine the correct `Dataset` type and return an instance.
@@ -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

roms_tools/setup/nesting.py

@@ -13,16 +13,16 @@ from roms_tools.plot import plot_nesting
 from roms_tools.setup.topography import clip_depth
 from roms_tools.setup.utils import (
     Timed,
+    check_and_set_boundaries,
     from_yaml,
     get_boundary_coords,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     pop_grid_data,
     to_dict,
-    wrap_longitudes,
     write_to_yaml,
 )
-from roms_tools.utils import save_datasets
+from roms_tools.utils import save_datasets, wrap_longitudes
 
 
 @dataclass(kw_only=True)
@@ -43,8 +43,8 @@ class ChildGrid(Grid):
     parent_grid : Grid
         The parent grid object, providing the reference for the topography and mask of the child grid.
     boundaries : Dict[str, bool]
-        Specifies which child grid boundaries (south, east, north, west) should be adjusted for topography/mask
-        and included in `ds_nesting`. Defaults to all `True`.
+        Specifies which child grid boundaries ('south', 'east', 'north', 'west') should be adjusted for topography/mask
+        and included in `ds_nesting`. If not provided, valid (non-land) boundaries are enabled automatically.
     metadata : Dict[str, Any]
         Dictionary configuring the boundary nesting process, including:
 
@@ -87,6 +87,10 @@ class ChildGrid(Grid):
             "=== Mapping the child grid boundary points onto the indices of the parent grid ===",
             verbose=verbose,
         ):
+            self.boundaries = check_and_set_boundaries(
+                self.boundaries, self.ds.mask_rho
+            )
+
             # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
             parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
 

roms_tools/setup/river_forcing.py

@@ -11,16 +11,16 @@ import xarray as xr
 
 from roms_tools import Grid
 from roms_tools.constants import MAX_DISTINCT_COLORS
+from roms_tools.datasets.river_datasets import (
+    DaiRiverDataset,
+    get_indices_of_nearest_grid_cell_for_rivers,
+)
 from roms_tools.plot import (
     assign_category_colors,
     get_projection,
     plot_2d_horizontal_field,
     plot_location,
 )
-from roms_tools.setup.river_datasets import (
-    DaiRiverDataset,
-    get_indices_of_nearest_grid_cell_for_rivers,
-)
 from roms_tools.setup.utils import (
     RawDataSource,
     add_time_info_to_ds,

roms_tools/setup/surface_forcing.py

@@ -8,9 +8,7 @@ import numpy as np
 import xarray as xr
 
 from roms_tools import Grid
-from roms_tools.plot import plot
-from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.lat_lon_datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCSurfaceForcingDataset,
     ERA5ARCODataset,
     ERA5Correction,
@@ -18,6 +16,8 @@ from roms_tools.setup.lat_lon_datasets import (
     LatLonDataset,
     UnifiedBGCSurfaceDataset,
 )
+from roms_tools.plot import plot
+from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.setup.utils import (
     RawDataSource,
     add_time_info_to_ds,
@@ -26,7 +26,6 @@ from roms_tools.setup.utils import (
     get_target_coords,
     get_variable_metadata,
     group_dataset,
-    interpolate_from_climatology,
     min_dist_to_land,
     nan_check,
     rotate_velocities,
@@ -34,7 +33,11 @@ from roms_tools.setup.utils import (
     to_dict,
     write_to_yaml,
 )
-from roms_tools.utils import save_datasets, transpose_dimensions
+from roms_tools.utils import (
+    interpolate_from_climatology,
+    save_datasets,
+    transpose_dimensions,
+)
 
 DEFAULT_ERA5_ARCO_PATH = (
     "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
@@ -461,8 +464,9 @@ class SurfaceForcing:
             corr_factor = xr.concat(
                 [
                     interpolate_from_climatology(
-                        correction_data.ds[correction_data.var_names["swr_corr"]],
-                        correction_data.dim_names["time"],
+                        field=correction_data.ds[correction_data.var_names["swr_corr"]],
+                        time_dim=correction_data.dim_names["time"],
+                        time_coord=correction_data.dim_names["time"],
                         time=time,
                     )
                     for time in radiation.time
@@ -472,8 +476,9 @@ class SurfaceForcing:
         else:
             # Interpolate across all time steps at once
             corr_factor = interpolate_from_climatology(
-                correction_data.ds[correction_data.var_names["swr_corr"]],
-                correction_data.dim_names["time"],
+                field=correction_data.ds[correction_data.var_names["swr_corr"]],
+                time_dim=correction_data.dim_names["time"],
+                time_coord=correction_data.dim_names["time"],
                 time=radiation.time,
             )
 

roms_tools/setup/tides.py

@@ -7,9 +7,9 @@ import numpy as np
 import xarray as xr
 
 from roms_tools import Grid
+from roms_tools.datasets.lat_lon_datasets import TPXOManager
 from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.lat_lon_datasets import TPXOManager
 from roms_tools.setup.utils import (
     from_yaml,
     get_target_coords,

roms_tools/setup/topography.py

@@ -7,8 +7,12 @@ import gcm_filters
 import numpy as np
 import xarray as xr
 
+from roms_tools.datasets.lat_lon_datasets import (
+    EMODDataset,
+    ETOPO5Dataset,
+    SRTM15Dataset,
+)
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.lat_lon_datasets import ETOPO5Dataset, SRTM15Dataset
 from roms_tools.setup.utils import handle_boundaries
 
 
@@ -115,9 +119,12 @@ def _get_topography_data(source):
     elif source["name"] == "SRTM15":
         kwargs["filename"] = source["path"]
         data = SRTM15Dataset(**kwargs)
+    elif source["name"] == "EMOD":
+        kwargs["filename"] = source["path"]
+        data = EMODDataset(**kwargs)
     else:
         raise ValueError(
-            'Only "ETOPO5" and "SRTM15" are valid options for topography_source["name"].'
+            'Only "ETOPO5", "SRTM15" and "EMOD" are valid options for topography_source["name"].'
         )
 
     return data
@@ -147,6 +154,7 @@ def _make_raw_topography(
     data.choose_subdomain(target_coords, buffer_points=3, verbose=verbose)
     # Enforce double precision to ensure reproducibility
     data.convert_to_float64()
+    data.apply_lateral_fill()
 
     if verbose:
         start_time = time.time()