xarray apply_ufunc: Could not broadcast input array from shape (2,) into shape (1,)

[ZZMitch]

Hello,

For the first time I am trying to use apply_ufunc to create my own xarry function. In this problem, I want to be able to apply find_peaks from scipy.signal (link) for each pixel across time in my 3D dataArray and return the number of identified peaks. The example dataArray looks like this:

I wrote a simple findPeaks function:

import xarray as xr
from scipy.signal import find_peaks

def findPeaks(arr):
    peaks, _ = find_peaks(arr) # More parameters here to specify types of peaks to find later
    return peaks

This returns the time index of any identified peaks. I can then get the length of the output to find the number of peaks. This works fine for a single pixel, since find_peaks expects a 1D array. However, when I run it on dataArray, if there number of peaks goes above 1 I get an error:

result = xr.apply_ufunc(findPeaks, dataArray, input_core_dims = [['time']], output_core_dims = [['peaks']], vectorize = True).count(dim = 'peaks')
result

ValueError: could not broadcast input array from shape (2,) into shape (1,)

In this case, findPeaks found 2 peaks, and I would hope the output value would be 2 for the given pixel. But I guess the way I have the code set up, there is an issue where apply_unfunc cannot accept an array of length > 1 (which I would hope to get the count of).

Any help would be appreciated! As I mentioned, this is my first time trying apply_unfunc and so I am just learning how to get parameters right to make these conversions from numpy functions to 3D xarray dataArrays.

Thanks!