-
-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Closed
Labels
Description
Discussed in #6762
Originally posted by tlsw231 July 7, 2022
I am trying to use map_blocks
to: ingest a multi-dimensional array as input, reduce along one dimension and add extra dimensions to the output. Is this possible? I am attaching a simple MRE below that gives me an zip argument #2 must support iteration
error. Any pointers on what I might be doing wrong?
[My real example is a 3d-dataset with (time,lat,lon)
dimensions and I am trying to reduce along time
while adding two new dimensions to the output. I tried so many things and got so many errors, including the one in the title, that I thought it is better to first understand how map_blocks
works!]
# The goal is to feed in a 2d array, reduce along one dimension and add two new dimensions to the output.
chunks={}
dummy = xr.DataArray(data=np.random.random([8,100]),dims=['dim1','dim2']).chunk(chunks)
def some_func(func):
dims=func.dims
n1 = len(func[func.dims[1]]) # This is 'dim2', we will average along 'dim1' below in the for loop
newdim1 = 2; newdim2 = 5;
output = xr.DataArray(np.nan*np.ones([n1,newdim1,newdim2]),dims=[dims[1],'new1','new2'])
for n in range(n1):
fmean = func.isel(dim2=n).mean(dims[0]).compute()
for i in range(newdim1):
for j in range(newdim2):
output[n,i,j] = fmean
return output
#out = some_func(dummy) # This works
template=xr.DataArray(np.nan*np.ones([len(dummy.dim2),2,5]),
dims=['dim2','new1','new2'])
out = xr.map_blocks(some_func,dummy,template=template).compute() # gives me the error message in the title
[Edit: Fixed a typo in the n1 = len(func[func.dims[1]])
line, of course getting the same error.]