tobac 1.6.2__py3-none-any.whl

tobac/tests/test_utils_coordinates.py

@@ -0,0 +1,328 @@
+"""Unit tests for tobac.utils.internal.general_internal.py"""
+
+import pytest
+import pandas as pd
+
+import tobac.utils.internal.coordinates as coord_utils
+
+
+def test_find_coord_in_dataframe_errors():
+    """Test that find_coord_in_dataframe raises errors correctly"""
+    defaults = ["x", "projection_x_coordinate", "__other_name"]
+
+    # Test no options raises ValueError:
+    with pytest.raises(ValueError, match="One of coord or defaults parameter*"):
+        coord_utils.find_coord_in_dataframe(pd.DataFrame(columns=["time", "x"]))
+
+    # Test coordinate specified not in dataframe raise ValueError:
+    with pytest.raises(ValueError, match="Coordinate*"):
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x"]), coord="projection_x_coordinate"
+        )
+
+    # Test no coordinates matching defaults:
+    with pytest.raises(ValueError, match="No coordinate found matching defaults*"):
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "y"]), defaults=defaults
+        )
+
+    # Test multiple matches with defaults:
+    with pytest.raises(ValueError, match="Multiple matching*"):
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x", "projection_x_coordinate"]),
+            defaults=defaults,
+        )
+
+    # Test that giving an object that is not a dataframe or series returns an error
+    with pytest.raises(ValueError, match="Input variable_dataframe is neither*"):
+        coord_utils.find_coord_in_dataframe("test_str", defaults=defaults)
+
+
+def test_find_coord_in_dataframe():
+    """Test that find_coord_in_dataframe returns correct results for both
+    default and specific coordinates
+    """
+    defaults = ["x", "projection_x_coordinate", "__other_name"]
+
+    # Now test correct returns:
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x", "projection_x_coordinate"]), coord="x"
+        )
+        == "x"
+    )
+
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x", "projection_x_coordinate"]),
+            coord="projection_x_coordinate",
+        )
+        == "projection_x_coordinate"
+    )
+
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x", "y"]), defaults=defaults
+        )
+        == "x"
+    )
+
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(
+                columns=["time", "projection_x_coordinate", "projection_y_coordinate"]
+            ),
+            defaults=defaults,
+        )
+        == "projection_x_coordinate"
+    )
+
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.DataFrame(columns=["time", "x", "projection_x_coordinate"]),
+            coord="x",
+            defaults=defaults,
+        )
+        == "x"
+    )
+
+    # Test pd.Series input:
+    assert (
+        coord_utils.find_coord_in_dataframe(
+            pd.Series(index=["time", "x", "projection_x_coordinate"]), coord="x"
+        )
+        == "x"
+    )
+
+
+def test_find_dataframe_vertical_coord_warning():
+    """Test the warning for coord="auto" in find_dataframe_vertical_coord"""
+    with pytest.warns(DeprecationWarning):
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["z"]), vertical_coord="auto"
+        )
+
+
+def test_find_dataframe_vertical_coord_error():
+    """Test find_dataframe_vertical_coord raises errors correctly"""
+    # Test the error for invalid coord input:
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["z"]), vertical_coord="__bad_coord_name"
+        )
+
+    # Test the error for no default coord found:
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_vertical_coord(pd.DataFrame(columns=["x"]))
+
+    # Test the error for multiple default coords found:
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["z", "geopotential_height"])
+        )
+
+
+def test_find_dataframe_vertical_coord():
+    """Test find_dataframe_vertical_coord provides correct results"""
+    # Test default coords
+    assert coord_utils.find_dataframe_vertical_coord(pd.DataFrame(columns=["z"])) == "z"
+    assert (
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["geopotential_height"])
+        )
+        == "geopotential_height"
+    )
+
+    # Test coord input
+    assert (
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["p"]), vertical_coord="p"
+        )
+        == "p"
+    )
+
+    # Test coord input when multiple default coords
+    assert (
+        coord_utils.find_dataframe_vertical_coord(
+            pd.DataFrame(columns=["z", "geopotential_height"]), vertical_coord="z"
+        )
+        == "z"
+    )
+
+
+def test_find_dataframe_horizontal_coords_error():
+    """Test find_dataframe_horizontal_coords raises errors correctly"""
+    # Test no matching coords
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "z"])
+        )
+
+    # Test hdim_1_coord or hdim_2_coord set but not coord_type
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "y"]), hdim1_coord="y"
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "y"]), hdim2_coord="x"
+        )
+
+    # Test one exists but not both:
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x"])
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "y"])
+        )
+
+    # Test one of each exists
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "lat"])
+        )
+
+    # Test failure to detect coords when hdim1_coord or hdim2_coord is specified:
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "lat"]), hdim1_coord="y", coord_type="xy"
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "y", "lon"]), hdim2_coord="x", coord_type="xy"
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "lon"]),
+            hdim1_coord="lat",
+            coord_type="latlon",
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "lat"]),
+            hdim1_coord="lon",
+            coord_type="latlon",
+        )
+
+
+def test_find_dataframe_horizontal_coords_error_coord_type():
+    """Test that find_dataframe_horizontal_coords raises errors correctly when
+    the specified coord_type does not match the coords present
+    """
+    # Check that if coord_type is specified that an error is raised even if the other type of coords are present
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "x", "y"]), coord_type="latlon"
+        )
+
+    with pytest.raises(ValueError):
+        coord_utils.find_dataframe_horizontal_coords(
+            pd.DataFrame(columns=["time", "lat", "lon"]), coord_type="xy"
+        )
+
+
+def test_find_dataframe_horizontal_coords_defaults_xy():
+    """Test find_dataframe_horizontal_coords for xy coords"""
+    # Test defaults xy:
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "x", "y"])
+    ) == ("y", "x", "xy")
+
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(
+            columns=["time", "projection_x_coordinate", "projection_y_coordinate"]
+        )
+    ) == ("projection_y_coordinate", "projection_x_coordinate", "xy")
+
+    # Test that xy take priority over latlon
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "x", "y", "lat", "lon"])
+    ) == ("y", "x", "xy")
+
+
+def test_find_dataframe_horizontal_coords_defaults_latlon():
+    """Test find_dataframe_horizontal_coords for lat/lon coords"""
+    # Test defaults latlon
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "lon", "lat"])
+    ) == ("lat", "lon", "latlon")
+
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "Longitude", "Latitude"])
+    ) == ("Latitude", "Longitude", "latlon")
+
+    # Test that if only one of xy take latlon instead
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "x", "lat", "lon"])
+    ) == ("lat", "lon", "latlon")
+
+    # Test that setting coord_type to latlon ignores xy coords
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "x", "y", "lat", "lon"]), coord_type="latlon"
+    ) == ("lat", "lon", "latlon")
+
+
+def test_find_dataframe_horizontal_coords_specific():
+    """Test find_dataframe_horizontal_coords when the coordinate name is
+    specified
+    """
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(
+            columns=[
+                "time",
+                "x",
+                "y",
+                "projection_x_coordinate",
+                "projection_y_coordinate",
+            ]
+        ),
+        hdim1_coord="y",
+        hdim2_coord="x",
+        coord_type="xy",
+    ) == ("y", "x", "xy")
+
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(
+            columns=[
+                "time",
+                "x",
+                "y",
+                "projection_x_coordinate",
+                "projection_y_coordinate",
+            ]
+        ),
+        hdim1_coord="projection_y_coordinate",
+        hdim2_coord="projection_x_coordinate",
+        coord_type="xy",
+    ) == ("projection_y_coordinate", "projection_x_coordinate", "xy")
+
+    # Check that order does not matter
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(
+            columns=[
+                "time",
+                "x",
+                "y",
+                "projection_x_coordinate",
+                "projection_y_coordinate",
+            ]
+        ),
+        hdim1_coord="x",
+        hdim2_coord="y",
+        coord_type="xy",
+    ) == ("x", "y", "xy")
+
+    # Check that coord_type can be set wrong
+    assert coord_utils.find_dataframe_horizontal_coords(
+        pd.DataFrame(columns=["time", "x", "y", "lat", "lon"]),
+        hdim1_coord="lat",
+        hdim2_coord="lon",
+        coord_type="xy",
+    ) == ("lat", "lon", "xy")

tobac/tests/test_utils_internal.py

@@ -0,0 +1,97 @@
+import tobac.utils.internal as internal_utils
+import tobac.testing as tbtest
+
+import pytest
+import numpy as np
+import xarray as xr
+
+
+@pytest.mark.parametrize(
+    "dset_type, time_axis, vertical_axis, expected_out",
+    [
+        ("iris", 0, 1, (2, 3)),
+        ("iris", -1, 0, (1, 2)),
+        ("iris", 0, -1, (1, 2)),
+        ("iris", 0, 2, (1, 3)),
+        ("iris", 3, 0, (1, 2)),
+        ("iris", 0, 3, (1, 2)),
+        ("iris", 1, 2, (0, 3)),
+        ("xarray", 0, 1, (2, 3)),
+        ("xarray", 0, 2, (1, 3)),
+        ("xarray", 3, 0, (1, 2)),
+        ("xarray", 0, 3, (1, 2)),
+        ("xarray", 1, 2, (0, 3)),
+    ],
+)
+def test_find_hdim_axes_3D(dset_type, time_axis, vertical_axis, expected_out):
+    """Tests tobac.utils.internal.file_hdim_axes_3D
+
+    Parameters
+    ----------
+    dset_type: str{"xarray" or "iris"}
+        type of the dataset to generate
+    time_axis: int
+        axis number of the time coordinate (or -1 to not have one)
+    vertical_axis: int
+        axis number of the vertical coordinate (or -1 to not have one)
+    expected_out: tuple (int, int)
+        expected output
+    """
+    ndims = 2 + (1 if time_axis >= 0 else 0) + (1 if vertical_axis >= 0 else 0)
+    test_dset_size = [2] * ndims
+
+    test_data = np.zeros(test_dset_size)
+
+    dset_opts = {
+        "in_arr": test_data,
+        "data_type": dset_type,
+    }
+    if time_axis >= 0:
+        dset_opts["time_dim_num"] = time_axis
+    if vertical_axis >= 0:
+        dset_opts["z_dim_num"] = vertical_axis
+        dset_opts["z_dim_name"] = "altitude"
+
+    y_set = False
+    for dim_number in range(ndims):
+        if time_axis != dim_number and vertical_axis != dim_number:
+            if not y_set:
+                dset_opts["y_dim_num"] = dim_number
+                y_set = True
+            else:
+                dset_opts["x_dim_num"] = dim_number
+
+    cube_test = tbtest.make_dataset_from_arr(**dset_opts)
+
+    out_coords = internal_utils.find_hdim_axes_3D(cube_test)
+
+    assert out_coords == expected_out
+
+
+@pytest.mark.parametrize(
+    "lat_name, lon_name, lat_name_test, lon_name_test, expected_result",
+    [
+        ("lat", "lon", None, None, ("lat", "lon")),
+        ("lat", "long", None, None, ("lat", "long")),
+        ("lat", "altitude", None, None, ("lat", None)),
+        ("lat", "longitude", "lat", "longitude", ("lat", "longitude")),
+    ],
+)
+def test_detect_latlon_coord_name(
+    lat_name, lon_name, lat_name_test, lon_name_test, expected_result
+):
+    """Tests tobac.utils.internal.detect_latlon_coord_name"""
+
+    in_arr = np.empty((50, 50))
+    lat_vals = np.empty(50)
+    lon_vals = np.empty(50)
+
+    in_xr = xr.Dataset(
+        {"data": ((lat_name, lon_name), in_arr)},
+        coords={lat_name: lat_vals, lon_name: lon_vals},
+    )
+    out_lat_name, out_lon_name = internal_utils.detect_latlon_coord_name(
+        in_xr["data"].to_iris(), lat_name_test, lon_name_test
+    )
+    assert out_lat_name == expected_result[0]
+    assert out_lon_name == expected_result[1]

tobac/tests/test_xarray_utils.py

@@ -0,0 +1,232 @@
+"""Tests for tobac.utils.internal_utils.xarray_utils"""
+
+from __future__ import annotations
+
+from typing import Union
+
+import pandas as pd
+import pytest
+import numpy as np
+import xarray as xr
+
+import tobac.utils.internal.xarray_utils as xr_utils
+import tobac.testing as tbtest
+import datetime
+
+
+@pytest.mark.parametrize(
+    "dim_names, coord_dim_map, coord_looking_for, expected_out, expected_raise",
+    [
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "time",  # coord_looking_for
+            0,
+            False,
+        ),
+        (
+            ("time", "time", "time", "time", "time"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+            },
+            "time",  # coord_looking_for
+            0,
+            True,
+        ),
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "altitude",  # coord_looking_for
+            1,
+            False,
+        ),
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "latitude",  # coord_looking_for
+            None,
+            True,
+        ),
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "x",  # coord_looking_for
+            2,
+            False,
+        ),
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "time": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "z",  # coord_looking_for
+            2,
+            True,
+        ),
+        (
+            ("time", "altitude", "x", "y"),  # dim_names
+            {  # coord_dim_map
+                "t": ("time",),
+                "latitude": ("x", "y"),
+                "longitude": ("x", "y"),
+                "altmsl": ("altitude", "x", "y"),
+            },
+            "t",  # coord_looking_for
+            0,
+            False,
+        ),
+    ],
+)
+def test_find_axis_from_dim_coord(
+    dim_names: tuple[str],
+    coord_dim_map: dict,
+    coord_looking_for: str,
+    expected_out: Union[int, None],
+    expected_raise: bool,
+):
+    """Tests tobac.utils.internal.file_hdim_axes_3D
+
+    Parameters
+    ----------
+    dim_names: tuple[str]
+        Names of the dimensions to have
+    coord_dim_map: dict[str : tuple[str],]
+        Mapping of coordinates (keys) to dimensions (values)
+    coord_looking_for: str
+        what coordinate/dimension to look for
+    expected_out: Union[int, None]
+        What the expected output is
+    expected_raise: bool
+        Whether or not we expect a raise
+    """
+
+    # size of the array per dimension
+    arr_sz = 4
+    arr_da = np.empty((arr_sz,) * len(dim_names))
+    coord_vals = {}
+    for coord_nm in coord_dim_map:
+        coord_vals[coord_nm] = (
+            coord_dim_map[coord_nm],
+            np.empty((arr_sz,) * len(coord_dim_map[coord_nm])),
+        )
+
+    xr_da = xr.DataArray(arr_da, dims=dim_names, coords=coord_vals)
+    if expected_raise:
+        with pytest.raises(ValueError):
+            _ = xr_utils.find_axis_from_dim_coord(xr_da, coord_looking_for)
+    else:
+        out_val = xr_utils.find_axis_from_dim_coord(xr_da, coord_looking_for)
+        if expected_out is not None:
+            assert out_val == expected_out
+        else:
+            assert out_val is None
+
+
+@pytest.mark.parametrize(
+    "dim_names, coord_dim_map, feature_pos, expected_vals",
+    [
+        (
+            ["time", "x", "y"],
+            {
+                "test_coord1": (tuple(), 1),
+                "test_coord_time": ("time", [5, 6, 7, 8, 9, 10]),
+            },
+            (1, 1),
+            {"test_coord1": (1, 1, 1), "test_coord_time": (5, 6, 7)},
+        ),
+        (
+            ["time", "x", "y"],
+            {
+                "test_coord_datetime": (
+                    "time",
+                    pd.date_range(
+                        datetime.datetime(2000, 1, 1),
+                        datetime.datetime(2000, 1, 1, 6),
+                        freq="1h",
+                        inclusive="left",
+                    ),
+                ),
+                "test_coord_time": ("time", [5, 6, 7, 8, 9, 10]),
+            },
+            (1, 1),
+            {
+                "test_coord_datetime": pd.date_range(
+                    datetime.datetime(2000, 1, 1),
+                    datetime.datetime(2000, 1, 1, 3),
+                    freq="1h",
+                    inclusive="left",
+                ),
+                "test_coord_time": (5, 6, 7),
+            },
+        ),
+    ],
+)
+def test_add_coordinates_to_features_interpolate_along_other_dims(
+    dim_names: tuple[str],
+    coord_dim_map: dict,
+    feature_pos: tuple[int],
+    expected_vals: dict[str, tuple],
+):
+    time_len: int = 6
+    if len(feature_pos) == 2:
+        all_feats = tbtest.generate_single_feature(
+            feature_pos[0],
+            feature_pos[1],
+            feature_num=1,
+            num_frames=3,
+            max_h1=100,
+            max_h2=100,
+        )
+        arr_size = (time_len, 5, 5)
+
+    elif len(feature_pos) == 3:
+        all_feats = tbtest.generate_single_feature(
+            feature_pos[1],
+            feature_pos[2],
+            start_v=feature_pos[0],
+            feature_num=1,
+            num_frames=3,
+            max_h1=100,
+            max_h2=100,
+        )
+        arr_size = (time_len, 1, 5, 5)
+    else:
+        raise ValueError("too many dimensions")
+    coord_dim_map["time"] = (
+        ("time",),
+        [
+            datetime.datetime(2000, 1, 1, 0) + datetime.timedelta(hours=x)
+            for x in range(time_len)
+        ],
+    )
+
+    test_xr_arr = xr.DataArray(np.empty(arr_size), dims=dim_names, coords=coord_dim_map)
+
+    resulting_df = xr_utils.add_coordinates_to_features(all_feats, test_xr_arr)
+    for coord in coord_dim_map:
+        assert coord in resulting_df
+        if coord != "time":
+            assert np.all(resulting_df[coord].values == expected_vals[coord])