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

roms_tools/__init__.py

@@ -20,7 +20,7 @@ from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
 from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
 from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
-from roms_tools.setup.datasets import get_glorys_bounds  # noqa: F401
+from roms_tools.datasets.lat_lon_datasets import get_glorys_bounds  # noqa: F401
 from roms_tools.tiling.join import open_partitions, join_netcdf  # noqa: F401
 
 

roms_tools/analysis/cdr_ensemble.py

@@ -59,7 +59,7 @@ class Ensemble:
         Parameters
         ----------
         ds : xr.Dataset
-            Dataset containing a "cdr_efficiency" variable and "abs_time" coordinate.
+            Dataset containing a "cdr_efficiency" variable and "time" coordinate.
 
         Returns
         -------
@@ -69,18 +69,20 @@ class Ensemble:
         Raises
         ------
         ValueError
-            If 'abs_time' coordinate is missing or there are no valid efficiency values.
+            If 'time' coordinate is missing or there are no valid efficiency values.
         """
         eff = ds["cdr_efficiency"]
 
-        # Check that abs_time exists
-        if "abs_time" in eff.coords:
-            abs_time = eff.coords["abs_time"]
-        elif "abs_time" in ds.data_vars:
-            abs_time = ds["abs_time"]
+        # Check that time coordinate exists
+        if "time" in eff.coords:
+            abs_time = eff.coords["time"]
+            eff = eff.drop_vars("time")
+        elif "time" in ds.data_vars:
+            abs_time = ds["time"]
+            eff = eff.drop_vars("time")
         else:
             raise ValueError(
-                "Dataset must contain an 'abs_time' coordinate or data variable."
+                "Dataset must contain a 'time' coordinate or data variable."
             )
 
         # Drop NaNs to find first valid time
@@ -97,11 +99,6 @@ class Ensemble:
         eff_rel = eff.assign_coords(time=time_rel)
         eff_rel.time.attrs["long_name"] = "time since release start"
 
-        if hasattr(eff_rel, "coords") and "abs_time" in eff_rel.coords:
-            eff_rel = eff_rel.drop_vars("abs_time")
-        elif hasattr(eff_rel, "variables") and "abs_time" in eff_rel.variables:
-            eff_rel = eff_rel.drop_vars("abs_time")
-
         return eff_rel
 
     def _align_times(self, effs: dict[str, xr.DataArray]) -> xr.Dataset:

roms_tools/analysis/roms_output.py

@@ -1,16 +1,11 @@
-import logging
-import re
-import warnings
-from dataclasses import dataclass, field
-from datetime import datetime, timedelta
-from pathlib import Path
+from dataclasses import dataclass
 
 import numpy as np
 import xarray as xr
 from matplotlib.axes import Axes
 
-from roms_tools import Grid
 from roms_tools.analysis.cdr_analysis import compute_cdr_metrics
+from roms_tools.datasets.roms_dataset import ROMSDataset
 from roms_tools.plot import plot, plot_uptake_efficiency
 from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from roms_tools.utils import (
@@ -18,23 +13,19 @@ from roms_tools.utils import (
     infer_nominal_horizontal_resolution,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    load_data,
-)
-from roms_tools.vertical_coordinate import (
-    compute_depth_coordinates,
 )
 
 
 @dataclass(kw_only=True)
-class ROMSOutput:
+class ROMSOutput(ROMSDataset):
     """Represents ROMS model output.
 
     Parameters
     ----------
+    path: str | Path | list[str | Path]
+        Filename, or list of filenames with model output.
     grid : Grid
         Object representing the grid information.
-    path : Union[str, Path, List[Union[str, Path]]]
-        Filename, or list of filenames with model output.
     model_reference_date : datetime, optional
         Reference date of ROMS simulation.
         If not specified, this is inferred from metadata of the model output
@@ -46,32 +37,6 @@ class ROMSOutput:
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
     """
 
-    grid: Grid
-    """Object representing the grid information."""
-    path: str | Path
-    """Filename, or list of filenames with model output."""
-    use_dask: bool = False
-    """Whether to use dask for processing."""
-    model_reference_date: datetime | None = None
-    """Reference date of ROMS simulation."""
-    adjust_depth_for_sea_surface_height: bool | None = False
-    """Whether to account for sea surface height variations when computing depth
-    coordinates."""
-
-    ds: xr.Dataset = field(init=False, repr=False)
-    """An xarray Dataset containing the ROMS output."""
-
-    def __post_init__(self):
-        ds = self._load_model_output()
-        self._infer_model_reference_date_from_metadata(ds)
-        self._check_vertical_coordinate(ds)
-        ds = self._add_absolute_time(ds)
-        ds = self._add_lat_lon_coords(ds)
-        self.ds = ds
-
-        # Dataset for depth coordinates
-        self.ds_depth_coords = xr.Dataset()
-
     def cdr_metrics(self) -> None:
         """
         Compute and plot Carbon Dioxide Removal (CDR) metrics.
@@ -379,272 +344,8 @@ class ROMSOutput:
         if regridded_datasets:
             ds = xr.merge(regridded_datasets)
 
-            with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", category=UserWarning)
-                ds = ds.rename({"abs_time": "time"}).set_index(time="time")
             ds["time"].attrs = {"long_name": "Time"}
             ds["lon"].attrs = {"long_name": "Longitude", "units": "Degrees East"}
             ds["lat"].attrs = {"long_name": "Latitude", "units": "Degrees North"}
 
             return ds
-
-    def _get_depth_coordinates(self, depth_type="layer", locations=["rho"]):
-        """Ensure depth coordinates are stored for a given location and depth type.
-
-        Calculates vertical depth coordinates (layer or interface) for specified locations (e.g., rho, u, v points)
-        and updates them in the dataset (`self.ds`).
-
-        Parameters
-        ----------
-        depth_type : str
-            The type of depth coordinate to compute. Valid options:
-            - "layer": Compute layer depth coordinates.
-            - "interface": Compute interface depth coordinates.
-        locations : list[str], optional
-            Locations for which to compute depth coordinates. Default is ["rho", "u", "v"].
-            Valid options include:
-            - "rho": Depth coordinates at rho points.
-            - "u": Depth coordinates at u points.
-            - "v": Depth coordinates at v points.
-
-        Updates
-        -------
-        self.ds_depth_coords : xarray.Dataset
-
-        Raises
-        ------
-        ValueError
-            If `adjust_depth_for_sea_surface_height` is enabled but `zeta` is missing from `self.ds`.
-
-        Notes
-        -----
-        - This method relies on the `compute_depth_coordinates` function to perform calculations.
-        - If `adjust_depth_for_sea_surface_height` is `True`, the method accounts for variations
-          in sea surface height (`zeta`).
-        """
-        if self.adjust_depth_for_sea_surface_height:
-            if "zeta" not in self.ds:
-                raise ValueError(
-                    "`zeta` is required in provided ROMS output when `adjust_depth_for_sea_surface_height` is enabled."
-                )
-            zeta = self.ds.zeta
-        else:
-            zeta = 0
-
-        for location in locations:
-            var_name = f"{depth_type}_depth_{location}"
-            if var_name not in self.ds_depth_coords:
-                self.ds_depth_coords[var_name] = compute_depth_coordinates(
-                    self.grid.ds, zeta, depth_type, location
-                )
-
-    def _load_model_output(self) -> xr.Dataset:
-        """Load the model output."""
-        # Load the dataset
-        ds = load_data(
-            self.path,
-            dim_names={"time": "time"},
-            use_dask=self.use_dask,
-            decode_times=False,
-            decode_timedelta=False,
-            time_chunking=True,
-            force_combine_nested=True,
-        )
-
-        return ds
-
-    def _infer_model_reference_date_from_metadata(self, ds: xr.Dataset) -> None:
-        """Infer and validate the model reference date from `ocean_time` metadata.
-
-        Parameters
-        ----------
-        ds : xr.Dataset
-            Dataset with an `ocean_time` variable and a `long_name` attribute
-            in the format `Time since YYYY/MM/DD`.
-
-        Raises
-        ------
-        ValueError
-            If `self.model_reference_date` is not set and the reference date cannot
-            be inferred, or if the inferred date does not match `self.model_reference_date`.
-
-        Warns
-        -----
-        UserWarning
-            If `self.model_reference_date` is set but the reference date cannot be inferred.
-        """
-        # Check if 'long_name' exists in the attributes of 'ocean_time'
-        if "long_name" in ds.ocean_time.attrs:
-            input_string = ds.ocean_time.attrs["long_name"]
-            match = re.search(r"(\d{4})/(\d{2})/(\d{2})", input_string)
-
-            if match:
-                # If a match is found, extract year, month, day and create the inferred date
-                year, month, day = map(int, match.groups())
-                inferred_date = datetime(year, month, day)
-
-                if hasattr(self, "model_reference_date") and self.model_reference_date:
-                    # Check if the inferred date matches the provided model reference date
-                    if self.model_reference_date != inferred_date:
-                        raise ValueError(
-                            f"Mismatch between `self.model_reference_date` ({self.model_reference_date}) "
-                            f"and inferred reference date ({inferred_date})."
-                        )
-                else:
-                    # Set the model reference date if not already set
-                    self.model_reference_date = inferred_date
-            else:
-                # Handle case where no match is found
-                if hasattr(self, "model_reference_date") and self.model_reference_date:
-                    logging.warning(
-                        "Could not infer the model reference date from the metadata. "
-                        "`self.model_reference_date` will be used.",
-                    )
-                else:
-                    raise ValueError(
-                        "Model reference date could not be inferred from the metadata, "
-                        "and `self.model_reference_date` is not set."
-                    )
-        else:
-            # Handle case where 'long_name' attribute doesn't exist
-            if hasattr(self, "model_reference_date") and self.model_reference_date:
-                logging.warning(
-                    "`long_name` attribute not found in ocean_time. "
-                    "`self.model_reference_date` will be used instead.",
-                )
-            else:
-                raise ValueError(
-                    "Model reference date could not be inferred from the metadata, "
-                    "and `self.model_reference_date` is not set."
-                )
-
-    def _check_vertical_coordinate(self, ds: xr.Dataset) -> None:
-        """Check that the vertical coordinate parameters in the dataset are consistent
-        with the model grid.
-
-        This method compares the vertical coordinate parameters (`theta_s`, `theta_b`, `hc`, `Cs_r`, `Cs_w`) in
-        the provided dataset (`ds`) with those in the model grid (`self.grid`). The first three parameters are
-        checked for exact equality, while the last two are checked for numerical closeness.
-
-        Parameters
-        ----------
-        ds : xarray.Dataset
-            The dataset containing vertical coordinate parameters in its attributes, such as `theta_s`, `theta_b`,
-            `hc`, `Cs_r`, and `Cs_w`.
-
-        Raises
-        ------
-        ValueError
-            If the vertical coordinate parameters do not match the expected values (based on exact or approximate equality).
-
-        Notes
-        -----
-        - Missing attributes trigger a warning instead of an exception.
-        - `theta_s`, `theta_b`, and `hc` are checked for exact equality using `np.array_equal`.
-        - `Cs_r` and `Cs_w` are checked for numerical closeness using `np.allclose`.
-        """
-        required_exact = ["theta_s", "theta_b", "hc"]
-        required_close = ["Cs_r", "Cs_w"]
-
-        # Check exact equality
-        for param in required_exact:
-            value = ds.attrs.get(param, None)
-            if value is None:
-                logging.warning(
-                    f"Dataset is missing attribute '{param}'. Skipping this check."
-                )
-                continue
-            if not np.array_equal(getattr(self.grid, param), value):
-                raise ValueError(
-                    f"{param} from grid ({getattr(self.grid, param)}) does not match dataset ({value})."
-                )
-
-        # Check numerical closeness
-        for param in required_close:
-            value = ds.attrs.get(param, None)
-            if value is None:
-                logging.warning(
-                    f"Dataset is missing attribute '{param}'. Skipping this check."
-                )
-                continue
-            grid_value = getattr(self.grid.ds, param)
-            if not np.allclose(grid_value, value):
-                raise ValueError(
-                    f"{param} from grid ({grid_value}) is not close to dataset ({value})."
-                )
-
-    def _add_absolute_time(self, ds: xr.Dataset) -> xr.Dataset:
-        """Add absolute time as a coordinate to the dataset.
-
-        Computes "abs_time" based on "ocean_time" and a reference date,
-        and adds it as a coordinate.
-
-        Parameters
-        ----------
-        ds : xarray.Dataset
-            Dataset containing "ocean_time" in seconds since the model reference date.
-
-        Returns
-        -------
-        xarray.Dataset
-            Dataset with "abs_time" added and "time" removed.
-        """
-        if self.model_reference_date is None:
-            raise ValueError(
-                "`model_reference_date` must be set before computing absolute time."
-            )
-
-        ocean_time_seconds = ds["ocean_time"].values
-
-        abs_time = np.array(
-            [
-                self.model_reference_date + timedelta(seconds=seconds)
-                for seconds in ocean_time_seconds
-            ]
-        )
-
-        abs_time = xr.DataArray(
-            abs_time, dims=["time"], coords={"time": ds["ocean_time"]}
-        )
-        abs_time.attrs["long_name"] = "absolute time"
-        ds = ds.assign_coords({"abs_time": abs_time})
-        ds = ds.drop_vars("time")
-
-        return ds
-
-    def _add_lat_lon_coords(self, ds: xr.Dataset) -> xr.Dataset:
-        """Add latitude and longitude coordinates to the dataset based on the grid.
-
-        This method assigns latitude and longitude coordinates from the grid to the dataset.
-        It always adds the "lat_rho" and "lon_rho" coordinates. If the dataset contains the
-        "xi_u" or "eta_v" dimensions, it also adds the corresponding "lat_u", "lon_u",
-        "lat_v", and "lon_v" coordinates.
-
-        Parameters
-        ----------
-        ds : xarray.Dataset
-            Input dataset to which latitude and longitude coordinates will be added.
-
-        Returns
-        -------
-        xarray.Dataset
-            Updated dataset with the appropriate latitude and longitude coordinates
-            assigned to "rho", "u", and "v" points if applicable.
-        """
-        coords_to_add = {
-            "lat_rho": self.grid.ds["lat_rho"],
-            "lon_rho": self.grid.ds["lon_rho"],
-        }
-
-        if "xi_u" in ds.dims:
-            coords_to_add.update(
-                {"lat_u": self.grid.ds["lat_u"], "lon_u": self.grid.ds["lon_u"]}
-            )
-        if "eta_v" in ds.dims:
-            coords_to_add.update(
-                {"lat_v": self.grid.ds["lat_v"], "lon_v": self.grid.ds["lon_v"]}
-            )
-
-        # Add all necessary coordinates in one go
-        ds = ds.assign_coords(coords_to_add)
-        return ds

roms_tools/{download.py → datasets/download.py}

@@ -90,6 +90,7 @@ pup_test_data = pooch.create(
         "GSHHS_l_L1.prj": "98aaf3d1c0ecadf1a424a4536de261c3daf4e373697cb86c40c43b989daf52eb",
         "GSHHS_l_L1.shp": "bc76f101f9b8671f90e734b4026da91c20066fc627cc8b5889ba22d90cbf97e9",
         "GSHHS_l_L1.shx": "72879354892d80d6c39c612f645661ec0edc75f3f9f8f74b19d9387ae0327377",
+        "EMODnet_C2_coarse100.nc": "4202a6a5877de726bf13f41de7a1edfea2db83278ce371fe7732eb4b6770ed6d",
     },
 )