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

roms_tools/setup/datasets.py

@@ -25,7 +25,7 @@ from roms_tools.setup.fill import LateralFill
 # lat-lon datasets
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class Dataset:
     """Represents forcing data on original grid.
 
@@ -35,10 +35,11 @@ class Dataset:
         The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
         or a list of strings or Path objects containing multiple files.
     start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
+        Start time for selecting relevant data. If not provided, no time-based filtering is applied.
     end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
+        End time for selecting relevant data. If not provided, the dataset selects the time entry
+        closest to `start_time` within the range `[start_time, start_time + 24 hours]`.
+        If `start_time` is also not provided, no time-based filtering is applied.
     dim_names: Dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset.
     var_names: Dict[str, str]
@@ -131,8 +132,8 @@ class Dataset:
         self.infer_horizontal_resolution(ds)
 
         # Check whether the data covers the entire globe
-        object.__setattr__(self, "is_global", self.check_if_global(ds))
-        object.__setattr__(self, "ds", ds)
+        self.is_global = self.check_if_global(ds)
+        self.ds = ds
 
         if self.apply_post_processing:
             self.post_process()
@@ -353,7 +354,7 @@ class Dataset:
         resolution = np.mean([lat_resolution, lon_resolution])
 
         # Set the computed resolution as an attribute
-        object.__setattr__(self, "resolution", resolution)
+        self.resolution = resolution
 
     def compute_minimal_grid_spacing(self, ds: xr.Dataset):
         """Compute the minimal grid spacing in a dataset based on latitude and longitude
@@ -510,6 +511,25 @@ class Dataset:
         """
         pass
 
+    def convert_to_float64(self) -> None:
+        """Convert all data variables in the dataset to float64.
+
+        This method updates the dataset by converting all of its data variables to the
+        `float64` data type, ensuring consistency for numerical operations that require
+        high precision. The dataset is modified in place.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        None
+            This method modifies the dataset in place and does not return anything.
+        """
+        ds = self.ds.astype({var: "float64" for var in self.ds.data_vars})
+        self.ds = ds
+
     def choose_subdomain(
         self,
         target_coords,
@@ -651,7 +671,7 @@ class Dataset:
         if return_copy:
             return Dataset.from_ds(self, subdomain)
         else:
-            object.__setattr__(self, "ds", subdomain)
+            self.ds = subdomain
 
     def apply_lateral_fill(self):
         """Apply lateral fill to variables using the dataset's mask and grid dimensions.
@@ -737,7 +757,7 @@ class Dataset:
         dataset = cls.__new__(cls)
 
         # Directly set the provided dataset as the 'ds' attribute
-        object.__setattr__(dataset, "ds", ds)
+        dataset.ds = ds
 
         # Copy all other attributes from the original data instance
         for attr in vars(original_dataset):
@@ -747,7 +767,7 @@ class Dataset:
         return dataset
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class TPXODataset(Dataset):
     """Represents tidal data on the original grid from the TPXO dataset.
 
@@ -838,15 +858,11 @@ class TPXODataset(Dataset):
             {"nx": "longitude", "ny": "latitude", self.dim_names["ntides"]: "ntides"}
         )
 
-        object.__setattr__(
-            self,
-            "dim_names",
-            {
-                "latitude": "latitude",
-                "longitude": "longitude",
-                "ntides": "ntides",
-            },
-        )
+        self.dim_names = {
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "ntides": "ntides",
+        }
 
         return ds
 
@@ -889,15 +905,15 @@ class TPXODataset(Dataset):
             ds["mask"] = mask
             ds = ds.drop_vars(["depth"])
 
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
         # Remove "depth" from var_names
         updated_var_names = {**self.var_names}  # Create a copy of the dictionary
         updated_var_names.pop("depth", None)  # Remove "depth" if it exists
-        object.__setattr__(self, "var_names", updated_var_names)
+        self.var_names = updated_var_names
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class GLORYSDataset(Dataset):
     """Represents GLORYS data on original grid.
 
@@ -976,7 +992,7 @@ class GLORYSDataset(Dataset):
         self.ds["mask_vel"] = mask_vel
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMDataset(Dataset):
     """Represents CESM data on original grid.
 
@@ -1057,12 +1073,12 @@ class CESMDataset(Dataset):
             # Update dimension names
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["time"] = "time"
-            object.__setattr__(self, "dim_names", updated_dim_names)
+            self.dim_names = updated_dim_names
 
         return ds
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMBGCDataset(CESMDataset):
     """Represents CESM BGC data on original grid.
 
@@ -1164,12 +1180,12 @@ class CESMBGCDataset(CESMDataset):
             if "z_t_150m" in ds.variables:
                 ds = ds.drop_vars("z_t_150m")
             # update dataset
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Update dim_names with "depth": "depth" key-value pair
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["depth"] = "depth"
-            object.__setattr__(self, "dim_names", updated_dim_names)
+            self.dim_names = updated_dim_names
 
         mask = xr.where(
             self.ds[self.var_names["PO4"]]
@@ -1182,7 +1198,7 @@ class CESMBGCDataset(CESMDataset):
         self.ds["mask"] = mask
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMBGCSurfaceForcingDataset(CESMDataset):
     """Represents CESM BGC surface forcing data on original grid.
 
@@ -1238,7 +1254,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 
         if "z_t" in self.ds.variables:
             ds = self.ds.drop_vars("z_t")
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
         mask = xr.where(
             self.ds[self.var_names["pco2_air"]]
@@ -1251,7 +1267,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
         self.ds["mask"] = mask
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ERA5Dataset(Dataset):
     """Represents ERA5 data on original grid.
 
@@ -1351,27 +1367,27 @@ class ERA5Dataset(Dataset):
             ds["qair"].attrs["long_name"] = "Absolute humidity at 2m"
             ds["qair"].attrs["units"] = "kg/kg"
             ds = ds.drop_vars([self.var_names["d2m"]])
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Update var_names dictionary
             var_names = {**self.var_names, "qair": "qair"}
             var_names.pop("d2m")
-            object.__setattr__(self, "var_names", var_names)
+            self.var_names = var_names
 
         if "mask" in self.var_names.keys():
             ds = self.ds
             mask = xr.where(self.ds[self.var_names["mask"]].isel(time=0).isnull(), 0, 1)
             ds["mask"] = mask
             ds = ds.drop_vars([self.var_names["mask"]])
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Remove mask from var_names dictionary
             var_names = self.var_names
             var_names.pop("mask")
-            object.__setattr__(self, "var_names", var_names)
+            self.var_names = var_names
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ERA5Correction(Dataset):
     """Global dataset to correct ERA5 radiation. The dataset contains multiplicative
     correction factors for the ERA5 shortwave radiation, obtained by comparing the
@@ -1473,10 +1489,10 @@ class ERA5Correction(Dataset):
             raise ValueError(
                 "The correction dataset does not contain all specified longitude values."
             )
-        object.__setattr__(self, "ds", subdomain)
+        self.ds = subdomain
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ETOPO5Dataset(Dataset):
     """Represents topography data on the original grid from the ETOPO5 dataset.
 
@@ -1533,7 +1549,7 @@ class ETOPO5Dataset(Dataset):
         return ds
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class SRTM15Dataset(Dataset):
     """Represents topography data on the original grid from the SRTM15 dataset.
 
@@ -1569,7 +1585,7 @@ class SRTM15Dataset(Dataset):
 
 
 # river datasets
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class RiverDataset:
     """Represents river data.
 
@@ -1627,7 +1643,7 @@ class RiverDataset:
 
         # Select relevant times
         ds = self.add_time_info(ds)
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def load_data(self) -> xr.Dataset:
         """Load dataset from the specified file.
@@ -1765,13 +1781,13 @@ class RiverDataset:
 
         ds = assign_dates_to_climatology(self.ds, "month")
         ds = ds.swap_dims({"month": "time"})
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
         updated_dim_names = {**self.dim_names}
         updated_dim_names["time"] = "time"
-        object.__setattr__(self, "dim_names", updated_dim_names)
+        self.dim_names = updated_dim_names
 
-        object.__setattr__(self, "climatology", True)
+        self.climatology = True
 
     def sort_by_river_volume(self, ds: xr.Dataset) -> xr.Dataset:
         """Sorts the dataset by river volume in descending order (largest rivers first),
@@ -1843,15 +1859,10 @@ class RiverDataset:
 
         Returns
         -------
-        indices : dict
+        indices : dict[str, list[tuple]]
             A dictionary containing the indices of the rivers that are within the threshold distance from
-            the target coordinates. The dictionary keys are:
-            - "station" : numpy.ndarray
-                The indices of the rivers that satisfy the distance threshold.
-            - "eta_rho" : numpy.ndarray
-                The indices of the `eta_rho` dimension corresponding to the selected stations.
-            - "xi_rho" : numpy.ndarray
-                The indices of the `xi_rho` dimension corresponding to the selected stations.
+            the target coordinates. The dictionary structure consists of river names as keys, and each value is a list of tuples. Each tuple represents
+            a pair of indices corresponding to the `eta_rho` and `xi_rho` grid coordinates of the river.
         """
 
         # Retrieve longitude and latitude of river mouths
@@ -1878,33 +1889,78 @@ class RiverDataset:
 
             # Find the indices of the closest grid cell to the river mouth
             indices = np.where(dist == dist_min)
+            stations = indices[0]
+            eta_rho_values = indices[1]
+            xi_rho_values = indices[2]
             names = (
-                self.ds[self.var_names["name"]]
-                .isel({self.dim_names["station"]: indices[0]})
+                ds[self.var_names["name"]]
+                .isel({self.dim_names["station"]: stations})
                 .values
             )
-            # Return the indices in a dictionary format
-            indices = {
-                "station": indices[0],
-                "eta_rho": indices[1],
-                "xi_rho": indices[2],
-                "name": names,
-            }
+            river_indices = {}
+            for i in range(len(stations)):
+                river_name = names[i]
+                river_indices[river_name] = [
+                    (int(eta_rho_values[i]), int(xi_rho_values[i]))
+                ]  # list of tuples
         else:
             ds = xr.Dataset()
-            indices = {
-                "station": [],
-                "eta_rho": [],
-                "xi_rho": [],
-                "name": [],
-            }
+            river_indices = {}
+
+        self.ds = ds
+
+        return river_indices
+
+    def extract_named_rivers(self, indices):
+        """Extracts a subset of the dataset based on the provided river names in the
+        indices dictionary.
+
+        This method filters the dataset to include only the rivers specified in the `indices` dictionary.
+        The resulting subset is stored in the `ds` attribute of the class.
 
-        object.__setattr__(self, "ds", ds)
+        Parameters
+        ----------
+        indices : dict
+            A dictionary where the keys are river names (strings) and the values are dictionaries
+            containing river-related data (e.g., river indices, coordinates).
+
+        Returns
+        -------
+        None
+            The method modifies the `self.ds` attribute in place, setting it to the filtered dataset
+            containing only the data related to the specified rivers.
+
+        Raises
+        ------
+        ValueError
+            - If `indices` is not a dictionary.
+            - If any of the requested river names are not found in the dataset.
+        """
+
+        if not isinstance(indices, dict):
+            raise ValueError("`indices` must be a dictionary.")
 
-        return indices
+        river_names = list(indices.keys())
+
+        # Ensure the dataset is filtered based on the provided river names
+        ds_filtered = self.ds.where(
+            self.ds[self.var_names["name"]].isin(river_names), drop=True
+        )
+
+        # Check that all requested rivers exist in the dataset
+        filtered_river_names = set(ds_filtered[self.var_names["name"]].values)
+        missing_rivers = set(river_names) - filtered_river_names
+
+        if missing_rivers:
+            raise ValueError(
+                f"The following rivers were not found in the dataset: {missing_rivers}"
+            )
+
+        # Set the filtered dataset as the new `ds`
+        self.ds = ds_filtered
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class DaiRiverDataset(RiverDataset):
     """Represents river data from the Dai river dataset.
 
@@ -2050,7 +2106,7 @@ def _check_dataset(
 
 
 def _select_relevant_times(
-    ds, time_dim, start_time=None, end_time=None, climatology=False
+    ds, time_dim, start_time, end_time=None, climatology=False
 ) -> xr.Dataset:
     """Select a subset of the dataset based on the specified time range.
 
@@ -2067,11 +2123,10 @@ def _select_relevant_times(
         The input dataset to be filtered. Must contain a time dimension.
     time_dim: str
         Name of time dimension.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
+    start_time : datetime
+        The start time for selecting relevant data.
     end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
+        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided.
     climatology : bool
         Indicates whether the dataset is climatological. Defaults to False.
 

roms_tools/setup/grid.py

@@ -18,12 +18,13 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_v,
     get_target_coords,
     gc_dist,
+    _pop_grid_data,
 )
 from roms_tools.setup.utils import extract_single_value
 from pathlib import Path
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class Grid:
     """A single ROMS grid, used for creating, plotting, and then saving a new ROMS
     domain grid.
@@ -112,7 +113,7 @@ class Grid:
         # Coarsen the dataset if needed
         self._coarsen()
 
-        # Topography and mask
+        # Topography
         self.update_topography(
             topography_source=self.topography_source,
             hmin=self.hmin,
@@ -130,7 +131,7 @@ class Grid:
 
     def _input_checks(self):
         if self.topography_source is None:
-            object.__setattr__(self, "topography_source", {"name": "ETOPO5"})
+            self.topography_source = {"name": "ETOPO5"}
 
         if "name" not in self.topography_source:
             raise ValueError(
@@ -156,7 +157,7 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def update_topography(
         self, topography_source=None, hmin=None, verbose=False
@@ -205,7 +206,7 @@ class Grid:
                 f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ==="
             )
 
-        # Add topography and mask to the dataset
+        # Add topography to the dataset
         ds = _add_topography(
             ds=self.ds,
             target_coords=target_coords,
@@ -222,9 +223,9 @@ class Grid:
             )
 
         # Update the grid's dataset and related attributes
-        object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "topography_source", topography_source)
-        object.__setattr__(self, "hmin", hmin)
+        self.ds = ds
+        self.topography_source = topography_source
+        self.hmin = hmin
 
     def update_vertical_coordinate(
         self, N=None, theta_s=None, theta_b=None, hc=None, verbose=False
@@ -320,11 +321,11 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "theta_s", theta_s)
-        object.__setattr__(self, "theta_b", theta_b)
-        object.__setattr__(self, "hc", hc)
-        object.__setattr__(self, "N", N)
+        self.ds = ds
+        self.theta_s = theta_s
+        self.theta_b = theta_b
+        self.hc = hc
+        self.N = N
 
     def _straddle(self) -> None:
         """Check if the Greenwich meridian goes through the domain.
@@ -341,9 +342,9 @@ class Grid:
             np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
             or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
         ):
-            object.__setattr__(self, "straddle", True)
+            self.straddle = True
         else:
-            object.__setattr__(self, "straddle", False)
+            self.straddle = False
 
     def _coarsen(self):
         """Update the grid by adding grid variables that are coarsened versions of the
@@ -390,7 +391,7 @@ class Grid:
             ds[coarse_var].attrs["long_name"] = f"{long_name} on coarsened grid"
             ds[coarse_var].attrs["units"] = ds[fine_var].attrs["units"]
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def plot(
         self, bathymetry: bool = True, title: str = None, with_dim_names: bool = False
@@ -460,7 +461,7 @@ class Grid:
         xi : int, optional
             The xi-index to plot. Default is None.
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots.
 
         Returns
         -------
@@ -580,14 +581,14 @@ class Grid:
         grid = cls.__new__(cls)
 
         # Set the dataset for the grid instance
-        object.__setattr__(grid, "ds", ds)
+        grid.ds = ds
 
         # Check if the Greenwich meridian goes through the domain.
         grid._straddle()
 
         if not all(coord in grid.ds for coord in ["lat_u", "lon_u", "lat_v", "lon_v"]):
             ds = _add_lat_lon_at_velocity_points(grid.ds, grid.straddle)
-            object.__setattr__(grid, "ds", ds)
+            grid.ds = ds
 
         # Coarsen the grid if necessary
         if not all(
@@ -605,7 +606,7 @@ class Grid:
         for var in ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]:
             if var not in ds.coords:
                 ds = grid.ds.set_coords(var)
-                object.__setattr__(grid, "ds", ds)
+                grid.ds = ds
 
         # Update vertical coordinate if necessary
         if not all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
@@ -619,14 +620,14 @@ class Grid:
                 N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=True
             )
         else:
-            object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
-            object.__setattr__(grid, "theta_b", ds.attrs["theta_b"].item())
-            object.__setattr__(grid, "hc", ds.attrs["hc"].item())
-            object.__setattr__(grid, "N", len(ds.s_rho))
+            grid.theta_s = ds.attrs["theta_s"].item()
+            grid.theta_b = ds.attrs["theta_b"].item()
+            grid.hc = ds.attrs["hc"].item()
+            grid.N = len(ds.s_rho)
 
         # Manually set the remaining attributes by extracting parameters from dataset
-        object.__setattr__(grid, "nx", ds.sizes["xi_rho"] - 2)
-        object.__setattr__(grid, "ny", ds.sizes["eta_rho"] - 2)
+        grid.nx = ds.sizes["xi_rho"] - 2
+        grid.ny = ds.sizes["eta_rho"] - 2
         if "center_lon" in ds.attrs:
             center_lon = ds.attrs["center_lon"]
         elif "tra_lon" in ds:
@@ -636,7 +637,7 @@ class Grid:
                 "Missing grid information: 'center_lon' attribute or 'tra_lon' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "center_lon", center_lon)
+        grid.center_lon = center_lon
         if "center_lat" in ds.attrs:
             center_lat = ds.attrs["center_lat"]
         elif "tra_lat" in ds:
@@ -646,7 +647,7 @@ class Grid:
                 "Missing grid information: 'center_lat' attribute or 'tra_lat' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "center_lat", center_lat)
+        grid.center_lat = center_lat
         if "rot" in ds.attrs:
             rot = ds.attrs["rot"]
         elif "rotate" in ds:
@@ -656,7 +657,7 @@ class Grid:
                 "Missing grid information: 'rot' attribute or 'rotate' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "rot", rot)
+        grid.rot = rot
 
         for attr in [
             "size_x",
@@ -688,9 +689,7 @@ class Grid:
         filepath = Path(filepath)
 
         data = asdict(self)
-        data.pop("ds", None)
-        data.pop("straddle", None)
-        data.pop("verbose", None)
+        data = _pop_grid_data(data)
 
         # Include the version of roms-tools
         try:
@@ -849,7 +848,7 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _add_global_metadata(self, ds):
         """Add global metadata and attributes to the dataset.