feat: initial release

Assisted-by: GLM 5.3 Flash

examples/spectral/main.go

@@ -0,0 +1,136 @@
+// 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] }