how to use metpy.calc.divergence with chunked Xarray?

[QINQINKONG]

Hi guys! I'm learning Metpy and try to calculate moisture flux divergence with the function metpy.calc.divergence. This function doesn't require to specify dx or dy if the input is xarray DataArray with latitude and longitude coordinate which works fine if the xarray DataArray is not chunked.

However, when I input chunked xarray DataArray, it seems that the metpy.calc.divergence cannot identify the latitude and longitude anymore (maybe the lat/lon coordinate is lost when each chunk is passed to the function), and the function asks me to input dx, dy.

Below is my code:

import xarray as xr
import numpy as np
import dask
import metpy.calc as mpcalc

u10=xr.open_dataset('/scratch/bell/kong97/ERA5/u10m/u10m1950.nc').metpy.parse_cf().u10.chunk({'time':10}).metpy.quantify()
v10=xr.open_dataset('/scratch/bell/kong97/ERA5/v10m/v10m1950.nc').metpy.parse_cf().v10.chunk({'time':10}).metpy.quantify()
q=xr.open_dataset('/scratch/bell/kong97/ERA5/huss/huss_1950.nc').metpy.parse_cf().huss.chunk({'time':10}).metpy.quantify()

Then I got u10, v10 and q as chunked xarray. Below is q for example:

Then I call mpcalc.divergence use xr.apply_ufunc:
div=xr.apply_ufunc(mpcalc.divergence,u10*q,v10*q,dask="parallelized",output_dtypes=[float])
which leads to a chunked xarray:

However, when I really want to make the calculation, error pops up and ask me to input dx/dy:

Is there a way to solve this problem?

Thank you very much!

[QINQINKONG]

One possible work around solution maybe to get the grid spacing in meter or kilometer from longitude and latitude coordinate of xarray DataArray. It seems that assign_y_x may be the way to achieve that. However, I couldn't get x, y successfully from latitude/longitude. Didn't find specific examples of assign_y_x as well.

But it would be much better if there is a way to let metpy.calc.divergence automatically recognize the latitude/longitude of chunked xarray!

Thanks!

[jthielen]

A key issue here is the following:

Then I call mpcalc.divergence use xr.apply_ufunc:
div=xr.apply_ufunc(mpcalc.divergence,u10*q,v10*q,dask="parallelized",output_dtypes=[float])

As explained in the documentation, xr.apply_ufunc applies a function for unlabeled arrays. So, when used this way, mpcalc.divergence doesn't see a xarray object, but instead just data inside, which means that none of MetPy's xarray coordinate-aware features will work. Also, you may have noticed that when you tried it this way, the pint.Quantity ended up inside the Dask array rather than the other way around, which comes from a failure of Dask to properly recognize upcast types. All-in-all, this approach won't work the way you want it too!

That being said, since your data are chunked with the horizontal being contiguous, there are several alternatives for making this work. (If instead your data were chunked in the horizontal, things get a lot more complicated due to the need to overlap blocks.)

1. Try if it "just works"

MetPy's Dask support is very provisional, however, there is a chance that this could work as-is:

div = mpcalc.divergence(u10 * q, v10 * q)

2. Use xr.map_blocks to handle Dask parallelization

If something in MetPy's internals isn't playing nicely with Dask, then a possible workaround is to let xarray split up the DataArrays based on the Dask chunks. This is a bit more complicated than it otherwise could be due to only one Dask-based object being permitted.

inputs = xr.merge((u10 * q).rename('x_component'), (v10 * q).rename('y_component'))
def wrapped_divergence(ds):
    return mpcalc.divergence(ds['x_component'], ds['y_component'])
div = xr.map_blocks(wrapped_divergence, inputs, template=q)
div.attrs = ... # will likely need to fix the attrs based on difference from template

Based on pydata/xarray#4208, this should work, but if there is a pint-dask interoperability issue then you may run into issues here too.

3. Drop out of xarray

If the xarray > pint > Dask structure isn't working with either of the two methods mentioned above, then, if you still wanted to try this with Dask, you may need to temporarily drop out of xarray data structures and do the calculation a bit more manually:

dx, dy = grid_deltas_from_dataarray(q, kind='actual')
div = mpcalc.divergence((u10 * q).metpy.unit_array, (v10 * q).metpy.unit_array, dx=dx, dy=dy)
div = xr.DataArray(div, dims=q.dims, coords=q.coords)

4. Use an alternative to Dask parallelization

There's a chance that having Dask arrays inside pint Quantities still won't play nicely with MetPy's calculations (since proper Dask support is still a work in progress and there are still some things being worked out in the ecosystem as a whole.

Since in this example problem, you are just chunking over time, you could design an alternative parallelization that loops the data over time, loading just one time at a time, saving to disk between steps (such as appending to a Zarr store), and then collecting the results upon completion.

[QINQINKONG]

Hi Jon,
Thanks for your answer! This is probably the MOST helpful answer I've ever seen. I tried your first 3 approaches. Approach 1 seems will load the whole data into memory which is too big. Approach 2 just perfectly solve the problem! A directly copy of your approach 3 seems will also load whole data into memory. But it works fine if I changed to the below:

sample=xr.open_dataset('/scratch/bell/kong97/ERA5/u10m/u10m1950.nc').metpy.parse_cf()
dx, dy = metpy.xarray.grid_deltas_from_dataarray(sample.u10, kind='actual')
def divergence(x,y):
    div=mpcalc.divergence(x,y,dx=dx,dy=dy)
    return div
div = xr.apply_ufunc(divergence,(u10 * q), (v10 * q),dask="parallelized",output_dtypes=[float])

I also learned lots of other things in your answer! Thanks for your help!