Detrending using apply_ufunc + scipy.signal.detrend

[tlsw231]

(Edit: I removed the task graph as there is no graph generated for dask='allowed' which is the setting used in the working MVE below.)

[This is related to another active thread I had started but feel like the question here is specific enough to warrant a separate thread. ]

What I am trying to do: Detrend an xarray dataarray using scipy.signal.detrend called within xr.apply_ufunc

What is the problem: The MVE below works for small sizes (10000,1000,3) but fails if I progressively increase the array size (10000,10000,3) even when I have more than enough memory. I get a 'kernel died' message right after the apply_ufunc step is finished.

# Import stuff
import xarray as xr
import numpy as np
import dask
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
import scipy
from scipy import signal

# Launch cluster with 280 gb memory
cluster = SLURMCluster(cores=8, processes=8, memory="280GB",queue="medium",
                       walltime="02:30:00")
print(cluster.job_script())
cluster.scale(8)
client = Client(cluster)

# Create array
data=np.random.random([10000,1000,3])

# Superpose a linear trend at two locations
data[:, 0, 0] = data[:, 0, 0] + np.arange(data.shape[0])
data[:,-1,-1] = data[:,-1,-1] + np.arange(data.shape[0])

# Create xarray dataarray using above array
arr = xr.DataArray(data,dims=['t','y','x'])

# Chunk to give ~110mb chunks
arr = arr.chunk(dict(t=-1,y=500))
# array size
arr.nbytes/1e9
# Output: 0.24

# apply_ufunc call
test_detrend = xr.apply_ufunc(scipy.signal.detrend,arr,kwargs=dict(axis=2),
                              input_core_dims=[['t']],output_core_dims=[['t']],
                              dask='allowed',vectorize=True)

# Plot below shows the detrending works correctly
arr.isel(x=-1,y=-1).plot()
test_detrend.isel(x=-1,y=-1).plot()

[tlsw231]

I tried changing dask='allowed' to dask='parallelized' leaving the other settings in the apply_ufunc call intact:

test_detrend = xr.apply_ufunc(scipy.signal.detrend,arr,kwargs={'axis':2},
                              input_core_dims=[["t"]],output_core_dims=[["t"]],
                              output_dtypes=[arr.dtype],
                              dask='parallelized',vectorize=True)

The compute step test_detrend.compute() fails with the message IndexError: tuple index out of range even for the smaller array size (10000,1000,3) that worked with dask='allowed'.

Here is the task graph:

[tlsw231]

I can get the smaller problem (10000,1000,3) to work with dask='parallelized' by changing kwargs=dict(axis=2) to kwargs=dict(axis= -1) in the apply_ufunc call. Not sure why that makes a difference as axis=2 and axis=-1 should be referring to the same axis ('t') in my example.

For the larger array (10000,10000,3), the test_detrend.compute() step with dask='parallelized' fails, giving me the same error message I got with dask='allowed', i.e., 'kernel died and is restarting'.

So now both dask='allowed' and dask='parallelized' work for the smaller array but fail for the larger array.

[OriolAbril]

It might not be possible (yet?). I have had what I believe are similar issues when wrapping scipy functions and applying them to dask arrays. There are cases where the scipy function tries to load everything into memory, so even if using dask="allowed" seems to work, calling compute inevitably ends up in a memory error and a dead kernel. Similarly, when trying to wrap functions from the einops library with apply_ufunc, dask="parallelized" was the one that seemed to work but ended up failing with strange errors on .compute().

I haven't had time to look into the root of either issue yet, but if you want to play and see for yourself, I have some of the experiments up on this github repo. If you try dask="allowed" in the scipy functions and track memory usage you'll see the whole ram gets used when it should not, and if you try the dask="parallelized" on the rearrange call you should get an error (even though iiuc, everything uses numpy and should work chunk-wise without problem)

[tlsw231]

Thanks for the reply! In my example, I am not sure why even loading everything into memory concurrently should fail as the job size is so small. I will check out that repo you linked to.