From 8b92e71bed50e72e5f91c65696eadcce0c2591db Mon Sep 17 00:00:00 2001 From: Mats Veldhuizen Date: Wed, 17 Jul 2024 16:27:51 +0200 Subject: [PATCH 1/8] import directly to xarray and rewrite utils_transformation to work with it --- ci/ci_test_env.yml | 1 + pysteps/converters.py | 202 ++++++++++ pysteps/decorators.py | 6 +- pysteps/io/importers.py | 4 +- pysteps/io/readers.py | 63 +-- pysteps/nowcasts/anvil.py | 5 +- pysteps/tests/helpers.py | 39 +- pysteps/tests/test_exporters.py | 14 +- pysteps/tests/test_utils_transformation.py | 424 +++++++++++++++------ pysteps/utils/conversion.py | 93 +++-- pysteps/utils/transformation.py | 220 +++++------ requirements.txt | 1 + requirements_dev.txt | 1 + 13 files changed, 754 insertions(+), 319 deletions(-) create mode 100644 pysteps/converters.py diff --git a/ci/ci_test_env.yml b/ci/ci_test_env.yml index 4c0b8732f..7857de61b 100644 --- a/ci/ci_test_env.yml +++ b/ci/ci_test_env.yml @@ -18,6 +18,7 @@ dependencies: - pillow - pyproj - scipy + - xarray # Optional dependencies - dask - pyfftw diff --git a/pysteps/converters.py b/pysteps/converters.py new file mode 100644 index 000000000..f27e1572e --- /dev/null +++ b/pysteps/converters.py @@ -0,0 +1,202 @@ +# -*- coding: utf-8 -*- +""" +pysteps.converters +================== + +Module with data converter functions. + +.. autosummary:: + :toctree: ../generated/ + + convert_to_xarray_dataset +""" + +import numpy as np +import pyproj +import xarray as xr + +from pysteps.utils.conversion import cf_parameters_from_unit + +# TODO(converters): Write methods for converting Proj.4 projection definitions +# into CF grid mapping attributes. Currently this has been implemented for +# the stereographic projection. +# The conversions implemented here are take from: +# https://github.com/cf-convention/cf-convention.github.io/blob/master/wkt-proj-4.md + + +def _convert_proj4_to_grid_mapping(proj4str): + tokens = proj4str.split("+") + + d = {} + for t in tokens[1:]: + t = t.split("=") + if len(t) > 1: + d[t[0]] = t[1].strip() + + params = {} + # TODO(exporters): implement more projection types here + if d["proj"] == "stere": + grid_mapping_var_name = "polar_stereographic" + grid_mapping_name = "polar_stereographic" + v = d["lon_0"] if d["lon_0"][-1] not in ["E", "W"] else d["lon_0"][:-1] + params["straight_vertical_longitude_from_pole"] = float(v) + v = d["lat_0"] if d["lat_0"][-1] not in ["N", "S"] else d["lat_0"][:-1] + params["latitude_of_projection_origin"] = float(v) + if "lat_ts" in list(d.keys()): + params["standard_parallel"] = float(d["lat_ts"]) + elif "k_0" in list(d.keys()): + params["scale_factor_at_projection_origin"] = float(d["k_0"]) + params["false_easting"] = float(d["x_0"]) + params["false_northing"] = float(d["y_0"]) + elif d["proj"] == "aea": # Albers Conical Equal Area + grid_mapping_var_name = "proj" + grid_mapping_name = "albers_conical_equal_area" + params["false_easting"] = float(d["x_0"]) if "x_0" in d else float(0) + params["false_northing"] = float(d["y_0"]) if "y_0" in d else float(0) + v = d["lon_0"] if "lon_0" in d else float(0) + params["longitude_of_central_meridian"] = float(v) + v = d["lat_0"] if "lat_0" in d else float(0) + params["latitude_of_projection_origin"] = float(v) + v1 = d["lat_1"] if "lat_1" in d else float(0) + v2 = d["lat_2"] if "lat_2" in d else float(0) + params["standard_parallel"] = (float(v1), float(v2)) + else: + print("unknown projection", d["proj"]) + return None, None, None + + return grid_mapping_var_name, grid_mapping_name, params + + +def convert_to_xarray_dataset( + precip: np.ndarray, + quality: np.ndarray | None, + metadata: dict[str, str | float | None], +) -> xr.Dataset: + """ + Read a precip, quality, metadata tuple as returned by the importers + (:py:mod:`pysteps.io.importers`) and return an xarray dataset containing + this data. + + Parameters + ---------- + precip: array + 2D array containing imported precipitation data. + quality: array, None + 2D array containing the quality values of the imported precipitation + data, can be None. + metadata: dict + Metadata dictionary containing the attributes described in the + documentation of :py:mod:`pysteps.io.importers`. + + Returns + ------- + out: Dataset + A CF compliant xarray dataset, which contains all data and metadata. + + """ + var_name, attrs = cf_parameters_from_unit(metadata["unit"]) + h, w = precip.shape + x_r = np.linspace(metadata["x1"], metadata["x2"], w + 1)[:-1] + x_r += 0.5 * (x_r[1] - x_r[0]) + y_r = np.linspace(metadata["y1"], metadata["y2"], h + 1)[:-1] + y_r += 0.5 * (y_r[1] - y_r[0]) + + # flip yr vector if yorigin is upper + if metadata["yorigin"] == "upper": + y_r = np.flip(y_r) + + x_2d, y_2d = np.meshgrid(x_r, y_r) + pr = pyproj.Proj(metadata["projection"]) + lon, lat = pr(x_2d.flatten(), y_2d.flatten(), inverse=True) + + ( + grid_mapping_var_name, + grid_mapping_name, + grid_mapping_params, + ) = _convert_proj4_to_grid_mapping(metadata["projection"]) + + data_vars = { + var_name: ( + ["y", "x"], + precip, + { + "units": attrs["units"], + "standard_name": attrs["standard_name"], + "long_name": attrs["long_name"], + "grid_mapping": "projection", + "transform": metadata["transform"], + "accutime": metadata["accutime"], + "threshold": metadata["threshold"], + "zerovalue": metadata["zerovalue"], + "zr_a": metadata["zr_a"], + "zr_b": metadata["zr_b"], + }, + ) + } + if quality is not None: + data_vars["quality"] = ( + ["y", "x"], + quality, + { + "units": "1", + "standard_name": "quality_flag", + "grid_mapping": "projection", + }, + ) + coords = { + "y": ( + ["y"], + y_r, + { + "axis": "Y", + "long_name": "y-coordinate in Cartesian system", + "standard_name": "projection_y_coordinate", + "units": metadata["cartesian_unit"], + }, + ), + "x": ( + ["x"], + x_r, + { + "axis": "X", + "long_name": "x-coordinate in Cartesian system", + "standard_name": "projection_x_coordinate", + "units": metadata["cartesian_unit"], + }, + ), + "lon": ( + ["y", "x"], + lon.reshape(precip.shape), + { + "long_name": "longitude coordinate", + "standard_name": "longitude", + # TODO(converters): Don't hard-code the unit. + "units": "degrees_east", + }, + ), + "lat": ( + ["y", "x"], + lat.reshape(precip.shape), + { + "long_name": "latitude coordinate", + "standard_name": "latitude", + # TODO(converters): Don't hard-code the unit. + "units": "degrees_north", + }, + ), + } + if grid_mapping_var_name is not None: + coords[grid_mapping_name] = ( + ( + [], + None, + {"grid_mapping_name": grid_mapping_name, **grid_mapping_params}, + ), + ) + attrs = { + "Conventions": "CF-1.7", + "institution": metadata["institution"], + "projection": metadata["projection"], + "precip_var": var_name, + } + return xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs) diff --git a/pysteps/decorators.py b/pysteps/decorators.py index 44fbaebdb..ee421977e 100644 --- a/pysteps/decorators.py +++ b/pysteps/decorators.py @@ -22,6 +22,8 @@ import numpy as np +from pysteps.converters import convert_to_xarray_dataset + def _add_extra_kwrds_to_docstrings(target_func, extra_kwargs_doc_text): """ @@ -66,7 +68,7 @@ def postprocess_import(fillna=np.nan, dtype="double"): def _postprocess_import(importer): @wraps(importer) def _import_with_postprocessing(*args, **kwargs): - precip, *other_args = importer(*args, **kwargs) + precip, quality, metadata = importer(*args, **kwargs) _dtype = kwargs.get("dtype", dtype) @@ -88,7 +90,7 @@ def _import_with_postprocessing(*args, **kwargs): mask = ~np.isfinite(precip) precip[mask] = _fillna - return (precip.astype(_dtype),) + tuple(other_args) + return convert_to_xarray_dataset(precip.astype(_dtype), quality, metadata) extra_kwargs_doc = """ Other Parameters diff --git a/pysteps/io/importers.py b/pysteps/io/importers.py index 1493b6d4b..5c2928203 100644 --- a/pysteps/io/importers.py +++ b/pysteps/io/importers.py @@ -89,12 +89,10 @@ from functools import partial import numpy as np - from matplotlib.pyplot import imread from pysteps.decorators import postprocess_import -from pysteps.exceptions import DataModelError -from pysteps.exceptions import MissingOptionalDependency +from pysteps.exceptions import DataModelError, MissingOptionalDependency from pysteps.utils import aggregate_fields try: diff --git a/pysteps/io/readers.py b/pysteps/io/readers.py index fcc6bda2e..7295b5ff7 100644 --- a/pysteps/io/readers.py +++ b/pysteps/io/readers.py @@ -12,13 +12,14 @@ """ import numpy as np +import xarray as xr -def read_timeseries(inputfns, importer, **kwargs): +def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None: """ Read a time series of input files using the methods implemented in the - :py:mod:`pysteps.io.importers` module and stack them into a 3d array of - shape (num_timesteps, height, width). + :py:mod:`pysteps.io.importers` module and stack them into a 3d xarray + dataset of shape (num_timesteps, height, width). Parameters ---------- @@ -32,50 +33,50 @@ def read_timeseries(inputfns, importer, **kwargs): Returns ------- - out: tuple - A three-element tuple containing the read data and quality rasters and + out: Dataset + A dataset containing the read data and quality rasters and associated metadata. If an input file name is None, the corresponding precipitation and quality fields are filled with nan values. If all input file names are None or if the length of the file name list is - zero, a three-element tuple containing None values is returned. + zero, None is returned. """ # check for missing data - precip_ref = None + dataset_ref = None if all(ifn is None for ifn in inputfns): - return None, None, None + return None else: if len(inputfns[0]) == 0: - return None, None, None + return None for ifn in inputfns[0]: if ifn is not None: - precip_ref, quality_ref, metadata = importer(ifn, **kwargs) + dataset_ref = importer(ifn, **kwargs) break - if precip_ref is None: - return None, None, None + if dataset_ref is None: + return None - precip = [] - quality = [] - timestamps = [] + startdate = min(inputfns[1]) + + datasets = [] for i, ifn in enumerate(inputfns[0]): if ifn is not None: - precip_, quality_, _ = importer(ifn, **kwargs) - precip.append(precip_) - quality.append(quality_) - timestamps.append(inputfns[1][i]) + dataset_ = importer(ifn, **kwargs) else: - precip.append(precip_ref * np.nan) - if quality_ref is not None: - quality.append(quality_ref * np.nan) - else: - quality.append(None) - timestamps.append(inputfns[1][i]) - - # Replace this with stack? - precip = np.concatenate([precip_[None, :, :] for precip_ in precip]) - # TODO: Q should be organized as R, but this is not trivial as Q_ can be also None or a scalar - metadata["timestamps"] = np.array(timestamps) + dataset_ = dataset_ref * np.nan + dataset_ = dataset_.expand_dims(dim="time", axis=0) + dataset_ = dataset_.assign_coords( + time=( + "time", + [inputfns[1][i]], + { + "long_name": "forecast time", + "units": f"seconds since {startdate:%Y-%m-%d %H:%M:%S}", + }, + ) + ) + datasets.append(dataset_) - return precip, quality, metadata + dataset = xr.concat(datasets, dim="time") + return dataset diff --git a/pysteps/nowcasts/anvil.py b/pysteps/nowcasts/anvil.py index 9da0fb47e..f5af038bb 100644 --- a/pysteps/nowcasts/anvil.py +++ b/pysteps/nowcasts/anvil.py @@ -19,12 +19,13 @@ """ import time + import numpy as np from scipy.ndimage import gaussian_filter -from pysteps import cascade, extrapolation + +from pysteps import cascade, extrapolation, utils from pysteps.nowcasts.utils import nowcast_main_loop from pysteps.timeseries import autoregression -from pysteps import utils try: import dask diff --git a/pysteps/tests/helpers.py b/pysteps/tests/helpers.py index 85bd861f5..2ae10026b 100644 --- a/pysteps/tests/helpers.py +++ b/pysteps/tests/helpers.py @@ -9,6 +9,7 @@ import numpy as np import pytest +import xarray as xr import pysteps as stp from pysteps import io, rcparams @@ -24,6 +25,32 @@ _reference_dates["mrms"] = datetime(2019, 6, 10, 0, 0) +def assert_dataset_equivalent(dataset1: xr.Dataset, dataset2: xr.Dataset) -> None: + xr.testing.assert_allclose(dataset1, dataset2) + assert np.isclose( + dataset1["precip_intensity"].attrs["threshold"], + dataset2["precip_intensity"].attrs["threshold"], + ) + assert ( + dataset1["precip_intensity"].attrs["units"] + == dataset2["precip_intensity"].attrs["units"] + ) + assert ( + dataset1["precip_intensity"].attrs["transform"] + == dataset2["precip_intensity"].attrs["transform"] + or dataset1["precip_intensity"].attrs["transform"] is None + and dataset2["precip_intensity"].attrs["transform"] is None + ) + assert ( + dataset1["precip_intensity"].attrs["accutime"] + == dataset2["precip_intensity"].attrs["accutime"] + ) + assert ( + dataset1["precip_intensity"].attrs["zerovalue"] + == dataset2["precip_intensity"].attrs["zerovalue"] + ) + + def get_precipitation_fields( num_prev_files=0, num_next_files=0, @@ -161,9 +188,7 @@ def get_precipitation_fields( # Read the radar composites importer = io.get_method(importer_name, "importer") - reference_field, __, ref_metadata = io.read_timeseries( - fns, importer, **_importer_kwargs - ) + ref_dataset = io.read_timeseries(fns, importer, **_importer_kwargs) if not return_raw: if (num_prev_files == 0) and (num_next_files == 0): @@ -171,14 +196,10 @@ def get_precipitation_fields( reference_field = np.squeeze(reference_field) # Convert to mm/h - reference_field, ref_metadata = stp.utils.to_rainrate( - reference_field, ref_metadata - ) + ref_dataset = stp.utils.to_rainrate(ref_dataset) # Clip domain - reference_field, ref_metadata = stp.utils.clip_domain( - reference_field, ref_metadata, clip - ) + ref_dataset = stp.utils.clip_domain(ref_dataset, clip) # Upscale data reference_field, ref_metadata = aggregate_fields_space( diff --git a/pysteps/tests/test_exporters.py b/pysteps/tests/test_exporters.py index 10e87d46e..dfe7e8ace 100644 --- a/pysteps/tests/test_exporters.py +++ b/pysteps/tests/test_exporters.py @@ -9,12 +9,14 @@ from numpy.testing import assert_array_almost_equal from pysteps.io import import_netcdf_pysteps -from pysteps.io.exporters import _get_geotiff_filename -from pysteps.io.exporters import close_forecast_files -from pysteps.io.exporters import export_forecast_dataset -from pysteps.io.exporters import initialize_forecast_exporter_netcdf -from pysteps.io.exporters import _convert_proj4_to_grid_mapping -from pysteps.tests.helpers import get_precipitation_fields, get_invalid_mask +from pysteps.io.exporters import ( + _convert_proj4_to_grid_mapping, + _get_geotiff_filename, + close_forecast_files, + export_forecast_dataset, + initialize_forecast_exporter_netcdf, +) +from pysteps.tests.helpers import get_invalid_mask, get_precipitation_fields # Test arguments exporter_arg_names = ( diff --git a/pysteps/tests/test_utils_transformation.py b/pysteps/tests/test_utils_transformation.py index 101e6b9d5..29e6e639c 100644 --- a/pysteps/tests/test_utils_transformation.py +++ b/pysteps/tests/test_utils_transformation.py @@ -1,190 +1,392 @@ # -*- coding: utf-8 -*- - import numpy as np import pytest -from numpy.testing import assert_array_almost_equal +import xarray as xr +from pysteps.tests.helpers import assert_dataset_equivalent from pysteps.utils import transformation # boxcox_transform -test_data = [ +test_data_boxcox_transform = [ ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": np.e, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), None, None, None, False, - np.array([0]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([0.0]), + { + "units": "mm/h", + "transform": "BoxCox", + "accutime": 5, + "threshold": 1, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "BoxCox", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": "BoxCox", + "accutime": 5, + "threshold": 1, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), None, None, None, True, - np.array([np.exp(1)]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([np.exp(1.0)]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": np.e, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": np.e, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), 1.0, None, None, False, - np.array([0]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([np.e - 2]), + { + "units": "mm/h", + "transform": "BoxCox", + "accutime": 5, + "threshold": np.e - 1, + "zerovalue": np.e - 2, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "BoxCox", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([np.e - 2]), + { + "units": "mm/h", + "transform": "BoxCox", + "accutime": 5, + "threshold": np.e - 1, + "zerovalue": np.e - 2, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), 1.0, None, None, True, - np.array([2.0]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([0.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": np.e, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ] @pytest.mark.parametrize( - "R, metadata, Lambda, threshold, zerovalue, inverse, expected", test_data + "dataset, Lambda, threshold, zerovalue, inverse, expected", + test_data_boxcox_transform, ) -def test_boxcox_transform(R, metadata, Lambda, threshold, zerovalue, inverse, expected): +def test_boxcox_transform(dataset, Lambda, threshold, zerovalue, inverse, expected): """Test the boxcox_transform.""" - assert_array_almost_equal( - transformation.boxcox_transform( - R, metadata, Lambda, threshold, zerovalue, inverse - )[0], - expected, + actual = transformation.boxcox_transform( + dataset, Lambda, threshold, zerovalue, inverse ) + assert_dataset_equivalent(actual, expected) # dB_transform -test_data = [ +test_data_dB_transform = [ ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + 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"}, + ), None, None, False, - np.array([0]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([0.0]), + { + "units": "mm/h", + "transform": "dB", + "accutime": 5, + "threshold": 0, + "zerovalue": -5, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([0.0]), + { + "units": "mm/h", + "transform": "dB", + "accutime": 5, + "threshold": 0, + "zerovalue": -5, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), None, None, True, - np.array([1.25892541]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ] @pytest.mark.parametrize( - "R, metadata, threshold, zerovalue, inverse, expected", test_data + "dataset, threshold, zerovalue, inverse, expected", test_data_dB_transform ) -def test_dB_transform(R, metadata, threshold, zerovalue, inverse, expected): +def test_dB_transform(dataset, threshold, zerovalue, inverse, expected): """Test the dB_transform.""" - assert_array_almost_equal( - transformation.dB_transform(R, metadata, threshold, zerovalue, inverse)[0], - expected, - ) + actual = transformation.dB_transform(dataset, threshold, zerovalue, inverse) + assert_dataset_equivalent(actual, expected) # NQ_transform -test_data = [ +test_data_NQ_transform = [ ( - np.array([1, 2]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0, 2.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 0, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), False, - np.array([-0.4307273, 0.4307273]), - ) + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([-0.4307273, 0.4307273]), + { + "units": "mm/h", + "transform": "NQT", + "accutime": 5, + "threshold": 0.4307273, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), + ), ] -@pytest.mark.parametrize("R, metadata, inverse, expected", test_data) -def test_NQ_transform(R, metadata, inverse, expected): +@pytest.mark.parametrize("dataset, inverse, expected", test_data_NQ_transform) +def test_NQ_transform(dataset, inverse, expected): """Test the NQ_transform.""" - assert_array_almost_equal( - transformation.NQ_transform(R, metadata, inverse)[0], expected - ) + actual = transformation.NQ_transform(dataset, inverse) + assert_dataset_equivalent(actual, expected) # sqrt_transform -test_data = [ +test_data_sqrt_transform = [ ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0, 4.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 4, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), False, - np.array([1]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0, 2.0]), + { + "units": "mm/h", + "transform": "sqrt", + "accutime": 5, + "threshold": 2, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0, 2.0]), + { + "units": "mm/h", + "transform": "sqrt", + "accutime": 5, + "threshold": 2, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), True, - np.array([1]), + xr.Dataset( + data_vars={ + "precip_intensity": ( + ["x"], + np.array([1.0, 4.0]), + { + "units": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 4, + "zerovalue": 0, + }, + ) + }, + attrs={"precip_var": "precip_intensity"}, + ), ), ] -@pytest.mark.parametrize("R, metadata, inverse, expected", test_data) -def test_sqrt_transform(R, metadata, inverse, expected): +@pytest.mark.parametrize("dataset, inverse, expected", test_data_sqrt_transform) +def test_sqrt_transform(dataset, inverse, expected): """Test the sqrt_transform.""" - assert_array_almost_equal( - transformation.sqrt_transform(R, metadata, inverse)[0], expected - ) + actual = transformation.sqrt_transform(dataset, inverse) + assert_dataset_equivalent(actual, expected) diff --git a/pysteps/utils/conversion.py b/pysteps/utils/conversion.py index f8dfae23b..527c06b9a 100644 --- a/pysteps/utils/conversion.py +++ b/pysteps/utils/conversion.py @@ -14,6 +14,9 @@ """ import warnings + +import xarray as xr + from . import transformation # TODO: This should not be done. Instead fix the code so that it doesn't @@ -22,18 +25,53 @@ warnings.filterwarnings("ignore", category=RuntimeWarning) -def to_rainrate(R, metadata, zr_a=None, zr_b=None): +def cf_parameters_from_unit(unit: str) -> tuple[str, dict[str, str | None]]: + if unit == "mm/h": + var_name = "precip_intensity" + var_standard_name = None + var_long_name = "instantaneous precipitation rate" + var_unit = "mm/h" + elif unit == "mm": + var_name = "precip_accum" + var_standard_name = None + var_long_name = "accumulated precipitation" + var_unit = "mm" + elif unit == "dBZ": + var_name = "reflectivity" + var_long_name = "equivalent reflectivity factor" + var_standard_name = "equivalent_reflectivity_factor" + var_unit = "dBZ" + else: + raise ValueError(f"unknown unit {unit}") + + return var_name, { + "standard_name": var_standard_name, + "long_name": var_long_name, + "units": var_unit, + } + + +def _change_unit(dataset: xr.Dataset, precip_var: str, new_unit: str) -> xr.Dataset: + new_var, new_attrs = cf_parameters_from_unit(new_unit) + dataset = dataset.rename_vars({precip_var: new_var}) + dataset.attrs["precip_var"] = new_var + + dataset[new_var].attrs = { + **dataset[new_var].attrs, + **new_attrs, + } + + return dataset + + +def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): """ Convert to rain rate [mm/h]. Parameters ---------- - R: array-like - Array of any shape to be (back-)transformed. - metadata: dict - Metadata dictionary containing the accutime, transform, unit, threshold - and zerovalue attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be (back-)transformed. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -45,46 +83,46 @@ def to_rainrate(R, metadata, zr_a=None, zr_b=None): Returns ------- - R: array-like - Array of any shape containing the converted units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the converted units. """ - R = R.copy() - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values if metadata["transform"] is not None: if metadata["transform"] == "dB": - R, metadata = transformation.dB_transform(R, metadata, inverse=True) + dataset = transformation.dB_transform(dataset, inverse=True) elif metadata["transform"] in ["BoxCox", "log"]: - R, metadata = transformation.boxcox_transform(R, metadata, inverse=True) + dataset = transformation.boxcox_transform(dataset, inverse=True) elif metadata["transform"] == "NQT": - R, metadata = transformation.NQ_transform(R, metadata, inverse=True) + dataset = transformation.NQ_transform(dataset, inverse=True) elif metadata["transform"] == "sqrt": - R, metadata = transformation.sqrt_transform(R, metadata, inverse=True) + dataset = transformation.sqrt_transform(dataset, inverse=True) else: - raise ValueError("Unknown transformation %s" % metadata["transform"]) + raise ValueError(f'Unknown transformation {metadata["transform"]}') - if metadata["unit"] == "mm/h": + if metadata["units"] == "mm/h": pass - elif metadata["unit"] == "mm": + elif metadata["units"] == "mm": threshold = metadata["threshold"] # convert the threshold, too zerovalue = metadata["zerovalue"] # convert the zerovalue, too - R = R / float(metadata["accutime"]) * 60.0 + precip_data = precip_data / float(metadata["accutime"]) * 60.0 threshold = threshold / float(metadata["accutime"]) * 60.0 zerovalue = zerovalue / float(metadata["accutime"]) * 60.0 metadata["threshold"] = threshold metadata["zerovalue"] = zerovalue - elif metadata["unit"] == "dBZ": + elif metadata["units"] == "dBZ": threshold = metadata["threshold"] # convert the threshold, too zerovalue = metadata["zerovalue"] # convert the zerovalue, too @@ -93,7 +131,7 @@ def to_rainrate(R, metadata, zr_a=None, zr_b=None): zr_a = metadata.get("zr_a", 200.0) # default to Marshall–Palmer if zr_b is None: zr_b = metadata.get("zr_b", 1.6) # default to Marshall–Palmer - R = (R / zr_a) ** (1.0 / zr_b) + precip_data = (precip_data / zr_a) ** (1.0 / zr_b) threshold = (threshold / zr_a) ** (1.0 / zr_b) zerovalue = (zerovalue / zr_a) ** (1.0 / zr_b) @@ -104,13 +142,12 @@ def to_rainrate(R, metadata, zr_a=None, zr_b=None): else: raise ValueError( - "Cannot convert unit %s and transform %s to mm/h" - % (metadata["unit"], metadata["transform"]) + f'Cannot convert unit {metadata["units"]} and transform {metadata["transform"]} to mm/h' ) - metadata["unit"] = "mm/h" - - return R, metadata + dataset[precip_var].data[:] = precip_data + dataset = _change_unit(dataset, precip_var, "mm/h") + return dataset def to_raindepth(R, metadata, zr_a=None, zr_b=None): diff --git a/pysteps/utils/transformation.py b/pysteps/utils/transformation.py index 87ac9adc7..1977583c6 100644 --- a/pysteps/utils/transformation.py +++ b/pysteps/utils/transformation.py @@ -14,9 +14,11 @@ sqrt_transform """ +import warnings + import numpy as np import scipy.stats as scipy_stats -import warnings +import xarray as xr from scipy.interpolate import interp1d warnings.filterwarnings( @@ -25,8 +27,8 @@ def boxcox_transform( - R, metadata=None, Lambda=None, threshold=None, zerovalue=None, inverse=False -): + dataset: xr.Dataset, Lambda=None, threshold=None, zerovalue=None, inverse=False +) -> xr.Dataset: """ The one-parameter Box-Cox transformation. @@ -39,12 +41,8 @@ def boxcox_transform( Parameters ---------- - R: array-like - Array of any shape to be transformed. - metadata: dict, optional - Metadata dictionary containing the transform, zerovalue and threshold - attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be transformed. Lambda: float, optional Parameter Lambda of the Box-Cox transformation. It is 0 by default, which produces the log transformation. @@ -52,7 +50,7 @@ def boxcox_transform( Choose Lambda < 1 for positively skewed data, Lambda > 1 for negatively skewed data. threshold: float, optional - The value that is used for thresholding with the same units as R. + The value that is used for thresholding with the same units as in the dataset. If None, the threshold contained in metadata is used. If no threshold is found in the metadata, a value of 0.1 is used as default. @@ -64,10 +62,8 @@ def boxcox_transform( Returns ------- - R: array-like - Array of any shape containing the (back-)transformed units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the (back-)transformed units. References ---------- @@ -76,20 +72,14 @@ def boxcox_transform( doi:10.1111/j.2517-6161.1964.tb00553.x """ - R = R.copy() - - if metadata is None: - if inverse: - metadata = {"transform": "BoxCox"} - else: - metadata = {"transform": None} - - else: - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values if not inverse: if metadata["transform"] == "BoxCox": - return R, metadata + return dataset if Lambda is None: Lambda = metadata.get("BoxCox_lambda", 0.0) @@ -97,21 +87,21 @@ def boxcox_transform( if threshold is None: threshold = metadata.get("threshold", 0.1) - zeros = R < threshold + zeros = precip_data < threshold # Apply Box-Cox transform if Lambda == 0.0: - R[~zeros] = np.log(R[~zeros]) + precip_data[~zeros] = np.log(precip_data[~zeros]) threshold = np.log(threshold) else: - R[~zeros] = (R[~zeros] ** Lambda - 1) / Lambda + precip_data[~zeros] = (precip_data[~zeros] ** Lambda - 1) / Lambda threshold = (threshold**Lambda - 1) / Lambda # Set value for zeros if zerovalue is None: zerovalue = threshold - 1 # TODO: set to a more meaningful value - R[zeros] = zerovalue + precip_data[zeros] = zerovalue metadata["transform"] = "BoxCox" metadata["BoxCox_lambda"] = Lambda @@ -120,7 +110,7 @@ def boxcox_transform( elif inverse: if metadata["transform"] not in ["BoxCox", "log"]: - return R, metadata + return precip_data, metadata if Lambda is None: Lambda = metadata.pop("BoxCox_lambda", 0.0) @@ -131,35 +121,35 @@ def boxcox_transform( # Apply inverse Box-Cox transform if Lambda == 0.0: - R = np.exp(R) + precip_data = np.exp(precip_data) threshold = np.exp(threshold) else: - R = np.exp(np.log(Lambda * R + 1) / Lambda) + precip_data = np.exp(np.log(Lambda * precip_data + 1) / Lambda) threshold = np.exp(np.log(Lambda * threshold + 1) / Lambda) - R[R < threshold] = zerovalue + precip_data[precip_data < threshold] = zerovalue metadata["transform"] = None metadata["zerovalue"] = zerovalue metadata["threshold"] = threshold - return R, metadata + dataset[precip_var].data[:] = precip_data + + return dataset -def dB_transform(R, metadata=None, threshold=None, zerovalue=None, inverse=False): +def dB_transform( + dataset: xr.Dataset, threshold=None, zerovalue=None, inverse=False +) -> xr.Dataset: """Methods to transform precipitation intensities to/from dB units. Parameters ---------- - R: array-like - Array of any shape to be (back-)transformed. - metadata: dict, optional - Metadata dictionary containing the transform, zerovalue and threshold - attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be (back-)transformed. threshold: float, optional - Optional value that is used for thresholding with the same units as R. + Optional value that is used for thresholding with the same units as in the dataset. If None, the threshold contained in metadata is used. If no threshold is found in the metadata, a value of 0.1 is used as default. @@ -171,82 +161,70 @@ def dB_transform(R, metadata=None, threshold=None, zerovalue=None, inverse=False Returns ------- - R: array-like - Array of any shape containing the (back-)transformed units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the (back-)transformed units. """ - R = R.copy() - - if metadata is None: - if inverse: - metadata = {"transform": "dB"} - else: - metadata = {"transform": None} - - else: - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values # to dB units if not inverse: if metadata["transform"] == "dB": - return R, metadata + return dataset if threshold is None: threshold = metadata.get("threshold", 0.1) - zeros = R < threshold + zeros = precip_data < threshold # Convert to dB - R[~zeros] = 10.0 * np.log10(R[~zeros]) + precip_data[~zeros] = 10.0 * np.log10(precip_data[~zeros]) threshold = 10.0 * np.log10(threshold) # Set value for zeros if zerovalue is None: zerovalue = threshold - 5 # TODO: set to a more meaningful value - R[zeros] = zerovalue + precip_data[zeros] = zerovalue metadata["transform"] = "dB" metadata["zerovalue"] = zerovalue metadata["threshold"] = threshold - return R, metadata - # from dB units elif inverse: if metadata["transform"] != "dB": - return R, metadata + return dataset if threshold is None: threshold = metadata.get("threshold", -10.0) if zerovalue is None: zerovalue = 0.0 - R = 10.0 ** (R / 10.0) + precip_data = 10.0 ** (precip_data / 10.0) threshold = 10.0 ** (threshold / 10.0) - R[R < threshold] = zerovalue + precip_data[precip_data < threshold] = zerovalue metadata["transform"] = None metadata["threshold"] = threshold metadata["zerovalue"] = zerovalue - return R, metadata + dataset[precip_var].data[:] = precip_data + + return dataset -def NQ_transform(R, metadata=None, inverse=False, **kwargs): +def NQ_transform(dataset: xr.Dataset, inverse: bool = False, **kwargs) -> xr.Dataset: """ The normal quantile transformation as in Bogner et al (2012). Zero rain vales are set to zero in norm space. Parameters ---------- - R: array-like - Array of any shape to be transformed. - metadata: dict, optional - Metadata dictionary containing the transform, zerovalue and threshold - attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be transformed. inverse: bool, optional If set to True, it performs the inverse transform. False by default. @@ -260,10 +238,8 @@ def NQ_transform(R, metadata=None, inverse=False, **kwargs): Returns ------- - R: array-like - Array of any shape containing the (back-)transformed units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the (back-)transformed units. References ---------- @@ -276,105 +252,95 @@ def NQ_transform(R, metadata=None, inverse=False, **kwargs): # defaults a = kwargs.get("a", 0.0) - R = R.copy() - shape0 = R.shape - R = R.ravel().astype(float) - idxNan = np.isnan(R) - R_ = R[~idxNan] - - if metadata is None: - if inverse: - metadata = {"transform": "NQT"} - else: - metadata = {"transform": None} - metadata["zerovalue"] = np.min(R_) + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values - else: - metadata = metadata.copy() + shape0 = precip_data.shape + precip_data = precip_data.ravel().astype(float) + idxNan = np.isnan(precip_data) + precip_data_ = precip_data[~idxNan] if not inverse: # Plotting positions # https://en.wikipedia.org/wiki/Q%E2%80%93Q_plot#Plotting_position - n = R_.size - Rpp = ((np.arange(n) + 1 - a) / (n + 1 - 2 * a)).reshape(R_.shape) + n = precip_data_.size + Rpp = ((np.arange(n) + 1 - a) / (n + 1 - 2 * a)).reshape(precip_data_.shape) # NQ transform Rqn = scipy_stats.norm.ppf(Rpp) - R__ = np.interp(R_, R_[np.argsort(R_)], Rqn) + precip_data__ = np.interp( + precip_data_, precip_data_[np.argsort(precip_data_)], Rqn + ) # set zero rain to 0 in norm space - R__[R[~idxNan] == metadata["zerovalue"]] = 0 + precip_data__[precip_data[~idxNan] == metadata["zerovalue"]] = 0 # build inverse transform metadata["inqt"] = interp1d( - Rqn, R_[np.argsort(R_)], bounds_error=False, fill_value=(R_.min(), R_.max()) + Rqn, + precip_data_[np.argsort(precip_data_)], + bounds_error=False, + fill_value=(precip_data_.min(), precip_data_.max()), ) metadata["transform"] = "NQT" metadata["zerovalue"] = 0 - metadata["threshold"] = R__[R__ > 0].min() + metadata["threshold"] = precip_data__[precip_data__ > 0].min() else: f = metadata.pop("inqt") - R__ = f(R_) + precip_data__ = f(precip_data_) metadata["transform"] = None - metadata["zerovalue"] = R__.min() - metadata["threshold"] = R__[R__ > R__.min()].min() + metadata["zerovalue"] = precip_data__.min() + metadata["threshold"] = precip_data__[precip_data__ > precip_data__.min()].min() - R[~idxNan] = R__ + precip_data[~idxNan] = precip_data__ - return R.reshape(shape0), metadata + dataset[precip_var].data[:] = precip_data.reshape(shape0) + return dataset -def sqrt_transform(R, metadata=None, inverse=False, **kwargs): + +def sqrt_transform(dataset: xr.Dataset, inverse: bool = False, **kwargs) -> xr.Dataset: """ Square-root transform. Parameters ---------- - R: array-like - Array of any shape to be transformed. - metadata: dict, optional - Metadata dictionary containing the transform, zerovalue and threshold - attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be transformed. inverse: bool, optional If set to True, it performs the inverse transform. False by default. Returns ------- - R: array-like - Array of any shape containing the (back-)transformed units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the (back-)transformed units. """ - R = R.copy() - - if metadata is None: - if inverse: - metadata = {"transform": "sqrt"} - else: - metadata = {"transform": None} - metadata["zerovalue"] = np.nan - metadata["threshold"] = np.nan - else: - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values if not inverse: # sqrt transform - R = np.sqrt(R) + precip_data = np.sqrt(precip_data) metadata["transform"] = "sqrt" metadata["zerovalue"] = np.sqrt(metadata["zerovalue"]) metadata["threshold"] = np.sqrt(metadata["threshold"]) else: # inverse sqrt transform - R = R**2 + precip_data = precip_data**2 metadata["transform"] = None metadata["zerovalue"] = metadata["zerovalue"] ** 2 metadata["threshold"] = metadata["threshold"] ** 2 - return R, metadata + dataset[precip_var].data[:] = precip_data + + return dataset diff --git a/requirements.txt b/requirements.txt index 1804df1d9..b5075ad35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ matplotlib jsmin jsonschema netCDF4 +xarray diff --git a/requirements_dev.txt b/requirements_dev.txt index 84cf372b1..2899e560d 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -9,6 +9,7 @@ matplotlib jsmin jsonschema netCDF4 +xarray # Optional dependencies dask From 23825dd54d8b5c3322b2ba8759459935071bb061 Mon Sep 17 00:00:00 2001 From: Mats Veldhuizen Date: Wed, 17 Jul 2024 18:10:46 +0200 Subject: [PATCH 2/8] convert to_rainrate function also --- pysteps/tests/test_utils_conversion.py | 381 +++++++++++++++++++------ pysteps/utils/conversion.py | 4 + 2 files changed, 299 insertions(+), 86 deletions(-) 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 From 9952bd3317a277546cbeeafb921327589c5c0bf8 Mon Sep 17 00:00:00 2001 From: Mats Veldhuizen Date: Thu, 18 Jul 2024 10:43:27 +0200 Subject: [PATCH 3/8] Rewrite more code to xarray --- pysteps/tests/helpers.py | 28 +- pysteps/tests/test_utils_conversion.py | 603 +++++++++++++++++-------- pysteps/utils/conversion.py | 121 ++--- 3 files changed, 480 insertions(+), 272 deletions(-) diff --git a/pysteps/tests/helpers.py b/pysteps/tests/helpers.py index 2ae10026b..68f3f7527 100644 --- a/pysteps/tests/helpers.py +++ b/pysteps/tests/helpers.py @@ -27,27 +27,25 @@ def assert_dataset_equivalent(dataset1: xr.Dataset, dataset2: xr.Dataset) -> None: xr.testing.assert_allclose(dataset1, dataset2) + precip_var = dataset1.attrs["precip_var"] + assert precip_var == dataset2.attrs["precip_var"] assert np.isclose( - dataset1["precip_intensity"].attrs["threshold"], - dataset2["precip_intensity"].attrs["threshold"], + dataset1[precip_var].attrs["threshold"], + dataset2[precip_var].attrs["threshold"], ) - assert ( - dataset1["precip_intensity"].attrs["units"] - == dataset2["precip_intensity"].attrs["units"] - ) - assert ( - dataset1["precip_intensity"].attrs["transform"] - == dataset2["precip_intensity"].attrs["transform"] - or dataset1["precip_intensity"].attrs["transform"] is None - and dataset2["precip_intensity"].attrs["transform"] is None + assert np.isclose( + dataset1[precip_var].attrs["zerovalue"], + dataset2[precip_var].attrs["zerovalue"], ) + assert dataset1[precip_var].attrs["units"] == dataset2[precip_var].attrs["units"] assert ( - dataset1["precip_intensity"].attrs["accutime"] - == dataset2["precip_intensity"].attrs["accutime"] + dataset1[precip_var].attrs["transform"] + == dataset2[precip_var].attrs["transform"] + or dataset1[precip_var].attrs["transform"] is None + and dataset2[precip_var].attrs["transform"] is None ) assert ( - dataset1["precip_intensity"].attrs["zerovalue"] - == dataset2["precip_intensity"].attrs["zerovalue"] + dataset1[precip_var].attrs["accutime"] == dataset2[precip_var].attrs["accutime"] ) diff --git a/pysteps/tests/test_utils_conversion.py b/pysteps/tests/test_utils_conversion.py index 80d052b77..de48c928d 100644 --- a/pysteps/tests/test_utils_conversion.py +++ b/pysteps/tests/test_utils_conversion.py @@ -104,7 +104,7 @@ "transform": None, "accutime": 5, "threshold": 1.25892541, - "zerovalue": 0, + "zerovalue": 0.0, }, ) }, @@ -138,7 +138,7 @@ "transform": None, "accutime": 5, "threshold": 15.10710494, - "zerovalue": 0, + "zerovalue": 0.0, }, ) }, @@ -172,7 +172,7 @@ "transform": None, "accutime": 5, "threshold": 0.04210719, - "zerovalue": 0, + "zerovalue": 0.0, }, ) }, @@ -206,7 +206,7 @@ "transform": None, "accutime": 5, "threshold": 2.71828183, - "zerovalue": 0, + "zerovalue": 0.0, }, ) }, @@ -240,7 +240,7 @@ "transform": None, "accutime": 5, "threshold": 32.61938194, - "zerovalue": 0, + "zerovalue": 0.0, }, ) }, @@ -273,8 +273,8 @@ "units": "mm/h", "transform": None, "accutime": 5, - "threshold": 1, - "zerovalue": 1, + "threshold": 1.0, + "zerovalue": 1.0, }, ) }, @@ -307,8 +307,8 @@ "units": "mm/h", "transform": None, "accutime": 5, - "threshold": 12, - "zerovalue": 12, + "threshold": 12.0, + "zerovalue": 12.0, }, ) }, @@ -326,220 +326,429 @@ def test_to_rainrate(dataset, expected): # to_raindepth -test_data = [ - ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.08333333]), - ), +test_data_to_raindepth = [ ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm", - "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_accum": ( + ["x"], + np.array([0.08333333]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 0.08333333, + "zerovalue": 0.08333333, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.10491045]), + 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_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm", - "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_accum": ( + ["x"], + np.array([0.10491045]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 0.10491045, + "zerovalue": 0.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "dBZ", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.00350893]), + 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_accum": ( + ["x"], + np.array([1.25892541]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 1.25892541, + "zerovalue": 0.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "log", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.22652349]), + 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_accum": ( + ["x"], + np.array([0.00350893]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 0.00350893, + "zerovalue": 0.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1.0]), - { - "accutime": 5, - "transform": "log", - "unit": "mm", - "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_accum": ( + ["x"], + np.array([0.22652349]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 0.22652349, + "zerovalue": 0.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([0.08333333]), + 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_accum": ( + ["x"], + np.array([0.08333333]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 0.08333333, + "zerovalue": 0.08333333, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ( - np.array([1.0]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([1.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_accum": ( + ["x"], + np.array([1.0]), + { + "units": "mm", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "zerovalue": 1.0, + }, + ) + }, + attrs={"precip_var": "precip_accum"}, + ), ), ] -@pytest.mark.parametrize("R, metadata, expected", test_data) -def test_to_raindepth(R, metadata, expected): +@pytest.mark.parametrize("dataset, expected", test_data_to_raindepth) +def test_to_raindepth(dataset, expected): """Test the to_raindepth.""" - assert_array_almost_equal(conversion.to_raindepth(R, metadata)[0], expected) + actual = conversion.to_raindepth(dataset) + assert_dataset_equivalent(actual, expected) # to_reflectivity -test_data = [ - ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([23.01029996]), - ), - ( - np.array([1]), - { - "accutime": 5, - "transform": None, - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([40.27719989]), - ), - ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([24.61029996]), - ), - ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([41.87719989]), - ), - ( - np.array([1]), - { - "accutime": 5, - "transform": "dB", - "unit": "dBZ", - "threshold": 0, - "zerovalue": 0, - }, - np.array([1]), - ), +test_data_to_reflectivity = [ ( - np.array([1]), - { - "accutime": 5, - "transform": "log", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([29.95901167]), - ), - ( - np.array([1.0]), - { - "accutime": 5, - "transform": "log", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([47.2259116]), - ), - ( - np.array([1]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm/h", - "threshold": 0, - "zerovalue": 0, - }, - np.array([23.01029996]), - ), - ( - np.array([1.0]), - { - "accutime": 5, - "transform": "sqrt", - "unit": "mm", - "threshold": 0, - "zerovalue": 0, - }, - np.array([40.27719989]), + 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={ + "reflectivity": ( + ["x"], + np.array([23.01029996]), + { + "units": "dBZ", + "transform": None, + "accutime": 5, + "threshold": 23.01029996, + "zerovalue": 23.01029996, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": None, + # "unit": "mm/h", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([23.01029996]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": None, + # "unit": "mm", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([40.27719989]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": "dB", + # "unit": "mm/h", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([24.61029996]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": "dB", + # "unit": "mm", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([41.87719989]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": "dB", + # "unit": "dBZ", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([1]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": "log", + # "unit": "mm/h", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([29.95901167]), + # ), + # ( + # np.array([1.0]), + # { + # "accutime": 5, + # "transform": "log", + # "unit": "mm", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([47.2259116]), + # ), + # ( + # np.array([1]), + # { + # "accutime": 5, + # "transform": "sqrt", + # "unit": "mm/h", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([23.01029996]), + # ), + # ( + # np.array([1.0]), + # { + # "accutime": 5, + # "transform": "sqrt", + # "unit": "mm", + # "threshold": 0, + # "zerovalue": 0, + # }, + # np.array([40.27719989]), + # ), ] -@pytest.mark.parametrize("R, metadata, expected", test_data) -def test_to_reflectivity(R, metadata, expected): +@pytest.mark.parametrize("dataset, expected", test_data_to_reflectivity) +def test_to_reflectivity(dataset, expected): """Test the to_reflectivity.""" - assert_array_almost_equal(conversion.to_reflectivity(R, metadata)[0], expected) + actual = conversion.to_reflectivity(dataset) + assert_dataset_equivalent(actual, expected) diff --git a/pysteps/utils/conversion.py b/pysteps/utils/conversion.py index 87fe22deb..21e32360b 100644 --- a/pysteps/utils/conversion.py +++ b/pysteps/utils/conversion.py @@ -90,7 +90,6 @@ def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): dataset = dataset.copy(deep=True) precip_var = dataset.attrs["precip_var"] metadata = dataset[precip_var].attrs - precip_data = dataset[precip_var].values if metadata["transform"] is not None: if metadata["transform"] == "dB": @@ -154,18 +153,14 @@ def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): return dataset -def to_raindepth(R, metadata, zr_a=None, zr_b=None): +def to_raindepth(dataset: xr.Dataset, zr_a=None, zr_b=None): """ Convert to rain depth [mm]. Parameters ---------- - R: array-like - Array of any shape to be (back-)transformed. - metadata: dict - Metadata dictionary containing the accutime, transform, unit, threshold - and zerovalue attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be (back-)transformed. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -177,46 +172,49 @@ def to_raindepth(R, metadata, zr_a=None, zr_b=None): Returns ------- - R: array-like - Array of any shape containing the converted units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the converted units. """ - R = R.copy() - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs if metadata["transform"] is not None: if metadata["transform"] == "dB": - R, metadata = transformation.dB_transform(R, metadata, inverse=True) + dataset = transformation.dB_transform(dataset, inverse=True) elif metadata["transform"] in ["BoxCox", "log"]: - R, metadata = transformation.boxcox_transform(R, metadata, inverse=True) + dataset = transformation.boxcox_transform(dataset, inverse=True) elif metadata["transform"] == "NQT": - R, metadata = transformation.NQ_transform(R, metadata, inverse=True) + dataset = transformation.NQ_transform(dataset, inverse=True) elif metadata["transform"] == "sqrt": - R, metadata = transformation.sqrt_transform(R, metadata, inverse=True) + dataset = transformation.sqrt_transform(dataset, inverse=True) else: - raise ValueError("Unknown transformation %s" % metadata["transform"]) + 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["unit"] == "mm" and metadata["transform"] is None: + if metadata["units"] == "mm" and metadata["transform"] is None: pass - elif metadata["unit"] == "mm/h": + elif metadata["units"] == "mm/h": threshold = metadata["threshold"] # convert the threshold, too zerovalue = metadata["zerovalue"] # convert the zerovalue, too - R = R / 60.0 * metadata["accutime"] + precip_data = precip_data / 60.0 * metadata["accutime"] threshold = threshold / 60.0 * metadata["accutime"] zerovalue = zerovalue / 60.0 * metadata["accutime"] metadata["threshold"] = threshold metadata["zerovalue"] = zerovalue - elif metadata["unit"] == "dBZ": + elif metadata["units"] == "dBZ": threshold = metadata["threshold"] # convert the threshold, too zerovalue = metadata["zerovalue"] # convert the zerovalue, too @@ -225,7 +223,7 @@ def to_raindepth(R, metadata, zr_a=None, zr_b=None): zr_a = metadata.get("zr_a", 200.0) # Default to Marshall–Palmer if zr_b is None: zr_b = metadata.get("zr_b", 1.6) # Default to Marshall–Palmer - R = (R / zr_a) ** (1.0 / zr_b) / 60.0 * metadata["accutime"] + precip_data = (precip_data / zr_a) ** (1.0 / zr_b) / 60.0 * metadata["accutime"] threshold = (threshold / zr_a) ** (1.0 / zr_b) / 60.0 * metadata["accutime"] zerovalue = (zerovalue / zr_a) ** (1.0 / zr_b) / 60.0 * metadata["accutime"] @@ -236,27 +234,22 @@ def to_raindepth(R, metadata, zr_a=None, zr_b=None): else: raise ValueError( - "Cannot convert unit %s and transform %s to mm" - % (metadata["unit"], metadata["transform"]) + f'Cannot convert unit {metadata["units"]} and transform {metadata["transform"]} to mm' ) - metadata["unit"] = "mm" - - return R, metadata + dataset[precip_var].data[:] = precip_data + dataset = _change_unit(dataset, precip_var, "mm") + return dataset -def to_reflectivity(R, metadata, zr_a=None, zr_b=None): +def to_reflectivity(dataset: xr.Dataset, zr_a=None, zr_b=None): """ Convert to reflectivity [dBZ]. Parameters ---------- - R: array-like - Array of any shape to be (back-)transformed. - metadata: dict - Metadata dictionary containing the accutime, transform, unit, threshold - and zerovalue attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. + dataset: Dataset + Dataset to be (back-)transformed. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -268,73 +261,81 @@ def to_reflectivity(R, metadata, zr_a=None, zr_b=None): Returns ------- - R: array-like - Array of any shape containing the converted units. - metadata: dict - The metadata with updated attributes. + dataset: Dataset + Dataset containing the converted units. """ - R = R.copy() - metadata = metadata.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs if metadata["transform"] is not None: if metadata["transform"] == "dB": - R, metadata = transformation.dB_transform(R, metadata, inverse=True) + dataset = transformation.dB_transform(dataset, inverse=True) elif metadata["transform"] in ["BoxCox", "log"]: - R, metadata = transformation.boxcox_transform(R, metadata, inverse=True) + dataset = transformation.boxcox_transform(dataset, inverse=True) elif metadata["transform"] == "NQT": - R, metadata = transformation.NQ_transform(R, metadata, inverse=True) + dataset = transformation.NQ_transform(dataset, inverse=True) elif metadata["transform"] == "sqrt": - R, metadata = transformation.sqrt_transform(R, metadata, inverse=True) + dataset = transformation.sqrt_transform(dataset, inverse=True) else: - raise ValueError("Unknown transformation %s" % metadata["transform"]) + raise ValueError(f'Unknown transformation {metadata["transform"]}') - if metadata["unit"] == "mm/h": + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values + + if metadata["units"] == "mm/h": # Z to R if zr_a is None: zr_a = metadata.get("zr_a", 200.0) # Default to Marshall–Palmer if zr_b is None: zr_b = metadata.get("zr_b", 1.6) # Default to Marshall–Palmer - R = zr_a * R**zr_b + precip_data = zr_a * precip_data**zr_b metadata["threshold"] = zr_a * metadata["threshold"] ** zr_b metadata["zerovalue"] = zr_a * metadata["zerovalue"] ** zr_b metadata["zr_a"] = zr_a metadata["zr_b"] = zr_b # Z to dBZ - R, metadata = transformation.dB_transform(R, metadata) + dataset = transformation.dB_transform(dataset) - elif metadata["unit"] == "mm": + elif metadata["units"] == "mm": # depth to rate - R, metadata = to_rainrate(R, metadata) + dataset = to_rainrate(dataset) # Z to R if zr_a is None: zr_a = metadata.get("zr_a", 200.0) # Default to Marshall-Palmer if zr_b is None: zr_b = metadata.get("zr_b", 1.6) # Default to Marshall-Palmer - R = zr_a * R**zr_b + precip_data = zr_a * precip_data**zr_b metadata["threshold"] = zr_a * metadata["threshold"] ** zr_b metadata["zerovalue"] = zr_a * metadata["zerovalue"] ** zr_b metadata["zr_a"] = zr_a metadata["zr_b"] = zr_b # Z to dBZ - R, metadata = transformation.dB_transform(R, metadata) + dataset = transformation.dB_transform(dataset) - elif metadata["unit"] == "dBZ": + elif metadata["units"] == "dBZ": # Z to dBZ - R, metadata = transformation.dB_transform(R, metadata) + dataset = transformation.dB_transform(dataset) else: raise ValueError( - "Cannot convert unit %s and transform %s to mm/h" - % (metadata["unit"], metadata["transform"]) + f'Cannot convert unit {metadata["units"]} and transform {metadata["transform"]} to dBZ' ) - metadata["unit"] = "dBZ" - return R, metadata + + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values + + dataset[precip_var].data[:] = precip_data + dataset = _change_unit(dataset, precip_var, "dBZ") + return dataset From 904bce88255ea2c03875d0b8116d044c98f20fb1 Mon Sep 17 00:00:00 2001 From: Mats Veldhuizen Date: Thu, 18 Jul 2024 12:11:20 +0200 Subject: [PATCH 4/8] conversion.py works on the xarray datamodel --- pysteps/tests/test_utils_conversion.py | 373 ++++++++++++++++++------- pysteps/utils/conversion.py | 17 +- 2 files changed, 282 insertions(+), 108 deletions(-) diff --git a/pysteps/tests/test_utils_conversion.py b/pysteps/tests/test_utils_conversion.py index de48c928d..bdf1fa42f 100644 --- a/pysteps/tests/test_utils_conversion.py +++ b/pysteps/tests/test_utils_conversion.py @@ -635,115 +635,288 @@ def test_to_raindepth(dataset, expected): np.array([23.01029996]), { "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 23.01029996, + "zerovalue": 18.01029996, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([40.27719989]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 40.27719989, + "zerovalue": 35.27719989, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([24.61029996]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 24.61029996, + "zerovalue": 19.61029996, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([41.87719989]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 41.87719989, + "zerovalue": 36.87719989, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([1.0]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 1.0, + "zerovalue": -4.0, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([29.95901167]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 29.95901167, + "zerovalue": 24.95901167, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([47.2259116]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 47.2259116, + "zerovalue": 42.2259116, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([23.01029996]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, "threshold": 23.01029996, - "zerovalue": 23.01029996, + "zerovalue": 18.01029996, + }, + ) + }, + attrs={"precip_var": "reflectivity"}, + ), + ), + ( + 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={ + "reflectivity": ( + ["x"], + np.array([40.27719989]), + { + "units": "dBZ", + "transform": "dB", + "accutime": 5, + "threshold": 40.27719989, + "zerovalue": 35.27719989, }, ) }, attrs={"precip_var": "reflectivity"}, ), ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": None, - # "unit": "mm/h", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([23.01029996]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": None, - # "unit": "mm", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([40.27719989]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": "dB", - # "unit": "mm/h", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([24.61029996]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": "dB", - # "unit": "mm", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([41.87719989]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": "dB", - # "unit": "dBZ", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([1]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": "log", - # "unit": "mm/h", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([29.95901167]), - # ), - # ( - # np.array([1.0]), - # { - # "accutime": 5, - # "transform": "log", - # "unit": "mm", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([47.2259116]), - # ), - # ( - # np.array([1]), - # { - # "accutime": 5, - # "transform": "sqrt", - # "unit": "mm/h", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([23.01029996]), - # ), - # ( - # np.array([1.0]), - # { - # "accutime": 5, - # "transform": "sqrt", - # "unit": "mm", - # "threshold": 0, - # "zerovalue": 0, - # }, - # np.array([40.27719989]), - # ), ] diff --git a/pysteps/utils/conversion.py b/pysteps/utils/conversion.py index 21e32360b..68228e981 100644 --- a/pysteps/utils/conversion.py +++ b/pysteps/utils/conversion.py @@ -302,13 +302,14 @@ def to_reflectivity(dataset: xr.Dataset, zr_a=None, zr_b=None): metadata["zr_a"] = zr_a metadata["zr_b"] = zr_b - # Z to dBZ - dataset = transformation.dB_transform(dataset) - elif metadata["units"] == "mm": # depth to rate dataset = to_rainrate(dataset) + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + precip_data = dataset[precip_var].values + # Z to R if zr_a is None: zr_a = metadata.get("zr_a", 200.0) # Default to Marshall-Palmer @@ -320,18 +321,18 @@ def to_reflectivity(dataset: xr.Dataset, zr_a=None, zr_b=None): metadata["zr_a"] = zr_a metadata["zr_b"] = zr_b - # Z to dBZ - dataset = transformation.dB_transform(dataset) - elif metadata["units"] == "dBZ": - # Z to dBZ - dataset = transformation.dB_transform(dataset) + pass else: raise ValueError( f'Cannot convert unit {metadata["units"]} and transform {metadata["transform"]} to dBZ' ) + dataset[precip_var].data[:] = precip_data + # Z to dBZ + dataset = transformation.dB_transform(dataset) + precip_var = dataset.attrs["precip_var"] metadata = dataset[precip_var].attrs precip_data = dataset[precip_var].values From b3aa00912ea172443851cbd526fcf633e8d9e880 Mon Sep 17 00:00:00 2001 From: Mats Veldhuizen Date: Thu, 18 Jul 2024 14:02:58 +0200 Subject: [PATCH 5/8] fix converters.py --- pysteps/converters.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/pysteps/converters.py b/pysteps/converters.py index f27e1572e..6c576c658 100644 --- a/pysteps/converters.py +++ b/pysteps/converters.py @@ -187,11 +187,9 @@ def convert_to_xarray_dataset( } if grid_mapping_var_name is not None: coords[grid_mapping_name] = ( - ( - [], - None, - {"grid_mapping_name": grid_mapping_name, **grid_mapping_params}, - ), + [], + None, + {"grid_mapping_name": grid_mapping_name, **grid_mapping_params}, ) attrs = { "Conventions": "CF-1.7", @@ -199,4 +197,5 @@ def convert_to_xarray_dataset( "projection": metadata["projection"], "precip_var": var_name, } - return xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs) + dataset = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs) + return dataset.sortby(["y", "x"]) From 016d28f3219443bf88c5b54015623af543971319 Mon Sep 17 00:00:00 2001 From: mats-knmi <145579783+mats-knmi@users.noreply.github.com> Date: Mon, 12 Aug 2024 08:57:19 +0200 Subject: [PATCH 6/8] make dimension.py xarray compatible (#397) * make dimension.py xarray compatible * convert final method in the dimension module * nanmin in stead of zerovalue in square domain method * make test steps skill run * undo accidental change * remove commented out code * The dataset can contain more than one dataarray * Address pull request comments * Add links to dataset documentation everywhere --- pysteps/converters.py | 18 +- pysteps/io/importers.py | 87 ++++ pysteps/tests/helpers.py | 40 +- pysteps/tests/test_nowcasts_steps.py | 20 +- pysteps/tests/test_utils_dimension.py | 230 ++++++---- pysteps/utils/conversion.py | 21 +- pysteps/utils/dimension.py | 618 ++++++++++---------------- pysteps/utils/transformation.py | 28 +- 8 files changed, 522 insertions(+), 540 deletions(-) diff --git a/pysteps/converters.py b/pysteps/converters.py index 6c576c658..2825af612 100644 --- a/pysteps/converters.py +++ b/pysteps/converters.py @@ -12,6 +12,7 @@ """ import numpy as np +import numpy.typing as npt import pyproj import xarray as xr @@ -67,6 +68,15 @@ def _convert_proj4_to_grid_mapping(proj4str): return grid_mapping_var_name, grid_mapping_name, params +def compute_lat_lon( + x_r: npt.ArrayLike, y_r: npt.ArrayLike, projection: str +) -> tuple[npt.ArrayLike, npt.ArrayLike]: + x_2d, y_2d = np.meshgrid(x_r, y_r) + pr = pyproj.Proj(projection) + lon, lat = pr(x_2d.flatten(), y_2d.flatten(), inverse=True) + return lat.reshape(x_2d.shape), lon.reshape(x_2d.shape) + + def convert_to_xarray_dataset( precip: np.ndarray, quality: np.ndarray | None, @@ -105,9 +115,7 @@ def convert_to_xarray_dataset( if metadata["yorigin"] == "upper": y_r = np.flip(y_r) - x_2d, y_2d = np.meshgrid(x_r, y_r) - pr = pyproj.Proj(metadata["projection"]) - lon, lat = pr(x_2d.flatten(), y_2d.flatten(), inverse=True) + lat, lon = compute_lat_lon(x_r, y_r, metadata["projection"]) ( grid_mapping_var_name, @@ -166,7 +174,7 @@ def convert_to_xarray_dataset( ), "lon": ( ["y", "x"], - lon.reshape(precip.shape), + lon, { "long_name": "longitude coordinate", "standard_name": "longitude", @@ -176,7 +184,7 @@ def convert_to_xarray_dataset( ), "lat": ( ["y", "x"], - lat.reshape(precip.shape), + lat, { "long_name": "latitude coordinate", "standard_name": "latitude", diff --git a/pysteps/io/importers.py b/pysteps/io/importers.py index 5c2928203..f61d4b25b 100644 --- a/pysteps/io/importers.py +++ b/pysteps/io/importers.py @@ -65,6 +65,93 @@ | zr_b | the Z-R exponent b in Z = a*R**b | +------------------+----------------------------------------------------------+ +The data and metadata is then postprocessed into an xarray dataset. This dataset will +always contain an x and y dimension, but can be extended with a time dimension and/or +an ensemble member dimension over the course of the process. + +The dataset can contain the following coordinate variables: + + +.. tabularcolumns:: |p{2cm}|L| + ++---------------+-------------------------------------------------------------------------------------------+ +| Coordinate | Description | ++===============+===========================================================================================+ +| y | y-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | ++---------------+-------------------------------------------------------------------------------------------+ +| x | x-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | ++---------------+-------------------------------------------------------------------------------------------+ +| lat | latitude coordinate in degrees | ++---------------+-------------------------------------------------------------------------------------------+ +| lon | longitude coordinate in degrees | ++---------------+-------------------------------------------------------------------------------------------+ +| time | forecast time in seconds since forecast start time | ++---------------+-------------------------------------------------------------------------------------------+ +| member | ensemble member number (integer) | ++---------------+-------------------------------------------------------------------------------------------+ + + +The dataset can contain the following data variables: + +.. tabularcolumns:: |p{2cm}|L| + ++-------------------+-----------------------------------------------------------------------------------------------------------+ +| Variable | Description | ++===================+===========================================================================================================+ +| precip_intensity, | precipitation data, based on the unit the data has it is stored in one of these 3 possible variables | +| precip_accum | precip_intensity if unit is ``mm/h``, precip_accum if unit is ``mm`` and reflectivity if unit is ``dBZ``, | +| or reflectivity | the attributes of this variable contain metadata relevant to this attribute (see below) | ++-------------------+-----------------------------------------------------------------------------------------------------------+ +| quality | value between 0 and 1 denoting the quality of the precipitation data, currently not used for anything | ++-------------------+-----------------------------------------------------------------------------------------------------------+ + +Some of the metadata in the metadata dictionary is not explicitely stored in the dataset, +but is still implicitly present. For example ``x1`` can easily be found by taking the first +value from the x coordinate variable. Metadata that is not implicitly present is explicitly +stored either in the datasets global attributes or as attributes of the precipitation variable. +Data that relates to the entire dataset is stored in the global attributes. The following data +is stored in the global attributes: + +.. tabularcolumns:: |p{2cm}|L| + ++------------------+----------------------------------------------------------+ +| Key | Value | ++==================+==========================================================+ +| projection | PROJ.4-compatible projection definition | ++------------------+----------------------------------------------------------+ +| institution | name of the institution who provides the data | ++------------------+----------------------------------------------------------+ +| precip_var | the name of the precipitation variable in this dataset | ++------------------+----------------------------------------------------------+ + +The following data is stored as attributes of the precipitation variable: + +.. tabularcolumns:: |p{2cm}|L| + ++------------------+----------------------------------------------------------+ +| Key | Value | ++==================+==========================================================+ +| units | the physical unit of the data: 'mm/h', 'mm' or 'dBZ' | ++------------------+----------------------------------------------------------+ +| transform | the transformation of the data: None, 'dB', 'Box-Cox' or | +| | others | ++------------------+----------------------------------------------------------+ +| accutime | the accumulation time in minutes of the data, float | ++------------------+----------------------------------------------------------+ +| threshold | the rain/no rain threshold with the same unit, | +| | transformation and accutime of the data. | ++------------------+----------------------------------------------------------+ +| zerovalue | the value assigned to the no rain pixels with the same | +| | unit, transformation and accutime of the data. | ++------------------+----------------------------------------------------------+ +| zr_a | the Z-R constant a in Z = a*R**b | ++------------------+----------------------------------------------------------+ +| zr_b | the Z-R exponent b in Z = a*R**b | ++------------------+----------------------------------------------------------+ + +Furthermore the dataset can contain some additional metadata to make the dataset +CF-compliant. + Available Importers ------------------- diff --git a/pysteps/tests/helpers.py b/pysteps/tests/helpers.py index 68f3f7527..24c58f8d1 100644 --- a/pysteps/tests/helpers.py +++ b/pysteps/tests/helpers.py @@ -14,6 +14,7 @@ import pysteps as stp from pysteps import io, rcparams from pysteps.utils import aggregate_fields_space +from pysteps.utils.dimension import clip_domain _reference_dates = dict() _reference_dates["bom"] = datetime(2018, 6, 16, 10, 0) @@ -53,7 +54,6 @@ def get_precipitation_fields( num_prev_files=0, num_next_files=0, return_raw=False, - metadata=False, upscale=None, source="mch", log_transform=True, @@ -100,9 +100,6 @@ def get_precipitation_fields( The pre-processing steps are: 1) Convert to mm/h, 2) Mask invalid values, 3) Log-transform the data [dBR]. - metadata: bool, optional - If True, also return file metadata. - upscale: float or None, optional Upscale fields in space during the pre-processing steps. If it is None, the precipitation field is not modified. @@ -127,8 +124,8 @@ def get_precipitation_fields( Returns ------- - reference_field : array - metadata : dict + dataset: xarray.Dataset + As described in the documentation of :py:mod:`pysteps.io.importers`. """ if source == "bom": @@ -186,41 +183,34 @@ def get_precipitation_fields( # Read the radar composites importer = io.get_method(importer_name, "importer") - ref_dataset = io.read_timeseries(fns, importer, **_importer_kwargs) + dataset = io.read_timeseries(fns, importer, **_importer_kwargs) if not return_raw: - if (num_prev_files == 0) and (num_next_files == 0): - # Remove time dimension - reference_field = np.squeeze(reference_field) + precip_var = dataset.attrs["precip_var"] # Convert to mm/h - ref_dataset = stp.utils.to_rainrate(ref_dataset) + dataset = stp.utils.to_rainrate(dataset) + precip_var = dataset.attrs["precip_var"] # Clip domain - ref_dataset = stp.utils.clip_domain(ref_dataset, clip) + dataset = clip_domain(dataset, clip) # Upscale data - reference_field, ref_metadata = aggregate_fields_space( - reference_field, ref_metadata, upscale - ) + dataset = aggregate_fields_space(dataset, upscale) # Mask invalid values - reference_field = np.ma.masked_invalid(reference_field) + valid_mask = np.isfinite(dataset[precip_var].values) if log_transform: # Log-transform the data [dBR] - reference_field, ref_metadata = stp.utils.dB_transform( - reference_field, ref_metadata, threshold=0.1, zerovalue=-15.0 - ) + dataset = stp.utils.dB_transform(dataset, threshold=0.1, zerovalue=-15.0) # Set missing values with the fill value - np.ma.set_fill_value(reference_field, ref_metadata["zerovalue"]) - reference_field.data[reference_field.mask] = ref_metadata["zerovalue"] - - if metadata: - return reference_field, ref_metadata + metadata = dataset[precip_var].attrs + zerovalue = metadata["zerovalue"] + dataset[precip_var].data[~valid_mask] = zerovalue - return reference_field + return dataset def smart_assert(actual_value, expected, tolerance=None): diff --git a/pysteps/tests/test_nowcasts_steps.py b/pysteps/tests/test_nowcasts_steps.py index 61af86ba5..adb6ea917 100644 --- a/pysteps/tests/test_nowcasts_steps.py +++ b/pysteps/tests/test_nowcasts_steps.py @@ -7,7 +7,6 @@ from pysteps import io, motion, nowcasts, verification from pysteps.tests.helpers import get_precipitation_fields - steps_arg_names = ( "n_ens_members", "n_cascade_levels", @@ -44,28 +43,29 @@ def test_steps_skill( ): """Tests STEPS nowcast skill.""" # inputs - precip_input, metadata = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=True, upscale=2000, ) - precip_input = precip_input.filled() - precip_obs = get_precipitation_fields( + dataset_obs = get_precipitation_fields( num_prev_files=0, num_next_files=3, return_raw=False, upscale=2000 - )[1:, :, :] - precip_obs = precip_obs.filled() + ).isel(time=slice(1, None, None)) + precip_var = dataset_input.attrs["precip_var"] + metadata = dataset_input[precip_var].attrs + precip_data = dataset_input[precip_var].values pytest.importorskip("cv2") oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_input) + retrieved_motion = oflow_method(precip_data) nowcast_method = nowcasts.get_method("steps") precip_forecast = nowcast_method( - precip_input, + precip_data, retrieved_motion, timesteps=timesteps, precip_thr=metadata["threshold"], @@ -86,7 +86,9 @@ def test_steps_skill( timesteps if isinstance(timesteps, int) else len(timesteps) ) - crps = verification.probscores.CRPS(precip_forecast[:, -1], precip_obs[-1]) + crps = verification.probscores.CRPS( + precip_forecast[:, -1], dataset_obs[precip_var].values[-1] + ) assert crps < max_crps, f"CRPS={crps:.2f}, required < {max_crps:.2f}" diff --git a/pysteps/tests/test_utils_dimension.py b/pysteps/tests/test_utils_dimension.py index ab753ed7d..2bbb63f58 100644 --- a/pysteps/tests/test_utils_dimension.py +++ b/pysteps/tests/test_utils_dimension.py @@ -4,63 +4,89 @@ import numpy as np import pytest +import xarray as xr from numpy.testing import assert_array_equal from pytest import raises +from pysteps.converters import convert_to_xarray_dataset from pysteps.utils import dimension +fillvalues_metadata = { + "x1": 0, + "x2": 4, + "y1": 0, + "y2": 4, + "xpixelsize": 1, + "ypixelsize": 1, + "zerovalue": 0, + "yorigin": "lower", + "unit": "mm/h", + "transform": None, + "accutime": 5, + "threshold": 1.0, + "projection": "+proj=stere +lat_0=90 +lon_0=0.0 +lat_ts=60.0 +a=6378.137 +b=6356.752 +x_0=0 +y_0=0", + "zr_a": 200, + "zr_b": 1.6, + "cartesian_unit": "km", + "institution": "KNMI", +} + test_data_not_trim = ( # "data, window_size, axis, method, expected" - (np.arange(6), 2, 0, "mean", np.array([0.5, 2.5, 4.5])), + ( + np.arange(12).reshape(2, 6), + 2, + "x", + "mean", + np.array([[0.5, 2.5, 4.5], [6.5, 8.5, 10.5]]), + ), ( np.arange(4 * 6).reshape(4, 6), (2, 3), - (0, 1), + ("y", "x"), "sum", np.array([[24, 42], [96, 114]]), ), ( np.arange(4 * 6).reshape(4, 6), (2, 2), - (0, 1), + ("y", "x"), "sum", np.array([[14, 22, 30], [62, 70, 78]]), ), ( np.arange(4 * 6).reshape(4, 6), 2, - (0, 1), + ("y", "x"), "sum", np.array([[14, 22, 30], [62, 70, 78]]), ), ( np.arange(4 * 6).reshape(4, 6), (2, 3), - (0, 1), + ("y", "x"), "mean", np.array([[4.0, 7.0], [16.0, 19.0]]), ), ( np.arange(4 * 6).reshape(4, 6), (2, 2), - (0, 1), + ("y", "x"), "mean", np.array([[3.5, 5.5, 7.5], [15.5, 17.5, 19.5]]), ), ( np.arange(4 * 6).reshape(4, 6), 2, - (0, 1), + ("y", "x"), "mean", np.array([[3.5, 5.5, 7.5], [15.5, 17.5, 19.5]]), ), ) -@pytest.mark.parametrize( - "data, window_size, axis, method, expected", test_data_not_trim -) -def test_aggregate_fields(data, window_size, axis, method, expected): +@pytest.mark.parametrize("data, window_size, dim, method, expected", test_data_not_trim) +def test_aggregate_fields(data, window_size, dim, method, expected): """ Test the aggregate_fields function. The windows size must divide exactly the data dimensions. @@ -68,23 +94,25 @@ def test_aggregate_fields(data, window_size, axis, method, expected): windows size does not divide the data dimensions. The length of each dimension should be larger than 2. """ + dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) - assert_array_equal( - dimension.aggregate_fields(data, window_size, axis=axis, method=method), - expected, - ) + actual = dimension.aggregate_fields(dataset, window_size, dim=dim, method=method) + assert_array_equal(actual["precip_intensity"].values, expected) # Test the trimming capabilities. - data = np.pad(data, (0, 1)) - assert_array_equal( - dimension.aggregate_fields( - data, window_size, axis=axis, method=method, trim=True - ), - expected, + if np.ndim(window_size) == 0: + data = np.pad(data, ((0, 0), (0, 1))) + else: + data = np.pad(data, (0, 1)) + dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) + + actual = dimension.aggregate_fields( + dataset, window_size, dim=dim, method=method, trim=True ) + assert_array_equal(actual["precip_intensity"].values, expected) with raises(ValueError): - dimension.aggregate_fields(data, window_size, axis=axis, method=method) + dimension.aggregate_fields(dataset, window_size, dim=dim, method=method) def test_aggregate_fields_errors(): @@ -93,80 +121,124 @@ def test_aggregate_fields_errors(): function. """ data = np.arange(4 * 6).reshape(4, 6) + dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) with raises(ValueError): - dimension.aggregate_fields(data, -1, axis=0) + dimension.aggregate_fields(dataset, -1, dim="y") with raises(ValueError): - dimension.aggregate_fields(data, 0, axis=0) + dimension.aggregate_fields(dataset, 0, dim="y") with raises(ValueError): - dimension.aggregate_fields(data, 1, method="invalid") + dimension.aggregate_fields(dataset, 1, method="invalid") with raises(TypeError): - dimension.aggregate_fields(data, (1, 1), axis=0) + dimension.aggregate_fields(dataset, (1, 1), dim="y") # aggregate_fields_time -timestamps = [dt.datetime.now() + dt.timedelta(minutes=t) for t in range(10)] -test_data = [ +now = dt.datetime.now() +timestamps = [now + dt.timedelta(minutes=t) for t in range(10)] +test_data_time = [ ( - np.ones((10, 1, 1)), + np.ones((2, 2)), {"unit": "mm/h", "timestamps": timestamps}, 2, False, - np.ones((5, 1, 1)), + np.ones((5, 2, 2)), ), ( - np.ones((10, 1, 1)), + np.ones((2, 2)), {"unit": "mm", "timestamps": timestamps}, 2, False, - 2 * np.ones((5, 1, 1)), + 2 * np.ones((5, 2, 2)), ), ] @pytest.mark.parametrize( - "R, metadata, time_window_min, ignore_nan, expected", test_data + "data, metadata, time_window_min, ignore_nan, expected", test_data_time ) -def test_aggregate_fields_time(R, metadata, time_window_min, ignore_nan, expected): +def test_aggregate_fields_time(data, metadata, time_window_min, ignore_nan, expected): """Test the aggregate_fields_time.""" + dataset_ref = convert_to_xarray_dataset( + data, None, {**fillvalues_metadata, **metadata} + ) + datasets = [] + for timestamp in metadata["timestamps"]: + dataset_ = dataset_ref.copy(deep=True) + dataset_ = dataset_.expand_dims(dim="time", axis=0) + dataset_ = dataset_.assign_coords(time=("time", [timestamp])) + datasets.append(dataset_) + + dataset = xr.concat(datasets, dim="time") assert_array_equal( - dimension.aggregate_fields_time(R, metadata, time_window_min, ignore_nan)[0], + dimension.aggregate_fields_time(dataset, time_window_min, ignore_nan)[ + "precip_intensity" if metadata["unit"] == "mm/h" else "precip_accum" + ].values, expected, ) # aggregate_fields_space -test_data = [ +test_data_space = [ ( - np.ones((1, 10, 10)), - {"unit": "mm/h", "xpixelsize": 1, "ypixelsize": 1}, + np.ones((10, 10)), + { + "unit": "mm/h", + "x1": 0, + "x2": 10, + "y1": 0, + "y2": 10, + "xpixelsize": 1, + "ypixelsize": 1, + }, 2, False, - np.ones((1, 5, 5)), + np.ones((5, 5)), ), ( - np.ones((1, 10, 10)), - {"unit": "mm", "xpixelsize": 1, "ypixelsize": 1}, + np.ones((10, 10)), + { + "unit": "mm", + "x1": 0, + "x2": 10, + "y1": 0, + "y2": 10, + "xpixelsize": 1, + "ypixelsize": 1, + }, 2, False, - np.ones((1, 5, 5)), + np.ones((5, 5)), ), ( - np.ones((1, 10, 10)), - {"unit": "mm/h", "xpixelsize": 1, "ypixelsize": 2}, - (2, 4), + np.ones((10, 10)), + { + "unit": "mm/h", + "x1": 0, + "x2": 10, + "y1": 0, + "y2": 20, + "xpixelsize": 1, + "ypixelsize": 2, + }, + (4, 2), False, - np.ones((1, 5, 5)), + np.ones((5, 5)), ), ] -@pytest.mark.parametrize("R, metadata, space_window, ignore_nan, expected", test_data) -def test_aggregate_fields_space(R, metadata, space_window, ignore_nan, expected): +@pytest.mark.parametrize( + "data, metadata, space_window, ignore_nan, expected", test_data_space +) +def test_aggregate_fields_space(data, metadata, space_window, ignore_nan, expected): """Test the aggregate_fields_space.""" + dataset = convert_to_xarray_dataset(data, None, {**fillvalues_metadata, **metadata}) assert_array_equal( - dimension.aggregate_fields_space(R, metadata, space_window, ignore_nan)[0], + dimension.aggregate_fields_space(dataset, space_window, ignore_nan)[ + "precip_intensity" if metadata["unit"] == "mm/h" else "precip_accum" + ].values, expected, ) @@ -174,64 +246,40 @@ def test_aggregate_fields_space(R, metadata, space_window, ignore_nan, expected) # clip_domain R = np.zeros((4, 4)) R[:2, :] = 1 -test_data = [ +test_data_clip_domain = [ ( R, - { - "x1": 0, - "x2": 4, - "y1": 0, - "y2": 4, - "xpixelsize": 1, - "ypixelsize": 1, - "zerovalue": 0, - "yorigin": "upper", - }, + {"yorigin": "lower"}, None, R, ), ( R, - { - "x1": 0, - "x2": 4, - "y1": 0, - "y2": 4, - "xpixelsize": 1, - "ypixelsize": 1, - "zerovalue": 0, - "yorigin": "lower", - }, + {"yorigin": "lower"}, (2, 4, 2, 4), np.zeros((2, 2)), ), ( R, - { - "x1": 0, - "x2": 4, - "y1": 0, - "y2": 4, - "xpixelsize": 1, - "ypixelsize": 1, - "zerovalue": 0, - "yorigin": "upper", - }, + {"yorigin": "upper"}, (2, 4, 2, 4), np.ones((2, 2)), ), ] -@pytest.mark.parametrize("R, metadata, extent, expected", test_data) +@pytest.mark.parametrize("R, metadata, extent, expected", test_data_clip_domain) def test_clip_domain(R, metadata, extent, expected): """Test the clip_domain.""" - assert_array_equal(dimension.clip_domain(R, metadata, extent)[0], expected) + dataset = convert_to_xarray_dataset(R, None, {**fillvalues_metadata, **metadata}) + assert_array_equal( + dimension.clip_domain(dataset, extent)["precip_intensity"].values, expected + ) # square_domain R = np.zeros((4, 2)) -test_data = [ +test_data_square = [ # square by padding ( R, @@ -258,7 +306,7 @@ def test_clip_domain(R, metadata, extent, expected): "y2": 4, "xpixelsize": 1, "ypixelsize": 1, - "orig_domain": (4, 2), + "orig_domain": (np.array([0.5, 1.5, 2.5, 3.5]), np.array([0.5, 1.5])), "square_method": "pad", }, "pad", @@ -275,7 +323,7 @@ def test_clip_domain(R, metadata, extent, expected): "y2": 3, "xpixelsize": 1, "ypixelsize": 1, - "orig_domain": (4, 2), + "orig_domain": (np.array([0.5, 1.5, 2.5, 3.5]), np.array([0.5, 1.5])), "square_method": "crop", }, "crop", @@ -285,9 +333,15 @@ def test_clip_domain(R, metadata, extent, expected): ] -@pytest.mark.parametrize("R, metadata, method, inverse, expected", test_data) -def test_square_domain(R, metadata, method, inverse, expected): +@pytest.mark.parametrize("data, metadata, method, inverse, expected", test_data_square) +def test_square_domain(data, metadata, method, inverse, expected): """Test the square_domain.""" + dataset = convert_to_xarray_dataset(data, None, {**fillvalues_metadata, **metadata}) + dataset["precip_intensity"].attrs = { + **dataset["precip_intensity"].attrs, + **metadata, + } assert_array_equal( - dimension.square_domain(R, metadata, method, inverse)[0], expected + dimension.square_domain(dataset, method, inverse)["precip_intensity"].values, + expected, ) diff --git a/pysteps/utils/conversion.py b/pysteps/utils/conversion.py index 68228e981..2ea6a3a12 100644 --- a/pysteps/utils/conversion.py +++ b/pysteps/utils/conversion.py @@ -70,8 +70,9 @@ def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): Parameters ---------- - dataset: Dataset - Dataset to be (back-)transformed. + dataset: xarray.Dataset + Dataset to be (back-)transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -83,7 +84,7 @@ def to_rainrate(dataset: xr.Dataset, zr_a=None, zr_b=None): Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the converted units. """ @@ -159,8 +160,9 @@ def to_raindepth(dataset: xr.Dataset, zr_a=None, zr_b=None): Parameters ---------- - dataset: Dataset - Dataset to be (back-)transformed. + dataset: xarray.Dataset + Dataset to be (back-)transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -172,7 +174,7 @@ def to_raindepth(dataset: xr.Dataset, zr_a=None, zr_b=None): Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the converted units. """ @@ -248,8 +250,9 @@ def to_reflectivity(dataset: xr.Dataset, zr_a=None, zr_b=None): Parameters ---------- - dataset: Dataset - Dataset to be (back-)transformed. + dataset: xarray.Dataset + Dataset to be (back-)transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. Additionally, in case of conversion to/from reflectivity units, the zr_a and zr_b attributes are also required, @@ -261,7 +264,7 @@ def to_reflectivity(dataset: xr.Dataset, zr_a=None, zr_b=None): Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the converted units. """ diff --git a/pysteps/utils/dimension.py b/pysteps/utils/dimension.py index 43b7e2ca5..efa459610 100644 --- a/pysteps/utils/dimension.py +++ b/pysteps/utils/dimension.py @@ -14,26 +14,34 @@ clip_domain square_domain """ - import numpy as np +import xarray as xr + +from pysteps.converters import compute_lat_lon _aggregation_methods = dict( sum=np.sum, mean=np.mean, nanmean=np.nanmean, nansum=np.nansum ) -def aggregate_fields_time(R, metadata, time_window_min, ignore_nan=False): +def aggregate_fields_time( + dataset: xr.Dataset, time_window_min, ignore_nan=False +) -> xr.Dataset: """Aggregate fields in time. + It attempts to aggregate the given dataset in the time direction in an integer + number of sections of length = ``time_window_min``. + If such a aggregation is not possible, an error is raised. + The data is aggregated by a method chosen based on the unit of the precipitation + data in the dataset. ``mean`` is used when the unit is ``mm/h`` and ``sum`` + is used when the unit is ``mm``. For other units an error is raised. + Parameters ---------- - R: array-like - Array of shape (t,m,n) or (l,t,m,n) containing - a time series of (ensemble) input fields. + dataset: xarray.Dataset + Dataset containing a time series of (ensemble) input fields + as described in the documentation of :py:mod:`pysteps.io.importers`. They must be evenly spaced in time. - metadata: dict - Metadata dictionary containing the timestamps and unit attributes as - described in the documentation of :py:mod:`pysteps.io.importers`. time_window_min: float or None The length in minutes of the time window that is used to aggregate the fields. @@ -45,12 +53,8 @@ def aggregate_fields_time(R, metadata, time_window_min, ignore_nan=False): Returns ------- - outputarray: array-like - The new array of aggregated fields of shape (k,m,n) or (l,k,m,n), where - k = t*delta/time_window_min and delta is the time interval between two - successive timestamps. - metadata: dict - The metadata with updated attributes. + dataset: xarray.Dataset + The new dataset. See also -------- @@ -58,40 +62,24 @@ def aggregate_fields_time(R, metadata, time_window_min, ignore_nan=False): pysteps.utils.dimension.aggregate_fields """ - R = R.copy() - metadata = metadata.copy() - if time_window_min is None: - return R, metadata - - unit = metadata["unit"] - timestamps = metadata["timestamps"] - if "leadtimes" in metadata: - leadtimes = metadata["leadtimes"] - - if len(R.shape) < 3: - raise ValueError("The number of dimension must be > 2") - if len(R.shape) == 3: - axis = 0 - if len(R.shape) == 4: - axis = 1 - if len(R.shape) > 4: - raise ValueError("The number of dimension must be <= 4") - - if R.shape[axis] != len(timestamps): - raise ValueError( - "The list of timestamps has length %i, " % len(timestamps) - + "but R contains %i frames" % R.shape[axis] - ) + return dataset + + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + + unit = metadata["units"] + + timestamps = dataset["time"].values # assumes that frames are evenly spaced - delta = (timestamps[1] - timestamps[0]).seconds / 60 + delta = (timestamps[1] - timestamps[0]) / np.timedelta64(1, "m") if delta == time_window_min: - return R, metadata - if (R.shape[axis] * delta) % time_window_min: - raise ValueError("time_window_size does not equally split R") + return dataset + if time_window_min % delta: + raise ValueError("time_window_size does not equally split dataset") - nframes = int(time_window_min / delta) + window_size = int(time_window_min / delta) # specify the operator to be used to aggregate # the values within the time window @@ -100,55 +88,47 @@ def aggregate_fields_time(R, metadata, time_window_min, ignore_nan=False): elif unit == "mm": method = "sum" else: - raise ValueError( - "can only aggregate units of 'mm/h' or 'mm'" + " not %s" % unit - ) + raise ValueError(f"can only aggregate units of 'mm/h' or 'mm' not {unit}") if ignore_nan: method = "".join(("nan", method)) - R = aggregate_fields(R, nframes, axis=axis, method=method) - - metadata["accutime"] = time_window_min - metadata["timestamps"] = timestamps[nframes - 1 :: nframes] - if "leadtimes" in metadata: - metadata["leadtimes"] = leadtimes[nframes - 1 :: nframes] + return aggregate_fields(dataset, window_size, dim="time", method=method) - return R, metadata - -def aggregate_fields_space(R, metadata, space_window, ignore_nan=False): +def aggregate_fields_space( + dataset: xr.Dataset, space_window, ignore_nan=False +) -> xr.Dataset: """ Upscale fields in space. + It attempts to aggregate the given dataset in y and x direction in an integer + number of sections of length = ``(window_size_y, window_size_x)``. + If such a aggregation is not possible, an error is raised. + The data is aggregated by computing the mean. Only datasets with precipitation + data in the ``mm`` or ``mm/h`` unit are currently supported. + Parameters ---------- - R: array-like - Array of shape (m,n), (t,m,n) or (l,t,m,n) containing a single field or - a time series of (ensemble) input fields. - metadata: dict - Metadata dictionary containing the xpixelsize, ypixelsize and unit - attributes as described in the documentation of + dataset: xarray.Dataset + Dataset containing a single field or + a time series of (ensemble) input fields as described in the documentation of :py:mod:`pysteps.io.importers`. space_window: float, tuple or None The length of the space window that is used to upscale the fields. If a float is given, the same window size is used for the x- and y-directions. Separate window sizes are used for x- and y-directions if - a two-element tuple is given. The space_window unit is the same used in - the geographical projection of R and hence the same as for the xpixelsize - and ypixelsize attributes. The space spanned by the n- and m-dimensions - of R must be a multiple of space_window. If set to None, the function - returns a copy of the original R and metadata. + a two-element tuple is given (y, x). The space_window unit is the same + as the unit of x and y in the input dataset. The space spanned by the + n- and m-dimensions of the dataset content must be a multiple of space_window. + If set to None, the function returns a copy of the original dataset. ignore_nan: bool, optional If True, ignore nan values. Returns ------- - outputarray: array-like - The new array of aggregated fields of shape (k,j), (t,k,j) or (l,t,k,j), - where k = m*ypixelsize/space_window[1] and j = n*xpixelsize/space_window[0]. - metadata: dict - The metadata with updated attributes. + dataset: xarray.Dataset + The new dataset. See also -------- @@ -156,110 +136,85 @@ def aggregate_fields_space(R, metadata, space_window, ignore_nan=False): pysteps.utils.dimension.aggregate_fields """ - R = R.copy() - metadata = metadata.copy() - if space_window is None: - return R, metadata - - unit = metadata["unit"] - ypixelsize = metadata["ypixelsize"] - xpixelsize = metadata["xpixelsize"] - - if len(R.shape) < 2: - raise ValueError("The number of dimensions must be >= 2") - if len(R.shape) == 2: - axes = [0, 1] - if len(R.shape) == 3: - axes = [1, 2] - if len(R.shape) == 4: - axes = [2, 3] - if len(R.shape) > 4: - raise ValueError("The number of dimensions must be <= 4") + return dataset + + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + + unit = metadata["units"] if np.isscalar(space_window): space_window = (space_window, space_window) # assumes that frames are evenly spaced - if ypixelsize == space_window[1] and xpixelsize == space_window[0]: - return R, metadata - - ysize = R.shape[axes[0]] * ypixelsize - xsize = R.shape[axes[1]] * xpixelsize - - if ( - abs(ysize / space_window[1] - round(ysize / space_window[1])) > 1e-10 - or abs(xsize / space_window[0] - round(xsize / space_window[0])) > 1e-10 - ): - raise ValueError("space_window does not equally split R") + ydelta = dataset["y"].values[1] - dataset["y"].values[0] + xdelta = dataset["x"].values[1] - dataset["x"].values[0] - nframes = [int(space_window[1] / ypixelsize), int(space_window[0] / xpixelsize)] + if space_window[0] % ydelta > 1e-10 or space_window[1] % xdelta > 1e-10: + raise ValueError("space_window does not equally split dataset") # specify the operator to be used to aggregate the values # within the space window if unit == "mm/h" or unit == "mm": method = "mean" else: - raise ValueError( - "can only aggregate units of 'mm/h' or 'mm' " + "not %s" % unit - ) + raise ValueError(f"can only aggregate units of 'mm/h' or 'mm' not {unit}") if ignore_nan: method = "".join(("nan", method)) - R = aggregate_fields(R, nframes[0], axis=axes[0], method=method) - R = aggregate_fields(R, nframes[1], axis=axes[1], method=method) - - metadata["ypixelsize"] = space_window[1] - metadata["xpixelsize"] = space_window[0] + window_size = (int(space_window[0] / ydelta), int(space_window[1] / xdelta)) - return R, metadata + return aggregate_fields(dataset, window_size, ["y", "x"], method) -def aggregate_fields(data, window_size, axis=0, method="mean", trim=False): +def aggregate_fields( + dataset: xr.Dataset, window_size, dim="x", method="mean", trim=False +) -> xr.Dataset: """Aggregate fields along a given direction. - It attempts to aggregate the given R axis in an integer number of sections + It attempts to aggregate the given dataset dim in an integer number of sections of length = ``window_size``. If such a aggregation is not possible, an error is raised unless ``trim`` - set to True, in which case the axis is trimmed (from the end) + set to True, in which case the dim is trimmed (from the end) to make it perfectly divisible". Parameters ---------- - data: array-like - Array of any shape containing the input fields. - window_size: int or tuple of ints + dataset: xarray.Dataset + Dataset containing the input fields as described in the documentation of + :py:mod:`pysteps.io.importers`. + window_size: int or array-like of ints The length of the window that is used to aggregate the fields. If a single integer value is given, the same window is used for - all the selected axis. + all the selected dim. If ``window_size`` is a 1D array-like, each element indicates the length of the window that is used - to aggregate the fields along each axis. In this case, + to aggregate the fields along each dim. In this case, the number of elements of 'window_size' must be the same as the elements - in the ``axis`` argument. - axis: int or array-like of ints - Axis or axes where to perform the aggregation. - If this is a tuple of ints, the aggregation is performed over multiple - axes, instead of a single axis + in the ``dim`` argument. + dim: str or array-like of strs + Dim or dims where to perform the aggregation. + If this is an array-like of strs, the aggregation is performed over multiple + dims, instead of a single dim method: string, optional Optional argument that specifies the operation to use to aggregate the values within the window. Default to mean operator. trim: bool In case that the ``data`` is not perfectly divisible by - ``window_size`` along the selected axis: + ``window_size`` along the selected dim: - trim=True: the data will be trimmed (from the end) along that - axis to make it perfectly divisible. + dim to make it perfectly divisible. - trim=False: a ValueError exception is raised. Returns ------- - new_array: array-like - The new aggregated array with shape[axis] = k, - where k = R.shape[axis] / window_size. + dataset: xarray.Dataset + The new dataset. See also -------- @@ -267,90 +222,60 @@ def aggregate_fields(data, window_size, axis=0, method="mean", trim=False): pysteps.utils.dimension.aggregate_fields_space """ - if np.ndim(axis) > 1: + if np.ndim(dim) > 1: raise TypeError( "Only integers or integer 1D arrays can be used for the " "'axis' argument." ) - if np.ndim(axis) == 1: - axis = np.asarray(axis) - if np.ndim(window_size) == 0: - window_size = (window_size,) * axis.size - - window_size = np.asarray(window_size, dtype="int") - - if window_size.shape != axis.shape: - raise ValueError( - "The 'window_size' and 'axis' shapes are incompatible." - f"window_size.shape: {str(window_size.shape)}, " - f"axis.shape: {str(axis.shape)}, " - ) - - new_data = data.copy() - for i in range(axis.size): - # Recursively call the aggregate_fields function - new_data = aggregate_fields( - new_data, window_size[i], axis=axis[i], method=method, trim=trim - ) - - return new_data + if np.ndim(dim) == 0: + dim = [dim] - if np.ndim(window_size) != 0: - raise TypeError( - "A single axis was selected for the aggregation but several" - f"of window_sizes were given: {str(window_size)}." - ) + if np.ndim(window_size) == 0: + window_size = [window_size for _ in dim] - data = np.asarray(data).copy() - orig_shape = data.shape + if len(window_size) != len(dim): + raise TypeError("The length of window size does not to match the length of dim") if method not in _aggregation_methods: raise ValueError( "Aggregation method not recognized. " f"Available methods: {str(list(_aggregation_methods.keys()))}" ) + for ws in window_size: + if ws <= 0: + raise ValueError("'window_size' must be strictly positive") - if window_size <= 0: - raise ValueError("'window_size' must be strictly positive") + for d, ws in zip(dim, window_size): + if (dataset.sizes[d] % ws) and (not trim): + raise ValueError( + f"Since 'trim' argument was set to False," + f"the 'window_size' {ws} must exactly divide" + f"the dimension along the selected axis:" + f"dataset.sizes[dim]={dataset.sizes[d]}" + ) - if (orig_shape[axis] % window_size) and (not trim): - raise ValueError( - f"Since 'trim' argument was set to False," - f"the 'window_size' {window_size} must exactly divide" - f"the dimension along the selected axis:" - f"data.shape[axis]={orig_shape[axis]}" + # FIXME: The aggregation method is applied to all DataArrays in the Dataset + # Fix to allow support for an aggregation method per DataArray + return ( + dataset.rolling(dict(zip(dim, window_size))) + .reduce(_aggregation_methods[method]) + .isel( + { + d: slice(ws - 1, dataset.sizes[d] - dataset.sizes[d] % ws, ws) + for d, ws in zip(dim, window_size) + } ) - - new_data = data.swapaxes(axis, 0) - if trim: - trim_size = data.shape[axis] % window_size - if trim_size > 0: - new_data = new_data[:-trim_size] - - new_data_shape = list(new_data.shape) - new_data_shape[0] //= window_size # Final shape - - new_data = new_data.reshape(new_data_shape[0], window_size, -1) - - new_data = _aggregation_methods[method](new_data, axis=1) - - new_data = new_data.reshape(new_data_shape).swapaxes(axis, 0) - - return new_data + ) -def clip_domain(R, metadata, extent=None): +def clip_domain(dataset: xr.Dataset, extent=None): """ Clip the field domain by geographical coordinates. Parameters ---------- - R: array-like - Array of shape (m,n) or (t,m,n) containing the input fields. - metadata: dict - Metadata dictionary containing the x1, x2, y1, y2, - xpixelsize, ypixelsize, - zerovalue and yorigin attributes as described in the documentation of + dataset: xarray.Dataset + Dataset containing the input fields as described in the documentation of :py:mod:`pysteps.io.importers`. extent: scalars (left, right, bottom, top), optional The extent of the bounding box in data coordinates to be used to clip @@ -362,107 +287,48 @@ def clip_domain(R, metadata, extent=None): Returns ------- - R: array-like - the clipped array - metadata: dict - the metadata with updated attributes. + dataset: xarray.Dataset + The clipped dataset """ + if extent is None: + return dataset + return dataset.sel(x=slice(extent[0], extent[1]), y=slice(extent[2], extent[3])) - R = R.copy() - R_shape = np.array(R.shape) - metadata = metadata.copy() - if extent is None: - return R, metadata - - if len(R.shape) < 2: - raise ValueError("The number of dimension must be > 1") - if len(R.shape) == 2: - R = R[None, None, :, :] - if len(R.shape) == 3: - R = R[None, :, :, :] - if len(R.shape) > 4: - raise ValueError("The number of dimension must be <= 4") - - # extract original domain coordinates - left = metadata["x1"] - right = metadata["x2"] - bottom = metadata["y1"] - top = metadata["y2"] - - # extract bounding box coordinates - left_ = extent[0] - right_ = extent[1] - bottom_ = extent[2] - top_ = extent[3] - - # compute its extent in pixels - dim_x_ = int((right_ - left_) / metadata["xpixelsize"]) - dim_y_ = int((top_ - bottom_) / metadata["ypixelsize"]) - R_ = np.ones((R.shape[0], R.shape[1], dim_y_, dim_x_)) * metadata["zerovalue"] - - # build set of coordinates for the original domain - y_coord = ( - np.linspace(bottom, top - metadata["ypixelsize"], R.shape[2]) - + metadata["ypixelsize"] / 2.0 - ) - x_coord = ( - np.linspace(left, right - metadata["xpixelsize"], R.shape[3]) - + metadata["xpixelsize"] / 2.0 +def _pad_domain( + dataset: xr.Dataset, dim_to_pad: str, idx_buffer: int, zerovalue: float +) -> xr.Dataset: + # assumes that frames are evenly spaced + delta = dataset[dim_to_pad].values[1] - dataset[dim_to_pad].values[0] + end_values = ( + dataset[dim_to_pad].values[0] - delta * idx_buffer, + dataset[dim_to_pad].values[-1] + delta * idx_buffer, ) - # build set of coordinates for the new domain - y_coord_ = ( - np.linspace(bottom_, top_ - metadata["ypixelsize"], R_.shape[2]) - + metadata["ypixelsize"] / 2.0 + dataset_ref = dataset + + # FIXME: The same zerovalue is used for all DataArrays in the Dataset + # Fix to allow support for a zerovalue per DataArray + dataset = dataset_ref.pad({dim_to_pad: idx_buffer}, constant_values=zerovalue) + dataset[dim_to_pad] = dataset_ref[dim_to_pad].pad( + {dim_to_pad: idx_buffer}, + mode="linear_ramp", + end_values={dim_to_pad: end_values}, ) - x_coord_ = ( - np.linspace(left_, right_ - metadata["xpixelsize"], R_.shape[3]) - + metadata["xpixelsize"] / 2.0 + dataset.lat.data[:], dataset.lon.data[:] = compute_lat_lon( + dataset.x.values, dataset.y.values, dataset.attrs["projection"] ) + return dataset - # origin='upper' reverses the vertical axes direction - if metadata["yorigin"] == "upper": - y_coord = y_coord[::-1] - y_coord_ = y_coord_[::-1] - - # extract original domain - idx_y = np.where(np.logical_and(y_coord < top_, y_coord > bottom_))[0] - idx_x = np.where(np.logical_and(x_coord < right_, x_coord > left_))[0] - - # extract new domain - idx_y_ = np.where(np.logical_and(y_coord_ < top, y_coord_ > bottom))[0] - idx_x_ = np.where(np.logical_and(x_coord_ < right, x_coord_ > left))[0] - - # compose the new array - R_[:, :, idx_y_[0] : (idx_y_[-1] + 1), idx_x_[0] : (idx_x_[-1] + 1)] = R[ - :, :, idx_y[0] : (idx_y[-1] + 1), idx_x[0] : (idx_x[-1] + 1) - ] - - # update coordinates - metadata["y1"] = bottom_ - metadata["y2"] = top_ - metadata["x1"] = left_ - metadata["x2"] = right_ - R_shape[-2] = R_.shape[-2] - R_shape[-1] = R_.shape[-1] - - return R_.reshape(R_shape), metadata - - -def square_domain(R, metadata, method="pad", inverse=False): +def square_domain(dataset: xr.Dataset, method="pad", inverse=False): """ Either pad or crop a field to obtain a square domain. Parameters ---------- - R: array-like - Array of shape (m,n) or (t,m,n) containing the input fields. - metadata: dict - Metadata dictionary containing the x1, x2, y1, y2, - xpixelsize, ypixelsize, - attributes as described in the documentation of + dataset: xarray.Dataset + Dataset containing the input fields as described in the documentation of :py:mod:`pysteps.io.importers`. method: {'pad', 'crop'}, optional Either pad or crop. @@ -477,123 +343,91 @@ def square_domain(R, metadata, method="pad", inverse=False): Returns ------- - R: array-like - the reshape dataset - metadata: dict - the metadata with updated attributes. + dataset: xarray.Dataset + the reshaped dataset """ - R = R.copy() - R_shape = np.array(R.shape) - metadata = metadata.copy() - - if not inverse: - if len(R.shape) < 2: - raise ValueError("The number of dimension must be > 1") - if len(R.shape) == 2: - R = R[None, None, :] - if len(R.shape) == 3: - R = R[None, :] - if len(R.shape) > 4: - raise ValueError("The number of dimension must be <= 4") - - if R.shape[2] == R.shape[3]: - return R.squeeze() - - orig_dim = R.shape - orig_dim_n = orig_dim[0] - orig_dim_t = orig_dim[1] - orig_dim_y = orig_dim[2] - orig_dim_x = orig_dim[3] + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip_data = dataset[precip_var].values + + x_len = len(dataset.x.values) + y_len = len(dataset.y.values) + + if inverse: + if "orig_domain" not in dataset.attrs or "square_method" not in dataset.attrs: + raise ValueError("Attempting to inverse a non squared dataset") + method = dataset.attrs.pop("square_method") + orig_domain = dataset.attrs.pop("orig_domain") if method == "pad": - new_dim = np.max(orig_dim[2:]) - R_ = np.ones((orig_dim_n, orig_dim_t, new_dim, new_dim)) * R.min() - - if orig_dim_x < new_dim: - idx_buffer = int((new_dim - orig_dim_x) / 2.0) - R_[:, :, :, idx_buffer : (idx_buffer + orig_dim_x)] = R - metadata["x1"] -= idx_buffer * metadata["xpixelsize"] - metadata["x2"] += idx_buffer * metadata["xpixelsize"] - - elif orig_dim_y < new_dim: - idx_buffer = int((new_dim - orig_dim_y) / 2.0) - R_[:, :, idx_buffer : (idx_buffer + orig_dim_y), :] = R - metadata["y1"] -= idx_buffer * metadata["ypixelsize"] - metadata["y2"] += idx_buffer * metadata["ypixelsize"] - - elif method == "crop": - new_dim = np.min(orig_dim[2:]) - R_ = np.zeros((orig_dim_n, orig_dim_t, new_dim, new_dim)) - - if orig_dim_x > new_dim: - idx_buffer = int((orig_dim_x - new_dim) / 2.0) - R_ = R[:, :, :, idx_buffer : (idx_buffer + new_dim)] - metadata["x1"] += idx_buffer * metadata["xpixelsize"] - metadata["x2"] -= idx_buffer * metadata["xpixelsize"] - - elif orig_dim_y > new_dim: - idx_buffer = int((orig_dim_y - new_dim) / 2.0) - R_ = R[:, :, idx_buffer : (idx_buffer + new_dim), :] - metadata["y1"] += idx_buffer * metadata["ypixelsize"] - metadata["y2"] -= idx_buffer * metadata["ypixelsize"] - - else: - raise ValueError("Unknown type") - - metadata["orig_domain"] = (orig_dim_y, orig_dim_x) - metadata["square_method"] = method - - R_shape[-2] = R_.shape[-2] - R_shape[-1] = R_.shape[-1] - - return R_.reshape(R_shape), metadata - - elif inverse: - if len(R.shape) < 2: - raise ValueError("The number of dimension must be > 2") - if len(R.shape) == 2: - R = R[None, None, :] - if len(R.shape) == 3: - R = R[None, :] - if len(R.shape) > 4: - raise ValueError("The number of dimension must be <= 4") - - method = metadata.pop("square_method") - shape = metadata.pop("orig_domain") - - if R.shape[2] == shape[0] and R.shape[3] == shape[1]: - return R.squeeze(), metadata - - R_ = np.zeros((R.shape[0], R.shape[1], shape[0], shape[1])) + if x_len > len(orig_domain[1]): + extent = ( + orig_domain[1].min(), + orig_domain[1].max(), + dataset.y.values.min(), + dataset.y.values.max(), + ) + elif y_len > len(orig_domain[0]): + extent = ( + dataset.x.values.min(), + dataset.x.values.max(), + orig_domain[0].min(), + orig_domain[0].max(), + ) + else: + return dataset + return clip_domain(dataset, extent) + + if method == "crop": + if x_len < len(orig_domain[1]): + dim_to_pad = "x" + idx_buffer = int((len(orig_domain[1]) - x_len) / 2.0) + elif y_len < len(orig_domain[0]): + dim_to_pad = "y" + idx_buffer = int((len(orig_domain[0]) - y_len) / 2.0) + else: + return dataset + return _pad_domain(dataset, dim_to_pad, idx_buffer, np.nanmin(precip_data)) + + raise ValueError(f"Unknown square method: {method}") + + else: + if "orig_domain" in dataset.attrs and "square_method" in dataset.attrs: + raise ValueError("Attempting to square an already squared dataset") + dataset.attrs["orig_domain"] = (dataset.y.values, dataset.x.values) + dataset.attrs["square_method"] = method if method == "pad": - if R.shape[2] == shape[0]: - idx_buffer = int((R.shape[3] - shape[1]) / 2.0) - R_ = R[:, :, :, idx_buffer : (idx_buffer + shape[1])] - metadata["x1"] += idx_buffer * metadata["xpixelsize"] - metadata["x2"] -= idx_buffer * metadata["xpixelsize"] - - elif R.shape[3] == shape[1]: - idx_buffer = int((R.shape[2] - shape[0]) / 2.0) - R_ = R[:, :, idx_buffer : (idx_buffer + shape[0]), :] - metadata["y1"] += idx_buffer * metadata["ypixelsize"] - metadata["y2"] -= idx_buffer * metadata["ypixelsize"] - - elif method == "crop": - if R.shape[2] == shape[0]: - idx_buffer = int((shape[1] - R.shape[3]) / 2.0) - R_[:, :, :, idx_buffer : (idx_buffer + R.shape[3])] = R - metadata["x1"] -= idx_buffer * metadata["xpixelsize"] - metadata["x2"] += idx_buffer * metadata["xpixelsize"] - - elif R.shape[3] == shape[1]: - idx_buffer = int((shape[0] - R.shape[2]) / 2.0) - R_[:, :, idx_buffer : (idx_buffer + R.shape[2]), :] = R - metadata["y1"] -= idx_buffer * metadata["ypixelsize"] - metadata["y2"] += idx_buffer * metadata["ypixelsize"] - - R_shape[-2] = R_.shape[-2] - R_shape[-1] = R_.shape[-1] - - return R_.reshape(R_shape), metadata + if x_len > y_len: + dim_to_pad = "y" + idx_buffer = int((x_len - y_len) / 2.0) + elif y_len > x_len: + dim_to_pad = "x" + idx_buffer = int((y_len - x_len) / 2.0) + else: + return dataset + return _pad_domain(dataset, dim_to_pad, idx_buffer, np.nanmin(precip_data)) + + if method == "crop": + if x_len > y_len: + idx_buffer = int((x_len - y_len) / 2.0) + extent = ( + dataset.x.values[idx_buffer], + dataset.x.values[-idx_buffer - 1], + dataset.y.values.min(), + dataset.y.values.max(), + ) + elif y_len > x_len: + idx_buffer = int((y_len - x_len) / 2.0) + extent = ( + dataset.x.values.min(), + dataset.x.values.max(), + dataset.y.values[idx_buffer], + dataset.y.values[-idx_buffer - 1], + ) + else: + return dataset + return clip_domain(dataset, extent) + + raise ValueError(f"Unknown square method: {method}") diff --git a/pysteps/utils/transformation.py b/pysteps/utils/transformation.py index 1977583c6..3e48fe0d8 100644 --- a/pysteps/utils/transformation.py +++ b/pysteps/utils/transformation.py @@ -41,8 +41,9 @@ def boxcox_transform( Parameters ---------- - dataset: Dataset - Dataset to be transformed. + dataset: xarray.Dataset + Dataset to be transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. Lambda: float, optional Parameter Lambda of the Box-Cox transformation. It is 0 by default, which produces the log transformation. @@ -62,7 +63,7 @@ def boxcox_transform( Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the (back-)transformed units. References @@ -146,8 +147,9 @@ def dB_transform( Parameters ---------- - dataset: Dataset - Dataset to be (back-)transformed. + dataset: xarray.Dataset + Dataset to be (back-)transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. threshold: float, optional Optional value that is used for thresholding with the same units as in the dataset. If None, the threshold contained in metadata is used. @@ -161,7 +163,7 @@ def dB_transform( Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the (back-)transformed units. """ @@ -223,8 +225,9 @@ def NQ_transform(dataset: xr.Dataset, inverse: bool = False, **kwargs) -> xr.Dat Parameters ---------- - dataset: Dataset - Dataset to be transformed. + dataset: xarray.Dataset + Dataset to be transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. inverse: bool, optional If set to True, it performs the inverse transform. False by default. @@ -238,7 +241,7 @@ def NQ_transform(dataset: xr.Dataset, inverse: bool = False, **kwargs) -> xr.Dat Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the (back-)transformed units. References @@ -309,14 +312,15 @@ def sqrt_transform(dataset: xr.Dataset, inverse: bool = False, **kwargs) -> xr.D Parameters ---------- - dataset: Dataset - Dataset to be transformed. + dataset: xarray.Dataset + Dataset to be transformed as described in the documentation of + :py:mod:`pysteps.io.importers`. inverse: bool, optional If set to True, it performs the inverse transform. False by default. Returns ------- - dataset: Dataset + dataset: xarray.Dataset Dataset containing the (back-)transformed units. """ From e7c081c1d71b3ae6093489b70fe1a7a79b526de1 Mon Sep 17 00:00:00 2001 From: mats-knmi <145579783+mats-knmi@users.noreply.github.com> Date: Mon, 2 Sep 2024 14:51:27 +0200 Subject: [PATCH 7/8] make all nowcast methods xarray compatible (#414) * make test steps skill run * undo accidental change * make steps nowcast xarray compatible * wrap all nowcasts in xarray * fix dimension.py tests * update dimension.py to work with new dataarrays * fix test_nowcast_utils tests * update docs and make xarray usage more explicit in nowcasts * update docs and make xarray usage in motion methods more explicit --- pysteps/decorators.py | 10 +- pysteps/io/importers.py | 43 +++-- pysteps/io/readers.py | 16 +- pysteps/motion/constant.py | 26 ++- pysteps/motion/darts.py | 29 ++-- pysteps/motion/lucaskanade.py | 61 +++---- pysteps/motion/proesmans.py | 44 +++-- pysteps/motion/vet.py | 39 +++-- pysteps/nowcasts/anvil.py | 60 ++++--- pysteps/nowcasts/extrapolation.py | 44 +++-- pysteps/nowcasts/lagrangian_probability.py | 46 ++--- pysteps/nowcasts/linda.py | 77 ++++----- pysteps/nowcasts/sprog.py | 62 +++---- pysteps/nowcasts/sseps.py | 78 +++++---- pysteps/nowcasts/steps.py | 57 +++---- pysteps/nowcasts/utils.py | 9 +- pysteps/tests/test_motion_lk.py | 16 +- pysteps/tests/test_nowcasts_anvil.py | 22 ++- .../test_nowcasts_lagrangian_probability.py | 63 +++++-- pysteps/tests/test_nowcasts_linda.py | 91 ++++++---- pysteps/tests/test_nowcasts_sprog.py | 20 +-- pysteps/tests/test_nowcasts_sseps.py | 28 +-- pysteps/tests/test_nowcasts_steps.py | 9 +- pysteps/tests/test_nowcasts_utils.py | 7 +- pysteps/tests/test_utils_dimension.py | 160 ++++++++++++++++-- pysteps/utils/dimension.py | 122 ++++++++++--- pysteps/{converters.py => xarray_helpers.py} | 67 +++++++- 27 files changed, 865 insertions(+), 441 deletions(-) rename pysteps/{converters.py => xarray_helpers.py} (76%) diff --git a/pysteps/decorators.py b/pysteps/decorators.py index ee421977e..69c9945bc 100644 --- a/pysteps/decorators.py +++ b/pysteps/decorators.py @@ -22,7 +22,7 @@ import numpy as np -from pysteps.converters import convert_to_xarray_dataset +from pysteps.xarray_helpers import convert_input_to_xarray_dataset def _add_extra_kwrds_to_docstrings(target_func, extra_kwargs_doc_text): @@ -90,7 +90,9 @@ def _import_with_postprocessing(*args, **kwargs): mask = ~np.isfinite(precip) precip[mask] = _fillna - return convert_to_xarray_dataset(precip.astype(_dtype), quality, metadata) + return convert_input_to_xarray_dataset( + precip.astype(_dtype), quality, metadata + ) extra_kwargs_doc = """ Other Parameters @@ -126,7 +128,9 @@ def new_function(*args, **kwargs): target motion_method_func function. """ - input_images = args[0] + dataset = args[0] + precip_var = dataset.attrs["precip_var"] + input_images = dataset[precip_var].values if input_images.ndim != 3: raise ValueError( "input_images dimension mismatch.\n" diff --git a/pysteps/io/importers.py b/pysteps/io/importers.py index f61d4b25b..71fe6dc78 100644 --- a/pysteps/io/importers.py +++ b/pysteps/io/importers.py @@ -74,22 +74,27 @@ .. tabularcolumns:: |p{2cm}|L| -+---------------+-------------------------------------------------------------------------------------------+ -| Coordinate | Description | -+===============+===========================================================================================+ -| y | y-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | -+---------------+-------------------------------------------------------------------------------------------+ -| x | x-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | -+---------------+-------------------------------------------------------------------------------------------+ -| lat | latitude coordinate in degrees | -+---------------+-------------------------------------------------------------------------------------------+ -| lon | longitude coordinate in degrees | -+---------------+-------------------------------------------------------------------------------------------+ -| time | forecast time in seconds since forecast start time | -+---------------+-------------------------------------------------------------------------------------------+ -| member | ensemble member number (integer) | -+---------------+-------------------------------------------------------------------------------------------+ - ++--------------------+-------------------------------------------------------------------------------------------+ +| Coordinate | Description | ++====================+===========================================================================================+ +| y | y-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | ++--------------------+-------------------------------------------------------------------------------------------+ +| x | x-coordinate in Cartesian system, with units determined by ``metadata["cartesian_unit"]`` | ++--------------------+-------------------------------------------------------------------------------------------+ +| lat | latitude coordinate in degrees | ++--------------------+-------------------------------------------------------------------------------------------+ +| lon | longitude coordinate in degrees | ++--------------------+-------------------------------------------------------------------------------------------+ +| time | forecast time in seconds since forecast start time | ++--------------------+-------------------------------------------------------------------------------------------+ +| ens_number | ensemble member number (integer) | ++--------------------+-------------------------------------------------------------------------------------------+ +| direction | used by proesmans to return the forward and backward advection and consistency fields | ++--------------------+-------------------------------------------------------------------------------------------+ + +The time, x and y dimensions all MUST be regularly spaced, with the stepsize included +in a ``stepsize`` attribute. The stepsize is given in the unit of the dimension (this +is alwyas seconds for the time dimension). The dataset can contain the following data variables: @@ -102,8 +107,14 @@ | precip_accum | precip_intensity if unit is ``mm/h``, precip_accum if unit is ``mm`` and reflectivity if unit is ``dBZ``, | | or reflectivity | the attributes of this variable contain metadata relevant to this attribute (see below) | +-------------------+-----------------------------------------------------------------------------------------------------------+ +| velocity_x | x-component of the advection field in cartesian_unit per timestep | ++-------------------+-----------------------------------------------------------------------------------------------------------+ +| velocity_y | y-component of the advection field in cartesian_unit per timestep | ++-------------------+-----------------------------------------------------------------------------------------------------------+ | quality | value between 0 and 1 denoting the quality of the precipitation data, currently not used for anything | +-------------------+-----------------------------------------------------------------------------------------------------------+ +| velocity_quality | value between 0 and 1 denoting the quality of the velocity data, currently only returned by proesmans | ++-------------------+-----------------------------------------------------------------------------------------------------------+ Some of the metadata in the metadata dictionary is not explicitely stored in the dataset, but is still implicitly present. For example ``x1`` can easily be found by taking the first diff --git a/pysteps/io/readers.py b/pysteps/io/readers.py index 7295b5ff7..30c4d4fc0 100644 --- a/pysteps/io/readers.py +++ b/pysteps/io/readers.py @@ -15,7 +15,7 @@ import xarray as xr -def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None: +def read_timeseries(inputfns, importer, timestep=None, **kwargs) -> xr.Dataset | None: """ Read a time series of input files using the methods implemented in the :py:mod:`pysteps.io.importers` module and stack them into a 3d xarray @@ -28,6 +28,9 @@ def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None: :py:mod:`pysteps.io.archive` module. importer: function A function implemented in the :py:mod:`pysteps.io.importers` module. + timestep: int, optional + The timestep in seconds, this value is optional if more than 1 inputfns + are given. kwargs: dict Optional keyword arguments for the importer. @@ -58,6 +61,16 @@ def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None: return None startdate = min(inputfns[1]) + sorted_dates = sorted(inputfns[1]) + timestep_dates = int((sorted_dates[1] - sorted_dates[0]).total_seconds()) + + if timestep is None: + timestep = timestep_dates + if timestep != timestep_dates: + raise ValueError("given timestep does not match inputfns") + for i in range(len(sorted_dates) - 1): + if int((sorted_dates[i + 1] - sorted_dates[i]).total_seconds()) != timestep: + raise ValueError("supplied dates are not evenly spaced") datasets = [] for i, ifn in enumerate(inputfns[0]): @@ -73,6 +86,7 @@ def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None: { "long_name": "forecast time", "units": f"seconds since {startdate:%Y-%m-%d %H:%M:%S}", + "stepsize": timestep, }, ) ) diff --git a/pysteps/motion/constant.py b/pysteps/motion/constant.py index a5c153616..a26831ac0 100644 --- a/pysteps/motion/constant.py +++ b/pysteps/motion/constant.py @@ -14,27 +14,32 @@ import numpy as np import scipy.optimize as op +import xarray as xr from scipy.ndimage import map_coordinates -def constant(R, **kwargs): +def constant(dataset: xr.Dataset, **kwargs): """ Compute a constant advection field by finding a translation vector that maximizes the correlation between two successive images. Parameters ---------- - R: array_like - Array of shape (T,m,n) containing a sequence of T two-dimensional input - images of shape (m,n). If T > 2, two last elements along axis 0 are used. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain a precipitation data variable. + The dataset has to have a time dimension. If the size of this dimension + is larger than 2, the last 2 entries of this dimension are used. Returns ------- - out: array_like - The constant advection field having shape (2, m, n), where out[0, :, :] - contains the x-components of the motion vectors and out[1, :, :] - contains the y-components. + out: xarray.Dataset + The input dataset with the constant advection field added in the ``velocity_x`` + and ``velocity_y`` data variables. """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + R = dataset[precip_var].values m, n = R.shape[1:] X, Y = np.meshgrid(np.arange(n), np.arange(m)) @@ -51,4 +56,7 @@ def f(v): options = {"initial_simplex": (np.array([(0, 1), (1, 0), (1, 1)]))} result = op.minimize(f, (1, 1), method="Nelder-Mead", options=options) - return np.stack([-result.x[0] * np.ones((m, n)), -result.x[1] * np.ones((m, n))]) + output = np.stack([-result.x[0] * np.ones((m, n)), -result.x[1] * np.ones((m, n))]) + dataset["velocity_x"] = (["y", "x"], output[0]) + dataset["velocity_y"] = (["y", "x"], output[1]) + return dataset diff --git a/pysteps/motion/darts.py b/pysteps/motion/darts.py index 4e5050d48..4aac80cd3 100644 --- a/pysteps/motion/darts.py +++ b/pysteps/motion/darts.py @@ -11,8 +11,10 @@ DARTS """ -import numpy as np import time + +import numpy as np +import xarray as xr from numpy.linalg import lstsq, svd from pysteps import utils @@ -20,16 +22,17 @@ @check_input_frames(just_ndim=True) -def DARTS(input_images, **kwargs): +def DARTS(dataset: xr.Dataset, **kwargs): """ Compute the advection field from a sequence of input images by using the DARTS method. :cite:`RCW2011` Parameters ---------- - input_images: array-like - Array of shape (T,m,n) containing a sequence of T two-dimensional input - images of shape (m,n). + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain a precipitation data variable. + The dataset has to have a time dimension. Other Parameters ---------------- @@ -67,13 +70,15 @@ def DARTS(input_images, **kwargs): Returns ------- - out: ndarray - Three-dimensional array (2,m,n) containing the dense x- and y-components - of the motion field in units of pixels / timestep as given by the input - array R. + out: xarray.Dataset + The input dataset with the advection field added in the ``velocity_x`` + and ``velocity_y`` data variables. """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + input_images = dataset[precip_var].values N_x = kwargs.get("N_x", 50) N_y = kwargs.get("N_y", 50) N_t = kwargs.get("N_t", 4) @@ -214,10 +219,14 @@ def DARTS(input_images, **kwargs): fft.ifft2(_fill(V, input_images.shape[0], input_images.shape[1], k_x, k_y)) ) + output = np.stack([U, V]) + dataset["velocity_x"] = (["y", "x"], output[0]) + dataset["velocity_y"] = (["y", "x"], output[1]) + if verbose: print("--- %s seconds ---" % (time.time() - t0)) - return np.stack([U, V]) + return dataset def _leastsq(A, B, y): diff --git a/pysteps/motion/lucaskanade.py b/pysteps/motion/lucaskanade.py index 133f860b7..b7a51a26b 100644 --- a/pysteps/motion/lucaskanade.py +++ b/pysteps/motion/lucaskanade.py @@ -22,22 +22,22 @@ dense_lucaskanade """ +import time + import numpy as np +import xarray as xr from numpy.ma.core import MaskedArray +from pysteps import feature, utils from pysteps.decorators import check_input_frames - -from pysteps import utils, feature from pysteps.tracking.lucaskanade import track_features from pysteps.utils.cleansing import decluster, detect_outliers from pysteps.utils.images import morph_opening -import time - @check_input_frames(2) def dense_lucaskanade( - input_images, + dataset: xr.Dataset, lk_kwargs=None, fd_method="shitomasi", fd_kwargs=None, @@ -73,18 +73,14 @@ def dense_lucaskanade( Parameters ---------- - input_images: ndarray_ or MaskedArray_ - Array of shape (T, m, n) containing a sequence of *T* two-dimensional - input images of shape (m, n). The indexing order in **input_images** is - assumed to be (time, latitude, longitude). - - *T* = 2 is the minimum required number of images. - With *T* > 2, all the resulting sparse vectors are pooled together for - the final interpolation on a regular grid. - - In case of ndarray_, invalid values (Nans or infs) are masked, - otherwise the mask of the MaskedArray_ is used. Such mask defines a - region where features are not detected for the tracking algorithm. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain a precipitation data variable. + The dataset has to have a time dimension. The size of the time dimension needs to + be at least 2. If it is larger than 2, all the resulting sparse vectors are pooled + together for the final interpolation on a regular grid. Invalid values (Nans or infs) + are masked. This mask defines a region where features are not detected for the tracking + algorithm. lk_kwargs: dict, optional Optional dictionary containing keyword arguments for the `Lucas-Kanade`_ @@ -151,14 +147,10 @@ def dense_lucaskanade( Returns ------- - out: ndarray_ or tuple - If **dense=True** (the default), return the advection field having shape - (2, m, n), where out[0, :, :] contains the x-components of the motion - vectors and out[1, :, :] contains the y-components. - The velocities are in units of pixels / timestep, where timestep is the - time difference between the two input images. - Return a zero motion field of shape (2, m, n) when no motion is - detected. + out: xarray.Dataset or tuple + If **dense=True** (the default), return the input dataset with the advection + field added in the ``velocity_x`` and ``velocity_y`` data variables. + Return a zero motion field when no motion is detected. If **dense=False**, it returns a tuple containing the 2-dimensional arrays **xy** and **uv**, where x, y define the vector locations, @@ -179,7 +171,9 @@ def dense_lucaskanade( Understanding Workshop, pp. 121–130, 1981. """ - input_images = input_images.copy() + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + input_images = dataset[precip_var].values if verbose: print("Computing the motion field with the Lucas-Kanade method.") @@ -244,7 +238,10 @@ def dense_lucaskanade( # return zero motion field is no sparse vectors are found if xy.shape[0] == 0: if dense: - return np.zeros((2, domain_size[0], domain_size[1])) + uvgrid = np.zeros((2, domain_size[0], domain_size[1])) + dataset["velocity_x"] = (["y", "x"], uvgrid[0]) + dataset["velocity_y"] = (["y", "x"], uvgrid[1]) + return dataset else: return xy, uv @@ -266,14 +263,20 @@ def dense_lucaskanade( # return zero motion field if no sparse vectors are left for interpolation if xy.shape[0] == 0: - return np.zeros((2, domain_size[0], domain_size[1])) + uvgrid = np.zeros((2, domain_size[0], domain_size[1])) + dataset["velocity_x"] = (["y", "x"], uvgrid[0]) + dataset["velocity_y"] = (["y", "x"], uvgrid[1]) + return dataset # interpolation xgrid = np.arange(domain_size[1]) ygrid = np.arange(domain_size[0]) uvgrid = interpolation_method(xy, uv, xgrid, ygrid, **interp_kwargs) + dataset["velocity_x"] = (["y", "x"], uvgrid[0]) + dataset["velocity_y"] = (["y", "x"], uvgrid[1]) + if verbose: print("--- total time: %.2f seconds ---" % (time.time() - t0)) - return uvgrid + return dataset diff --git a/pysteps/motion/proesmans.py b/pysteps/motion/proesmans.py index 8760092ba..4b122a620 100644 --- a/pysteps/motion/proesmans.py +++ b/pysteps/motion/proesmans.py @@ -12,6 +12,7 @@ """ import numpy as np +import xarray as xr from scipy.ndimage import gaussian_filter from pysteps.decorators import check_input_frames @@ -20,7 +21,7 @@ @check_input_frames(2, 2) def proesmans( - input_images, + dataset: xr.Dataset, lam=50.0, num_iter=100, num_levels=6, @@ -34,8 +35,11 @@ def proesmans( Parameters ---------- - input_images: array_like - Array of shape (2, m, n) containing the first and second input image. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain a precipitation data variable. + The dataset has to have a time dimension. The size of this dimension + has to be 2. lam: float Multiplier of the smoothness term. Smaller values give a smoother motion field. @@ -49,22 +53,20 @@ def proesmans( verbose: bool, optional Verbosity enabled if True (default). full_output: bool, optional - If True, the output is a two-element tuple containing the - forward-backward advection and consistency fields. The first element - is shape (2, 2, m, n), where the index along the first dimension refers - to the forward and backward advection fields. The second element is an - array of shape (2, m, n), where the index along the first dimension - refers to the forward and backward consistency fields. - Default: False. + If True, both the forward and backwards advection fields are returned + and the consistency fields are returned as well in the ``velocity_quality`` + data variable. Returns ------- out: ndarray - If full_output=False, the advection field having shape (2, m, n), where - out[0, :, :] contains the x-components of the motion vectors and - out[1, :, :] contains the y-components. The velocities are in units of - pixels / timestep, where timestep is the time difference between the - two input images. + The input dataset with the advection field added in the ``velocity_x`` + and ``velocity_y`` data variables. + + If full_output=True, a ``velocity_direction`` dimension + is added to the dataset, so that the velocity data can be returned containing + the forward and backwards advection fields. Also the ``velocity_quality`` data + coordinate is present containing the forward and backward consistency fields. References ---------- @@ -73,6 +75,9 @@ def proesmans( """ del verbose # Not used + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + input_images = dataset[precip_var].values im1 = input_images[-2, :, :].copy() im2 = input_images[-1, :, :].copy() @@ -89,6 +94,11 @@ def proesmans( advfield, quality = _compute_advection_field(im, lam, num_iter, num_levels) if not full_output: - return advfield[0] + dataset["velocity_x"] = (["y", "x"], advfield[0, 0]) + dataset["velocity_y"] = (["y", "x"], advfield[0, 1]) else: - return advfield, quality + dataset["velocity_x"] = (["direction", "y", "x"], advfield[:, 0]) + dataset["velocity_y"] = (["direction", "y", "x"], advfield[:, 1]) + dataset["velocity_quality"] = (["direction", "y", "x"], quality) + + return dataset diff --git a/pysteps/motion/vet.py b/pysteps/motion/vet.py index 391ebe189..f30703bee 100644 --- a/pysteps/motion/vet.py +++ b/pysteps/motion/vet.py @@ -35,12 +35,13 @@ """ import numpy +import xarray as xr from numpy.ma.core import MaskedArray from scipy.ndimage import zoom from scipy.optimize import minimize from pysteps.decorators import check_input_frames -from pysteps.motion._vet import _warp, _cost_function +from pysteps.motion._vet import _cost_function, _warp def round_int(scalar): @@ -301,7 +302,7 @@ def vet_cost_function( @check_input_frames(2, 3) def vet( - input_images, + dataset: xr.Dataset, sectors=((32, 16, 4, 2), (32, 16, 4, 2)), smooth_gain=1e6, first_guess=None, @@ -366,15 +367,13 @@ def vet( Parameters ---------- - input_images: ndarray_ or MaskedArray - Input images, sequence of 2D arrays, or 3D arrays. - The first dimension represents the images time dimension. - - The template_image (first element in first dimensions) denotes the + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain a precipitation data variable. + The dataset has to have a time dimension. The size of this dimension + has to be 2. The first element in the time dimension denotes the reference image used to obtain the displacement (2D array). The second is the target image. - - The expected dimensions are (2,ni,nj). sectors: list or array, optional Number of sectors on each dimension used in the scaling procedure. If dimension is 1, the same sectors will be used both image dimensions @@ -411,13 +410,11 @@ def vet( Returns ------- - displacement_field: ndarray_ - Displacement Field (2D array representing the transformation) that - warps the template image into the input image. - The dimensions are (2,ni,nj), where the first - dimension indicates the displacement along x (0) or y (1) in units of - pixels / timestep as given by the input_images array. - intermediate_steps: list of ndarray_ + out: xarray.Dataset + The input dataset with the displacement field that + warps the template image into the input image added in the ``velocity_x`` + and ``velocity_y`` data variables. + intermediate_steps: list of ndarray_, optional List with the first guesses obtained during the scaling procedure. References @@ -437,6 +434,9 @@ def vet( Nocedal, J, and S J Wright. 2006. Numerical Optimization. Springer New York. """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + input_images = dataset[precip_var].values if verbose: def debug_print(*args, **kwargs): @@ -642,7 +642,10 @@ def debug_print(*args, **kwargs): if padding > 0: first_guess = first_guess[:, padding:-padding, padding:-padding] + dataset["velocity_x"] = (["y", "x"], first_guess[0]) + dataset["velocity_y"] = (["y", "x"], first_guess[1]) + if intermediate_steps: - return first_guess, scaling_guesses + return dataset, scaling_guesses - return first_guess + return dataset diff --git a/pysteps/nowcasts/anvil.py b/pysteps/nowcasts/anvil.py index f5af038bb..88ed6b0af 100644 --- a/pysteps/nowcasts/anvil.py +++ b/pysteps/nowcasts/anvil.py @@ -21,11 +21,13 @@ import time import numpy as np +import xarray as xr from scipy.ndimage import gaussian_filter from pysteps import cascade, extrapolation, utils from pysteps.nowcasts.utils import nowcast_main_loop from pysteps.timeseries import autoregression +from pysteps.xarray_helpers import convert_output_to_xarray_dataset try: import dask @@ -36,10 +38,8 @@ def forecast( - vil, - velocity, + dataset: xr.Dataset, timesteps, - rainrate=None, n_cascade_levels=6, extrap_method="semilagrangian", ar_order=2, @@ -70,22 +70,21 @@ def forecast( Parameters ---------- - vil: array_like - Array of shape (ar_order+2,m,n) containing the input fields ordered by - timestamp from oldest to newest. The inputs are expected to contain VIL - or rain rate. The time steps between the inputs are assumed to be regular. - velocity: array_like - Array of shape (2,m,n) containing the x- and y-components of the - advection field. The velocities are assumed to represent one time step - between the inputs. All values are required to be finite. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as either VIL values in the + ``precip_accum`` data variable or rainrate in the ``precip_intensity`` + data variable. The time dimension of the dataset has to be size + ``ar_order + 2`` and the precipitation variable has to have this dimension. + When VIL values are supplied, optionally ``precip_accum`` can be supplied + as well without a time dimension, containing the most recently observed rain + rate field. If not supplied, no R(VIL) conversion is done and the outputs + are in the same units as the inputs. timesteps: int or list of floats Number of time steps to forecast or a list of time steps for which the forecasts are computed (relative to the input time step). The elements of the list are required to be in ascending order. - rainrate: array_like - Array of shape (m,n) containing the most recently observed rain rate - field. If set to None, no R(VIL) conversion is done and the outputs - are in the same units as the inputs. n_cascade_levels: int, optional The number of cascade levels to use. Defaults to 6, see issue #385 on GitHub. @@ -128,18 +127,28 @@ def forecast( Returns ------- - out: ndarray - A three-dimensional array of shape (num_timesteps,m,n) containing a time - series of forecast precipitation fields. The time series starts from - t0+timestep, where timestep is taken from the input VIL/rain rate - fields. If measure_time is True, the return value is a three-element - tuple containing the nowcast array, the initialization time of the - nowcast generator and the time used in the main loop (seconds). + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast + precipitation fields. Otherwise, a None value + is returned. The time series starts from t0+timestep, where timestep is + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). References ---------- :cite:`PCLH2020` """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + vil = dataset[precip_var].values + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) + rainrate = None + if precip_var == "precip_intensity" and "precip_accum" in dataset: + rainrate = dataset["precip_accum"].values + _check_inputs(vil, rainrate, velocity, timesteps, ar_order) if extrap_kwargs is None: @@ -292,8 +301,6 @@ def worker(vil, i): print("Starting nowcast computation.") - rainrate_f = [] - extrap_kwargs["return_displacement"] = True state = {"vil_dec": vil_dec} @@ -323,10 +330,11 @@ def worker(vil, i): if measure_time: rainrate_f, mainloop_time = rainrate_f + output_dataset = convert_output_to_xarray_dataset(dataset, timesteps, rainrate_f) if measure_time: - return np.stack(rainrate_f), init_time, mainloop_time + return output_dataset, init_time, mainloop_time else: - return np.stack(rainrate_f) + return output_dataset def _check_inputs(vil, rainrate, velocity, timesteps, ar_order): diff --git a/pysteps/nowcasts/extrapolation.py b/pysteps/nowcasts/extrapolation.py index 143a39d7c..a70b6985c 100644 --- a/pysteps/nowcasts/extrapolation.py +++ b/pysteps/nowcasts/extrapolation.py @@ -11,14 +11,16 @@ """ import time + import numpy as np +import xarray as xr from pysteps import extrapolation +from pysteps.xarray_helpers import convert_output_to_xarray_dataset def forecast( - precip, - velocity, + dataset: xr.Dataset, timesteps, extrap_method="semilagrangian", extrap_kwargs=None, @@ -32,13 +34,11 @@ def forecast( Parameters ---------- - precip: array-like - Two-dimensional array of shape (m,n) containing the input precipitation - field. - velocity: array-like - Array of shape (2,m,n) containing the x- and y-components of the - advection field. The velocities are assumed to represent one time step - between the inputs. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as any pecipitation data variable. + It should contain a time dimension of size 1. timesteps: int or list of floats Number of time steps to forecast or a list of time steps for which the forecasts are computed (relative to the input time step). The elements @@ -54,18 +54,25 @@ def forecast( Returns ------- - out: ndarray_ - Three-dimensional array of shape (num_timesteps, m, n) containing a time - series of nowcast precipitation fields. The time series starts from - t0 + timestep, where timestep is taken from the advection field velocity. - If *measure_time* is True, the return value is a two-element tuple - containing this array and the computation time (seconds). + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast + precipitation fields. Otherwise, a None value + is returned. The time series starts from t0+timestep, where timestep is + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). See also -------- pysteps.extrapolation.interface """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip = dataset[precip_var].values[0] + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) _check_inputs(precip, velocity, timesteps) if extrap_kwargs is None: @@ -95,10 +102,13 @@ def forecast( computation_time = time.time() - start_time print(f"{computation_time:.2f} seconds.") + output_dataset = convert_output_to_xarray_dataset( + dataset, timesteps, precip_forecast + ) if measure_time: - return precip_forecast, computation_time + return output_dataset, computation_time else: - return precip_forecast + return output_dataset def _check_inputs(precip, velocity, timesteps): diff --git a/pysteps/nowcasts/lagrangian_probability.py b/pysteps/nowcasts/lagrangian_probability.py index 727e94806..7bae440cc 100644 --- a/pysteps/nowcasts/lagrangian_probability.py +++ b/pysteps/nowcasts/lagrangian_probability.py @@ -12,20 +12,20 @@ """ import numpy as np +import xarray as xr from scipy.signal import convolve from pysteps.nowcasts import extrapolation def forecast( - precip, - velocity, + dataset: xr.Dataset, timesteps, threshold, extrap_method="semilagrangian", extrap_kwargs=None, slope=5, -): +) -> xr.Dataset: """ Generate a probability nowcast by a local lagrangian approach. The ouput is the probability of exceeding a given intensity threshold, i.e. @@ -33,13 +33,11 @@ def forecast( Parameters ---------- - precip: array_like - Two-dimensional array of shape (m,n) containing the input precipitation - field. - velocity: array_like - Array of shape (2,m,n) containing the x- and y-components of the - advection field. The velocities are assumed to represent one time step - between the inputs. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as any pecipitation data variable. + It should contain a time dimension of size 1. timesteps: int or list of floats Number of time steps to forecast or a sorted list of time steps for which the forecasts are computed (relative to the input time step). @@ -54,10 +52,15 @@ def forecast( Returns ------- - out: ndarray - Three-dimensional array of shape (num_timesteps, m, n) containing a time - series of nowcast exceedence probabilities. The time series starts from - t0 + timestep, where timestep is taken from the advection field velocity. + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast + precipitation fields. Otherwise, a None value + is returned. The time series starts from t0+timestep, where timestep is + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). References ---------- @@ -68,16 +71,14 @@ def forecast( """ # Compute deterministic extrapolation forecast if isinstance(timesteps, int) and timesteps > 0: - timesteps = np.arange(1, timesteps + 1) + timesteps = list(range(1, timesteps + 1)) elif not isinstance(timesteps, list): raise ValueError(f"invalid value for argument 'timesteps': {timesteps}") - precip_forecast = extrapolation.forecast( - precip, - velocity, - timesteps, - extrap_method, - extrap_kwargs, + dataset_forecast = extrapolation.forecast( + dataset, timesteps, extrap_method, extrap_kwargs ) + precip_var = dataset_forecast.attrs["precip_var"] + precip_forecast = dataset_forecast[precip_var].values # Ignore missing values nanmask = np.isnan(precip_forecast) @@ -104,7 +105,8 @@ def forecast( precip_forecast[i, ...] /= kernel_sum precip_forecast = np.clip(precip_forecast, 0, 1) precip_forecast[nanmask] = np.nan - return precip_forecast + dataset_forecast[precip_var].data[:] = precip_forecast + return dataset_forecast def _get_kernel(size): diff --git a/pysteps/nowcasts/linda.py b/pysteps/nowcasts/linda.py index 7d737eaa9..2fc5dfd70 100644 --- a/pysteps/nowcasts/linda.py +++ b/pysteps/nowcasts/linda.py @@ -40,6 +40,8 @@ import time import warnings +from pysteps.xarray_helpers import convert_output_to_xarray_dataset + try: import dask @@ -47,28 +49,19 @@ except ImportError: DASK_IMPORTED = False import numpy as np +import xarray as xr +from scipy import optimize as opt +from scipy import stats from scipy.integrate import nquad from scipy.interpolate import interp1d -from scipy import optimize as opt from scipy.signal import convolve -from scipy import stats from pysteps import extrapolation, feature, noise -from pysteps.decorators import deprecate_args from pysteps.nowcasts.utils import nowcast_main_loop -@deprecate_args( - { - "precip_fields": "precip", - "advection_field": "velocity", - "num_ens_members": "n_ens_members", - }, - "1.8.0", -) def forecast( - precip, - velocity, + dataset: xr.Dataset, timesteps, feature_method="blob", max_num_features=25, @@ -100,15 +93,13 @@ def forecast( Parameters ---------- - precip: array_like - Array of shape (ari_order + 2, m, n) containing the input rain rate - or reflectivity fields (in linear scale) ordered by timestamp from - oldest to newest. The time steps between the inputs are assumed to be - regular. - velocity: array_like - Array of shape (2, m, n) containing the x- and y-components of the - advection field. The velocities are assumed to represent one time step - between the inputs. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as either reflectivity values in the + ``reflectivity`` data variable (in linear scale) or rainrate in the ``precip_intensity`` + data variable. The time dimension of the dataset has to be size + ``ari_order + 2`` and the precipitation variable has to have this dimension. timesteps: int Number of time steps to forecast. feature_method: {'blob', 'domain' 'shitomasi'} @@ -202,16 +193,15 @@ def forecast( Returns ------- - out: numpy.ndarray - A four-dimensional array of shape (n_ens_members, timesteps, m, n) - containing a time series of forecast precipitation fields for each - ensemble member. If add_perturbations is False, the first dimension is - dropped. The time series starts from t0 + timestep, where timestep is - taken from the input fields. If measure_time is True, the return value - is a three-element tuple containing the nowcast array, the initialization - time of the nowcast generator and the time used in the main loop - (seconds). If return_output is set to False, a single None value is - returned instead. + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast + precipitation fields for each ensemble member. Otherwise, a None value + is returned. The time series starts from t0+timestep, where timestep is + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). Notes ----- @@ -224,6 +214,10 @@ def forecast( variable OMP_NUM_THREADS to 1. This avoids slowdown caused by too many simultaneous threads. """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip = dataset[precip_var].values + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) _check_inputs(precip, velocity, timesteps, ari_order) if feature_kwargs is None: @@ -363,14 +357,21 @@ def forecast( callback, ) - if return_output: - if measure_time: - return precip_forecast[0], init_time, precip_forecast[1] - else: - return precip_forecast - else: + if not return_output: return None + if measure_time: + precip_forecast, mainloop_time = precip_forecast + + output_dataset = convert_output_to_xarray_dataset( + dataset, timesteps, precip_forecast + ) + + if measure_time: + return output_dataset, init_time, mainloop_time + else: + return output_dataset + def _check_inputs(precip, velocity, timesteps, ari_order): if ari_order not in [1, 2]: diff --git a/pysteps/nowcasts/sprog.py b/pysteps/nowcasts/sprog.py index 86c840dcb..2ebcfde41 100644 --- a/pysteps/nowcasts/sprog.py +++ b/pysteps/nowcasts/sprog.py @@ -10,17 +10,17 @@ forecast """ -import numpy as np import time -from pysteps import cascade -from pysteps import extrapolation -from pysteps import utils -from pysteps.decorators import deprecate_args +import numpy as np +import xarray as xr + +from pysteps import cascade, extrapolation, utils from pysteps.nowcasts import utils as nowcast_utils +from pysteps.nowcasts.utils import compute_percentile_mask, nowcast_main_loop from pysteps.postprocessing import probmatching from pysteps.timeseries import autoregression, correlation -from pysteps.nowcasts.utils import compute_percentile_mask, nowcast_main_loop +from pysteps.xarray_helpers import convert_output_to_xarray_dataset try: import dask @@ -30,10 +30,8 @@ DASK_IMPORTED = False -@deprecate_args({"R": "precip", "V": "velocity", "R_thr": "precip_thr"}, "1.8.0") def forecast( - precip, - velocity, + dataset: xr.Dataset, timesteps, precip_thr=None, n_cascade_levels=6, @@ -55,15 +53,13 @@ def forecast( Parameters ---------- - precip: array-like - Array of shape (ar_order+1,m,n) containing the input precipitation fields - ordered by timestamp from oldest to newest. The time steps between - the inputs are assumed to be regular. - velocity: array-like - Array of shape (2,m,n) containing the x- and y-components of the - advection field. - The velocities are assumed to represent one time step between the - inputs. All values are required to be finite. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as any precipitation data variable. + The time dimension of the dataset has to be size + ``ar_order + 1`` and the precipitation variable has to have this dimension. All + velocity values are required to be finite. timesteps: int or list of floats Number of time steps to forecast or a list of time steps for which the forecasts are computed (relative to the input time step). The elements @@ -120,13 +116,15 @@ def forecast( Returns ------- - out: ndarray - A three-dimensional array of shape (num_timesteps,m,n) containing a time - series of forecast precipitation fields. The time series starts from - t0+timestep, where timestep is taken from the input precipitation fields - precip. If measure_time is True, the return value is a three-element - tuple containing the nowcast array, the initialization time of the - nowcast generator and the time used in the main loop (seconds). + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast + precipitation fields. Otherwise, a None value + is returned. The time series starts from t0+timestep, where timestep is + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). See also -------- @@ -137,6 +135,10 @@ def forecast( :cite:`Seed2003`, :cite:`PCH2019a` """ + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip = dataset[precip_var].values + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) _check_inputs(precip, velocity, timesteps, ar_order) if extrap_kwargs is None: @@ -327,8 +329,6 @@ def f(precip, i): print("Starting nowcast computation.") - precip_forecast = [] - state = {"precip_cascades": precip_cascades, "precip_decomp": precip_decomp} params = { "domain": domain, @@ -358,12 +358,14 @@ def f(precip, i): if measure_time: precip_forecast, mainloop_time = precip_forecast - precip_forecast = np.stack(precip_forecast) + output_dataset = convert_output_to_xarray_dataset( + dataset, timesteps, precip_forecast + ) if measure_time: - return precip_forecast, init_time, mainloop_time + return output_dataset, init_time, mainloop_time else: - return precip_forecast + return output_dataset def _check_inputs(precip, velocity, timesteps, ar_order): diff --git a/pysteps/nowcasts/sseps.py b/pysteps/nowcasts/sseps.py index a8848d3e3..1d083c04b 100644 --- a/pysteps/nowcasts/sseps.py +++ b/pysteps/nowcasts/sseps.py @@ -18,18 +18,17 @@ forecast """ -import numpy as np import time -from scipy.ndimage import generate_binary_structure, iterate_structure +import numpy as np +import xarray as xr +from scipy.ndimage import generate_binary_structure, iterate_structure -from pysteps import cascade -from pysteps import extrapolation -from pysteps import noise -from pysteps.decorators import deprecate_args +from pysteps import cascade, extrapolation, noise from pysteps.nowcasts import utils as nowcast_utils from pysteps.postprocessing import probmatching from pysteps.timeseries import autoregression, correlation +from pysteps.xarray_helpers import convert_output_to_xarray_dataset try: import dask @@ -39,11 +38,8 @@ dask_imported = False -@deprecate_args({"R": "precip", "V": "velocity"}, "1.8.0") def forecast( - precip, - metadata, - velocity, + dataset: xr.Dataset, timesteps, n_ens_members=24, n_cascade_levels=6, @@ -78,18 +74,14 @@ def forecast( Parameters ---------- - precip: array-like - Array of shape (ar_order+1,m,n) containing the input precipitation fields - ordered by timestamp from oldest to newest. The time steps between the inputs - are assumed to be regular, and the inputs are required to have finite values. - metadata: dict - Metadata dictionary containing the accutime, xpixelsize, threshold and - zerovalue attributes as described in the documentation of - :py:mod:`pysteps.io.importers`. xpixelsize is assumed to be in meters. - velocity: array-like - Array of shape (2,m,n) containing the x- and y-components of the advection - field. The velocities are assumed to represent one time step between the - inputs. All values are required to be finite. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as any precipitation data variable. + The units and stepsize of ``y`` and ``x`` have to be the same and the only supported + units are meters and kilometers. The time dimension of the dataset has to be size + ``ar_order + 1`` and the precipitation variable has to have this dimension. All + velocity values are required to be finite. win_size: int or two-element sequence of ints Size-length of the localization window. overlap: float [0,1[ @@ -181,12 +173,15 @@ def forecast( Returns ------- - out: ndarray - If return_output is True, a four-dimensional array of shape - (n_ens_members,num_timesteps,m,n) containing a time series of forecast + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast precipitation fields for each ensemble member. Otherwise, a None value is returned. The time series starts from t0+timestep, where timestep is - taken from the input precipitation fields. + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the + initialization time of the nowcast generator and the time used in the + main loop (seconds). See also -------- @@ -201,7 +196,20 @@ def forecast( ---------- :cite:`Seed2003`, :cite:`BPS2006`, :cite:`SPN2013`, :cite:`NBSG2017` """ - + timesteps_in = timesteps + x_units = dataset["x"].attrs["units"] + y_units = dataset["y"].attrs["units"] + x_stepsize = dataset["x"].attrs["stepsize"] + y_stepsize = dataset["y"].attrs["stepsize"] + if x_units != y_units or x_stepsize != y_stepsize: + raise ValueError("units and stepsize needs to be the same for x and y") + if x_units not in ["m", "km"]: + raise ValueError("only m and km supported as x and y units") + + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip = dataset[precip_var].values + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) _check_inputs(precip, velocity, timesteps, ar_order) if extrap_kwargs is None: @@ -237,8 +245,10 @@ def forecast( else: win_size = tuple([int(win_size[i]) for i in range(2)]) - timestep = metadata["accutime"] - kmperpixel = metadata["xpixelsize"] / 1000 + timestep = dataset["time"].attrs["stepsize"] / 60 + kmperpixel = x_stepsize + if x_units == "m": + kmperpixel = kmperpixel / 1000 print("Computing SSEPS nowcast") print("-----------------------") @@ -292,8 +302,8 @@ def forecast( f"velocity perturbations, perpendicular: {vp_perp[0]},{vp_perp[1]},{vp_perp[2]}" ) - precip_thr = metadata["threshold"] - precip_min = metadata["zerovalue"] + precip_thr = dataset[precip_var].attrs["threshold"] + precip_min = dataset[precip_var].attrs["zerovalue"] num_ensemble_workers = n_ens_members if num_workers > n_ens_members else num_workers @@ -911,10 +921,12 @@ def worker(j): if return_output: outarr = np.stack([np.stack(precip_forecast[j]) for j in range(n_ens_members)]) + output_dataset = convert_output_to_xarray_dataset(dataset, timesteps_in, outarr) + if measure_time: - return outarr, init_time, mainloop_time + return output_dataset, init_time, mainloop_time else: - return outarr + return output_dataset else: return None diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index 50c6a5d22..366f32f6d 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -11,19 +11,18 @@ forecast """ +import time + import numpy as np +import xarray as xr from scipy.ndimage import generate_binary_structure, iterate_structure -import time -from pysteps import cascade -from pysteps import extrapolation -from pysteps import noise -from pysteps import utils -from pysteps.decorators import deprecate_args +from pysteps import cascade, extrapolation, noise, utils from pysteps.nowcasts import utils as nowcast_utils +from pysteps.nowcasts.utils import compute_percentile_mask, nowcast_main_loop from pysteps.postprocessing import probmatching from pysteps.timeseries import autoregression, correlation -from pysteps.nowcasts.utils import compute_percentile_mask, nowcast_main_loop +from pysteps.xarray_helpers import convert_output_to_xarray_dataset try: import dask @@ -33,10 +32,8 @@ DASK_IMPORTED = False -@deprecate_args({"R": "precip", "V": "velocity", "R_thr": "precip_thr"}, "1.8.0") def forecast( - precip, - velocity, + dataset: xr.Dataset, timesteps, n_ens_members=24, n_cascade_levels=6, @@ -72,14 +69,13 @@ def forecast( Parameters ---------- - precip: array-like - Array of shape (ar_order+1,m,n) containing the input precipitation fields - ordered by timestamp from oldest to newest. The time steps between the - inputs are assumed to be regular. - velocity: array-like - Array of shape (2,m,n) containing the x- and y-components of the advection - field. The velocities are assumed to represent one time step between the - inputs. All values are required to be finite. + dataset: xarray.Dataset + Input dataset as described in the documentation of + :py:mod:`pysteps.io.importers`. It has to contain the ``velocity_x`` and + ``velocity_y`` data variables, as well as any precipitation data variable. + The time dimension of the dataset has to be size + ``ar_order + 1`` and the precipitation variable has to have this dimension. All + velocity values are required to be finite. timesteps: int or list of floats Number of time steps to forecast or a list of time steps for which the forecasts are computed (relative to the input time step). The elements @@ -241,13 +237,13 @@ def forecast( Returns ------- - out: ndarray - If return_output is True, a four-dimensional array of shape - (n_ens_members,num_timesteps,m,n) containing a time series of forecast + out: xarray.Dataset + If return_output is True, a dataset as described in the documentation of + :py:mod:`pysteps.io.importers` is returned containing a time series of forecast precipitation fields for each ensemble member. Otherwise, a None value is returned. The time series starts from t0+timestep, where timestep is - taken from the input precipitation fields. If measure_time is True, the - return value is a three-element tuple containing the nowcast array, the + taken from the metadata of the time coordinate. If measure_time is True, the + return value is a three-element tuple containing the nowcast dataset, the initialization time of the nowcast generator and the time used in the main loop (seconds). @@ -261,6 +257,11 @@ def forecast( :cite:`Seed2003`, :cite:`BPS2006`, :cite:`SPN2013`, :cite:`PCH2019b` """ + timesteps_in = timesteps + dataset = dataset.copy(deep=True) + precip_var = dataset.attrs["precip_var"] + precip = dataset[precip_var].values + velocity = np.stack([dataset["velocity_x"], dataset["velocity_y"]]) _check_inputs(precip, velocity, timesteps, ar_order) if extrap_kwargs is None: @@ -574,8 +575,6 @@ def f(precip, i): else: velocity_perturbators = None - precip_forecast = [[] for _ in range(n_ens_members)] - if probmatching_method == "mean": mu_0 = np.mean(precip[-1, :, :][precip[-1, :, :] >= precip_thr]) else: @@ -680,13 +679,13 @@ def f(precip, i): precip_forecast, mainloop_time = precip_forecast if return_output: - precip_forecast = np.stack( - [np.stack(precip_forecast[j]) for j in range(n_ens_members)] + output_dataset = convert_output_to_xarray_dataset( + dataset, timesteps_in, precip_forecast ) if measure_time: - return precip_forecast, init_time, mainloop_time + return output_dataset, init_time, mainloop_time else: - return precip_forecast + return output_dataset else: return None diff --git a/pysteps/nowcasts/utils.py b/pysteps/nowcasts/utils.py index fd111e28d..fed1c2f96 100644 --- a/pysteps/nowcasts/utils.py +++ b/pysteps/nowcasts/utils.py @@ -17,6 +17,7 @@ """ import time + import numpy as np from scipy.ndimage import binary_dilation, generate_binary_structure @@ -412,10 +413,10 @@ def worker2(i): if not ensemble: precip_forecast_out = precip_forecast_out[0, :] - if measure_time: - return precip_forecast_out, time.time() - starttime_total - else: - return precip_forecast_out + if measure_time: + return precip_forecast_out, time.time() - starttime_total + else: + return precip_forecast_out def print_ar_params(phi): diff --git a/pysteps/tests/test_motion_lk.py b/pysteps/tests/test_motion_lk.py index 871dcd98b..a8f640533 100644 --- a/pysteps/tests/test_motion_lk.py +++ b/pysteps/tests/test_motion_lk.py @@ -3,8 +3,8 @@ """ """ -import pytest import numpy as np +import pytest from pysteps import motion, verification from pysteps.tests.helpers import get_precipitation_fields @@ -61,19 +61,19 @@ def test_lk( pytest.importorskip("pandas") # inputs - precip, metadata = get_precipitation_fields( + dataset = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=True, upscale=2000, ) - precip = precip.filled() + precip_var = dataset.attrs["precip_var"] # Retrieve motion field oflow_method = motion.get_method("LK") - output = oflow_method( - precip, + output_dataset = oflow_method( + dataset, lk_kwargs=lk_kwargs, fd_method=fd_method, dense=dense, @@ -86,13 +86,17 @@ def test_lk( # Check format of ouput if dense: + output = np.stack( + [output_dataset["velocity_x"].values, output_dataset["velocity_y"].values] + ) assert isinstance(output, np.ndarray) assert output.ndim == 3 assert output.shape[0] == 2 - assert output.shape[1:] == precip[0].shape + assert output.shape[1:] == dataset[precip_var].values[0].shape if nr_std_outlier == 0: assert output.sum() == 0 else: + output = output_dataset assert isinstance(output, tuple) assert len(output) == 2 assert isinstance(output[0], np.ndarray) diff --git a/pysteps/tests/test_nowcasts_anvil.py b/pysteps/tests/test_nowcasts_anvil.py index 14a130fb1..35d84e0f3 100644 --- a/pysteps/tests/test_nowcasts_anvil.py +++ b/pysteps/tests/test_nowcasts_anvil.py @@ -31,31 +31,28 @@ def test_anvil_rainrate( ): """Tests ANVIL nowcast using rain rate precipitation fields.""" # inputs - precip_input = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=4, num_next_files=0, return_raw=False, metadata=False, upscale=2000, ) - precip_input = precip_input.filled() - precip_obs = get_precipitation_fields( + dataset_obs = get_precipitation_fields( num_prev_files=0, num_next_files=3, return_raw=False, upscale=2000 - )[1:, :, :] - precip_obs = precip_obs.filled() + ).isel(time=slice(1, None, None)) + precip_var = dataset_input.attrs["precip_var"] pytest.importorskip("cv2") oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_input) + dataset_w_motion = oflow_method(dataset_input) nowcast_method = nowcasts.get_method("anvil") output = nowcast_method( - precip_input[-(ar_order + 2) :], - retrieved_motion, + dataset_w_motion.isel(time=slice(-(ar_order + 2), None, None)), timesteps=timesteps, - rainrate=None, # no R(VIL) conversion is done n_cascade_levels=n_cascade_levels, ar_order=ar_order, ar_window_radius=ar_window_radius, @@ -63,9 +60,10 @@ def test_anvil_rainrate( measure_time=measure_time, ) if measure_time: - precip_forecast, __, __ = output + dataset_forecast, __, __ = output else: - precip_forecast = output + dataset_forecast = output + precip_forecast = dataset_forecast[precip_var].values assert precip_forecast.ndim == 3 assert precip_forecast.shape[0] == ( @@ -73,7 +71,7 @@ def test_anvil_rainrate( ) result = verification.det_cat_fct( - precip_forecast[-1], precip_obs[-1], thr=0.1, scores="CSI" + precip_forecast[-1], dataset_obs[precip_var].values[-1], thr=0.1, scores="CSI" )["CSI"] assert result > min_csi, f"CSI={result:.2f}, required > {min_csi:.2f}" diff --git a/pysteps/tests/test_nowcasts_lagrangian_probability.py b/pysteps/tests/test_nowcasts_lagrangian_probability.py index 1ec352b0b..d75b29e87 100644 --- a/pysteps/tests/test_nowcasts_lagrangian_probability.py +++ b/pysteps/tests/test_nowcasts_lagrangian_probability.py @@ -1,10 +1,13 @@ # -*- coding: utf-8 -*- +from datetime import datetime, timezone + import numpy as np import pytest +import xarray as xr +from pysteps.motion.lucaskanade import dense_lucaskanade from pysteps.nowcasts.lagrangian_probability import forecast from pysteps.tests.helpers import get_precipitation_fields -from pysteps.motion.lucaskanade import dense_lucaskanade def test_numerical_example(): @@ -12,12 +15,23 @@ def test_numerical_example(): precip = np.zeros((20, 20)) precip[5:10, 5:10] = 1 velocity = np.zeros((2, *precip.shape)) + now = datetime.now(tz=timezone.utc).replace(tzinfo=None) + dataset_input = xr.Dataset( + data_vars={ + "precip_intensity": (["time", "y", "x"], [precip]), + "velocity_x": (["y", "x"], velocity[0]), + "velocity_y": (["y", "x"], velocity[1]), + }, + coords={"time": (["time"], [now], {"stepsize": 300})}, + attrs={"precip_var": "precip_intensity"}, + ) timesteps = 4 thr = 0.5 slope = 1 # pixels / timestep # compute probability forecast - fct = forecast(precip, velocity, timesteps, thr, slope=slope) + dataset_forecast = forecast(dataset_input, timesteps, thr, slope=slope) + fct = dataset_forecast["precip_intensity"].values assert fct.ndim == 3 assert fct.shape[0] == timesteps @@ -26,7 +40,8 @@ def test_numerical_example(): assert fct.min() >= 0.0 # slope = 0 should return a binary field - fct = forecast(precip, velocity, timesteps, thr, slope=0) + dataset_forecast = forecast(dataset_input, timesteps, thr, slope=0) + fct = dataset_forecast["precip_intensity"].values ref = (np.repeat(precip[None, ...], timesteps, axis=0) >= thr).astype(float) assert np.allclose(fct, fct.astype(bool)) assert np.allclose(fct, ref) @@ -37,12 +52,23 @@ def test_numerical_example_with_float_slope_and_float_list_timesteps(): precip = np.zeros((20, 20)) precip[5:10, 5:10] = 1 velocity = np.zeros((2, *precip.shape)) + now = datetime.now(tz=timezone.utc).replace(tzinfo=None) + dataset_input = xr.Dataset( + data_vars={ + "precip_intensity": (["time", "y", "x"], [precip]), + "velocity_x": (["y", "x"], velocity[0]), + "velocity_y": (["y", "x"], velocity[1]), + }, + coords={"time": (["time"], [now], {"stepsize": 300})}, + attrs={"precip_var": "precip_intensity"}, + ) timesteps = [1.0, 2.0, 5.0, 12.0] thr = 0.5 slope = 1.0 # pixels / timestep # compute probability forecast - fct = forecast(precip, velocity, timesteps, thr, slope=slope) + dataset_forecast = forecast(dataset_input, timesteps, thr, slope=slope) + fct = dataset_forecast["precip_intensity"].values assert fct.ndim == 3 assert fct.shape[0] == len(timesteps) @@ -56,16 +82,18 @@ def test_real_case(): pytest.importorskip("cv2") # inputs - precip, metadata = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=True, upscale=2000, ) + precip_var = dataset_input.attrs["precip_var"] + metadata = dataset_input[precip_var].attrs # motion - motion = dense_lucaskanade(precip) + dataset_w_motion = dense_lucaskanade(dataset_input) # parameters timesteps = [1, 2, 3] @@ -74,13 +102,18 @@ def test_real_case(): # compute probability forecast extrap_kwargs = dict(allow_nonfinite_values=True) - fct = forecast( - precip[-1], motion, timesteps, thr, slope=slope, extrap_kwargs=extrap_kwargs + dataset_forecast = forecast( + dataset_w_motion.isel(time=slice(-1, None, None)), + timesteps, + thr, + slope=slope, + extrap_kwargs=extrap_kwargs, ) + fct = dataset_forecast["precip_intensity"].values assert fct.ndim == 3 assert fct.shape[0] == len(timesteps) - assert fct.shape[1:] == precip.shape[1:] + assert fct.shape[1:] == dataset_input[precip_var].values.shape[1:] assert np.nanmax(fct) <= 1.0 assert np.nanmin(fct) >= 0.0 @@ -89,11 +122,19 @@ def test_wrong_inputs(): # dummy inputs precip = np.zeros((3, 3)) velocity = np.zeros((2, *precip.shape)) + dataset_input = xr.Dataset( + data_vars={ + "precip_intensity": (["y", "x"], precip), + "velocity_x": (["y", "x"], velocity[0]), + "velocity_y": (["y", "x"], velocity[1]), + }, + attrs={"precip_var": "precip_intensity"}, + ) # timesteps must be > 0 with pytest.raises(ValueError): - forecast(precip, velocity, 0, 1) + forecast(dataset_input, 0, 1) # timesteps must be a sorted list with pytest.raises(ValueError): - forecast(precip, velocity, [2, 1], 1) + forecast(dataset_input, [2, 1], 1) diff --git a/pysteps/tests/test_nowcasts_linda.py b/pysteps/tests/test_nowcasts_linda.py index 2d5f03b71..a5b60611b 100644 --- a/pysteps/tests/test_nowcasts_linda.py +++ b/pysteps/tests/test_nowcasts_linda.py @@ -1,13 +1,13 @@ -from datetime import timedelta import os +from datetime import timedelta + import numpy as np import pytest +import xarray as xr from pysteps import io, motion, nowcasts, verification -from pysteps.nowcasts.linda import forecast from pysteps.tests.helpers import get_precipitation_fields - linda_arg_names = ( "add_perturbations", "kernel_type", @@ -42,7 +42,7 @@ def test_linda( pytest.importorskip("skimage") # inputs - precip_input, metadata = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=2, num_next_files=0, metadata=True, @@ -51,20 +51,23 @@ def test_linda( log_transform=False, ) - precip_obs = get_precipitation_fields( + dataset_obs = get_precipitation_fields( num_prev_files=0, num_next_files=3, clip=(354000, 866000, -96000, 416000), upscale=4000, log_transform=False, - )[1:, :, :] + ).isel(time=slice(1, None, None)) + precip_var = dataset_input.attrs["precip_var"] + metadata = dataset_input[precip_var].attrs oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_input) + dataset_w_motion = oflow_method(dataset_input) - precip_forecast = forecast( - precip_input, - retrieved_motion, + nowcast_method = nowcasts.get_method("linda") + + dataset_forecast = nowcast_method( + dataset_w_motion, 3, kernel_type=kernel_type, vel_pert_method=vel_pert_method, @@ -78,68 +81,82 @@ def test_linda( seed=42, ) if measure_time: - assert len(precip_forecast) == 3 - assert isinstance(precip_forecast[1], float) - precip_forecast = precip_forecast[0] + assert len(dataset_forecast) == 3 + assert isinstance(dataset_forecast[1], float) + dataset_forecast = dataset_forecast[0] + + precip_forecast = dataset_forecast[precip_var].values if not add_perturbations: assert precip_forecast.ndim == 3 assert precip_forecast.shape[0] == 3 - assert precip_forecast.shape[1:] == precip_input.shape[1:] + assert precip_forecast.shape[1:] == dataset_input[precip_var].values.shape[1:] csi = verification.det_cat_fct( - precip_forecast[-1], precip_obs[-1], thr=1.0, scores="CSI" + precip_forecast[-1], + dataset_obs[precip_var].values[-1], + thr=1.0, + scores="CSI", )["CSI"] assert csi > min_csi, f"CSI={csi:.1f}, required > {min_csi:.1f}" else: assert precip_forecast.ndim == 4 assert precip_forecast.shape[0] == 5 assert precip_forecast.shape[1] == 3 - assert precip_forecast.shape[2:] == precip_input.shape[1:] + assert precip_forecast.shape[2:] == dataset_input[precip_var].values.shape[1:] - crps = verification.probscores.CRPS(precip_forecast[:, -1], precip_obs[-1]) + crps = verification.probscores.CRPS( + precip_forecast[:, -1], dataset_obs[precip_var].values[-1] + ) assert crps < max_crps, f"CRPS={crps:.2f}, required < {max_crps:.2f}" def test_linda_wrong_inputs(): # dummy inputs - precip = np.zeros((3, 3, 3)) - velocity = np.zeros((2, 3, 3)) + dataset_input = xr.Dataset( + data_vars={ + "precip_intensity": (["time", "y", "x"], np.zeros((3, 3, 3))), + "velocity_x": (["y", "x"], np.zeros((3, 3))), + "velocity_y": (["y", "x"], np.zeros((3, 3))), + }, + attrs={"precip_var": "precip_intensity"}, + ) + dataset_input_4d = xr.Dataset( + data_vars={ + "precip_intensity": ( + ["ens_number", "time", "y", "x"], + np.zeros((3, 3, 3, 3)), + ), + "velocity_x": (["y", "x"], np.zeros((3, 3))), + "velocity_y": (["y", "x"], np.zeros((3, 3))), + }, + attrs={"precip_var": "precip_intensity"}, + ) + + nowcast_method = nowcasts.get_method("linda") # vel_pert_method is set but kmperpixel is None with pytest.raises(ValueError): - forecast(precip, velocity, 1, vel_pert_method="bps", kmperpixel=None) + nowcast_method(dataset_input, 1, vel_pert_method="bps", kmperpixel=None) # vel_pert_method is set but timestep is None with pytest.raises(ValueError): - forecast( - precip, velocity, 1, vel_pert_method="bps", kmperpixel=1, timestep=None + nowcast_method( + dataset_input, 1, vel_pert_method="bps", kmperpixel=1, timestep=None ) # fractional time steps not yet implemented # timesteps is not an integer with pytest.raises(ValueError): - forecast(precip, velocity, [1.0, 2.0]) + nowcast_method(dataset_input, [1.0, 2.0]) # ari_order 1 or 2 required with pytest.raises(ValueError): - forecast(precip, velocity, 1, ari_order=3) + nowcast_method(dataset_input, 1, ari_order=3) # precip_fields must be a three-dimensional array with pytest.raises(ValueError): - forecast(np.zeros((3, 3, 3, 3)), velocity, 1) - - # precip_fields.shape[0] < ari_order+2 - with pytest.raises(ValueError): - forecast(np.zeros((2, 3, 3)), velocity, 1, ari_order=1) - - # advection_field must be a three-dimensional array - with pytest.raises(ValueError): - forecast(precip, velocity[0], 1) - - # dimension mismatch between precip_fields and advection_field - with pytest.raises(ValueError): - forecast(np.zeros((3, 2, 3)), velocity, 1) + nowcast_method(dataset_input_4d, 1) def test_linda_callback(tmp_path): diff --git a/pysteps/tests/test_nowcasts_sprog.py b/pysteps/tests/test_nowcasts_sprog.py index 1077c3edd..f64900cd8 100644 --- a/pysteps/tests/test_nowcasts_sprog.py +++ b/pysteps/tests/test_nowcasts_sprog.py @@ -30,29 +30,28 @@ def test_sprog( ): """Tests SPROG nowcast.""" # inputs - precip_input, metadata = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=True, upscale=2000, ) - precip_input = precip_input.filled() - precip_obs = get_precipitation_fields( + dataset_obs = get_precipitation_fields( num_prev_files=0, num_next_files=3, return_raw=False, upscale=2000 - )[1:, :, :] - precip_obs = precip_obs.filled() + ).isel(time=slice(1, None, None)) + precip_var = dataset_input.attrs["precip_var"] + metadata = dataset_input[precip_var].attrs pytest.importorskip("cv2") oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_input) + dataset_w_motion = oflow_method(dataset_input) nowcast_method = nowcasts.get_method("sprog") - precip_forecast = nowcast_method( - precip_input, - retrieved_motion, + dataset_forecast = nowcast_method( + dataset_w_motion, timesteps=timesteps, precip_thr=metadata["threshold"], n_cascade_levels=n_cascade_levels, @@ -60,6 +59,7 @@ def test_sprog( probmatching_method=probmatching_method, domain=domain, ) + precip_forecast = dataset_forecast[precip_var].values assert precip_forecast.ndim == 3 assert precip_forecast.shape[0] == ( @@ -67,7 +67,7 @@ def test_sprog( ) result = verification.det_cat_fct( - precip_forecast[-1], precip_obs[-1], thr=0.1, scores="CSI" + precip_forecast[-1], dataset_obs[precip_var].values[-1], thr=0.1, scores="CSI" )["CSI"] assert result > min_csi, f"CSI={result:.1f}, required > {min_csi:.1f}" diff --git a/pysteps/tests/test_nowcasts_sseps.py b/pysteps/tests/test_nowcasts_sseps.py index 4d89fd33a..6d3a3c9c0 100644 --- a/pysteps/tests/test_nowcasts_sseps.py +++ b/pysteps/tests/test_nowcasts_sseps.py @@ -17,8 +17,8 @@ ) sseps_arg_values = [ - (5, 6, 2, "incremental", "cdf", 200, 3, 0.60), - (5, 6, 2, "incremental", "cdf", 200, [3], 0.60), + (5, 6, 2, "incremental", "cdf", 200, 3, 0.62), + (5, 6, 2, "incremental", "cdf", 200, [3], 0.62), ] @@ -35,32 +35,29 @@ def test_sseps( ): """Tests SSEPS nowcast.""" # inputs - precip_input, metadata = get_precipitation_fields( + dataset_input = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=True, upscale=2000, ) - precip_input = precip_input.filled() + precip_var = dataset_input.attrs["precip_var"] - precip_obs = get_precipitation_fields( + dataset_obs = get_precipitation_fields( num_prev_files=0, num_next_files=3, return_raw=False, upscale=2000 - )[1:, :, :] - precip_obs = precip_obs.filled() + ).isel(time=slice(1, None, None)) pytest.importorskip("cv2") oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_input) + dataset_w_motion = oflow_method(dataset_input) nowcast_method = nowcasts.get_method("sseps") - precip_forecast = nowcast_method( - precip_input, - metadata, - retrieved_motion, + dataset_forecast = nowcast_method( + dataset_w_motion, + timesteps, win_size=win_size, - timesteps=timesteps, n_ens_members=n_ens_members, n_cascade_levels=n_cascade_levels, ar_order=ar_order, @@ -68,6 +65,7 @@ def test_sseps( mask_method=mask_method, probmatching_method=probmatching_method, ) + precip_forecast = dataset_forecast[precip_var].values assert precip_forecast.ndim == 4 assert precip_forecast.shape[0] == n_ens_members @@ -75,7 +73,9 @@ def test_sseps( timesteps if isinstance(timesteps, int) else len(timesteps) ) - crps = verification.probscores.CRPS(precip_forecast[:, -1], precip_obs[-1]) + crps = verification.probscores.CRPS( + precip_forecast[:, -1], dataset_obs[precip_var].values[-1] + ) assert crps < max_crps, f"CRPS={crps:.2f}, required < {max_crps:.2f}" diff --git a/pysteps/tests/test_nowcasts_steps.py b/pysteps/tests/test_nowcasts_steps.py index adb6ea917..16d2e4de5 100644 --- a/pysteps/tests/test_nowcasts_steps.py +++ b/pysteps/tests/test_nowcasts_steps.py @@ -56,17 +56,15 @@ def test_steps_skill( ).isel(time=slice(1, None, None)) precip_var = dataset_input.attrs["precip_var"] metadata = dataset_input[precip_var].attrs - precip_data = dataset_input[precip_var].values pytest.importorskip("cv2") oflow_method = motion.get_method("LK") - retrieved_motion = oflow_method(precip_data) + dataset_w_motion = oflow_method(dataset_input) nowcast_method = nowcasts.get_method("steps") - precip_forecast = nowcast_method( - precip_data, - retrieved_motion, + dataset_forecast = nowcast_method( + dataset_w_motion, timesteps=timesteps, precip_thr=metadata["threshold"], kmperpixel=2.0, @@ -79,6 +77,7 @@ def test_steps_skill( probmatching_method=probmatching_method, domain=domain, ) + precip_forecast = dataset_forecast[precip_var].values assert precip_forecast.ndim == 4 assert precip_forecast.shape[0] == n_ens_members diff --git a/pysteps/tests/test_nowcasts_utils.py b/pysteps/tests/test_nowcasts_utils.py index 075225427..1dfeb27a9 100644 --- a/pysteps/tests/test_nowcasts_utils.py +++ b/pysteps/tests/test_nowcasts_utils.py @@ -26,17 +26,18 @@ def test_nowcast_main_loop( timesteps, ensemble, num_ensemble_members, velocity_perturbations ): """Test the nowcast_main_loop function.""" - precip = get_precipitation_fields( + dataset = get_precipitation_fields( num_prev_files=2, num_next_files=0, return_raw=False, metadata=False, upscale=2000, ) - precip = precip.filled() oflow_method = motion.get_method("LK") - velocity = oflow_method(precip) + dataset = oflow_method(dataset) + precip = dataset["precip_intensity"].values + velocity = np.stack([dataset["velocity_x"].values, dataset["velocity_y"].values]) precip = precip[-1] diff --git a/pysteps/tests/test_utils_dimension.py b/pysteps/tests/test_utils_dimension.py index 2bbb63f58..038b725a0 100644 --- a/pysteps/tests/test_utils_dimension.py +++ b/pysteps/tests/test_utils_dimension.py @@ -5,19 +5,17 @@ import numpy as np import pytest import xarray as xr -from numpy.testing import assert_array_equal +from numpy.testing import assert_array_almost_equal, assert_array_equal from pytest import raises -from pysteps.converters import convert_to_xarray_dataset from pysteps.utils import dimension +from pysteps.xarray_helpers import convert_input_to_xarray_dataset fillvalues_metadata = { "x1": 0, "x2": 4, "y1": 0, "y2": 4, - "xpixelsize": 1, - "ypixelsize": 1, "zerovalue": 0, "yorigin": "lower", "unit": "mm/h", @@ -32,7 +30,6 @@ } test_data_not_trim = ( - # "data, window_size, axis, method, expected" ( np.arange(12).reshape(2, 6), 2, @@ -94,7 +91,7 @@ def test_aggregate_fields(data, window_size, dim, method, expected): windows size does not divide the data dimensions. The length of each dimension should be larger than 2. """ - dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) + dataset = convert_input_to_xarray_dataset(data, None, fillvalues_metadata) actual = dimension.aggregate_fields(dataset, window_size, dim=dim, method=method) assert_array_equal(actual["precip_intensity"].values, expected) @@ -104,7 +101,7 @@ def test_aggregate_fields(data, window_size, dim, method, expected): data = np.pad(data, ((0, 0), (0, 1))) else: data = np.pad(data, (0, 1)) - dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) + dataset = convert_input_to_xarray_dataset(data, None, fillvalues_metadata) actual = dimension.aggregate_fields( dataset, window_size, dim=dim, method=method, trim=True @@ -115,13 +112,85 @@ def test_aggregate_fields(data, window_size, dim, method, expected): dimension.aggregate_fields(dataset, window_size, dim=dim, method=method) +test_data_agg_w_velocity = ( + ( + np.arange(12).reshape(2, 6), + np.arange(12).reshape(2, 6), + np.arange(12).reshape(2, 6), + np.arange(0, 1.2, 0.1).reshape(2, 6), + 2, + "x", + "mean", + "mean", + np.array([[0.5, 2.5, 4.5], [6.5, 8.5, 10.5]]), + np.array([[0.5, 2.5, 4.5], [6.5, 8.5, 10.5]]), + np.array([[0, 0.2, 0.4], [0.6, 0.8, 1]]), + ), + ( + np.arange(4 * 6).reshape(4, 6), + np.arange(4 * 6).reshape(4, 6), + np.arange(4 * 6).reshape(4, 6), + np.arange(0, 1.2, 0.05).reshape(4, 6), + (2, 3), + ("y", "x"), + "mean", + "sum", + np.array([[4, 7], [16, 19]]), + np.array([[24, 42], [96, 114]]), + np.array([[0, 0.15], [0.6, 0.75]]), + ), +) + + +@pytest.mark.parametrize( + "data, data_vx, data_vy, data_qual, window_size, dim, method, velocity_method, expected, expected_v, expected_qual", + test_data_agg_w_velocity, +) +def test_aggregate_fields_w_velocity( + data, + data_vx, + data_vy, + data_qual, + window_size, + dim, + method, + velocity_method, + expected, + expected_v, + expected_qual, +): + """ + Test the aggregate_fields function for dataset with velocity information. + The windows size must divide exactly the data dimensions. + Internally, additional test are generated for situations where the + windows size does not divide the data dimensions. + The length of each dimension should be larger than 2. + """ + dataset = convert_input_to_xarray_dataset(data, None, fillvalues_metadata) + dataset = dataset.assign( + { + "velocity_x": (("y", "x"), data_vx), + "velocity_y": (("y", "x"), data_vy), + "quality": (("y", "x"), data_qual), + } + ) + + actual = dimension.aggregate_fields( + dataset, window_size, dim=dim, method=method, velocity_method=velocity_method + ) + assert_array_equal(actual["precip_intensity"].values, expected) + assert_array_equal(actual["velocity_x"].values, expected_v) + assert_array_equal(actual["velocity_y"].values, expected_v) + assert_array_almost_equal(actual["quality"].values, expected_qual) + + def test_aggregate_fields_errors(): """ Test that the errors are correctly captured in the aggregate_fields function. """ data = np.arange(4 * 6).reshape(4, 6) - dataset = convert_to_xarray_dataset(data, None, fillvalues_metadata) + dataset = convert_input_to_xarray_dataset(data, None, fillvalues_metadata) with raises(ValueError): dimension.aggregate_fields(dataset, -1, dim="y") @@ -160,7 +229,7 @@ def test_aggregate_fields_errors(): ) def test_aggregate_fields_time(data, metadata, time_window_min, ignore_nan, expected): """Test the aggregate_fields_time.""" - dataset_ref = convert_to_xarray_dataset( + dataset_ref = convert_input_to_xarray_dataset( data, None, {**fillvalues_metadata, **metadata} ) datasets = [] @@ -234,7 +303,9 @@ def test_aggregate_fields_time(data, metadata, time_window_min, ignore_nan, expe ) def test_aggregate_fields_space(data, metadata, space_window, ignore_nan, expected): """Test the aggregate_fields_space.""" - dataset = convert_to_xarray_dataset(data, None, {**fillvalues_metadata, **metadata}) + dataset = convert_input_to_xarray_dataset( + data, None, {**fillvalues_metadata, **metadata} + ) assert_array_equal( dimension.aggregate_fields_space(dataset, space_window, ignore_nan)[ "precip_intensity" if metadata["unit"] == "mm/h" else "precip_accum" @@ -271,7 +342,9 @@ def test_aggregate_fields_space(data, metadata, space_window, ignore_nan, expect @pytest.mark.parametrize("R, metadata, extent, expected", test_data_clip_domain) def test_clip_domain(R, metadata, extent, expected): """Test the clip_domain.""" - dataset = convert_to_xarray_dataset(R, None, {**fillvalues_metadata, **metadata}) + dataset = convert_input_to_xarray_dataset( + R, None, {**fillvalues_metadata, **metadata} + ) assert_array_equal( dimension.clip_domain(dataset, extent)["precip_intensity"].values, expected ) @@ -336,12 +409,67 @@ def test_clip_domain(R, metadata, extent, expected): @pytest.mark.parametrize("data, metadata, method, inverse, expected", test_data_square) def test_square_domain(data, metadata, method, inverse, expected): """Test the square_domain.""" - dataset = convert_to_xarray_dataset(data, None, {**fillvalues_metadata, **metadata}) - dataset["precip_intensity"].attrs = { - **dataset["precip_intensity"].attrs, - **metadata, - } + dataset = convert_input_to_xarray_dataset( + data, None, {**fillvalues_metadata, **metadata} + ) + if "square_method" in metadata: + dataset.attrs["square_method"] = metadata["square_method"] + if "orig_domain" in metadata: + dataset.attrs["orig_domain"] = metadata["orig_domain"] assert_array_equal( dimension.square_domain(dataset, method, inverse)["precip_intensity"].values, expected, ) + + +# square_domain +R = np.ones((4, 2)) +test_data_square_w_velocity = [ + # square by padding + ( + R, + {"x1": 0, "x2": 2, "y1": 0, "y2": 4, "xpixelsize": 1, "ypixelsize": 1}, + "pad", + False, + np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]), + np.array([[0, 1, 1, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 1, 1, 0]]), + ) +] + + +@pytest.mark.parametrize( + "data, metadata, method, inverse, expected, expected_velqual", + test_data_square_w_velocity, +) +def test_square_w_velocity(data, metadata, method, inverse, expected, expected_velqual): + """Test the square_domain.""" + dataset = convert_input_to_xarray_dataset( + data, None, {**fillvalues_metadata, **metadata} + ) + dataset = dataset.assign( + { + "velocity_x": (("y", "x"), data), + "velocity_y": (("y", "x"), data), + "quality": (("y", "x"), data), + } + ) + if "square_method" in metadata: + dataset.attrs["square_method"] = metadata["square_method"] + if "orig_domain" in metadata: + dataset.attrs["orig_domain"] = metadata["orig_domain"] + assert_array_equal( + dimension.square_domain(dataset, method, inverse)["precip_intensity"].values, + expected, + ) + assert_array_equal( + dimension.square_domain(dataset, method, inverse)["velocity_x"].values, + expected_velqual, + ) + assert_array_equal( + dimension.square_domain(dataset, method, inverse)["velocity_y"].values, + expected_velqual, + ) + assert_array_equal( + dimension.square_domain(dataset, method, inverse)["quality"].values, + expected_velqual, + ) diff --git a/pysteps/utils/dimension.py b/pysteps/utils/dimension.py index efa459610..d039d8ef0 100644 --- a/pysteps/utils/dimension.py +++ b/pysteps/utils/dimension.py @@ -14,14 +14,23 @@ clip_domain square_domain """ +from typing import Any, Callable + import numpy as np import xarray as xr -from pysteps.converters import compute_lat_lon +from pysteps.xarray_helpers import compute_lat_lon -_aggregation_methods = dict( - sum=np.sum, mean=np.mean, nanmean=np.nanmean, nansum=np.nansum -) +_aggregation_methods: dict[str, Callable[..., Any]] = { + "sum": np.sum, + "mean": np.mean, + "min": np.min, + "max": np.max, + "nanmean": np.nanmean, + "nansum": np.nansum, + "nanmin": np.nanmin, + "nanmax": np.nanmax, +} def aggregate_fields_time( @@ -93,7 +102,9 @@ def aggregate_fields_time( if ignore_nan: method = "".join(("nan", method)) - return aggregate_fields(dataset, window_size, dim="time", method=method) + return aggregate_fields( + dataset, window_size, dim="time", method=method, velocity_method="sum" + ) def aggregate_fields_space( @@ -147,9 +158,8 @@ def aggregate_fields_space( if np.isscalar(space_window): space_window = (space_window, space_window) - # assumes that frames are evenly spaced - ydelta = dataset["y"].values[1] - dataset["y"].values[0] - xdelta = dataset["x"].values[1] - dataset["x"].values[0] + ydelta = dataset["y"].attrs["stepsize"] + xdelta = dataset["x"].attrs["stepsize"] if space_window[0] % ydelta > 1e-10 or space_window[1] % xdelta > 1e-10: raise ValueError("space_window does not equally split dataset") @@ -166,11 +176,16 @@ def aggregate_fields_space( window_size = (int(space_window[0] / ydelta), int(space_window[1] / xdelta)) - return aggregate_fields(dataset, window_size, ["y", "x"], method) + return aggregate_fields(dataset, window_size, ["y", "x"], method, "mean") def aggregate_fields( - dataset: xr.Dataset, window_size, dim="x", method="mean", trim=False + dataset: xr.Dataset, + window_size, + dim="x", + method="mean", + velocity_method="mean", + trim=False, ) -> xr.Dataset: """Aggregate fields along a given direction. @@ -201,7 +216,11 @@ def aggregate_fields( dims, instead of a single dim method: string, optional Optional argument that specifies the operation to use - to aggregate the values within the window. + to aggregate the precipitation values within the window. + Default to mean operator. + velocity_method: string, optional + Optional argument that specifies the operation to use + to aggregate the velocity values within the window. Default to mean operator. trim: bool In case that the ``data`` is not perfectly divisible by @@ -254,9 +273,8 @@ def aggregate_fields( f"dataset.sizes[dim]={dataset.sizes[d]}" ) - # FIXME: The aggregation method is applied to all DataArrays in the Dataset - # Fix to allow support for an aggregation method per DataArray - return ( + dataset_ref = dataset + dataset = ( dataset.rolling(dict(zip(dim, window_size))) .reduce(_aggregation_methods[method]) .isel( @@ -266,6 +284,50 @@ def aggregate_fields( } ) ) + if "velocity_x" in dataset_ref: + dataset["velocity_x"] = ( + dataset_ref["velocity_x"] + .rolling(dict(zip(dim, window_size))) + .reduce(_aggregation_methods[velocity_method]) + .isel( + { + d: slice( + ws - 1, dataset_ref.sizes[d] - dataset_ref.sizes[d] % ws, ws + ) + for d, ws in zip(dim, window_size) + } + ) + ) + if "velocity_y" in dataset_ref: + dataset["velocity_y"] = ( + dataset_ref["velocity_y"] + .rolling(dict(zip(dim, window_size))) + .reduce(_aggregation_methods[velocity_method]) + .isel( + { + d: slice( + ws - 1, dataset_ref.sizes[d] - dataset_ref.sizes[d] % ws, ws + ) + for d, ws in zip(dim, window_size) + } + ) + ) + if "quality" in dataset_ref: + dataset["quality"] = ( + dataset_ref["quality"] + .rolling(dict(zip(dim, window_size))) + .reduce(_aggregation_methods["min"]) + .isel( + { + d: slice( + ws - 1, dataset_ref.sizes[d] - dataset_ref.sizes[d] % ws, ws + ) + for d, ws in zip(dim, window_size) + } + ) + ) + + return dataset def clip_domain(dataset: xr.Dataset, extent=None): @@ -298,8 +360,7 @@ def clip_domain(dataset: xr.Dataset, extent=None): def _pad_domain( dataset: xr.Dataset, dim_to_pad: str, idx_buffer: int, zerovalue: float ) -> xr.Dataset: - # assumes that frames are evenly spaced - delta = dataset[dim_to_pad].values[1] - dataset[dim_to_pad].values[0] + delta = dataset[dim_to_pad].attrs["stepsize"] end_values = ( dataset[dim_to_pad].values[0] - delta * idx_buffer, dataset[dim_to_pad].values[-1] + delta * idx_buffer, @@ -307,8 +368,6 @@ def _pad_domain( dataset_ref = dataset - # FIXME: The same zerovalue is used for all DataArrays in the Dataset - # Fix to allow support for a zerovalue per DataArray dataset = dataset_ref.pad({dim_to_pad: idx_buffer}, constant_values=zerovalue) dataset[dim_to_pad] = dataset_ref[dim_to_pad].pad( {dim_to_pad: idx_buffer}, @@ -318,6 +377,24 @@ def _pad_domain( dataset.lat.data[:], dataset.lon.data[:] = compute_lat_lon( dataset.x.values, dataset.y.values, dataset.attrs["projection"] ) + if "velocity_x" in dataset_ref: + dataset["velocity_x"].data = ( + dataset_ref["velocity_x"] + .pad({dim_to_pad: idx_buffer}, constant_values=0.0) + .values + ) + if "velocity_y" in dataset_ref: + dataset["velocity_y"].data = ( + dataset_ref["velocity_y"] + .pad({dim_to_pad: idx_buffer}, constant_values=0.0) + .values + ) + if "quality" in dataset_ref: + dataset["quality"].data = ( + dataset_ref["quality"] + .pad({dim_to_pad: idx_buffer}, constant_values=0.0) + .values + ) return dataset @@ -332,14 +409,17 @@ def square_domain(dataset: xr.Dataset, method="pad", inverse=False): :py:mod:`pysteps.io.importers`. method: {'pad', 'crop'}, optional Either pad or crop. - If pad, an equal number of zeros is added to both ends of its shortest - side in order to produce a square domain. + If pad, an equal number of pixels + filled with the minimum value of the precipitation + field is added to both ends of the precipitation fields shortest + side in order to produce a square domain. The quality and velocity fields + are always padded with zeros. If crop, an equal number of pixels is removed to both ends of its longest side in order to produce a square domain. Note that the crop method involves an irreversible loss of data. inverse: bool, optional Perform the inverse method to recover the original domain shape. - After a crop, the inverse is performed by padding the field with zeros. + After a crop, the inverse is performed by doing a pad. Returns ------- diff --git a/pysteps/converters.py b/pysteps/xarray_helpers.py similarity index 76% rename from pysteps/converters.py rename to pysteps/xarray_helpers.py index 2825af612..a1049f32f 100644 --- a/pysteps/converters.py +++ b/pysteps/xarray_helpers.py @@ -3,7 +3,7 @@ pysteps.converters ================== -Module with data converter functions. +Module with xarray helper functions. .. autosummary:: :toctree: ../generated/ @@ -11,6 +11,8 @@ convert_to_xarray_dataset """ +from datetime import datetime, timedelta + import numpy as np import numpy.typing as npt import pyproj @@ -77,7 +79,7 @@ def compute_lat_lon( return lat.reshape(x_2d.shape), lon.reshape(x_2d.shape) -def convert_to_xarray_dataset( +def convert_input_to_xarray_dataset( precip: np.ndarray, quality: np.ndarray | None, metadata: dict[str, str | float | None], @@ -111,6 +113,21 @@ def convert_to_xarray_dataset( y_r = np.linspace(metadata["y1"], metadata["y2"], h + 1)[:-1] y_r += 0.5 * (y_r[1] - y_r[0]) + if "xpixelsize" in metadata: + xpixelsize = metadata["xpixelsize"] + else: + xpixelsize = x_r[1] - x_r[0] + + if "ypixelsize" in metadata: + ypixelsize = metadata["ypixelsize"] + else: + ypixelsize = y_r[1] - y_r[0] + + if x_r[1] - x_r[0] != xpixelsize: + raise ValueError("xpixelsize does not match x1, x2 and array shape") + if y_r[1] - y_r[0] != ypixelsize: + raise ValueError("ypixelsize does not match y1, y2 and array shape") + # flip yr vector if yorigin is upper if metadata["yorigin"] == "upper": y_r = np.flip(y_r) @@ -160,6 +177,7 @@ def convert_to_xarray_dataset( "long_name": "y-coordinate in Cartesian system", "standard_name": "projection_y_coordinate", "units": metadata["cartesian_unit"], + "stepsize": ypixelsize, }, ), "x": ( @@ -170,6 +188,7 @@ def convert_to_xarray_dataset( "long_name": "x-coordinate in Cartesian system", "standard_name": "projection_x_coordinate", "units": metadata["cartesian_unit"], + "stepsize": xpixelsize, }, ), "lon": ( @@ -178,7 +197,6 @@ def convert_to_xarray_dataset( { "long_name": "longitude coordinate", "standard_name": "longitude", - # TODO(converters): Don't hard-code the unit. "units": "degrees_east", }, ), @@ -188,7 +206,6 @@ def convert_to_xarray_dataset( { "long_name": "latitude coordinate", "standard_name": "latitude", - # TODO(converters): Don't hard-code the unit. "units": "degrees_north", }, ), @@ -207,3 +224,45 @@ def convert_to_xarray_dataset( } dataset = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs) return dataset.sortby(["y", "x"]) + + +def convert_output_to_xarray_dataset( + dataset: xr.Dataset, timesteps: int | list[int], output: np.ndarray +) -> xr.Dataset: + precip_var = dataset.attrs["precip_var"] + metadata = dataset[precip_var].attrs + + last_timestamp = ( + dataset["time"][-1].values.astype("datetime64[us]").astype(datetime) + ) + time_metadata = dataset["time"].attrs + timestep_seconds = dataset["time"].attrs["stepsize"] + dataset = dataset.drop_vars([precip_var]).drop_dims(["time"]) + if isinstance(timesteps, int): + timesteps = list(range(1, timesteps + 1)) + next_timestamps = [ + last_timestamp + timedelta(seconds=timestep_seconds * i) for i in timesteps + ] + dataset = dataset.assign_coords( + {"time": (["time"], next_timestamps, time_metadata)} + ) + + if output.ndim == 4: + dataset = dataset.assign_coords( + { + "ens_number": ( + ["ens_number"], + list(range(1, output.shape[0] + 1)), + { + "long_name": "ensemble member", + "standard_name": "realization", + "units": "", + }, + ) + } + ) + dataset[precip_var] = (["ens_number", "time", "y", "x"], output, metadata) + else: + dataset[precip_var] = (["time", "y", "x"], output, metadata) + + return dataset From c2ae9db5a86c3f41e1ff390a27f75c2ab27ad55a Mon Sep 17 00:00:00 2001 From: gjm174 <56946945+gjm174@users.noreply.github.com> Date: Mon, 23 Sep 2024 15:34:40 +0200 Subject: [PATCH 8/8] Added member and time dimension (#432) --- pysteps/xarray_helpers.py | 75 +++++++++++++++++++++++++++++++++------ 1 file changed, 65 insertions(+), 10 deletions(-) diff --git a/pysteps/xarray_helpers.py b/pysteps/xarray_helpers.py index a1049f32f..33ec2f40c 100644 --- a/pysteps/xarray_helpers.py +++ b/pysteps/xarray_helpers.py @@ -83,6 +83,7 @@ def convert_input_to_xarray_dataset( precip: np.ndarray, quality: np.ndarray | None, metadata: dict[str, str | float | None], + startdate: datetime | None = None, ) -> xr.Dataset: """ Read a precip, quality, metadata tuple as returned by the importers @@ -99,6 +100,8 @@ def convert_input_to_xarray_dataset( metadata: dict Metadata dictionary containing the attributes described in the documentation of :py:mod:`pysteps.io.importers`. + startdate: datetime, None + Datetime object containing the start date and time for the nowcast Returns ------- @@ -107,7 +110,31 @@ def convert_input_to_xarray_dataset( """ var_name, attrs = cf_parameters_from_unit(metadata["unit"]) - h, w = precip.shape + + dims = None + timesteps = None + ens_number = None + + if precip.ndim == 4: + ens_number, timesteps, h, w = precip.shape + dims = ["ens_number", "time", "y", "x"] + + if startdate is None: + raise Exception("startdate missing") + + elif precip.ndim == 3: + timesteps, h, w = precip.shape + dims = ["time", "y", "x"] + + if startdate is None: + raise Exception("startdate missing") + + elif precip.ndim == 2: + h, w = precip.shape + dims = ["y", "x"] + else: + raise Exception(f"Precip field shape: {precip.shape} not supported") + x_r = np.linspace(metadata["x1"], metadata["x2"], w + 1)[:-1] x_r += 0.5 * (x_r[1] - x_r[0]) y_r = np.linspace(metadata["y1"], metadata["y2"], h + 1)[:-1] @@ -142,25 +169,33 @@ def convert_input_to_xarray_dataset( data_vars = { var_name: ( - ["y", "x"], + dims, precip, { "units": attrs["units"], "standard_name": attrs["standard_name"], "long_name": attrs["long_name"], "grid_mapping": "projection", - "transform": metadata["transform"], - "accutime": metadata["accutime"], - "threshold": metadata["threshold"], - "zerovalue": metadata["zerovalue"], - "zr_a": metadata["zr_a"], - "zr_b": metadata["zr_b"], }, ) } + + metadata_keys = [ + "transform", + "accutime", + "threshold", + "zerovalue", + "zr_a", + "zr_b", + ] + + for metadata_field in metadata_keys: + if metadata_field in metadata: + data_vars[var_name][2][metadata_field] = metadata[metadata_field] + if quality is not None: data_vars["quality"] = ( - ["y", "x"], + dims, quality, { "units": "1", @@ -210,6 +245,26 @@ def convert_input_to_xarray_dataset( }, ), } + + if ens_number is not None: + coords["ens_number"] = ( + ["ens_number"], + list(range(1, ens_number + 1, 1)), + { + "long_name": "ensemble member", + "standard_name": "realization", + "units": "", + }, + ) + + if timesteps is not None: + startdate_str = datetime.strftime(startdate, "%Y-%m-%d %H:%M:%S") + + coords["time"] = ( + ["time"], + list(range(1, timesteps + 1, 1)), + {"long_name": "forecast time", "units": "seconds since %s" % startdate_str}, + ) if grid_mapping_var_name is not None: coords[grid_mapping_name] = ( [], @@ -223,7 +278,7 @@ def convert_input_to_xarray_dataset( "precip_var": var_name, } dataset = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs) - return dataset.sortby(["y", "x"]) + return dataset.sortby(dims) def convert_output_to_xarray_dataset(