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

roms_tools/setup/mixins.py

@@ -123,9 +123,9 @@ class ROMSToolsMixins:
         if vars_3d:
             # 3d interpolation
             coords = {
+                data.dim_names["depth"]: self.grid.ds["layer_depth_rho"],
                 data.dim_names["latitude"]: lat,
                 data.dim_names["longitude"]: lon,
-                data.dim_names["depth"]: self.grid.ds["layer_depth_rho"],
             }
         # extrapolate deepest value all the way to bottom ("flooding")
         for var in vars_3d:
@@ -149,36 +149,50 @@ class ROMSToolsMixins:
             if data.dim_names["time"] != "time":
                 data_vars[var] = data_vars[var].rename({data.dim_names["time"]: "time"})
 
+            # transpose to correct order (time, s_rho, eta_rho, xi_rho)
+            data_vars[var] = data_vars[var].transpose(
+                "time", "s_rho", "eta_rho", "xi_rho"
+            )
+
         return data_vars
 
-    def process_velocities(self, data_vars, angle, interpolate=True):
+    def process_velocities(self, data_vars, angle, uname, vname, interpolate=True):
         """
-        Processes and rotates velocity components, and interpolates them to the appropriate grid points.
+        Process and rotate velocity components to align with the grid orientation and optionally interpolate
+        them to the appropriate grid points.
 
         This method performs the following steps:
-        1. Rotates the velocity components to align with the grid orientation using the provided angle.
-        2. Optionally interpolates the rotated velocities to the u- and v-points of the grid.
-        3. If the velocities are 3D (with vertical coordinates), computes barotropic (depth-averaged) velocities.
+
+        1. **Rotation**: Rotates the velocity components (e.g., `u`, `v`) to align with the grid orientation
+           using the provided angle data.
+        2. **Interpolation**: Optionally interpolates the rotated velocities from rho-points to u- and v-points
+           of the grid.
+        3. **Barotropic Velocity Calculation**: If the velocity components are 3D (with vertical coordinates),
+           computes the barotropic (depth-averaged) velocities.
 
         Parameters
         ----------
         data_vars : dict of str: xarray.DataArray
-            Dictionary containing the velocity components to be processed. Must include keys "u" and "v"
-            or "uwnd" and "vwnd".
+            Dictionary containing the velocity components to be processed. The dictionary should include keys
+            corresponding to the velocity component names (e.g., `uname`, `vname`).
         angle : xarray.DataArray
-            DataArray containing the angle used for rotating the velocity components to the grid orientation.
+            DataArray containing the grid angle values used to rotate the velocity components to the correct
+            orientation on the grid.
+        uname : str
+            The key corresponding to the zonal (east-west) velocity component in `data_vars`.
+        vname : str
+            The key corresponding to the meridional (north-south) velocity component in `data_vars`.
         interpolate : bool, optional
-            If True, interpolates the velocities to the u- and v-points. Defaults to True.
+            If True, interpolates the rotated velocity components to the u- and v-points of the grid.
+            Defaults to True.
 
         Returns
         -------
         dict of str: xarray.DataArray
-            Dictionary of processed velocity components. Includes "ubar" and "vbar" if the velocity components
-            have vertical coordinates and are processed for barotropic (depth-averaged) velocities.
+            A dictionary of the processed velocity components. The returned dictionary includes the rotated and,
+            if applicable, interpolated velocity components. If the input velocities are 3D (having a vertical
+            dimension), the dictionary also includes the barotropic (depth-averaged) velocities (`ubar` and `vbar`).
         """
-        # Determine the correct variable names based on the keys in data_vars
-        uname = "u" if "u" in data_vars else "uwnd"
-        vname = "v" if "v" in data_vars else "vwnd"
 
         # Rotate velocities to grid orientation
         u_rot = data_vars[uname] * np.cos(angle) + data_vars[vname] * np.sin(angle)
@@ -230,6 +244,26 @@ class ROMSToolsMixins:
         """
 
         d = {
+            "ssh_Re": {"long_name": "Tidal elevation, real part", "units": "m"},
+            "ssh_Im": {"long_name": "Tidal elevation, complex part", "units": "m"},
+            "pot_Re": {"long_name": "Tidal potential, real part", "units": "m"},
+            "pot_Im": {"long_name": "Tidal potential, complex part", "units": "m"},
+            "u_Re": {
+                "long_name": "Tidal velocity in x-direction, real part",
+                "units": "m/s",
+            },
+            "u_Im": {
+                "long_name": "Tidal velocity in x-direction, complex part",
+                "units": "m/s",
+            },
+            "v_Re": {
+                "long_name": "Tidal velocity in y-direction, real part",
+                "units": "m/s",
+            },
+            "v_Im": {
+                "long_name": "Tidal velocity in y-direction, complex part",
+                "units": "m/s",
+            },
             "uwnd": {"long_name": "10 meter wind in x-direction", "units": "m/s"},
             "vwnd": {"long_name": "10 meter wind in y-direction", "units": "m/s"},
             "swrad": {
@@ -323,18 +357,18 @@ class ROMSToolsMixins:
         return d
 
     def get_boundary_info(self):
-        """
-        Provides boundary coordinate information and renaming conventions for grid boundaries.
 
-        This method returns two dictionaries: one specifying the boundary coordinates for different types of
-        grid variables (e.g., "rho", "u", "v"), and another specifying how to rename dimensions for these boundaries.
+        """
+        This method provides information about the boundary points for the rho, u, and v
+        variables on the grid, specifying the indices for the south, east, north, and west
+        boundaries.
 
         Returns
         -------
-        tuple of (dict, dict)
-            - A dictionary mapping variable types and directions to boundary coordinates.
-            - A dictionary mapping variable types and directions to new dimension names.
-
+        dict
+            A dictionary where keys are variable types ("rho", "u", "v"), and values
+            are nested dictionaries mapping directions ("south", "east", "north", "west")
+            to the corresponding boundary coordinates.
         """
 
         # Boundary coordinates
@@ -359,27 +393,4 @@ class ROMSToolsMixins:
             },
         }
 
-        # How to rename the dimensions
-
-        rename = {
-            "rho": {
-                "south": {"xi_rho": "xi_rho_south"},
-                "east": {"eta_rho": "eta_rho_east"},
-                "north": {"xi_rho": "xi_rho_north"},
-                "west": {"eta_rho": "eta_rho_west"},
-            },
-            "u": {
-                "south": {"xi_u": "xi_u_south"},
-                "east": {"eta_rho": "eta_u_east"},
-                "north": {"xi_u": "xi_u_north"},
-                "west": {"eta_rho": "eta_u_west"},
-            },
-            "v": {
-                "south": {"xi_rho": "xi_v_south"},
-                "east": {"eta_v": "eta_v_east"},
-                "north": {"xi_rho": "xi_v_north"},
-                "west": {"eta_v": "eta_v_west"},
-            },
-        }
-
-        return bdry_coords, rename
+        return bdry_coords

roms_tools/setup/plot.py

@@ -8,7 +8,6 @@ def _plot(
     field=None,
     depth_contours=False,
     straddle=False,
-    coarse_grid=False,
     c="red",
     title="",
     kwargs={},
@@ -21,29 +20,8 @@ def _plot(
     else:
 
         field = field.squeeze()
-
-        if coarse_grid:
-
-            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
-            field = field.where(grid_ds.mask_coarse)
-            lon_deg = field.lon
-            lat_deg = field.lat
-
-        else:
-            if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-                field = field.where(grid_ds.mask_rho)
-                lon_deg = grid_ds["lon_rho"]
-                lat_deg = grid_ds["lat_rho"]
-            elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-                field = field.where(grid_ds.mask_u)
-                lon_deg = grid_ds["lon_u"]
-                lat_deg = grid_ds["lat_u"]
-            elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-                field = field.where(grid_ds.mask_v)
-                lon_deg = grid_ds["lon_v"]
-                lat_deg = grid_ds["lat_v"]
-            else:
-                ValueError("provided field does not have two horizontal dimension")
+        lon_deg = field.lon
+        lat_deg = field.lat
 
         # check if North or South pole are in domain
         if lat_deg.max().values > 89 or lat_deg.min().values < -89:
@@ -96,23 +74,7 @@ def _plot(
         plt.colorbar(p, label=f"{field.long_name} [{field.units}]")
 
     if depth_contours:
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            if "layer_depth_rho" in field.coords:
-                depth = field.layer_depth_rho
-            else:
-                depth = field.interface_depth_rho
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            if "layer_depth_u" in field.coords:
-                depth = field.layer_depth_u
-            else:
-                depth = field.interface_depth_u
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            if "layer_depth_v" in field.coords:
-                depth = field.layer_depth_v
-            else:
-                depth = field.interface_depth_v
-
-        cs = ax.contour(lon_deg, lat_deg, depth, transform=proj, colors="k")
+        cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
     return fig
@@ -135,12 +97,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}):
         )
 
     depths_to_check = [
-        "layer_depth_rho",
-        "layer_depth_u",
-        "layer_depth_v",
-        "interface_depth_rho",
-        "interface_depth_u",
-        "interface_depth_v",
+        "layer_depth",
+        "interface_depth",
     ]
     try:
         depth_label = next(

roms_tools/setup/surface_forcing.py

@@ -95,7 +95,9 @@ class SurfaceForcing(ROMSToolsMixins):
         vars_2d = ["uwnd", "vwnd", "swrad", "lwrad", "Tair", "qair", "rain"]
         vars_3d = []
         data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super().process_velocities(data_vars, angle, interpolate=False)
+        data_vars = super().process_velocities(
+            data_vars, angle, "uwnd", "vwnd", interpolate=False
+        )
 
         if self.correct_radiation:
             correction_data = self._get_correction_data()
@@ -235,18 +237,19 @@ class SurfaceForcing(ROMSToolsMixins):
 
         if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-            bgc_data = CESMBGCSurfaceForcingDataset(
+            data = CESMBGCSurfaceForcingDataset(
                 filename=self.bgc_source["path"],
                 start_time=self.start_time,
                 end_time=self.end_time,
                 climatology=self.bgc_source["climatology"],
             )
+            data.post_process()
         else:
             raise ValueError(
                 'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
             )
 
-        return bgc_data
+        return data
 
     def _write_into_dataset(self, data, d_meta):
 
@@ -259,7 +262,6 @@ class SurfaceForcing(ROMSToolsMixins):
             ds[var].attrs["units"] = d_meta[var]["units"]
 
         if self.use_coarse_grid:
-            ds = ds.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
             ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
 
         # Preserve absolute time coordinate for readability
@@ -295,6 +297,10 @@ class SurfaceForcing(ROMSToolsMixins):
         if data.climatology:
             ds["time"].attrs["cycle_length"] = 365.25
 
+        variables_to_drop = ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]
+        existing_vars = [var for var in variables_to_drop if var in ds]
+        ds = ds.drop_vars(existing_vars)
+
         return ds
 
     def _write_into_datatree(self, data, bgc_data, d_meta):
@@ -392,6 +398,15 @@ class SurfaceForcing(ROMSToolsMixins):
         field = ds[varname].isel(time=time).load()
         title = field.long_name
 
+        # assign lat / lon
+        if self.use_coarse_grid:
+            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
+            field = field.where(self.grid.ds.mask_coarse)
+        else:
+            field = field.where(self.grid.ds.mask_rho)
+
+        field = field.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
+
         # choose colorbar
         if varname in ["uwnd", "vwnd"]:
             vmax = max(field.max().values, -field.min().values)
@@ -412,7 +427,6 @@ class SurfaceForcing(ROMSToolsMixins):
             self.grid.ds,
             field=field,
             straddle=self.grid.straddle,
-            coarse_grid=self.use_coarse_grid,
             title=title,
             kwargs=kwargs,
             c="g",

roms_tools/setup/tides.py

@@ -15,11 +15,12 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
 )
+from roms_tools.setup.mixins import ROMSToolsMixins
 import matplotlib.pyplot as plt
 
 
 @dataclass(frozen=True, kw_only=True)
-class TidalForcing:
+class TidalForcing(ROMSToolsMixins):
     """
     Represents tidal forcing data used in ocean modeling.
 
@@ -59,30 +60,13 @@ class TidalForcing:
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
-        if "name" not in self.source.keys():
-            raise ValueError("`source` must include a 'name'.")
-        if "path" not in self.source.keys():
-            raise ValueError("`source` must include a 'path'.")
-        if self.source["name"] == "TPXO":
-            data = TPXODataset(filename=self.source["path"])
-        else:
-            raise ValueError('Only "TPXO" is a valid option for source["name"].')
+
+        self._input_checks()
+        lon, lat, angle, straddle = super().get_target_lon_lat()
+
+        data = self._get_data()
 
         data.check_number_constituents(self.ntides)
-        # operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in longitude away from Greenwich meridian
-        lon = self.grid.ds.lon_rho
-        lat = self.grid.ds.lat_rho
-        angle = self.grid.ds.angle
-
-        lon = xr.where(lon > 180, lon - 360, lon)
-        straddle = True
-        if not self.grid.straddle and abs(lon).min() > 5:
-            lon = xr.where(lon < 0, lon + 360, lon)
-            straddle = False
-
-        # Restrict data to relevant subdomain to achieve better performance and to avoid discontinuous longitudes introduced by converting
-        # to a different longitude range (+- 360 degrees). Discontinues longitudes can lead to artifacts in the interpolation process that
-        # would not be detected by the nan_check function.
         data.choose_subdomain(
             latitude_range=[lat.min().values, lat.max().values],
             longitude_range=[lon.min().values, lon.max().values],
@@ -121,56 +105,64 @@ class TidalForcing:
                 method="linear",
             )
 
-        # Rotate to grid orientation
-        u_Re = data_vars["u_Re"] * np.cos(angle) + data_vars["v_Re"] * np.sin(angle)
-        v_Re = data_vars["v_Re"] * np.cos(angle) - data_vars["u_Re"] * np.sin(angle)
-        u_Im = data_vars["u_Im"] * np.cos(angle) + data_vars["v_Im"] * np.sin(angle)
-        v_Im = data_vars["v_Im"] * np.cos(angle) - data_vars["u_Im"] * np.sin(angle)
+        data_vars = super().process_velocities(
+            data_vars, angle, "u_Re", "v_Re", interpolate=False
+        )
+        data_vars = super().process_velocities(
+            data_vars, angle, "u_Im", "v_Im", interpolate=False
+        )
 
         # Convert to barotropic velocity
-        u_Re = u_Re / self.grid.ds.h
-        v_Re = v_Re / self.grid.ds.h
-        u_Im = u_Im / self.grid.ds.h
-        v_Im = v_Im / self.grid.ds.h
+        for varname in ["u_Re", "v_Re", "u_Im", "v_Im"]:
+            data_vars[varname] = data_vars[varname] / self.grid.ds.h
 
         # Interpolate from rho- to velocity points
-        u_Re = interpolate_from_rho_to_u(u_Re)
-        v_Re = interpolate_from_rho_to_v(v_Re)
-        u_Im = interpolate_from_rho_to_u(u_Im)
-        v_Im = interpolate_from_rho_to_v(v_Im)
+        for uname in ["u_Re", "u_Im"]:
+            data_vars[uname] = interpolate_from_rho_to_u(data_vars[uname])
+        for vname in ["v_Re", "v_Im"]:
+            data_vars[vname] = interpolate_from_rho_to_v(data_vars[vname])
+
+        d_meta = super().get_variable_metadata()
+        ds = self._write_into_dataset(data_vars, d_meta)
+        ds["omega"] = tides["omega"]
+
+        ds = self._add_global_metadata(ds)
+
+        for var in ["ssh_Re", "u_Re", "v_Im"]:
+            nan_check(ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+
+        object.__setattr__(self, "ds", ds)
+
+    def _input_checks(self):
+
+        if "name" not in self.source.keys():
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source.keys():
+            raise ValueError("`source` must include a 'path'.")
+
+    def _get_data(self):
+
+        if self.source["name"] == "TPXO":
+            data = TPXODataset(filename=self.source["path"])
+        else:
+            raise ValueError('Only "TPXO" is a valid option for source["name"].')
+        return data
+
+    def _write_into_dataset(self, data_vars, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        # ds["omega"] = tides["omega"]
-
-        ds["ssh_Re"] = data_vars["ssh_Re"].astype(np.float32)
-        ds["ssh_Im"] = data_vars["ssh_Im"].astype(np.float32)
-        ds["ssh_Re"].attrs["long_name"] = "Tidal elevation, real part"
-        ds["ssh_Im"].attrs["long_name"] = "Tidal elevation, complex part"
-        ds["ssh_Re"].attrs["units"] = "m"
-        ds["ssh_Im"].attrs["units"] = "m"
-
-        ds["pot_Re"] = data_vars["pot_Re"].astype(np.float32)
-        ds["pot_Im"] = data_vars["pot_Im"].astype(np.float32)
-        ds["pot_Re"].attrs["long_name"] = "Tidal potential, real part"
-        ds["pot_Im"].attrs["long_name"] = "Tidal potential, complex part"
-        ds["pot_Re"].attrs["units"] = "m"
-        ds["pot_Im"].attrs["units"] = "m"
-
-        ds["u_Re"] = u_Re.astype(np.float32)
-        ds["u_Im"] = u_Im.astype(np.float32)
-        ds["u_Re"].attrs["long_name"] = "Tidal velocity in x-direction, real part"
-        ds["u_Im"].attrs["long_name"] = "Tidal velocity in x-direction, complex part"
-        ds["u_Re"].attrs["units"] = "m/s"
-        ds["u_Im"].attrs["units"] = "m/s"
-
-        ds["v_Re"] = v_Re.astype(np.float32)
-        ds["v_Im"] = v_Im.astype(np.float32)
-        ds["v_Re"].attrs["long_name"] = "Tidal velocity in y-direction, real part"
-        ds["v_Im"].attrs["long_name"] = "Tidal velocity in y-direction, complex part"
-        ds["v_Re"].attrs["units"] = "m/s"
-        ds["v_Im"].attrs["units"] = "m/s"
+        for var in data_vars.keys():
+            ds[var] = data_vars[var].astype(np.float32)
+            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
+            ds[var].attrs["units"] = d_meta[var]["units"]
+
+        ds = ds.drop_vars(["lat_rho", "lon_rho"])
+
+        return ds
+
+    def _add_global_metadata(self, ds):
 
         ds.attrs["title"] = "ROMS tidal forcing created by ROMS-Tools"
         # Include the version of roms-tools
@@ -185,10 +177,7 @@ class TidalForcing:
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
         ds.attrs["allan_factor"] = self.allan_factor
 
-        object.__setattr__(self, "ds", ds)
-
-        for var in ["ssh_Re", "u_Re", "v_Im"]:
-            nan_check(self.ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+        return ds
 
     def plot(self, varname, ntides=0) -> None:
         """
@@ -228,6 +217,25 @@ class TidalForcing:
         """
 
         field = self.ds[varname].isel(ntides=ntides).compute()
+        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
+            field = field.where(self.grid.ds.mask_rho)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
+            field = field.where(self.grid.ds.mask_u)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
+            field = field.where(self.grid.ds.mask_v)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
+            )
+        else:
+            ValueError("provided field does not have two horizontal dimension")
 
         title = "%s, ntides = %i" % (field.long_name, self.ds[varname].ntides[ntides])
 

roms_tools/setup/vertical_coordinate.py

@@ -114,4 +114,9 @@ def compute_depth(zeta, h, hc, cs, sigma):
     s = (hc * sigma + h * cs) / (hc + h)
     z = zeta + (zeta + h) * s
 
+    if "s_rho" in z.dims:
+        z = z.transpose("s_rho", "eta_rho", "xi_rho")
+    elif "s_w" in z.dims:
+        z = z.transpose("s_w", "eta_rho", "xi_rho")
+
     return z