125 lines
3.5 KiB
Go
125 lines
3.5 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
package signal
|
|
|
|
import "sourcedock.dev/petrbalvin/tensor/internal/core"
|
|
|
|
import (
|
|
"math"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
// naiveDFT computes the transform quadratically, the reference the fast
|
|
// paths are verified against.
|
|
func naiveDFT(vals []complex128, sign float64) []complex128 {
|
|
n := len(vals)
|
|
out := make([]complex128, n)
|
|
for k := range n {
|
|
var s complex128
|
|
for j, x := range vals {
|
|
angle := sign * 2 * math.Pi * float64(j*k%n) / float64(n)
|
|
s += x * complex(math.Cos(angle), math.Sin(angle))
|
|
}
|
|
out[k] = s
|
|
}
|
|
return out
|
|
}
|
|
|
|
func assertCplxApprox(t *testing.T, name string, got []complex128, want []complex128) {
|
|
t.Helper()
|
|
for i := range want {
|
|
if math.Abs(real(got[i])-real(want[i])) > 1e-9 || math.Abs(imag(got[i])-imag(want[i])) > 1e-9 {
|
|
t.Fatalf("%s[%d]: got %v, want %v", name, i, got[i], want[i])
|
|
}
|
|
}
|
|
}
|
|
|
|
func fftValues(t *testing.T, a *core.Array) []complex128 {
|
|
t.Helper()
|
|
out, err := FFT(a)
|
|
if err != nil {
|
|
t.Fatalf("FFT: %v", err)
|
|
}
|
|
vals := make([]complex128, out.Len())
|
|
for i := range vals {
|
|
vals[i], _ = core.ComplexAt(out, i)
|
|
}
|
|
return vals
|
|
}
|
|
|
|
func TestFFTKnownValues(t *testing.T) {
|
|
// FFT of [1, 2, 3, 4] is the textbook [10, -2+2i, -2, -2-2i].
|
|
a := mustFromInts(t, []int64{1, 2, 3, 4}, 4)
|
|
got := fftValues(t, a)
|
|
assertCplxApprox(t, "FFT", got, []complex128{10, complex(-2, 2), -2, complex(-2, -2)})
|
|
|
|
// A constant input concentrates in the DC bin.
|
|
c := mustFromFloats(t, []float64{2.5, 2.5, 2.5, 2.5}, 4)
|
|
got = fftValues(t, c)
|
|
assertCplxApprox(t, "FFT DC", got[:1], []complex128{10})
|
|
for i := 1; i < 4; i++ {
|
|
if math.Abs(real(got[i])) > 1e-9 || math.Abs(imag(got[i])) > 1e-9 {
|
|
t.Fatalf("FFT constant input leaked into bin %d: %v", i, got[i])
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestFFTAgainstNaive(t *testing.T) {
|
|
// Powers of two take the radix-2 path; the odd and mixed lengths take
|
|
// Bluestein. Both must agree with the quadratic reference.
|
|
lengths := []int{1, 2, 4, 8, 16, 3, 5, 6, 7, 12, 13, 20}
|
|
for _, n := range lengths {
|
|
vals := make([]complex128, n)
|
|
g := core.NewGenerator(int64(n))
|
|
for i := range vals {
|
|
f, _ := core.Floats(g, 2)
|
|
re, _ := core.FloatAt(f, 0)
|
|
im, _ := core.FloatAt(f, 1)
|
|
vals[i] = complex(re*10-5, im*10-5)
|
|
}
|
|
src, err := core.FromComplexes(vals, n)
|
|
if err != nil {
|
|
t.Fatalf("FromComplexes(%d): %v", n, err)
|
|
}
|
|
assertCplxApprox(t, "FFT radix/bluestein", fftValues(t, src), naiveDFT(vals, -1))
|
|
}
|
|
}
|
|
|
|
func TestIFFTRoundTrip(t *testing.T) {
|
|
all := []float64{0.5, -1.25, 3, -0.5, 2, 1, -2, 4, 0.25, -3, 1.5, 2.75}
|
|
for _, n := range []int{8, 5, 12} {
|
|
src := mustFromFloats(t, all[:n], n)
|
|
fwd, err := FFT(src)
|
|
if err != nil {
|
|
t.Fatalf("FFT(%d): %v", n, err)
|
|
}
|
|
back, err := IFFT(fwd)
|
|
if err != nil {
|
|
t.Fatalf("IFFT(%d): %v", n, err)
|
|
}
|
|
for i := range n {
|
|
want, _ := core.FloatAt(src, i)
|
|
got, _ := core.ComplexAt(back, i)
|
|
if math.Abs(real(got)-want) > 1e-9 || math.Abs(imag(got)) > 1e-9 {
|
|
t.Fatalf("IFFT(%d)[%d]: got %v, want %v", n, i, got, want)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestFFTErrors(t *testing.T) {
|
|
empty := mustFromInts(t, nil, 0)
|
|
if _, err := FFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
|
|
t.Fatalf("FFT empty: %v", err)
|
|
}
|
|
if _, err := IFFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
|
|
t.Fatalf("IFFT empty: %v", err)
|
|
}
|
|
m := mustFromInts(t, []int64{1, 2, 3, 4}, 2, 2)
|
|
if _, err := FFT(m); err == nil || !strings.Contains(err.Error(), "needs a 1-D array") {
|
|
t.Fatalf("FFT 2-D: %v", err)
|
|
}
|
|
}
|