@@ -293,3 +293,64 @@ def cf_to_points(ds: xr.Dataset):
293293 j += n
294294
295295 return xr .DataArray (geoms , dims = node_count .dims , coords = node_count .coords )
296+
297+
298+ def grid_bounds_to_polygons (ds : xr .Dataset ) -> xr .DataArray :
299+ """
300+ Converts a regular 2D lat/lon grid to a 2D array of shapely polygons.
301+
302+ Modified from https://notebooksharing.space/view/c6c1f3a7d0c260724115eaa2bf78f3738b275f7f633c1558639e7bbd75b31456.
303+
304+ Parameters
305+ ----------
306+ ds : xr.Dataset
307+ Dataset with "latitude" and "longitude" variables as well as their bounds variables.
308+ 1D "latitude" and "longitude" variables are supported. This function will automatically
309+ broadcast them against each other.
310+
311+ Returns
312+ -------
313+ DataArray
314+ DataArray with shapely polygon per grid cell.
315+ """
316+ from shapely import Polygon
317+
318+ grid = ds .cf [["latitude" , "longitude" ]].load ().reset_coords ()
319+ bounds = ds .cf .bounds
320+
321+ assert "latitude" in bounds
322+ assert "longitude" in bounds
323+ (lon_bounds ,) = bounds ["longitude" ]
324+ (lat_bounds ,) = bounds ["latitude" ]
325+
326+ (points ,) = xr .broadcast (grid )
327+
328+ bounds_dim = grid .cf .get_bounds_dim_name ("latitude" )
329+ points = points .transpose (..., bounds_dim )
330+ assert points .sizes [bounds_dim ] == 2
331+
332+ lonbnd = points [lon_bounds ].data
333+ latbnd = points [lat_bounds ].data
334+
335+ # geopandas needs this
336+ expanded_lon = lonbnd [..., [0 , 0 , 1 , 1 ]]
337+ mask = expanded_lon [..., 0 ] >= 180
338+ expanded_lon [mask , :] = expanded_lon [mask , :] - 360
339+
340+ # these magic numbers are OK :
341+ # - 4 corners to a polygon, and
342+ # - 2 from stacking lat, lon along the last axis
343+ # flatten here to make iteration easier. It would be nice to avoid that.
344+ # potentially with just np.vectorize. The polygon creation is the real slow bit.
345+ # Shapely's MultiPolygon also iterates over a list in Python...
346+ blocked = np .stack ([expanded_lon , latbnd [..., [0 , 1 , 1 , 0 ]]], axis = - 1 ).reshape (
347+ - 1 , 4 , 2
348+ )
349+ polyarray = np .array (
350+ [Polygon (blocked [i , ...]) for i in range (blocked .shape [0 ])], dtype = "O"
351+ )
352+ newshape = latbnd .shape [:- 1 ]
353+ polyarray = polyarray .reshape (newshape )
354+ boxes = points [lon_bounds ][..., 0 ].copy (data = polyarray )
355+
356+ return boxes
0 commit comments