diff --git a/pysteps/tests/test_utils_conversion.py b/pysteps/tests/test_utils_conversion.py index 169cdb50e..80d052b77 100644 --- a/pysteps/tests/test_utils_conversion.py +++ b/pysteps/tests/test_utils_conversion.py @@ -1,119 +1,328 @@ # -*- coding: utf-8 -*- - import numpy as np import pytest +import xarray as xr from numpy.testing import assert_array_almost_equal +from pysteps.tests.helpers import assert_dataset_equivalent from pysteps.utils import conversion # to_rainrate -test_data = [ +test_data_to_rainrate = [ ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([1]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([12]), + xr.Dataset( + data_vars={ + "precip_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([12.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 12.0, + "zerovalue": 12.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([1.25892541]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": "dB", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.25892541]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1.25892541, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([15.10710494]), + xr.Dataset( + data_vars={ + "precip_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": "dB", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([15.10710494]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 15.10710494, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "dBZ", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.04210719]), + xr.Dataset( + data_vars={ + "reflectivity": ( + ["x"], + np.array([1.0]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([0.04210719]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 0.04210719, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "log", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([2.71828183]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": "log", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([2.71828183]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 2.71828183, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1.0]), - { - "accutime": 5, - "transform": "log", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([32.61938194]), + xr.Dataset( + data_vars={ + "precip_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": "log", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([32.61938194]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 32.61938194, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([1]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": "sqrt", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1, + "zerovalue": 1, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1.0]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([12.0]), + xr.Dataset( + data_vars={ + "precip_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": "sqrt", + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([12.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 12, + "zerovalue": 12, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ] -@pytest.mark.parametrize("R, metadata, expected", test_data) -def test_to_rainrate(R, metadata, expected): +@pytest.mark.parametrize("dataset, expected", test_data_to_rainrate) +def test_to_rainrate(dataset, expected): """Test the to_rainrate.""" - assert_array_almost_equal(conversion.to_rainrate(R, metadata)[0], expected) + actual = conversion.to_rainrate(dataset) + assert_dataset_equivalent(actual, expected) # to_raindepth diff --git a/pysteps/utils/conversion.py b/pysteps/utils/conversion.py index 527c06b9a..87fe22deb 100644 --- a/pysteps/utils/conversion.py +++ b/pysteps/utils/conversion.py @@ -108,6 +108,10 @@ def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): else: raise ValueError(f'Unknown transformation {metadata["transform"]}') + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values + if metadata["units"] == "mm/h": pass