Coverage for mlair/helpers/geofunctions.py: 100%

1"""Functions related to geo locations etc.""" 

2__author__ = 'Felix Kleinert' 

3__date__ = '2021-02-16' 

4 

5import dask.array as da 

6import numpy as np 

7import xarray as xr 

8 

9from mlair.helpers.helpers import convert2xrda 

10from typing import Union, Tuple 

11 

12xr_int_float = Union[xr.DataArray, xr.Dataset, np.ndarray, int, float] 

13tuple_of_2xr_int_float = Tuple[xr_int_float, xr_int_float] 

14 

15 

16def deg2rad_all_points(lat1: xr_int_float, lon1: xr_int_float, 

17 lat2: xr_int_float, lon2: xr_int_float) -> Tuple[tuple_of_2xr_int_float, tuple_of_2xr_int_float]: 

18 """ 

19 Converts coordinates provided in lat1, lon1, lat2, and lon2 from deg to rad. In fact this method just calls 

20 dasks deg2rad method on all inputs and returns a tuple of tuples. 

21 

22 :param lat1: Latitude(-s) of first location 

23 :type lat1: 

24 :param lon1: Longitude(-s) of first location 

25 :type lon1: 

26 :param lat2: Latitude(-s) of second location 

27 :type lat2: 

28 :param lon2: Longitude(-s) of second location 

29 :type lon2: 

30 :return: Lats and lons in radians ((lat1, lon1), (lat2, lon2)) 

31 :rtype: 

32 """ 

33 lat1, lon1, lat2, lon2 = da.deg2rad(lat1), da.deg2rad(lon1), da.deg2rad(lat2), da.deg2rad(lon2) 

34 return (lat1, lon1), (lat2, lon2) 

35 

36 

37def haversine_dist(lat1: xr_int_float, lon1: xr_int_float, 

38 lat2: xr_int_float, lon2: xr_int_float, 

39 to_radians: bool = True, earth_radius: float = 6371.,) -> xr.DataArray: 

40 """ 

41 Calculate the great circle distance between two points 

42 on the Earth (specified in decimal degrees or in radians) 

43 

44 Reference: ToBeAdded 

45 (First implementation provided by M. Langguth) 

46 

47 :param lat1: Latitude(-s) of first location 

48 :param lon1: Longitude(-s) of first location 

49 :param lat2: Latitude(-s) of second location 

50 :param lon2: Longitude(-s) of second location 

51 :param to_radians: Flag if conversion from degree to radiant is required 

52 :param earth_radius: Earth radius in kilometers 

53 :return: Distance between locations in kilometers 

54 """ 

55 

56 if to_radians: 

57 (lat1, lon1), (lat2, lon2) = deg2rad_all_points(lat1, lon1, lat2, lon2) 

58 

59 lat1 = convert2xrda(lat1, use_1d_default=True) 

60 lon1 = convert2xrda(lon1, use_1d_default=True) 

61 lat2 = convert2xrda(lat2, use_1d_default=True) 

62 lon2 = convert2xrda(lon2, use_1d_default=True) 

63 

64 assert lat1.shape == lon1.shape 

65 assert lat2.shape == lon2.shape 

66 assert isinstance(lat1, xr.DataArray) 

67 assert isinstance(lon1, xr.DataArray) 

68 assert isinstance(lat2, xr.DataArray) 

69 assert isinstance(lon2, xr.DataArray) 

70 assert len(lat1.shape) >= len(lat2.shape) 

71 

72 # broadcast lats and lons to calculate distances in a vectorized manner. 

73 lat1, lat2 = xr.broadcast(lat1, lat2) 

74 lon1, lon2 = xr.broadcast(lon1, lon2) 

75 

76 a = da.sin((lat2 - lat1) / 2.0) ** 2 + \ 

77 da.cos(lat1) * da.cos(lat2) * da.sin((lon2 - lon1) / 2.0) ** 2 

78 

79 dist = earth_radius * 2. * da.arcsin(da.sqrt(a)) 

80 

81 return dist 

82