roms-tools 0.0.6__py3-none-any.whl → 0.20__py3-none-any.whl

roms_tools/tests/test_grid.py

@@ -0,0 +1,226 @@
+import pytest
+import numpy as np
+import numpy.testing as npt
+from roms_tools import Grid
+import os
+import tempfile
+import importlib.metadata
+import textwrap
+
+
+def test_simple_regression():
+    grid = Grid(nx=1, ny=1, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=0)
+
+    expected_lat = np.array(
+        [
+            [-8.99249453e-01, -8.99249453e-01, -8.99249453e-01],
+            [0.0, 0.0, 0.0],
+            [8.99249453e-01, 8.99249453e-01, 8.99249453e-01],
+        ]
+    )
+    expected_lon = np.array(
+        [
+            [339.10072286, 340.0, 340.89927714],
+            [339.10072286, 340.0, 340.89927714],
+            [339.10072286, 340.0, 340.89927714],
+        ]
+    )
+
+    # TODO: adapt tolerances according to order of magnitude of respective fields
+    npt.assert_allclose(grid.ds["lat_rho"], expected_lat, atol=1e-8)
+    npt.assert_allclose(grid.ds["lon_rho"], expected_lon, atol=1e-8)
+
+
+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_roundtrip_netcdf():
+    """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",
+        smooth_factor=2,
+        hmin=5.0,
+        rmax=0.2,
+    )
+
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+        filepath = tmpfile.name
+
+    try:
+        # Save the grid to a file
+        grid_init.save(filepath)
+
+        # Load the grid from the file
+        grid_from_file = Grid.from_file(filepath)
+
+        # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
+        assert grid_init == grid_from_file
+
+    finally:
+        os.remove(filepath)
+
+
+def test_roundtrip_yaml():
+    """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",
+        smooth_factor=2,
+        hmin=5.0,
+        rmax=0.2,
+    )
+
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+        filepath = tmpfile.name
+
+    try:
+        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
+
+    finally:
+        os.remove(filepath)
+
+
+def test_from_yaml_missing_version():
+
+    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
+      smooth_factor: 8
+      hmin: 5.0
+      rmax: 0.2
+    """
+    )
+    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+        yaml_filepath = tmp_file.name
+        tmp_file.write(yaml_content.encode())
+
+    try:
+        with pytest.raises(
+            ValueError, match="Version of ROMS-Tools not found in the YAML file."
+        ):
+            Grid.from_yaml(yaml_filepath)
+    finally:
+        os.remove(yaml_filepath)
+
+
+def test_from_yaml_missing_grid():
+    roms_tools_version = importlib.metadata.version("roms-tools")
+
+    yaml_content = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+        yaml_filepath = tmp_file.name
+        tmp_file.write(yaml_content.encode())
+
+    try:
+        with pytest.raises(
+            ValueError, match="No Grid configuration found in the YAML file."
+        ):
+            Grid.from_yaml(yaml_filepath)
+    finally:
+        os.remove(yaml_filepath)
+
+
+def test_from_yaml_version_mismatch():
+    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
+    """
+    )
+
+    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+        yaml_filepath = tmp_file.name
+        tmp_file.write(yaml_content.encode())
+
+    try:
+        with pytest.warns(
+            UserWarning,
+            match="Current roms-tools version.*does not match the version in the YAML header.*",
+        ):
+            Grid.from_yaml(yaml_filepath)
+    finally:
+        os.remove(yaml_filepath)

roms_tools/tests/test_initial_conditions.py

@@ -0,0 +1,300 @@
+import pytest
+from datetime import datetime
+from roms_tools import InitialConditions, Grid, VerticalCoordinate
+import xarray as xr
+import numpy as np
+import tempfile
+import os
+import textwrap
+from roms_tools.setup.datasets import download_test_data
+
+
+@pytest.fixture
+def example_grid():
+    """
+    Fixture for creating a Grid object.
+    """
+    grid = Grid(
+        nx=2, ny=2, size_x=500, size_y=1000, center_lon=0, center_lat=55, rot=10
+    )
+
+    return grid
+
+
+@pytest.fixture
+def example_vertical_coordinate(example_grid):
+    """
+    Fixture for creating a VerticalCoordinate object.
+    """
+    vertical_coordinate = VerticalCoordinate(
+        grid=example_grid,
+        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 vertical_coordinate
+
+
+@pytest.fixture
+def initial_conditions(example_grid, example_vertical_coordinate):
+    """
+    Fixture for creating a dummy InitialConditions object.
+    """
+
+    fname = download_test_data("GLORYS_test_data.nc")
+
+    return InitialConditions(
+        grid=example_grid,
+        vertical_coordinate=example_vertical_coordinate,
+        ini_time=datetime(2021, 6, 29),
+        filename=fname,
+    )
+
+
+def test_initial_conditions_creation(initial_conditions):
+    """
+    Test the creation of the InitialConditions object.
+    """
+    assert initial_conditions.ini_time == datetime(2021, 6, 29)
+    assert initial_conditions.filename == download_test_data("GLORYS_test_data.nc")
+    assert initial_conditions.source == "GLORYS"
+
+
+def test_initial_conditions_ds_attribute(initial_conditions):
+    """
+    Test the ds attribute of the InitialConditions object.
+    """
+    assert isinstance(initial_conditions.ds, xr.Dataset)
+    assert "temp" in initial_conditions.ds
+    assert "salt" in initial_conditions.ds
+    assert "u" in initial_conditions.ds
+    assert "v" in initial_conditions.ds
+    assert "zeta" in initial_conditions.ds
+
+
+def test_initial_conditions_data_consistency_plot_save(initial_conditions, tmp_path):
+    """
+    Test that the data within the InitialConditions object remains consistent.
+    Also test plot and save methods in the same test since we dask arrays are already computed.
+    """
+    initial_conditions.ds.load()
+
+    # Define the expected data
+    expected_temp = np.array(
+        [
+            [
+                [
+                    [16.84414, 16.905312, 16.967817],
+                    [18.088203, 18.121834, 18.315424],
+                    [18.431192, 18.496748, 18.718002],
+                    [19.294329, 19.30358, 19.439777],
+                ],
+                [
+                    [12.639833, 13.479691, 14.426711],
+                    [15.712767, 15.920951, 16.10028],
+                    [13.02848, 14.5227165, 15.05175],
+                    [18.633307, 18.637077, 18.667465],
+                ],
+                [
+                    [11.027701, 11.650267, 12.200586],
+                    [12.302642, 12.646921, 13.150708],
+                    [8.143677, 11.435992, 13.356925],
+                    [8.710737, 11.25943, 13.111585],
+                ],
+                [
+                    [10.233599, 10.546486, 10.671082],
+                    [10.147331, 10.502733, 10.68275],
+                    [10.458557, 11.209945, 11.377164],
+                    [9.20282, 10.667074, 11.752404],
+                ],
+            ]
+        ],
+        dtype=np.float32,
+    )
+    expected_salt = np.array(
+        [
+            [
+                [
+                    [33.832672, 33.77759, 33.633846],
+                    [32.50002, 32.48105, 32.154694],
+                    [30.922323, 30.909824, 30.508572],
+                    [28.337738, 28.335176, 28.067144],
+                ],
+                [
+                    [34.8002, 34.691143, 34.382282],
+                    [32.43797, 32.369576, 32.027843],
+                    [34.885834, 34.82964, 34.775684],
+                    [29.237692, 29.232145, 29.09444],
+                ],
+                [
+                    [34.046825, 33.950684, 33.87148],
+                    [33.892323, 33.84102, 33.69169],
+                    [34.964134, 34.91892, 34.91941],
+                    [34.975933, 34.48586, 32.729057],
+                ],
+                [
+                    [35.21593, 35.209476, 35.20767],
+                    [35.224304, 35.209522, 35.20715],
+                    [35.299217, 35.31244, 35.31555],
+                    [34.25124, 33.828175, 33.234303],
+                ],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    expected_zeta = np.array(
+        [
+            [
+                [-0.30468762, -0.29416865, -0.30391693, -0.32985148],
+                [-0.34336275, -0.29455253, -0.3718359, -0.36176518],
+                [-0.3699948, -0.34693155, -0.41338325, -0.40663475],
+                [-0.5534979, -0.5270749, -0.45107934, -0.40699923],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    expected_u = np.array(
+        [
+            [
+                [[-0.0, -0.0, -0.0], [-0.0, -0.0, -0.0], [0.0, -0.0, -0.0]],
+                [[0.0, -0.0, -0.0], [-0.0, -0.0, -0.0], [-0.0, -0.0, -0.0]],
+                [
+                    [0.0, 0.0, -0.0],
+                    [0.0, 0.0, -0.0],
+                    [0.06979556, 0.06167743, -0.02247071],
+                ],
+                [
+                    [0.04268532, 0.03889201, 0.03351666],
+                    [0.04645353, 0.04914769, 0.03673013],
+                    [0.0211786, 0.03679834, 0.0274788],
+                ],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    expected_v = np.array(
+        [
+            [
+                [
+                    [0.0, 0.0, 0.0],
+                    [0.0, 0.0, -0.0],
+                    [-0.0, -0.0, -0.0],
+                    [-0.0, -0.0, -0.0],
+                ],
+                [
+                    [-0.0, -0.0, -0.0],
+                    [-0.0, -0.0, -0.0],
+                    [-0.03831354, -0.02400788, -0.03179555],
+                    [-0.0, -0.0, -0.0],
+                ],
+                [
+                    [-0.00951457, -0.00576979, -0.02147919],
+                    [-0.0, -0.0, -0.0],
+                    [0.01915873, 0.02625698, 0.01757628],
+                    [-0.06720348, -0.08354441, -0.13835917],
+                ],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    expected_ubar = np.array(
+        [
+            [
+                [0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.04028399],
+                [0.03866891, 0.04446249, 0.02812303],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    expected_vbar = np.array(
+        [
+            [
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, -0.03169237, 0.0],
+                [-0.01189703, 0.0, 0.02102064, -0.09326097],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
+    # Check the values in the dataset
+    assert np.allclose(initial_conditions.ds["temp"].values, expected_temp)
+    assert np.allclose(initial_conditions.ds["salt"].values, expected_salt)
+    assert np.allclose(initial_conditions.ds["zeta"].values, expected_zeta)
+    assert np.allclose(initial_conditions.ds["u"].values, expected_u)
+    assert np.allclose(initial_conditions.ds["v"].values, expected_v)
+    assert np.allclose(initial_conditions.ds["ubar"].values, expected_ubar)
+    assert np.allclose(initial_conditions.ds["vbar"].values, expected_vbar)
+
+    initial_conditions.plot(varname="temp", s=0)
+    initial_conditions.plot(varname="temp", eta=0)
+    initial_conditions.plot(varname="temp", xi=0)
+    initial_conditions.plot(varname="temp", s=0, xi=0)
+    initial_conditions.plot(varname="temp", eta=0, xi=0)
+    initial_conditions.plot(varname="zeta")
+
+    filepath = tmp_path / "initial_conditions.nc"
+    initial_conditions.save(filepath)
+    assert filepath.exists()
+
+
+def test_roundtrip_yaml(initial_conditions):
+    """Test that creating an InitialConditions object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+        filepath = tmpfile.name
+
+    try:
+        initial_conditions.to_yaml(filepath)
+
+        initial_conditions_from_file = InitialConditions.from_yaml(filepath)
+
+        assert initial_conditions == initial_conditions_from_file
+
+    finally:
+        os.remove(filepath)
+
+
+def test_from_yaml_missing_initial_conditions():
+    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
+    """
+    )
+
+    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+        yaml_filepath = tmp_file.name
+        tmp_file.write(yaml_content.encode())
+
+    try:
+        with pytest.raises(
+            ValueError,
+            match="No InitialConditions configuration found in the YAML file.",
+        ):
+            InitialConditions.from_yaml(yaml_filepath)
+    finally:
+        os.remove(yaml_filepath)