Need to do a lowpass filter, best way to use xarray effeciently?

[elhoff00]

I'm sort of new to xarray, and most of the things I've had to do I've been able to do with built in functions that work easily with dask. Right now, however, I need to use a lowpass filter on a large dataset, dsl (a few decades of daily time steps, lat, lon). Below is code I think works the way I want it to on smaller datasets, but here it crashes my kernel due to the size of dsl. I tried apply_ufunc, but it's new to me and I'm getting an error I'm not sure how to resolve, also below.

fs = 1
cut_off = 1/60
order = 6

nyq = 0.5*fs
norm = cut_off/nyq
b, a = butter(order, norm, btype='lowpass')

# lrclo = lfilter( b, a, dsl, 0)  # crashes

# lrclo = xr.apply_ufunc(lfilter, b, a, dsl, kwargs={'axis': 0}, dask='parallelized')  # fails on loopdim2
lrclo = xr.apply_ufunc(lfilter, b, a, dsl, input_core_dims=[[],[],['time']], kwargs={'axis': -1}, dask='parallelized')  # fails on loopdim1

ValueError Traceback (most recent call last)
Cell In[17], line 12
9 b, a = butter(order, norm, btype='lowpass')
11 # lrclo = lfilter( b, a, dsl, dim=0)
---> 12 lrclo = xr.apply_ufunc(lfilter, b, a, dsl, input_core_dims=[[],[],['time']], kwargs={'axis': -1}, dask='parallelized')
13 lrclo

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/xarray/core/computation.py:1197, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1195 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1196 elif any(isinstance(a, DataArray) for a in args):
-> 1197 return apply_dataarray_vfunc(
1198 variables_vfunc,
1199 *args,
1200 signature=signature,
1201 join=join,
1202 exclude_dims=exclude_dims,
1203 keep_attrs=keep_attrs,
1204 )
1205 # feed Variables directly through apply_variable_ufunc
1206 elif any(isinstance(a, Variable) for a in args):

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/xarray/core/computation.py:304, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
299 result_coords, result_indexes = build_output_coords_and_indexes(
300 args, signature, exclude_dims, combine_attrs=keep_attrs
301 )
303 data_vars = [getattr(a, "variable", a) for a in args]
--> 304 result_var = func(*data_vars)
306 out: tuple[DataArray, ...] | DataArray
307 if signature.num_outputs > 1:

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/xarray/core/computation.py:761, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
756 if vectorize:
757 func = _vectorize(
758 func, signature, output_dtypes=output_dtypes, exclude_dims=exclude_dims
759 )
--> 761 result_data = func(*input_data)
763 if signature.num_outputs == 1:
764 result_data = (result_data,)

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/xarray/core/computation.py:738, in apply_variable_ufunc..func(*arrays)
737 def func(*arrays):
--> 738 res = chunkmanager.apply_gufunc(
739 numpy_func,
740 signature.to_gufunc_string(exclude_dims),
741 *arrays,
742 vectorize=vectorize,
743 output_dtypes=output_dtypes,
744 **dask_gufunc_kwargs,
745 )
747 return res

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/xarray/core/daskmanager.py:117, in DaskManager.apply_gufunc(self, func, signature, axes, axis, keepdims, output_dtypes, output_sizes, vectorize, allow_rechunk, meta, *args, **kwargs)
100 def apply_gufunc(
101 self,
102 func: Callable,
(...)
113 **kwargs,
114 ):
115 from dask.array.gufunc import apply_gufunc
--> 117 return apply_gufunc(
118 func,
119 signature,
120 *args,
121 axes=axes,
122 axis=axis,
123 keepdims=keepdims,
124 output_dtypes=output_dtypes,
125 output_sizes=output_sizes,
126 vectorize=vectorize,
127 allow_rechunk=allow_rechunk,
128 meta=meta,
129 **kwargs,
130 )

File ~/.conda/envs/TPA_O2_env/lib/python3.12/site-packages/dask/array/gufunc.py:431, in apply_gufunc(func, signature, axes, axis, keepdims, output_dtypes, output_sizes, vectorize, allow_rechunk, meta, *args, **kwargs)
428 for dim, sizes in dimsizess.items():
429 #### Check that the arrays have same length for same dimensions or dimension 1
430 if set(sizes) | {1} != {1, max(sizes)}:
--> 431 raise ValueError(f"Dimension '{dim}' with different lengths in arrays")
432 if not allow_rechunk:
433 chunksizes = chunksizess[dim]

ValueError: Dimension '__loopdim1__' with different lengths in arrays

If this has been answered before or should be asked somewhere else please let me know. Some google searching and browsing this discussion board hasn't helped me so far, on both the general approach and my specific errors. I did find the xr-scipy package but it's also creating errors I can't figure out how to resolve with its existing documentation.

[elhoff00]

Ok, so winning setting combination:

lrclo = xr.apply_ufunc(lfilter, b, a, dsl, 
                       input_core_dims=[[],[],['time']], 
                       output_core_dims=[['time']],
                       kwargs={'axis': -1}, 
                       dask='allowed')

dask='allowed' instead of 'parallelize' got rid of the error. If I used vectorize=True (something I thought might be an issue) it crashed, without it it only almost crashed (got up to ~31.5 of my 32 GB server in the middle of running). If anyone has thoughts on how to improve this code, however, please chime in. I still have to run this with a higher resolution product (1/10 degree rather than 1 degree resolution, many more lats and lons) and I can try on a bigger server, but not 100x bigger :/