roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/tests/test_setup/test_grid.py

@@ -0,0 +1,377 @@
+import pytest
+import xarray as xr
+from roms_tools import Grid
+import importlib.metadata
+import textwrap
+from roms_tools.setup.download import download_test_data
+from conftest import calculate_file_hash
+from pathlib import Path
+
+
+def test_grid_creation(grid):
+
+    assert grid.nx == 1
+    assert grid.ny == 1
+    assert grid.size_x == 100
+    assert grid.size_y == 100
+    assert grid.center_lon == -20
+    assert grid.center_lat == 0
+    assert grid.rot == 0
+    assert isinstance(grid.ds, xr.Dataset)
+
+
+def test_plot_save_methods(tmp_path):
+
+    grid = Grid(
+        nx=20, ny=20, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=0
+    )
+
+    grid.plot(bathymetry=True)
+
+    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),
+        ]:  # test for Path object and str
+
+            # Test saving without partitioning
+            saved_filenames = grid.save(filepath)
+            # Check if the .nc file was created
+            filepath = Path(filepath).with_suffix(".nc")
+
+            assert saved_filenames == [filepath]
+            assert filepath.exists()
+            # Clean up the .nc file
+            filepath.unlink()
+
+            # Test saving with partitioning
+            saved_filenames = grid.save(filepath, np_eta=2, np_xi=5)
+
+            filepath_str = str(filepath.with_suffix(""))
+            expected_filepath_list = [
+                Path(filepath_str + f".{index}.nc") for index in range(10)
+            ]
+            assert saved_filenames == expected_filepath_list
+            for expected_filepath in expected_filepath_list:
+                assert expected_filepath.exists()
+                expected_filepath.unlink()
+
+
+def test_raise_if_domain_too_large():
+    with pytest.raises(ValueError, match="Domain size has to be smaller"):
+        Grid(nx=3, ny=3, size_x=30000, size_y=30000, center_lon=0, center_lat=51.5)
+
+    # test grid with reasonable domain size
+    grid = Grid(
+        nx=3,
+        ny=3,
+        size_x=1800,
+        size_y=2400,
+        center_lon=-21,
+        center_lat=61,
+        rot=20,
+    )
+    assert isinstance(grid, Grid)
+
+
+def test_grid_straddle_crosses_meridian():
+    grid = Grid(
+        nx=3,
+        ny=3,
+        size_x=100,
+        size_y=100,
+        center_lon=0,
+        center_lat=61,
+        rot=20,
+    )
+    assert grid.straddle
+
+    grid = Grid(
+        nx=3,
+        ny=3,
+        size_x=100,
+        size_y=100,
+        center_lon=180,
+        center_lat=61,
+        rot=20,
+    )
+    assert not grid.straddle
+
+
+def test_compatability_with_matlab_grid(tmp_path):
+
+    fname = download_test_data("grid_created_with_matlab.nc")
+
+    grid = Grid.from_file(fname)
+
+    assert not grid.straddle
+    assert grid.theta_s == 5.0
+    assert grid.theta_b == 2.0
+    assert grid.hc == 300.0
+    assert grid.N == 100
+    assert grid.nx == 24
+    assert grid.ny == 24
+    assert grid.center_lon == -4.1
+    assert grid.center_lat == 52.4
+    assert grid.rot == 0.0
+
+    expected_coords = set(
+        [
+            "lat_rho",
+            "lon_rho",
+            "lat_u",
+            "lon_u",
+            "lat_v",
+            "lon_v",
+            "lat_coarse",
+            "lon_coarse",
+            "layer_depth_rho",
+            "layer_depth_u",
+            "layer_depth_v",
+            "interface_depth_rho",
+            "interface_depth_u",
+            "interface_depth_v",
+        ]
+    )
+    actual_coords = set(grid.ds.coords.keys())
+    assert actual_coords == expected_coords
+
+    grid.plot(bathymetry=True)
+
+    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),
+        ]:  # test for Path object and str
+
+            # Test saving without partitioning
+            _ = grid.save(filepath)
+
+            filepath = Path(filepath)
+
+            # Load the grid from the file
+            grid_from_file = Grid.from_file(filepath.with_suffix(".nc"))
+
+            # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
+            assert grid == grid_from_file
+
+            # Clean up the .nc file
+            (filepath.with_suffix(".nc")).unlink()
+
+
+def test_roundtrip_netcdf(tmp_path):
+    """Test that creating a grid, saving it to file, and re-opening it is the same as just creating it."""
+
+    # Initialize a Grid object using the initializer
+    grid_init = Grid(
+        nx=10,
+        ny=15,
+        size_x=100.0,
+        size_y=150.0,
+        center_lon=0.0,
+        center_lat=0.0,
+        rot=0.0,
+        topography_source="ETOPO5",
+        hmin=5.0,
+    )
+
+    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),
+        ]:  # test for Path object and str
+
+            grid_init.save(filepath)
+
+            filepath = Path(filepath)
+
+            # Load the grid from the file
+            grid_from_file = Grid.from_file(filepath.with_suffix(".nc"))
+
+            # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
+            assert grid_init == grid_from_file
+
+            # Clean up the .nc file
+            (filepath.with_suffix(".nc")).unlink()
+
+
+def test_roundtrip_yaml(tmp_path):
+    """Test that creating a grid, saving its parameters to yaml file, and re-opening yaml file creates the same grid."""
+
+    # Initialize a Grid object using the initializer
+    grid_init = Grid(
+        nx=10,
+        ny=15,
+        size_x=100.0,
+        size_y=150.0,
+        center_lon=0.0,
+        center_lat=0.0,
+        rot=0.0,
+        topography_source="ETOPO5",
+        hmin=5.0,
+    )
+
+    # 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
+
+        grid_init.to_yaml(filepath)
+
+        grid_from_file = Grid.from_yaml(filepath)
+
+        # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
+        assert grid_init == grid_from_file
+
+        filepath = Path(filepath)
+        filepath.unlink()
+
+
+def test_files_have_same_hash(tmp_path):
+
+    # Initialize a Grid object using the initializer
+    grid_init = Grid(
+        nx=10,
+        ny=15,
+        size_x=100.0,
+        size_y=150.0,
+        center_lon=0.0,
+        center_lat=0.0,
+        rot=0.0,
+        topography_source="ETOPO5",
+        hmin=5.0,
+    )
+
+    yaml_filepath = tmp_path / "test_yaml"
+    filepath1 = tmp_path / "test1.nc"
+    filepath2 = tmp_path / "test2.nc"
+
+    grid_init.to_yaml(yaml_filepath)
+    grid_init.save(filepath1)
+
+    grid_from_file = Grid.from_yaml(yaml_filepath)
+    grid_from_file.save(filepath2)
+
+    hash1 = calculate_file_hash(filepath1)
+    hash2 = calculate_file_hash(filepath2)
+
+    assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
+
+    yaml_filepath.unlink()
+    filepath1.unlink()
+    filepath2.unlink()
+
+
+def test_from_yaml_missing_version(tmp_path):
+
+    yaml_content = textwrap.dedent(
+        """\
+    Grid:
+      nx: 100
+      ny: 100
+      size_x: 1800
+      size_y: 2400
+      center_lon: -10
+      center_lat: 61
+      rot: -20
+      topography_source: ETOPO5
+      hmin: 5.0
+    """
+    )
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for yaml_filepath in [
+        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)
+        else:
+            with open(yaml_filepath, "w") as f:
+                f.write(yaml_content)
+
+        with pytest.raises(
+            ValueError, match="Version of ROMS-Tools not found in the YAML file."
+        ):
+            Grid.from_yaml(yaml_filepath)
+
+        yaml_filepath = Path(yaml_filepath)
+        yaml_filepath.unlink()
+
+
+def test_from_yaml_missing_grid(tmp_path):
+    roms_tools_version = importlib.metadata.version("roms-tools")
+
+    yaml_content = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for yaml_filepath in [
+        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)
+        else:
+            with open(yaml_filepath, "w") as f:
+                f.write(yaml_content)
+
+        with pytest.raises(
+            ValueError, match="No Grid configuration found in the YAML file."
+        ):
+            Grid.from_yaml(yaml_filepath)
+
+        yaml_filepath = Path(yaml_filepath)
+        yaml_filepath.unlink()
+
+
+def test_from_yaml_version_mismatch(tmp_path):
+    yaml_content = textwrap.dedent(
+        """\
+    ---
+    roms_tools_version: 0.0.0
+    ---
+    Grid:
+      nx: 100
+      ny: 100
+      size_x: 1800
+      size_y: 2400
+      center_lon: -10
+      center_lat: 61
+      rot: -20
+      topography_source: ETOPO5
+      hmin: 5.0
+    """
+    )
+
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for yaml_filepath in [
+        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)
+        else:
+            with open(yaml_filepath, "w") as f:
+                f.write(yaml_content)
+
+        with pytest.warns(
+            UserWarning,
+            match="Current roms-tools version.*does not match the version in the YAML header.*",
+        ):
+            Grid.from_yaml(yaml_filepath)
+
+        yaml_filepath = Path(yaml_filepath)
+        yaml_filepath.unlink()

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -0,0 +1,338 @@
+import pytest
+from datetime import datetime
+from roms_tools import InitialConditions, Grid
+import xarray as xr
+import numpy as np
+import textwrap
+from roms_tools.setup.download import download_test_data
+from roms_tools.setup.datasets import CESMBGCDataset
+from pathlib import Path
+from conftest import calculate_file_hash
+
+
+@pytest.mark.parametrize(
+    "ic_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions_with_bgc",
+        "initial_conditions_with_bgc_from_climatology",
+    ],
+)
+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)
+    assert ic.source == {
+        "name": "GLORYS",
+        "path": download_test_data("GLORYS_coarse_test_data.nc"),
+        "climatology": False,
+    }
+    assert isinstance(ic.ds, xr.Dataset)
+    assert "temp" in ic.ds
+    assert "salt" in ic.ds
+    assert "u" in ic.ds
+    assert "v" in ic.ds
+    assert "zeta" in ic.ds
+    assert ic.ds.coords["ocean_time"].attrs["units"] == "seconds"
+
+
+@pytest.fixture
+def example_grid():
+
+    grid = Grid(
+        nx=2,
+        ny=2,
+        size_x=500,
+        size_y=1000,
+        center_lon=0,
+        center_lat=55,
+        rot=10,
+        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
+    )
+
+    return grid
+
+
+# Test initialization with missing 'name' in source
+def test_initial_conditions_missing_physics_name(example_grid, use_dask):
+    with pytest.raises(ValueError, match="`source` must include a 'name'."):
+        InitialConditions(
+            grid=example_grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"path": "physics_data.nc"},
+            use_dask=use_dask,
+        )
+
+
+# Test initialization with missing 'path' in source
+def test_initial_conditions_missing_physics_path(example_grid, use_dask):
+    with pytest.raises(ValueError, match="`source` must include a 'path'."):
+        InitialConditions(
+            grid=example_grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"name": "GLORYS"},
+            use_dask=use_dask,
+        )
+
+
+# Test initialization with missing 'name' in bgc_source
+def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    with pytest.raises(
+        ValueError, match="`bgc_source` must include a 'name' if it is provided."
+    ):
+        InitialConditions(
+            grid=example_grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"name": "GLORYS", "path": fname},
+            bgc_source={"path": "bgc_data.nc"},
+            use_dask=use_dask,
+        )
+
+
+# Test initialization with missing 'path' in bgc_source
+def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    with pytest.raises(
+        ValueError, match="`bgc_source` must include a 'path' if it is provided."
+    ):
+        InitialConditions(
+            grid=example_grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"name": "GLORYS", "path": fname},
+            bgc_source={"name": "CESM_REGRIDDED"},
+            use_dask=use_dask,
+        )
+
+
+# Test default climatology value
+def test_initial_conditions_default_climatology(example_grid, use_dask):
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+
+    initial_conditions = InitialConditions(
+        grid=example_grid,
+        ini_time=datetime(2021, 6, 29),
+        source={"name": "GLORYS", "path": fname},
+        use_dask=use_dask,
+    )
+
+    assert initial_conditions.source["climatology"] is False
+    assert initial_conditions.bgc_source is None
+
+
+def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname_bgc = download_test_data("CESM_regional_test_data_one_time_slice.nc")
+
+    initial_conditions = InitialConditions(
+        grid=example_grid,
+        ini_time=datetime(2021, 6, 29),
+        source={"name": "GLORYS", "path": fname},
+        bgc_source={"name": "CESM_REGRIDDED", "path": fname_bgc},
+        use_dask=use_dask,
+    )
+
+    assert initial_conditions.bgc_source["climatology"] is False
+
+
+def test_interpolation_from_climatology(
+    initial_conditions_with_bgc_from_climatology, use_dask
+):
+
+    fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+    ds = xr.open_dataset(fname_bgc)
+
+    # check if interpolated value for Jan 15 is indeed January value from climatology
+    bgc_data = CESMBGCDataset(
+        filename=fname_bgc,
+        start_time=datetime(2012, 1, 15),
+        climatology=True,
+        use_dask=use_dask,
+        apply_post_processing=False,
+    )
+    assert np.allclose(ds["ALK"].sel(month=1), bgc_data.ds["ALK"], equal_nan=True)
+
+    # check if interpolated value for Jan 30 is indeed average of January and February value from climatology
+    bgc_data = CESMBGCDataset(
+        filename=fname_bgc,
+        start_time=datetime(2012, 1, 30),
+        climatology=True,
+        use_dask=use_dask,
+        apply_post_processing=False,
+    )
+    assert np.allclose(
+        0.5 * (ds["ALK"].sel(month=1) + ds["ALK"].sel(month=2)),
+        bgc_data.ds["ALK"],
+        equal_nan=True,
+    )
+
+
+def test_initial_conditions_plot_save(
+    initial_conditions_with_bgc_from_climatology, tmp_path
+):
+    """
+    Test plot and save methods.
+    """
+
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0)
+    initial_conditions_with_bgc_from_climatology.plot(
+        varname="temp", s=0, depth_contours=True
+    )
+    initial_conditions_with_bgc_from_climatology.plot(
+        varname="temp", eta=0, layer_contours=True
+    )
+    initial_conditions_with_bgc_from_climatology.plot(
+        varname="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(
+        varname="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(
+        varname="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)
+
+    for file_str in ["test_ic", "test_ic.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 partitioning
+            saved_filenames = initial_conditions_with_bgc_from_climatology.save(
+                filepath
+            )
+            # Check if the .nc file was created
+            filepath = Path(filepath).with_suffix(".nc")
+            assert saved_filenames == [filepath]
+            assert filepath.exists()
+            # Clean up the .nc file
+            filepath.unlink()
+
+            # Test saving with partitioning
+            saved_filenames = initial_conditions_with_bgc_from_climatology.save(
+                filepath, np_eta=2
+            )
+
+            filepath_str = str(filepath.with_suffix(""))
+            expected_filepath_list = [
+                Path(filepath_str + f".{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_roundtrip_yaml(initial_conditions, 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."""
+
+    # 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
+
+        initial_conditions.to_yaml(filepath)
+
+        initial_conditions_from_file = InitialConditions.from_yaml(
+            filepath, use_dask=use_dask
+        )
+
+        assert initial_conditions == initial_conditions_from_file
+
+        filepath = Path(filepath)
+        filepath.unlink()
+
+
+def test_files_have_same_hash(initial_conditions, tmp_path, use_dask):
+
+    yaml_filepath = tmp_path / "test_yaml.yaml"
+    filepath1 = tmp_path / "test1.nc"
+    filepath2 = tmp_path / "test2.nc"
+
+    initial_conditions.to_yaml(yaml_filepath)
+    initial_conditions.save(filepath1)
+    ic_from_file = InitialConditions.from_yaml(yaml_filepath, use_dask)
+    ic_from_file.save(filepath2)
+
+    hash1 = calculate_file_hash(filepath1)
+    hash2 = calculate_file_hash(filepath2)
+
+    assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
+
+    yaml_filepath.unlink()
+    filepath1.unlink()
+    filepath2.unlink()
+
+
+def test_from_yaml_missing_initial_conditions(tmp_path, use_dask):
+    yaml_content = textwrap.dedent(
+        """\
+    ---
+    roms_tools_version: 0.0.0
+    ---
+    Grid:
+      nx: 100
+      ny: 100
+      size_x: 1800
+      size_y: 2400
+      center_lon: -10
+      center_lat: 61
+      rot: -20
+      topography_source: ETOPO5
+      smooth_factor: 8
+      hmin: 5.0
+      rmax: 0.2
+    """
+    )
+
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for yaml_filepath in [
+        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)
+        else:
+            with open(yaml_filepath, "w") as f:
+                f.write(yaml_content)
+
+        with pytest.raises(
+            ValueError,
+            match="No InitialConditions configuration found in the YAML file.",
+        ):
+            InitialConditions.from_yaml(yaml_filepath, use_dask)
+
+        yaml_filepath = Path(yaml_filepath)
+        yaml_filepath.unlink()