119 lines
2.3 KiB
Go
119 lines
2.3 KiB
Go
package extractfreq
|
|
|
|
import "gonum.org/v1/gonum/floats"
|
|
|
|
type mmItem struct {
|
|
Val float64
|
|
MaxVal float64
|
|
}
|
|
|
|
type moveMax struct {
|
|
pushStack []mmItem
|
|
popStack []mmItem
|
|
}
|
|
|
|
func newMoveMax() *moveMax {
|
|
return &moveMax{
|
|
pushStack: []mmItem{},
|
|
popStack: []mmItem{},
|
|
}
|
|
}
|
|
|
|
func (mm *moveMax) Max() float64 {
|
|
pushTop := len(mm.pushStack) - 1
|
|
popTop := len(mm.popStack) - 1
|
|
|
|
if pushTop == -1 {
|
|
item := mm.popStack[popTop]
|
|
return item.MaxVal
|
|
} else if popTop == -1 {
|
|
item := mm.pushStack[pushTop]
|
|
return item.MaxVal
|
|
}
|
|
|
|
item1 := mm.pushStack[pushTop]
|
|
item2 := mm.popStack[popTop]
|
|
if item1.MaxVal > item2.MaxVal {
|
|
return item1.MaxVal
|
|
} else {
|
|
return item2.MaxVal
|
|
}
|
|
}
|
|
|
|
func (mm *moveMax) push(stack []mmItem, val float64) []mmItem {
|
|
item := mmItem{
|
|
Val: val,
|
|
MaxVal: val,
|
|
}
|
|
|
|
// Ensure max is set correctly.
|
|
pushTop := len(stack) - 1
|
|
if pushTop > -1 {
|
|
top := stack[pushTop]
|
|
if top.MaxVal > item.MaxVal {
|
|
item.MaxVal = top.MaxVal
|
|
}
|
|
}
|
|
|
|
// Push.
|
|
return append(stack, item)
|
|
}
|
|
|
|
func (mm *moveMax) Push(val float64) {
|
|
mm.pushStack = mm.push(mm.pushStack, val)
|
|
}
|
|
|
|
func (mm *moveMax) Pop() {
|
|
if len(mm.popStack) == 0 {
|
|
// Move push stack onto pop stack.
|
|
for len(mm.pushStack) > 0 {
|
|
item := mm.pushStack[len(mm.pushStack)-1]
|
|
mm.pushStack = mm.pushStack[:len(mm.pushStack)-1]
|
|
mm.popStack = mm.push(mm.popStack, item.Val)
|
|
}
|
|
}
|
|
mm.popStack = mm.popStack[:len(mm.popStack)-1]
|
|
}
|
|
|
|
func getPeaks(nHalfWin int, data []float64) (inds []int, amps []float64) {
|
|
mm := newMoveMax()
|
|
nWin := 2*nHalfWin + 1
|
|
|
|
if len(data) < nWin {
|
|
nWin = len(data)
|
|
}
|
|
|
|
// Push nWin samples.
|
|
for i := range data[:nWin] {
|
|
mm.Push(data[i])
|
|
}
|
|
|
|
for idxPush := nWin + 1; idxPush < len(data); idxPush++ {
|
|
mm.Pop()
|
|
mm.Push(data[idxPush])
|
|
idx := idxPush - nHalfWin
|
|
max := mm.Max()
|
|
if data[idx] == max && data[idx+1] != max {
|
|
inds = append(inds, idx)
|
|
amps = append(amps, data[idx])
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func sortPeaks(inds []int, amps []float64) ([]int, []float64) {
|
|
sort := make([]int, len(amps))
|
|
out := make([]int, len(amps))
|
|
floats.Argsort(amps, sort)
|
|
for i, idx := range sort {
|
|
out[i] = inds[idx]
|
|
}
|
|
|
|
for i := 0; i < len(inds)/2; i++ {
|
|
out[i], out[len(out)-1-i] = out[len(out)-1-i], out[i]
|
|
amps[i], amps[len(amps)-1-i] = amps[len(amps)-1-i], amps[i]
|
|
}
|
|
return out, amps
|
|
}
|