Memory errors with detrending + rolling operations

[tlsw231]

I am trying to perform these two operations in succession on a 3d array with dimensions (t,y,x):
(i) detrend along t
(ii) rolling mean along t
I get memory related failures if the size along the t dimension becomes too large even though I think I have sufficient memory. That is, the progress bars in the dask dashboard show many "erred" processes and the kernel dies midway. Here is some stand-alone code that mimics what I am trying to with comparable array sizes:

# Import a bunch of stuff
import sys
import os 

import math
import xarray as xr
import numpy as np
import dask
from dask.distributed import Client
from dask_jobqueue import SLURMCluster

# Launch the cluster: 8 cores, 340 Gb
cluster = SLURMCluster(cores=8, processes=8, memory="340GB",queue="short",
                       walltime="01:00:00")
print(cluster.job_script())
cluster.scale(8)
client = Client(cluster)

# Array size. My actual application has nt ~ 12000 but this code fails even for the value below
(nt,ny,nx)=(2000,489,655)

# Define the data array
dummy=xr.DataArray(data=np.random.randn(nt,ny,nx),dims=['t','y','x'])

# The size in Gb
dummy.nbytes/1e9
5.12472

# Function to detrend
# Source: https://gist.github.com/rabernat/1ea82bb067c3273a6166d1b1f77d490f 
def detrend_dim(da, dim, deg=1):
    # detrend along a single dimension
    p = da.polyfit(dim=dim, deg=deg)
    fit = xr.polyval(da[dim], p.polyfit_coefficients)
    return da - fit

# Detrend
chunks='auto'
field = dummy.chunk(chunks)
dim = 't'
field_detrend = detrend_dim(field,dim,deg=1)

# Rolling mean
half_win = 8;
field_roll = field_detrend.chunk('auto').rolling(t=2*half_win+1,center=True).mean(skipna=True)

# -----------Kernel dies while trying to execute this step----------------
# Slice along the time dimension to throw out the NaN values 
# from the rolling operation and persist the result
out = ( field_roll.chunk('auto').isel(t=slice(half_win,-half_win)) ).persist()

Any suggestions on why the above fails? Thanks.

[dcherian]

I think you're running into #5629 . There is a reshape operation that is extremely inefficient. It might work better if your chunk sizes are small along time and large in space.

[dcherian]

I looked in to this a bit: https://github.com/NCAR/esds/blob/fcc6909289b8aa4931ff4d2a996c6d0a64ed2f5a/posts/2022/dask-debug-detrend-1.ipynb

Unfortunately I think the problem is a bad dask regression: dask/dask#6916

[dcherian]

The dask bug is now fixed and things seem to work OK (https://ncar-esds--166.org.readthedocs.build/en/166/posts/2022/dask-debug-detrend-1.html) That new polyval + that blockwise reshape trick should work decently. @tlsw231 this could be a good time to retry it.

[tlsw231]

Thanks so much for pursuing this! I read your earlier message too, apologies for not responding sooner. The link in your latest message is not working for me.

Even without the new reshape routine, I am able to process my data after having made some changes (my dataset is 3-hourly output for 8 years). [I am having some issues with the reshape routine for my dataset even though it works perfectly well for the test cases I created. The routine itself works but the detrended field is having ridiculously large magnitudes like 10^15 etc. But this is something for me to sort out.]. Here is what I am doing presently:

1. Resample everything to daily-averaged values and work with those.
2. I found that using isel to slice in space and time simultaneously was giving me large uneven chunks in time, like what you also see in your notebook that you shared. I found that using isel twice, first on spatial coordinates and then again on the time coordinate removed this problem. The latter option gives me even chunks along all axes without one very large chunk, as was happening earlier.
3. Use the setting below at the start of my notebook, am not sure if this always works as intended.
   dask.config.set(**{'array.slicing.split_large_chunks': True})

If I can sort out my problem with the reshape routine perhaps I can even work directly with the 3-hour data.

[dcherian]

It's here now: https://ncar.github.io/esds/posts/2022/dask-debug-detrend/

[tlsw231]

I hope I am not violating any policies wrt posting here but I had also raised a related query on trying to use apply_ufunc for the detrending. Is there any reason to expect the apply_ufunc option might be faster? [I could not get apply_ufunc to work if I progressively increased the data size, as described in that thread.]