roms-tools 1.6.1__py3-none-any.whl → 1.6.2__py3-none-any.whl

roms_tools/setup/plot.py

@@ -12,6 +12,32 @@ def _plot(
     title="",
     kwargs={},
 ):
+    """Plots a grid or field on a map with optional depth contours.
+
+    This function plots a map using Cartopy projections. It supports plotting a grid, a field, and adding depth contours if desired.
+    The projection can be customized, and the grid can be adjusted for domains straddling the 180° meridian.
+
+    Parameters
+    ----------
+    grid_ds : xarray.Dataset
+        The grid dataset containing coordinates (`lon_rho`, `lat_rho`).
+    field : xarray.DataArray, optional
+        The field to plot. If None, only the grid is plotted.
+    depth_contours : bool, optional
+        If True, adds depth contours to the plot.
+    straddle : bool, optional
+        If True, adjusts longitude values to straddle across the 180° meridian.
+    c : str, optional
+        Color for the boundary plot (default is 'red').
+    title : str, optional
+        Title of the plot.
+    kwargs : dict, optional
+        Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
+
+    Notes
+    -----
+    The function raises a `NotImplementedError` if the domain contains the North or South Pole.
+    """
 
     if field is None:
         lon_deg = grid_ds["lon_rho"]
@@ -84,12 +110,11 @@ def _plot(
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
-    return fig
 
+def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
-def _section_plot(field, interface_depth=None, title="", kwargs={}):
-
-    fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(9, 5))
 
     dims_to_check = ["eta_rho", "eta_u", "eta_v", "xi_rho", "xi_u", "xi_v"]
     try:
@@ -132,7 +157,23 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}):
     ax.set_title(title)
 
 
-def _profile_plot(field, title=""):
+def _profile_plot(field, title="", ax=None):
+    """Plots a profile of the given field against depth.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Raises
+    ------
+    ValueError
+        If no expected depth coordinate is found in the field.
+    """
 
     depths_to_check = [
         "layer_depth_rho",
@@ -153,16 +194,33 @@ def _profile_plot(field, title=""):
             "None of the expected coordinates (layer_depth_rho, layer_depth_u, layer_depth_v, interface_depth_rho, interface_depth_u, interface_depth_v) found in field.coords"
         )
 
-    fig, ax = plt.subplots(1, 1, figsize=(4, 7))
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(4, 7))
     kwargs = {"y": depth_label, "yincrease": False}
     field.plot(**kwargs)
     ax.set_title(title)
     ax.grid()
 
 
-def _line_plot(field, title=""):
-
-    fig, ax = plt.subplots(1, 1, figsize=(7, 4))
-    field.plot()
+def _line_plot(field, title="", ax=None):
+    """Plots a line graph of the given field.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Returns
+    -------
+    None
+        Modifies the plot in-place.
+    """
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(7, 4))
+    field.plot(ax=ax)
     ax.set_title(title)
     ax.grid()

roms_tools/setup/surface_forcing.py

@@ -133,14 +133,7 @@ class SurfaceForcing:
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
 
-        if self.use_coarse_grid:
-            mask = self.grid.ds["mask_coarse"].rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            )
-        else:
-            mask = self.grid.ds["mask_rho"]
-
-        self._validate(ds, mask)
+        self._validate(ds, target_coords["mask"], variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -233,28 +226,34 @@ class SurfaceForcing:
         # Define a dictionary for variable names and their associated information
         if self.type == "physics":
             variable_info = {
-                "swrad": default_info,
-                "lwrad": default_info,
-                "Tair": default_info,
-                "qair": default_info,
-                "rain": default_info,
+                "swrad": {**default_info, "validate": True},
+                "lwrad": {**default_info, "validate": False},
+                "Tair": {**default_info, "validate": False},
+                "qair": {**default_info, "validate": True},
+                "rain": {**default_info, "validate": False},
                 "uwnd": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "vwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
                 "vwnd": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "uwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
             }
         elif self.type == "bgc":
             variable_info = {}
-            for var in data.var_names.keys():
-                variable_info[var] = default_info
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
+                if var_name == "pco2_air":
+                    variable_info[var_name] = {**default_info, "validate": True}
+                else:
+                    variable_info[var_name] = {**default_info, "validate": False}
 
         return variable_info
 
@@ -300,9 +299,6 @@ class SurfaceForcing:
             ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
             ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
-        if self.use_coarse_grid:
-            ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-
         ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
@@ -369,7 +365,7 @@ class SurfaceForcing:
 
         return ds
 
-    def _validate(self, ds, mask):
+    def _validate(self, ds, mask, variable_info):
         """Validates the dataset by checking for NaN values at wet points, which would
         indicate missing raw data coverage over the target domain.
 
@@ -379,6 +375,10 @@ class SurfaceForcing:
             The dataset to validate.
         mask : xarray.DataArray
             Land mask (1=ocean, 0=land) to determine wet points in the domain.
+        variable_info : dict
+            A dictionary containing metadata about each variable (e.g., location,
+            whether it's a 3D variable, etc.). Used to retrieve information for
+            validating each variable.
 
         Raises
         ------
@@ -392,7 +392,9 @@ class SurfaceForcing:
         """
 
         for var_name in ds.data_vars:
-            nan_check(ds[var_name].isel(time=0), mask)
+            # Only validate variables based on "validate" flag if use_dask is False
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
 
@@ -471,12 +473,7 @@ class SurfaceForcing:
 
         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.where(self.target_coords["mask"])
 
         field = field.assign_coords(
             {"lon": self.target_coords["lon"], "lat": self.target_coords["lat"]}

roms_tools/setup/tides.py

@@ -129,7 +129,7 @@ class TidalForcing:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        self._validate(ds, variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -176,33 +176,37 @@ class TidalForcing:
 
         # Define a dictionary for variable names and their associated information
         variable_info = {
-            "ssh_Re": default_info,
-            "ssh_Im": default_info,
-            "pot_Re": default_info,
-            "pot_Im": default_info,
+            "ssh_Re": {**default_info, "validate": True},
+            "ssh_Im": {**default_info, "validate": False},
+            "pot_Re": {**default_info, "validate": False},
+            "pot_Im": {**default_info, "validate": False},
             "u_Re": {
                 "location": "u",
                 "is_vector": True,
                 "vector_pair": "v_Re",
                 "is_3d": False,
+                "validate": True,
             },
             "v_Re": {
                 "location": "v",
                 "is_vector": True,
                 "vector_pair": "u_Re",
                 "is_3d": False,
+                "validate": True,
             },
             "u_Im": {
                 "location": "u",
                 "is_vector": True,
                 "vector_pair": "v_Im",
                 "is_3d": False,
+                "validate": False,
             },
             "v_Im": {
                 "location": "v",
                 "is_vector": True,
                 "vector_pair": "u_Im",
                 "is_3d": False,
+                "validate": False,
             },
         }
 
@@ -239,14 +243,17 @@ class TidalForcing:
 
         return ds
 
-    def _validate(self, ds):
-        """Validates the dataset by checking for NaN values at wet points, which would
-        indicate missing raw data coverage over the target domain.
+    def _validate(self, ds, variable_info):
+        """Validates the dataset by checking for NaN values at wet points for specified
+        variables, which would indicate missing raw data coverage over the target
+        domain.
 
         Parameters
         ----------
         ds : xarray.Dataset
             The dataset to validate, containing tidal variables and a mask for wet points.
+        variable_info : dict
+            A dictionary containing metadata about the variables, including whether to validate them.
 
         Raises
         ------
@@ -260,7 +267,9 @@ class TidalForcing:
         The method utilizes `self.grid.ds.mask_rho` to determine the wet points in the domain.
         """
         for var_name in ds.data_vars:
-            nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
+            # only validate variables based on "validate" flag if use_dask is false
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
 
     def plot(self, var_name, ntides=0) -> None:
         """Plot the specified tidal forcing variable for a given tidal constituent.

roms_tools/setup/utils.py

@@ -724,23 +724,31 @@ def get_target_coords(grid, use_coarse_grid=False):
     Returns
     -------
     dict
-        Dictionary containing the longitude, latitude, and angle arrays, along with a boolean indicating
+        Dictionary containing the longitude, latitude, angle and mask arrays, along with a boolean indicating
         if the grid straddles the meridian.
     """
     # Select grid variables based on whether the coarse grid is used
     if use_coarse_grid:
-        lat, lon, angle = (
-            grid.ds.lat_coarse,
-            grid.ds.lon_coarse,
-            grid.ds.angle_coarse,
+        lat, lon, angle, mask = (
+            grid.ds.lat_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
+            grid.ds.lon_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
+            grid.ds.angle_coarse.rename(
+                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+            ),
+            grid.ds.mask_coarse.rename(
+                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+            ),
         )
+
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
+
     else:
-        lat, lon, angle = (
+        lat, lon, angle, mask = (
             grid.ds.lat_rho,
             grid.ds.lon_rho,
             grid.ds.angle,
+            grid.ds.mask_rho,
         )
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
@@ -763,6 +771,7 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lat_psi": lat_psi,
         "lon_psi": lon_psi,
         "angle": angle,
+        "mask": mask,
         "straddle": straddle,
     }
 

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -1,17 +1,19 @@
 import pytest
 from datetime import datetime
-from roms_tools import BoundaryForcing
+import xarray as xr
+from roms_tools import Grid, BoundaryForcing
 import textwrap
 from roms_tools.setup.download import download_test_data
 from conftest import calculate_file_hash
 from pathlib import Path
+import logging
 
 
 @pytest.mark.parametrize(
     "boundary_forcing_fixture",
     [
         "boundary_forcing",
-        "boundary_forcing_with_2d_fill",
+        # "boundary_forcing_with_2d_fill",
     ],
 )
 def test_boundary_forcing_creation(boundary_forcing_fixture, request):
@@ -51,7 +53,7 @@ def test_boundary_forcing_creation(boundary_forcing_fixture, request):
     "boundary_forcing_fixture",
     [
         "bgc_boundary_forcing_from_climatology",
-        "bgc_boundary_forcing_from_climatology_with_2d_fill",
+        # "bgc_boundary_forcing_from_climatology_with_2d_fill",
     ],
 )
 def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
@@ -85,6 +87,99 @@ def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
     assert hasattr(boundary_forcing.ds, "climatology")
 
 
+def test_unsuccessful_boundary_forcing_creation_with_1d_fill(use_dask):
+
+    grid = Grid(
+        nx=2,
+        ny=2,
+        size_x=500,
+        size_y=1000,
+        center_lon=0,
+        center_lat=55,
+        rot=10,
+        N=3,  # number of vertical levels
+        theta_s=5.0,  # surface control parameter
+        theta_b=2.0,  # bottom control parameter
+        hc=250.0,  # critical depth
+    )
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+
+    with pytest.raises(ValueError, match="consists entirely of NaNs"):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"name": "GLORYS", "path": fname},
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+
+    fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+
+    with pytest.raises(ValueError, match="consists entirely of NaNs"):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"path": fname_bgc, "name": "CESM_REGRIDDED", "climatology": True},
+            type="bgc",
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+
+
+def test_boundary_divided_by_land_warning(caplog, use_dask):
+
+    # Iceland intersects the western boundary of the following grid
+    grid = Grid(
+        nx=5, ny=5, size_x=500, size_y=500, center_lon=-10, center_lat=65, rot=0
+    )
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+
+    with caplog.at_level(logging.WARNING):
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"path": fname, "name": "GLORYS", "climatology": False},
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert "the western boundary is divided by land" in caplog.text
+
+
+def test_1d_and_2d_fill_coincide_if_no_land(use_dask):
+
+    # this grid lies entirely over open ocean
+    grid = Grid(nx=5, ny=5, size_x=300, size_y=300, center_lon=-5, center_lat=65, rot=0)
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+
+    kwargs = {
+        "grid": grid,
+        "start_time": datetime(2021, 6, 29),
+        "end_time": datetime(2021, 6, 29),
+        "source": {"path": fname, "name": "GLORYS", "climatology": False},
+        "use_dask": use_dask,
+    }
+
+    bf_1d_fill = BoundaryForcing(
+        **kwargs,
+        apply_2d_horizontal_fill=False,
+    )
+    bf_2d_fill = BoundaryForcing(
+        **kwargs,
+        apply_2d_horizontal_fill=True,
+    )
+
+    xr.testing.assert_allclose(bf_1d_fill.ds, bf_2d_fill.ds, rtol=1.0e-4)
+
+
 def test_boundary_forcing_plot(boundary_forcing):
     """Test plot."""