'scipy.signal.detrend within apply_ufunc failing for large array sizes

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'])

# Rechunk to unchunk along core dimension and to get ~100mb 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()

Validation of detrending calculation



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source