Add grid_bounds_to_polygons

cf_xarray/geometry.py

@@ -293,3 +293,58 @@
  j += n
 
  return xr.DataArray(geoms, dims=node_count.dims, coords=node_count.coords)
+
+
+def grid_to_polygons(ds: xr.Dataset) -> xr.DataArray:
+ """
+ Converts a regular 2D lat/lon grid to a 2D array of shapely polygons.
+
+ Modified from https://notebooksharing.space/view/c6c1f3a7d0c260724115eaa2bf78f3738b275f7f633c1558639e7bbd75b31456.
+
+ Parameters
+ ----------
+ ds : xr.Dataset
+ Dataset with "latitude" and "longitude" variables as well as their bounds variables.
+ 1D "latitude" and "longitude" variables are supported. This function will automatically
+ broadcast them against each other.
+
+ Returns
+ -------
+ DataArray
+ DataArray with shapely polygon per grid cell.
+ """
+ import shapely
+
+ grid = ds.cf[["latitude", "longitude"]].load().reset_coords()
+ bounds = ds.cf.bounds
+
+ assert "latitude" in bounds
+ assert "longitude" in bounds
+ (lon_bounds,) = bounds["longitude"]
+ (lat_bounds,) = bounds["latitude"]
+
+ (points,) = xr.broadcast(grid)
+
+ bounds_dim = grid.cf.get_bounds_dim_name("latitude")
+ points = points.transpose(..., bounds_dim)
+ lonbnd = points[lon_bounds].data
+ latbnd = points[lat_bounds].data
+
+ if points.sizes[bounds_dim] == 2:
+ # geopandas needs this
+ lonbnd = lonbnd[..., [0, 0, 1, 1]]
+ mask = lonbnd[..., 0] >= 180
+ lonbnd[mask, :] = lonbnd[mask, :] - 360
+ latbnd = latbnd[..., [0, 1, 1, 0]]
+
+ elif points.sizes[bounds_dim] == 4:
+ raise NotImplementedError
+ else:
+ raise ValueError(
+ f"The size of the detected bounds or vertex dimension {bounds_dim} is not 2 or 4."
+ )
+
+ polyarray = shapely.polygons(shapely.linearrings(lonbnd, latbnd))
+ boxes = points[lon_bounds][..., 0].copy(data=polyarray)
+
+ return boxes