Skip to content
Snippets Groups Projects
Commit ec681b05 authored by felix kleinert's avatar felix kleinert
Browse files

Merge branch 'felix_issue280-refac-make-geo-helpers-available' into 'develop'

Resolve "Refac: Make geo-helpers available"

See merge request toar/mlair!249
parents 33ecf03f 7be3126e
Branches
Tags
5 merge requests!264Merge develop into felix_issue287_tech-wrf-datahandler-should-inherit-from-singlestationdatahandler,!259Draft: Resolve "WRF-Datahandler should inherit from SingleStationDatahandler",!253include current develop,!252Resolve "release v1.3.0",!249Resolve "Refac: Make geo-helpers available"
Pipeline #60683 passed
"""Functions related to geo locations etc."""
__author__ = 'Felix Kleinert'
__date__ = '2021-02-16'
import dask.array as da
import numpy as np
import xarray as xr
from mlair.helpers.helpers import convert2xrda
from typing import Union, Tuple
xr_int_float = Union[xr.DataArray, xr.Dataset, np.ndarray, int, float]
tuple_of_2xr_int_float = Tuple[xr_int_float, xr_int_float]
def deg2rad_all_points(lat1: xr_int_float, lon1: xr_int_float,
lat2: xr_int_float, lon2: xr_int_float) -> Tuple[tuple_of_2xr_int_float, tuple_of_2xr_int_float]:
"""
Converts coordinates provided in lat1, lon1, lat2, and lon2 from deg to rad. In fact this method just calls
dasks deg2rad method on all inputs and returns a tuple of tuples.
:param lat1: Latitude(-s) of first location
:type lat1:
:param lon1: Longitude(-s) of first location
:type lon1:
:param lat2: Latitude(-s) of second location
:type lat2:
:param lon2: Longitude(-s) of second location
:type lon2:
:return: Lats and lons in radians ((lat1, lon1), (lat2, lon2))
:rtype:
"""
lat1, lon1, lat2, lon2 = da.deg2rad(lat1), da.deg2rad(lon1), da.deg2rad(lat2), da.deg2rad(lon2)
return (lat1, lon1), (lat2, lon2)
def haversine_dist(lat1: xr_int_float, lon1: xr_int_float,
lat2: xr_int_float, lon2: xr_int_float,
to_radians: bool = True, earth_radius: float = 6371.,) -> xr.DataArray:
"""
Calculate the great circle distance between two points
on the Earth (specified in decimal degrees or in radians)
Reference: ToBeAdded
(First implementation provided by M. Langguth)
:param lat1: Latitude(-s) of first location
:param lon1: Longitude(-s) of first location
:param lat2: Latitude(-s) of second location
:param lon2: Longitude(-s) of second location
:param to_radians: Flag if conversion from degree to radiant is required
:param earth_radius: Earth radius in kilometers
:return: Distance between locations in kilometers
"""
if to_radians:
(lat1, lon1), (lat2, lon2) = deg2rad_all_points(lat1, lon1, lat2, lon2)
lat1 = convert2xrda(lat1, use_1d_default=True)
lon1 = convert2xrda(lon1, use_1d_default=True)
lat2 = convert2xrda(lat2, use_1d_default=True)
lon2 = convert2xrda(lon2, use_1d_default=True)
assert lat1.shape == lon1.shape
assert lat2.shape == lon2.shape
assert isinstance(lat1, xr.DataArray)
assert isinstance(lon1, xr.DataArray)
assert isinstance(lat2, xr.DataArray)
assert isinstance(lon2, xr.DataArray)
assert len(lat1.shape) >= len(lat2.shape)
# broadcast lats and lons to calculate distances in a vectorized manner.
lat1, lat2 = xr.broadcast(lat1, lat2)
lon1, lon2 = xr.broadcast(lon1, lon2)
a = da.sin((lat2 - lat1) / 2.0) ** 2 + \
da.cos(lat1) * da.cos(lat2) * da.sin((lon2 - lon1) / 2.0) ** 2
dist = earth_radius * 2. * da.arcsin(da.sqrt(a))
return dist
...@@ -5,7 +5,9 @@ __date__ = '2019-10-21' ...@@ -5,7 +5,9 @@ __date__ = '2019-10-21'
import inspect import inspect
import math import math
import numpy as np
import xarray as xr import xarray as xr
import dask.array as da
from typing import Dict, Callable, Union, List, Any from typing import Dict, Callable, Union, List, Any
...@@ -125,3 +127,54 @@ def extract_value(encapsulated_value): ...@@ -125,3 +127,54 @@ def extract_value(encapsulated_value):
"entry is not supported by this function.") "entry is not supported by this function.")
except TypeError: except TypeError:
return encapsulated_value return encapsulated_value
def is_xarray(arr) -> bool:
"""
Returns True if arr is xarray.DataArray or xarray.Dataset.
:param arr: variable in question
:type arr: Any
:return:
:rtype: bool
"""
return isinstance(arr, xr.DataArray) or isinstance(arr, xr.Dataset)
def convert2xrda(arr: Union[xr.DataArray, xr.Dataset, np.ndarray, int, float],
use_1d_default: bool = False, **kwargs) -> Union[xr.DataArray, xr.Dataset]:
"""
Converts np.array, int or float object to xarray.DataArray.
If a xarray.DataArray or xarray.Dataset is passed, returns that unchanged.
:param arr:
:type arr: xr.DataArray, xr.Dataset, np.ndarray, int, float
:param use_1d_default:
:type use_1d_default: bool
:param kwargs: Any additional kwargs which are accepted by xr.DataArray()
:type kwargs:
:return:
:rtype: xr.DataArray, xr.DataSet
"""
if is_xarray(arr):
return arr
else:
if use_1d_default:
if isinstance(arr, da.core.Array):
raise TypeError(f"`use_1d_default=True' is used with `arr' of type da.array. For da.arrays please "
f"pass `use_1d_default=False' and specify keywords for xr.DataArray via kwargs.")
dims = kwargs.pop('dims', 'points')
coords = kwargs.pop('coords', None)
try:
if coords is None:
coords = {dims: range(arr.shape[0])}
except (AttributeError, IndexError):
if isinstance(arr, int) or isinstance(arr, float):
coords = kwargs.pop('coords', {dims: range(1)})
dims = to_list(dims)
else:
raise TypeError(f"`arr' must be arry-like, int or float. But is of type {type(arr)}")
kwargs.update({'dims': dims, 'coords': coords})
return xr.DataArray(arr, **kwargs)
import pytest
import dask.array as da
import numpy as np
import xarray as xr
from mlair.helpers.geofunctions import deg2rad_all_points, haversine_dist
class TestDeg2RadAllPoints:
@pytest.fixture
def custom_np_arrays(self):
return np.array([0., 0.]), np.array([30., 30.]), np.array([60., 60.]), np.array([90., 90.])
@pytest.fixture
def custom_da_arrays(self, custom_np_arrays):
return (da.array(i) for i in custom_np_arrays)
@pytest.fixture
def custom_xr_arrays(self, custom_np_arrays):
return (xr.DataArray(i) for i in custom_np_arrays)
@pytest.mark.parametrize("value", ((0., 30., 60., 90.), (0, 30, 60, 90),
pytest.lazy_fixture('custom_np_arrays')))
def test_deg2rad_all_points_scalar_inputs(self, value):
(lat1, lon1), (lat2, lon2) = deg2rad_all_points(*value)
assert lat1 == pytest.approx(0.0)
assert lon1 == pytest.approx(np.pi/6.)
assert lat2 == pytest.approx(np.pi/3.)
assert lon2 == pytest.approx(np.pi/2)
def test_deg2rad_all_points_xr_arr_inputs(self, custom_xr_arrays):
(lat1, lon1), (lat2, lon2) = deg2rad_all_points(*custom_xr_arrays)
assert (lat1 == pytest.approx(0.0)).all()
assert (lon1 == pytest.approx(np.pi / 6.)).all()
assert (lat2 == pytest.approx(np.pi / 3.)).all()
assert (lon2 == pytest.approx(np.pi / 2)).all()
def test_deg2rad_all_points_xr_arr_inputs(self, custom_da_arrays):
(lat1, lon1), (lat2, lon2) = deg2rad_all_points(*custom_da_arrays)
assert lat1.compute() == pytest.approx(0.0)
assert lon1.compute() == pytest.approx(np.pi / 6.)
assert lat2.compute() == pytest.approx(np.pi / 3.)
assert lon2.compute() == pytest.approx(np.pi / 2)
class TestHaversineDist:
@pytest.mark.parametrize("lat1,lon1,lat2,lon2,to_radians,expected_dist",
((90., 0., -90., 0., True, np.pi),
(np.pi / 2., 0., np.pi / -2., 0., False, np.pi),
(0., 0., 0., 180., True, np.pi),
(0., 0., 0., np.pi, False, np.pi),
(0., 0., -45., 0, True, np.pi / 4.),
(0., 0., np.pi / -4., 0, False, np.pi / 4.),
))
def test_haversine_dist_on_unit_sphere_scalars(self, lat1, lon1, lat2, lon2, to_radians, expected_dist):
dist = haversine_dist(lat1=lat1, lon1=lon1, lat2=lat2, lon2=lon2, to_radians=to_radians, earth_radius=1.)
assert dist == pytest.approx(expected_dist)
@pytest.mark.parametrize("lat1,lon1,lat2,lon2,to_radians,expected_dist",
(
(xr.DataArray(np.array([90., 0, 45.]), dims='first', coords={'first': range(3)}),
xr.DataArray(np.array([0., 0, 0]), dims='first', coords={'first': range(3)}),
-90., 0., True,
np.array([[np.pi], [np.pi / 2.], [3./4.*np.pi]])),
(xr.DataArray(np.array([90., 0, 45.]), dims='first', coords={'first': range(3)}),
xr.DataArray(np.array([0., 0, 0]), dims='first', coords={'first': range(3)}),
np.array([-90., 0., 45.]),
np.array([0., 0., 0.]), True,
np.array([[np.pi, np.pi/2., np.pi/4.],
[np.pi/2., 0., np.pi/4.],
[3./4.*np.pi, np.pi/4., 0.]])),
))
def test_haversine_dist_on_unit_sphere_fields_and_scalars(self, lat1, lon1, lat2, lon2, to_radians, expected_dist):
dist = haversine_dist(lat1=lat1, lon1=lon1, lat2=lat2, lon2=lon2, to_radians=to_radians, earth_radius=1.)
assert (dist == pytest.approx(expected_dist)).all()
@pytest.mark.parametrize("lat1,lon1,lat2,lon2,to_radians",
(
(np.array([0., 0.]), 0., 0., 0., True),
(0., np.array([0., 0.]), 0., 0., True),
(0., 0., np.array([0., 0.]), 0., True),
(0., 0., 0., np.array([0., 0.]), True),
))
def test_haversine_dist_on_unit_sphere_missmatch_dimensions(self, lat1, lon1, lat2, lon2, to_radians):
with pytest.raises(AssertionError) as e:
dist = haversine_dist(lat1=lat1, lon1=lon1, lat2=lat2, lon2=lon2, to_radians=to_radians, earth_radius=1.)
import numpy as np import numpy as np
import xarray as xr import xarray as xr
import dask.array as da
import datetime as dt import datetime as dt
import logging import logging
...@@ -9,10 +10,12 @@ import os ...@@ -9,10 +10,12 @@ import os
import mock import mock
import pytest import pytest
import string
from mlair.helpers import to_list, dict_to_xarray, float_round, remove_items, extract_value, select_from_dict from mlair.helpers import to_list, dict_to_xarray, float_round, remove_items, extract_value, select_from_dict
from mlair.helpers import PyTestRegex from mlair.helpers import PyTestRegex
from mlair.helpers import Logger, TimeTracking from mlair.helpers import Logger, TimeTracking
from mlair.helpers.helpers import is_xarray, convert2xrda
class TestToList: class TestToList:
...@@ -312,3 +315,116 @@ class TestExtractValue: ...@@ -312,3 +315,116 @@ class TestExtractValue:
extract_value([1, 2, 3]) extract_value([1, 2, 3])
assert "Trying to extract an encapsulated value from objects with more than a single entry is not supported " \ assert "Trying to extract an encapsulated value from objects with more than a single entry is not supported " \
"by this function." in e.value.args[0] "by this function." in e.value.args[0]
class TestIsXarray:
@pytest.fixture
def custom_xr_data(self):
return xr.DataArray(np.array(range(5)))
def test_is_xarray_xr_input(self, custom_xr_data):
# data_array = xr.DataArray(np.array(range(5)))
assert is_xarray(custom_xr_data) is True
assert is_xarray(xr.Dataset({'test': custom_xr_data})) is True
def test_is_xarray_other_input(self, custom_xr_data):
assert is_xarray(1) is False
assert is_xarray(1.) is False
assert is_xarray([1, 2.]) is False
assert is_xarray([custom_xr_data]) is False
class TestConvert2xrDa:
@pytest.fixture
def custom_1d_npdata(self):
return np.array(range(9))
@pytest.fixture()
def custom_2d_npdata(self, custom_1d_npdata):
return np.stack([custom_1d_npdata, 2 * custom_1d_npdata])
@pytest.fixture
def custom_xr_dataarray(self, custom_1d_npdata):
return xr.DataArray(custom_1d_npdata)
@pytest.fixture
def custom_xr_dataset(self, custom_xr_dataarray):
return xr.Dataset({'test_1': custom_xr_dataarray})
@pytest.fixture
def custom_1d_daarray(self, custom_1d_npdata):
return da.array(custom_1d_npdata)
def test_convert2xrda_xrdata_in(self, custom_xr_dataarray, custom_xr_dataset):
assert (convert2xrda(custom_xr_dataarray) == custom_xr_dataarray).all()
assert (convert2xrda(custom_xr_dataset) == custom_xr_dataset).all()
def test_convert2xrda_npdata_in_nokwargs(self, custom_1d_npdata, custom_2d_npdata):
converted_data = convert2xrda(custom_1d_npdata)
assert isinstance(converted_data, xr.DataArray)
assert (converted_data.values == custom_1d_npdata).all()
assert converted_data.dims == ('dim_0',)
assert converted_data.dim_0.size == custom_1d_npdata.shape[0]
# Feed in a 2D-np.array without additional kwargs
converted_data = convert2xrda(custom_2d_npdata)
assert isinstance(converted_data, xr.DataArray)
assert (converted_data.values == custom_2d_npdata).all()
assert converted_data.dims == ('dim_0', 'dim_1')
assert converted_data.dim_0.size == custom_2d_npdata.shape[0]
assert converted_data.dim_1.size == custom_2d_npdata.shape[1]
def test_convert2xrda_npdata_in_nokwargs_default_true(self, custom_1d_npdata, custom_2d_npdata):
converted_data = convert2xrda(custom_1d_npdata, use_1d_default=True)
assert isinstance(converted_data, xr.DataArray)
assert (converted_data.values == custom_1d_npdata).all()
assert converted_data.dims == ('points',)
assert converted_data.points.size == custom_1d_npdata.shape[0]
# Feed in a 2D-np.array without additional kwargs
with pytest.raises(ValueError) as e:
converted_data = convert2xrda(custom_2d_npdata, use_1d_default=True)
assert "different number of dimensions on data and dims: 2 vs 1" in e.value.args[0]
@pytest.mark.parametrize("use_1d_default", (False, True))
def test_convert2xrda_npdata_in_kwargs(self, custom_1d_npdata, custom_2d_npdata, use_1d_default):
converted_data = convert2xrda(custom_1d_npdata, use_1d_default=use_1d_default, dims='other_points')
assert isinstance(converted_data, xr.DataArray)
assert (converted_data.values == custom_1d_npdata).all()
assert converted_data.dims == ('other_points',)
assert converted_data.other_points.size == custom_1d_npdata.shape[0]
# Feed in a 2D-np.array with correct additional kwargs
converted_data = convert2xrda(custom_2d_npdata, use_1d_default=use_1d_default,
dims=['test_dim_0', 'test_dim_1'],
coords={'test_dim_0': list(string.ascii_lowercase[:custom_2d_npdata.shape[0]]),
'test_dim_1': list(string.ascii_lowercase[:custom_2d_npdata.shape[1]]),
},
)
assert isinstance(converted_data, xr.DataArray)
assert (converted_data.values == custom_2d_npdata).all()
assert converted_data.dims == ('test_dim_0', 'test_dim_1')
assert (converted_data.coords['test_dim_0'].values == np.array(['a', 'b'])).all()
@pytest.mark.parametrize("scalar", (1, 2.))
def test_convert2xrda_int_float_in_nokwargs_default_true(self, scalar):
converted_data = convert2xrda(scalar, use_1d_default=True)
assert isinstance(converted_data, xr.DataArray)
assert converted_data.values == np.array([scalar])
assert converted_data.dims == ('points',)
@pytest.mark.parametrize("wrong_input", ({1: 'b'}, [1], 'abc'))
def test_convert2xrda_wrong_type_in_default_true_nokwargs(self, wrong_input):
with pytest.raises(TypeError) as e:
converted_data = convert2xrda(wrong_input, use_1d_default=True)
assert f"`arr' must be arry-like, int or float. But is of type {type(wrong_input)}" in e.value.args[0]
def test_convert2xrda_dask_in_default_true_nokwargs(self, custom_1d_daarray):
with pytest.raises(TypeError) as e:
convert2xrda(custom_1d_daarray, True)
assert "`use_1d_default=True' is used with `arr' of type da.array. For da.arrays please pass `use_1d_default=False' and specify keywords for xr.DataArray via kwargs." in \
e.value.args[0]
assert "`use_1d_default=True' is used with `arr' of type da.array. For da.arrays please pass" + \
" `use_1d_default=False' and specify keywords for xr.DataArray via kwargs." in e.value.args[0]
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment