Different answers using .sum() within and outside map_blocks

[tlsw231]

Hi all,
I am trying to use map_blocks to sum up an array along a specified dimension. The data is from a global climate simulation. I created a function to add up an xarray dataarray along dimensions specified by the user. When I try this function within map_blocks on a toy example, things work fine. When I try the same function on the climate model output the map_blocks step goes through without any errors. The output, however, is quite different from that obtained using just the .sum() function outside map_blocks. I have taken some of the relevant output from my notebook and put them into an html file at this link below. Has anybody seen this problem before in their calculations?

https://www.dropbox.com/s/a2kmr8l0tbumrr6/debug_mapblocks.html?dl=0

Thanks.

PS: There was some issue with the link, have fixed it now.

[tlsw231]

With help from a colleague, I figured what I was doing wrong. Because map_blocks operates on chunk by chunk, it was not aggregating the sum across the entire record along the specified dimension but only printing the value for the last chunk.

I have some very basic questions about how to use map_blocks. Here is some simple code:

import xarray as xr

nt=4
da=xr.DataArray(np.arange(nt),coords={'t':np.arange(nt)},dims=['t'])

def getsum_da(da,sumdims):
    sumda=da.sum(dim=sumdims,skipna=True)
    return sumda

da.sum(dim='t').compute() # prints 6 = 0 + 1 + 2 + 3

result=xr.map_blocks(getsum_da,da.chunk(chunks={'t':-1}),args=['t']) # no chunking along summation index
result.compute() # prints 6

result=xr.map_blocks(getsum_da,da.chunk(chunks={'t':1}),args=['t']) # with chunking along summation index
result.compute() # prints 3, the value for the last of 4 chunks

Suppose I was trying to do the same thing but with a dataset of size (24,62,2400,3600) where t is the first dimension and is equal to the number of months (24 months = 2 years). That is, I want to sum over t and obtain a 3d array of size (62,2400,3600). These numbers are closer to what I am working with. As in the toy example above, forcing xarray not to chunk along t gives me the correct answer which agrees with what I get without using map_blocks (in the code below, da is the 4d array):

result=xr.map_blocks(getsum_da,da.chunk(chunks={'t':-1}),args=['t']) # no chunking along summation index

Now onto my question. I have decades worth of data (>20 years) and so my dataArray size is at least(240,62,2400,3600). The command above no longer works as I run out of memory trying to use map_blocks. Without chunking along t, the summation index, the chunk size keeps increasing with the addition of more files until eventually the memory requirements become too high. I tried making the chunks along the non-summation indices smaller but that does not help.

How do I continue to use map_blocks for such large datasets?

[dcherian]

Did you try using smaller ~100-200MB chunk sizes , in particular providing these when reading in the data?

[tlsw231]

Thanks, just adding chunks='auto' to the open_mfdataset step makes a massive difference! Earlier, I was rechunking with chunks='auto' after reading in the data which was also giving me around 100 mb size chunks. But the performance was much poorer. A follow-up question:

I am computing the sum for many 4d arrays in succession. I persist one calculation, wait for it to complete and then move to the next variable and so on. There is no spillage to disk for the earlier variables but for the later ones, the calculation starts to spill to disk (6-8 gb). The spillage is still much less than I was seeing before (tens of gb) when the calculation would not even run. Is there any way to prevent this spillage to disk? Or is some amount of spillage inevitable for such calculations?

[update: came across this related thread also reporting very different performance depending on when you chunk the dataset.]

[dcherian]

chunks='auto' after reading in the data which was also giving me around 100 mb size chunks.

Yeah this loads the entire file then breaks it up into smaller chunks (I think).

I am computing the sum for many 4d arrays in succession. I persist one calculation,

Why are you persisting? Persist saves to distributed RAM; if you don't "cancel" the persisted datasets, they just use up memory which is why your later calculations spill.

What is your end-goal with the reduced datasets?

[tlsw231]

I thought the alternative to .persist() would be to use .compute(). Wouldn't that make things worse? The end-goal of my analysis is to go from the 4d array to globally averaged vertical profiles that are 1d darrays. I am doing it in this sequence:

1. Compute the time-mean to reduce the 4d to 3d arrays.
2. Save all the 3d arrays to disk because step 1 is expensive and I do not want to repeat it every time I need to plot something.
3. Compute the lateral mean of the 3d array over lon/lat to obtain the 1d array.

How does one "cancel" a persisted dataset and release that memory, assuming it is even possible?

[tlsw231]

Actually, I do some more xgcm-type operations on the 3d arrays between steps 1 and 3 before I compute the lateral mean of those derived variables.