roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

roms_tools/setup/plot.py

@@ -10,6 +10,7 @@ def _plot(
     straddle=False,
     c="red",
     title="",
+    with_dim_names=False,
     kwargs={},
 ):
     """Plots a grid or field on a map with optional depth contours.
@@ -65,22 +66,136 @@ def _plot(
 
     fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
 
-    _add_plot_to_ax(
-        ax, lon_deg, lat_deg, trans, field, depth_contours, c, title, kwargs=kwargs
+    if c is not None:
+        _add_boundary_to_ax(
+            ax, lon_deg, lat_deg, trans, c, with_dim_names=with_dim_names
+        )
+
+    if field is not None:
+        _add_field_to_ax(ax, lon_deg, lat_deg, field, depth_contours, kwargs=kwargs)
+
+    ax.coastlines(
+        resolution="50m", linewidth=0.5, color="black"
+    )  # add map of coastlines
+
+    # Add gridlines with labels for latitude and longitude
+    gridlines = ax.gridlines(
+        draw_labels=True, linewidth=0.5, color="gray", alpha=0.7, linestyle="--"
     )
+    gridlines.top_labels = False  # Hide top labels
+    gridlines.right_labels = False  # Hide right labels
+    gridlines.xlabel_style = {
+        "size": 10,
+        "color": "black",
+    }  # Customize longitude label style
+    gridlines.ylabel_style = {
+        "size": 10,
+        "color": "black",
+    }  # Customize latitude label style
+
+    ax.set_title(title)
 
 
-def _add_plot_to_ax(
+def _add_boundary_to_ax(
+    ax, lon_deg, lat_deg, trans, c="red", label="", with_dim_names=False
+):
+    """Plots a grid or field on a map with optional depth contours.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes._axes.Axes
+        The axes on which to plot the data (Cartopy axis with projection).
+
+    lon_deg : np.ndarray
+        Longitude values in degrees.
+
+    lat_deg : np.ndarray
+        Latitude values in degrees.
+
+    trans : cartopy.crs.Projection
+        The projection for transforming coordinates.
+
+    c : str, optional
+        Color of the grid boundary (default is 'red').
+    """
+    proj = ccrs.PlateCarree()
+
+    # find corners
+    corners = [
+        (lon_deg[0, 0], lat_deg[0, 0]),
+        (lon_deg[0, -1], lat_deg[0, -1]),
+        (lon_deg[-1, -1], lat_deg[-1, -1]),
+        (lon_deg[-1, 0], lat_deg[-1, 0]),
+    ]
+
+    # transform coordinates to projected space
+    transformed_corners = [trans.transform_point(lo, la, proj) for lo, la in corners]
+    transformed_lons, transformed_lats = zip(*transformed_corners)
+
+    ax.plot(
+        list(transformed_lons) + [transformed_lons[0]],
+        list(transformed_lats) + [transformed_lats[0]],
+        "o-",
+        c=c,
+        label=label,
+    )
+
+    if with_dim_names:
+        for i in range(len(corners)):
+            if i in [0, 2]:
+                dim_name = r"$\xi$"
+            else:
+                dim_name = r"$\eta$"
+            # Define start and end points for each edge
+            start_lon, start_lat = transformed_corners[i]
+            end_lon, end_lat = transformed_corners[(i + 1) % len(corners)]
+
+            # Compute midpoint
+            mid_lon = (start_lon + end_lon) / 2
+            mid_lat = (start_lat + end_lat) / 2
+
+            # Compute vector direction for arrow
+            arrow_dx = (end_lon - start_lon) * 0.4  # Scale arrow size
+            arrow_dy = (end_lat - start_lat) * 0.4
+
+            # Reverse arrow direction for edges 2 and 3
+            if i in [2, 3]:
+                arrow_dx *= -1
+                arrow_dy *= -1
+
+            # Add arrow
+            ax.annotate(
+                "",
+                xy=(mid_lon + arrow_dx, mid_lat + arrow_dy),
+                xytext=(mid_lon - arrow_dx, mid_lat - arrow_dy),
+                arrowprops=dict(arrowstyle="->", color=c, lw=1.5),
+            )
+
+            ax.text(
+                mid_lon,
+                mid_lat,
+                dim_name,
+                color=c,
+                fontsize=10,
+                ha="center",
+                va="center",
+                bbox=dict(
+                    facecolor="white",
+                    edgecolor="none",
+                    alpha=0.7,
+                    boxstyle="round,pad=0.2",
+                ),
+            )
+
+
+def _add_field_to_ax(
     ax,
     lon_deg,
     lat_deg,
-    trans,
-    field=None,
+    field,
     depth_contours=False,
-    c="red",
-    title="",
     add_colorbar=True,
-    kwargs=None,
+    kwargs={},
 ):
     """Plots a grid or field on a map with optional depth contours.
 
@@ -95,21 +210,12 @@ def _add_plot_to_ax(
     lat_deg : np.ndarray
         Latitude values in degrees.
 
-    trans : cartopy.crs.Projection
-        The projection for transforming coordinates.
-
     field : xarray.DataArray, optional
         Field data to plot (e.g., temperature, salinity). If None, only the grid is plotted.
 
     depth_contours : bool, optional
         If True, adds depth contours to the plot.
 
-    c : str, optional
-        Color of the grid boundary (default is 'red').
-
-    title : str, optional
-        Title of the plot.
-
     add_colorbar : bool, optional
         If True, add colobar.
 
@@ -118,48 +224,20 @@ def _add_plot_to_ax(
 
     Notes
     -----
-    - If `field` is provided, a colorbar is added.
     - If `depth_contours` is True, the field’s `layer_depth` is used to add contours.
     """
     proj = ccrs.PlateCarree()
 
-    # find corners
-    corners = [
-        (lon_deg[0, 0], lat_deg[0, 0]),
-        (lon_deg[0, -1], lat_deg[0, -1]),
-        (lon_deg[-1, -1], lat_deg[-1, -1]),
-        (lon_deg[-1, 0], lat_deg[-1, 0]),
-    ]
-
-    # transform coordinates to projected space
-    transformed_corners = [trans.transform_point(lo, la, proj) for lo, la in corners]
-    transformed_lons, transformed_lats = zip(*transformed_corners)
-
-    if c is not None:
-        ax.plot(
-            list(transformed_lons) + [transformed_lons[0]],
-            list(transformed_lats) + [transformed_lats[0]],
-            "o-",
-            c=c,
-        )
-
-    ax.coastlines(
-        resolution="50m", linewidth=0.5, color="black"
-    )  # add map of coastlines
-    ax.gridlines()
-    ax.set_title(title)
-
-    if field is not None:
-        p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-        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 = ""
-        if add_colorbar:
-            plt.colorbar(p, label=label)
+    p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
+    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 = ""
+    if add_colorbar:
+        plt.colorbar(p, label=label)
 
     if depth_contours:
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
@@ -238,18 +316,14 @@ def _profile_plot(field, title="", ax=None):
     """
 
     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(
             depth_label
-            for depth_label in depths_to_check
-            if depth_label in field.coords
+            for depth_label in field.coords
+            if any(depth_label.startswith(prefix) for prefix in depths_to_check)
         )
     except StopIteration:
         raise ValueError(
@@ -286,3 +360,84 @@ def _line_plot(field, title="", ax=None):
     field.plot(ax=ax)
     ax.set_title(title)
     ax.grid()
+
+
+def _plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=False):
+
+    parent_lon_deg = parent_grid_ds["lon_rho"]
+    parent_lat_deg = parent_grid_ds["lat_rho"]
+
+    child_lon_deg = child_grid_ds["lon_rho"]
+    child_lat_deg = child_grid_ds["lat_rho"]
+
+    if parent_straddle:
+        parent_lon_deg = xr.where(
+            parent_lon_deg > 180, parent_lon_deg - 360, parent_lon_deg
+        )
+        child_lon_deg = xr.where(
+            child_lon_deg > 180, child_lon_deg - 360, child_lon_deg
+        )
+
+    trans = _get_projection(parent_lon_deg, parent_lat_deg)
+
+    parent_lon_deg = parent_lon_deg.values
+    parent_lat_deg = parent_lat_deg.values
+    child_lon_deg = child_lon_deg.values
+    child_lat_deg = child_lat_deg.values
+
+    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
+
+    _add_boundary_to_ax(
+        ax,
+        parent_lon_deg,
+        parent_lat_deg,
+        trans,
+        c="r",
+        label="parent grid",
+        with_dim_names=with_dim_names,
+    )
+
+    _add_boundary_to_ax(
+        ax,
+        child_lon_deg,
+        child_lat_deg,
+        trans,
+        c="g",
+        label="child grid",
+        with_dim_names=with_dim_names,
+    )
+
+    vmax = 3
+    vmin = 0
+    cmap = plt.colormaps.get_cmap("Blues")
+    kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
+
+    _add_field_to_ax(
+        ax,
+        parent_lon_deg,
+        parent_lat_deg,
+        parent_grid_ds.mask_rho,
+        add_colorbar=False,
+        kwargs=kwargs,
+    )
+
+    ax.coastlines(
+        resolution="50m", linewidth=0.5, color="black"
+    )  # add map of coastlines
+
+    # Add gridlines with labels for latitude and longitude
+    gridlines = ax.gridlines(
+        draw_labels=True, linewidth=0.5, color="gray", alpha=0.7, linestyle="--"
+    )
+    gridlines.top_labels = False  # Hide top labels
+    gridlines.right_labels = False  # Hide right labels
+    gridlines.xlabel_style = {
+        "size": 10,
+        "color": "black",
+    }  # Customize longitude label style
+    gridlines.ylabel_style = {
+        "size": 10,
+        "color": "black",
+    }  # Customize latitude label style
+
+    ax.legend(loc="best")

roms_tools/setup/regrid.py

@@ -28,20 +28,22 @@ class LateralRegrid:
             source_dim_names["longitude"]: target_coords["lon"],
         }
 
-    def apply(self, da):
+    def apply(self, da, method="linear"):
         """Fills missing values and regrids the variable.
 
         Parameters
         ----------
         da : xarray.DataArray
             Input data to fill and regrid.
+        method : str
+            Interpolation method to use.
 
         Returns
         -------
         xarray.DataArray
             Regridded data with filled values.
         """
-        regridded = da.interp(self.coords, method="linear").drop_vars(
+        regridded = da.interp(self.coords, method=method).drop_vars(
             list(self.coords.keys())
         )
         return regridded

roms_tools/setup/river_forcing.py

@@ -16,8 +16,9 @@ from roms_tools.setup.utils import (
     save_datasets,
     _to_yaml,
     _from_yaml,
+    get_variable_metadata,
 )
-from roms_tools.setup.plot import _get_projection, _add_plot_to_ax
+from roms_tools.setup.plot import _get_projection, _add_field_to_ax
 import cartopy.crs as ccrs
 
 
@@ -52,6 +53,8 @@ class RiverForcing:
           - "never": Do not compute climatology.
           - "always": Compute climatology for all rivers, regardless of missing data.
 
+    include_bgc : bool, optional
+        Whether to include BGC tracers. Defaults to `False`.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
 
@@ -68,6 +71,7 @@ class RiverForcing:
     end_time: datetime
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]] = None
     convert_to_climatology: str = "if_any_missing"
+    include_bgc: bool = False
     model_reference_date: datetime = datetime(2000, 1, 1)
 
     ds: xr.Dataset = field(init=False, repr=False)
@@ -87,7 +91,7 @@ class RiverForcing:
         object.__setattr__(self, "original_indices", original_indices)
 
         if len(original_indices["station"]) > 0:
-            self.move_rivers_to_closest_coast(target_coords, data)
+            self._move_rivers_to_closest_coast(target_coords, data)
             ds = self._create_river_forcing(data)
             self._validate(ds)
 
@@ -138,13 +142,13 @@ class RiverForcing:
         return data
 
     def _create_river_forcing(self, data):
-        """Create river forcing data for volume flux and tracers (temperature and
-        salinity).
+        """Create river forcing data for volume flux and tracers (temperature, salinity,
+        BGC tracers).
 
-        This method computes the river volume flux and associated tracers (temperature and salinity)
+        This method computes the river volume flux and associated tracers (temperature, salinity, BGC tracers)
         based on the provided input data. It generates a new `xarray.Dataset` that contains:
         - `river_volume`: The river volume flux, calculated as the product of river flux and a specified ratio, with units of m³/s.
-        - `river_tracer`: A tracer array containing temperature and salinity values at each river station over time.
+        - `river_tracer`: A tracer array containing temperature, salinity, and BGC tracer values for each river over time.
 
         The method also handles climatological adjustments for missing or incomplete data, depending on the `convert_to_climatology` setting.
 
@@ -161,7 +165,7 @@ class RiverForcing:
         xr.Dataset
             A new `xarray.Dataset` containing the computed river forcing data. The dataset includes:
             - `river_volume`: A `DataArray` representing the river volume flux (m³/s).
-            - `river_tracer`: A `DataArray` representing tracer data for temperature and salinity at each river station over time.
+            - `river_tracer`: A `DataArray` representing tracer data for temperature, salinity and BGC tracers (if specified) for each river over time.
         """
         if self.source["climatology"]:
             object.__setattr__(self, "climatology", True)
@@ -199,18 +203,64 @@ class RiverForcing:
         river_volume.coords["river_name"] = name
         ds["river_volume"] = river_volume
 
+        if self.include_bgc:
+            ntracers = 2 + 32
+        else:
+            ntracers = 2
         tracer_data = np.zeros(
-            (len(ds.river_time), 2, len(ds.nriver)), dtype=np.float32
+            (len(ds.river_time), ntracers, len(ds.nriver)), dtype=np.float32
         )
         tracer_data[:, 0, :] = 17.0
         tracer_data[:, 1, :] = 1.0
+        tracer_data[:, 2:, :] = 0.0
 
         river_tracer = xr.DataArray(
             tracer_data, dims=("river_time", "ntracers", "nriver")
         )
         river_tracer.attrs["long_name"] = "River tracer data"
-        river_tracer.attrs["units"] = "degrees C [temperature]; psu [salinity]"
-        tracer_names = xr.DataArray(["temperature", "salinity"], dims="ntracers")
+
+        if self.include_bgc:
+            tracer_names = xr.DataArray(
+                [
+                    "temp",
+                    "salt",
+                    "PO4",
+                    "NO3",
+                    "SiO3",
+                    "NH4",
+                    "Fe",
+                    "Lig",
+                    "O2",
+                    "DIC",
+                    "DIC_ALT_CO2",
+                    "ALK",
+                    "ALK_ALT_CO2",
+                    "DOC",
+                    "DON",
+                    "DOP",
+                    "DOPr",
+                    "DONr",
+                    "DOCr",
+                    "zooC",
+                    "spChl",
+                    "spC",
+                    "spP",
+                    "spFe",
+                    "spCaCO3",
+                    "diatChl",
+                    "diatC",
+                    "diatP",
+                    "diatFe",
+                    "diatSi",
+                    "diazChl",
+                    "diazC",
+                    "diazP",
+                    "diazFe",
+                ],
+                dims="ntracers",
+            )
+        else:
+            tracer_names = xr.DataArray(["temp", "salt"], dims="ntracers")
         tracer_names.attrs["long_name"] = "Tracer name"
         river_tracer.coords["tracer_name"] = tracer_names
         ds["river_tracer"] = river_tracer
@@ -225,7 +275,7 @@ class RiverForcing:
 
         return ds
 
-    def move_rivers_to_closest_coast(self, target_coords, data):
+    def _move_rivers_to_closest_coast(self, target_coords, data):
         """Move river mouths to the closest coastal grid cell.
 
         This method computes the closest coastal grid point to each river mouth
@@ -295,10 +345,10 @@ class RiverForcing:
             "xi_rho": indices[2],
             "name": names,
         }
-        self.write_indices_into_grid_file(indices)
+        self._write_indices_into_grid_file(indices)
         object.__setattr__(self, "updated_indices", indices)
 
-    def write_indices_into_grid_file(self, indices):
+    def _write_indices_into_grid_file(self, indices):
         """Writes river location indices into the grid dataset as the "river_flux"
         variable.
 
@@ -396,16 +446,28 @@ class RiverForcing:
         )
 
         for ax in axs:
-            _add_plot_to_ax(
+            _add_field_to_ax(
                 ax,
                 lon_deg,
                 lat_deg,
-                trans,
                 field,
-                c=None,
                 add_colorbar=False,
                 kwargs=kwargs,
             )
+            # Add gridlines with labels for latitude and longitude
+            gridlines = ax.gridlines(
+                draw_labels=True, linewidth=0.5, color="gray", alpha=0.7, linestyle="--"
+            )
+            gridlines.top_labels = False  # Hide top labels
+            gridlines.right_labels = False  # Hide right labels
+            gridlines.xlabel_style = {
+                "size": 10,
+                "color": "black",
+            }  # Customize longitude label style
+            gridlines.ylabel_style = {
+                "size": 10,
+                "color": "black",
+            }  # Customize latitude label style
 
         for ax, indices in zip(axs, [self.original_indices, self.updated_indices]):
             for i in range(len(indices["name"])):
@@ -445,9 +507,43 @@ class RiverForcing:
         ----------
         var_name : str, optional
             The variable to plot. It can be one of the following:
-            - 'river_volume' : Plot the river volume flux.
-            - 'river_temperature' : Plot the river temperature (from the river_tracer).
-            - 'river_salinity' : Plot the river salinity (from the river_tracer).
+
+            - 'river_volume' : river volume flux.
+            - 'river_temp' : river temperature (from river_tracer).
+            - 'river_salt' : river salinity (from river_tracer).
+            - 'river_PO4' : river PO4 (from river_tracer).
+            - 'river_NO3' : river NO3 (from river_tracer).
+            - 'river_SiO3' : river SiO3 (from river_tracer).
+            - 'river_NH4' : river NH4 (from river_tracer).
+            - 'river_Fe' : river Fe (from river_tracer).
+            - 'river_Lig' : river Lig (from river_tracer).
+            - 'river_O2' : river O2 (from river_tracer).
+            - 'river_DIC' : river DIC (from river_tracer).
+            - 'river_DIC_ALT_CO2' : river DIC_ALT_CO2 (from river_tracer).
+            - 'river_ALK' : river ALK (from river_tracer).
+            - 'river_ALK_ALT_CO2' : river ALK_ALT_CO2 (from river_tracer).
+            - 'river_DOC' : river DOC (from river_tracer).
+            - 'river_DON' : river DON (from river_tracer).
+            - 'river_DOP' : river DOP (from river_tracer).
+            - 'river_DOPr' : river DOPr (from river_tracer).
+            - 'river_DONr' : river DONr (from river_tracer).
+            - 'river_DOCr' : river DOCr (from river_tracer).
+            - 'river_zooC' : river zooC (from river_tracer).
+            - 'river_spChl' : river sphChl (from river_tracer).
+            - 'river_spC' : river spC (from river_tracer).
+            - 'river_spP' : river spP (from river_tracer).
+            - 'river_spFe' : river spFe (from river_tracer).
+            - 'river_spCaCO3' : river spCaCO3 (from river_tracer).
+            - 'river_diatChl' : river diatChl (from river_tracer).
+            - 'river_diatC' : river diatC (from river_tracer).
+            - 'river_diatP' : river diatP (from river_tracer).
+            - 'river_diatFe' : river diatFe (from river_tracer).
+            - 'river_diatSi' : river diatSi (from river_tracer).
+            - 'river_diazChl' : river diazChl (from river_tracer).
+            - 'river_diazC' : river diazC (from river_tracer).
+            - 'river_diazP' : river diazP (from river_tracer).
+            - 'river_diazFe' : river diazFe (from river_tracer).
+
             The default is 'river_volume'.
         """
         fig, ax = plt.subplots(1, 1, figsize=(9, 5))
@@ -463,14 +559,14 @@ class RiverForcing:
             field = self.ds[var_name]
             units = f"${self.ds.river_volume.units}$"
             long_name = self.ds[var_name].long_name
-        elif var_name == "river_temperature":
-            field = self.ds["river_tracer"].isel(ntracers=0)
-            units = "degrees C"
-            long_name = "River temperature"
-        elif var_name == "river_salinity":
-            field = self.ds["river_tracer"].isel(ntracers=1)
-            units = "psu"
-            long_name = "River salinity"
+        else:
+            d = get_variable_metadata()
+            var_name_wo_river = var_name.split("_")[1]
+            field = self.ds["river_tracer"].isel(
+                ntracers=self.ds.tracer_name == var_name_wo_river
+            )
+            units = d[var_name_wo_river]["units"]
+            long_name = f"River {d[var_name_wo_river]['long_name']}"
 
         for i in range(len(self.ds.nriver)):
 
@@ -519,10 +615,10 @@ class RiverForcing:
         ----------
         filepath : Union[str, Path]
             The base path and filename for the output files. The filenames will include the specified path and the `.nc` extension.
-            If partitioning is used, additional indices will be appended to the filenames, e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
+            If partitioning is used, additional indices will be appended to the filenames, e.g., `"filepath.0.nc"`, `"filepath.1.nc"`, etc.
 
         filepath_grid : Union[str, Path]
-            The base path and filename for saving the grid file. This file is essential for including the `river_flux` field.
+            The base path and filename for saving the grid file.
 
         np_eta : int, optional
             The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed along the `eta` axis.