From 597c27e615de773594623bef6a283fe0420a28f4 Mon Sep 17 00:00:00 2001
From: Igor Dunaev <idunaev71@gmail.com>
Date: Fri, 3 May 2024 17:04:08 +0300
Subject: [PATCH] Improved peak finding algorithm

---
 .../kscience/kmath/series/peakFinding.kt      | 89 +++++++++++++++++++
 1 file changed, 89 insertions(+)
 create mode 100644 kmath-stat/src/commonMain/kotlin/space/kscience/kmath/series/peakFinding.kt

diff --git a/kmath-stat/src/commonMain/kotlin/space/kscience/kmath/series/peakFinding.kt b/kmath-stat/src/commonMain/kotlin/space/kscience/kmath/series/peakFinding.kt
new file mode 100644
index 000000000..cd340b044
--- /dev/null
+++ b/kmath-stat/src/commonMain/kotlin/space/kscience/kmath/series/peakFinding.kt
@@ -0,0 +1,89 @@
+/*
+ * Copyright 2018-2024 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package space.kscience.kmath.series
+
+
+public enum class PlateauEdgePolicy  {
+    /**
+     * Midpoints of plateau are returned, edges not belonging to a plateau are ignored.
+     *
+     * A midpoint is the index closest to the average of indices of the left and right edges
+     * of the plateau:
+     *
+     * `val midpoint = ((leftEdge + rightEdge) / 2).toInt`
+     */
+     AVERAGE,
+
+    /**
+     * Both left and right edges are returned.
+     */
+    KEEP_ALL_EDGES,
+
+    /**
+     * Only right edges are returned.
+     */
+    KEEP_RIGHT_EDGES,
+
+    /**
+     * Only left edges are returned.
+     */
+    KEEP_LEFT_EDGES,
+
+    /**
+     * Ignore plateau, only peaks (troughs) with values strictly greater (less)
+     * than values of the adjacent points are returned.
+     */
+    IGNORE
+}
+
+
+public fun <T: Comparable<T>> Series<T>.peaks(
+    plateauEdgePolicy: PlateauEdgePolicy = PlateauEdgePolicy.AVERAGE
+): List<Int> = findPeaks(plateauEdgePolicy, { other -> this > other },  { other -> this >= other })
+
+public fun <T: Comparable<T>> Series<T>.troughs(
+    plateauEdgePolicy: PlateauEdgePolicy = PlateauEdgePolicy.AVERAGE
+): List<Int> = findPeaks(plateauEdgePolicy, { other -> this < other }, { other -> this <= other })
+
+
+private fun <T: Comparable<T>> Series<T>.findPeaks(
+    plateauPolicy: PlateauEdgePolicy = PlateauEdgePolicy.AVERAGE,
+    cmpStrong: T.(T) -> Boolean,
+    cmpWeak: T.(T) -> Boolean
+): List<Int> {
+    require(size >= 3) { "Expected series with at least 3 elements, but got $size elements" }
+    if (plateauPolicy == PlateauEdgePolicy.AVERAGE) return peaksWithPlateau(cmpStrong)
+    fun peakCriterion(left: T, middle: T, right: T): Boolean = when(plateauPolicy) {
+        PlateauEdgePolicy.KEEP_LEFT_EDGES -> middle.cmpStrong(left) && middle.cmpWeak(right)
+        PlateauEdgePolicy.KEEP_RIGHT_EDGES -> middle.cmpWeak(left) && middle.cmpStrong(right)
+        PlateauEdgePolicy.KEEP_ALL_EDGES ->
+            (middle.cmpStrong(left) && middle.cmpWeak(right)) || (middle.cmpWeak(left) && middle.cmpStrong(right))
+        else -> middle.cmpStrong(right) && middle.cmpStrong(left)
+    }
+    val indices = mutableListOf<Int>()
+    for (index in 1 .. size - 2) {
+        val left = this[index - 1]
+        val middle = this[index]
+        val right = this[index + 1]
+        if (peakCriterion(left, middle, right)) indices.add(index)
+    }
+    return indices
+}
+
+
+private fun <T: Comparable<T>> Series<T>.peaksWithPlateau(cmpStrong: T.(T) -> Boolean): List<Int> {
+    val peaks = mutableListOf<Int>()
+    tailrec fun peaksPlateauInner(index: Int) {
+        val nextUnequal = (index + 1 ..< size).firstOrNull { this[it] != this[index] } ?: (size - 1)
+        val newIndex = if (this[index].cmpStrong(this[index - 1]) && this[index].cmpStrong(this[nextUnequal])) {
+            peaks.add((index + nextUnequal) / 2)
+            nextUnequal
+        } else index + 1
+        if (newIndex < size - 1) peaksPlateauInner(newIndex)
+    }
+    peaksPlateauInner(1)
+    return peaks
+}
\ No newline at end of file