tobac 1.6.2__py3-none-any.whl

tobac/tests/test_convert.py

@@ -0,0 +1,265 @@
+"""Tests for the iris/xarray conversion decorators"""
+
+import pytest
+import tobac
+import tobac.testing
+import xarray
+from iris.cube import Cube
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from copy import deepcopy
+from tobac.utils.decorators import (
+    xarray_to_iris,
+    iris_to_xarray,
+    xarray_to_irispandas,
+    irispandas_to_xarray,
+)
+
+
+@pytest.mark.parametrize(
+    "decorator, input_types, expected_internal_types, expected_output_type",
+    [
+        (
+            xarray_to_iris,
+            [xarray.DataArray, xarray.DataArray],
+            [Cube, Cube],
+            xarray.DataArray,
+        ),
+        (xarray_to_iris, [Cube, Cube], [Cube, Cube], Cube),
+        (xarray_to_iris, [Cube, xarray.DataArray], [Cube, Cube], xarray.DataArray),
+        (xarray_to_iris, [xarray.DataArray, Cube], [Cube, Cube], xarray.DataArray),
+        (iris_to_xarray, [Cube, Cube], [xarray.DataArray, xarray.DataArray], Cube),
+        (
+            iris_to_xarray,
+            [xarray.DataArray, xarray.DataArray],
+            [xarray.DataArray, xarray.DataArray],
+            xarray.DataArray,
+        ),
+        (
+            iris_to_xarray,
+            [xarray.DataArray, Cube],
+            [xarray.DataArray, xarray.DataArray],
+            Cube,
+        ),
+        (
+            iris_to_xarray,
+            [Cube, xarray.DataArray],
+            [xarray.DataArray, xarray.DataArray],
+            Cube,
+        ),
+        (
+            xarray_to_irispandas,
+            [xarray.DataArray, xarray.DataArray],
+            [Cube, Cube],
+            xarray.DataArray,
+        ),
+        (xarray_to_irispandas, [Cube, Cube], [Cube, Cube], Cube),
+        (
+            xarray_to_irispandas,
+            [Cube, xarray.DataArray],
+            [Cube, Cube],
+            xarray.DataArray,
+        ),
+        (
+            xarray_to_irispandas,
+            [xarray.DataArray, Cube],
+            [Cube, Cube],
+            xarray.DataArray,
+        ),
+        (
+            xarray_to_irispandas,
+            [xarray.Dataset, xarray.Dataset],
+            [pd.DataFrame, pd.DataFrame],
+            xarray.Dataset,
+        ),
+        (
+            xarray_to_irispandas,
+            [pd.DataFrame, pd.DataFrame],
+            [pd.DataFrame, pd.DataFrame],
+            pd.DataFrame,
+        ),
+        (
+            xarray_to_irispandas,
+            [xarray.Dataset, pd.DataFrame],
+            [pd.DataFrame, pd.DataFrame],
+            xarray.Dataset,
+        ),
+        (
+            xarray_to_irispandas,
+            [pd.DataFrame, xarray.Dataset],
+            [pd.DataFrame, pd.DataFrame],
+            xarray.Dataset,
+        ),
+        (
+            xarray_to_irispandas,
+            [xarray.Dataset, xarray.DataArray],
+            [pd.DataFrame, Cube],
+            xarray.Dataset,
+        ),
+        (
+            irispandas_to_xarray,
+            [Cube, Cube],
+            [xarray.DataArray, xarray.DataArray],
+            Cube,
+        ),
+        (
+            irispandas_to_xarray,
+            [xarray.DataArray, xarray.DataArray],
+            [xarray.DataArray, xarray.DataArray],
+            xarray.DataArray,
+        ),
+        (
+            irispandas_to_xarray,
+            [xarray.DataArray, Cube],
+            [xarray.DataArray, xarray.DataArray],
+            Cube,
+        ),
+        (
+            irispandas_to_xarray,
+            [Cube, xarray.DataArray],
+            [xarray.DataArray, xarray.DataArray],
+            Cube,
+        ),
+        (
+            irispandas_to_xarray,
+            [pd.DataFrame, pd.DataFrame],
+            [xarray.Dataset, xarray.Dataset],
+            pd.DataFrame,
+        ),
+        (
+            irispandas_to_xarray,
+            [xarray.Dataset, xarray.Dataset],
+            [xarray.Dataset, xarray.Dataset],
+            xarray.Dataset,
+        ),
+        (
+            irispandas_to_xarray,
+            [pd.DataFrame, xarray.Dataset],
+            [xarray.Dataset, xarray.Dataset],
+            pd.DataFrame,
+        ),
+        (
+            irispandas_to_xarray,
+            [xarray.Dataset, pd.DataFrame],
+            [xarray.Dataset, xarray.Dataset],
+            pd.DataFrame,
+        ),
+        (
+            irispandas_to_xarray,
+            [pd.DataFrame, Cube],
+            [xarray.Dataset, xarray.DataArray],
+            pd.DataFrame,
+        ),
+    ],
+)
+def test_converting(
+    decorator, input_types, expected_internal_types, expected_output_type
+):
+    """Testing the conversions of the decorators internally and for the output"""
+
+    def test_function_kwarg(test_input, kwarg=None):
+        assert (
+            type(test_input) == expected_internal_types[0]
+        ), "Expected internal type {}, got {} for {}".format(
+            expected_internal_types[0], type(test_input), decorator.__name__
+        )
+        assert (
+            type(kwarg) == expected_internal_types[1]
+        ), "Expected internal type {}, got {} for {} as keyword argument".format(
+            expected_internal_types[1], type(kwarg), decorator.__name__
+        )
+        return test_input
+
+    def test_function_tuple_output(test_input, kwarg=None):
+        return (test_input, test_input)
+
+    decorator_i = decorator()
+    decorated_function_kwarg = decorator_i(test_function_kwarg)
+    decorated_function_tuple = decorator_i(test_function_tuple_output)
+
+    if input_types[0] == xarray.DataArray:
+        data = xarray.DataArray.from_iris(tobac.testing.make_simple_sample_data_2D())
+    elif input_types[0] == Cube:
+        data = tobac.testing.make_simple_sample_data_2D()
+    elif input_types[0] == xarray.Dataset:
+        data = tobac.testing.generate_single_feature(
+            1, 1, max_h1=100, max_h2=100
+        ).to_xarray()
+    elif input_types[0] == pd.DataFrame:
+        data = tobac.testing.generate_single_feature(1, 1, max_h1=100, max_h2=100)
+
+    if input_types[1] == xarray.DataArray:
+        kwarg = xarray.DataArray.from_iris(tobac.testing.make_simple_sample_data_2D())
+    elif input_types[1] == Cube:
+        kwarg = tobac.testing.make_simple_sample_data_2D()
+    elif input_types[1] == xarray.Dataset:
+        kwarg = tobac.testing.generate_single_feature(
+            1, 1, max_h1=100, max_h2=100
+        ).to_xarray()
+    elif input_types[1] == pd.DataFrame:
+        kwarg = tobac.testing.generate_single_feature(1, 1, max_h1=100, max_h2=100)
+
+    output = decorated_function_kwarg(data, kwarg=kwarg)
+    tuple_output = decorated_function_tuple(data, kwarg=kwarg)
+
+    assert (
+        type(output) == expected_output_type
+    ), "Expected output type {}, got {} for {}".format(
+        expected_output_type, type(output), decorator.__name__
+    )
+    assert (
+        type(tuple_output[0]) == expected_output_type
+    ), "Expected output type {}, but got {} for {} (1st tuple output(".format(
+        expected_output_type, type(tuple_output[0]), decorator.__name__
+    )
+    assert (
+        type(tuple_output[1]) == expected_output_type
+    ), "Expected output type {}, but got {} for {} (2nd tuple output(".format(
+        expected_output_type, type(tuple_output[1]), decorator.__name__
+    )
+
+
+def test_xarray_workflow():
+    """Test comparing the outputs of the standard functions of tobac for a test dataset
+    with the output of the same functions decorated with tobac.utils.xarray_to_iris"""
+
+    data = tobac.testing.make_sample_data_2D_3blobs()
+    data_xarray = xarray.DataArray.from_iris(deepcopy(data))
+
+    # Testing the get_spacings utility
+    xarray_to_iris_i = xarray_to_iris()
+    get_spacings_xarray = xarray_to_iris_i(tobac.utils.get_spacings)
+    dxy, dt = tobac.utils.get_spacings(data)
+    dxy_xarray, dt_xarray = get_spacings_xarray(data_xarray)
+
+    assert dxy == dxy_xarray
+    assert dt == dt_xarray
+
+    # Testing feature detection
+    feature_detection_xarray = xarray_to_iris_i(
+        tobac.feature_detection.feature_detection_multithreshold
+    )
+    features = tobac.feature_detection.feature_detection_multithreshold(
+        data, dxy, threshold=1.0
+    )
+    features_xarray = feature_detection_xarray(data_xarray, dxy_xarray, threshold=1.0)
+
+    assert_frame_equal(features, features_xarray)
+
+    # Testing the segmentation
+    segmentation_xarray = xarray_to_iris_i(tobac.segmentation.segmentation)
+    mask, features = tobac.segmentation.segmentation(features, data, dxy, threshold=1.0)
+    mask_xarray, features_xarray = segmentation_xarray(
+        features_xarray, data_xarray, dxy_xarray, threshold=1.0
+    )
+
+    assert (mask.data == mask_xarray.to_iris().data).all()
+
+    # testing tracking
+    tracking_xarray = xarray_to_iris_i(tobac.tracking.linking_trackpy)
+    track = tobac.tracking.linking_trackpy(features, data, dt, dxy, v_max=100.0)
+    track_xarray = tracking_xarray(
+        features_xarray, data_xarray, dt_xarray, dxy_xarray, v_max=100.0
+    )
+
+    assert_frame_equal(track, track_xarray)

tobac/tests/test_datetime.py

@@ -0,0 +1,216 @@
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+import cftime
+import pytest
+
+import tobac.utils.datetime as datetime_utils
+
+
+def test_to_cftime():
+    """Test conversion of datetime types to cftime calendars"""
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_cftime(date, "standard") == cftime.datetime(2000, 1, 1)
+        assert datetime_utils.to_cftime(date, "gregorian") == cftime.DatetimeGregorian(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime(date, "360_day") == cftime.Datetime360Day(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime(date, "365_day") == cftime.DatetimeNoLeap(
+            2000, 1, 1
+        )
+
+    # Test array-like input
+    for date in test_dates:
+        assert datetime_utils.to_cftime([date], "standard")[0] == cftime.datetime(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime([date], "gregorian")[
+            0
+        ] == cftime.DatetimeGregorian(2000, 1, 1)
+        assert datetime_utils.to_cftime([date], "360_day")[0] == cftime.Datetime360Day(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime([date], "365_day")[0] == cftime.DatetimeNoLeap(
+            2000, 1, 1
+        )
+
+
+def test_to_timestamp():
+    """Test conversion of various datetime types to pandas timestamps"""
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_timestamp(date) == pd.to_datetime("2000-01-01")
+
+    # Test array input
+    for date in test_dates:
+        assert datetime_utils.to_timestamp([date])[0] == pd.to_datetime("2000-01-01")
+
+
+def test_to_datetime():
+    """Test conversion of various datetime types to python datetime"""
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_datetime(date) == datetime(2000, 1, 1)
+
+    # Test array input
+    for date in test_dates:
+        assert datetime_utils.to_datetime([date])[0] == datetime(2000, 1, 1)
+
+
+def test_to_datetime64():
+    """Test conversion of various datetime types to numpy datetime64"""
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_datetime64(date) == np.datetime64(
+            "2000-01-01 00:00:00.000000000"
+        )
+
+    # Test array input
+    for date in test_dates:
+        assert datetime_utils.to_datetime64([date])[0] == np.datetime64(
+            "2000-01-01 00:00:00.000000000"
+        )
+
+
+def test_to_datestr():
+    """Test conversion of various datetime types to ISO format datestring"""
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert (
+            datetime_utils.to_datestr(date) == "2000-01-01T00:00:00.000000000"
+            or datetime_utils.to_datestr(date) == "2000-01-01T00:00:00"
+        )
+
+
+def test_to_datestr_array():
+    """Test conversion of arrays of various datetime types to ISO format
+    datestring
+    """
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+    for date in test_dates:
+        assert datetime_utils.to_datestr([date]) == [
+            "2000-01-01T00:00:00.000000000"
+        ] or datetime_utils.to_datestr([date]) == ["2000-01-01T00:00:00"]
+
+
+def test_match_datetime_format():
+    """Test match_datetime_format for various datetime-like combinations"""
+    test_dates = [
+        "2000-01-01T00:00:00.000000000",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for target in test_dates:
+        for date in test_dates:
+            assert datetime_utils.match_datetime_format(date, target) == target
+
+
+def test_match_datetime_format_array():
+    """Test match_datetime_format for various datetime-like combinations with
+    array input
+    """
+    test_dates = [
+        "2000-01-01T00:00:00.000000000",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for target in test_dates:
+        for date in test_dates:
+            assert datetime_utils.match_datetime_format([date], [target]) == np.array(
+                [target]
+            )
+
+
+def test_match_datetime_format_error():
+    """Test that if a non datetime-like object is provided as target to
+    match_datetime_format that a ValueError is raised:
+    """
+    with pytest.raises(ValueError, match="Target is not a valid datetime*"):
+        datetime_utils.match_datetime_format(datetime(2000, 1, 1), 1.5)

tobac/tests/test_decorators.py

@@ -0,0 +1,148 @@
+"""
+Tests for tobac.utils.decorators
+"""
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+import iris
+
+from tobac.utils import decorators
+
+
+def test_convert_cube_to_dataarray():
+    test_da_float = xr.DataArray(np.arange(15, dtype=float).reshape(3, 5) + 0.5)
+    test_da_int = xr.DataArray(np.arange(15, dtype=int).reshape(3, 5))
+
+    assert np.all(
+        decorators.convert_cube_to_dataarray(test_da_float.to_iris())
+        == test_da_float.values
+    )
+    assert np.all(
+        decorators.convert_cube_to_dataarray(test_da_int.to_iris())
+        == test_da_int.values
+    )
+
+
+def test_conv_kwargs_iris_to_xarray():
+    assert decorators._conv_kwargs_iris_to_xarray({}) == {}
+    assert decorators._conv_kwargs_iris_to_xarray(dict(test_int=1)) == dict(test_int=1)
+
+    test_da = xr.DataArray(np.arange(5))
+
+    test_xr_kwarg = decorators._conv_kwargs_iris_to_xarray(dict(test_xr=test_da))
+    assert isinstance(test_xr_kwarg["test_xr"], xr.DataArray)
+
+    test_iris_kwarg = decorators._conv_kwargs_iris_to_xarray(
+        dict(test_iris=test_da.to_iris())
+    )
+    assert isinstance(test_iris_kwarg["test_iris"], xr.DataArray)
+
+
+def test_conv_kwargs_irispandas_to_xarray():
+    assert decorators._conv_kwargs_irispandas_to_xarray({}) == {}
+    assert decorators._conv_kwargs_irispandas_to_xarray(dict(test_int=1)) == dict(
+        test_int=1
+    )
+
+    test_da = xr.DataArray(np.arange(5))
+
+    test_xr_kwarg = decorators._conv_kwargs_irispandas_to_xarray(dict(test_xr=test_da))
+    assert isinstance(test_xr_kwarg["test_xr"], xr.DataArray)
+
+    test_iris_kwarg = decorators._conv_kwargs_irispandas_to_xarray(
+        dict(test_iris=test_da.to_iris())
+    )
+    assert isinstance(test_iris_kwarg["test_iris"], xr.DataArray)
+
+    test_ds = xr.Dataset({"test": test_da})
+    test_ds_kwarg = decorators._conv_kwargs_irispandas_to_xarray(dict(test_xr=test_ds))
+    assert isinstance(test_ds_kwarg["test_xr"], xr.Dataset)
+
+    test_pd_kwarg = decorators._conv_kwargs_irispandas_to_xarray(
+        dict(test_pd=test_ds.to_pandas())
+    )
+    assert isinstance(test_pd_kwarg["test_pd"], xr.Dataset)
+
+
+def test_conv_kwargs_xarray_to_iris():
+    assert decorators._conv_kwargs_xarray_to_iris({}) == {}
+    assert decorators._conv_kwargs_xarray_to_iris(dict(test_int=1)) == dict(test_int=1)
+
+    test_da = xr.DataArray(np.arange(5))
+
+    test_xr_kwarg = decorators._conv_kwargs_xarray_to_iris(dict(test_xr=test_da))
+    assert isinstance(test_xr_kwarg["test_xr"], iris.cube.Cube)
+
+    test_iris_kwarg = decorators._conv_kwargs_xarray_to_iris(
+        dict(test_iris=test_da.to_iris())
+    )
+    assert isinstance(test_iris_kwarg["test_iris"], iris.cube.Cube)
+
+
+def test_conv_kwargs_xarray_to_irispandas():
+    assert decorators._conv_kwargs_xarray_to_irispandas({}) == {}
+    assert decorators._conv_kwargs_xarray_to_irispandas(dict(test_int=1)) == dict(
+        test_int=1
+    )
+
+    test_da = xr.DataArray(np.arange(5))
+
+    test_xr_kwarg = decorators._conv_kwargs_xarray_to_irispandas(dict(test_xr=test_da))
+    assert isinstance(test_xr_kwarg["test_xr"], iris.cube.Cube)
+
+    test_iris_kwarg = decorators._conv_kwargs_xarray_to_irispandas(
+        dict(test_iris=test_da.to_iris())
+    )
+    assert isinstance(test_iris_kwarg["test_iris"], iris.cube.Cube)
+
+    test_ds = xr.Dataset({"test": test_da})
+    test_ds_kwarg = decorators._conv_kwargs_xarray_to_irispandas(dict(test_xr=test_ds))
+    assert isinstance(test_ds_kwarg["test_xr"], pd.DataFrame)
+
+    test_pd_kwarg = decorators._conv_kwargs_xarray_to_irispandas(
+        dict(test_pd=test_ds.to_pandas())
+    )
+    assert isinstance(test_pd_kwarg["test_pd"], pd.DataFrame)
+
+
+@decorators.iris_to_xarray(save_iris_info=True)
+def _test_iris_to_xarray(*args, **kwargs):
+    return kwargs["converted_from_iris"]
+
+
+def test_iris_to_xarray():
+    test_da = xr.DataArray(np.arange(5))
+
+    assert _test_iris_to_xarray(test_da) == False
+    assert _test_iris_to_xarray(kwarg_xr=test_da) == False
+
+    assert _test_iris_to_xarray(test_da.to_iris()) == True
+    assert _test_iris_to_xarray(kwarg_ir=test_da.to_iris()) == True
+
+
+@decorators.irispandas_to_xarray(save_iris_info=True)
+def _test_irispandas_to_xarray(*args, **kwargs):
+    return kwargs["converted_from_iris"]
+
+
+def test_irispandas_to_xarray():
+    test_da = xr.DataArray(np.arange(5))
+
+    assert _test_irispandas_to_xarray(test_da) == False
+    assert _test_irispandas_to_xarray(kwarg_xr=test_da) == False
+
+    assert _test_irispandas_to_xarray(test_da.to_iris()) == True
+    assert _test_irispandas_to_xarray(kwarg_ir=test_da.to_iris()) == True
+
+
+@decorators.xarray_to_irispandas()
+def _test_xarray_to_irispandas(*args, **kwargs):
+    return args, kwargs
+
+
+def test_xarray_to_irispandas():
+    test_da = xr.DataArray(np.arange(5, dtype=float))
+
+    assert isinstance(_test_xarray_to_irispandas(test_da)[0][0], iris.cube.Cube)
+    assert _test_xarray_to_irispandas(test_da)[1] == {}