Rotate vectors - get crs

CHANGELOG.rst

@@ -15,6 +15,8 @@ New features and enhancements
 * Generalize ``xs.regrid.create_bounds_gridmapping`` to include dataset with `crs`. (:pull:`628`, :issue:`627`).
 * Ability to return multiple periods if passed multiple warming levels in ``xs.extract.get_period_from_warming_level``. (:pull:`630`, :issue:`629`).
 * Update xscen to xclim 0.58 (:pull:`634`).
+* New function ``xs.spatial.rotate_vectors`` to rotate vectors from/to their native grid axes to/from real west-east/south-north axes. (:pull:`635`).
+* New function ``xs.spatial.get_crs`` to get a cartopy crs from a grid mapping variable (only Rotated Pole and Oblique Mercator) (:pull:`635`).
 
 Bug fixes
 ^^^^^^^^^

src/xscen/regrid.py

@@ -22,7 +22,7 @@
     Regridder = "xesmf.Regridder"
 
 from .config import parse_config
-from .spatial import get_grid_mapping
+from .spatial import get_crs, get_grid_mapping
 
 __all__ = ["create_bounds_gridmapping", "create_mask", "regrid_dataset"]
 
@@ -411,35 +411,11 @@ def create_bounds_gridmapping(ds: xr.Dataset, gridmap: str | None = None) -> xr.
         f"{xname}_vertices", f"{yname}_vertices"
     )
 
-    # Some CRS have additional attributes according to CF conventions
-    def _get_opt_attr_as_float(da: xr.DataArray, attr: str) -> float | None:
-        return float(da.attrs[attr]) if attr in da.attrs else None
-
-    if ds[gridmap].attrs["grid_mapping_name"] == "rotated_latitude_longitude":
-        # Get cartopy's crs for the projection
-        RP = ccrs.RotatedPole(
-            pole_longitude=float(ds[gridmap].grid_north_pole_longitude),
-            pole_latitude=float(ds[gridmap].grid_north_pole_latitude),
-            central_rotated_longitude=_get_opt_attr_as_float(
-                ds[gridmap], "north_pole_grid_longitude"
-            ),
-        )
-    elif ds[gridmap].attrs["grid_mapping_name"] == "oblique_mercator":
-        RP = ccrs.ObliqueMercator(
-            central_longitude=float(ds[gridmap].longitude_of_projection_origin),
-            central_latitude=float(ds[gridmap].latitude_of_projection_origin),
-            false_easting=_get_opt_attr_as_float(ds[gridmap], "false_easting"),
-            false_northing=_get_opt_attr_as_float(ds[gridmap], "false_northing"),
-            scale_factor=float(ds[gridmap].scale_factor_at_projection_origin),
-            azimuth=float(ds[gridmap].azimuth_of_central_line),
-        )
-    else:
-        raise NotImplementedError(f"Grid mapping {gridmap} not yet implemented.")
-
-    PC = ccrs.PlateCarree()
+    crs = get_crs(ds[gridmap])
+    PC = ccrs.PlateCarree(globe=crs.globe)
 
     # Project points
-    pts = PC.transform_points(RP, xv.values, yv.values)
+    pts = PC.transform_points(crs, xv.values, yv.values)
     lonv = xv.copy(data=pts[..., 0]).rename("lon_vertices")
     latv = yv.copy(data=pts[..., 1]).rename("lat_vertices")
 

src/xscen/spatial.py

@@ -7,6 +7,7 @@
 from collections.abc import Sequence
 from pathlib import Path
 
+import cartopy.crs
 import clisops.core.subset
 import dask
 import geopandas as gpd
@@ -23,7 +24,9 @@
 __all__ = [
     "creep_fill",
     "creep_weights",
+    "get_crs",
     "get_grid_mapping",
+    "rotate_vectors",
     "subset",
 ]
 
@@ -158,6 +161,73 @@ def _dot(arr, wei):
     )
 
 
+def rotate_vectors(
+    uu: xr.DataArray,
+    vv: xr.DataArray,
+    crs: cartopy.crs.Projection | None = None,
+    reverse: bool = False,
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """
+    Rotate a vector field defined in its grid space to the real south-north / west-east axes.
+
+    Parameters
+    ----------
+    uu : xr.DataArray
+      The u component of the vector (along the X grid axis).
+      Must be CF-compliant, so its coordinates are correctly found.
+    vv : xr.DataArray
+      The v component of the vector (along the Y grid axis)
+      Must be CF-compliant, so its coordinates are correctly found.
+    crs : pyproj.CRS, optional
+      The projection of the UU and VV grid. If not given, the grid mapping attribute from UU is read.
+    reverse: bool
+      If True, the opposite rotation is performed : inputs are understood as in the real S-N, W-E axes and are
+      rotated back to the grid axes.
+
+    Returns
+    -------
+    uuc : xr.DataArray
+      The west-east component of the vector. Attributes from uu are copied.
+    vvc : xr.DataArray
+      The south-north component of the vector. Attributes from vv are copied.
+    """
+    if crs is None:
+        crs = get_crs(uu.rename("uu").to_dataset())
+
+    lalo = cartopy.crs.PlateCarree(globe=crs.globe)
+    geod = lalo.get_geod()
+
+    X = uu.cf["X"]
+    Y = uu.cf["Y"]
+    # ensure everything has the same dim order
+    YY, XX = xr.broadcast(Y, X)
+    XXarr = XX.values
+    YYarr = YY.values
+    # a very small increment along the y axis
+    if X.dims == Y.dims:
+        # we got a list of points, can't infer the grid resolution
+        dy = 1e-5
+    else:
+        dy = (Y[1] - Y[0]).item() / 1000
+
+    # Get lat lon of center points
+    crds_c = lalo.transform_points(crs, XXarr, YYarr)
+    # Get lat lon of points just slightly above along the y axis
+    crds_y = lalo.transform_points(crs, XXarr, YYarr + dy)
+    # The azimuth of this short vector
+    az, _, _ = geod.inv(crds_c[..., 0], crds_c[..., 1], crds_y[..., 0], crds_y[..., 1])
+    if reverse:
+        az = -az
+    # The rotation angle : opposite of the azimuth, in rad
+    th = xr.DataArray(-np.deg2rad(az), dims=YY.dims, coords=YY.coords)
+    # Rotation matrix like in Algèbre linéaire et géométrie vectorielle
+    c = np.cos(th)
+    s = np.sin(th)
+    uuc = uu * c - vv * s
+    vvc = uu * s + vv * c
+    return uuc.assign_attrs(uu.attrs), vvc.assign_attrs(vv.attrs)
+
+
 def subset(
     ds: xr.Dataset,
     method: str,
@@ -475,6 +545,57 @@ def get_grid_mapping(ds: xr.Dataset) -> str:
     return gridmap[0] if gridmap else ""
 
 
+def get_crs(gridmap: xr.Dataset | xr.DataArray) -> cartopy.crs.Projection:
+    """
+    Get the cartopy CRS.
+
+    Parameters
+    ----------
+    gridmap: xr.Dataset or xr.DataArray
+      Either a dataset that has a grid mapping variable or that grid mapping variable directly.
+
+    Returns
+    -------
+    cartopy.crs.Projection
+      The cartopy crs. Only RotatedPole and ObliqueMercator are supported.
+    """
+    if isinstance(gridmap, xr.Dataset):
+        gridmap = gridmap[get_grid_mapping(gridmap)]
+    cf_params = gridmap.attrs
+    globe = cartopy.crs.Globe(
+        datum=cf_params.get("horizontal_datum_name"),
+        ellipse=cf_params.get("reference_ellipsoid_name", "WGS84"),
+        semimajor_axis=(
+            cf_params.get("earth_radius") or cf_params.get("semi_major_axis")
+        ),
+        semiminor_axis=cf_params.get("semi_minor_axis"),
+        inverse_flattening=cf_params.get("inverse_flattening"),
+        towgs84=cf_params.get("towgs84"),
+    )
+    if cf_params["grid_mapping_name"] == "rotated_latitude_longitude":
+        crs = cartopy.crs.RotatedPole(
+            pole_longitude=cf_params["grid_north_pole_longitude"],
+            pole_latitude=cf_params["grid_north_pole_latitude"],
+            central_rotated_longitude=cf_params.get("north_pole_grid_longitude", 0),
+            globe=globe,
+        )
+    elif cf_params["grid_mapping_name"] == "oblique_mercator":
+        crs = cartopy.crs.ObliqueMercator(
+            central_longitude=cf_params["longitude_of_projection_origin"],
+            central_latitude=cf_params["latitude_of_projection_origin"],
+            scale_factor=cf_params["scale_factor_at_projection_origin"],
+            azimuth=cf_params["azimuth_of_central_line"],
+            false_easting=cf_params.get("false_easting", 0),
+            false_northing=cf_params.get("false_northing", 0),
+            globe=globe,
+        )
+    else:
+        raise NotImplementedError(
+            f"Grid mapping {cf_params['grid_mapping_name']} not implemented."
+        )
+    return crs
+
+
 def _estimate_grid_resolution(ds: xr.Dataset) -> tuple[float, float]:
     # Since this is to compute a buffer, we take the maximum difference as an approximation.
     # Estimate the grid resolution

src/xscen/utils.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 from types import ModuleType
 
+import cartopy.crs as ccrs
 import cftime
 import flox.xarray
 import numpy as np

tests/test_spatial.py

@@ -447,3 +447,62 @@ def test_estimate_res_2d(lon_res, lat_res):
     lon_res_est, lat_res_est = _estimate_grid_resolution(ds)
     np.testing.assert_allclose(lon_res_est, ds.lon.diff("rlon").max())
     np.testing.assert_allclose(lat_res_est, ds.lat.diff("rlat").max())
+
+
+def test_rotate_vectors():
+    # Test data from CaSR 3.1, original rotation done by ECCC using librmn
+    rlon = xr.DataArray(
+        [20.042786, 20.042786, -29.997223, -29.997223],
+        dims=("x",),
+        name="rlon",
+        attrs={"axis": "X", "standard_name": "grid_longitude"},
+    )
+    rlat = xr.DataArray(
+        [-29.970001, 19.980001, -29.970001, 19.980001],
+        dims=("x",),
+        name="rlat",
+        attrs={"axis": "Y", "standard_name": "grid_latitude"},
+    )
+    crs = xr.DataArray(
+        attrs={
+            "grid_mapping_name": "rotated_latitude_longitude",
+            "semi_major_axis": 6370997,
+            "semi_minor_axis": 6370997,
+            "grid_north_pole_latitude": 31.758312454493154,
+            "grid_north_pole_longitude": 87.59703130293302,
+            "north_pole_grid_longitude": 0,
+            "reference_ellipsoid_name": "sphere",
+        }
+    )
+    UU = xr.DataArray(
+        [0.01779366, 10.668184, -7.882597, 1.4650593],
+        dims=("x",),
+        attrs={"grid_mapping": "crs"},
+    )
+    VV = xr.DataArray(
+        [8.246281, -6.699032, -8.255672, -5.027157],
+        dims=("x",),
+        attrs={"grid_mapping": "crs"},
+    )
+    UUC = xr.DataArray(
+        [2.6767654, 1.1289139, -3.2187424, 5.0918045],
+        dims=("x",),
+        attrs={"grid_mapping": "crs"},
+    )
+    VVC = xr.DataArray(
+        [7.800007, -12.545696, -10.951946, -1.2234306],
+        dims=("x",),
+        attrs={"grid_mapping": "crs"},
+    )
+    ds = xr.Dataset(
+        data_vars={"uu": UU, "vv": VV, "uuc": UUC, "vvc": VVC},
+        coords={"rlon": rlon, "rlat": rlat, "crs": crs},
+    )
+
+    myuuc, myvvc = xs.spatial.rotate_vectors(ds.uu, ds.vv)
+    np.testing.assert_allclose(myuuc, ds.uuc, atol=1e-3)
+    np.testing.assert_allclose(myvvc, ds.vvc, atol=1e-3)
+
+    myuu, myvv = xs.spatial.rotate_vectors(ds.uuc, ds.vvc, reverse=True)
+    np.testing.assert_allclose(myuu, ds.uu, atol=1e-3)
+    np.testing.assert_allclose(myvv, ds.vv, atol=1e-3)