include vector rotation from Lambert conformal to geographic lat/lon

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]
@@ -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
@@ -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.
-
-