roms-tools 2.4.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -893,13 +893,9 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     if hasattr(forcing_object, "grid") and forcing_object.grid is not None:
         grid_data = asdict(forcing_object.grid)
         grid_yaml_data = {"Grid": _pop_grid_data(grid_data)}
-    else:
-        parent_grid_data = asdict(forcing_object.parent_grid)
-        parent_grid_yaml_data = {"ParentGrid": _pop_grid_data(parent_grid_data)}
-        child_grid_data = asdict(forcing_object.child_grid)
-        child_grid_yaml_data = {"ChildGrid": _pop_grid_data(child_grid_data)}
-
-        grid_yaml_data = {**parent_grid_yaml_data, **child_grid_yaml_data}
+    elif hasattr(forcing_object, "parent_grid"):
+        grid_data = asdict(forcing_object.parent_grid)
+        grid_yaml_data = {"ParentGrid": _pop_grid_data(grid_data)}
 
     # Step 2: Ensure Paths are Strings
     def ensure_paths_are_strings(obj, key):
@@ -933,11 +929,12 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         not in (
             "grid",
             "parent_grid",
-            "child_grid",
             "ds",
             "use_dask",
-            "bypass_validation",
             "climatology",
+            "verbose",
+            "straddle",
+            "indices",
         )
     ]
 
@@ -952,18 +949,37 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         # Add the field and its value to the forcing_data dictionary
         forcing_data[field_name] = value
 
-    # Step 5: Combine Grid and Forcing Data into a single dictionary for the final YAML content
+    # Step 5: Serialize `indices` data (conditionally)
+    indices_data = getattr(forcing_object, "indices", None)
+    if indices_data is not None:
+        serialized_indices = {
+            "_convention": "eta_rho, xi_rho"
+        }  # Add convention metadata
+
+        for key, value in indices_data.items():
+            serialized_indices[key] = [
+                f"{tup[0]}, {tup[1]}" for tup in value
+            ]  # Comma-separated string
+
+        forcing_data["indices"] = serialized_indices
+
+    # Step 6: Combine Grid and Forcing Data into a single dictionary for the final YAML content
     yaml_data = {
         **grid_yaml_data,  # Add the grid data to the final YAML structure
         forcing_object.__class__.__name__: forcing_data,  # Include the serialized forcing object data
     }
 
-    # Step 6: Write to YAML file
+    # Step 7: Write to YAML file
     with filepath.open("w") as file:
         # Write the header first
         file.write(header)
         # Write the serialized YAML data
-        yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+        yaml.dump(
+            yaml_data,
+            file,
+            default_flow_style=False,
+            sort_keys=False,
+        )
 
 
 def _pop_grid_data(grid_data):

roms_tools/tests/test_analysis/test_roms_output.py

@@ -18,7 +18,6 @@ def roms_output_from_restart_file(use_dask):
     return ROMSOutput(
         grid=grid,
         path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
-        type="restart",
         use_dask=use_dask,
     )
 
@@ -28,20 +27,6 @@ def test_load_model_output_file(roms_output_from_restart_file, use_dask):
     assert isinstance(roms_output_from_restart_file.ds, xr.Dataset)
 
 
-def test_load_model_output_directory(use_dask):
-    fname_grid = Path(download_test_data("epac25km_grd.nc"))
-    grid = Grid.from_file(fname_grid)
-
-    # Download at least two files, so these will be found within the pooch directory
-    _ = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
-    _ = Path(download_test_data("eastpac25km_rst.19980126000000.nc"))
-
-    # Directory
-    directory = os.path.dirname(download_test_data("eastpac25km_rst.19980106000000.nc"))
-    output = ROMSOutput(grid=grid, path=directory, type="restart", use_dask=use_dask)
-    assert isinstance(output.ds, xr.Dataset)
-
-
 def test_load_model_output_file_list(use_dask):
     fname_grid = Path(download_test_data("epac25km_grd.nc"))
     grid = Grid.from_file(fname_grid)
@@ -49,24 +34,23 @@ def test_load_model_output_file_list(use_dask):
     # List of files
     file1 = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
     file2 = Path(download_test_data("eastpac25km_rst.19980126000000.nc"))
-    output = ROMSOutput(
-        grid=grid, path=[file1, file2], type="restart", use_dask=use_dask
-    )
+    output = ROMSOutput(grid=grid, path=[file1, file2], use_dask=use_dask)
     assert isinstance(output.ds, xr.Dataset)
 
 
-def test_invalid_type(use_dask):
+def test_load_model_output_with_wildcard(use_dask):
     fname_grid = Path(download_test_data("epac25km_grd.nc"))
     grid = Grid.from_file(fname_grid)
 
-    # Invalid type
-    with pytest.raises(ValueError, match="Invalid type 'invalid_type'"):
-        ROMSOutput(
-            grid=grid,
-            path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
-            type="invalid_type",
-            use_dask=use_dask,
-        )
+    # Download at least two files, so these will be found within the pooch directory
+    Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
+    Path(download_test_data("eastpac25km_rst.19980126000000.nc"))
+    directory = Path(
+        os.path.dirname(download_test_data("eastpac25km_rst.19980106000000.nc"))
+    )
+
+    output = ROMSOutput(grid=grid, path=directory / "*rst*.nc", use_dask=use_dask)
+    assert isinstance(output.ds, xr.Dataset)
 
 
 def test_invalid_path(use_dask):
@@ -78,16 +62,6 @@ def test_invalid_path(use_dask):
         ROMSOutput(
             grid=grid,
             path=Path("/path/to/nonexistent/file.nc"),
-            type="restart",
-            use_dask=use_dask,
-        )
-
-    # Non-existent directory
-    with pytest.raises(FileNotFoundError):
-        ROMSOutput(
-            grid=grid,
-            path=Path("/path/to/nonexistent/directory"),
-            type="restart",
             use_dask=use_dask,
         )
 
@@ -99,7 +73,6 @@ def test_set_correct_model_reference_date(use_dask):
     output = ROMSOutput(
         grid=grid,
         path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
-        type="restart",
         use_dask=use_dask,
     )
     assert output.model_reference_date == datetime(1995, 1, 1)
@@ -117,7 +90,6 @@ def test_model_reference_date_mismatch(use_dask):
         ROMSOutput(
             grid=grid,
             path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
-            type="restart",
             model_reference_date=model_ref_date,
             use_dask=use_dask,
         )
@@ -138,7 +110,7 @@ def test_model_reference_date_no_metadata(use_dask, tmp_path, caplog):
             expected_exception,
             match="Model reference date could not be inferred from the metadata",
         ):
-            ROMSOutput(grid=grid, path=fname_mod, type="restart", use_dask=use_dask)
+            ROMSOutput(grid=grid, path=fname_mod, use_dask=use_dask)
 
         # Test case 2: When a model reference date is explicitly set, verify the warning
         with caplog.at_level(logging.WARNING):
@@ -146,7 +118,6 @@ def test_model_reference_date_no_metadata(use_dask, tmp_path, caplog):
                 grid=grid,
                 path=fname_mod,
                 model_reference_date=datetime(1995, 1, 1),
-                type="restart",
                 use_dask=use_dask,
             )
 
@@ -175,20 +146,53 @@ def test_model_reference_date_no_metadata(use_dask, tmp_path, caplog):
 def test_compute_depth_coordinates(use_dask):
     fname_grid = Path(download_test_data("epac25km_grd.nc"))
     grid = Grid.from_file(fname_grid)
+    fname_restart1 = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
+
+    for adjust_depth_for_sea_surface_height in [True, False]:
+        output = ROMSOutput(
+            grid=grid,
+            path=fname_restart1,
+            use_dask=use_dask,
+            adjust_depth_for_sea_surface_height=adjust_depth_for_sea_surface_height,
+        )
+
+        # Before calling get_vertical_coordinates, check if the dataset doesn't already have depth coordinates
+        assert "layer_depth_rho" not in output.ds_depth_coords.data_vars
+
+        # Call the method to get vertical coordinates
+        output._get_depth_coordinates(depth_type="layer")
+
+        # Check if the depth coordinates were added
+        assert "layer_depth_rho" in output.ds_depth_coords.data_vars
+
+
+def test_missing_zeta_gets_raised(use_dask):
+    """Test that a ValueError is raised when `zeta` is missing from the dataset and
+    `adjust_depth_for_sea_surface_height` is enabled."""
+    # Load the grid
+    fname_grid = Path(download_test_data("epac25km_grd.nc"))
+    grid = Grid.from_file(fname_grid)
 
+    # Load the ROMS output
     fname_restart1 = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
-    output = ROMSOutput(
-        grid=grid, path=fname_restart1, type="restart", use_dask=use_dask
+    roms_output = ROMSOutput(
+        grid=grid,
+        path=fname_restart1,
+        use_dask=use_dask,
+        adjust_depth_for_sea_surface_height=True,
     )
 
-    # Before calling get_vertical_coordinates, check if the dataset doesn't already have depth coordinates
-    assert "layer_depth_rho" not in output.ds.data_vars
-
-    # Call the method to get vertical coordinates
-    output.compute_depth_coordinates(depth_type="layer")
+    # Remove `zeta` from the dataset
+    object.__setattr__(
+        roms_output, "ds", roms_output.ds.drop_vars("zeta", errors="ignore")
+    )
 
-    # Check if the depth coordinates were added
-    assert "layer_depth_rho" in output.ds.data_vars
+    # Expect ValueError when calling `_get_depth_coordinates`
+    with pytest.raises(
+        ValueError,
+        match="`zeta` is required in provided ROMS output when `adjust_depth_for_sea_surface_height` is enabled.",
+    ):
+        roms_output._get_depth_coordinates()
 
 
 def test_check_vertical_coordinate_mismatch(use_dask):
@@ -196,9 +200,7 @@ def test_check_vertical_coordinate_mismatch(use_dask):
     grid = Grid.from_file(fname_grid)
 
     fname_restart1 = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
-    output = ROMSOutput(
-        grid=grid, path=fname_restart1, type="restart", use_dask=use_dask
-    )
+    output = ROMSOutput(grid=grid, path=fname_restart1, use_dask=use_dask)
 
     # create a mock dataset with inconsistent vertical coordinate parameters
     ds_mock = output.ds.copy()
@@ -226,9 +228,7 @@ def test_that_coordinates_are_added(use_dask):
     grid = Grid.from_file(fname_grid)
 
     fname_restart1 = Path(download_test_data("eastpac25km_rst.19980106000000.nc"))
-    output = ROMSOutput(
-        grid=grid, path=fname_restart1, type="restart", use_dask=use_dask
-    )
+    output = ROMSOutput(grid=grid, path=fname_restart1, use_dask=use_dask)
 
     assert "abs_time" in output.ds.coords
     assert "lat_rho" in output.ds.coords
@@ -239,23 +239,186 @@ def test_plot(roms_output_from_restart_file, use_dask):
 
     kwargs = {}
     for var_name in ["temp", "u", "v"]:
-        roms_output_from_restart_file.plot(var_name, time=0, s=-1, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, eta=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, xi=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, eta=0, xi=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, s=-1, eta=0, **kwargs)
+        for include_boundary in [False, True]:
+            roms_output_from_restart_file.plot(
+                var_name, time=0, s=-1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, eta=0, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, eta=1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, xi=0, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, xi=1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=0,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=0,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=1,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=1,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                eta=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                eta=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
 
     kwargs = {"depth_contours": True, "layer_contours": True}
     for var_name in ["temp", "u", "v"]:
-        roms_output_from_restart_file.plot(var_name, time=0, s=-1, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, eta=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, xi=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, eta=0, xi=0, **kwargs)
-        roms_output_from_restart_file.plot(var_name, time=0, s=-1, eta=0, **kwargs)
-
-    roms_output_from_restart_file.plot("zeta", time=0, **kwargs)
-    roms_output_from_restart_file.plot("zeta", time=0, eta=0, **kwargs)
-    roms_output_from_restart_file.plot("zeta", time=0, xi=0, **kwargs)
+        for include_boundary in [False, True]:
+            roms_output_from_restart_file.plot(
+                var_name, time=0, s=-1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, eta=0, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, eta=1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, xi=0, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name, time=0, xi=1, **kwargs, include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=0,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=0,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=1,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                eta=1,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                eta=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                eta=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                xi=0,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+            roms_output_from_restart_file.plot(
+                var_name,
+                time=0,
+                s=-1,
+                xi=1,
+                **kwargs,
+                include_boundary=include_boundary
+            )
+
+    for include_boundary in [False, True]:
+        roms_output_from_restart_file.plot(
+            "zeta", time=0, **kwargs, include_boundary=include_boundary
+        )
+        roms_output_from_restart_file.plot(
+            "zeta", time=0, eta=0, **kwargs, include_boundary=include_boundary
+        )
+        roms_output_from_restart_file.plot(
+            "zeta", time=0, eta=1, **kwargs, include_boundary=include_boundary
+        )
+        roms_output_from_restart_file.plot(
+            "zeta", time=0, xi=0, **kwargs, include_boundary=include_boundary
+        )
+        roms_output_from_restart_file.plot(
+            "zeta", time=0, xi=1, **kwargs, include_boundary=include_boundary
+        )
 
 
 def test_plot_errors(roms_output_from_restart_file, use_dask):

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -139,6 +139,69 @@ def test_unsuccessful_boundary_forcing_creation_with_1d_fill(use_dask):
         )
 
 
+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
+    with pytest.raises(
+        ValueError, match="Both `start_time` and `end_time` must be provided together"
+    ):
+        BoundaryForcing(
+            grid=None,
+            start_time=datetime(2022, 1, 1),
+            end_time=None,  # end_time is None, should raise an error
+            source={"name": "GLORYS", "path": "glorys_data.nc"},
+            use_dask=use_dask,
+        )
+
+    # Case 2: Only end_time provided
+    with pytest.raises(
+        ValueError, match="Both `start_time` and `end_time` must be provided together"
+    ):
+        BoundaryForcing(
+            grid=None,
+            start_time=None,  # start_time is None, should raise an error
+            end_time=datetime(2022, 1, 2),
+            source={"name": "GLORYS", "path": "glorys_data.nc"},
+            use_dask=use_dask,
+        )
+
+
+def test_start_time_end_time_warning(use_dask, caplog):
+    """Test that a warning is triggered when both start_time and end_time are None."""
+    # Catching the warning during test
+    grid = Grid(
+        nx=3,
+        ny=3,
+        size_x=400,
+        size_y=400,
+        center_lon=-8,
+        center_lat=58,
+        rot=0,
+        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
+    )
+
+    fname1 = Path(download_test_data("GLORYS_NA_20120101.nc"))
+    fname2 = Path(download_test_data("GLORYS_NA_20121231.nc"))
+
+    with caplog.at_level(logging.INFO):
+        BoundaryForcing(
+            grid=grid,
+            start_time=None,
+            end_time=None,
+            source={"name": "GLORYS", "path": [fname1, fname2]},
+            use_dask=use_dask,
+        )
+
+    # Verify the warning message in the log
+    assert (
+        "Both `start_time` and `end_time` are None. No time filtering will be applied to the source data."
+        in caplog.text
+    )
+
+
 def test_boundary_divided_by_land_warning(caplog, use_dask):
 
     # Iceland intersects the western boundary of the following grid

roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zattrs

@@ -1,9 +1,11 @@
 {
+    "adjust_depth_for_sea_surface_height": "False",
+    "apply_2d_horizontal_fill": "True",
     "climatology": "True",
     "end_time": "2021-06-30 00:00:00",
     "hc": 250.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev169+dirty",
+    "roms_tools_version": "2.1.1.dev3+gef240eb.d20250116",
     "source": "CESM_REGRIDDED",
     "start_time": "2021-06-29 00:00:00",
     "theta_b": 2.0,

roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zmetadata

@@ -1,11 +1,13 @@
 {
     "metadata": {
         ".zattrs": {
+            "adjust_depth_for_sea_surface_height": "False",
+            "apply_2d_horizontal_fill": "True",
             "climatology": "True",
             "end_time": "2021-06-30 00:00:00",
             "hc": 250.0,
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev169+dirty",
+            "roms_tools_version": "2.1.1.dev3+gef240eb.d20250116",
             "source": "CESM_REGRIDDED",
             "start_time": "2021-06-29 00:00:00",
             "theta_b": 2.0,