Working in xarray, I'm trying to sort the third dimension (wavenumber) of a 6x53x9001 array by mapping a sorting function across the first two dimensions. I've accomplished this by nesting two for loops, but dask's parallel functionality doesn't seem to allow for this type of indexing:
def sort(o_depth):
for i in range(0,6):
for j in range(0,53):
odepth_sorted[i,j] = o_depth[i,j].sortby(o_depth[i,j]).values
return odepth_sorted
odepth_sorted = sort(o_depth)
The sort is dependent on which [i,j] column of 9001 wavenumbers I'm looking at, which is where I'm getting stuck trying to apply/map a function. I can reduce my indices to one "i", which then lets me apply a function as such:
def sort(o_depth):
return o_depth.sortby(o_depth).values
odepth_sorted = o_depth[0].groupby("level").map(sort)
However, I have 5 other i's to perform this over and am still not sure how best to parallelize this type of action with dask - delaying this particular function and running .compute() gives the error "Delayed objects of unspecified length are not iterable". Any thoughts or advice? Thanks in advance.