restructure rotation class. Baseclass contains all methods which are not explicitly rotation

mlair/helpers/geofunctions.py

@@ -177,6 +177,20 @@ class VectorRotate:
     RAD_PER_DEG = np.pi / 180.
     DEG_PER_RAD = 180. / np.pi
 
+    def __init__(self, xlat='XLAT', xlong='XLONG'):
+        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)}")
+
     @staticmethod
     def interpolate_to_grid_center(data: xr.DataArray, interpolation_dim: str, **kwargs) -> da.array:
         """
@@ -212,6 +226,98 @@ class VectorRotate:
             d = da.transpose(d, idx_to_backward_transpose)
         return d
 
+    def ugrd2ull(self, ugrd, vgrd):
+        """
+        Rotate u from to defined system to geographic lat/lon.
+
+        :param ugrd: Wind's u-component on grid cell center in to defined system
+        :type ugrd:
+        :param vgrd: Wind's v-component on grid cell center in to defined system
+        :type vgrd:
+        :return:
+        :rtype:
+        """
+        raise NotImplementedError
+
+    def vgrd2vll(self, ugrd, vgrd):
+        """
+        Rotate v from to defined system to geographic lat/lon.
+
+        :param ugrd: Wind's u-component on grid cell center in to defined system
+        :type ugrd:
+        :param vgrd: Wind's v-component on grid cell center in to defined system
+        :type vgrd:
+        :return:
+        :rtype:
+        """
+        raise NotImplementedError
+
+    def ugrd_vgrd2ull_vll(self, ugrd, vgrd):
+        """
+        Rotate u, v from defined system to u, v in geographic lat/lon.
+
+        :param ugrd: Wind's u-component on grid cell center in defined system
+        :type ugrd:
+        :param vgrd: Wind's v-component on grid cell center in defined system
+        :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 defined system to u, v
+        in geographic lat/lon.
+
+        :param ustg: Wind's u-component on staged ("half") grid in defined system
+        :type ustg:
+        :param vstg: Wind's v-component on staged ("half") grid in defined system
+        :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 cell 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 defined system to wdspd, wdir in geographic lat/lon.
+
+        :param ugrd: Wind's u-component on grid center in defined system
+        :type ugrd:
+        :param vgrd: Wind's v-component on grid center in defined system
+        :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 defined system to
+        wdspd, wdir in geographic lat/lon.
+
+        :param ustg: Wind's u-component on staged ("half") grid in defined system
+        :type ustg:
+        :param vstg: Wind's v-component on staged ("half") grid in defined system
+        :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)
+
     @staticmethod
     def ull_vll2wspd(ull, vll):
         """
@@ -262,24 +368,14 @@ class VectorRotateLambertConformal2latlon(VectorRotate):
     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):
+    def __init__(self, cen_lon=CEN_LON, cen_lat=CEN_LAT, truelat1=TRUELAT1,
+                 truelat2=TRUELAT2, stand_lon=STAND_LON, **kwargs):
+        super().__init__(**kwargs)
         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):
@@ -335,71 +431,7 @@ class VectorRotateLambertConformal2latlon(VectorRotate):
         """
         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: Wind's u-component on grid cell center in Lambert conformal
-        :type ugrd:
-        :param vgrd: Wind's v-component on grid cell center in Lambert conformal
-        :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 cell 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)