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

roms_tools/setup/utils.py

@@ -3,7 +3,6 @@ import numpy as np
 from typing import Union, Any, Dict, Type
 import pandas as pd
 import cftime
-from roms_tools.utils import partition
 from pathlib import Path
 from datetime import datetime
 from dataclasses import fields, asdict
@@ -131,23 +130,36 @@ def interpolate_from_climatology(
     time_dim_name: str,
     time: Union[xr.DataArray, pd.DatetimeIndex],
 ) -> Union[xr.DataArray, xr.Dataset]:
-    """Interpolates the given field temporally based on the specified time points.
+    """Temporally interpolates a field based on specified time points.
 
-    If `field` is an xarray.Dataset, this function applies the interpolation to all data variables in the dataset.
+    This function performs temporal interpolation on the input `field` to match the provided `time` values.
+    If the input `field` is an `xarray.Dataset`, the interpolation is applied to all its data variables individually.
 
     Parameters
     ----------
     field : xarray.DataArray or xarray.Dataset
-        The field data to be interpolated. Can be a single DataArray or a Dataset.
+        The input field to be interpolated.
+        - If `field` is an `xarray.DataArray`, it should have a time dimension identified by `time_dim_name`.
+        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along the specified time dimension.
+        The time dimension is assumed to represent `day_of_year` for climatological purposes.
     time_dim_name : str
-        The name of the dimension in `field` that represents time.
+        The name of the time dimension in the `field`. This dimension is used for interpolation.
     time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation.
+        The target time points for interpolation. The time values should be compatible with the time format used in the `field`.
 
     Returns
     -------
     xarray.DataArray or xarray.Dataset
-        The field values interpolated to the specified time points. The type matches the input type.
+        The interpolated field, with the same type as the input (`xarray.DataArray` or `xarray.Dataset`),
+        but aligned to the specified `time` values.
+
+    Notes
+    -----
+    - The interpolation assumes the time dimension in `field` corresponds to `day_of_year`.
+      If the input time values are in a datetime format, ensure they are converted to `day_of_year` before calling this function.
+      For example, you can preprocess the time as follows:
+
+      >>> field["time"] = field["time"].dt.dayofyear
     """
 
     def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
@@ -157,11 +169,11 @@ def interpolate_from_climatology(
             day_of_year = time.dt.dayofyear
         else:
             if np.size(time) == 1:
-                day_of_year = time.timetuple().tm_yday
+                # Convert single datetime64 object to pandas.Timestamp
+                day_of_year = pd.Timestamp(time).dayofyear
             else:
-                day_of_year = np.array([t.timetuple().tm_yday for t in time])
-
-        data_array[time_dim_name] = data_array[time_dim_name].dt.days
+                # Convert each datetime64 object in the array to pandas.Timestamp
+                day_of_year = np.array([pd.Timestamp(t).dayofyear for t in time])
 
         # Concatenate across the beginning and end of the year
         time_concat = xr.concat(
@@ -201,6 +213,7 @@ def interpolate_from_climatology(
             for var, data_array in field.data_vars.items()
         }
         return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
+
     else:
         raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
 
@@ -568,59 +581,6 @@ def group_by_year(ds, filepath):
     return dataset_list, output_filenames
 
 
-def save_datasets(dataset_list, output_filenames, np_eta=None, np_xi=None):
-    """Save the list of datasets to netCDF4 files, with optional spatial partitioning.
-
-    Parameters
-    ----------
-    dataset_list : list
-        List of datasets to be saved.
-    output_filenames : list
-        List of filenames for the output files.
-    np_eta : int, optional
-        The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
-    np_xi : int, optional
-        The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
-
-    Returns
-    -------
-    List[Path]
-        A list of Path objects for the filenames that were saved.
-    """
-
-    saved_filenames = []
-
-    if np_eta is None and np_xi is None:
-        # Save the dataset as a single file
-        output_filenames = [f"{filename}.nc" for filename in output_filenames]
-        xr.save_mfdataset(dataset_list, output_filenames)
-
-        saved_filenames.extend(Path(f) for f in output_filenames)
-
-    else:
-        # Partition the dataset and save each partition as a separate file
-        np_eta = np_eta or 1
-        np_xi = np_xi or 1
-
-        partitioned_datasets = []
-        partitioned_filenames = []
-        for dataset, base_filename in zip(dataset_list, output_filenames):
-            partition_indices, partitions = partition(
-                dataset, np_eta=np_eta, np_xi=np_xi
-            )
-            partition_filenames = [
-                f"{base_filename}.{index}.nc" for index in partition_indices
-            ]
-            partitioned_datasets.extend(partitions)
-            partitioned_filenames.extend(partition_filenames)
-
-        xr.save_mfdataset(partitioned_datasets, partitioned_filenames)
-
-        saved_filenames.extend(Path(f) for f in partitioned_filenames)
-
-    return saved_filenames
-
-
 def get_target_coords(grid, use_coarse_grid=False):
     """Retrieves longitude and latitude coordinates from the grid, adjusting them based
     on longitude range.
@@ -652,8 +612,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         mask = grid.ds.get("mask_coarse")
         if mask is not None:
             mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-            mask_u = interpolate_from_rho_to_u(mask, method="multiplicative")
-            mask_v = interpolate_from_rho_to_v(mask, method="multiplicative")
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
@@ -663,8 +621,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         lon = grid.ds.lon_rho
         angle = grid.ds.angle
         mask = grid.ds.get("mask_rho")
-        mask_u = grid.ds.get("mask_u")
-        mask_v = grid.ds.get("mask_v")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -687,8 +643,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lon_psi": lon_psi,
         "angle": angle,
         "mask": mask,
-        "mask_u": mask_u,
-        "mask_v": mask_v,
         "straddle": straddle,
     }
 

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -1,10 +1,11 @@
 import pytest
 from datetime import datetime
 import xarray as xr
+import numpy as np
 from roms_tools import Grid, BoundaryForcing
 import textwrap
 from roms_tools.download import download_test_data
-from conftest import calculate_file_hash
+from conftest import calculate_data_hash
 from pathlib import Path
 import logging
 
@@ -13,19 +14,23 @@ import logging
     "boundary_forcing_fixture",
     [
         "boundary_forcing",
-        # "boundary_forcing_with_2d_fill",
+        "boundary_forcing_adjusted_for_zeta",
+        "boundary_forcing_with_2d_fill",
+        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
     ],
 )
 def test_boundary_forcing_creation(boundary_forcing_fixture, request):
     """Test the creation of the BoundaryForcing object."""
 
-    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
-    assert boundary_forcing.start_time == datetime(2021, 6, 29)
-    assert boundary_forcing.end_time == datetime(2021, 6, 30)
+
+    fname1 = Path(download_test_data("GLORYS_NA_20120101.nc"))
+    fname2 = Path(download_test_data("GLORYS_NA_20121231.nc"))
+    assert boundary_forcing.start_time == datetime(2012, 1, 1)
+    assert boundary_forcing.end_time == datetime(2012, 12, 31)
     assert boundary_forcing.source == {
         "name": "GLORYS",
-        "path": fname,
+        "path": [fname1, fname2],
         "climatology": False,
     }
     assert boundary_forcing.model_reference_date == datetime(2000, 1, 1)
@@ -44,16 +49,17 @@ def test_boundary_forcing_creation(boundary_forcing_fixture, request):
         assert f"v_{direction}" in boundary_forcing.ds
         assert f"zeta_{direction}" in boundary_forcing.ds
 
-    assert len(boundary_forcing.ds.bry_time) == 1
+    assert len(boundary_forcing.ds.bry_time) == 2
     assert boundary_forcing.ds.coords["bry_time"].attrs["units"] == "days"
     assert not hasattr(boundary_forcing.ds, "climatology")
+    assert hasattr(boundary_forcing.ds, "adjust_depth_for_sea_surface_height")
+    assert hasattr(boundary_forcing.ds, "apply_2d_horizontal_fill")
 
 
 @pytest.mark.parametrize(
     "boundary_forcing_fixture",
     [
         "bgc_boundary_forcing_from_climatology",
-        # "bgc_boundary_forcing_from_climatology_with_2d_fill",
     ],
 )
 def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
@@ -155,6 +161,113 @@ def test_boundary_divided_by_land_warning(caplog, use_dask):
     assert "the western boundary is divided by land" in caplog.text
 
 
+def test_info_depth(caplog, use_dask):
+
+    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=datetime(2012, 1, 1),
+            end_time=datetime(2012, 12, 31),
+            source={"name": "GLORYS", "path": [fname1, fname2]},
+            adjust_depth_for_sea_surface_height=True,
+            use_dask=use_dask,
+        )
+
+    # Verify the warning message in the log
+    assert "Sea surface height will be used to adjust depth coordinates." in caplog.text
+
+    # Clear the log before the next test
+    caplog.clear()
+
+    with caplog.at_level(logging.INFO):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2012, 1, 1),
+            end_time=datetime(2012, 12, 31),
+            source={"name": "GLORYS", "path": [fname1, fname2]},
+            adjust_depth_for_sea_surface_height=False,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert (
+        "Sea surface height will NOT be used to adjust depth coordinates."
+        in caplog.text
+    )
+
+
+def test_info_fill(caplog, use_dask):
+
+    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=datetime(2012, 1, 1),
+            end_time=datetime(2012, 12, 31),
+            source={"name": "GLORYS", "path": [fname1, fname2]},
+            apply_2d_horizontal_fill=True,
+            use_dask=use_dask,
+        )
+
+    # Verify the warning message in the log
+    assert (
+        "Applying 2D horizontal fill to the source data before regridding."
+        in caplog.text
+    )
+
+    # Clear the log before the next test
+    caplog.clear()
+
+    with caplog.at_level(logging.INFO):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2012, 1, 1),
+            end_time=datetime(2012, 12, 31),
+            source={"name": "GLORYS", "path": [fname1, fname2]},
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert (
+        "Applying 1D horizontal fill separately to each regridded boundary."
+        in caplog.text
+    )
+
+
 def test_1d_and_2d_fill_coincide_if_no_land(use_dask):
 
     # this grid lies entirely over open ocean
@@ -182,23 +295,95 @@ def test_1d_and_2d_fill_coincide_if_no_land(use_dask):
     xr.testing.assert_allclose(bf_1d_fill.ds, bf_2d_fill.ds, rtol=1.0e-4)
 
 
-def test_boundary_forcing_plot(boundary_forcing):
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "boundary_forcing_adjusted_for_zeta",
+        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
+    ],
+)
+def test_correct_depth_coords_adjusted_for_zeta(
+    boundary_forcing_fixture, request, use_dask
+):
+
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
+
+    for direction in ["south", "east", "north", "west"]:
+
+        # Test that uppermost interface coincides with sea surface height
+        assert np.allclose(
+            boundary_forcing.ds_depth_coords[f"interface_depth_rho_{direction}"]
+            .isel(s_w=-1)
+            .values,
+            -boundary_forcing.ds[f"zeta_{direction}"].values,
+            atol=1e-6,
+        )
+
+
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "boundary_forcing",
+        "boundary_forcing_with_2d_fill",
+    ],
+)
+def test_correct_depth_coords_zero_zeta(boundary_forcing_fixture, request, use_dask):
+
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
+
+    for direction in ["south", "east", "north", "west"]:
+
+        # Test that uppermost interface coincides with sea surface height
+        assert np.allclose(
+            boundary_forcing.ds_depth_coords[f"interface_depth_rho_{direction}"]
+            .isel(s_w=-1)
+            .values,
+            0 * boundary_forcing.ds[f"zeta_{direction}"].values,
+            atol=1e-6,
+        )
+
+
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "boundary_forcing",
+        "boundary_forcing_with_2d_fill",
+        "boundary_forcing_adjusted_for_zeta",
+        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
+    ],
+)
+def test_boundary_forcing_plot(boundary_forcing_fixture, request):
     """Test plot."""
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
-    boundary_forcing.plot(var_name="temp_south", layer_contours=True)
-    boundary_forcing.plot(var_name="temp_east", layer_contours=True)
-    boundary_forcing.plot(var_name="temp_north", layer_contours=True)
-    boundary_forcing.plot(var_name="temp_west", layer_contours=True)
-    boundary_forcing.plot(var_name="zeta_south")
-    boundary_forcing.plot(var_name="zeta_east")
-    boundary_forcing.plot(var_name="zeta_north")
-    boundary_forcing.plot(var_name="zeta_west")
-    boundary_forcing.plot(var_name="vbar_north")
-    boundary_forcing.plot(var_name="ubar_west")
+    for direction in ["south", "east", "north", "west"]:
+        for layer_contours in [False, True]:
+            boundary_forcing.plot(
+                var_name=f"temp_{direction}", layer_contours=layer_contours
+            )
+            boundary_forcing.plot(
+                var_name=f"u_{direction}", layer_contours=layer_contours
+            )
+            boundary_forcing.plot(
+                var_name=f"v_{direction}", layer_contours=layer_contours
+            )
+        boundary_forcing.plot(var_name=f"zeta_{direction}")
+        boundary_forcing.plot(var_name=f"vbar_{direction}")
+        boundary_forcing.plot(var_name=f"ubar_{direction}")
 
 
-def test_boundary_forcing_save(boundary_forcing, tmp_path):
+@pytest.mark.parametrize(
+    "boundary_forcing_fixture",
+    [
+        "boundary_forcing",
+        "boundary_forcing_with_2d_fill",
+        "boundary_forcing_adjusted_for_zeta",
+        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
+    ],
+)
+def test_boundary_forcing_save(boundary_forcing_fixture, request, tmp_path):
     """Test save method."""
+    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
     for file_str in ["test_bf", "test_bf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -207,8 +392,8 @@ def test_boundary_forcing_save(boundary_forcing, tmp_path):
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning and grouping
-            saved_filenames = boundary_forcing.save(filepath)
+            # Test saving without grouping
+            saved_filenames = boundary_forcing.save(filepath, group=False)
 
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}.nc")
@@ -217,28 +402,16 @@ def test_boundary_forcing_save(boundary_forcing, tmp_path):
             assert expected_filepath.exists()
             expected_filepath.unlink()
 
-            # Test saving without partitioning but with grouping
+            # Test saving 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")
+            expected_filepath = Path(f"{filepath_str}_2012.nc")
 
             assert saved_filenames == [expected_filepath]
             assert expected_filepath.exists()
             expected_filepath.unlink()
 
-            # 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)
-            ]
-
-            assert saved_filenames == expected_filepath_list
-
-            for expected_filepath in expected_filepath_list:
-                assert expected_filepath.exists()
-                expected_filepath.unlink()
-
 
 def test_bgc_boundary_forcing_plot(bgc_boundary_forcing_from_climatology):
     """Test plot method."""
@@ -264,7 +437,9 @@ def test_bgc_boundary_forcing_save(bgc_boundary_forcing_from_climatology, tmp_pa
         ]:  # test for Path object and str
 
             # Test saving without partitioning and grouping
-            saved_filenames = bgc_boundary_forcing_from_climatology.save(filepath)
+            saved_filenames = bgc_boundary_forcing_from_climatology.save(
+                filepath, group=False
+            )
 
             filepath_str = str(Path(filepath).with_suffix(""))
             expected_filepath = Path(f"{filepath_str}.nc")
@@ -283,20 +458,6 @@ def test_bgc_boundary_forcing_save(bgc_boundary_forcing_from_climatology, tmp_pa
             assert expected_filepath.exists()
             expected_filepath.unlink()
 
-            # Test saving with partitioning
-            saved_filenames = bgc_boundary_forcing_from_climatology.save(
-                filepath, np_xi=2, group=True
-            )
-
-            expected_filepath_list = [
-                Path(filepath_str + f"_clim.{index}.nc") for index in range(2)
-            ]
-            assert saved_filenames == expected_filepath_list
-
-            for expected_filepath in expected_filepath_list:
-                assert expected_filepath.exists()
-                expected_filepath.unlink()
-
 
 @pytest.mark.parametrize(
     "bdry_forcing_fixture",
@@ -341,11 +502,12 @@ def test_files_have_same_hash(boundary_forcing, tmp_path, use_dask):
 
     filepath_str1 = str(Path(filepath1).with_suffix(""))
     filepath_str2 = str(Path(filepath2).with_suffix(""))
-    expected_filepath1 = f"{filepath_str1}_202106.nc"
-    expected_filepath2 = f"{filepath_str2}_202106.nc"
+    expected_filepath1 = f"{filepath_str1}_2012.nc"
+    expected_filepath2 = f"{filepath_str2}_2012.nc"
 
-    hash1 = calculate_file_hash(expected_filepath1)
-    hash2 = calculate_file_hash(expected_filepath2)
+    # Only compare hash of datasets because metadata is non-deterministic with dask
+    hash1 = calculate_data_hash(expected_filepath1)
+    hash2 = calculate_data_hash(expected_filepath2)
 
     assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
 
@@ -372,8 +534,9 @@ def test_files_have_same_hash_clim(
     expected_filepath1 = f"{filepath_str1}_clim.nc"
     expected_filepath2 = f"{filepath_str2}_clim.nc"
 
-    hash1 = calculate_file_hash(expected_filepath1)
-    hash2 = calculate_file_hash(expected_filepath2)
+    # Only compare hash of datasets because metadata is non-deterministic with dask
+    hash1 = calculate_data_hash(expected_filepath1)
+    hash2 = calculate_data_hash(expected_filepath2)
 
     assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
 

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

@@ -1,10 +1,12 @@
 {
-    "end_time": "2021-06-30 00:00:00",
+    "adjust_depth_for_sea_surface_height": "False",
+    "apply_2d_horizontal_fill": "False",
+    "end_time": "2012-12-31 00:00:00",
     "hc": 250.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev157+dirty",
+    "roms_tools_version": "0.1.dev177",
     "source": "GLORYS",
-    "start_time": "2021-06-29 00:00:00",
+    "start_time": "2012-01-01 00:00:00",
     "theta_b": 2.0,
     "theta_s": 5.0,
     "title": "ROMS boundary forcing file created by ROMS-Tools"