在列表中查找距离最小的 N 个最大元素

Question

I want to extract from a list the N largest elements, but I want that for any two elements x[i] and x[j] , abs(ij) > min_distance .

scipy.signal.find_peaks(x, distance=min_distance) offers this functionality. However I need to repeat this operation millions of times and I was trying to speed up a bit the operation.

I noticed that find_peaks does not accept a parameter N to indicate how many peaks you want to extract. It also doesn't allow to return the peaks from largest to smallest, requiring an additional call to l.sort() and l = l[:N] .

I tried to code a lazy sorter that only looks for the N largest elements without sorting the rest of the list.

Following the results obtained here I opted for heapq . This is my try:

import heapq

def new_find_peaks(x, N, min_distance=0):
    x = enumerate(x)

    x = [(-val,i) for (i,val) in x]
    heapq.heapify(x)

    val, pos = heapq.heappop(x)
    peaks = [(-val, pos,)]

    while len(peaks)<N:

        while True:
            val, pos = heapq.heappop(x)
            d = min([abs(pos - pos_i) for _,pos_i in peaks])
            if d >= min_distance:
                break

        peaks.append((-val, pos,))

    return map(list, zip(*peaks)) #Transpose peaks into 2 lists

This is, however, still 20 times slower than find_peaks , probably due to find_peaks CPython implementation. Also, I noticed that almost half of the time is spent on

x = [(-val,i) for (i,val) in x]

Do you have any better idea to speed up this operation?

--- minimal reproducible example ---

For example:

x = [-8.11, -7.33, -7.48, -5.77, -8.73, -8.73, -7.02, -7.02,
 -7.80, -10.92, -9.36, -9.83, -10.14, -10.77, -11.23, -9.20,
 -9.52, -9.67, -11.23, -9.98, -7.95, -9.83, -8.89, -7.33,
 -4.20, -4.05, -6.70, -7.02, -9.20, -9.21]

new_find_peaks(x, N=3, min_distance=5)

>> [[-4.05, -5.77, -7.8], [25, 3, 8]]

Note that x[24] is -4.2, but since x[25] is greater and 25-24 < min_distance , this is discarded. Also note that x[8] is not a real peak, as x[7] is bigger, but this is discarded due to the distance with x[3] . This is an intended behavior.

Answer 1

Improving your code in python could possibly give you some improvement, but as your code seems clean and the algorithm's idea sound, I don't think you will beat find_peaks with a python approach.

Hence I suggest you write your own library in a language that is closer to the metal, and write your own python wrapper if you need the result in python. For instance, you could use Swift. Here is an implementation of the heap queue in Swift and here you find described a way to interface with python.

Connecting the dots is left as an exercise. ;)

Answer 2

In order to speed up the calculations, I wrote a second algorithm that does not use heapq. This avoid the whole list to be reshaped in a heap queue.

The new algorithm looks like this

from heapq import nlargest

def find_peaks_fast(x, N, min_distance=0):

    peaks = []
    last_i = 0
    last_peak = x[0]

    for i, val in enumerate(x[1:], 1):

        if i - last_i == min_distance:
            # Store peak
            peaks.append(last_peak)

            # Store the new item and move on
            last_peak = val
            last_i = i

        elif val > last_peak:
            last_peak = val
            last_i = i

    return nlargest(N,peaks)

The algorithm scans through the list once and extract all the samples that are higher of both the N samples before and the N samples after. These are then stored in a list from where only the nlargest elements are extracted using heapq.nlargest

This, on its own, brings the execution time down to 3.7 ms. Quick, but still almost 4 times slower than scipy's find_peaks.

However, this can be changes using the package numba . This aims to "compile" the python code on-the-fly and to execute a compiled version to improve speed. And it improves it a lot !!

from numba import njit
from heapq import nlargest

from numba.errors import NumbaPendingDeprecationWarning
import warnings

warnings.simplefilter('ignore', category=NumbaPendingDeprecationWarning)

@njit
def find_peaks_fast(x, N, min_distance=0):

    peaks = []
    last_i = 0
    last_peak = x[0]

    for i, val in enumerate(x[1:], 1):

        if i - last_i == min_distance:
            # Store peak
            peaks.append(last_peak)

            # Store the new item and move on
            last_peak = val
        last_i = i

    elif val > last_peak:
        last_peak = val
        last_i = i

return nlargest(N,peaks)

And testing it

from scipy.signal import find_peaks

from timeit import repeat
from numpy.random import randn

from numba import njit
from heapq import nlargest


def new_find_peaks2(x, N, min_distance=0):

    peaks = []
    last_i = 0
    last_peak = x[0]

    for i, val in enumerate(x[1:], 1):

        if i - last_i == min_distance:
            # Store peak
            peaks.append(last_peak)

            # Store the new item and move on
            last_peak = val
            last_i = i

        elif val > last_peak:
            last_peak = val
            last_i = i

    return nlargest(N,peaks)

@njit
def new_find_peaks2_jit(x, N, min_distance=0):

    peaks = []
    last_i = 0
    last_peak = x[0]

    for i, val in enumerate(x[1:], 1):

        if i - last_i == min_distance:
            # Store peak
            peaks.append(last_peak)

            # Store the new item and move on
            last_peak = val
            last_i = i

        elif val > last_peak:
            last_peak = val
            last_i = i

    return nlargest(N,peaks)


num = 500
rep = 10

N = 20
x = randn(20000)
sep = 10

code1 = '''
i_pks, _  = find_peaks(x, distance=sep)
pks = x[i_pks]
pks[::-1].sort()
pks = pks[:N]
'''

code2 = '''
_ = new_find_peaks2(x, N=N, min_distance=sep)
'''

code2_jit = '''
_ = new_find_peaks2_jit(x, N=N, min_distance=sep)
'''

i_pks, _  = find_peaks(x, distance=sep)
pks = x[i_pks]
pks[::-1].sort()
pks1 = pks[:N]
pks2 = new_find_peaks2(x, N=N, min_distance=sep)

print(pks1==pks2)

t = min(repeat(stmt=code1, globals=globals(), number=num, repeat=rep))/num
print(f'np.find_peaks:\t\t{t*1000} [ms]')

t = min(repeat(stmt=code2, globals=globals(), number=num, repeat=rep))/num
print(f'new_find_peaks2:\t{t*1000} [ms]')

t = min(repeat(stmt=code2_jit, globals=globals(), number=num, repeat=rep))/num
print(f'new_find_peaks2_jit:\t{t*1000} [ms]')

Leads to the result:

np.find_peaks:          1.1234994470141828 [ms]
new_find_peaks2:      3.565517600043677 [ms]
new_find_peaks2_jit:  0.10387242998695001 [ms]

That is a x10 speed up!

Conclusions:

• The algorithm could be sped up
• numba.njit proved to be an incredible wrapper, increasing the execution of the same function of almost 35 times!