Using Xarray's map_blocks() with MetPy functions?

[DanielAdriaansen]

I know that there's several limitations to processing gridded data through functions that rely on get_layer() (e.g. #1533), but I was wondering if a possible workaround is to use xr.map_blocks() for these functions onto chunked arrays, or whether the limitation in get_layer() would prevent this also?

Here is a somewhat simple example of what I am trying:

from siphon.catalog import TDSCatalog
from metpy import calc as mpcalc
from metpy.units import units
import xarray as xr
import dask

latest_gfs = TDSCatalog('https://thredds.ucar.edu/thredds/catalog/grib/NCEP/GFS/Global_0p25deg_ana/latest.xml')
gfs_ds = latest_gfs.datasets[0].remote_access(use_xarray=True)
gfs_ds = gfs_ds[['Temperature_isobaric',\
                 'Specific_humidity_isobaric',\
                 'isobaric']].sel(lat=slice(60,15),lon=slice(230,300)).isel(time=0)
gfs_ds['TD'] = mpcalc.dewpoint_from_specific_humidity(gfs_ds['isobaric']*units('pascal'),
                                                      gfs_ds['Temperature_isobaric']*units('kelvin'),\
                                                      gfs_ds['Specific_humidity_isobaric']*units('dimensionless'))
gfs_ds['P3D'] = xr.DataArray(gfs_ds.isobaric.values,dims=['isobaric'],\
                            coords={'isobaric':gfs_ds.isobaric.values},\
                            attrs={'units':'pascal'}).broadcast_like(gfs_ds['Temperature_isobaric'])

gfs_ds = gfs_ds.chunk(chunks={'lat':50,'lon':50})
gfs_ds.map_blocks(mpcalc.precipitable_water,\
                  args=[gfs_ds['P3D'],gfs_ds['TD']],\
                  kwargs={'bottom':(1000.0*units('hectopascal')),'top':(100.0*units('hectopascal'))},\
                  template=gfs_ds['TD'].isel(isobaric=0)).compute()

I end up with a huge traceback that ultimately ends with:

   3107 else:
   3108     if param.kind in (_VAR_KEYWORD, _KEYWORD_ONLY):
   3109         # Looks like we have no parameter for this positional
   3110         # argument
-> 3111         raise TypeError(
   3112             'too many positional arguments') from None
   3114     if param.kind == _VAR_POSITIONAL:
   3115         # We have an '*args'-like argument, let's fill it with
   3116         # all positional arguments we have left and move on to
   3117         # the next phase
   3118         values = [arg_val]

TypeError: too many positional arguments

So either what I am doing isn't possible based upon the design of precipitable_water(), or I am not using xr.map_blocks() correctly.

Does anyone have any experience or thoughts on this?

[dopplershift]

Try this instead:

xr.map_blocks(mpcalc.precipitable_water,
              args=[gfs_ds['P3D'],gfs_ds['TD']],
              kwargs={'bottom':(1000.0*units('hectopascal')),'top':(100.0*units('hectopascal'))},
              template=gfs_ds['TD'].isel(isobaric=0)).compute()

If you use the dataset version (like you do originally):

func (callable()) – User-provided function that accepts a Dataset as its first parameter. The function will receive a subset or ‘block’ of this Dataset (see below), corresponding to one chunk along each chunked dimension. func will be executed as func(subset_dataset, *subset_args, **kwargs).
This function must return either a single DataArray or a single Dataset.

I'm not sure this will work, but it should address the original error.

[DanielAdriaansen]

Thanks @dopplershift. I had to add the original obj (the Dataset) back in the command you suggested:

xr.map_blocks(mpcalc.precipitable_water,gfs_ds,\
              args=[gfs_ds['P3D'],gfs_ds['TD']],
              kwargs={'bottom':(1000.0*units('hectopascal')),'top':(100.0*units('hectopascal'))},
              template=gfs_ds['TD'].isel(isobaric=0)).compute()

But I am still getting the same error:

   3107 else:
   3108     if param.kind in (_VAR_KEYWORD, _KEYWORD_ONLY):
   3109         # Looks like we have no parameter for this positional
   3110         # argument
-> 3111         raise TypeError(
   3112             'too many positional arguments') from None
   3114     if param.kind == _VAR_POSITIONAL:
   3115         # We have an '*args'-like argument, let's fill it with
   3116         # all positional arguments we have left and move on to
   3117         # the next phase
   3118         values = [arg_val]

TypeError: too many positional arguments

Further up in the Traceback, I see this (the only MetPy reference in the full traceback):

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/metpy/xarray.py:1284, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1282 @functools.wraps(func)
   1283 def wrapper(*args, **kwargs):
-> 1284     bound_args = sig.bind(*args, **kwargs)
   1286     # Auto-broadcast select xarray arguments, and update bound_args
   1287     if broadcast is not **None:**

And, here is the full traceback for transparency:

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/xarray/core/dataarray.py:1101, in DataArray.compute(self, **kwargs)
   1082 """Manually trigger loading of this array's data from disk or a
   1083 remote source into memory and return a new array. The original is
   1084 left unaltered.
   (...)
   1098 dask.compute
   1099 """
   1100 new = self.copy(deep=False)
-> 1101 return new.load(**kwargs)

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/xarray/core/dataarray.py:1075, in DataArray.load(self, **kwargs)
   1057 def load(self: T_DataArray, **kwargs) -> T_DataArray:
   1058     """Manually trigger loading of this array's data from disk or a
   1059     remote source into memory and return this array.
   1060 
   (...)
   1073     dask.compute
   1074     """
-> 1075     ds = self._to_temp_dataset().load(**kwargs)
   1076     new = self._from_temp_dataset(ds)
   1077     self._variable = new._variable

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/xarray/core/dataset.py:747, in Dataset.load(self, **kwargs)
    744 import dask.array as da
    746 # evaluate all the dask arrays simultaneously
--> 747 evaluated_data = da.compute(*lazy_data.values(), **kwargs)
    749 for k, data in zip(lazy_data, evaluated_data):
    750     self.variables[k].data = data

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/base.py:599, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    596     keys.append(x.__dask_keys__())
    597     postcomputes.append(x.__dask_postcompute__())
--> 599 results = schedule(dsk, keys, **kwargs)
    600 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/threaded.py:89, in get(dsk, keys, cache, num_workers, pool, **kwargs)
     86     elif isinstance(pool, multiprocessing.pool.Pool):
     87         pool = MultiprocessingPoolExecutor(pool)
---> 89 results = get_async(
     90     pool.submit,
     91     pool._max_workers,
     92     dsk,
     93     keys,
     94     cache=cache,
     95     get_id=_thread_get_id,
     96     pack_exception=pack_exception,
     97     **kwargs,
     98 )
    100 # Cleanup pools associated to dead threads
    101 with pools_lock:

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/local.py:511, in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    509         _execute_task(task, data)  # Re-execute locally
    510     else:
--> 511         raise_exception(exc, tb)
    512 res, worker_id = loads(res_info)
    513 state["cache"][key] = res

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/local.py:319, in reraise(exc, tb)
    317 if exc.__traceback__ is not tb:
    318     raise exc.with_traceback(tb)
--> 319 raise exc

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/local.py:224, in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    222 try:
    223     task, data = loads(task_info)
--> 224     result = _execute_task(task, data)
    225     id = get_id()
    226     result = dumps((result, id))

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/core.py:119, in _execute_task(arg, cache, dsk)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/xarray/core/parallel.py:277, in map_blocks.<locals>._wrapper(func, args, kwargs, arg_is_array, expected)
    267 """
    268 Wrapper function that receives datasets in args; converts to dataarrays when necessary;
    269 passes these to the user function `func` and checks returned objects for expected shapes/sizes/etc.
    270 """
    272 converted_args = [
    273     dataset_to_dataarray(arg) if is_array else arg
    274     for is_array, arg in zip(arg_is_array, args)
    275 ]
--> 277 result = func(*converted_args, **kwargs)
    279 # check all dims are present
    280 missing_dimensions = set(expected["shapes"]) - set(result.sizes)

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/metpy/xarray.py:1284, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1282 @functools.wraps(func)
   1283 def wrapper(*args, **kwargs):
-> 1284     bound_args = sig.bind(*args, **kwargs)
   1286     # Auto-broadcast select xarray arguments, and update bound_args
   1287     if broadcast is not None:

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/inspect.py:3185, in Signature.bind(self, *args, **kwargs)
   3180 def bind(self, /, *args, **kwargs):
   3181     """Get a BoundArguments object, that maps the passed `args`
   3182     and `kwargs` to the function's signature.  Raises `TypeError`
   3183     if the passed arguments can not be bound.
   3184     """
-> 3185     return self._bind(args, kwargs)

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/inspect.py:3111, in Signature._bind(self, args, kwargs, partial)
   3107 else:
   3108     if param.kind in (_VAR_KEYWORD, _KEYWORD_ONLY):
   3109         # Looks like we have no parameter for this positional
   3110         # argument
-> 3111         raise TypeError(
   3112             'too many positional arguments') from None
   3114     if param.kind == _VAR_POSITIONAL:
   3115         # We have an '*args'-like argument, let's fill it with
   3116         # all positional arguments we have left and move on to
   3117         # the next phase
   3118         values = [arg_val]

TypeError: too many positional arguments

I don't quite understand the nuances of xr_map_blocks by passing the Dataset as obj, vs. just calling ds.map_blocks directly? I know this is deviating towards Xarray questions too, so I can hop over there too if that would be more useful.

[dopplershift]

Ack, I messed up. Try this one:

xr.map_blocks(mpcalc.precipitable_water, gfs_ds['P3D'],\
              args=[,gfs_ds['TD']],
              kwargs={'bottom':(1000.0*units('hectopascal')),'top':(100.0*units('hectopascal'))},
              template=gfs_ds['TD'].isel(isobaric=0)).compute()

Essentially I'm trying to get the Dataset out of the equation since MetPy doesn't know what to do with it. That's the benefit of moving away from the Dataset.map_blocks() method in favor of the xarray.map_blocks() function. The former assumes (since you're calling it through a Dataset) that you want to be working with it.

[DanielAdriaansen]

OK thanks! This makes some sense now. One question regarding MetPy and Xarray- I see there is preprocess_and_wrap, has there ever been discussion of postprocess_and_wrap such that if an Xarray object is passed in, an Xarray object is returned out (i.e. Pint array inside Xarray DataArray, then return a pint array inside a DataArray rather than a straight pint array)? I see related or adjacent topics in #1601 and #2364, so maybe I could comment on one of those.

I tried hacking this to get it to work but really got bogged down deep in Xarray and Dask. First, inside precipitable_water() since I am using model (and not sounding) data, I just commented out these lines (they failed otherwise) since I know pressure is ordered the way it should be:

MetPy/src/metpy/calc/indices.py

Lines 62 to 65 in 66c01eb

	# Sort pressure and dewpoint to be in decreasing pressure order (increasing height)
	sort_inds = np.argsort(pressure)[::-1]
	pressure = pressure[sort_inds]
	dewpoint = dewpoint[sort_inds]

Then, map_blocks() inside Xarray requires that a Xarray object is returned from func (precipitable_water()), so I hacked in a DataArray creation from the pint object returned from MetPy something like this:

result = DataArray(func(*converted_args, **kwargs).m,coords=converted_args[0].coords,dims=converted_args[0].dims)

But then I got some weird error in Dask, where something is going awry with the chunks:

File ~/opt/miniconda3/envs/pagkickoff/lib/python3.10/site-packages/dask/array/core.py:5317, in concatenate3(arrays)
   5315         while arr.ndim < ndim:
   5316             arr = arr[None, ...]
-> 5317     result[idx] = arr
   5319 return result

ValueError: could not broadcast input array from shape (41,50,50) into shape (41,50)

This somehow has to do with the result needing to have shape (50,50) and not (41,50,50) like the input (the input chunk is 3D and the output chunk reduces to 2D). I am providing template to map_blocks(), but the problem here is that the chunks of the template are actually (50,50), not (41,50,50) but dask doesn't actually do anything with the template as far as I can tell, it's only used in xr.map_blocks().

It seems like not a unique situation that a function may return data with a different shape, but perhaps dask really isn't designed for that situation, I'm not entirely sure there. Either way, I may have gone as far as I intend to take this, although I still feel that there is a fundamental question here about xr.map_blocks() and when the template shape does not match the input shape, how to handle that in the chunking/etc. that's passed on to dask from inside Xarray. I'll report back if I get anywhere in those channels. Thanks for your help @dopplershift!

[dopplershift]

I forgot about #1533 (which you opened originally) where precipitable_water doesn't work on grids due to get_layer, so this is currently doomed to failure.

In general, part of preprocess_and_wrap's mission in life is to handle returning xarray data when it's supplied. If you pass wrap_like to preprocess_and_wrap(), the decorated function will make sure the returned data matches the type of the given data--and for xarray it uses this source to get coordinate information since at the lowest level MetPy isn't exploiting data arrays directly. Therein lies the problem, in that it becomes problematic to match the coordinates after some dimensionally reductive operation (e.g. the vertical integration in precipitable_water). It's certainly something we want to support better, though.

[DanielAdriaansen]

Yup! I'm still busy with precipitable_water even since #1533. Since I opened that, I hacked together a multiprocessing approach in #2259 but then as I have been learning more about advanced Dask and Xarray features lately it seemed like xr.map_blocks() on chunked data was similar to my approach in #2259 but I do think there's some difference, as you note almost certainly related to get_layer. Thanks for playing along with me as I explored this approach, I do appreciate your help!

And thanks for the details regarding preprocess_and_wrap. There's a lot of details here I could stand to learn more so that's really helpful. Again, really appreciate the engagement as always!