roms-tools 2.7.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -1,18 +1,22 @@
-import pytest
+import logging
+import textwrap
 from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import pytest
 import xarray as xr
+
+from conftest import calculate_data_hash
 from roms_tools import Grid, SurfaceForcing
 from roms_tools.download import download_test_data
-import textwrap
-from pathlib import Path
-import logging
-from conftest import calculate_data_hash
 
 
 @pytest.fixture
 def grid_that_straddles_dateline():
     """Fixture for creating a domain that straddles the dateline and lies within the
-    bounds of the regional ERA5 data."""
+    bounds of the regional ERA5 data.
+    """
     grid = Grid(
         nx=20,
         ny=20,
@@ -75,7 +79,6 @@ def grid_that_lies_east_of_dateline_less_than_five_degrees_away():
     We care about the 5 degree mark because it decides whether the code handles the
     longitudes as straddling the dateline or not.
     """
-
     grid = Grid(
         nx=5,
         ny=5,
@@ -118,7 +121,6 @@ def grid_that_lies_west_of_dateline_less_than_five_degrees_away():
     We care about the 5 degree mark because it decides whether the code handles the
     longitudes as straddling the dateline or not.
     """
-
     grid = Grid(
         nx=5,
         ny=5,
@@ -140,7 +142,6 @@ def grid_that_lies_west_of_dateline_more_than_five_degrees_away():
     We care about the 5 degree mark because it decides whether the code handles the
     longitudes as straddling the dateline or not.
     """
-
     grid = Grid(
         nx=5,
         ny=5,
@@ -154,6 +155,60 @@ def grid_that_lies_west_of_dateline_more_than_five_degrees_away():
     return grid
 
 
+def _test_successful_initialization(
+    grid: Grid,
+    start_time: datetime,
+    end_time: datetime,
+    source: dict[str, str],
+    coarse_grid_mode: str,
+    use_dask: bool,
+    caplog,
+):
+    with caplog.at_level(logging.INFO):
+        sfc_forcing = SurfaceForcing(
+            grid=grid,
+            start_time=start_time,
+            end_time=end_time,
+            source=source,
+            correct_radiation=True,
+            wind_dropoff=True,
+            coarse_grid_mode=coarse_grid_mode,
+            use_dask=use_dask,
+        )
+
+        assert sfc_forcing.grid == grid
+        assert sfc_forcing.start_time == start_time
+        assert sfc_forcing.end_time == end_time
+        assert sfc_forcing.type == "physics"
+        assert sfc_forcing.source == source
+        assert sfc_forcing.correct_radiation
+        assert sfc_forcing.wind_dropoff
+
+        if coarse_grid_mode == "always":
+            assert sfc_forcing.use_coarse_grid
+            assert (
+                "Data will be interpolated onto grid coarsened by factor 2."
+                in caplog.text
+            )
+        elif coarse_grid_mode == "never":
+            assert not sfc_forcing.use_coarse_grid
+            assert "Data will be interpolated onto fine grid." in caplog.text
+
+        assert isinstance(sfc_forcing.ds, xr.Dataset)
+        assert "uwnd" in sfc_forcing.ds
+        assert "vwnd" in sfc_forcing.ds
+        assert "swrad" in sfc_forcing.ds
+        assert "lwrad" in sfc_forcing.ds
+        assert "Tair" in sfc_forcing.ds
+        assert "qair" in sfc_forcing.ds
+        assert "rain" in sfc_forcing.ds
+        assert sfc_forcing.ds.coords["time"].attrs["units"] == "days"
+
+        assert sfc_forcing.ds.attrs["source"] == source["name"]
+        assert sfc_forcing.ds.attrs["correct_radiation"] == "True"
+        assert sfc_forcing.ds.attrs["wind_dropoff"] == "True"
+
+
 @pytest.mark.parametrize(
     "grid_fixture",
     [
@@ -165,7 +220,7 @@ def grid_that_lies_west_of_dateline_more_than_five_degrees_away():
     ],
 )
 def test_successful_initialization_with_regional_data(
-    grid_fixture, request, caplog, use_dask
+    grid_fixture, request, use_dask, caplog
 ):
     """Test the initialization of SurfaceForcing with regional ERA5 data.
 
@@ -173,58 +228,56 @@ def test_successful_initialization_with_regional_data(
     - First with the default fine grid.
     - Then with the coarse grid enabled.
     """
-
-    start_time = datetime(2020, 1, 31)
-    end_time = datetime(2020, 2, 2)
-
     fname = Path(download_test_data("ERA5_regional_test_data.nc"))
-
     grid = request.getfixturevalue(grid_fixture)
 
     for coarse_grid_mode in ["always", "never"]:
-        with caplog.at_level(logging.INFO):
-            sfc_forcing = SurfaceForcing(
-                grid=grid,
-                start_time=start_time,
-                end_time=end_time,
-                source={"name": "ERA5", "path": fname},
-                correct_radiation=True,
-                coarse_grid_mode=coarse_grid_mode,
-                use_dask=use_dask,
-            )
+        _test_successful_initialization(
+            grid=grid,
+            start_time=datetime(2020, 1, 31),
+            end_time=datetime(2020, 2, 2),
+            source={"name": "ERA5", "path": fname, "climatology": False},
+            coarse_grid_mode=coarse_grid_mode,
+            use_dask=use_dask,
+            caplog=caplog,
+        )
 
-        assert sfc_forcing.ds is not None
-        assert "uwnd" in sfc_forcing.ds
-        assert "vwnd" in sfc_forcing.ds
-        assert "swrad" in sfc_forcing.ds
-        assert "lwrad" in sfc_forcing.ds
-        assert "Tair" in sfc_forcing.ds
-        assert "qair" in sfc_forcing.ds
-        assert "rain" in sfc_forcing.ds
 
-        assert sfc_forcing.start_time == start_time
-        assert sfc_forcing.end_time == end_time
-        assert sfc_forcing.type == "physics"
-        assert sfc_forcing.source == {
-            "name": "ERA5",
-            "path": fname,
-            "climatology": False,
-        }
-        assert sfc_forcing.ds.coords["time"].attrs["units"] == "days"
+@pytest.mark.parametrize(
+    "grid_fixture",
+    [
+        "grid_that_straddles_dateline",
+        "grid_that_lies_east_of_dateline_less_than_five_degrees_away",
+        "grid_that_lies_east_of_dateline_more_than_five_degrees_away",
+        "grid_that_lies_west_of_dateline_less_than_five_degrees_away",
+        "grid_that_lies_west_of_dateline_more_than_five_degrees_away",
+        "grid_that_straddles_dateline_but_is_too_big_for_regional_test_data",
+        "another_grid_that_straddles_dateline_but_is_too_big_for_regional_test_data",
+        "grid_that_straddles_180_degree_meridian",
+    ],
+)
+def test_successful_initialization_with_global_data(
+    grid_fixture, request, use_dask, caplog
+):
+    """Test initialization of SurfaceForcing with global data.
 
-        if coarse_grid_mode == "always":
-            assert sfc_forcing.use_coarse_grid
-            assert (
-                "Data will be interpolated onto grid coarsened by factor 2."
-                in caplog.text
-            )
-        elif coarse_grid_mode == "never":
-            assert not sfc_forcing.use_coarse_grid
-            assert "Data will be interpolated onto fine grid." in caplog.text
+    Verifies that the SurfaceForcing object is correctly initialized with global data,
+    including the correct handling of the grid and physics data. Checks both coarse and
+    fine grid initialization.
+    """
+    fname = Path(download_test_data("ERA5_global_test_data.nc"))
+    grid = request.getfixturevalue(grid_fixture)
 
-        sfc_forcing.plot("uwnd", time=0)
-        sfc_forcing.plot("vwnd", time=0)
-        sfc_forcing.plot("rain", time=0)
+    for coarse_grid_mode in ["always", "never"]:
+        _test_successful_initialization(
+            grid=grid,
+            start_time=datetime(2020, 1, 31),
+            end_time=datetime(2020, 2, 2),
+            source={"name": "ERA5", "path": fname, "climatology": False},
+            coarse_grid_mode=coarse_grid_mode,
+            use_dask=use_dask,
+            caplog=caplog,
+        )
 
 
 @pytest.mark.parametrize(
@@ -250,7 +303,6 @@ def test_nan_detection_initialization_with_regional_data(
 
     for coarse_grid_mode in ["always", "never"]:
         with pytest.raises(ValueError, match="NaN values found"):
-
             SurfaceForcing(
                 grid=grid,
                 coarse_grid_mode=coarse_grid_mode,
@@ -279,7 +331,6 @@ def test_no_longitude_intersection_initialization_with_regional_data(
         with pytest.raises(
             ValueError, match="Selected longitude range does not intersect with dataset"
         ):
-
             SurfaceForcing(
                 grid=grid_that_straddles_180_degree_meridian,
                 coarse_grid_mode=coarse_grid_mode,
@@ -290,67 +341,6 @@ def test_no_longitude_intersection_initialization_with_regional_data(
             )
 
 
-@pytest.mark.parametrize(
-    "grid_fixture",
-    [
-        "grid_that_straddles_dateline",
-        "grid_that_lies_east_of_dateline_less_than_five_degrees_away",
-        "grid_that_lies_east_of_dateline_more_than_five_degrees_away",
-        "grid_that_lies_west_of_dateline_less_than_five_degrees_away",
-        "grid_that_lies_west_of_dateline_more_than_five_degrees_away",
-        "grid_that_straddles_dateline_but_is_too_big_for_regional_test_data",
-        "another_grid_that_straddles_dateline_but_is_too_big_for_regional_test_data",
-        "grid_that_straddles_180_degree_meridian",
-    ],
-)
-def test_successful_initialization_with_global_data(grid_fixture, request, use_dask):
-    """Test initialization of SurfaceForcing with global data.
-
-    Verifies that the SurfaceForcing object is correctly initialized with global data,
-    including the correct handling of the grid and physics data. Checks both coarse and
-    fine grid initialization.
-    """
-    start_time = datetime(2020, 1, 31)
-    end_time = datetime(2020, 2, 2)
-
-    fname = Path(download_test_data("ERA5_global_test_data.nc"))
-
-    grid = request.getfixturevalue(grid_fixture)
-
-    for coarse_grid_mode in ["always", "never"]:
-        sfc_forcing = SurfaceForcing(
-            grid=grid,
-            coarse_grid_mode=coarse_grid_mode,
-            start_time=start_time,
-            end_time=end_time,
-            source={"name": "ERA5", "path": fname},
-            use_dask=use_dask,
-        )
-        assert sfc_forcing.start_time == start_time
-        assert sfc_forcing.end_time == end_time
-        assert sfc_forcing.type == "physics"
-        assert sfc_forcing.source == {
-            "name": "ERA5",
-            "path": fname,
-            "climatology": False,
-        }
-
-        assert "uwnd" in sfc_forcing.ds
-        assert "vwnd" in sfc_forcing.ds
-        assert "swrad" in sfc_forcing.ds
-        assert "lwrad" in sfc_forcing.ds
-        assert "Tair" in sfc_forcing.ds
-        assert "qair" in sfc_forcing.ds
-        assert "rain" in sfc_forcing.ds
-        assert sfc_forcing.ds.attrs["source"] == "ERA5"
-        assert sfc_forcing.ds.coords["time"].attrs["units"] == "days"
-
-        if coarse_grid_mode == "always":
-            assert sfc_forcing.use_coarse_grid
-        elif coarse_grid_mode == "never":
-            assert not sfc_forcing.use_coarse_grid
-
-
 def test_start_time_end_time_error(use_dask):
     """Test error when start_time and end_time are not both provided or both None."""
     # Case 1: Only start_time provided
@@ -415,7 +405,6 @@ def test_nans_filled_in(
     - First with the default fine grid.
     - Then with the coarse grid enabled.
     """
-
     start_time = datetime(2020, 1, 31)
     end_time = datetime(2020, 2, 2)
 
@@ -451,8 +440,8 @@ def test_nans_filled_in(
 )
 def test_time_attr_climatology(bgc_surface_forcing_fixture, request):
     """Test that the 'cycle_length' attribute is present in the time coordinate of the
-    BGC dataset when using climatology data."""
-
+    BGC dataset when using climatology data.
+    """
     bgc_surface_forcing = request.getfixturevalue(bgc_surface_forcing_fixture)
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert hasattr(
@@ -464,7 +453,8 @@ def test_time_attr_climatology(bgc_surface_forcing_fixture, request):
 
 def test_time_attr(bgc_surface_forcing):
     """Test that the 'cycle_length' attribute is not present in the time coordinate of
-    the BGC dataset when not using climatology data."""
+    the BGC dataset when not using climatology data.
+    """
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert not hasattr(
             bgc_surface_forcing.ds[time_coord],
@@ -505,7 +495,6 @@ def test_surface_forcing_creation(
     Ensures that expected variables are present in the dataset and that attributes match
     the given configurations.
     """
-
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
     assert sfc_forcing.ds is not None
@@ -525,8 +514,6 @@ def test_surface_forcing_creation(
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert sfc_forcing.ds.coords[time_coord].attrs["units"] == "days"
 
-    sfc_forcing.plot("pco2_air", time=0)
-
 
 @pytest.mark.parametrize(
     "sfc_forcing_fixture",
@@ -538,7 +525,6 @@ def test_surface_forcing_creation(
 )
 def test_surface_forcing_pco2_replication(sfc_forcing_fixture, request):
     """Test whether pco2_air and pco2_air_alt is the same after processing."""
-
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
     xr.testing.assert_allclose(
@@ -546,6 +532,19 @@ def test_surface_forcing_pco2_replication(sfc_forcing_fixture, request):
     )
 
 
+def test_computed_missing_optional_fields(bgc_surface_forcing_from_unified_climatology):
+    ds = bgc_surface_forcing_from_unified_climatology.ds
+
+    # Use tight tolerances because 'nox' and 'nhy' can have values order 1e-12
+
+    # 'nhy' was missing in the source data and should have been filled with a constant default value
+    assert np.allclose(ds.nhy.std(), 0.0, rtol=1e-13, atol=1e-13), (
+        "NHy should be constant across space and time"
+    )
+    # 'nox' was present in the source data and should show spatial or temporal variability
+    assert ds.nox.std() > 1e-13, "NOx should vary across space and time"
+
+
 def test_determine_usage_coarse_grid():
     # ERA5 data with 1/4 degree resolution spanning [-50E, 30E] and [40N, 80N]
     fname = download_test_data("ERA5_regional_test_data.nc")
@@ -597,26 +596,69 @@ def test_determine_usage_coarse_grid():
     assert surface_forcing.use_coarse_grid
 
 
+def test_apply_wind_correction(surface_forcing):
+    uwnd = surface_forcing.ds["uwnd"]
+    vwnd = surface_forcing.ds["vwnd"]
+
+    prev_coords = surface_forcing.target_coords.copy()
+    uwnd_corr, vwnd_corr = surface_forcing._apply_wind_correction(uwnd, vwnd)
+
+    assert isinstance(uwnd_corr, xr.DataArray)
+    assert isinstance(vwnd_corr, xr.DataArray)
+
+    # sanity check that the degrees conversion doesn't change the coords on our xarray
+    assert prev_coords == surface_forcing.target_coords
+
+    # Wind correction should not increase magnitude
+    assert (abs(uwnd_corr) <= abs(uwnd)).all()
+    assert (abs(vwnd_corr) <= abs(vwnd)).all()
+
+    # Direction (sign) should be preserved
+    assert (np.sign(uwnd_corr) == np.sign(uwnd)).all()
+    assert (np.sign(vwnd_corr) == np.sign(vwnd)).all()
+
+    # the ratio should be 1 far away, and 0.6 over/near land
+    assert np.isclose((uwnd_corr / uwnd).max(), 1.0)
+    assert np.isclose((vwnd_corr / vwnd).max(), 1.0)
+
+    assert np.isclose((uwnd_corr / uwnd).min(), 0.6)
+    assert np.isclose((vwnd_corr / vwnd).min(), 0.6)
+
+
 @pytest.mark.parametrize(
     "sfc_forcing_fixture",
     [
         "surface_forcing",
         "corrected_surface_forcing",
+        "surface_forcing_with_wind_dropoff",
         "coarse_surface_forcing",
     ],
 )
-def test_surface_forcing_save(sfc_forcing_fixture, request, tmp_path):
+def test_surface_forcing_plot(sfc_forcing_fixture, request, tmp_path):
     """Test save method."""
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
     sfc_forcing.plot(var_name="uwnd", time=0)
 
+
+@pytest.mark.parametrize(
+    "sfc_forcing_fixture",
+    [
+        "surface_forcing",
+        "corrected_surface_forcing",
+        "surface_forcing_with_wind_dropoff",
+        "coarse_surface_forcing",
+    ],
+)
+def test_surface_forcing_save(sfc_forcing_fixture, request, tmp_path):
+    """Test save method."""
+    sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
+
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             # Test saving without grouping
             saved_filenames = sfc_forcing.save(filepath, group=False)
             filepath_str = str(Path(filepath).with_suffix(""))
@@ -634,22 +676,28 @@ def test_surface_forcing_save(sfc_forcing_fixture, request, tmp_path):
             expected_filepath.unlink()
 
 
-def test_surface_forcing_bgc_plot(bgc_surface_forcing):
+@pytest.mark.parametrize(
+    "sfc_forcing_fixture",
+    [
+        "bgc_surface_forcing",
+        "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
+    ],
+)
+def test_surface_forcing_bgc_plot(sfc_forcing_fixture, request):
     """Test plot method."""
-
-    bgc_surface_forcing.plot(var_name="pco2_air", time=0)
+    sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
+    sfc_forcing.plot(var_name="pco2_air", time=0)
 
 
 def test_surface_forcing_bgc_save(bgc_surface_forcing, tmp_path):
     """Test save method."""
-
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             # Test saving without grouping
             saved_filenames = bgc_surface_forcing.save(filepath, group=False)
             filepath_str = str(Path(filepath).with_suffix(""))
@@ -676,7 +724,6 @@ def test_surface_forcing_bgc_save(bgc_surface_forcing, tmp_path):
 )
 def test_surface_forcing_bgc_from_clim_save(sfc_forcing_fixture, tmp_path, request):
     """Test save method."""
-
     bgc_surface_forcing_from_climatology = request.getfixturevalue(sfc_forcing_fixture)
 
     for file_str in ["test_sf", "test_sf.nc"]:
@@ -685,7 +732,6 @@ def test_surface_forcing_bgc_from_clim_save(sfc_forcing_fixture, tmp_path, reque
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             # Test saving without grouping
             saved_filenames = bgc_surface_forcing_from_climatology.save(
                 filepath, group=False
@@ -713,6 +759,7 @@ def test_surface_forcing_bgc_from_clim_save(sfc_forcing_fixture, tmp_path, reque
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
+        "surface_forcing_with_wind_dropoff",
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
         "bgc_surface_forcing_from_unified_climatology",
@@ -720,8 +767,8 @@ def test_surface_forcing_bgc_from_clim_save(sfc_forcing_fixture, tmp_path, reque
 )
 def test_roundtrip_yaml(sfc_forcing_fixture, request, tmp_path, use_dask):
     """Test that creating an SurfaceForcing object, saving its parameters to yaml file,
-    and re-opening yaml file creates the same object."""
-
+    and re-opening yaml file creates the same object.
+    """
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
     # Create a temporary filepath using the tmp_path fixture
@@ -730,7 +777,6 @@ def test_roundtrip_yaml(sfc_forcing_fixture, request, tmp_path, use_dask):
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         sfc_forcing.to_yaml(filepath)
 
         sfc_forcing_from_file = SurfaceForcing.from_yaml(filepath, use_dask)
@@ -746,12 +792,12 @@ def test_roundtrip_yaml(sfc_forcing_fixture, request, tmp_path, use_dask):
     [
         "surface_forcing",
         "corrected_surface_forcing",
+        "surface_forcing_with_wind_dropoff",
         "coarse_surface_forcing",
         "bgc_surface_forcing",
     ],
 )
 def test_files_have_same_hash(sfc_forcing_fixture, request, tmp_path, use_dask):
-
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
     yaml_filepath = tmp_path / "test_yaml.yaml"
@@ -787,7 +833,6 @@ def test_files_have_same_hash(sfc_forcing_fixture, request, tmp_path, use_dask):
     ],
 )
 def test_files_have_same_hash_clim(sfc_forcing_fixture, tmp_path, use_dask, request):
-
     bgc_surface_forcing_from_climatology = request.getfixturevalue(sfc_forcing_fixture)
 
     yaml_filepath = tmp_path / "test_yaml"
@@ -841,7 +886,6 @@ def test_from_yaml_missing_surface_forcing(tmp_path, use_dask):
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         # Write YAML content to file
         if isinstance(yaml_filepath, Path):
             yaml_filepath.write_text(yaml_content)
@@ -856,3 +900,35 @@ def test_from_yaml_missing_surface_forcing(tmp_path, use_dask):
             SurfaceForcing.from_yaml(yaml_filepath, use_dask=use_dask)
         yaml_filepath = Path(yaml_filepath)
         yaml_filepath.unlink()
+
+
+@pytest.mark.stream
+def test_surface_forcing_arco(surface_forcing_arco, tmp_path):
+    """One big integration test for cloud-based ERA5 data because the streaming takes a
+    long time.
+    """
+    # Test plotting
+    surface_forcing_arco.plot(var_name="uwnd", time=0)
+
+    # Roundtrip yaml
+    yaml_filepath = tmp_path / "test_yaml.yaml"
+    surface_forcing_arco.to_yaml(yaml_filepath)
+    sfc_forcing_from_yaml = SurfaceForcing.from_yaml(yaml_filepath, use_dask=True)
+    assert surface_forcing_arco == sfc_forcing_from_yaml
+
+    # Compare hashes
+    filepath1 = tmp_path / "test1.nc"
+    filepath2 = tmp_path / "test2.nc"
+    surface_forcing_arco.save(filepath1, group=True)
+    sfc_forcing_from_yaml.save(filepath2, group=True)
+    filepath_str1 = str(Path(filepath1).with_suffix(""))
+    filepath_str2 = str(Path(filepath2).with_suffix(""))
+    expected_filepath1 = f"{filepath_str1}_202002.nc"
+    expected_filepath2 = f"{filepath_str2}_202002.nc"
+    hash1 = calculate_data_hash(expected_filepath1)
+    hash2 = calculate_data_hash(expected_filepath2)
+    assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
+
+    yaml_filepath.unlink()
+    Path(expected_filepath1).unlink()
+    Path(expected_filepath2).unlink()