# Optimizing climatology calculation with Xarray and Dask

**URL:** https://discourse.pangeo.io/t/optimizing-climatology-calculation-with-xarray-and-dask/2453
**Category:** Science
**Created:** [May 8, 2022, 2:35pm UTC](https://discourse.pangeo.io/t/optimizing-climatology-calculation-with-xarray-and-dask/2453 "2022-05-08T14:35:06Z")
**Posts on this page:** 1
**Showing post:** 22

<div class="post-metadata">

### Author: ![rabernat](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/rabernat/32/22_2.png) [@rabernat](https://discourse.pangeo.io/u/rabernat)
#### Post date: [May 11, 2022, 10:22pm UTC](https://discourse.pangeo.io/t/optimizing-climatology-calculation-with-xarray-and-dask/2453/22 "2022-05-11T22:22:31Z")

</div>

Update on the rechunked version

### Rechunking Configuration

```auto
import zarr
import os
from rechunker import rechunk
import gcsfs

url = f"{os.environ['SCRATCH_BUCKET']}/ERA5_HiRes_Hourly.zarr"
zg = zarr.open_group(url)

fs = gcsfs.GCSFileSystem(skip_instance_cache=True, use_listings_cache=False)

temp_path = f"{os.environ['SCRATCH_BUCKET']}/ERA5_HiRes_Hourly_rechunk/temp.zarr"
target_path = f"{os.environ['SCRATCH_BUCKET']}/ERA5_HiRes_Hourly_rechunk/target.zarr"

temp_store = zarr.storage.FSStore(temp_path)
target_store = zarr.storage.FSStore(target_path)

target_chunks = {
    'tp': (7305, 103, 10),
    'time': None,
    'longitude': None,
    'latitude': None,
}

max_mem = '8GB'

r = rechunk(zg, target_chunks, max_mem, target_store, temp_store=temp_store)

# done with a 20 worker dask cluster w/ 40 GB memory each
r.execute()

```

This takes a long time (like an hour), but it is pretty stable and reliable.

### Now do flox groupby with map-reduce

I open the data and specify even longer chunks in time such that the time axis is totally contiguous.

```python
target_path = f"{os.environ['SCRATCH_BUCKET']}/ERA5_HiRes_Hourly_rechunk/target.zarr"
dsr = xr.open_dataset(
    target_path, engine="zarr", consolidated=False,
    chunks={'time': -1, 'latitude': 103, 'longitude': 10}
)

```

 ![image](https://canada1.discourse-cdn.com/flex030/uploads/pangeo/original/2X/2/243b4dbea68f216e72c5bff1a7eaed360178b796.png)

The groupby now executes flawlessly with map-reduce.

```python
method = "map-reduce"
tpmr = flox.xarray.xarray_reduce(dsr.tp, dsr.time.dt.hour, 
                                func="mean", 
                                method=method)

```

I used the same configuration with 40GB of memory per worker, but the memory usage never rose above about 10GB per worker, meaning I could have got by much cheaper

This is what my dask task stream looked like. It took about 4 minutes to do 678 GB of data. 🕶

 ![image](https://canada1.discourse-cdn.com/flex030/uploads/pangeo/original/2X/1/123a47900854e20a14133fcb975172d72cf42e08.png)

So I think this illustrates two key takeaways.

- Rechunking is really, really helpful for solving these performance. That was the outcome of the epic discussion in [Best practices to go from 1000s of netcdf files to analyses on a HPC cluster?](https://discourse.pangeo.io/t/best-practices-to-go-from-1000s-of-netcdf-files-to-analyses-on-a-hpc-cluster/588/) which led to the creation of the `rechunker` pakage.
- @dcherian’s flox package performs flawlessly when given properly chunked data.

Neither of these packages existed two years ago. I am tempted to say we have nailed it. 🚀 However, I think there is more work to do to make rechunker faster and more memory efficient.

@fmaussion - you should try rechunking the data on your HPC system and see how it goes.

---

_[View the full topic](https://discourse.pangeo.io/t/optimizing-climatology-calculation-with-xarray-and-dask/2453)._