diff --git a/maxima.py b/maxima.py index ed03ef6..6a3e500 100644 --- a/maxima.py +++ b/maxima.py @@ -1,3 +1,5 @@ +import numpy as np + def find_maxima(x): """Find local maxima of x. @@ -14,8 +16,26 @@ def find_maxima(x): """ idx = [] + bfr = [] + up = True for i in range(len(x)): # `i` is a local maximum if the signal decreases before and after it - if x[i-1] < x[i] and x[i+1] < x[i]: - idx.append(i) + + # starting or going up + if i == 0 or x[i] > x[i-1]: + up = True + bfr = [i] + # plateau + elif x[i] == x[i-1]: + if up: + bfr.append(i) + # going down + elif x[i] < x[i-1]: + up = False + idx += bfr + bfr = [] + # end + if i == len(x) - 1: + if up: + idx += bfr return idx diff --git a/test_maxima.py b/test_maxima.py new file mode 100644 index 0000000..9bc67d5 --- /dev/null +++ b/test_maxima.py @@ -0,0 +1,28 @@ +from maxima import find_maxima +import numpy as np + +def test_1(): + assert find_maxima([0, 1, 2, 1, 2, 1, 0]) == [2, 4] + +def test_2(): + assert find_maxima([2,1,0,1,4]) == [0, 4] + +def test_3(): + assert find_maxima([np.sin(2*alpha) for alpha in np.linspace(0.0, 5.0, 100)]) == [16, 78] + +def test_4(): + assert find_maxima([1, 2, 2, 3, 1]) == [3] + + +def test_5(): + assert find_maxima([1, 3, 2, 2, 1]) == [1] + + +def test_6(): + assert find_maxima([3, 2, 2, 3]) == [0, 3] + +def test_7(): + assert find_maxima([1, 2, 2, 1]) == [1, 2] + +def test_8(): + assert find_maxima([1, 2, 2, 2]) == [1, 2, 3]