Conservative Region Aggregation with Xarray, Geopandas and Sparse¶

Goal: Regrid a global precipitation dataset into countries conservatively, i.e. by exactly partitioning each grid cell into the precise region boundaries.

Meta Goal: Demonstrate that we don't necessarily need a package for this workflow and showcase some of the new capabilities of GeoPandas along the way.

Approach: We take a three step approach:

Represent both the original grid and target grid as GeoSeries with Polygon geometry
Compute their area overlay and turn it into a sparse matrix
Perform matrix multiplication on the full Xarray dataset (with a time dimension)

It is quite fast and transparent.

In [1]:

import xarray as xr
import geopandas as gp
import pandas as pd
import sparse
%xmode minimal
import hvplot.pandas
import hvplot.xarray
import dask
import cf_xarray

Exception reporting mode: Minimal

Load Region Data¶

In [2]:

#! wget https://www.naturalearthdata.com/http//www.naturalearthdata.com/download/50m/cultural/ne_50m_admin_0_countries.zip
#! unzip ne_50m_admin_0_countries.zip

Load with geopandas:

In [3]:

regions_df = gp.read_file("ne_50m_admin_0_countries.shp")

In [4]:

regions_df.head(1)

Out[4]:

	featurecla	scalerank	LABELRANK	SOVEREIGNT	SOV_A3	ADM0_DIF	LEVEL	TYPE	TLC	ADMIN	...	FCLASS_TR	FCLASS_ID	FCLASS_PL	FCLASS_GR	FCLASS_IT	FCLASS_NL	FCLASS_SE	FCLASS_BD	FCLASS_UA	geometry
0	Admin-0 country	1	3	Zimbabwe	ZWE	0	2	Sovereign country	1	Zimbabwe	...	None	None	None	None	None	None	None	None	None	POLYGON ((31.28789 -22.40205, 31.19727 -22.344...

1 rows × 169 columns

In [5]:

region_name = 'SOVEREIGNT'

Check it:

In [6]:

regions_df.plot(column=region_name, figsize=(10,4))

Out[6]:

<AxesSubplot:>

In [7]:

bbox = tuple(regions_df.total_bounds)
bbox

Out[7]:

(-180.0, -89.99892578125002, 180.0, 83.599609375)

All geodataframes should have a coordinate reference system. This is important (and sometimes unfamiliar to users coming from the global climate world).

In [8]:

crs_orig = f'EPSG:{regions_df.crs.to_epsg()}'
crs_orig

Out[8]:

'EPSG:4326'

We will now transform to an area preserving projection. This is imporant because we want to do area-weighted regridding.

In [9]:

# use an area preserving projections
crs_area = "ESRI:53034"
regions_df = regions_df.to_crs(crs_area)

Load Precipitation Data¶

This is the NASA / NOAA Global Precipitation Climatology Project processed and stored by Pangeo Forge: https://pangeo-forge.org/dashboard/feedstock/42. We want to aggregate the precipitation into countries defined in the shapefile.

Note that we have the entire dataset in one object. We convert the longitude from [0,360] to [-180,180] because that's what holoviz expects.

In [10]:

store = 'https://ncsa.osn.xsede.org/Pangeo/pangeo-forge/gpcp-feedstock/gpcp.zarr'
with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ds = xr.open_dataset(store, engine='zarr', chunks={})
    lon = ds.cf['X'].name
    lat = ds.cf['Y'].name
    lon_attrs = ds.cf['X'].attrs # assign_coords drops attrs, so copy...
    ds = ds.cf.assign_coords(X=(((ds[lon] + 180) % 360) - 180)).sortby(lon)
    ds[lon].attrs = lon_attrs # ... and paste
    ds[lon].attrs['valid_range'] = [-180,180]

Subset gridded model results to bounds of spatial dataframe to save on memory and computation. More useful when the regions_df is much smaller than the footprint of the gridded model

In [11]:

ds = ds.cf.sel(X=slice(bbox[0],bbox[2]), Y=slice(bbox[1],bbox[3]))

Now we extract just the horizontal grid information. The dataset has information about the lat and lon bounds of each cell, which we need to create the polygons.

In [12]:

grid = ds.drop(['time', 'time_bounds',  'precip']).reset_coords().load()
grid

In [13]:

ds

Out[13]:

<xarray.Dataset>
Dimensions:      (latitude: 173, nv: 2, longitude: 360, time: 9226)
Coordinates:
    lat_bounds   (latitude, nv) float32 dask.array<chunksize=(173, 2), meta=np.ndarray>
  * latitude     (latitude) float32 -89.0 -88.0 -87.0 -86.0 ... 81.0 82.0 83.0
    lon_bounds   (longitude, nv) float32 dask.array<chunksize=(360, 2), meta=np.ndarray>
  * longitude    (longitude) float32 -180.0 -179.0 -178.0 ... 177.0 178.0 179.0
  * time         (time) datetime64[ns] 1996-10-01 1996-10-02 ... 2021-12-31
    time_bounds  (time, nv) datetime64[ns] dask.array<chunksize=(200, 2), meta=np.ndarray>
Dimensions without coordinates: nv
Data variables:
    precip       (time, latitude, longitude) float32 dask.array<chunksize=(200, 173, 360), meta=np.ndarray>
Attributes: (12/45)
    Conventions:                CF-1.6, ACDD 1.3
    Metadata_Conventions:       CF-1.6, Unidata Dataset Discovery v1.0, NOAA ...
    acknowledgment:             This project was supported in part by a grant...
    cdm_data_type:              Grid
    cdr_program:                NOAA Climate Data Record Program for satellit...
    cdr_variable:               precipitation
    ...                         ...
    standard_name_vocabulary:   CF Standard Name Table (v41, 22 February 2017)
    summary:                    Global Precipitation Climatology Project (GPC...
    time_coverage_duration:     P1D
    time_coverage_end:          1996-10-01T23:59:59Z
    time_coverage_start:        1996-10-01T00:00:00Z
    title:                      Global Precipitation Climatatology Project (G...

Now we "stack" the data into a single 1D array. This is the first step towards transitioning to pandas.

In [14]:

points = grid.stack(point=(lat, lon))
points

This function creates geometries for a single pair of bounds. It is not fast, but it is fast enough here. Perhaps could be vectorized using pygeos...

In [15]:

from shapely.geometry import Polygon

def bounds_to_poly(lon_bounds, lat_bounds):
    if lon_bounds[0] >= 180:
        # geopandas needs this
        lon_bounds = lon_bounds - 360
    return Polygon([
        (lon_bounds[0], lat_bounds[0]),
        (lon_bounds[0], lat_bounds[1]),
        (lon_bounds[1], lat_bounds[1]),
        (lon_bounds[1], lat_bounds[0])
    ])

We apply this function to each grid cell.

In [16]:

%%time
import numpy as np
boxes = xr.apply_ufunc(
    bounds_to_poly,
    points.lon_bounds,
    points.lat_bounds,
    input_core_dims=[("nv",),  ("nv",)],
    output_dtypes=[np.dtype('O')],
    vectorize=True
)
boxes

CPU times: user 1.05 s, sys: 60.7 ms, total: 1.11 s
Wall time: 1.04 s

Out[16]:

<xarray.DataArray (point: 62280)>
array([<shapely.geometry.polygon.Polygon object at 0x7fb5f4356040>,
       <shapely.geometry.polygon.Polygon object at 0x7fb5f43a44c0>,
       <shapely.geometry.polygon.Polygon object at 0x7fb5f43a46a0>, ...,
       <shapely.geometry.polygon.Polygon object at 0x7fb5edc1dfa0>,
       <shapely.geometry.polygon.Polygon object at 0x7fb5edc1dfd0>,
       <shapely.geometry.polygon.Polygon object at 0x7fb5edc24040>],
      dtype=object)
Coordinates:
  * point      (point) object MultiIndex
  * latitude   (point) float32 -89.0 -89.0 -89.0 -89.0 ... 83.0 83.0 83.0 83.0
  * longitude  (point) float32 -180.0 -179.0 -178.0 -177.0 ... 177.0 178.0 179.0

Finally, we convert to a GeoDataframe, specifying the original CRS

In [17]:

%%time
grid_df= gp.GeoDataFrame(
    data={"geometry": boxes.values, "latitude": boxes[lat], "longitude": boxes[lon]},
    index=boxes.indexes["point"],
    crs=crs_orig
)
grid_df

CPU times: user 53 ms, sys: 8.07 ms, total: 61.1 ms
Wall time: 59.7 ms

Out[17]:

		geometry	latitude	longitude
latitude	longitude
-89.0	-180.0	POLYGON ((-180.00000 -89.00000, -180.00000 -88...	-89.0	-180.0
	-179.0	POLYGON ((-179.00000 -89.00000, -179.00000 -88...	-89.0	-179.0
	-178.0	POLYGON ((-178.00000 -89.00000, -178.00000 -88...	-89.0	-178.0
	-177.0	POLYGON ((-177.00000 -89.00000, -177.00000 -88...	-89.0	-177.0
	-176.0	POLYGON ((-176.00000 -89.00000, -176.00000 -88...	-89.0	-176.0
...	...	...	...	...
83.0	175.0	POLYGON ((175.00000 83.00000, 175.00000 84.000...	83.0	175.0
	176.0	POLYGON ((176.00000 83.00000, 176.00000 84.000...	83.0	176.0
	177.0	POLYGON ((177.00000 83.00000, 177.00000 84.000...	83.0	177.0
	178.0	POLYGON ((178.00000 83.00000, 178.00000 84.000...	83.0	178.0
	179.0	POLYGON ((179.00000 83.00000, 179.00000 84.000...	83.0	179.0

62280 rows × 3 columns

In [18]:

grid_df.crs

Out[18]:

<Geographic 2D CRS: EPSG:4326>
Name: WGS 84
Axis Info [ellipsoidal]:
- Lat[north]: Geodetic latitude (degree)
- Lon[east]: Geodetic longitude (degree)
Area of Use:
- name: World.
- bounds: (-180.0, -90.0, 180.0, 90.0)
Datum: World Geodetic System 1984 ensemble
- Ellipsoid: WGS 84
- Prime Meridian: Greenwich

Now we transform to the area-preserving CRS.

In [19]:

grid_df = grid_df.to_crs(crs_area)
grid_df.crs

Out[19]:

<Derived Projected CRS: ESRI:53034>
Name: Sphere_Cylindrical_Equal_Area
Axis Info [cartesian]:
- E[east]: Easting (metre)
- N[north]: Northing (metre)
Area of Use:
- name: World.
- bounds: (-180.0, -90.0, 180.0, 90.0)
Coordinate Operation:
- name: Sphere_Cylindrical_Equal_Area
- method: Lambert Cylindrical Equal Area (Spherical)
Datum: Not specified (based on Authalic Sphere)
- Ellipsoid: Sphere
- Prime Meridian: Greenwich

Key Step: Overlay the two geometries¶

This is the magic of geopandas; it can calculate the overlap between the original grid and the regions. It is expensive because it has to compare 64800 grid boxes with 242 regions.

In this dataframe, the latitude and longitude values are from the grid, while all the other columns are from the regions.

In [20]:

%time overlay = grid_df.overlay(regions_df, keep_geom_type=True)
overlay

CPU times: user 12.7 s, sys: 93.1 ms, total: 12.8 s
Wall time: 12.8 s

Out[20]:

	latitude	longitude	featurecla	scalerank	LABELRANK	SOVEREIGNT	SOV_A3	ADM0_DIF	LEVEL	TYPE	...	FCLASS_TR	FCLASS_ID	FCLASS_PL	FCLASS_GR	FCLASS_IT	FCLASS_NL	FCLASS_SE	FCLASS_BD	FCLASS_UA	geometry
0	-89.0	-180.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-20015086.796 -6369062.742, -2001508...
1	-89.0	-179.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19903891.869 -6367118.959, -1979269...
2	-89.0	-178.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19792696.943 -6367118.959, -1968150...
3	-89.0	-177.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19681502.016 -6367118.959, -1957030...
4	-89.0	-176.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19570307.089 -6367118.959, -1945911...
...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...
27274	66.0	-19.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-2001508.680 5820198.111, -202...
27275	66.0	-18.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-1890313.753 5820198.111, -194...
27276	66.0	-17.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	POLYGON ((-1890313.753 5828182.772, -1886925.7...
27277	66.0	-16.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	POLYGON ((-1779118.826 5843385.413, -1777495.4...
27278	66.0	-15.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-1667923.900 5835808.462, -166...

27279 rows × 171 columns

This is essentially already a sparse matrix mapping one grid space to the other. How sparse?

In [21]:

sparsity = len(overlay) / (len(grid_df) * len(regions_df))
sparsity

Out[21]:

0.0018099412411025652

In [22]:

overlay

Out[22]:

	latitude	longitude	featurecla	scalerank	LABELRANK	SOVEREIGNT	SOV_A3	ADM0_DIF	LEVEL	TYPE	...	FCLASS_TR	FCLASS_ID	FCLASS_PL	FCLASS_GR	FCLASS_IT	FCLASS_NL	FCLASS_SE	FCLASS_BD	FCLASS_UA	geometry
0	-89.0	-180.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-20015086.796 -6369062.742, -2001508...
1	-89.0	-179.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19903891.869 -6367118.959, -1979269...
2	-89.0	-178.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19792696.943 -6367118.959, -1968150...
3	-89.0	-177.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19681502.016 -6367118.959, -1957030...
4	-89.0	-176.0	Admin-0 country	3	4	Antarctica	ATA	0	2	Indeterminate	...	None	None	None	None	None	None	None	None	None	POLYGON ((-19570307.089 -6367118.959, -1945911...
...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...
27274	66.0	-19.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-2001508.680 5820198.111, -202...
27275	66.0	-18.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-1890313.753 5820198.111, -194...
27276	66.0	-17.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	POLYGON ((-1890313.753 5828182.772, -1886925.7...
27277	66.0	-16.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	POLYGON ((-1779118.826 5843385.413, -1777495.4...
27278	66.0	-15.0	Admin-0 country	1	3	Iceland	ISL	0	2	Sovereign country	...	None	None	None	None	None	None	None	None	None	MULTIPOLYGON (((-1667923.900 5835808.462, -166...

27279 rows × 171 columns

Let's explore these overlays a little bit

In [23]:

overlay[overlay[region_name] == "Italy"].geometry.plot(edgecolor='k')

Out[23]:

<AxesSubplot:>

We can verify that each basin's area is preserved in the overlay operation.

In [24]:

overlay.geometry.area.groupby(overlay[region_name]).sum().nlargest(10)/1e6  # km2

Out[24]:

SOVEREIGNT
Russia                      1.687963e+07
Antarctica                  1.221779e+07
Canada                      9.872058e+06
United States of America    9.449645e+06
China                       9.372241e+06
Brazil                      8.499337e+06
Australia                   7.706575e+06
India                       3.155821e+06
Argentina                   2.782670e+06
Kazakhstan                  2.712770e+06
dtype: float64

In [25]:

regions_df.geometry.area.groupby(regions_df[region_name]).sum().nlargest(10)

Out[25]:

SOVEREIGNT
Russia                      1.687963e+13
Antarctica                  1.225664e+13
Canada                      9.872058e+12
United States of America    9.449645e+12
China                       9.372241e+12
Brazil                      8.499337e+12
Australia                   7.706575e+12
India                       3.155821e+12
Argentina                   2.782670e+12
Kazakhstan                  2.712770e+12
dtype: float64

Calculate the area fraction for each region¶

This is another key step. This transform tells us how much of a country's total area comes from each of the grid cells. This is accurate because we used an area-preserving CRS.

In [26]:

grid_cell_fraction = overlay.geometry.area.groupby(overlay[region_name]).transform(lambda x: x / x.sum())
grid_cell_fraction

Out[26]:

0        0.000026
1        0.000026
2        0.000026
3        0.000026
4        0.000026
           ...   
27274    0.005459
27275    0.004817
27276    0.016131
27277    0.015022
27278    0.001937
Length: 27279, dtype: float64

We can verify that these all sum up to one.

In [27]:

grid_cell_fraction.groupby(overlay[region_name]).sum()

Out[27]:

SOVEREIGNT
Afghanistan       1.0
Albania           1.0
Algeria           1.0
Andorra           1.0
Angola            1.0
                 ... 
Western Sahara    1.0
Yemen             1.0
Zambia            1.0
Zimbabwe          1.0
eSwatini          1.0
Length: 201, dtype: float64

Turn this into a sparse Xarray DataArray¶

The first step is making a MultiIndex

In [28]:

multi_index = overlay.set_index([lat, lon, region_name]).index
df_weights = pd.DataFrame({"weights": grid_cell_fraction.values}, index=multi_index)
df_weights

Out[28]:

			weights
latitude	longitude	SOVEREIGNT
-89.0	-180.0	Antarctica	0.000026
	-179.0	Antarctica	0.000026
	-178.0	Antarctica	0.000026
	-177.0	Antarctica	0.000026
	-176.0	Antarctica	0.000026
...	...	...	...
66.0	-19.0	Iceland	0.005459
	-18.0	Iceland	0.004817
	-17.0	Iceland	0.016131
	-16.0	Iceland	0.015022
	-15.0	Iceland	0.001937

27279 rows × 1 columns

We can bring this directly into xarray as a 1D Dataset.

In [29]:

import xarray as xr
ds_weights = xr.Dataset(df_weights)
ds_weights

Now we unstack it into a sparse array.

In [30]:

weights_sparse = ds_weights.unstack(sparse=True, fill_value=0.).weights
weights_sparse

Out[30]:

<xarray.DataArray 'weights' (latitude: 170, longitude: 360, SOVEREIGNT: 201)>
<COO: shape=(170, 360, 201), dtype=float64, nnz=27267, fill_value=0.0>
Coordinates:
  * latitude    (latitude) float64 -89.0 -88.0 -87.0 -86.0 ... 81.0 82.0 83.0
  * longitude   (longitude) float64 -180.0 -179.0 -178.0 ... 177.0 178.0 179.0
  * SOVEREIGNT  (SOVEREIGNT) object 'Afghanistan' 'Albania' ... 'eSwatini'

Here we can clearly see that this is a sparse matrix, mapping the input space (lat, lon) to the output space (SOVEREIGNT).

Perform Matrix Multiplication¶

To regrid the data, we just have to multiply the original precip dataset by this matrix. There are many different ways to do this. The simplest one would just be:

In [31]:

regridded = xr.dot(ds.precip, weights_sparse)
regridded

/home/conda/users/266bc1f046057293eb778c8b95277aae62235b8cac1c12e1334002c7a7656436-20221021-195048-630919-211-pangeo/lib/python3.9/site-packages/xarray/core/indexing.py:1374: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]

Out[31]:

<xarray.DataArray (time: 9226, SOVEREIGNT: 201)>
dask.array<sum-aggregate, shape=(9226, 201), dtype=float64, chunksize=(200, 201), chunktype=numpy.ndarray>
Coordinates:
  * time        (time) datetime64[ns] 1996-10-01 1996-10-02 ... 2021-12-31
  * SOVEREIGNT  (SOVEREIGNT) object 'Afghanistan' 'Albania' ... 'eSwatini'

Unfortunately, that doesn't work out of the box, because sparse doesn't implement einsum (see https://github.com/pydata/sparse/issues/31).

In [32]:

# regridded[0].compute()  # fails

There are many ways to work around this, starting from just using a dense matrix instead. Below is the fastest thing I came up with after some extensive experimentation.

Sparse does implement matmul, so we can use that. But we have to do some reshaping to make it work with our data.

In [33]:

def apply_weights_matmul_sparse(weights, data):

    assert isinstance(weights, sparse.SparseArray)
    assert isinstance(data, np.ndarray)
    data = sparse.COO.from_numpy(data)
    data_shape = data.shape
    # k = nlat * nlon
    n, k = data_shape[0], data_shape[1] * data_shape[2]
    data = data.reshape((n, k))
    weights_shape = weights.shape
    k_, m = weights_shape[0] * weights_shape[1], weights_shape[2]
    assert k == k_
    weights_data = weights.reshape((k, m))

    regridded = sparse.matmul(data, weights_data)
    assert regridded.shape == (n, m)
    return regridded.todense()

Before applying this to the data, let's load it into memory and then chunk it again. This is not necessary (we could just stream the data from the cloud), but it is a cleaner benchmark. Chunking again allows us to leverage dask parallelism. We also eliminate some know corrupted values via a mask.

In [34]:

%%time
mask = (ds.precip >= 0) & (ds.precip < 3000)
precip = ds.precip.where(mask)
precip_in_mem = precip.compute().chunk({"time": "10MB"})
precip_in_mem

CPU times: user 12.3 s, sys: 3.14 s, total: 15.4 s
Wall time: 23.4 s

Out[34]:

<xarray.DataArray 'precip' (time: 9226, latitude: 173, longitude: 360)>
dask.array<xarray-<this-array>, shape=(9226, 173, 360), dtype=float32, chunksize=(40, 173, 360), chunktype=numpy.ndarray>
Coordinates:
  * latitude   (latitude) float32 -89.0 -88.0 -87.0 -86.0 ... 81.0 82.0 83.0
  * longitude  (longitude) float32 -180.0 -179.0 -178.0 ... 177.0 178.0 179.0
  * time       (time) datetime64[ns] 1996-10-01 1996-10-02 ... 2021-12-31
Attributes:
    cell_methods:   area: mean time: mean
    long_name:      NOAA Climate Data Record (CDR) of Daily GPCP Satellite-Ga...
    standard_name:  lwe_precipitation_rate
    units:          mm/day
    valid_range:    [0.0, 100.0]

xarray.DataArray

'precip'

time: 9226
latitude: 173
longitude: 360

dask.array<chunksize=(40, 173, 360), meta=np.ndarray>

	Array	Chunk
Bytes	2.14 GiB	9.50 MiB
Shape	(9226, 173, 360)	(40, 173, 360)
Count	231 Tasks	231 Chunks
Type	float32	numpy.ndarray

Coordinates: (3)

latitude

(latitude)

float32

-89.0 -88.0 -87.0 ... 82.0 83.0

axis :: Y
long_name :: Latitude
standard_name :: latitude
units :: degrees_north
valid_range :: [-90.0, 90.0]

array([-89., -88., -87., -86., -85., -84., -83., -82., -81., -80., -79., -78.,
       -77., -76., -75., -74., -73., -72., -71., -70., -69., -68., -67., -66.,
       -65., -64., -63., -62., -61., -60., -59., -58., -57., -56., -55., -54.,
       -53., -52., -51., -50., -49., -48., -47., -46., -45., -44., -43., -42.,
       -41., -40., -39., -38., -37., -36., -35., -34., -33., -32., -31., -30.,
       -29., -28., -27., -26., -25., -24., -23., -22., -21., -20., -19., -18.,
       -17., -16., -15., -14., -13., -12., -11., -10.,  -9.,  -8.,  -7.,  -6.,
        -5.,  -4.,  -3.,  -2.,  -1.,   0.,   1.,   2.,   3.,   4.,   5.,   6.,
         7.,   8.,   9.,  10.,  11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,
        19.,  20.,  21.,  22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,  30.,
        31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,  42.,
        43.,  44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,  54.,
        55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,  66.,
        67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,  77.,  78.,
        79.,  80.,  81.,  82.,  83.], dtype=float32)

longitude
(longitude)
float32
-180.0 -179.0 ... 178.0 179.0
axis :
X
long_name :
Longitude
standard_name :
longitude
units :
degrees_east
valid_range :
[-180, 180]
```
array([-180., -179., -178., ...,  177.,  178.,  179.], dtype=float32)
```

time

(time)

datetime64[ns]

1996-10-01 ... 2021-12-31

axis :: T
long_name :: time
standard_name :: time

array(['1996-10-01T00:00:00.000000000', '1996-10-02T00:00:00.000000000',
       '1996-10-03T00:00:00.000000000', ..., '2021-12-29T00:00:00.000000000',
       '2021-12-30T00:00:00.000000000', '2021-12-31T00:00:00.000000000'],
      dtype='datetime64[ns]')

Attributes: (5)
cell_methods :
area: mean time: mean
long_name :
NOAA Climate Data Record (CDR) of Daily GPCP Satellite-Gauge Combined Precipitation
standard_name :
lwe_precipitation_rate
units :
mm/day
valid_range :
[0.0, 100.0]

Take a look at mean precip over CONUS

In [35]:

da = precip_in_mem.mean(dim='time')

In [36]:

da.hvplot.image(x=ds.cf['X'].name, y=ds.cf['Y'].name, rasterize=True, geo=True, cmap='viridis', frame_width=600, tiles='OSM')

Out[36]:

In [37]:

with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    precip_regridded = xr.apply_ufunc(
        apply_weights_matmul_sparse,
        weights_sparse,
        precip_in_mem,
        join="left",
        input_core_dims=[[lat, lon, region_name], [lat, lon]],
        output_core_dims=[[region_name]],
        dask="parallelized",
        dask_gufunc_kwargs=dict(meta=[np.ndarray((0,))])
    )

    precip_regridded

Finally, we compute it!

In [38]:

%%time
from dask.diagnostics import ProgressBar

with ProgressBar():
    precip_regridded.load()

[#                                       ] | 3% Completed |  3.0s

/home/conda/users/266bc1f046057293eb778c8b95277aae62235b8cac1c12e1334002c7a7656436-20221021-195048-630919-211-pangeo/lib/python3.9/site-packages/sparse/_common.py:232: RuntimeWarning: Nan will not be propagated in matrix multiplication
  warnings.warn(

[########################################] | 100% Completed | 22.2s
CPU times: user 41.1 s, sys: 556 ms, total: 41.6 s
Wall time: 22.5 s

In [39]:

precip_regridded

Out[39]:

<xarray.DataArray (time: 9226, SOVEREIGNT: 201)>
array([[3.82323997e-05, 0.00000000e+00, 6.48175116e-06, ...,
        1.74498871e-02, 0.00000000e+00, 0.00000000e+00],
       [4.26562427e-01, 1.12849841e-02, 8.86701773e-02, ...,
        4.27312567e-02, 0.00000000e+00, 0.00000000e+00],
       [8.62065596e-01, 9.50921967e+00, 1.33509326e-02, ...,
        7.96630231e-04, 0.00000000e+00, 0.00000000e+00],
       ...,
       [2.08564464e-01, 1.53430027e+00, 0.00000000e+00, ...,
        5.72802686e-01, 8.95664055e+00, 1.99760773e+00],
       [7.82498343e-02, 2.20751937e-02, 4.20740792e-03, ...,
        1.99925202e+00, 1.33976102e+00, 1.72203755e+01],
       [8.46889840e+00, 0.00000000e+00, 0.00000000e+00, ...,
        4.77679192e+00, 3.91162785e+00, 1.33407013e+01]])
Coordinates:
  * SOVEREIGNT  (SOVEREIGNT) object 'Afghanistan' 'Albania' ... 'eSwatini'
  * time        (time) datetime64[ns] 1996-10-01 1996-10-02 ... 2021-12-31

Explore timeseries data by region:¶

In [40]:

ds_precip = precip_regridded.sel(SOVEREIGNT=['Italy', 'Morocco']).resample(time="MS").mean().to_dataset(region_name)

In [41]:

ds_precip.hvplot(x='time', grid=True, frame_width=1000)

Out[41]:

Explore the mean precip by region¶

In [42]:

df_mean = precip_regridded.to_pandas().mean()
df_mean.name = "precip"
df_mean = pd.DataFrame(df_mean).reset_index()

In [43]:

merged = pd.merge(regions_df, df_mean)

Convert back to geographic coordinates for plotting

In [44]:

merged_geo = merged.to_crs(crs_orig)

Matplotlib:

In [45]:

merged_geo.plot(column='precip', edgecolor='k', figsize=(12,5), legend=True)

Out[45]:

<AxesSubplot:>

Holoviz:

In [46]:

merged_geo.hvplot(c='precip', geo=True, cmap='viridis', frame_width=600, tiles='OSM', 
                  title='Mean Precip (mm/day): NASA/NOAA Global Precip Climatology [1996-10-01/2021-12-31]')

/home/conda/users/266bc1f046057293eb778c8b95277aae62235b8cac1c12e1334002c7a7656436-20221021-195048-630919-211-pangeo/lib/python3.9/site-packages/geoviews/operation/projection.py:79: ShapelyDeprecationWarning: Iteration over multi-part geometries is deprecated and will be removed in Shapely 2.0. Use the `geoms` property to access the constituent parts of a multi-part geometry.
  polys = [g for g in geom if g.area > 1e-15]

Out[46]:

Format	coo
Data Type	float64
Shape	(170, 360, 201)
nnz	27267
Density	0.00221661301336455
Read-only	True
Size	372.8K
Storage ratio	0.0