How to apply a 1D median filter to a 3D DataArray using xarray.apply_ufunc()

Question

I have 3-dimensional DataArray (using xarray). I would like to apply a 1-dimensional to it along a certain dimension. Specifically, I want to apply the scipy.signal.medfilt() function but, again, it should be 1-dimensional.

So far I've successfully implemented this the following way:

for sample in data_raw.coords["sample"]:
    for experiment in data_raw.coords["experiment"]:
        data_filtered.loc[sample,experiment,:] = signal.medfilt(data_raw.loc[sample,experiment,:], 15)

(My data array has dimensions "sample", "experiment" and "wave_number. This code applies the filter along the "wave_number" dimension)

The problem with this is that it takes rather long to calculate and my intuition tells me that looping though coordinates like this is an inefficient way to do it. So I'm thinking about using the xarray.apply_ufunc() function, especially since I've used it in a similar fashion in the same code:

xr.apply_ufunc(np.linalg.norm, data, kwargs={"axis": 2}, input_core_dims=[["wave_number"]])

(This calculates the length of the vector along the "wave_number" dimension.)

I originally also had this loop through the coordinates just like the first code here.

The problem is when I try

xr.apply_ufunc(signal.medfilt, data_smooth, kwargs={"kernel_size": 15})

it returns a data array full of zeroes, presumably because it applies a 3D median filter and the data array contains NaN entries. I realize that the problem here is that I need to feed the scipy.signal.medfilt() function a 1D array but unfortunately there is no way to specify an axis along which to apply the filter (unlike numpy.linalg.norm()).

SO, how do I apply a 1D median filter without looping through coordinates?

JulianGiles JulianGiles · Accepted Answer · 2019-08-26T20:03:54

If I understood correctly, you should use it like this:

xr.apply_ufunc(signal.medfilt, data_smooth, kwargs={"kernel_size": 15}, input_core_dims = [['wave_number']], vectorize=True)

with vectorize = True you vectorize your input function to be applied to slices of your array defined to preserve the core dimensions.

Nonetheless, as stated in the documentation:

This option exists for convenience, but is almost always slower than supplying a pre-vectorized function

because the implementation is essentially a for loop. However I still got faster results than by making my own loops.

How to apply a 1D median filter to a 3D DataArray using xarray.apply_ufunc()

1 Answers