diff --git a/test/test_helpers/test_helpers.py b/test/test_helpers/test_helpers.py index a5aaa707c83a65c3e10f76fdbfcd8142d3615e2e..7e6b1b3b77af69c68e4de9369b29b12b7f501878 100644 --- a/test/test_helpers/test_helpers.py +++ b/test/test_helpers/test_helpers.py @@ -9,10 +9,13 @@ import os import mock import pytest +import string from mlair.helpers import to_list, dict_to_xarray, float_round, remove_items, extract_value, select_from_dict from mlair.helpers import PyTestRegex from mlair.helpers import Logger, TimeTracking +from mlair.helpers.helpers import is_xarray, convert2xrda, convert_1d_np2xr + class TestToList: @@ -312,3 +315,105 @@ class TestExtractValue: extract_value([1, 2, 3]) assert "Trying to extract an encapsulated value from objects with more than a single entry is not supported " \ "by this function." in e.value.args[0] + + +class TestIsXarray: + + @pytest.fixture + def custom_xr_data(self): + return xr.DataArray(np.array(range(5))) + + def test_is_xarray_xr_input(self, custom_xr_data): + # data_array = xr.DataArray(np.array(range(5))) + assert is_xarray(custom_xr_data) is True + assert is_xarray(xr.Dataset({'test': custom_xr_data})) is True + + def test_is_xarray_other_input(self, custom_xr_data): + assert is_xarray(1) is False + assert is_xarray(1.) is False + assert is_xarray([1, 2.]) is False + assert is_xarray([custom_xr_data]) is False + + +class TestConvert2xrDa: + + @pytest.fixture + def custom_1d_npdata(self): + return np.array(range(9)) + + @pytest.fixture() + def custom_2d_npdata(self, custom_1d_npdata): + return np.stack([custom_1d_npdata, 2 * custom_1d_npdata]) + + @pytest.fixture + def custom_xr_dataarray(self, custom_1d_npdata): + return xr.DataArray(custom_1d_npdata) + + @pytest.fixture + def custom_xr_dataset(self, custom_xr_dataarray): + return xr.Dataset({'test_1': custom_xr_dataarray}) + + def test_convert2xrda_xrdata_in(self, custom_xr_dataarray, custom_xr_dataset): + assert (convert2xrda(custom_xr_dataarray) == custom_xr_dataarray).all() + assert (convert2xrda(custom_xr_dataset) == custom_xr_dataset).all() + + def test_convert2xrda_npdata_in_nokwargs(self, custom_1d_npdata, custom_2d_npdata): + converted_data = convert2xrda(custom_1d_npdata) + assert isinstance(converted_data, xr.DataArray) + assert (converted_data.values == custom_1d_npdata).all() + assert converted_data.dims == ('dim_0',) + assert converted_data.dim_0.size == custom_1d_npdata.shape[0] + + # Feed in a 2D-np.array without additional kwargs + converted_data = convert2xrda(custom_2d_npdata) + assert isinstance(converted_data, xr.DataArray) + assert (converted_data.values == custom_2d_npdata).all() + assert converted_data.dims == ('dim_0', 'dim_1') + assert converted_data.dim_0.size == custom_2d_npdata.shape[0] + assert converted_data.dim_1.size == custom_2d_npdata.shape[1] + + def test_convert2xrda_npdata_in_nokwargs_default_true(self, custom_1d_npdata, custom_2d_npdata): + converted_data = convert2xrda(custom_1d_npdata, use_1d_default=True) + assert isinstance(converted_data, xr.DataArray) + assert (converted_data.values == custom_1d_npdata).all() + assert converted_data.dims == ('points',) + assert converted_data.points.size == custom_1d_npdata.shape[0] + + # Feed in a 2D-np.array without additional kwargs + with pytest.raises(ValueError) as e: + converted_data = convert2xrda(custom_2d_npdata, use_1d_default=True) + assert "different number of dimensions on data and dims: 2 vs 1" in e.value.args[0] + + @pytest.mark.parametrize("use_1d_default", (False, True)) + def test_convert2xrda_npdata_in_kwargs(self, custom_1d_npdata, custom_2d_npdata, use_1d_default): + converted_data = convert2xrda(custom_1d_npdata, use_1d_default=use_1d_default, dims='other_points') + assert isinstance(converted_data, xr.DataArray) + assert (converted_data.values == custom_1d_npdata).all() + assert converted_data.dims == ('other_points',) + assert converted_data.other_points.size == custom_1d_npdata.shape[0] + + # Feed in a 2D-np.array with correct additional kwargs + converted_data = convert2xrda(custom_2d_npdata, use_1d_default=use_1d_default, + dims=['test_dim_0', 'test_dim_1'], + coords={'test_dim_0': list(string.ascii_lowercase[:custom_2d_npdata.shape[0]]), + 'test_dim_1': list(string.ascii_lowercase[:custom_2d_npdata.shape[1]]), + }, + ) + assert isinstance(converted_data, xr.DataArray) + assert (converted_data.values == custom_2d_npdata).all() + assert converted_data.dims == ('test_dim_0', 'test_dim_1') + assert (converted_data.coords['test_dim_0'].values == np.array(['a', 'b'])).all() + + @pytest.mark.parametrize("scalar", (1, 2.)) + def test_convert2xrda_int_float_in_nokwargs_default_true(self, scalar): + converted_data = convert2xrda(scalar, use_1d_default=True) + assert isinstance(converted_data, xr.DataArray) + assert converted_data.values == np.array([scalar]) + assert converted_data.dims == ('points',) + + @pytest.mark.parametrize("wrong_input", ({1: 'b'}, [1], 'abc')) + def test_convert2xrda_wrong_type_in_default_true_nokwargs(self, wrong_input): + with pytest.raises(TypeError) as e: + converted_data = convert2xrda(wrong_input, use_1d_default=True) + assert f"`arr' must be arry-like, int or float. But is of type {type(wrong_input)}" in e.value.args[0] +