137 lines
3.7 KiB
Go
137 lines
3.7 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
// Command spectral estimates the frequency content of a signal two
|
|
// ways: Welch's averaged periodogram on evenly sampled data, and the
|
|
// Lomb-Scargle periodogram on the same signal observed at irregular
|
|
// times, where an FFT cannot run at all. Both must find the two
|
|
// buried sinusoids.
|
|
//
|
|
// Usage: go run ./examples/spectral
|
|
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"math"
|
|
|
|
"sourcedock.dev/petrbalvin/tensor"
|
|
)
|
|
|
|
func main() {
|
|
const (
|
|
fs = 100.0
|
|
seconds = 4.0
|
|
f1 = 5.0
|
|
f2 = 13.0
|
|
)
|
|
n := int(fs * seconds)
|
|
gen := tensor.NewGenerator(11)
|
|
|
|
// The signal: two sinusoids plus noise.
|
|
t := make([]float64, n)
|
|
x := make([]float64, n)
|
|
for i := range n {
|
|
t[i] = float64(i) / fs
|
|
x[i] = math.Sin(2*math.Pi*f1*t[i]) + 0.6*math.Sin(2*math.Pi*f2*t[i]) + 0.4*gen.NormalUnit()
|
|
}
|
|
xArr, err := tensor.FromFloats(x, n)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
// Welch: average periodograms over Hann-windowed segments, the
|
|
// variance-suppressed estimate an FFT alone cannot give.
|
|
freqs, psd, err := tensor.WelchPSD(xArr, fs, 256, 128, "hann")
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
wf1, wf2 := twoPeaks(peakFrequencies(freqs, psd, 2))
|
|
fmt.Printf("welch peaks at %.2f Hz and %.2f Hz (want %.1f and %.1f)\n", wf1, wf2, f1, f2)
|
|
|
|
// Lomb-Scargle: keep every second sample at jittered times, the
|
|
// uneven regime the DFT does not define. The mean rate stays at
|
|
// 50 Hz, comfortably above both sources' Nyquist needs, while the
|
|
// jitter is what makes the ordinary FFT inapplicable.
|
|
times := make([]float64, 0, n/2)
|
|
values := make([]float64, 0, n/2)
|
|
for i := 0; i < n; i += 2 {
|
|
jitter := 0.6 * gen.Unit() / fs
|
|
times = append(times, t[i]+jitter)
|
|
values = append(values, x[i])
|
|
}
|
|
tArr, err := tensor.FromFloats(times, len(times))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
vArr, err := tensor.FromFloats(values, len(values))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
lsFreqs, power, err := tensor.LombScargle(tArr, vArr, 1.0, 30.0, 3000)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
lf1, lf2 := twoPeaks(peakFrequencies(lsFreqs, power, 2))
|
|
fmt.Printf("lomb-scargle peaks at %.2f Hz and %.2f Hz (want %.1f and %.1f)\n", lf1, lf2, f1, f2)
|
|
|
|
for _, got := range []float64{wf1, wf2, lf1, lf2} {
|
|
if math.Abs(got-f1) > 0.3 && math.Abs(got-f2) > 0.3 {
|
|
log.Fatalf("a peak landed at %.2f Hz, away from both sources", got)
|
|
}
|
|
}
|
|
}
|
|
|
|
// peakFrequencies returns the abscissae of the count largest local
|
|
// maxima of a periodogram, descending by height and kept at least
|
|
// 1.5 Hz apart so a sidelobe of a tall peak cannot shadow a real one.
|
|
func peakFrequencies(freqs, power *tensor.Array, count int) []float64 {
|
|
n := freqs.Len()
|
|
// Three-point boxcar smooth: the periodogram's noise is white, a
|
|
// genuine peak is not.
|
|
smooth := make([]float64, n)
|
|
for i := range n {
|
|
lo := max(i-1, 0)
|
|
hi := min(i+1, n-1)
|
|
s := 0.0
|
|
for j := lo; j <= hi; j++ {
|
|
s += power.FloatAt(j)
|
|
}
|
|
smooth[i] = s / float64(hi-lo+1)
|
|
}
|
|
type peak struct {
|
|
f, h float64
|
|
}
|
|
var peaks []peak
|
|
for i := 1; i < n-1; i++ {
|
|
if smooth[i] > smooth[i-1] && smooth[i] >= smooth[i+1] {
|
|
peaks = append(peaks, peak{freqs.FloatAt(i), smooth[i]})
|
|
}
|
|
}
|
|
for i := 1; i < len(peaks); i++ {
|
|
for j := i; j > 0 && peaks[j-1].h < peaks[j].h; j-- {
|
|
peaks[j-1], peaks[j] = peaks[j], peaks[j-1]
|
|
}
|
|
}
|
|
out := make([]float64, 0, count)
|
|
for _, p := range peaks {
|
|
if len(out) == count {
|
|
break
|
|
}
|
|
far := true
|
|
for _, f := range out {
|
|
if math.Abs(p.f-f) < 1.5 {
|
|
far = false
|
|
break
|
|
}
|
|
}
|
|
if far {
|
|
out = append(out, p.f)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// twoPeaks unpacks the two-element result of peakFrequencies.
|
|
func twoPeaks(fs []float64) (float64, float64) { return fs[0], fs[1] }
|