roms-tools 2.6.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

ci/environment-with-xesmf.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   - xesmf
   # testing
   - zarr

ci/environment.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   # testing
   - zarr
   - pytest

roms_tools/__init__.py

@@ -1,5 +1,5 @@
+import logging
 from importlib.metadata import version as _version
-import logging  # noqa: F811
 
 try:
     __version__ = _version("roms_tools")
@@ -7,16 +7,18 @@ except ImportError:  # pragma: no cover
     # Local copy or not installed with setuptools
     __version__ = "9999"
 
-
-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.initial_conditions import InitialConditions  # noqa: F401
+# grid must be imported first
+from roms_tools.setup.grid import Grid  # noqa: I001, F401
+from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
-from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
+from roms_tools.setup.cdr_forcing import CDRForcing  # noqa: F401
+from roms_tools.setup.cdr_release import TracerPerturbation, VolumeRelease  # noqa: F401
+from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.nesting import ChildGrid  # noqa: F401
+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.analysis.roms_output import ROMSOutput  # noqa: F401
 
 # Configure logging when the package is imported
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")

roms_tools/analysis/roms_output.py

@@ -1,22 +1,30 @@
-import xarray as xr
-import numpy as np
-import matplotlib.pyplot as plt
-from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
-from roms_tools.utils import _load_data
-from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
-from dataclasses import dataclass, field
-from typing import Union, Optional
-from pathlib import Path
-import re
 import logging
+import re
 import warnings
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
+from roms_tools.analysis.utils import _validate_plot_inputs
+from roms_tools.plot import _line_plot, _plot, _profile_plot, _section_plot
+from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
+from roms_tools.utils import (
+    _generate_coordinate_range,
+    _load_data,
+    _remove_edge_nans,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    normalize_longitude,
+)
 from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
 )
-from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
-from roms_tools.analysis.utils import _validate_plot_inputs, _generate_coordinate_range
 
 
 @dataclass(kw_only=True)
@@ -131,7 +139,7 @@ class ROMSOutput:
             By default, depth contours are not shown (False).
 
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots. Default is None.
+            The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots. Default is None.
 
         save_path : str, optional
             Path to save the generated plot. If None, the plot is shown interactively.
@@ -207,6 +215,9 @@ class ROMSOutput:
         lon_deg = self.grid.ds[f"lon_{loc}"]
         if self.grid.straddle:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        if lon is not None:
+            lon = normalize_longitude(lon, self.grid.straddle)
+
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
         # Mask the field
@@ -301,7 +312,7 @@ class ROMSOutput:
                 lats = [lat]
                 title = title + f", lat = {lat}°N"
             else:
-                resolution = self._infer_nominal_horizontal_resolution()
+                resolution = self.grid._infer_nominal_horizontal_resolution()
                 lats = _generate_coordinate_range(
                     field.lat.min().values, field.lat.max().values, resolution
                 )
@@ -311,7 +322,7 @@ class ROMSOutput:
                 lons = [lon]
                 title = title + f", lon = {lon}°E"
             else:
-                resolution = self._infer_nominal_horizontal_resolution(lat)
+                resolution = self.grid._infer_nominal_horizontal_resolution(lat)
                 lons = _generate_coordinate_range(
                     field.lon.min().values, field.lon.max().values, resolution
                 )
@@ -327,33 +338,6 @@ class ROMSOutput:
         if compute_layer_depth:
             field = field.assign_coords({"layer_depth": layer_depth})
 
-        def _remove_edge_nans(field, xdim, layer_depth=None):
-            """Removes NaNs from the edges along the specified dimension."""
-            if xdim in field.dims:
-                if layer_depth is not None:
-                    nan_mask = layer_depth.isnull().sum(
-                        dim=[dim for dim in layer_depth.dims if dim != xdim]
-                    )
-                else:
-                    nan_mask = field.isnull().sum(
-                        dim=[dim for dim in field.dims if dim != xdim]
-                    )
-
-                # Find the valid indices where the sum of the nans is 0
-                valid_indices = np.where(nan_mask.values == 0)[0]
-
-                if len(valid_indices) > 0:
-                    first_valid = valid_indices[0]
-                    last_valid = valid_indices[-1]
-
-                    field = field.isel({xdim: slice(first_valid, last_valid + 1)})
-                    if layer_depth is not None:
-                        layer_depth = layer_depth.isel(
-                            {xdim: slice(first_valid, last_valid + 1)}
-                        )
-
-            return field, layer_depth
-
         if lat is not None:
             field, layer_depth = _remove_edge_nans(
                 field, "lon", layer_depth if "layer_depth" in locals() else None
@@ -459,7 +443,7 @@ class ROMSOutput:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
 
         if horizontal_resolution is None:
-            horizontal_resolution = self._infer_nominal_horizontal_resolution()
+            horizontal_resolution = self.grid._infer_nominal_horizontal_resolution()
         lons = _generate_coordinate_range(
             lon_deg.min().values, lon_deg.max().values, horizontal_resolution
         )
@@ -472,7 +456,7 @@ class ROMSOutput:
 
         # Prepare vertical regrid
         if depth_levels is None:
-            depth_levels, _ = self._compute_exponential_depth_levels()
+            depth_levels, _ = self.grid._compute_exponential_depth_levels()
 
         # Ensure depth_levels is an xarray.DataArray
         if not isinstance(depth_levels, xr.DataArray):
@@ -808,96 +792,3 @@ class ROMSOutput:
         # Add all necessary coordinates in one go
         ds = ds.assign_coords(coords_to_add)
         return ds
-
-    def _infer_nominal_horizontal_resolution(self, lat=None):
-        """Estimate the nominal horizontal resolution of the grid in degrees at a
-        specified latitude.
-
-        This method calculates the nominal horizontal resolution of the grid by first
-        determining the average grid spacing in meters. The spacing is then converted
-        to degrees, accounting for the Earth's curvature, and the latitude where the
-        resolution is being computed.
-
-        Parameters
-        ----------
-        lat : float, optional
-            Latitude (in degrees) at which to estimate the horizontal resolution.
-            If not provided, the resolution is calculated at the average latitude of
-            the grid (`lat_rho`).
-
-        Returns
-        -------
-        float
-            The estimated horizontal resolution in degrees, adjusted for the Earth's curvature.
-        """
-        # Earth radius in meters
-        r_earth = 6371315.0
-
-        if lat is None:
-            # Center latitude in degrees
-            lat = (self.grid.ds.lat_rho.max() + self.grid.ds.lat_rho.min()) / 2
-
-        # Convert latitude to radians
-        lat_rad = np.deg2rad(lat)
-
-        # Mean resolution in meters
-        resolution_in_m = (
-            (1 / self.grid.ds.pm).mean() + (1 / self.grid.ds.pn).mean()
-        ) / 2
-
-        # Meters per degree at the equator
-        meters_per_degree = 2 * np.pi * r_earth / 360
-
-        # Correct for latitude by multiplying by cos(latitude) for longitude
-        resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
-
-        return resolution_in_degrees
-
-    def _compute_exponential_depth_levels(self, Nz=None, depth=None, h=None):
-        """Compute vertical grid center and face depths using an exponential profile.
-
-        Parameters
-        ----------
-        Nz : int, optional
-            Number of vertical levels. Defaults to `len(self.ds.s_rho)`.
-
-        depth : float, optional
-            Total depth of the domain. Defaults to `grid.ds.h.max().values`.
-
-        h : float, optional
-            Scaling parameter for the exponential profile. Defaults to `Nz / 4.5`.
-
-        Returns
-        -------
-        tuple of numpy.ndarray
-            Depth values at the vertical grid centers (`z_centers`) and grid faces (`z_faces`),
-            both rounded to two decimal places.
-        """
-        if Nz is None:
-            Nz = len(self.ds.s_rho)
-        if depth is None:
-            depth = self.grid.ds.h.max().values
-        if h is None:
-            h = Nz / 4.5
-
-        k = np.arange(1, Nz + 2)
-
-        # Define the exponential profile function
-        def exponential_profile(k, Nz, h):
-            return np.exp(k / h)
-
-        z_faces = np.vectorize(exponential_profile)(k, Nz, h)
-
-        # Normalize
-        z_faces -= z_faces[0]
-        z_faces *= depth / z_faces[-1]
-        z_faces[0] = 0.0
-
-        # Calculate center depths (average between adjacent face depths)
-        z_centers = (z_faces[:-1] + z_faces[1:]) / 2
-
-        # Round both z_faces and z_centers to two decimal places
-        z_faces = np.round(z_faces, 2)
-        z_centers = np.round(z_centers, 2)
-
-        return z_centers, z_faces

roms_tools/analysis/utils.py

@@ -1,6 +1,3 @@
-import numpy as np
-
-
 def _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary):
     """Validate input parameters for the plot method.
 
@@ -104,66 +101,3 @@ def _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary):
                     f"Invalid xi index: {xi} lies on the boundary, which is excluded when `include_boundary = False`. "
                     "Either set `include_boundary = True`, or adjust eta to avoid boundary values."
                 )
-
-
-def _generate_coordinate_range(min, max, resolution):
-    """Generate an array of target coordinates (e.g., latitude or longitude) within a
-    specified range, with a resolution that is rounded to the nearest value of the form
-    `1/n` (or integer).
-
-    This method generates an array of target coordinates between the provided `min` and `max`
-    values, ensuring that both `min` and `max` are included in the resulting range. The resolution
-    is rounded to the nearest fraction of the form `1/n` or an integer, based on the input.
-
-    Parameters
-    ----------
-    min : float
-        The minimum value (in degrees) of the coordinate range (inclusive).
-
-    max : float
-        The maximum value (in degrees) of the coordinate range (inclusive).
-
-    resolution : float
-        The spacing (in degrees) between each coordinate in the array. The resolution will
-        be rounded to the nearest value of the form `1/n` or an integer, depending on the size
-        of the resolution value.
-
-    Returns
-    -------
-    numpy.ndarray
-        An array of target coordinates generated from the specified range, with the resolution
-        rounded to a suitable fraction (e.g., `1/n`) or integer, depending on the input resolution.
-    """
-
-    # Find the closest fraction of the form 1/n or integer to match the resolution
-    resolution_rounded = None
-    min_diff = float("inf")  # Initialize the minimum difference as infinity
-
-    # Search for the best fraction or integer approximation to the resolution
-    for n in range(1, 1000):  # Try fractions 1/n, where n ranges from 1 to 999
-        if resolution <= 1:
-            fraction = (
-                1 / n
-            )  # For small resolutions (<= 1), consider fractions of the form 1/n
-        else:
-            fraction = n  # For larger resolutions (>1), consider integers (n)
-
-        diff = abs(
-            fraction - resolution
-        )  # Calculate the difference between the fraction and the resolution
-
-        if diff < min_diff:  # If the current fraction is a better approximation
-            min_diff = diff
-            resolution_rounded = fraction  # Update the best fraction (or integer) found
-
-    # Adjust the start and end of the range to include integer values
-    start_int = np.floor(min)  # Round the minimum value down to the nearest integer
-    end_int = np.ceil(max)  # Round the maximum value up to the nearest integer
-
-    # Generate the array of target coordinates, including both the min and max values
-    target = np.arange(start_int, end_int + resolution_rounded, resolution_rounded)
-
-    # Truncate any values that exceed max (including small floating point errors)
-    target = target[target <= end_int + 1e-10]
-
-    return target.astype(np.float32)

roms_tools/constants.py

@@ -0,0 +1,5 @@
+R_EARTH = 6371315.0  # Earth radius in meters
+MAXIMUM_GRID_SIZE = 25000  # in km
+UPPER_BOUND_THETA_S = 10  # upper bound for surface vertical stretching parameter
+UPPER_BOUND_THETA_B = 10  # upper bound for bottom vertical stretching parameter
+NUM_TRACERS = 34  # Number of tracers (temperature, salinity, BGC tracers)

roms_tools/download.py

@@ -22,7 +22,7 @@ correction_data = pooch.create(
     },
 )
 
-# Create a Pooch object to manage the global topography data
+# Create a Pooch object to manage the global river data
 river_data = pooch.create(
     # Use the default cache folder for the operating system
     path=pooch.os_cache("roms-tools"),
@@ -32,6 +32,18 @@ river_data = pooch.create(
         "dai_trenberth_may2019.nc": "sha256:793849e6aa60d1f6bdb480c345515fb2453d903c0a30599241b3d752f53715ab",
     },
 )
+
+# Create a Pooch object to manage the global SAL TPXO data
+sal_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "sal_tpxo9.v2a.nc": "sha256:5343d745b4374170a069bfc2c67dcdd9f1dc4eb0df7c0e6de5c004432c903f40",
+    },
+)
+
 # Create a Pooch object to manage the test data
 pup_test_data = pooch.create(
     # Use the default cache folder for the operating system
@@ -46,8 +58,19 @@ pup_test_data = pooch.create(
         "GLORYS_NA_20121231.nc": "03c1155087195deff76ad3f136d6a7f35bc01ccae3402f3d95557a2886d39e71",
         "ERA5_regional_test_data.nc": "bd12ce3b562fbea2a80a3b79ba74c724294043c28dc98ae092ad816d74eac794",
         "ERA5_global_test_data.nc": "8ed177ab64c02caf509b9fb121cf6713f286cc603b1f302f15f3f4eb0c21dc4f",
-        "TPXO_global_test_data.nc": "457bfe87a7b247ec6e04e3c7d3e741ccf223020c41593f8ae33a14f2b5255e60",
-        "TPXO_regional_test_data.nc": "11739245e2286d9c9d342dce5221e6435d2072b50028bef2e86a30287b3b4032",
+        "global_grid_tpxo10.v2.nc": "26eb97cd135cd6f2b4e894c5f11bf7f860ff19cec8dbaa9190e37d30ee6e744e",
+        "global_h_tpxo10.v2.nc": "ef60fae6d52fa514dcc59a737435d74aa798dc114b57f01b123aa39dbaffc592",
+        "global_u_tpxo10.v2.nc": "022e57e6287e51f52eb1e5296614b1086e0e22ecd0bd57c9fd8d0e155babf5c3",
+        "regional_grid_tpxo10v2a.nc": "c5022bfe93ead7cd46e836578645bd87cb5be63c736e660937c7f5703c968cbc",
+        "regional_h_tpxo10.v2.nc": "202fd0c197490ac460af12cd9fa1156aa40023c0023c705f145c596de5b5ad3d",
+        "regional_grid_tpxo10v2.nc": "0789b6a24ecb2ced522481dfcfb7282e32f999984747b9b9f46f044a8898d0ac",
+        "regional_grid_tpxo9v5a.nc": "497a2ae9e6adc7e4b06408dadb57734e2ad24afaa3f0e2e4fd90ebc6eafc2557",
+        "regional_h_tpxo10v2a.nc": "2df2f181f748a960e4072f975226f6f98f6a6c4d5b77da23057946585152d59c",
+        "regional_h_tpxo10v2.nc": "202fd0c197490ac460af12cd9fa1156aa40023c0023c705f145c596de5b5ad3d",
+        "regional_h_tpxo9v5a.nc": "c7e4d9ab73bc11dcb415f88c48131531488e1aed5113df5797e80d3d374607fc",
+        "regional_u_tpxo10v2a.nc": "2d1680ecd53242e858281a762221d91827999967f8e1f3cb7de3d23b47efe8c8",
+        "regional_u_tpxo10v2.nc": "3b0849473cbb7f9076ca907e4fc39eceda3c7d64659c121fa0692024d59dcdb3",
+        "regional_u_tpxo9v5a.nc": "b0cc5f6934d2e212549c7120d458d61a4963ba73d17055e67cc9e4312901b041",
         "CESM_BGC_coarse_global_clim.nc": "20806e4e99285d6de168d3236e2d9245f4e9106474b1464beaa266a73e6ef79f",
         "CESM_BGC_2012.nc": "e374d5df3c1be742d564fd26fd861c2d40af73be50a432c51d258171d5638eb6",
         "CESM_regional_test_data_one_time_slice.nc": "43b578ecc067c85f95d6b97ed7b9dc8da7846f07c95331c6ba7f4a3161036a17",
@@ -56,7 +79,7 @@ pup_test_data = pooch.create(
         "CESM_BGC_SURFACE_2012.nc": "3c4d156adca97909d0fac36bf50b99583ab37d8020d7a3e8511e92abf2331b38",
         "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
         "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
-        "coarsened_UNIFIED_bgc_dataset.nc": "58b510dfd9125e8bfd92342af8cbf7c2639ef9a42c21f3ee74fc31b361e3750c",
+        "coarsened_UNIFIED_bgc_dataset.nc": "c5eaeca6f7640f2358afe5c46ace99ff4879948df8a778a2f21967de034f8d0a",
         "grid_created_with_matlab.nc": "fd537ef8159fabb18e38495ec8d44e2fa1b7fb615fcb1417dd4c0e1bb5f4e41d",
         "etopo5_coarsened_and_shifted.nc": "9a5cb4b38c779d22ddb0ad069b298b9722db34ca85a89273eccca691e89e6f96",
         "srtm15_coarsened.nc": "48bc8f4beecfdca9c192b13f4cbeef1455f49d8261a82563aaec5757e100dff9",
@@ -127,6 +150,26 @@ def download_correction_data(filename: str) -> str:
     return fname
 
 
+def download_sal_data(filename: str) -> str:
+    """Download the SAL data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "sal_tpxo9.v2a.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = sal_data.fetch(filename)
+
+    return fname
+
+
 def download_test_data(filename: str) -> str:
     """Download the test data file.
 

roms_tools/plot.py

@@ -1,7 +1,7 @@
 import cartopy.crs as ccrs
 import matplotlib.pyplot as plt
-import xarray as xr
 import numpy as np
+import xarray as xr
 
 
 def _plot(
@@ -11,7 +11,9 @@ def _plot(
     title="",
     with_dim_names=False,
     plot_data=True,
+    add_colorbar=True,
     kwargs={},
+    ax=None,
 ):
     """Plots a grid or field on a map with optional depth contours.
 
@@ -30,13 +32,17 @@ def _plot(
     plot_data : bool, optional
         If True, plots the provided field data on the map. If False, only the grid
         boundaries and optional depth contours are plotted. Default is True.
+    add_colorbar : bool, optional
+        If True, add colobar.
     kwargs : dict, optional
         Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
+    ax : matplotlib.axes.Axes, optional
+        Pre-existing axes to draw the plot on. If None, a new figure and axes are created.
 
     Returns
     -------
-    matplotlib.figure.Figure
-        The generated figure with the plotted data.
+    matplotlib.figure.Figure, optional
+        The generated figure with the plotted data, only returned if `ax` is None.
 
     Raises
     ------
@@ -56,18 +62,27 @@ def _plot(
 
     trans = _get_projection(lon_deg, lat_deg)
 
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
+
     lon_deg = lon_deg.values
     lat_deg = lat_deg.values
 
-    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
-
     if c is not None:
         _add_boundary_to_ax(
             ax, lon_deg, lat_deg, trans, c, with_dim_names=with_dim_names
         )
 
     if plot_data:
-        _add_field_to_ax(ax, lon_deg, lat_deg, field, depth_contours, kwargs=kwargs)
+        _add_field_to_ax(
+            ax,
+            lon_deg,
+            lat_deg,
+            field,
+            depth_contours,
+            add_colorbar=add_colorbar,
+            kwargs=kwargs,
+        )
 
     ax.coastlines(
         resolution="50m", linewidth=0.5, color="black"
@@ -90,7 +105,9 @@ def _plot(
 
     ax.set_title(title)
 
-    return fig
+    # Only return fig if it was created inside the function (i.e., ax was not provided)
+    if ax is None:
+        return fig
 
 
 def _plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=False):
@@ -215,8 +232,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
     Returns
     -------
-    matplotlib.figure.Figure
-        The generated figure with the plotted section.
+    matplotlib.figure.Figure, optional
+        The generated figure with the plotted section, only returned if `ax` is None.
 
     Raises
     ------
@@ -287,7 +304,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
         xlabel = xdim
     ax.set_xlabel(xlabel)
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _profile_plot(field, title="", ax=None):
@@ -309,7 +327,7 @@ def _profile_plot(field, title="", ax=None):
     Returns
     -------
     matplotlib.figure.Figure
-        The generated figure with the plotted profile.
+        The generated figure with the plotted profile, only returned if `ax` is None.
 
     Raises
     ------
@@ -344,7 +362,8 @@ def _profile_plot(field, title="", ax=None):
     ax.set_ylabel("Depth [m]")
     ax.grid()
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _line_plot(field, title="", ax=None):
@@ -363,7 +382,8 @@ def _line_plot(field, title="", ax=None):
     Returns
     -------
     matplotlib.figure.Figure
-        The generated figure with the plotted data and highlighted NaN regions.
+        The generated figure with the plotted data and highlighted NaN regions, only
+        returned if `ax` is None.
 
     Raises
     ------
@@ -430,7 +450,8 @@ def _line_plot(field, title="", ax=None):
         xlabel = xdim
     ax.set_xlabel(xlabel)
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _add_boundary_to_ax(

roms_tools/regrid.py

@@ -1,7 +1,8 @@
-import xgcm
-import xarray as xr
 import warnings
 
+import xarray as xr
+import xgcm
+
 
 class LateralRegridToROMS:
     """Handles lateral regridding of data onto a new spatial grid."""

roms_tools/setup/boundary_forcing.py

@@ -1,39 +1,42 @@
-import xarray as xr
-import numpy as np
-from scipy.ndimage import label
-import logging
 import importlib.metadata
-from typing import Dict, Union, List, Optional
+import logging
 from dataclasses import dataclass, field
 from datetime import datetime
-import matplotlib.pyplot as plt
 from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+from scipy.ndimage import label
+
 from roms_tools import Grid
+from roms_tools.plot import _line_plot, _section_plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.utils import save_datasets
-from roms_tools.vertical_coordinate import compute_depth
-from roms_tools.plot import _section_plot, _line_plot
-from roms_tools.utils import (
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-    transpose_dimensions,
-)
-from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset, UnifiedBGCDataset
+from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
 from roms_tools.setup.utils import (
-    get_variable_metadata,
-    group_dataset,
-    get_target_coords,
-    rotate_velocities,
+    _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    add_time_info_to_ds,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
-    one_dim_fill,
+    get_boundary_coords,
+    get_target_coords,
+    get_variable_metadata,
+    group_dataset,
     nan_check,
+    one_dim_fill,
+    rotate_velocities,
     substitute_nans_by_fillvalue,
-    convert_to_roms_time,
-    get_boundary_coords,
-    _to_yaml,
-    _from_yaml,
 )
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    save_datasets,
+    transpose_dimensions,
+)
+from roms_tools.vertical_coordinate import compute_depth
 
 
 @dataclass(kw_only=True)
@@ -686,8 +689,7 @@ class BoundaryForcing:
         ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
         ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
-        # Convert the time coordinate to the format expected by ROMS
-        ds, bry_time = convert_to_roms_time(
+        ds, bry_time = add_time_info_to_ds(
             ds, self.model_reference_date, data.climatology
         )
 
@@ -1003,7 +1005,8 @@ class BoundaryForcing:
             The path to the YAML file where the parameters will be saved.
         """
 
-        _to_yaml(self, filepath)
+        forcing_dict = _to_dict(self)
+        _write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(