使用xarray map_blocks处理维度变化时遇错求助

问题背景 #

我有一个维度为('time', 'lat', 'lon')的xarray.Dataset，需要用xr.map_blocks对30x30的网格块计算特定属性，预期输出维度为('block', 'location', 'time')。

最小复现代码 #

数据与输出模板定义 #

import numpy as np
import xarray as xr

data = xr.Dataset(
    data_vars = {'rain': (['time','lat','lon'], np.random.rand(12, 180, 360))}, 
    coords = {'time': np.arange(12), 'lat': np.arange(90,-90,-1)-0.5, 'lon': np.arange(-180,180)+0.5}
).chunk({'time': -1, 'lat': 30, 'lon': 30})

n_blocks = np.prod([len(dim_chunks) for dim, dim_chunks in data.chunks.items()])
output_template = xr.Dataset(
    data_vars = {'rain': (['block','location','time'], np.empty(shape=(n_blocks, 20, 12)))}, 
    coords = {'time': np.arange(12), 'location': np.arange(20), 'block': np.arange(n_blocks)}
).chunk({'block': 1, 'time': -1, 'location': -1})

处理函数与调用 #

def group_fn(block):
    return xr.Dataset(
        data_vars = {'rain': (['time','location'], np.random.rand(12, 20))}, 
        coords = {'time': np.arange(12), 'location': np.arange(20)}
    )
mapped = xr.map_blocks(group_fn, data, template=output_template)
computed = mapped.compute()

首次错误信息 #

运行后抛出ValueError：

ValueError                                Traceback (most recent call last)
/tmp/ipykernel_20583/1876666684.py in <module>
     18 
     19 mapped = xr.map_blocks(group_fn, data, template=output_template)
---> 20 computed = mapped.compute()

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/dataset.py in compute(self,** kwargs)
   1027         new = self.copy(deep=False)
-> 1028         return new.load(**kwargs)
   1029 
   1030     def _persist_inplace(self,** kwargs) -> "Dataset":

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/dataset.py in load(self,** kwargs)
    860 
    861             # evaluate all the dask arrays simultaneously
-> 862             evaluated_data = da.compute(*lazy_data.values(), **kwargs)
    863 
    864             for k, data in zip(lazy_data, evaluated_data):

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/base.py in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    571         postcomputes.append(x.__dask_postcompute__())
    572 
-> 573         results = schedule(dsk, keys, **kwargs)
    574         return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    575 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     87         get_id=_thread_get_id,
     88         pack_exception=pack_exception,
-> 89         **kwargs,
     90     )
     91 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    504                             _execute_task(task, data)  # Re-execute locally
    505                         else:
-> 506                             raise_exception(exc, tb)
    507                     res, worker_id = loads(res_info)
    508                     state["cache"][key] = res

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in reraise(exc, tb)
    312     if exc.__traceback__ is not tb:
    313         raise exc.with_traceback(tb)
-> 314     raise exc
    315 
    316 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    217     try:
    218         task, data = loads(task_info)
-> 219         result = _execute_task(task, data)
    220         id = get_id()
    221         result = dumps((result, id))

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/core.py in <genexpr>(.0)
    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

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/parallel.py in _wrapper(func, args, kwargs, arg_is_array, expected)
    312 
    313         # check that all expected variables were returned
-> 314         check_result_variables(result, expected, "coords")
    315         if isinstance(result, Dataset):
    316             check_result_variables(result, expected, "data_vars")

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/parallel.py in check_result_variables(result, expected, kind)
     66     if missing:
     67         raise ValueError(
-> 68             "Result from applying user function does not contain "
     69             f"{nice_str} variables {missing}."
     70         )

ValueError: Result from applying user function does not contain coordinate variables {'block'}.

尝试修改后的代码及新错误 #

按照提示在函数中添加block坐标：

def group_fn(block):
    return xr.Dataset(
        data_vars = {'rain': (['time','location'], np.random.rand(12, 20))},
        coords = {'time': np.arange(12), 'location': np.arange(20), 'block': np.arange(n_blocks)}
    )
mapped = xr.map_blocks(group_fn, data, template=output_template)
computed = mapped.compute()

运行后抛出KeyError：

KeyError                                  Traceback (most recent call last)
/tmp/ipykernel_20583/908260062.py in <module>
     18 
     19 mapped = xr.map_blocks(group_fn, data, template=output_template)
---> 20 computed = mapped.compute()

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/dataset.py in compute(self,** kwargs)
   1027         new = self.copy(deep=False)
-> 1028         return new.load(**kwargs)
   1029 
   1030     def _persist_inplace(self,** kwargs) -> "Dataset":

~/anaconda3/envs/science/lib/python3.7/site-packages/xarray/core/dataset.py in load(self,** kwargs)
    860 
    861             # evaluate all the dask arrays simultaneously
-> 862             evaluated_data = da.compute(*lazy_data.values(), **kwargs)
    863 
    864             for k, data in zip(lazy_data, evaluated_data):

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/base.py in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    571         postcomputes.append(x.__dask_postcompute__())
    572 
-> 573         results = schedule(dsk, keys, **kwargs)
    574         return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    575 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool,** kwargs)
     87         get_id=_thread_get_id,
     88         pack_exception=pack_exception,
-> 89         **kwargs,
     90     )
     91 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    518                     finish(dsk, state, not succeeded)
    519 
-> 520         return nested_get(result, state["cache"])
    521 
    522 

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in nested_get(ind, coll)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in <genexpr>(.0)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in nested_get(ind, coll)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in <genexpr>(.0)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in nested_get(ind, coll)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in <genexpr>(.0)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in nested_get(ind, coll)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in <genexpr>(.0)
    296     if isinstance(ind, list):
-> 297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
    299         return coll[ind]

~/anaconda3/envs/science/lib/python3.7/site-packages/dask/local.py in nested_get(ind, coll)
    297         return tuple(nested_get(i, coll) for i in ind)
    298     else:
-> 299         return coll[ind]
    300 
    301 

KeyError: ('rain-group_fn-1a1e7d36744015031d014d432cc570f6', 1, 0, 0)

补充说明 #

简化问题后（移除location维度，输入维度('time', 'lat', 'lon')，输出维度('block', 'time')），仍会出现相同错误。

内容的提问来源于stack exchange，提问作者Agustín Begue