'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()
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
| Solution | Source |
|---|

