roms-tools 2.3.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

roms_tools/setup/surface_forcing.py

@@ -5,8 +5,10 @@ from datetime import datetime
 import numpy as np
 import matplotlib.pyplot as plt
 from pathlib import Path
-from typing import Dict, Union, List
+import logging
+from typing import Dict, Union, List, Optional
 from roms_tools import Grid
+from roms_tools.utils import save_datasets
 from roms_tools.regrid import LateralRegrid
 from roms_tools.plot import _plot
 from roms_tools.setup.datasets import (
@@ -21,7 +23,6 @@ from roms_tools.setup.utils import (
     interpolate_from_climatology,
     get_variable_metadata,
     group_dataset,
-    save_datasets,
     rotate_velocities,
     convert_to_roms_time,
     _to_yaml,
@@ -37,10 +38,14 @@ class SurfaceForcing:
     ----------
     grid : Grid
         Object representing the grid information.
-    start_time : datetime
-        Start time of the desired surface forcing data.
-    end_time : datetime
-        End time of the desired surface forcing data.
+    start_time : datetime, optional
+        The start time of the desired surface forcing data. This time is used to filter the dataset
+        to include only records on or after this time, with a single record at or before this time.
+        If no time filtering is desired, set it to None. Default is None.
+    end_time : datetime, optional
+        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.
     source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
         Dictionary specifying the source of the surface forcing data. Keys include:
 
@@ -59,9 +64,14 @@ class SurfaceForcing:
           - "bgc": for biogeochemical forcing.
 
     correct_radiation : bool
-        Whether to correct shortwave radiation. Default is False.
-    use_coarse_grid: bool
-        Whether to interpolate to coarsened grid. Default is False.
+        Whether to correct shortwave radiation. Default is True.
+    coarse_grid_mode : str, optional
+        Specifies whether to interpolate onto grid coarsened by a factor of two. Options are:
+
+          - "auto" (default): Automatically decide based on the comparison of source and target spatial resolutions.
+          - "always": Always interpolate onto the coarse grid.
+          - "never": Never use the coarse grid; interpolate onto the fine grid instead.
+
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
     use_dask: bool, optional
@@ -84,12 +94,12 @@ class SurfaceForcing:
     """
 
     grid: Grid
-    start_time: datetime
-    end_time: datetime
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     type: str = "physics"
-    correct_radiation: bool = False
-    use_coarse_grid: bool = False
+    correct_radiation: bool = True
+    coarse_grid_mode: str = "auto"
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
     bypass_validation: bool = False
@@ -99,14 +109,29 @@ class SurfaceForcing:
     def __post_init__(self):
 
         self._input_checks()
+        data = self._get_data()
+
+        if self.coarse_grid_mode == "always":
+            use_coarse_grid = True
+        elif self.coarse_grid_mode == "never":
+            use_coarse_grid = False
+        elif self.coarse_grid_mode == "auto":
+            use_coarse_grid = self._determine_coarse_grid_usage(data)
+        if use_coarse_grid:
+            logging.info("Data will be interpolated onto grid coarsened by factor 2.")
+        else:
+            logging.info("Data will be interpolated onto fine grid.")
+        object.__setattr__(self, "use_coarse_grid", use_coarse_grid)
+
         target_coords = get_target_coords(self.grid, self.use_coarse_grid)
         object.__setattr__(self, "target_coords", target_coords)
 
-        data = self._get_data()
         data.choose_subdomain(
             target_coords,
             buffer_points=20,  # lateral fill needs some buffer from data margin
         )
+        # Enforce double precision to ensure reproducibility
+        data.convert_to_float64()
 
         data.apply_lateral_fill()
 
@@ -149,6 +174,18 @@ class SurfaceForcing:
         object.__setattr__(self, "ds", ds)
 
     def _input_checks(self):
+        # Check that start_time and end_time are both None or none of them is
+        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"]:
             raise ValueError("`type` must be either 'physics' or 'bgc'.")
@@ -166,6 +203,47 @@ class SurfaceForcing:
             {**self.source, "climatology": self.source.get("climatology", False)},
         )
 
+        # Validate 'coarse_grid_mode'
+        valid_modes = ["auto", "always", "never"]
+        if self.coarse_grid_mode not in valid_modes:
+            raise ValueError(
+                f"`coarse_grid_mode` must be one of {valid_modes}, but got '{self.coarse_grid_mode}'."
+            )
+
+    def _determine_coarse_grid_usage(self, data):
+        """Determine if coarse grid interpolation should be used based on the resolution
+        of the dataset and the target grid.
+
+        Parameters
+        ----------
+        data : object
+            The dataset object containing the data to be analyzed for grid spacing.
+
+        Returns
+        -------
+        use_coarse_grid : bool
+            Whether to use the coarse grid or not.
+        """
+        # Get the target coordinates and select the subdomain of the data
+        target_coords = get_target_coords(self.grid, use_coarse_grid=False)
+        data_coords = data.choose_subdomain(
+            target_coords, buffer_points=1, return_coords_only=True
+        )
+
+        # Compute minimal grid spacing in the data subdomain
+        min_grid_spacing_data = data.compute_minimal_grid_spacing(data_coords)
+
+        # Compute the maximum grid spacing in the ROMS grid
+        max_grid_spacing = max((1 / self.grid.ds.pm).max(), (1 / self.grid.ds.pn).max())
+
+        # Determine whether to use coarse grid based on grid spacing comparison
+        if 2 * max_grid_spacing < min_grid_spacing_data:
+            use_coarse_grid = True
+        else:
+            use_coarse_grid = False
+
+        return use_coarse_grid
+
     def _get_data(self):
 
         data_dict = {
@@ -238,14 +316,14 @@ class SurfaceForcing:
                 "qair": {**default_info, "validate": True},
                 "rain": {**default_info, "validate": False},
                 "uwnd": {
-                    "location": "u",
+                    "location": "rho",
                     "is_vector": True,
                     "vector_pair": "vwnd",
                     "is_3d": False,
                     "validate": True,
                 },
                 "vwnd": {
-                    "location": "v",
+                    "location": "rho",
                     "is_vector": True,
                     "vector_pair": "uwnd",
                     "is_3d": False,
@@ -266,7 +344,7 @@ class SurfaceForcing:
     def _apply_correction(self, processed_fields, data):
 
         correction_data = self._get_correction_data()
-        # choose same subdomain as forcing data so that we can use same mask
+        # Match subdomain to forcing data to reuse the mask
         coords_correction = {
             "lat": data.ds[data.dim_names["latitude"]],
             "lon": data.ds[data.dim_names["longitude"]],
@@ -275,23 +353,39 @@ class SurfaceForcing:
             coords_correction, straddle=self.target_coords["straddle"]
         )
         correction_data.ds["mask"] = data.ds["mask"]  # use mask from ERA5 data
+        correction_data.ds["time"] = correction_data.ds["time"].dt.days
+
         correction_data.apply_lateral_fill()
-        # regrid
-        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
-        corr_factor = lateral_regrid.apply(
-            correction_data.ds[correction_data.var_names["swr_corr"]]
-        )
 
-        # temporal interpolation
-        corr_factor = interpolate_from_climatology(
-            corr_factor,
-            correction_data.dim_names["time"],
-            time=processed_fields["swrad"].time,
-        )
+        # Temporal interpolation: Perform before spatial regridding for better performance
+        if self.use_dask:
+            # Perform temporal interpolation for each time slice to enforce chunking in time.
+            # This reduces memory usage by processing one time step at a time.
+            # The interpolated slices are then concatenated along the "time" dimension.
+            corr_factor = xr.concat(
+                [
+                    interpolate_from_climatology(
+                        correction_data.ds[correction_data.var_names["swr_corr"]],
+                        correction_data.dim_names["time"],
+                        time=time,
+                    )
+                    for time in processed_fields["swrad"].time
+                ],
+                dim="time",
+            )
+        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"],
+                time=processed_fields["swrad"].time,
+            )
 
-        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
+        # Spatial regridding
+        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
+        corr_factor = lateral_regrid.apply(corr_factor)
 
-        del corr_factor
+        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
 
         return processed_fields
 
@@ -357,12 +451,8 @@ class SurfaceForcing:
 
         for var_name in ds.data_vars:
             if self.variable_info[var_name]["validate"]:
-                if self.variable_info[var_name]["location"] == "rho":
-                    mask = self.target_coords["mask"]
-                elif self.variable_info[var_name]["location"] == "u":
-                    mask = self.target_coords["mask_u"]
-                elif self.variable_info[var_name]["location"] == "v":
-                    mask = self.target_coords["mask_v"]
+                # all variables are at rho-points
+                mask = self.target_coords["mask"]
                 nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
@@ -476,42 +566,26 @@ class SurfaceForcing:
     def save(
         self,
         filepath: Union[str, Path],
-        np_eta: int = None,
-        np_xi: int = None,
-        group: bool = False,
+        group: bool = True,
     ) -> None:
         """Save the surface forcing fields to one or more netCDF4 files.
 
-        This method saves the dataset either as a single file or as multiple files depending on the partitioning and grouping options.
-        The dataset can be saved in two modes:
-
-        1. **Single File Mode (default)**:
-            - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file.
-            - The file is named based on the `filepath`, with `.nc` automatically appended.
-
-        2. **Partitioned Mode**:
-            - If either `np_eta` or `np_xi` is specified, the dataset is partitioned into spatial tiles along the `eta` and `xi` axes.
-            - Each tile is saved as a separate netCDF4 file, and filenames are modified with an index (e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`).
-
-        Additionally, if `group` is set to `True`, the dataset is first grouped into temporal subsets, resulting in multiple grouped files before partitioning and saving.
+        This method saves the dataset to disk as either a single netCDF4 file or multiple files, depending on the `group` parameter.
+        If `group` is `True`, the dataset is divided into subsets (e.g., monthly or yearly) based on the temporal frequency
+        of the data, and each subset is saved to a separate file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
-            The base path and filename for the output files. The format of the filenames depends on whether partitioning is used
-            and the temporal range of the data. For partitioned datasets, files will be named with an additional index, e.g.,
-            `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
-        group: bool, optional
-            If `True`, groups the dataset into multiple files based on temporal data frequency. Defaults to `False`.
+            The base path and filename for the output file(s). If `group` is `True`, the filenames will include additional
+            time-based information (e.g., year or month) to distinguish the subsets.
+        group : bool, optional
+            Whether to divide the dataset into multiple files based on temporal frequency. Defaults to `True`.
 
         Returns
         -------
         List[Path]
-            A list of Path objects for the filenames that were saved.
+            A list of `Path` objects representing the filenames of the saved file(s).
         """
 
         # Ensure filepath is a Path object
@@ -521,12 +595,6 @@ class SurfaceForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds.load()
-
         if group:
             dataset_list, output_filenames = group_dataset(self.ds, str(filepath))
         else:
@@ -534,7 +602,7 @@ class SurfaceForcing:
             output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+            dataset_list, output_filenames, use_dask=self.use_dask
         )
 
         return saved_filenames
@@ -556,7 +624,6 @@ class SurfaceForcing:
         cls,
         filepath: Union[str, Path],
         use_dask: bool = False,
-        bypass_validation: bool = False,
     ) -> "SurfaceForcing":
         """Create an instance of the SurfaceForcing class from a YAML file.
 
@@ -566,10 +633,6 @@ class SurfaceForcing:
             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
         -------
@@ -581,6 +644,4 @@ class SurfaceForcing:
         grid = Grid.from_yaml(filepath)
         params = _from_yaml(cls, filepath)
 
-        return cls(
-            grid=grid, **params, use_dask=use_dask, bypass_validation=bypass_validation
-        )
+        return cls(grid=grid, **params, use_dask=use_dask)

roms_tools/setup/tides.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass, field
 from roms_tools import Grid
 from roms_tools.plot import _plot
 from roms_tools.regrid import LateralRegrid
+from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
     nan_check,
@@ -16,7 +17,6 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     get_variable_metadata,
-    save_datasets,
     get_target_coords,
     rotate_velocities,
     get_vector_pairs,
@@ -84,6 +84,9 @@ class TidalForcing:
             target_coords,
             buffer_points=20,
         )
+        # Enforce double precision to ensure reproducibility
+        data.convert_to_float64()
+
         # select desired number of constituents
         object.__setattr__(data, "ds", data.ds.isel(ntides=slice(None, self.ntides)))
         self._correct_tides(data)
@@ -368,36 +371,19 @@ class TidalForcing:
             kwargs=kwargs,
         )
 
-    def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
-    ) -> None:
+    def save(self, filepath: Union[str, Path]) -> None:
         """Save the tidal forcing information to a netCDF4 file.
 
-        This method supports saving the dataset in two modes:
-
-          1. **Single File Mode (default)**:
-
-            If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file
-            with the base filename specified by `filepath.nc`.
-
-          2. **Partitioned Mode**:
-
-            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-            - Each spatial tile is saved as a separate netCDF4 file.
-
         Parameters
         ----------
         filepath : Union[str, Path]
-            The base path or filename where the dataset should be saved.
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
+            The path or filename where the dataset will be saved. If a directory is specified,
+            the file will be saved with a default name within that directory.
 
         Returns
         -------
-        List[Path]
-            A list of Path objects for the filenames that were saved.
+        Path
+            A `Path` object representing the location of the saved file.
         """
 
         # Ensure filepath is a Path object
@@ -407,17 +393,11 @@ class TidalForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds.load()
-
         dataset_list = [self.ds]
         output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+            dataset_list, output_filenames, use_dask=self.use_dask
         )
 
         return saved_filenames
@@ -439,7 +419,6 @@ class TidalForcing:
         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.
 
@@ -449,10 +428,6 @@ 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
         -------
@@ -467,7 +442,6 @@ class TidalForcing:
             grid=grid,
             **tidal_forcing_params,
             use_dask=use_dask,
-            bypass_validation=bypass_validation,
         )
 
     def _correct_tides(self, data):

roms_tools/setup/topography.py

@@ -145,6 +145,8 @@ def _make_raw_topography(
         The regridded topography data with the sign flipped (bathymetry positive).
     """
     data.choose_subdomain(target_coords, buffer_points=3, verbose=verbose)
+    # Enforce double precision to ensure reproducibility
+    data.convert_to_float64()
 
     if verbose:
         start_time = time.time()
@@ -233,7 +235,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         rmax_log = 0.0
 
     # Apply hmin threshold
-    h = xr.where(h < hmin, hmin, h)
+    h = _clip_depth(h, hmin)
 
     # Perform logarithmic transformation of the height field
     h_log = np.log(h / hmin)
@@ -316,7 +318,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         h = hmin * np.exp(h_log)
 
         # Apply hmin threshold again
-        h = xr.where(h < hmin, hmin, h)
+        h = _clip_depth(h, hmin)
 
         # Compute maximum slope parameter r
         r_eta, r_xi = _compute_rfactor(h)
@@ -327,6 +329,27 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     return h
 
 
+def _clip_depth(h: xr.DataArray, hmin: float) -> xr.DataArray:
+    """Ensures that depth values do not fall below a minimum threshold.
+
+    This function replaces all depth values in `h` that are less than `hmin` with `hmin`,
+    ensuring a minimum depth constraint.
+
+    Parameters
+    ----------
+    h : xr.DataArray
+        The depth (bathymetry) array.
+    hmin : float
+        The minimum allowable depth value.
+
+    Returns
+    -------
+    xr.DataArray
+        The modified depth array with values clipped at `hmin`.
+    """
+    return xr.where(h < hmin, hmin, h)
+
+
 def _compute_rfactor(h):
     """Computes the slope parameter (r-factor) in both horizontal directions.