From 34718173bd337ec7bcae149fc3226488027c6f64 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Thu, 28 Aug 2025 18:29:14 -0400 Subject: [PATCH 1/4] Rotate vectors - get crs --- CHANGELOG.rst | 2 + src/xscen/regrid.py | 30 ++--------- src/xscen/spatial.py | 121 ++++++++++++++++++++++++++++++++++++++++++ src/xscen/utils.py | 1 + tests/test_spatial.py | 59 ++++++++++++++++++++ 5 files changed, 186 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a3f8c050..6b693fa8 100644 --- a/CHANGELOG.rst +++ b/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 ^^^^^^^^^ diff --git a/src/xscen/regrid.py b/src/xscen/regrid.py index 04012628..a4b51776 100644 --- a/src/xscen/regrid.py +++ b/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,32 +411,8 @@ 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) diff --git a/src/xscen/spatial.py b/src/xscen/spatial.py index 1de715df..22157d14 100644 --- a/src/xscen/spatial.py +++ b/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 coordinate are correctly found. + vv : xr.DataArray + The v component of the vector (along the Y grid axis) + Must be CF-compliant, so its coordinate 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 + vvc : xr.DataArray + The south-north component of the vector + """ + 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 vectoriel + c = np.cos(th) + s = np.sin(th) + uuc = uu * c - vv * s + vvc = uu * s + vv * c + return uuc, vvc + + 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 diff --git a/src/xscen/utils.py b/src/xscen/utils.py index bbef4ec5..8a3e001a 100644 --- a/src/xscen/utils.py +++ b/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 diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 23eaf8f6..4e460945 100644 --- a/tests/test_spatial.py +++ b/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) From 53e504e85af70a4c34b0b21f58bd482dd40e2037 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Thu, 28 Aug 2025 18:55:34 -0400 Subject: [PATCH 2/4] forgot a name change --- src/xscen/regrid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/xscen/regrid.py b/src/xscen/regrid.py index a4b51776..9cb9fda5 100644 --- a/src/xscen/regrid.py +++ b/src/xscen/regrid.py @@ -415,7 +415,7 @@ def create_bounds_gridmapping(ds: xr.Dataset, gridmap: str | None = None) -> xr. 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") From 305e15899496915fa36ed3f4543ebd84be6f9927 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 2 Sep 2025 12:00:18 -0400 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: RondeauG <38501935+RondeauG@users.noreply.github.com> --- src/xscen/spatial.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/xscen/spatial.py b/src/xscen/spatial.py index 22157d14..742e5b10 100644 --- a/src/xscen/spatial.py +++ b/src/xscen/spatial.py @@ -174,10 +174,10 @@ def rotate_vectors( ---------- uu : xr.DataArray The u component of the vector (along the X grid axis). - Must be CF-compliant, so its coordinate are correctly found. + 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 coordinate are correctly found. + 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 @@ -220,7 +220,7 @@ def rotate_vectors( 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 vectoriel + # 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 From c91d5ded9b3e38d37f39a7a6d3a4fcfb22f61e93 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 2 Sep 2025 13:35:31 -0400 Subject: [PATCH 4/4] copy attributes --- src/xscen/spatial.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/xscen/spatial.py b/src/xscen/spatial.py index 742e5b10..06f42936 100644 --- a/src/xscen/spatial.py +++ b/src/xscen/spatial.py @@ -187,9 +187,9 @@ def rotate_vectors( Returns ------- uuc : xr.DataArray - The west-east component of the vector + The west-east component of the vector. Attributes from uu are copied. vvc : xr.DataArray - The south-north component of the vector + The south-north component of the vector. Attributes from vv are copied. """ if crs is None: crs = get_crs(uu.rename("uu").to_dataset()) @@ -225,7 +225,7 @@ def rotate_vectors( s = np.sin(th) uuc = uu * c - vv * s vvc = uu * s + vv * c - return uuc, vvc + return uuc.assign_attrs(uu.attrs), vvc.assign_attrs(vv.attrs) def subset(