diff --git a/mlair/helpers/geofunctions.py b/mlair/helpers/geofunctions.py
index 2be519c50b99db7f44bdb49bb3eefd3540950812..d1a3b71236763e17b4ed5bf29d88961c29030569 100644
--- a/mlair/helpers/geofunctions.py
+++ b/mlair/helpers/geofunctions.py
@@ -3,12 +3,11 @@ __author__ = 'Felix Kleinert'
__date__ = '2021-02-16'
import dask.array as da
-import numpy
import numpy as np
import xarray as xr
from mlair.helpers.helpers import convert2xrda
-from typing import Union, Tuple, List
+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]
@@ -37,7 +36,7 @@ def deg2rad_all_points(lat1: xr_int_float, lon1: xr_int_float,
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:
+ 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)
@@ -83,8 +82,8 @@ def haversine_dist(lat1: xr_int_float, lon1: xr_int_float,
def bearing_angle(lat1: xr_int_float, lon1: xr_int_float,
- lat2: xr_int_float, lon2: xr_int_float,
- to_radians: bool = True, return_deg: bool = True) -> xr.DataArray:
+ lat2: xr_int_float, lon2: xr_int_float,
+ to_radians: bool = True, return_deg: bool = True) -> xr.DataArray:
"""
Calculate initial bearing angle (forward azimuth) between multiple points.
Implemented by following
@@ -150,7 +149,7 @@ class WindSector:
deg_per_sector = 360. / number_of_sectores
edges_per_sector = np.array([- deg_per_sector / 2. + n * deg_per_sector for n in range(number_of_sectores + 1)])
# convert to range 0,360
- edges_per_sector = (edges_per_sector + 360) % 360.
+ edges_per_sector = (edges_per_sector + 360.) % 360.
wind_sector_edges = {sect: [edges_per_sector[i], edges_per_sector[i + 1]] for i, sect in
enumerate(wind_sectores)}
self.edges_per_sector = edges_per_sector
@@ -160,7 +159,7 @@ class WindSector:
def _is_value_in_sector(value, left_edge, right_edge) -> bool:
if left_edge < right_edge: # all sectors which does not cross 360°/0°
- return da.logical_and(left_edge <= value, value < right_edge)
+ return da.logical_and(left_edge <= value, value < right_edge)
elif left_edge > right_edge: # (noth) sector which crosses 360°/0°
return da.logical_or(
da.logical_and(left_edge <= value, value < 360.), # value between left_edge and 360°
@@ -174,6 +173,239 @@ class WindSector:
return self._is_value_in_sector(value, left_edge, right_edge)
+class VectorRotate:
+ RAD_PER_DEG = np.pi / 180.
+ DEG_PER_RAD = 180. / np.pi
+
+ @staticmethod
+ def interpolate_to_grid_center(data: xr.DataArray, interpolation_dim: str, **kwargs) -> da.array:
+ """
+ Interpolate from staged grid cells to grid cell center.
+
+ :param data: variable on staged ("half-level") grid
+ :type data:
+ :param interpolation_dim: Dimension name of staged dimension
+ :type interpolation_dim:
+ :param constant_values:
+ :type constant_values:
+ :return:
+ :rtype:
+ """
+ constant_values = kwargs.pop('constant_values', np.nan)
+
+ # store order of orig dims
+ orig_dims = data.dims
+
+ # transpose to ensure that staged dimension is located at -1
+ data = data.transpose(..., interpolation_dim)
+ new_dims = data.dims
+
+ # create idx list for backward transposition to original dimension order
+ idx_to_backward_transpose = [list(orig_dims).index(i) for i in new_dims]
+
+ # pad data i) left, ii) right for interpolation (mean between two staged values)
+ d1 = data.pad({interpolation_dim: (1, 0)}, constant_values=constant_values,)
+ d2 = data.pad({interpolation_dim: (0, 1)}, constant_values=constant_values,)
+ d = (d1.data + d2.data) * .5
+
+ # drop padding boundary
+ d = d[..., 1:-1]
+ if orig_dims != new_dims:
+ d = da.transpose(d, idx_to_backward_transpose)
+ return d
+
+ @staticmethod
+ def ull_vll2wspd(ull, vll):
+ """
+ Convert u, v to wind speed.
+
+ :param ull: Wind's u-component
+ :type ull:
+ :param vll: Wind's v-component
+ :type vll:
+ :return:
+ :rtype:
+ """
+ return da.sqrt(ull ** 2 + vll ** 2)
+
+ def ull_vll2wdir(self, ull, vll):
+ """
+ Conbert u, v to wind direction.
+
+ :param ull: Wind's u-component
+ :type ull:
+ :param vll: Wind's v-component
+ :type vll:
+ :return:
+ :rtype:
+ """
+ wdir = da.arctan2(-ull, -vll) * self.DEG_PER_RAD
+ wdir = (wdir + 360.) % 360.
+ return wdir
+
+ def ull_vll2wspd_wdir(self, ull, vll):
+ """
+ Convert u, v to wind speed and direction (wspd, wdir).
+
+ :param ull: Wind's u-component
+ :type ull:
+ :param vll: Wind's v-component
+ :type vll:
+ :return:
+ :rtype:
+ """
+ return self.ull_vll2wspd(ull, vll), self.ull_vll2wdir(ull, vll)
+
+
+class VectorRotateLambertConformal2latlon(VectorRotate):
+ CEN_LON = 12.
+ CEN_LAT = 52.5
+ TRUELAT1 = 30.
+ TRUELAT2 = 60.
+ STAND_LON = 12.
+
+ def __init__(self, uu, vv, xlat='XLAT', xlong='XLONG', cen_lon=CEN_LON, cen_lat=CEN_LAT, truelat1=TRUELAT1, truelat2=TRUELAT2, stand_lon=STAND_LON):
+ self.cen_lon = cen_lon
+ self.cen_lat = cen_lat
+ self.truelat1 = truelat1
+ self.truelat2 = truelat2
+ self.stand_lon = stand_lon
+ self.xlat = self._set_llcoords(None, xlat)
+ self.xlong = self._set_llcoords(None, xlong)
+
+ @staticmethod
+ def _set_llcoords(wind_grd, xll):
+ if isinstance(xll, str):
+ return wind_grd[xll]
+ elif isinstance(xll, np.ndarray) or isinstance(xll, da.core.Array) or isinstance(xll, xr.DataArray):
+ return xll
+ else:
+ raise TypeError(
+ f"`xll' must be name of xll coord of wind_grd as str or np/da/xr array. But is of type {type(xll)}")
+
+ @property
+ def cone(self):
+ """
+ Rotation formulas taken from WRFSI's module_map_util.F.
+
+ :return:
+ :rtype:
+ """
+ # case (1) ! Lambert conformal
+ if np.abs(self.truelat1 - self.truelat2) > 0.1: # !secant projection
+ return (da.log(da.cos(self.truelat1 * self.RAD_PER_DEG)) - da.log(
+ da.cos(self.truelat2 * self.RAD_PER_DEG))) / (
+ da.log(da.tan((90. - np.abs(self.truelat1)) * self.RAD_PER_DEG * .5)) - da.log(
+ da.tan((90. - np.abs(self.truelat2)) * self.RAD_PER_DEG * .5)))
+ else: # !tangent projection
+ return da.sin(np.abs(self.truelat1)*self.RAD_PER_DEG)
+
+ @property
+ def _alpha(self):
+ return self._difflong * self.cone * self.RAD_PER_DEG
+
+ @property
+ def _difflong(self):
+ dlong = self.xlong - self.stand_lon
+ dlong = da.where(dlong > +180., dlong - 360., dlong)
+ dlong = da.where(dlong < -180., dlong + 360., dlong)
+ return dlong
+
+ def ugrd2ull(self, ugrd, vgrd):
+ """
+ Rotate u from Lambert conformal to geographic lat/lon.
+
+ :param ugrd:
+ :type ugrd:
+ :param vgrd:
+ :type vgrd:
+ :return:
+ :rtype:
+ """
+ return vgrd * da.sin(self._alpha) + ugrd * da.cos(self._alpha)
+
+ def vgrd2vll(self, ugrd, vgrd):
+ """
+ Rotate v from Lambert conformal to geographic lat/lon.
+
+ :param ugrd:
+ :type ugrd:
+ :param vgrd:
+ :type vgrd:
+ :return:
+ :rtype:
+ """
+ return vgrd * da.cos(self._alpha) - ugrd * da.sin(self._alpha)
+
+ def ugrd_vgrd2ull_vll(self, ugrd, vgrd):
+ """
+ Rotate u, v from Lambert conformal to u, v in geographic lat/lon.
+
+ :param ugrd:
+ :type ugrd:
+ :param vgrd:
+ :type vgrd:
+ :return:
+ :rtype:
+ """
+ return self.ugrd2ull(ugrd, vgrd), self.vgrd2vll(ugrd, vgrd)
+
+ def ustg_vstg2ull_vll(self, ustg, vstg, ustg_dim: str, vstg_dim: str):
+ """
+ First interpolate u, v from staged grid cells to grid center, than rotate u, v from Lambert conformal to u, v
+ in geographic lat/lon.
+ :param ustg: wind's u-component on staged ("half") grid in Lambert conformal
+ :type ustg:
+ :param vstg: wind's v-component on staged ("half") grid in Lambert conformal
+ :type vstg:
+ :param ustg_dim: Dimension name of staged u-component
+ :type ustg_dim:
+ :param vstg_dim: Dimension name of staged v-component
+ :type vstg_dim:
+ :return: u, v on grid center in geographic lat/lon direction
+ :rtype:
+ """
+ ugrd = self.interpolate_to_grid_center(ustg, ustg_dim)
+ vgrd = self.interpolate_to_grid_center(vstg, vstg_dim)
+ return self.ugrd_vgrd2ull_vll(ugrd, vgrd)
+
+ def ugrd_vgrd2wspd_wdir(self, ugrd, vgrd):
+ """
+ Convert u, v from Lambert conformal to wdspd, wdir in geographic lat/lon.
+
+ :param ugrd: wind's u-component on grid center in Lambert conformal
+ :type ugrd:
+ :param vgrd: wind's v-component on grid center in Lambert conformal
+ :type vgrd:
+ :return:
+ :rtype:
+ """
+ ull, vll = self.ugrd_vgrd2ull_vll(ugrd, vgrd)
+ return self.ull_vll2wspd_wdir(ull, vll)
+
+ def ustg_vstg2wspd_wdir(self, ustg, vstg, ustg_dim: str, vstg_dim: str):
+ """
+ First interpolate u, v from staged grid cells to grid center, than rotate u, v from Lambert conformal to
+ wdspd, wdir in geographic lat/lon.
+
+ :param ustg: wind's u-component on staged ("half") grid in Lambert conformal
+ :type ustg:
+ :param vstg: wind's v-component on staged ("half") grid in Lambert conformal
+ :type vstg:
+ :param ustg_dim: Dimension name of staged u-component
+ :type ustg_dim:
+ :param vstg_dim: Dimension name of staged u-component
+ :type vstg_dim:
+ :return:
+ :rtype:
+ """
+ ull, vll = self.ustg_vstg2ull_vll(ustg, vstg, ustg_dim, vstg_dim)
+ return self.ull_vll2wspd_wdir(ull, vll)
+
+
+
+
+
if __name__ == '__main__':
kansas_StLouis = bearing_angle(lat1=39.099912, lon1=-94.581213,
lat2=38.627089, lon2=-90.200203,
@@ -187,5 +419,3 @@ if __name__ == '__main__':
t1_exp = 0.
t2 = bearing_angle(lat1=0., lon1=0., lat2=0., lon2=90., to_radians=True, return_deg=True)
t2_exp = 90.
-
-