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

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

@@ -2,7 +2,7 @@
     "correct_radiation": "False",
     "end_time": "2020-02-02 00:00:00",
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev152+dirty",
+    "roms_tools_version": "0.1.dev157+dirty",
     "source": "ERA5",
     "start_time": "2020-01-31 00:00:00",
     "title": "ROMS surface forcing file created by ROMS-Tools",

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

@@ -4,7 +4,7 @@
             "correct_radiation": "False",
             "end_time": "2020-02-02 00:00:00",
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev152+dirty",
+            "roms_tools_version": "0.1.dev157+dirty",
             "source": "ERA5",
             "start_time": "2020-01-31 00:00:00",
             "title": "ROMS surface forcing file created by ROMS-Tools",

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

@@ -1,7 +1,7 @@
 {
     "allan_factor": 2.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev152+dirty",
+    "roms_tools_version": "0.1.dev157+dirty",
     "source": "TPXO",
     "title": "ROMS tidal forcing created by ROMS-Tools"
 }

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

@@ -3,7 +3,7 @@
         ".zattrs": {
             "allan_factor": 2.0,
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev152+dirty",
+            "roms_tools_version": "0.1.dev157+dirty",
             "source": "TPXO",
             "title": "ROMS tidal forcing created by ROMS-Tools"
         },
@@ -33,7 +33,9 @@
         "omega/.zattrs": {
             "_ARRAY_DIMENSIONS": [
                 "ntides"
-            ]
+            ],
+            "long_name": "Tidal frequencies",
+            "units": "1/s"
         },
         "pot_Im/.zarray": {
             "chunks": [

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/omega/.zattrs

@@ -1,5 +1,7 @@
 {
     "_ARRAY_DIMENSIONS": [
         "ntides"
-    ]
+    ],
+    "long_name": "Tidal frequencies",
+    "units": "1/s"
 }

roms_tools/tests/test_setup/test_datasets.py

@@ -259,6 +259,42 @@ def test_warnings_times(global_dataset, tmp_path, use_dask):
         )
 
 
+def test_from_ds(global_dataset, global_dataset_with_noon_times, use_dask, tmp_path):
+    """Test the from_ds method of the Dataset class."""
+
+    start_time = datetime(2022, 1, 1)
+
+    filepath = tmp_path / "test.nc"
+    global_dataset.to_netcdf(filepath)
+
+    dataset = Dataset(
+        filename=filepath,
+        var_names={"var": "var"},
+        dim_names={
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "time": "time",
+            "depth": "depth",
+        },
+        start_time=start_time,
+        use_dask=use_dask,
+    )
+
+    new_dataset = Dataset.from_ds(dataset, global_dataset_with_noon_times)
+
+    assert isinstance(new_dataset, Dataset)  # Check that the new instance is a Dataset
+    assert new_dataset.ds.equals(
+        global_dataset_with_noon_times
+    )  # Verify the new ds attribute is set correctly
+
+    # Ensure other attributes are copied correctly
+    for attr in vars(dataset):
+        if attr != "ds":
+            assert getattr(new_dataset, attr) == getattr(
+                dataset, attr
+            ), f"Attribute {attr} not copied correctly."
+
+
 def test_reverse_latitude_reverse_depth_choose_subdomain(
     global_dataset, tmp_path, use_dask
 ):
@@ -288,31 +324,37 @@ def test_reverse_latitude_reverse_depth_choose_subdomain(
     assert np.all(np.diff(dataset.ds["depth"]) > 0)
 
     # test choosing subdomain for domain that straddles the dateline
-    ds = dataset.choose_subdomain(
-        latitude_range=(-10, 10),
-        longitude_range=(-10, 10),
-        margin=1,
-        straddle=True,
-        return_subdomain=True,
+    target_coords = {
+        "lat": xr.DataArray(np.linspace(-10, 10, 100)),
+        "lon": xr.DataArray(np.linspace(-10, 10, 100)),
+        "straddle": True,
+    }
+    sub_dataset = dataset.choose_subdomain(
+        target_coords,
+        buffer_points=1,
+        return_copy=True,
     )
 
-    assert -11 <= ds["latitude"].min() <= 11
-    assert -11 <= ds["latitude"].max() <= 11
-    assert -11 <= ds["longitude"].min() <= 11
-    assert -11 <= ds["longitude"].max() <= 11
-
-    ds = dataset.choose_subdomain(
-        latitude_range=(-10, 10),
-        longitude_range=(10, 20),
-        margin=1,
-        straddle=False,
-        return_subdomain=True,
+    assert -11 <= sub_dataset.ds["latitude"].min() <= 11
+    assert -11 <= sub_dataset.ds["latitude"].max() <= 11
+    assert -11 <= sub_dataset.ds["longitude"].min() <= 11
+    assert -11 <= sub_dataset.ds["longitude"].max() <= 11
+
+    target_coords = {
+        "lat": xr.DataArray(np.linspace(-10, 10, 100)),
+        "lon": xr.DataArray(np.linspace(10, 20, 100)),
+        "straddle": False,
+    }
+    sub_dataset = dataset.choose_subdomain(
+        target_coords,
+        buffer_points=1,
+        return_copy=True,
     )
 
-    assert -11 <= ds["latitude"].min() <= 11
-    assert -11 <= ds["latitude"].max() <= 11
-    assert 9 <= ds["longitude"].min() <= 21
-    assert 9 <= ds["longitude"].max() <= 21
+    assert -11 <= sub_dataset.ds["latitude"].min() <= 11
+    assert -11 <= sub_dataset.ds["latitude"].max() <= 11
+    assert 9 <= sub_dataset.ds["longitude"].min() <= 21
+    assert 9 <= sub_dataset.ds["longitude"].max() <= 21
 
 
 def test_check_if_global_with_global_dataset(global_dataset, tmp_path, use_dask):
@@ -475,3 +517,19 @@ def test_climatology_error(use_dask):
             climatology=False,
             use_dask=use_dask,
         )
+
+
+@pytest.mark.parametrize(
+    "data_fixture, expected_resolution",
+    [
+        ("era5_data", 0.25),
+        ("glorys_data", 1 / 12),
+        ("tpxo_data", 1 / 6),
+        ("cesm_bgc_data", 1.0),
+        ("cesm_surface_bgc_data", 1.0),
+    ],
+)
+def test_horizontal_resolution(data_fixture, expected_resolution, request):
+
+    data = request.getfixturevalue(data_fixture)
+    assert np.isclose(data.resolution, expected_resolution)

roms_tools/tests/test_setup/test_fill.py

@@ -1,134 +1,58 @@
 import pytest
-from roms_tools.setup.datasets import (
-    GLORYSDataset,
-    ERA5Dataset,
-    CESMBGCDataset,
-    CESMBGCSurfaceForcingDataset,
-    TPXODataset,
-)
-from roms_tools.setup.download import download_test_data
 from roms_tools.setup.fill import LateralFill
-from roms_tools.setup.utils import extrapolate_deepest_to_bottom
-from datetime import datetime
 import numpy as np
 import xarray as xr
 
 
-@pytest.fixture()
-def era5_data(request, use_dask):
-    fname = download_test_data("ERA5_regional_test_data.nc")
-    data = ERA5Dataset(
-        filename=fname,
-        start_time=datetime(2020, 1, 31),
-        end_time=datetime(2020, 2, 2),
-        use_dask=use_dask,
-    )
-    data.post_process()
-
-    return data
-
-
-@pytest.fixture()
-def glorys_data(request, use_dask):
-    fname = download_test_data("GLORYS_NA_2012.nc")
-
-    data = GLORYSDataset(
-        filename=fname,
-        start_time=datetime(2012, 1, 1),
-        end_time=datetime(2013, 1, 1),
-        use_dask=use_dask,
-    )
-    data.post_process()
-
-    # extrapolate deepest value to bottom so all levels can use the same surface mask
-    for var in data.var_names:
-        if var != "zeta":
-            data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
-                data.ds[data.var_names[var]], data.dim_names["depth"]
-            )
-
-    return data
-
-
-@pytest.fixture()
-def tpxo_data(request, use_dask):
-    fname = download_test_data("TPXO_regional_test_data.nc")
-
-    data = TPXODataset(
-        filename=fname,
-        use_dask=use_dask,
-    )
-    data.post_process()
-
-    return data
-
-
-@pytest.fixture()
-def cesm_bgc_data(request, use_dask):
-    fname = download_test_data("CESM_BGC_2012.nc")
-
-    data = CESMBGCDataset(
-        filename=fname,
-        start_time=datetime(2012, 1, 1),
-        end_time=datetime(2013, 1, 1),
-        climatology=False,
-        use_dask=use_dask,
-    )
-    data.post_process()
-
-    # extrapolate deepest value to bottom so all levels can use the same surface mask
-    for var in data.var_names:
-        data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
-            data.ds[data.var_names[var]], data.dim_names["depth"]
-        )
-
-    return data
-
-
-@pytest.fixture()
-def cesm_surface_bgc_data(request, use_dask):
-    fname = download_test_data("CESM_BGC_SURFACE_2012.nc")
-
-    data = CESMBGCSurfaceForcingDataset(
-        filename=fname,
-        start_time=datetime(2012, 1, 1),
-        end_time=datetime(2013, 1, 1),
-        climatology=False,
-        use_dask=use_dask,
-    )
-    data.post_process()
-
-    return data
-
-
 @pytest.mark.parametrize(
     "data_fixture",
-    ["era5_data", "glorys_data", "tpxo_data", "cesm_bgc_data", "cesm_surface_bgc_data"],
+    [
+        "era5_data",
+        "glorys_data",
+        "tpxo_data",
+        "coarsened_cesm_bgc_data",
+        "cesm_surface_bgc_data",
+    ],
 )
-def test_lateral_fill_no_nans(data_fixture, request, use_dask):
+def test_lateral_fill_no_nans(data_fixture, request):
     data = request.getfixturevalue(data_fixture)
     lateral_fill = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
 
     for var in data.var_names:
-        filled = lateral_fill.apply(data.ds[data.var_names[var]].astype(np.float64))
+        if var in ["u", "v"]:
+            filled = lateral_fill_vel.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
+        else:
+            filled = lateral_fill.apply(data.ds[data.var_names[var]].astype(np.float64))
         assert not filled.isnull().any()
 
 
-def test_lateral_fill_correct_order_of_magnitude(cesm_bgc_data):
+def test_lateral_fill_correct_order_of_magnitude(coarsened_cesm_bgc_data):
 
     lateral_fill = LateralFill(
-        cesm_bgc_data.ds["mask"],
-        [cesm_bgc_data.dim_names["latitude"], cesm_bgc_data.dim_names["longitude"]],
+        coarsened_cesm_bgc_data.ds["mask"],
+        [
+            coarsened_cesm_bgc_data.dim_names["latitude"],
+            coarsened_cesm_bgc_data.dim_names["longitude"],
+        ],
     )
 
-    ALK = cesm_bgc_data.ds["ALK"]
+    ALK = coarsened_cesm_bgc_data.ds["ALK"]
 
     # zero out alkalinity field in all depth levels but the uppermost
-    ALK = xr.where(cesm_bgc_data.ds.ALK.depth > 25, 0, cesm_bgc_data.ds.ALK)
-    ALK = ALK.where(cesm_bgc_data.ds.mask)
+    ALK = xr.where(
+        coarsened_cesm_bgc_data.ds.ALK.depth > 25, 0, coarsened_cesm_bgc_data.ds.ALK
+    )
+    ALK = ALK.where(coarsened_cesm_bgc_data.ds.mask)
 
     filled = lateral_fill.apply(ALK.astype(np.float64))
 
@@ -150,28 +74,46 @@ def test_lateral_fill_correct_order_of_magnitude(cesm_bgc_data):
         "glorys_data",
     ],
 )
-def test_lateral_fill_reproducibility(data_fixture, request, use_dask):
+def test_lateral_fill_reproducibility(data_fixture, request):
 
     data = request.getfixturevalue(data_fixture)
     lateral_fill0 = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel0 = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
     lateral_fill1 = LateralFill(
         data.ds["mask"],
         [data.dim_names["latitude"], data.dim_names["longitude"]],
     )
+    if "mask_vel" in data.ds.data_vars:
+        lateral_fill_vel1 = LateralFill(
+            data.ds["mask_vel"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
 
     ds0 = data.ds.copy()
     ds1 = data.ds.copy()
 
     for var in data.var_names:
+        if var in ["u", "v"]:
+            ds0[data.var_names[var]] = lateral_fill_vel0.apply(
+                ds0[data.var_names[var]].astype(np.float64)
+            )
+            ds1[data.var_names[var]] = lateral_fill_vel1.apply(
+                ds1[data.var_names[var]].astype(np.float64)
+            )
 
-        ds0[data.var_names[var]] = lateral_fill0.apply(
-            ds0[data.var_names[var]].astype(np.float64)
-        )
-        ds1[data.var_names[var]] = lateral_fill1.apply(
-            ds1[data.var_names[var]].astype(np.float64)
-        )
+        else:
+            ds0[data.var_names[var]] = lateral_fill0.apply(
+                ds0[data.var_names[var]].astype(np.float64)
+            )
+            ds1[data.var_names[var]] = lateral_fill1.apply(
+                ds1[data.var_names[var]].astype(np.float64)
+            )
 
     assert ds0.equals(ds1)

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -181,37 +181,37 @@ def test_initial_conditions_plot_save(
 ):
     """Test plot and save methods."""
 
-    initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", s=0)
     initial_conditions_with_bgc_from_climatology.plot(
-        varname="temp", s=0, depth_contours=True
+        var_name="temp", s=0, depth_contours=True
     )
     initial_conditions_with_bgc_from_climatology.plot(
-        varname="temp", eta=0, layer_contours=True
+        var_name="temp", eta=0, layer_contours=True
     )
     initial_conditions_with_bgc_from_climatology.plot(
-        varname="temp", xi=0, layer_contours=True
+        var_name="temp", xi=0, layer_contours=True
     )
-    initial_conditions_with_bgc_from_climatology.plot(varname="temp", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="temp", xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0, xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="temp", eta=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", eta=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", s=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", eta=0, xi=0)
     initial_conditions_with_bgc_from_climatology.plot(
-        varname="u", s=0, layer_contours=True
+        var_name="u", s=0, layer_contours=True
     )
-    initial_conditions_with_bgc_from_climatology.plot(varname="u", s=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="u", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="u", xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="u", s=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="u", eta=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="u", xi=0)
     initial_conditions_with_bgc_from_climatology.plot(
-        varname="v", s=0, layer_contours=True
+        var_name="v", s=0, layer_contours=True
     )
-    initial_conditions_with_bgc_from_climatology.plot(varname="v", s=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="v", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="v", xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="zeta")
-    initial_conditions_with_bgc_from_climatology.plot(varname="ubar")
-    initial_conditions_with_bgc_from_climatology.plot(varname="vbar")
-    initial_conditions_with_bgc_from_climatology.plot(varname="ALK", s=0, xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(varname="ALK", eta=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="v", s=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="v", eta=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="v", xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="zeta")
+    initial_conditions_with_bgc_from_climatology.plot(var_name="ubar")
+    initial_conditions_with_bgc_from_climatology.plot(var_name="vbar")
+    initial_conditions_with_bgc_from_climatology.plot(var_name="ALK", s=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(var_name="ALK", eta=0, xi=0)
 
     for file_str in ["test_ic", "test_ic.nc"]:
         # Create a temporary filepath using the tmp_path fixture

roms_tools/tests/test_setup/test_regrid.py

@@ -0,0 +1,53 @@
+import pytest
+import numpy as np
+import xarray as xr
+from roms_tools.setup.regrid import VerticalRegrid
+
+
+def vertical_regridder(depth_values, layer_depth_rho_values):
+    class DataContainer:
+        """Mock class for holding data and dimension names."""
+
+        def __init__(self, ds):
+            self.ds = ds
+            self.dim_names = {"depth": "depth"}
+
+    class Grid:
+        """Mock class representing the grid object with layer depth information."""
+
+        def __init__(self, ds):
+            self.ds = ds
+
+    target_depth = xr.DataArray(data=layer_depth_rho_values, dims=["s_rho"])
+    source_depth = xr.DataArray(data=depth_values, dims=["depth"])
+
+    return VerticalRegrid(target_depth, source_depth)
+
+
+@pytest.mark.parametrize(
+    "depth_values, layer_depth_rho_values, temp_data",
+    [
+        ([5, 50, 100, 150], [130, 100, 70, 30, 10], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [130, 100, 70, 30, 2], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [200, 100, 70, 30, 10], [30, 25, 10, 2]),
+        ([5, 50, 100, 150], [200, 100, 70, 30, 1], [30, 25, 10, 2]),
+    ],
+)
+def test_vertical_regrid(request, depth_values, layer_depth_rho_values, temp_data):
+
+    vertical_regrid = vertical_regridder(
+        depth_values=depth_values, layer_depth_rho_values=layer_depth_rho_values
+    )
+    data = xr.Dataset({"temp_data": (["depth"], temp_data)})
+
+    # without filling in NaNs
+    regridded = vertical_regrid.apply(data.temp_data, fill_nans=False)
+    expected = np.interp(
+        layer_depth_rho_values, depth_values, temp_data, left=np.nan, right=np.nan
+    )
+    assert np.allclose(expected, regridded.data, equal_nan=True)
+
+    # with filling in NaNs
+    regridded = vertical_regrid.apply(data.temp_data, fill_nans=True)
+    expected = np.interp(layer_depth_rho_values, depth_values, temp_data)
+    assert np.allclose(expected, regridded.data, equal_nan=True)

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -1,5 +1,6 @@
 import pytest
 from datetime import datetime
+import xarray as xr
 from roms_tools import Grid, SurfaceForcing
 from roms_tools.setup.download import download_test_data
 import textwrap
@@ -487,6 +488,23 @@ def test_surface_forcing_creation(
     sfc_forcing.plot("pco2_air", time=0)
 
 
+@pytest.mark.parametrize(
+    "sfc_forcing_fixture",
+    [
+        "bgc_surface_forcing",
+        "bgc_surface_forcing_from_climatology",
+    ],
+)
+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(
+        sfc_forcing.ds.pco2_air, sfc_forcing.ds.pco2_air_alt, rtol=1.0e-5
+    )
+
+
 @pytest.mark.parametrize(
     "sfc_forcing_fixture",
     [
@@ -498,7 +516,7 @@ def test_surface_forcing_creation(
 def test_surface_forcing_plot_save(sfc_forcing_fixture, request, tmp_path):
     """Test plot and save methods."""
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
-    sfc_forcing.plot(varname="uwnd", time=0)
+    sfc_forcing.plot(var_name="uwnd", time=0)
 
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -531,7 +549,7 @@ def test_surface_forcing_bgc_plot_save(bgc_surface_forcing, tmp_path):
     """Test plot and save methods."""
 
     # Check the values in the dataset
-    bgc_surface_forcing.plot(varname="pco2_air", time=0)
+    bgc_surface_forcing.plot(var_name="pco2_air", time=0)
 
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
@@ -566,7 +584,7 @@ def test_surface_forcing_bgc_from_clim_plot_save(
     """Test plot and save methods."""
 
     # Check the values in the dataset
-    bgc_surface_forcing_from_climatology.plot(varname="pco2_air", time=0)
+    bgc_surface_forcing_from_climatology.plot(var_name="pco2_air", time=0)
 
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture

roms_tools/tests/test_setup/test_tides.py

@@ -174,7 +174,7 @@ def test_tidal_forcing_plot_save(tidal_forcing, tmp_path):
     computed."""
     tidal_forcing.ds.load()
 
-    tidal_forcing.plot(varname="ssh_Re", ntides=0)
+    tidal_forcing.plot(var_name="ssh_Re", ntides=0)
 
     for file_str in ["test_tides", "test_tides.nc"]:
         # Create a temporary filepath using the tmp_path fixture