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

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.
@@ -300,7 +309,14 @@ class TidalForcing:
         >>> tidal_forcing.plot("ssh_Re", nc=0)
         """
 
-        field = self.ds[var_name].isel(ntides=ntides).compute()
+        field = self.ds[var_name].isel(ntides=ntides)
+
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                field = field.load()
+
         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(
@@ -378,7 +394,13 @@ class TidalForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        dataset_list = [self.ds.load()]
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds.load()
+
+        dataset_list = [self.ds]
         output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
@@ -419,8 +441,8 @@ class TidalForcing:
             "TidalForcing": {
                 "source": self.source,
                 "ntides": self.ntides,
-                "model_reference_date": self.model_reference_date.isoformat(),
                 "allan_factor": self.allan_factor,
+                "model_reference_date": self.model_reference_date.isoformat(),
             }
         }
 
@@ -431,7 +453,7 @@ class TidalForcing:
             # Write header
             file.write(header)
             # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
+            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
     def from_yaml(

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,8 +87,101 @@ def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
     assert hasattr(boundary_forcing.ds, "climatology")
 
 
-def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
-    """Test plot and save methods."""
+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."""
 
     boundary_forcing.plot(var_name="temp_south", layer_contours=True)
     boundary_forcing.plot(var_name="temp_east", layer_contours=True)
@@ -99,6 +194,10 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
     boundary_forcing.plot(var_name="vbar_north")
     boundary_forcing.plot(var_name="ubar_west")
 
+
+def test_boundary_forcing_save(boundary_forcing, tmp_path):
+    """Test save method."""
+
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
@@ -106,9 +205,19 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
+            # Test saving without partitioning and grouping
             saved_filenames = boundary_forcing.save(filepath)
 
+            filepath_str = str(Path(filepath).with_suffix(""))
+            expected_filepath = Path(f"{filepath_str}.nc")
+
+            assert saved_filenames == [expected_filepath]
+            assert expected_filepath.exists()
+            expected_filepath.unlink()
+
+            # Test saving without partitioning but with grouping
+            saved_filenames = boundary_forcing.save(filepath, group=True)
+
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}_202106.nc")
 
@@ -116,8 +225,8 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
             assert expected_filepath.exists()
             expected_filepath.unlink()
 
-            # Test saving with partitioning
-            saved_filenames = boundary_forcing.save(filepath, np_eta=2)
+            # Test saving with partitioning and grouping
+            saved_filenames = boundary_forcing.save(filepath, np_eta=2, group=True)
             expected_filepath_list = [
                 Path(filepath_str + f"_202106.{index}.nc") for index in range(2)
             ]
@@ -129,16 +238,18 @@ def test_boundary_forcing_plot_save(boundary_forcing, tmp_path):
                 expected_filepath.unlink()
 
 
-def test_bgc_boundary_forcing_plot_save(
-    bgc_boundary_forcing_from_climatology, tmp_path
-):
-    """Test plot and save methods."""
+def test_bgc_boundary_forcing_plot(bgc_boundary_forcing_from_climatology):
+    """Test plot method."""
 
     bgc_boundary_forcing_from_climatology.plot(var_name="ALK_south")
     bgc_boundary_forcing_from_climatology.plot(var_name="ALK_east")
     bgc_boundary_forcing_from_climatology.plot(var_name="ALK_north")
     bgc_boundary_forcing_from_climatology.plot(var_name="ALK_west")
 
+
+def test_bgc_boundary_forcing_save(bgc_boundary_forcing_from_climatology, tmp_path):
+    """Test save method."""
+
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
@@ -146,9 +257,20 @@ def test_bgc_boundary_forcing_plot_save(
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
+            # Test saving without partitioning and grouping
             saved_filenames = bgc_boundary_forcing_from_climatology.save(filepath)
 
+            filepath_str = str(Path(filepath).with_suffix(""))
+            expected_filepath = Path(f"{filepath_str}.nc")
+            assert saved_filenames == [expected_filepath]
+            assert expected_filepath.exists()
+            expected_filepath.unlink()
+
+            # Test saving without partitioning but with grouping
+            saved_filenames = bgc_boundary_forcing_from_climatology.save(
+                filepath, group=True
+            )
+
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}_clim.nc")
             assert saved_filenames == [expected_filepath]
@@ -157,7 +279,7 @@ def test_bgc_boundary_forcing_plot_save(
 
             # Test saving with partitioning
             saved_filenames = bgc_boundary_forcing_from_climatology.save(
-                filepath, np_xi=2
+                filepath, np_xi=2, group=True
             )
 
             expected_filepath_list = [
@@ -207,9 +329,9 @@ def test_files_have_same_hash(boundary_forcing, tmp_path, use_dask):
     filepath2 = tmp_path / "test2.nc"
 
     boundary_forcing.to_yaml(yaml_filepath)
-    boundary_forcing.save(filepath1)
+    boundary_forcing.save(filepath1, group=True)
     bdry_forcing_from_file = BoundaryForcing.from_yaml(yaml_filepath, use_dask=use_dask)
-    bdry_forcing_from_file.save(filepath2)
+    bdry_forcing_from_file.save(filepath2, group=True)
 
     filepath_str1 = str(Path(filepath1).with_suffix(""))
     filepath_str2 = str(Path(filepath2).with_suffix(""))
@@ -235,9 +357,9 @@ def test_files_have_same_hash_clim(
     filepath2 = tmp_path / "test2.nc"
 
     bgc_boundary_forcing_from_climatology.to_yaml(yaml_filepath)
-    bgc_boundary_forcing_from_climatology.save(filepath1)
+    bgc_boundary_forcing_from_climatology.save(filepath1, group=True)
     bdry_forcing_from_file = BoundaryForcing.from_yaml(yaml_filepath, use_dask=use_dask)
-    bdry_forcing_from_file.save(filepath2)
+    bdry_forcing_from_file.save(filepath2, group=True)
 
     filepath_str1 = str(Path(filepath1).with_suffix(""))
     filepath_str2 = str(Path(filepath2).with_suffix(""))