roms-tools 1.0.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

roms_tools/__init__.py

@@ -11,6 +11,5 @@ except ImportError:  # pragma: no cover
 from roms_tools.setup.grid import Grid  # noqa: F401
 from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
-from roms_tools.setup.vertical_coordinate import VerticalCoordinate  # noqa: F401
 from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401

roms_tools/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "1.0.1"
+__version__ = "1.2.0"

roms_tools/setup/boundary_forcing.py

@@ -7,7 +7,6 @@ import importlib.metadata
 from typing import Dict, Union, Optional
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.vertical_coordinate import VerticalCoordinate
 from roms_tools.setup.mixins import ROMSToolsMixins
 from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
@@ -29,8 +28,6 @@ class BoundaryForcing(ROMSToolsMixins):
     ----------
     grid : Grid
         Object representing the grid information.
-    vertical_coordinate: VerticalCoordinate
-        Object representing the vertical coordinate information.
     start_time : datetime
         Start time of the desired boundary forcing data.
     end_time : datetime
@@ -59,7 +56,6 @@ class BoundaryForcing(ROMSToolsMixins):
     --------
     >>> boundary_forcing = BoundaryForcing(
     ...     grid=grid,
-    ...     vertical_coordinate=vertical_coordinate,
     ...     boundaries={"south": True, "east": True, "north": False, "west": True},
     ...     start_time=datetime(2022, 1, 1),
     ...     end_time=datetime(2022, 1, 2),
@@ -73,7 +69,6 @@ class BoundaryForcing(ROMSToolsMixins):
     """
 
     grid: Grid
-    vertical_coordinate: VerticalCoordinate
     start_time: datetime
     end_time: datetime
     boundaries: Dict[str, bool] = field(
@@ -106,7 +101,7 @@ class BoundaryForcing(ROMSToolsMixins):
         vars_2d = ["zeta"]
         vars_3d = ["temp", "salt", "u", "v"]
         data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super().process_velocities(data_vars, angle)
+        data_vars = super().process_velocities(data_vars, angle, "u", "v")
         object.__setattr__(data, "data_vars", data_vars)
 
         if self.bgc_source is not None:
@@ -126,9 +121,9 @@ class BoundaryForcing(ROMSToolsMixins):
             bgc_data = None
 
         d_meta = super().get_variable_metadata()
-        bdry_coords, rename = super().get_boundary_info()
+        bdry_coords = super().get_boundary_info()
 
-        ds = self._write_into_datatree(data, bgc_data, d_meta, bdry_coords, rename)
+        ds = self._write_into_datatree(data, bgc_data, d_meta, bdry_coords)
 
         for direction in ["south", "east", "north", "west"]:
             if self.boundaries[direction]:
@@ -207,7 +202,7 @@ class BoundaryForcing(ROMSToolsMixins):
 
         return data
 
-    def _write_into_dataset(self, data, d_meta, bdry_coords, rename):
+    def _write_into_dataset(self, data, d_meta, bdry_coords):
 
         # save in new dataset
         ds = xr.Dataset()
@@ -220,21 +215,18 @@ class BoundaryForcing(ROMSToolsMixins):
                         ds[f"{var}_{direction}"] = (
                             data.data_vars[var]
                             .isel(**bdry_coords["u"][direction])
-                            .rename(**rename["u"][direction])
                             .astype(np.float32)
                         )
                     elif var in ["v", "vbar"]:
                         ds[f"{var}_{direction}"] = (
                             data.data_vars[var]
                             .isel(**bdry_coords["v"][direction])
-                            .rename(**rename["v"][direction])
                             .astype(np.float32)
                         )
                     else:
                         ds[f"{var}_{direction}"] = (
                             data.data_vars[var]
                             .isel(**bdry_coords["rho"][direction])
-                            .rename(**rename["rho"][direction])
                             .astype(np.float32)
                         )
                     ds[f"{var}_{direction}"].attrs[
@@ -293,127 +285,63 @@ class BoundaryForcing(ROMSToolsMixins):
 
         return ds
 
-    def _write_into_datatree(self, data, bgc_data, d_meta, bdry_coords, rename):
+    def _write_into_datatree(self, data, bgc_data, d_meta, bdry_coords):
 
         ds = self._add_global_metadata()
-        ds["sc_r"] = self.vertical_coordinate.ds["sc_r"]
-        ds["Cs_r"] = self.vertical_coordinate.ds["Cs_r"]
+        ds["sc_r"] = self.grid.ds["sc_r"]
+        ds["Cs_r"] = self.grid.ds["Cs_r"]
 
         ds = DataTree(name="root", data=ds)
 
-        ds_physics = self._write_into_dataset(data, d_meta, bdry_coords, rename)
-        ds_physics = self._add_coordinates(bdry_coords, rename, ds_physics)
+        ds_physics = self._write_into_dataset(data, d_meta, bdry_coords)
         ds_physics = self._add_global_metadata(ds_physics)
         ds_physics.attrs["physics_source"] = self.physics_source["name"]
 
         ds_physics = DataTree(name="physics", parent=ds, data=ds_physics)
 
         if bgc_data:
-            ds_bgc = self._write_into_dataset(bgc_data, d_meta, bdry_coords, rename)
-            ds_bgc = self._add_coordinates(bdry_coords, rename, ds_bgc)
+            ds_bgc = self._write_into_dataset(bgc_data, d_meta, bdry_coords)
             ds_bgc = self._add_global_metadata(ds_bgc)
             ds_bgc.attrs["bgc_source"] = self.bgc_source["name"]
             ds_bgc = DataTree(name="bgc", parent=ds, data=ds_bgc)
 
         return ds
 
-    def _add_coordinates(self, bdry_coords, rename, ds=None):
-
-        if ds is None:
-            ds = xr.Dataset()
-
-        for direction in ["south", "east", "north", "west"]:
-
-            if self.boundaries[direction]:
+    def _get_coordinates(self, direction, point):
+        """
+        Retrieve layer and interface depth coordinates for a specified grid boundary.
 
-                lat_rho = self.grid.ds.lat_rho.isel(
-                    **bdry_coords["rho"][direction]
-                ).rename(**rename["rho"][direction])
-                lon_rho = self.grid.ds.lon_rho.isel(
-                    **bdry_coords["rho"][direction]
-                ).rename(**rename["rho"][direction])
-                layer_depth_rho = (
-                    self.vertical_coordinate.ds["layer_depth_rho"]
-                    .isel(**bdry_coords["rho"][direction])
-                    .rename(**rename["rho"][direction])
-                )
-                interface_depth_rho = (
-                    self.vertical_coordinate.ds["interface_depth_rho"]
-                    .isel(**bdry_coords["rho"][direction])
-                    .rename(**rename["rho"][direction])
-                )
+        This method extracts the layer depth and interface depth coordinates along
+        a specified boundary (north, south, east, or west) and for a specified point
+        type (rho, u, or v) from the grid dataset.
 
-                lat_u = self.grid.ds.lat_u.isel(**bdry_coords["u"][direction]).rename(
-                    **rename["u"][direction]
-                )
-                lon_u = self.grid.ds.lon_u.isel(**bdry_coords["u"][direction]).rename(
-                    **rename["u"][direction]
-                )
-                layer_depth_u = (
-                    self.vertical_coordinate.ds["layer_depth_u"]
-                    .isel(**bdry_coords["u"][direction])
-                    .rename(**rename["u"][direction])
-                )
-                interface_depth_u = (
-                    self.vertical_coordinate.ds["interface_depth_u"]
-                    .isel(**bdry_coords["u"][direction])
-                    .rename(**rename["u"][direction])
-                )
+        Parameters
+        ----------
+        direction : str
+            The direction of the boundary to retrieve coordinates for. Valid options
+            are "north", "south", "east", and "west".
+        point : str
+            The type of grid point to retrieve coordinates for. Valid options are
+            "rho" for the grid's central points, "u" for the u-flux points, and "v"
+            for the v-flux points.
 
-                lat_v = self.grid.ds.lat_v.isel(**bdry_coords["v"][direction]).rename(
-                    **rename["v"][direction]
-                )
-                lon_v = self.grid.ds.lon_v.isel(**bdry_coords["v"][direction]).rename(
-                    **rename["v"][direction]
-                )
-                layer_depth_v = (
-                    self.vertical_coordinate.ds["layer_depth_v"]
-                    .isel(**bdry_coords["v"][direction])
-                    .rename(**rename["v"][direction])
-                )
-                interface_depth_v = (
-                    self.vertical_coordinate.ds["interface_depth_v"]
-                    .isel(**bdry_coords["v"][direction])
-                    .rename(**rename["v"][direction])
-                )
+        Returns
+        -------
+        xarray.DataArray, xarray.DataArray
+            The layer depth and interface depth coordinates for the specified grid
+            boundary and point type.
+        """
 
-                ds = ds.assign_coords(
-                    {
-                        f"layer_depth_rho_{direction}": layer_depth_rho,
-                        f"layer_depth_u_{direction}": layer_depth_u,
-                        f"layer_depth_v_{direction}": layer_depth_v,
-                        f"interface_depth_rho_{direction}": interface_depth_rho,
-                        f"interface_depth_u_{direction}": interface_depth_u,
-                        f"interface_depth_v_{direction}": interface_depth_v,
-                        f"lat_rho_{direction}": lat_rho,
-                        f"lat_u_{direction}": lat_u,
-                        f"lat_v_{direction}": lat_v,
-                        f"lon_rho_{direction}": lon_rho,
-                        f"lon_u_{direction}": lon_u,
-                        f"lon_v_{direction}": lon_v,
-                    }
-                )
+        bdry_coords = super().get_boundary_info()
 
-        # Gracefully handle dropping variables that might not be present
-        variables_to_drop = [
-            "s_rho",
-            "layer_depth_rho",
-            "layer_depth_u",
-            "layer_depth_v",
-            "interface_depth_rho",
-            "interface_depth_u",
-            "interface_depth_v",
-            "lat_rho",
-            "lon_rho",
-            "lat_u",
-            "lon_u",
-            "lat_v",
-            "lon_v",
-        ]
-        existing_vars = [var for var in variables_to_drop if var in ds]
-        ds = ds.drop_vars(existing_vars)
+        layer_depth = self.grid.ds[f"layer_depth_{point}"].isel(
+            **bdry_coords[point][direction]
+        )
+        interface_depth = self.grid.ds[f"interface_depth_{point}"].isel(
+            **bdry_coords[point][direction]
+        )
 
-        return ds
+        return layer_depth, interface_depth
 
     def _add_global_metadata(self, ds=None):
 
@@ -430,10 +358,9 @@ class BoundaryForcing(ROMSToolsMixins):
         ds.attrs["end_time"] = str(self.end_time)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
 
-        ds.attrs["theta_s"] = self.vertical_coordinate.ds["theta_s"].item()
-        ds.attrs["theta_b"] = self.vertical_coordinate.ds["theta_b"].item()
-        ds.attrs["Tcline"] = self.vertical_coordinate.ds["Tcline"].item()
-        ds.attrs["hc"] = self.vertical_coordinate.ds["hc"].item()
+        ds.attrs["theta_s"] = self.grid.ds.attrs["theta_s"]
+        ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
+        ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
         return ds
 
@@ -492,8 +419,9 @@ class BoundaryForcing(ROMSToolsMixins):
         time : int, optional
             The time index to plot. Default is 0.
         layer_contours : bool, optional
-            Whether to include layer contours in the plot. This can help visualize the depth levels
-            of the field. Default is False.
+            If True, contour lines representing the boundaries between vertical layers will
+            be added to the plot. For clarity, the number of layer
+            contours displayed is limited to a maximum of 10. Default is False.
 
         Returns
         -------
@@ -519,6 +447,19 @@ class BoundaryForcing(ROMSToolsMixins):
         field = ds[varname].isel(bry_time=time).load()
         title = field.long_name
 
+        if "s_rho" in field.dims:
+            if varname.startswith(("u_", "ubar_")):
+                point = "u"
+            elif varname.startswith(("v_", "vbar_")):
+                point = "v"
+            else:
+                point = "rho"
+            direction = varname.split("_")[-1]
+
+            layer_depth, interface_depth = self._get_coordinates(direction, point)
+
+            field = field.assign_coords({"layer_depth": layer_depth})
+
         # chose colorbar
         if varname.startswith(("u", "v", "ubar", "vbar", "zeta")):
             vmax = max(field.max().values, -field.min().values)
@@ -536,27 +477,6 @@ class BoundaryForcing(ROMSToolsMixins):
 
         if len(field.dims) == 2:
             if layer_contours:
-                depths_to_check = [
-                    "interface_depth_rho",
-                    "interface_depth_u",
-                    "interface_depth_v",
-                ]
-                try:
-                    interface_depth = next(
-                        ds[depth_label]
-                        for depth_label in ds.coords
-                        if any(
-                            depth_label.startswith(prefix) for prefix in depths_to_check
-                        )
-                        and (
-                            set(ds[depth_label].dims) - {"s_w"}
-                            == set(field.dims) - {"s_rho"}
-                        )
-                    )
-                except StopIteration:
-                    raise ValueError(
-                        f"None of the expected depths ({', '.join(depths_to_check)}) have dimensions matching field.dims"
-                    )
                 # restrict number of layer_contours to 10 for the sake of plot clearity
                 nr_layers = len(interface_depth["s_w"])
                 selected_layers = np.linspace(
@@ -664,11 +584,6 @@ class BoundaryForcing(ROMSToolsMixins):
         grid_data.pop("ds", None)  # Exclude non-serializable fields
         grid_data.pop("straddle", None)
 
-        # Serialize VerticalCoordinate data
-        vertical_coordinate_data = asdict(self.vertical_coordinate)
-        vertical_coordinate_data.pop("ds", None)  # Exclude non-serializable fields
-        vertical_coordinate_data.pop("grid", None)  # Exclude non-serializable fields
-
         # Include the version of roms-tools
         try:
             roms_tools_version = importlib.metadata.version("roms-tools")
@@ -679,7 +594,6 @@ class BoundaryForcing(ROMSToolsMixins):
         header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
 
         grid_yaml_data = {"Grid": grid_data}
-        vertical_coordinate_yaml_data = {"VerticalCoordinate": vertical_coordinate_data}
 
         boundary_forcing_data = {
             "BoundaryForcing": {
@@ -694,7 +608,6 @@ class BoundaryForcing(ROMSToolsMixins):
 
         yaml_data = {
             **grid_yaml_data,
-            **vertical_coordinate_yaml_data,
             **boundary_forcing_data,
         }
 
@@ -745,13 +658,10 @@ class BoundaryForcing(ROMSToolsMixins):
                 boundary_forcing_data[date_string]
             )
 
-        # Create VerticalCoordinate instance from the YAML file
-        vertical_coordinate = VerticalCoordinate.from_yaml(filepath)
-        grid = vertical_coordinate.grid
+        grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of InitialConditions
         return cls(
             grid=grid,
-            vertical_coordinate=vertical_coordinate,
             **boundary_forcing_data,
         )

roms_tools/setup/datasets.py

@@ -584,9 +584,14 @@ class TPXODataset(Dataset):
                 "ntides": self.dim_names["ntides"],
             },
         )
+        self.check_dataset(ds)
+
         # Select relevant fields
         ds = super().select_relevant_fields(ds)
 
+        # Make sure that latitude is ascending
+        ds = super().ensure_latitude_ascending(ds)
+
         # Check whether the data covers the entire globe
         object.__setattr__(self, "is_global", super().check_if_global(ds))
 
@@ -769,6 +774,8 @@ class CESMDataset(Dataset):
             ds = assign_dates_to_climatology(ds, time_dim)
             # rename dimension
             ds = ds.swap_dims({time_dim: "time"})
+            if time_dim in ds.variables:
+                ds = ds.drop_vars(time_dim)
             # Update dimension names
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["time"] = "time"
@@ -872,9 +879,9 @@ class CESMBGCDataset(CESMDataset):
             ds["depth"].attrs["long_name"] = "Depth"
             ds["depth"].attrs["units"] = "m"
             ds = ds.swap_dims({"z_t": "depth"})
-            if "z_t" in ds:
+            if "z_t" in ds.variables:
                 ds = ds.drop_vars("z_t")
-            if "z_t_150m" in ds:
+            if "z_t_150m" in ds.variables:
                 ds = ds.drop_vars("z_t_150m")
             # update dataset
             object.__setattr__(self, "ds", ds)
@@ -932,6 +939,19 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 
     climatology: Optional[bool] = False
 
+    def post_process(self):
+        """
+        Perform post-processing on the dataset to remove specific variables.
+
+        This method checks if the variable "z_t" exists in the dataset. If it does,
+        the variable is removed from the dataset. The modified dataset is then
+        reassigned to the `ds` attribute of the object.
+        """
+
+        if "z_t" in self.ds.variables:
+            ds = self.ds.drop_vars("z_t")
+            object.__setattr__(self, "ds", ds)
+
 
 @dataclass(frozen=True, kw_only=True)
 class ERA5Dataset(Dataset):