Parallelisation of Dataset.rio.clip

[vboussange]

Hi all,
I am new to xarray, and would gladly appreciate your help.
I have a list of polygons, and an xarray.Dataset, call it dataset. For each polygon, I want to extract the dataset values which coordinates intersect with the polygon. For this task, I loop over each polygon, and use dataset.rio.clip. Here is a MWE, which works nicely.

import xarray as xr
import numpy as np
from tqdm import tqdm
from shapely.geometry import Polygon

def create_random_triangles(num_polygons, bounds):
    polygons = []
    for _ in range(num_polygons):
        points = np.random.rand(3, 2) * np.array(bounds)
        polygon = Polygon(points)
        polygons.append(polygon)
    return polygons

def create_synthetic_dataset(size, num_vars):
    data = {f'var{i}': (["x", "y"], np.random.rand(size, size)) for i in range(num_vars)}
    coords = {'x': np.arange(size), 'y': np.arange(size)}
    ds = xr.Dataset(data, coords=coords)
    ds.rio.set_spatial_dims(x_dim="x", y_dim="y", inplace=True)
    ds.rio.write_crs("epsg:4326", inplace=True)
    return ds

def process_gdf(polygons, dataset):
    env_pred = []
    for pol in tqdm(polygons):
        cropped_dataset = dataset.rio.clip([pol], dataset.rio.crs, all_touched=True)
        numpy_array = np.stack([
            cropped_dataset[var].values.reshape(-1) 
            for var in cropped_dataset.variables 
            if var not in cropped_dataset.coords
        ])
        env_pred.append(numpy_array)

    return env_pred

if __name__ == "__main__":
    num_polygons = 100
    bounds = (100, 100)
    dataset_size = 1000
    num_vars = 20
    
    polygons = create_random_triangles(num_polygons, bounds)
    dataset = create_synthetic_dataset(dataset_size, num_vars)

    env_pred = process_gdf(polygons, dataset)

However, because my list polygons can be (very) long, I would like to parallelise this code. Any ideas on how to do it? Does a vectorised function exist somewhere? My attempts with joblib or dask have proven unsuccessful. Thanks in advance for your input!

[mathause]

You can take a look at regionmask (e.g. mask_3D). (Disclaimer: I am the author). However, the all_touched option is not yet directly available (regionmask/regionmask#455).

You could use rasterio (see rasterio.features.rasterize) directly, but have to be careful with overlapping regions. (So you might be better off using the private internal regionmask function: regionmask.core.mask._mask_rasterize_3D_internal(lon, lat, polygons, all_touched=True)). You'll want the newest gdal version (v3.8.3), though (there have been some bug fixes).

[scottyhq]

My attempts with joblib or dask have proven unsuccessful

@vboussange not sure what you've tried so far. But so long as your dataset is not very large and therefore backed by a numpy array, you can use dask to parallelize your for-loop (https://tutorial.dask.org/03_dask.delayed.html#A-Typical-Workflow). Performance improvement will depend on your local resources (e.g. number of cores) as dask uses a single machine threaded scheduler by default (https://docs.dask.org/en/stable/scheduling.html).

In any case, this pattern would work with any library if you choose to try other options for the logic within process_gdf. See below:

import dask
from dask.diagnostics import ProgressBar

@dask.delayed
def process_gdf(pol, dataset):
    cropped_dataset = dataset.rio.clip([pol], dataset.rio.crs, all_touched=True)
    numpy_array = np.stack([
        cropped_dataset[var].values.reshape(-1) 
        for var in cropped_dataset.variables 
        if var not in cropped_dataset.coords
    ])
    return numpy_array

lazy_results = []
for pol in polygons:
    lazy_results.append(process_gdf(pol, dataset))

with ProgressBar():
    env_pred = dask.compute(*lazy_results)

[vboussange]

Hey @scottyhq, your solution is elegant and works beautifully! However, there is a memory leak with this Dask option. It seems that the memory usage is dramatically increasing over time, and this increase is not directly related to the list env_pred being appended, since it does not occur in the non-parallelised version. Any idea on where could this come from?