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

roms_tools/setup/river_forcing.py

@@ -9,13 +9,13 @@ from pathlib import Path
 import matplotlib.pyplot as plt
 from roms_tools import Grid
 from roms_tools.plot import _get_projection, _add_field_to_ax
+from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import DaiRiverDataset
 from roms_tools.setup.utils import (
     get_target_coords,
     gc_dist,
     substitute_nans_by_fillvalue,
     convert_to_roms_time,
-    save_datasets,
     _to_yaml,
     _from_yaml,
     get_variable_metadata,
@@ -91,8 +91,9 @@ class RiverForcing:
         object.__setattr__(self, "original_indices", original_indices)
 
         if len(original_indices["station"]) > 0:
-            self._move_rivers_to_closest_coast(target_coords, data)
             ds = self._create_river_forcing(data)
+            self._move_rivers_to_closest_coast(target_coords, data)
+            ds = self._write_indices_into_dataset(ds)
             self._validate(ds)
 
             for var_name in ds.data_vars:
@@ -345,49 +346,44 @@ class RiverForcing:
             "xi_rho": indices[2],
             "name": names,
         }
-        self._write_indices_into_grid_file(indices)
         object.__setattr__(self, "updated_indices", indices)
 
-    def _write_indices_into_grid_file(self, indices):
-        """Writes river location indices into the grid dataset as the "river_flux"
-        variable.
+    def _write_indices_into_dataset(self, ds):
+        """Adds river location indices to the dataset as the "river_location" variable.
 
-        This method checks if the variable "river_flux" already exists in the grid dataset.
-        If it does, the method removes it. Then, it creates a new variable called "river_flux"
-        based on the provided `indices` and assigns it to the dataset. The `indices` dictionary
-        provides information about the river station locations and their corresponding grid
-        cell indices (eta_rho and xi_rho).
+        This method checks if the "river_location" variable already exists in the dataset.
+        If it does, the method removes it. Then, it creates a new "river_location" variable
+        using river station indices from `self.updated_indices` and assigns it to the dataset.
+        The indices specify the river station locations in terms of eta_rho and xi_rho grid cell indices.
 
-        Parameters:
-        -----------
-        indices : dict
-            A dictionary containing information about river station locations.
-            It must contain the following keys:
-            - "name" (list or array): Names of the river stations.
-            - "station" (list or array): River station identifiers.
-            - "eta_rho" (list or array): The eta (row) index for each river location.
-            - "xi_rho" (list or array): The xi (column) index for each river location.
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to which the "river_location" variable will be added.
 
-        Returns:
-        --------
-        None
-            This method modifies the grid's dataset in-place by adding the "river_flux" variable.
+        Returns
+        -------
+        xarray.Dataset
+            The modified dataset with the "river_location" variable added.
         """
 
-        if "river_flux" in self.grid.ds:
-            ds = self.grid.ds.drop_vars("river_flux")
-            object.__setattr__(self.grid, "ds", ds)
+        if "river_location" in ds:
+            ds = ds.drop_vars("river_location")
 
         river_locations = xr.zeros_like(self.grid.ds.mask_rho)
-        for i in range(len(indices["name"])):
-            station = indices["station"][i]
-            eta_index = indices["eta_rho"][i]
-            xi_index = indices["xi_rho"][i]
+        for i in range(len(self.updated_indices["name"])):
+            station = self.updated_indices["station"][i]
+            eta_index = self.updated_indices["eta_rho"][i]
+            xi_index = self.updated_indices["xi_rho"][i]
             river_locations[eta_index, xi_index] = station + 2
 
-        river_locations.attrs["long_name"] = "River volume flux partition"
+        river_locations.attrs["long_name"] = "River index plus local volume fraction"
         river_locations.attrs["units"] = "none"
-        self.grid.ds["river_flux"] = river_locations
+        ds["river_location"] = river_locations
+
+        ds = ds.drop_vars(["lat_rho", "lon_rho"])
+
+        return ds
 
     def _validate(self, ds):
         """Validates the dataset by checking for NaN values in river forcing data.
@@ -592,37 +588,13 @@ class RiverForcing:
     def save(
         self,
         filepath: Union[str, Path],
-        filepath_grid: Union[str, Path],
-        np_eta: int = None,
-        np_xi: int = None,
     ) -> None:
-        """Save the river forcing and grid file to netCDF4 files. The grid file is
-        required because a new field `river_flux` has been added.
-
-        This method allows saving the river forcing and grid data either each as a single file or each partitioned into multiple files, based on the provided 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 provided `filepath`, with `.nc` automatically appended to the filename.
-
-        2. **Partitioned Mode**:
-            - If either `np_eta` or `np_xi` is specified, the dataset is partitioned spatially along the `eta` and `xi` axes into tiles.
-            - Each tile is saved as a separate netCDF4 file. Filenames will be modified with an index to represent each partition, e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
+        """Save the river forcing to netCDF4 file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
-            The base path and filename for the output files. The filenames will include the specified path and the `.nc` extension.
-            If partitioning is used, additional indices will be appended to the filenames, e.g., `"filepath.0.nc"`, `"filepath.1.nc"`, etc.
-
-        filepath_grid : Union[str, Path]
-            The base path and filename for saving the grid file.
-
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed along the `eta` axis.
-
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed along the `xi` axis.
+            The base path and filename for the output files.
 
         Returns
         -------
@@ -632,20 +604,15 @@ class RiverForcing:
 
         # Ensure filepath is a Path object
         filepath = Path(filepath)
-        filepath_grid = Path(filepath_grid)
 
         # Remove ".nc" suffix if present
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
-        if filepath_grid.suffix == ".nc":
-            filepath_grid = filepath_grid.with_suffix("")
 
-        dataset_list = [self.ds, self.grid.ds]
-        output_filenames = [str(filepath), str(filepath_grid)]
+        dataset_list = [self.ds]
+        output_filenames = [str(filepath)]
 
-        saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
-        )
+        saved_filenames = save_datasets(dataset_list, output_filenames)
 
         return saved_filenames
 

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
+import logging
 from typing import Dict, Union, List
 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,
@@ -59,9 +60,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
@@ -88,8 +94,8 @@ class SurfaceForcing:
     end_time: datetime
     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,10 +105,23 @@ 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
@@ -166,6 +185,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 +298,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 +326,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 +335,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 +433,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 +548,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 +577,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 +584,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

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,
@@ -368,36 +368,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 +390,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