roms-tools 1.1.0__py3-none-any.whl → 1.3.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)
@@ -193,13 +207,6 @@ class ROMSToolsMixins:
             data_vars[vname] = v_rot
 
         if "s_rho" in data_vars[uname].dims and "s_rho" in data_vars[vname].dims:
-            # 3D masks for ROMS domain
-            umask = self.grid.ds.mask_u.expand_dims({"s_rho": data_vars[uname].s_rho})
-            vmask = self.grid.ds.mask_v.expand_dims({"s_rho": data_vars[vname].s_rho})
-
-            data_vars[uname] = data_vars[uname] * umask
-            data_vars[vname] = data_vars[vname] * vmask
-
             # Compute barotropic velocity
             dz = -self.grid.ds["interface_depth_rho"].diff(dim="s_w")
             dz = dz.rename({"s_w": "s_rho"})
@@ -214,172 +221,3 @@ class ROMSToolsMixins:
             ).transpose("time", "eta_v", "xi_rho")
 
         return data_vars
-
-    def get_variable_metadata(self):
-        """
-        Retrieves metadata for commonly used variables in the dataset.
-
-        This method returns a dictionary containing the metadata for various variables, including long names
-        and units for each variable.
-
-        Returns
-        -------
-        dict of str: dict
-            Dictionary where keys are variable names and values are dictionaries with "long_name" and "units" keys.
-
-        """
-
-        d = {
-            "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": {
-                "long_name": "downward short-wave (solar) radiation",
-                "units": "W/m^2",
-            },
-            "lwrad": {
-                "long_name": "downward long-wave (thermal) radiation",
-                "units": "W/m^2",
-            },
-            "Tair": {"long_name": "air temperature at 2m", "units": "degrees Celsius"},
-            "qair": {"long_name": "absolute humidity at 2m", "units": "kg/kg"},
-            "rain": {"long_name": "total precipitation", "units": "cm/day"},
-            "temp": {"long_name": "potential temperature", "units": "degrees Celsius"},
-            "salt": {"long_name": "salinity", "units": "PSU"},
-            "zeta": {"long_name": "sea surface height", "units": "m"},
-            "u": {"long_name": "u-flux component", "units": "m/s"},
-            "v": {"long_name": "v-flux component", "units": "m/s"},
-            "w": {"long_name": "w-flux component", "units": "m/s"},
-            "ubar": {
-                "long_name": "vertically integrated u-flux component",
-                "units": "m/s",
-            },
-            "vbar": {
-                "long_name": "vertically integrated v-flux component",
-                "units": "m/s",
-            },
-            "PO4": {"long_name": "dissolved inorganic phosphate", "units": "mmol/m^3"},
-            "NO3": {"long_name": "dissolved inorganic nitrate", "units": "mmol/m^3"},
-            "SiO3": {"long_name": "dissolved inorganic silicate", "units": "mmol/m^3"},
-            "NH4": {"long_name": "dissolved ammonia", "units": "mmol/m^3"},
-            "Fe": {"long_name": "dissolved inorganic iron", "units": "mmol/m^3"},
-            "Lig": {"long_name": "iron binding ligand", "units": "mmol/m^3"},
-            "O2": {"long_name": "dissolved oxygen", "units": "mmol/m^3"},
-            "DIC": {"long_name": "dissolved inorganic carbon", "units": "mmol/m^3"},
-            "DIC_ALT_CO2": {
-                "long_name": "dissolved inorganic carbon, alternative CO2",
-                "units": "mmol/m^3",
-            },
-            "ALK": {"long_name": "alkalinity", "units": "meq/m^3"},
-            "ALK_ALT_CO2": {
-                "long_name": "alkalinity, alternative CO2",
-                "units": "meq/m^3",
-            },
-            "DOC": {"long_name": "dissolved organic carbon", "units": "mmol/m^3"},
-            "DON": {"long_name": "dissolved organic nitrogen", "units": "mmol/m^3"},
-            "DOP": {"long_name": "dissolved organic phosphorus", "units": "mmol/m^3"},
-            "DOCr": {
-                "long_name": "refractory dissolved organic carbon",
-                "units": "mmol/m^3",
-            },
-            "DONr": {
-                "long_name": "refractory dissolved organic nitrogen",
-                "units": "mmol/m^3",
-            },
-            "DOPr": {
-                "long_name": "refractory dissolved organic phosphorus",
-                "units": "mmol/m^3",
-            },
-            "zooC": {"long_name": "zooplankton carbon", "units": "mmol/m^3"},
-            "spChl": {
-                "long_name": "small phytoplankton chlorophyll",
-                "units": "mg/m^3",
-            },
-            "spC": {"long_name": "small phytoplankton carbon", "units": "mmol/m^3"},
-            "spP": {
-                "long_name": "small phytoplankton phosphorous",
-                "units": "mmol/m^3",
-            },
-            "spFe": {"long_name": "small phytoplankton iron", "units": "mmol/m^3"},
-            "spCaCO3": {"long_name": "small phytoplankton CaCO3", "units": "mmol/m^3"},
-            "diatChl": {"long_name": "diatom chloropyll", "units": "mg/m^3"},
-            "diatC": {"long_name": "diatom carbon", "units": "mmol/m^3"},
-            "diatP": {"long_name": "diatom phosphorus", "units": "mmol/m^3"},
-            "diatFe": {"long_name": "diatom iron", "units": "mmol/m^3"},
-            "diatSi": {"long_name": "diatom silicate", "units": "mmol/m^3"},
-            "diazChl": {"long_name": "diazotroph chloropyll", "units": "mg/m^3"},
-            "diazC": {"long_name": "diazotroph carbon", "units": "mmol/m^3"},
-            "diazP": {"long_name": "diazotroph phosphorus", "units": "mmol/m^3"},
-            "diazFe": {"long_name": "diazotroph iron", "units": "mmol/m^3"},
-            "pco2_air": {"long_name": "atmospheric pCO2", "units": "ppmv"},
-            "pco2_air_alt": {
-                "long_name": "atmospheric pCO2, alternative CO2",
-                "units": "ppmv",
-            },
-            "iron": {"long_name": "iron decomposition", "units": "nmol/cm^2/s"},
-            "dust": {"long_name": "dust decomposition", "units": "kg/m^2/s"},
-            "nox": {"long_name": "NOx decomposition", "units": "kg/m^2/s"},
-            "nhy": {"long_name": "NHy decomposition", "units": "kg/m^2/s"},
-        }
-        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.
-
-        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.
-
-        """
-
-        # Boundary coordinates
-        bdry_coords = {
-            "rho": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_rho": 0},
-            },
-            "u": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_u": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_u": 0},
-            },
-            "v": {
-                "south": {"eta_v": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_v": -1},
-                "west": {"xi_rho": 0},
-            },
-        }
-
-        # 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

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:
@@ -93,26 +71,17 @@ def _plot(
 
     if field is not None:
         p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-        plt.colorbar(p, label=f"{field.long_name} [{field.units}]")
+        if hasattr(field, "long_name"):
+            label = f"{field.long_name} [{field.units}]"
+        elif hasattr(field, "Long_name"):
+            # this is the case for matlab generated grids
+            label = f"{field.Long_name} [{field.units}]"
+        else:
+            label = ""
+        plt.colorbar(p, label=label)
 
     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 +104,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(