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

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -1,14 +1,17 @@
-import pytest
+import logging
+import textwrap
 from datetime import datetime
-from roms_tools import InitialConditions, Grid
-import xarray as xr
+from pathlib import Path
+
+import matplotlib.pyplot as plt
 import numpy as np
-import textwrap
-import logging
+import pytest
+import xarray as xr
+
+from conftest import calculate_data_hash
+from roms_tools import Grid, InitialConditions
 from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import CESMBGCDataset, UnifiedBGCDataset
-from pathlib import Path
-from conftest import calculate_data_hash
 
 
 @pytest.mark.parametrize(
@@ -24,7 +27,6 @@ from conftest import calculate_data_hash
 )
 def test_initial_conditions_creation(ic_fixture, request):
     """Test the creation of the InitialConditions object."""
-
     ic = request.getfixturevalue(ic_fixture)
 
     assert ic.ini_time == datetime(2021, 6, 29)
@@ -43,6 +45,43 @@ def test_initial_conditions_creation(ic_fixture, request):
     assert ic.ds.coords["ocean_time"].attrs["units"] == "seconds"
 
 
+def test_initial_conditions_creation_with_duplicates(use_dask: bool) -> None:
+    """Test the creation of the InitialConditions object with duplicates in source data
+    works as expected.
+    """
+    fname1 = Path(download_test_data("GLORYS_NA_20120101.nc"))
+    fname2 = Path(download_test_data("GLORYS_NA_20121231.nc"))
+
+    grid = Grid(
+        nx=3,
+        ny=3,
+        size_x=400,
+        size_y=400,
+        center_lon=-8,
+        center_lat=58,
+        rot=0,
+        N=3,
+    )
+
+    initial_conditions = InitialConditions(
+        grid=grid,
+        ini_time=datetime(2012, 1, 1),
+        source={"path": [fname1, fname2], "name": "GLORYS"},
+        use_dask=use_dask,
+    )
+
+    initial_conditions_with_duplicates_in_source_data = InitialConditions(
+        grid=grid,
+        ini_time=datetime(2012, 1, 1),
+        source={"path": [fname1, fname1, fname2], "name": "GLORYS"},
+        use_dask=use_dask,
+    )
+
+    assert initial_conditions.ds.identical(
+        initial_conditions_with_duplicates_in_source_data.ds
+    )
+
+
 @pytest.mark.parametrize(
     "ic_fixture",
     [
@@ -54,7 +93,6 @@ def test_initial_conditions_creation(ic_fixture, request):
 )
 def test_initial_condition_creation_with_bgc(ic_fixture, request):
     """Test the creation of the BoundaryForcing object."""
-
     ic = request.getfixturevalue(ic_fixture)
 
     expected_bgc_variables = [
@@ -98,7 +136,6 @@ def test_initial_condition_creation_with_bgc(ic_fixture, request):
 
 @pytest.fixture
 def example_grid():
-
     grid = Grid(
         nx=2,
         ny=2,
@@ -140,7 +177,6 @@ def test_initial_conditions_missing_physics_path(example_grid, use_dask):
 
 # Test initialization with missing 'name' in bgc_source
 def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
-
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     with pytest.raises(
         ValueError, match="`bgc_source` must include a 'name' if it is provided."
@@ -156,7 +192,6 @@ def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
 
 # Test initialization with missing 'path' in bgc_source
 def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
-
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     with pytest.raises(
         ValueError, match="`bgc_source` must include a 'path' if it is provided."
@@ -172,7 +207,6 @@ def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
 
 # Test initialization with missing ini_time
 def test_initial_conditions_missing_ini_time(example_grid, use_dask):
-
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     with pytest.raises(
         ValueError,
@@ -188,7 +222,6 @@ def test_initial_conditions_missing_ini_time(example_grid, use_dask):
 
 # Test default climatology value
 def test_initial_conditions_default_climatology(example_grid, use_dask):
-
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     initial_conditions = InitialConditions(
@@ -203,7 +236,6 @@ def test_initial_conditions_default_climatology(example_grid, use_dask):
 
 
 def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
-
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     fname_bgc = Path(download_test_data("CESM_regional_test_data_one_time_slice.nc"))
 
@@ -219,7 +251,6 @@ def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
 
 
 def test_info_depth(caplog, use_dask):
-
     grid = Grid(
         nx=2,
         ny=2,
@@ -237,7 +268,6 @@ def test_info_depth(caplog, use_dask):
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     with caplog.at_level(logging.INFO):
-
         InitialConditions(
             grid=grid,
             ini_time=datetime(2021, 6, 29),
@@ -252,7 +282,6 @@ def test_info_depth(caplog, use_dask):
     caplog.clear()
 
     with caplog.at_level(logging.INFO):
-
         InitialConditions(
             grid=grid,
             ini_time=datetime(2021, 6, 29),
@@ -277,7 +306,6 @@ def test_info_depth(caplog, use_dask):
 def test_correct_depth_coords_adjusted_for_zeta(
     initial_conditions_fixture, request, use_dask
 ):
-
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     # compute interface depth at rho-points and write it into .ds_depth_coords
@@ -315,7 +343,6 @@ def test_correct_depth_coords_adjusted_for_zeta(
     ],
 )
 def test_correct_depth_coords_zero_zeta(initial_conditions_fixture, request, use_dask):
-
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     # compute interface depth at rho-points and write it into .ds_depth_coords
@@ -342,7 +369,6 @@ def test_correct_depth_coords_zero_zeta(initial_conditions_fixture, request, use
 
 
 def test_interpolation_from_climatology(use_dask):
-
     # CESM climatology
     fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
     ds = xr.open_dataset(fname_bgc)
@@ -401,6 +427,21 @@ def test_interpolation_from_climatology(use_dask):
     )
 
 
+def test_computed_missing_optional_fields(
+    initial_conditions_with_unified_bgc_from_climatology,
+):
+    ds = initial_conditions_with_unified_bgc_from_climatology.ds
+
+    # Use tight tolerances because 'DOC' and 'DOCr' can have values order 1e-6
+
+    # 'DOCr' was missing in the source data and should have been filled with a constant default value
+    assert np.allclose(ds.DOCr.std(), 0.0, rtol=1e-10, atol=1e-10), (
+        "DOCr should be constant across space and time"
+    )
+    # 'DOC' was present in the source data and should show spatial or temporal variability
+    assert ds.DOC.std() > 1e-10, "DOC should vary across space and time"
+
+
 @pytest.mark.parametrize(
     "initial_conditions_fixture",
     [
@@ -411,7 +452,6 @@ def test_interpolation_from_climatology(use_dask):
 )
 def test_initial_conditions_plot(initial_conditions_fixture, request):
     """Test plot method."""
-
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     # horizontal slices plots
@@ -441,6 +481,13 @@ def test_initial_conditions_plot(initial_conditions_fixture, request):
     initial_conditions.plot(var_name="ubar")
     initial_conditions.plot(var_name="vbar")
 
+    # Test that passing a matplotlib.axes.Axes works
+    fig, ax = plt.subplots(1, 1)
+    initial_conditions.plot(var_name="temp", s=0, ax=ax)
+    initial_conditions.plot(var_name="temp", eta=0, ax=ax)
+    initial_conditions.plot(var_name="temp", s=0, xi=0, ax=ax)
+    initial_conditions.plot(var_name="zeta", ax=ax)
+
 
 @pytest.mark.parametrize(
     "initial_conditions_fixture",
@@ -453,7 +500,6 @@ def test_initial_conditions_plot(initial_conditions_fixture, request):
 )
 def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
     """Test save method."""
-
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     for file_str in ["test_ic", "test_ic.nc"]:
@@ -462,7 +508,6 @@ def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             saved_filenames = initial_conditions.save(filepath)
             # Check if the .nc file was created
             filepath = Path(filepath).with_suffix(".nc")
@@ -483,8 +528,8 @@ def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
 )
 def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask):
     """Test that creating an InitialConditions object, saving its parameters to yaml
-    file, and re-opening yaml file creates the same object."""
-
+    file, and re-opening yaml file creates the same object.
+    """
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     # Create a temporary filepath using the tmp_path fixture
@@ -493,7 +538,6 @@ def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask)
         tmp_path / file_str,
         str(tmp_path / file_str),
     ]:  # test for Path object and str
-
         initial_conditions.to_yaml(filepath)
 
         initial_conditions_from_file = InitialConditions.from_yaml(
@@ -516,7 +560,6 @@ def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask)
     ],
 )
 def test_files_have_same_hash(initial_conditions_fixture, request, tmp_path, use_dask):
-
     initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     yaml_filepath = tmp_path / "test_yaml.yaml"
@@ -565,7 +608,6 @@ def test_from_yaml_missing_initial_conditions(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)

roms_tools/tests/test_setup/test_nesting.py

@@ -1,17 +1,19 @@
-import pytest
-import xarray as xr
-import numpy as np
 import logging
 from pathlib import Path
-from roms_tools import Grid, ChildGrid
-from roms_tools.setup.utils import get_boundary_coords
+
+import numpy as np
+import pytest
+import xarray as xr
+
 from conftest import calculate_file_hash
+from roms_tools import ChildGrid, Grid
 from roms_tools.setup.nesting import (
+    compute_boundary_distance,
     interpolate_indices,
     map_child_boundaries_onto_parent_grid_indices,
-    compute_boundary_distance,
     modify_child_topography_and_mask,
 )
+from roms_tools.setup.utils import get_boundary_coords
 
 
 @pytest.fixture()
@@ -87,7 +89,6 @@ class TestInterpolateIndices:
     )
     def test_correct_indices_of_same_grid(self, grid, caplog, request):
         """Verify boundary indices are correctly interpolated for the same grid."""
-
         grid = request.getfixturevalue(grid)
 
         bdry_coords_dict = get_boundary_coords()
@@ -138,7 +139,6 @@ class TestInterpolateIndices:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Ensure interpolated indices fall within the parent grid's bounds."""
-
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -166,7 +166,6 @@ class TestInterpolateIndices:
 class TestMapChildBoundaries:
     def test_update_indices_does_nothing_if_no_parent_land(self, small_grid, baby_grid):
         """Verify no change in indices when parent grid has no land at boundaries."""
-
         ds_without_updated_indices = map_child_boundaries_onto_parent_grid_indices(
             small_grid.ds, baby_grid.ds, update_land_indices=False
         )
@@ -187,7 +186,6 @@ class TestMapChildBoundaries:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Check updated indices map to wet points on the parent grid."""
-
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -221,8 +219,8 @@ class TestMapChildBoundaries:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Test that child boundary indices are monotonically increasing or decreasing
-        in both the xi and eta directions, for all boundaries and locations."""
-
+        in both the xi and eta directions, for all boundaries and locations.
+        """
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -260,7 +258,6 @@ class TestBoundaryDistance:
         self, grid_fixture, request
     ):
         """Ensure boundary distance is zero for grids without land along boundaries."""
-
         grid = request.getfixturevalue(grid_fixture)
         alpha = compute_boundary_distance(grid.ds.mask_rho)
 
@@ -280,7 +277,8 @@ class TestBoundaryDistance:
 
     def test_boundary_distance_for_grid_with_land_along_boundary(self, big_grid):
         """Test that there are 1s along the boundary of alpha if the grid has land along
-        the boundary."""
+        the boundary.
+        """
         alpha = compute_boundary_distance(big_grid.ds.mask_rho)
         assert (alpha.isel(eta_rho=0) == 1).any()
         assert (alpha.isel(eta_rho=-1) == 1).any()
@@ -293,7 +291,6 @@ class TestModifyChid:
         self, small_grid, baby_grid
     ):
         """Confirm child mask remains unchanged if no parent land is at boundaries."""
-
         mask_original = baby_grid.ds.mask_rho.copy()
         modified_baby_grid_ds = modify_child_topography_and_mask(
             small_grid.ds, baby_grid.ds
@@ -310,7 +307,6 @@ class TestModifyChid:
     )
     def test_no_modification_if_parent_and_child_coincide(self, grid_fixture, request):
         """Ensure no changes occur when parent and child grids coincide."""
-
         grid = request.getfixturevalue(grid_fixture)
 
         h_original = grid.ds.h.copy()
@@ -322,8 +318,8 @@ class TestModifyChid:
 
     def test_modification_only_along_boundaries(self, big_grid, small_grid):
         """Test that modifications to the child grid's topography and mask occur only
-        along the boundaries, leaving the interior unchanged."""
-
+        along the boundaries, leaving the interior unchanged.
+        """
         # Make copies of original data for comparison
         h_original = small_grid.ds.h.copy()
         mask_original = small_grid.ds.mask_rho.copy()
@@ -338,15 +334,15 @@ class TestModifyChid:
         # Assert that the center values remain the same
         assert mask_original.isel(
             eta_rho=eta_center, xi_rho=xi_center
-        ) == modified_ds.mask_rho.isel(
-            eta_rho=eta_center, xi_rho=xi_center
-        ), "Mask at the grid center was modified."
+        ) == modified_ds.mask_rho.isel(eta_rho=eta_center, xi_rho=xi_center), (
+            "Mask at the grid center was modified."
+        )
 
         assert h_original.isel(
             eta_rho=eta_center, xi_rho=xi_center
-        ) == modified_ds.h.isel(
-            eta_rho=eta_center, xi_rho=xi_center
-        ), "Topography at the grid center was modified."
+        ) == modified_ds.h.isel(eta_rho=eta_center, xi_rho=xi_center), (
+            "Topography at the grid center was modified."
+        )
 
 
 class TestNesting:
@@ -426,11 +422,9 @@ class TestNesting:
 
     def test_save(self, child_grid, tmp_path):
         """Test save methods."""
-
         for file_str in ["test_grid", "test_grid.nc"]:
             # Create a temporary filepath using the tmp_path fixture
             for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
-
                 saved_filenames = child_grid.save(filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -442,7 +436,6 @@ class TestNesting:
         for file_str in ["test_nesting", "test_nesting.nc"]:
             # Create a temporary filepath using the tmp_path fixture
             for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
-
                 saved_filenames = child_grid.save_nesting(filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -461,15 +454,14 @@ class TestNesting:
 
     def test_roundtrip_yaml(self, child_grid, tmp_path):
         """Test that creating a ChildGrid 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.
+        """
         # Create a temporary filepath using the tmp_path fixture
         file_str = "test_yaml"
         for filepath in [
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             child_grid.to_yaml(filepath)
 
             child_grid_from_file = ChildGrid.from_yaml(filepath)
@@ -480,7 +472,6 @@ class TestNesting:
             filepath.unlink()
 
     def test_files_have_same_hash(self, child_grid, tmp_path):
-
         yaml_filepath = tmp_path / "test_yaml.yaml"
         filepath1 = tmp_path / "test1.nc"
         filepath2 = tmp_path / "test2.nc"