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

roms_tools/setup/initial_conditions.py

@@ -1,8 +1,10 @@
 import importlib.metadata
 import logging
+from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
+from typing import Literal
 
 import numpy as np
 import xarray as xr
@@ -11,7 +13,14 @@ 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 CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
+from roms_tools.setup.datasets import (
+    CESMBGCDataset,
+    Dataset,
+    GLORYSDataset,
+    GLORYSDefaultDataset,
+    RawDataSource,
+    UnifiedBGCDataset,
+)
 from roms_tools.setup.utils import (
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
@@ -25,7 +34,6 @@ from roms_tools.setup.utils import (
     write_to_yaml,
 )
 from roms_tools.utils import (
-    get_dask_chunks,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     save_datasets,
@@ -48,7 +56,7 @@ class InitialConditions:
     ini_time : datetime
         The date and time at which the initial conditions are set.
         If no exact match is found, the closest time entry to `ini_time` within the time range [ini_time, ini_time + 24 hours] is selected.
-    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+    source : RawDataSource
 
         Dictionary specifying the source of the physical initial condition data. Keys include:
 
@@ -57,10 +65,12 @@ class InitialConditions:
 
             - A single string (with or without wildcards).
             - A single Path object.
-            - A list of strings or Path objects containing multiple files.
+            - A list of strings or Path objects.
+            If omitted, the data will be streamed via the Copernicus Marine Toolkit.
+            Note: streaming is currently not recommended due to performance limitations.
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
-    bgc_source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+    bgc_source : RawDataSource, optional
         Dictionary specifying the source of the biogeochemical (BGC) initial condition data. Keys include:
 
           - "name" (str): Name of the data source (e.g., "CESM_REGRIDDED").
@@ -78,6 +88,13 @@ class InitialConditions:
         The reference date for the model. Defaults to January 1, 2000.
     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.
+    allow_flex_time: bool, optional
+        Controls how strictly `ini_time` is handled:
+
+        - If False (default): requires an exact match to `ini_time`. Raises a ValueError if no match exists.
+        - If True: allows a +24h search window after `ini_time` and selects the closest available
+          time entry within that window. Raises a ValueError if none are found.
+
     horizontal_chunk_size : int, optional
         The chunk size used for horizontal partitioning for the vertical regridding when `use_dask = True`. Defaults to 50.
         A larger number results in a bigger memory footprint but faster computations.
@@ -105,9 +122,9 @@ class InitialConditions:
     """Object representing the grid information."""
     ini_time: datetime
     """The date and time at which the initial conditions are set."""
-    source: dict[str, str | Path | list[str | Path]]
+    source: RawDataSource
     """Dictionary specifying the source of the physical initial condition data."""
-    bgc_source: dict[str, str | Path | list[str | Path]] | None = None
+    bgc_source: RawDataSource | None = None
     """Dictionary specifying the source of the biogeochemical (BGC) initial condition
     data."""
     model_reference_date: datetime = datetime(2000, 1, 1)
@@ -115,6 +132,8 @@ class InitialConditions:
     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
     """Whether to use dask for processing."""
     horizontal_chunk_size: int = 50
@@ -161,14 +180,10 @@ class InitialConditions:
     def _process_data(self, processed_fields, type="physics"):
         target_coords = get_target_coords(self.grid)
 
-        if type == "physics":
-            data = self._get_data()
-        else:
-            data = self._get_bgc_data()
+        data = self._get_data(forcing_type=type)
 
         data.choose_subdomain(
             target_coords,
-            buffer_points=20,  # lateral fill needs good buffer from data margin
         )
         # Enforce double precision to ensure reproducibility
         data.convert_to_float64()
@@ -255,7 +270,7 @@ class InitialConditions:
                         field = processed_fields[var_name]
                         if self.use_dask:
                             field = field.chunk(
-                                get_dask_chunks(location, self.horizontal_chunk_size)
+                                _set_dask_chunks(location, self.horizontal_chunk_size)
                             )
                         processed_fields[var_name] = vertical_regrid.apply(field)
 
@@ -280,7 +295,11 @@ class InitialConditions:
         if "name" not in self.source.keys():
             raise ValueError("`source` must include a 'name'.")
         if "path" not in self.source.keys():
-            raise ValueError("`source` must include a 'path'.")
+            if self.source["name"] != "GLORYS":
+                raise ValueError("`source` must include a 'path'.")
+
+            self.source["path"] = GLORYSDefaultDataset.dataset_name
+
         # set self.source["climatology"] to False if not provided
         self.source = {
             **self.source,
@@ -307,40 +326,63 @@ class InitialConditions:
                 "Sea surface height will NOT be used to adjust depth coordinates."
             )
 
-    def _get_data(self):
-        if self.source["name"] == "GLORYS":
-            data = GLORYSDataset(
-                filename=self.source["path"],
-                start_time=self.ini_time,
-                climatology=self.source["climatology"],
-                use_dask=self.use_dask,
-            )
-        else:
-            raise ValueError('Only "GLORYS" is a valid option for source["name"].')
-        return data
-
-    def _get_bgc_data(self):
-        if self.bgc_source["name"] == "CESM_REGRIDDED":
-            data = CESMBGCDataset(
-                filename=self.bgc_source["path"],
-                start_time=self.ini_time,
-                climatology=self.bgc_source["climatology"],
-                use_dask=self.use_dask,
-            )
-        elif self.bgc_source["name"] == "UNIFIED":
-            data = UnifiedBGCDataset(
-                filename=self.bgc_source["path"],
-                start_time=self.ini_time,
-                climatology=self.bgc_source["climatology"],
-                use_dask=self.use_dask,
-            )
+    def _get_data(self, forcing_type=Literal["physics", "bgc"]) -> Dataset:
+        """Determine the correct `Dataset` type and return an instance.
 
-        else:
-            raise ValueError(
-                'Only "CESM_REGRIDDED" and "UNIFIED" are valid options for bgc_source["name"].'
+        forcing_type : str
+            Specifies the type of forcing data. Options are:
+
+            - "physics": for physical atmospheric forcing.
+            - "bgc": for biogeochemical forcing.
+        Returns
+        -------
+        Dataset
+            The `Dataset` instance
+        """
+        dataset_map: dict[str, dict[str, dict[str, type[Dataset]]]] = {
+            "physics": {
+                "GLORYS": {
+                    "external": GLORYSDataset,
+                    "default": GLORYSDefaultDataset,
+                },
+            },
+            "bgc": {
+                "CESM_REGRIDDED": defaultdict(lambda: CESMBGCDataset),
+                "UNIFIED": defaultdict(lambda: UnifiedBGCDataset),
+            },
+        }
+
+        source_dict = self.source if forcing_type == "physics" else self.bgc_source
+
+        if source_dict is None:
+            raise ValueError(f"{forcing_type} source is not set")
+
+        source_name = str(source_dict["name"])
+        if source_name not in dataset_map[forcing_type]:
+            tpl = 'Valid options for source["name"] for type {} include: {}'
+            msg = tpl.format(
+                forcing_type, " and ".join(dataset_map[forcing_type].keys())
             )
+            raise ValueError(msg)
 
-        return data
+        has_no_path = "path" not in source_dict
+        has_default_path = source_dict.get("path") == GLORYSDefaultDataset.dataset_name
+        use_default = has_no_path or has_default_path
+
+        variant = "default" if use_default else "external"
+
+        data_type = dataset_map[forcing_type][source_name][variant]
+
+        if isinstance(source_dict["path"], bool):
+            raise ValueError('source["path"] cannot be a boolean here')
+
+        return data_type(
+            filename=source_dict["path"],
+            start_time=self.ini_time,
+            climatology=source_dict["climatology"],  # type: ignore
+            allow_flex_time=self.allow_flex_time,
+            use_dask=self.use_dask,
+        )
 
     def _set_variable_info(self, data, type="physics"):
         """Sets up a dictionary with metadata for variables based on the type.
@@ -483,10 +525,10 @@ class InitialConditions:
                     zeta = interpolate_from_rho_to_v(zeta)
 
             if self.use_dask:
-                h = h.chunk(get_dask_chunks(location, self.horizontal_chunk_size))
-                if self.adjust_depth_for_sea_surface_height:
+                h = h.chunk(_set_dask_chunks(location, self.horizontal_chunk_size))
+                if isinstance(zeta, xr.DataArray):
                     zeta = zeta.chunk(
-                        get_dask_chunks(location, self.horizontal_chunk_size)
+                        _set_dask_chunks(location, self.horizontal_chunk_size)
                     )
             depth = compute_depth(zeta, h, self.grid.ds.attrs["hc"], Cs, sigma)
             self.ds_depth_coords[key] = depth
@@ -816,3 +858,26 @@ class InitialConditions:
             **initial_conditions_params,
             use_dask=use_dask,
         )
+
+
+def _set_dask_chunks(location: str, chunk_size: int):
+    """Returns the appropriate Dask chunking dictionary based on grid location.
+
+    Parameters
+    ----------
+    location : str
+        The grid location, one of "rho", "u", or "v".
+    chunk_size : int
+        The chunk size to apply.
+
+    Returns
+    -------
+    dict
+        Dictionary specifying the chunking strategy.
+    """
+    chunk_mapping = {
+        "rho": {"eta_rho": chunk_size, "xi_rho": chunk_size},
+        "u": {"eta_rho": chunk_size, "xi_u": chunk_size},
+        "v": {"eta_v": chunk_size, "xi_rho": chunk_size},
+    }
+    return chunk_mapping.get(location, {})

roms_tools/setup/mask.py

@@ -1,5 +1,8 @@
+import logging
 import warnings
+from pathlib import Path
 
+import geopandas as gpd
 import numpy as np
 import regionmask
 import xarray as xr
@@ -12,7 +15,7 @@ from roms_tools.setup.utils import (
 )
 
 
-def _add_mask(ds):
+def add_mask(ds: xr.Dataset, shapefile: str | Path | None = None) -> xr.Dataset:
     """Adds a land/water mask to the dataset at rho-points.
 
     Parameters
@@ -20,20 +23,47 @@ def _add_mask(ds):
     ds : xarray.Dataset
         Input dataset containing latitude and longitude coordinates at rho-points.
 
+    shapefile: str or Path | None
+        Path to a coastal shapefile to determine the land mask. If None, NaturalEarth 10m is used.
+
     Returns
     -------
     xarray.Dataset
         The original dataset with an added 'mask_rho' variable, representing land/water mask.
     """
-    land = regionmask.defined_regions.natural_earth_v5_0_0.land_10
-
     # Suppress specific warning
     with warnings.catch_warnings():
         warnings.filterwarnings(
             "ignore", message="No gridpoint belongs to any region.*"
         )
-        land_mask = land.mask(ds["lon_rho"], ds["lat_rho"])
-    mask = land_mask.isnull()
+
+        if shapefile:
+            coast = gpd.read_file(shapefile)
+
+            try:
+                # 3D method: returns a boolean array for each region, then take max along the region dimension
+                # Pros: more memory-efficient for high-res grids if number of regions isn't extreme
+                mask = ~regionmask.mask_3D_geopandas(
+                    coast, ds["lon_rho"], ds["lat_rho"]
+                ).max(dim="region")
+
+            except MemoryError:
+                logging.info(
+                    "MemoryError encountered with 3D mask; falling back to 2D method."
+                )
+                # 2D method: returns a single array with integer codes for each region, using np.nan for points not in any region
+                # Pros: works well for small/medium grids
+                # Cons: can use a large float64 array internally for very high-resolution grids
+                mask_2d = regionmask.mask_geopandas(coast, ds["lon_rho"], ds["lat_rho"])
+                mask = mask_2d.isnull()
+
+        else:
+            # Use Natural Earth 10m land polygons if no shapefile is provided
+            land = regionmask.defined_regions.natural_earth_v5_0_0.land_10
+            land_mask = land.mask(ds["lon_rho"], ds["lat_rho"])
+            mask = land_mask.isnull()
+
+    ds = _add_coastlines_metadata(ds, shapefile)
 
     # fill enclosed basins with land
     mask = _fill_enclosed_basins(mask.values)
@@ -45,7 +75,33 @@ def _add_mask(ds):
         "long_name": "Mask at rho-points",
         "units": "land/water (0/1)",
     }
-    ds = _add_velocity_masks(ds)
+
+    ds = add_velocity_masks(ds)
+
+    return ds
+
+
+def _add_coastlines_metadata(
+    ds: xr.Dataset,
+    shapefile: str | Path | None = None,
+) -> xr.Dataset:
+    """
+    Add coastline metadata to a dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset to be updated.
+    shapefile : str or pathlib.Path or None, optional
+        Path to the shapefile used for land/ocean masking.
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset with updated coastline-related metadata.
+    """
+    if shapefile is not None:
+        ds.attrs["mask_shapefile"] = str(shapefile)
 
     return ds
 
@@ -85,7 +141,7 @@ def _fill_enclosed_basins(mask) -> np.ndarray:
     return mask
 
 
-def _add_velocity_masks(ds):
+def add_velocity_masks(ds):
     """Adds velocity masks for u- and v-points based on the rho-point mask.
 
     This function generates masks for u- and v-points by interpolating the rho-point land/water mask.

roms_tools/setup/nesting.py

@@ -9,8 +9,9 @@ from scipy.interpolate import griddata, interp1d
 
 from roms_tools import Grid
 from roms_tools.plot import plot_nesting
-from roms_tools.setup.topography import _clip_depth
+from roms_tools.setup.topography import clip_depth
 from roms_tools.setup.utils import (
+    Timed,
     from_yaml,
     get_boundary_coords,
     interpolate_from_rho_to_u,
@@ -48,6 +49,8 @@ class ChildGrid(Grid):
 
         - `"prefix"` (str): Prefix for variable names in `ds_nesting`. Defaults to `"child"`.
         - `"period"` (float): Temporal resolution for boundary outputs in seconds. Defaults to 3600 (hourly).
+    verbose: bool, optional
+        Indicates whether to print grid generation steps with timing. Defaults to False.
     """
 
     parent_grid: Grid
@@ -67,6 +70,8 @@ class ChildGrid(Grid):
         default_factory=lambda: {"prefix": "child", "period": 3600.0}
     )
     """Dictionary configuring the boundary nesting process."""
+    verbose: bool = False
+    """Whether to print grid generation steps with timing."""
 
     ds: xr.Dataset = field(init=False, repr=False)
     """An xarray Dataset containing child grid variables aligned with the
@@ -77,44 +82,48 @@ class ChildGrid(Grid):
 
     def __post_init__(self):
         super().__post_init__()
-        self._map_child_boundaries_onto_parent_grid_indices()
-        self._modify_child_topography_and_mask()
+        self._map_child_boundaries_onto_parent_grid_indices(verbose=self.verbose)
+        self._modify_child_topography_and_mask(verbose=self.verbose)
 
-    def _map_child_boundaries_onto_parent_grid_indices(self):
+    def _map_child_boundaries_onto_parent_grid_indices(self, verbose: bool = False):
         """Maps child grid boundary points onto absolute indices of the parent grid."""
-        # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
-        parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
-
-        # Map child boundaries onto parent grid indices
-        ds_nesting = map_child_boundaries_onto_parent_grid_indices(
-            parent_grid_ds,
-            child_grid_ds,
-            self.boundaries,
-            self.metadata["prefix"],
-            self.metadata["period"],
-        )
+        with Timed(
+            "=== Mapping the child grid boundary points onto the indices of the parent grid ===",
+            verbose=verbose,
+        ):
+            # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
+            parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
+
+            # Map child boundaries onto parent grid indices
+            ds_nesting = map_child_boundaries_onto_parent_grid_indices(
+                parent_grid_ds,
+                child_grid_ds,
+                self.boundaries,
+                self.metadata["prefix"],
+                self.metadata["period"],
+            )
 
-        self.ds_nesting = ds_nesting
+            self.ds_nesting = ds_nesting
 
-    def _modify_child_topography_and_mask(self):
+    def _modify_child_topography_and_mask(self, verbose: bool = False):
         """Adjust the topography and mask of the child grid to align with the parent grid.
 
         Uses a weighted sum based on boundary distance and clips depth values to a
         minimum.
         """
-        # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
-        parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
-
-        child_grid_ds = modify_child_topography_and_mask(
-            parent_grid_ds, child_grid_ds, self.boundaries, self.hmin
-        )
+        with Timed("=== Modifying child topography and mask ===", verbose=verbose):
+            # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
+            parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
+            child_grid_ds = modify_child_topography_and_mask(
+                parent_grid_ds, child_grid_ds, self.boundaries, self.hmin
+            )
 
-        # Finalize grid datasets by adjusting longitudes back to [0, 360] range
-        parent_grid_ds, child_grid_ds = self._finalize_grid_datasets(
-            parent_grid_ds, child_grid_ds
-        )
+            # Finalize grid datasets by adjusting longitudes back to [0, 360] range
+            parent_grid_ds, child_grid_ds = self._finalize_grid_datasets(
+                parent_grid_ds, child_grid_ds
+            )
 
-        self.ds = child_grid_ds
+            self.ds = child_grid_ds
 
     def update_topography(
         self, topography_source=None, hmin=None, verbose=False
@@ -155,7 +164,7 @@ class ChildGrid(Grid):
         )
 
         # Modify child topography and mask to match the parent grid
-        self._modify_child_topography_and_mask()
+        self._modify_child_topography_and_mask(verbose=verbose)
 
     def plot_nesting(self, with_dim_names=False) -> None:
         """Plot the parent and child grids in a single figure.
@@ -216,13 +225,19 @@ class ChildGrid(Grid):
         write_to_yaml(forcing_dict, filepath)
 
     @classmethod
-    def from_yaml(cls, filepath: str | Path) -> "ChildGrid":
+    def from_yaml(
+        cls, filepath: str | Path, verbose: bool = False, **kwargs: Any
+    ) -> "ChildGrid":
         """Create an instance of the ChildGrid class from a YAML file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
+        verbose : bool, optional
+            Indicates whether to print grid generation steps with timing. Defaults to False.
+        **kwargs : Any
+            Additional keyword arguments passed to Grid.from_yaml.
 
         Returns
         -------
@@ -231,10 +246,12 @@ class ChildGrid(Grid):
         """
         filepath = Path(filepath)
 
-        parent_grid = Grid.from_yaml(filepath, "ParentGrid")
+        parent_grid = Grid.from_yaml(
+            filepath, verbose=verbose, section_name="ParentGrid"
+        )
         params = from_yaml(cls, filepath)
 
-        return cls(parent_grid=parent_grid, **params)
+        return cls(parent_grid=parent_grid, **params, verbose=verbose)
 
     def _prepare_grid_datasets(self) -> tuple[xr.Dataset, xr.Dataset]:
         """Prepare parent and child grid datasets by adjusting longitudes for dateline
@@ -258,7 +275,7 @@ class ChildGrid(Grid):
 
     def _finalize_grid_datasets(
         self, parent_grid_ds: xr.Dataset, child_grid_ds: xr.Dataset
-    ) -> None:
+    ) -> tuple[xr.Dataset, xr.Dataset]:
         """Finalize the grid datasets by converting longitudes back to the [0, 360]
         range.
 
@@ -269,13 +286,28 @@ class ChildGrid(Grid):
 
         child_grid_ds : xr.Dataset
             The child grid dataset after modifications.
+
+        Returns
+        -------
+        tuple[xr.Dataset, xr.Dataset]
+            The finalized parent and child grid datasets with longitudes wrapped to [0, 360].
+
         """
         parent_grid_ds = wrap_longitudes(parent_grid_ds, straddle=False)
         child_grid_ds = wrap_longitudes(child_grid_ds, straddle=False)
+
         return parent_grid_ds, child_grid_ds
 
     @classmethod
-    def from_file(cls, filepath: str | Path, verbose: bool = False) -> "ChildGrid":
+    def from_file(
+        cls,
+        filepath: str | Path,
+        theta_s: float | None = None,
+        theta_b: float | None = None,
+        hc: float | None = None,
+        N: int | None = None,
+        verbose: bool = False,
+    ) -> "ChildGrid":
         """This method is disabled in this subclass.
 
         .. noindex::
@@ -410,13 +442,6 @@ def map_child_boundaries_onto_parent_grid_indices(
     }
     ds = ds.rename(dims_to_rename)
 
-    ds = ds.assign_coords(
-        {
-            "indices_rho": ("two", ["xi", "eta"]),
-            "indices_vel": ("three", ["xi", "eta", "angle"]),
-        }
-    )
-
     return ds
 
 
@@ -643,7 +668,7 @@ def modify_child_topography_and_mask(
         alpha * child_grid_ds["h"] + (1 - alpha) * h_parent_interpolated
     )
     # Clip depth on modified child topography
-    child_grid_ds["h"] = _clip_depth(child_grid_ds["h"], hmin)
+    child_grid_ds["h"] = clip_depth(child_grid_ds["h"], hmin)
 
     child_mask = (
         alpha * child_grid_ds["mask_rho"] + (1 - alpha) * mask_parent_interpolated