248 lines
6.6 KiB
Go
248 lines
6.6 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"
|
||
)
|
||
|
||
func TestFFT2(t *testing.T) {
|
||
// 2×2 identity-like: shift.
|
||
a := mustFromFloats(t, []float64{
|
||
1, 0,
|
||
0, 0,
|
||
}, 2, 2)
|
||
got, err := FFT2(a)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got.Shape()[0] != 2 || got.Shape()[1] != 2 {
|
||
t.Errorf("FFT2 shape: %v", got.Shape())
|
||
}
|
||
if got.Dtype() != core.Complex {
|
||
t.Errorf("FFT2 dtype: %s", got.Dtype())
|
||
}
|
||
// Round-trip: IFFT2(FFT2(x)) ≈ x.
|
||
back, err := IFFT2(got)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for i := range 4 {
|
||
v, _ := core.ComplexAt(back, i/2, i%2)
|
||
orig, _ := core.FloatAt(a, i/2, i%2)
|
||
if math.Abs(real(v)-orig) > 1e-9 {
|
||
t.Errorf("FFT2 round-trip [%d]: got %v, want %v", i, real(v), orig)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestFFT3(t *testing.T) {
|
||
a := mustFromFloats(t, []float64{
|
||
1, 0, 0, 0,
|
||
0, 0, 0, 0,
|
||
0, 0, 0, 0,
|
||
}, 3, 2, 2)
|
||
got, err := FFT3(a)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got.Shape()[0] != 3 || got.Shape()[1] != 2 || got.Shape()[2] != 2 {
|
||
t.Errorf("FFT3 shape: %v", got.Shape())
|
||
}
|
||
// Round-trip.
|
||
back, err := IFFT3(got)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
v0, _ := core.ComplexAt(back, 0, 0, 0)
|
||
if math.Abs(real(v0)-1) > 1e-9 {
|
||
t.Errorf("FFT3 round-trip [0,0,0]: got %v, want 1", real(v0))
|
||
}
|
||
}
|
||
|
||
func TestFFTN(t *testing.T) {
|
||
a := mustFromFloats(t, []float64{
|
||
1, 2, 3, 4,
|
||
5, 6, 7, 8,
|
||
}, 2, 4)
|
||
got, err := FFTN(a, nil)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got.Dtype() != core.Complex {
|
||
t.Errorf("FFTN dtype: %s", got.Dtype())
|
||
}
|
||
// Round-trip.
|
||
back, err := IFFTN(got, nil)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for i := range a.Len() {
|
||
v, _ := core.ComplexAt(back, i/4, i%4)
|
||
orig, _ := core.FloatAt(a, i/4, i%4)
|
||
if math.Abs(real(v)-orig) > 1e-9 {
|
||
t.Errorf("FFTN round-trip [%d]: got %v, want %v", i, real(v), orig)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestRFFT(t *testing.T) {
|
||
// Real input of length 8 gives a spectrum of length 5.
|
||
a := mustFromFloats(t, []float64{1, 0, 0, 0, 0, 0, 0, 0}, 8)
|
||
spec, err := RFFT(a)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if spec.Len() != 5 {
|
||
t.Errorf("RFFT spectrum length: %d, want 5", spec.Len())
|
||
}
|
||
// IRFFT round-trip.
|
||
back, err := IRFFT(spec, 8)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for i := range 8 {
|
||
v, _ := core.FloatAt(back, i)
|
||
orig, _ := core.FloatAt(a, i)
|
||
if math.Abs(v-orig) > 1e-9 {
|
||
t.Errorf("RFFT round-trip [%d]: got %v, want %v", i, v, orig)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestFFTFreq(t *testing.T) {
|
||
f := FFTFreq(8, 1.0)
|
||
if f.Len() != 8 {
|
||
t.Errorf("FFTFreq length: %d", f.Len())
|
||
}
|
||
// Frequencies: [0, 1/8, 2/8, 3/8, -4/8, -3/8, -2/8, -1/8] = [0, 0.125, 0.25, 0.375, -0.5, -0.375, -0.25, -0.125]
|
||
want := []float64{0, 0.125, 0.25, 0.375, -0.5, -0.375, -0.25, -0.125}
|
||
for i, w := range want {
|
||
v, _ := core.FloatAt(f, i)
|
||
if math.Abs(v-w) > 1e-9 {
|
||
t.Errorf("FFTFreq [%d]: got %v, want %v", i, v, w)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestIRFFTLengthOneSpectrum pins the degenerate spectrum: a single
|
||
// DC bin defaults to n = 1 instead of n = 0, an empty spectrum is an
|
||
// error, and a non-positive explicit n is an error.
|
||
func TestIRFFTLengthOneSpectrum(t *testing.T) {
|
||
spec := mustComplexes(t, []complex128{complex(5, 0)}, 1)
|
||
back, err := IRFFT(spec, 0)
|
||
if err != nil {
|
||
t.Fatalf("IRFFT default n: %v", err)
|
||
}
|
||
if back.Len() != 1 || math.Abs(back.FloatAt(0)-5) > 1e-12 {
|
||
t.Fatalf("IRFFT of a DC-only spectrum = %v, want [5]", back)
|
||
}
|
||
back, err = IRFFT(spec, 1)
|
||
if err != nil {
|
||
t.Fatalf("IRFFT n=1: %v", err)
|
||
}
|
||
if back.Len() != 1 || math.Abs(back.FloatAt(0)-5) > 1e-12 {
|
||
t.Fatalf("IRFFT n=1 = %v, want [5]", back)
|
||
}
|
||
if _, err := IRFFT(spec, 3); err == nil {
|
||
t.Fatal("expected an error for a spectrum length that mismatches n")
|
||
}
|
||
if _, err := IRFFT(mustComplexes(t, nil, 0), 0); err == nil {
|
||
t.Fatal("expected an error for an empty spectrum")
|
||
}
|
||
if _, err := IRFFT(spec, -2); err == nil {
|
||
t.Fatal("expected an error for a negative n")
|
||
}
|
||
}
|
||
|
||
// TestFFTEmptyComplexIsError pins the empty guard for complex input,
|
||
// which bypasses the real-to-complex conversion where the check used
|
||
// to live.
|
||
func TestFFTEmptyComplexIsError(t *testing.T) {
|
||
empty := mustComplexes(t, nil, 0)
|
||
if _, err := FFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
|
||
t.Fatalf("FFT of an empty complex array: %v", err)
|
||
}
|
||
if _, err := IFFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
|
||
t.Fatalf("IFFT of an empty complex array: %v", err)
|
||
}
|
||
}
|
||
|
||
// TestIFFTNSubsetDimsRoundTrip pins the inverse scaling: IFFTN divides
|
||
// by the product of the transformed extents, so inverting a subset of
|
||
// the dimensions restores the input exactly (scaling by the total
|
||
// length would return a/2 here).
|
||
func TestIFFTNSubsetDimsRoundTrip(t *testing.T) {
|
||
src := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 2, 3)
|
||
fwd, err := FFTN(src, []int{1})
|
||
if err != nil {
|
||
t.Fatalf("FFTN: %v", err)
|
||
}
|
||
back, err := IFFTN(fwd, []int{1})
|
||
if err != nil {
|
||
t.Fatalf("IFFTN: %v", err)
|
||
}
|
||
for i := range 6 {
|
||
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("subset round-trip [%d]: got %v, want %v", i, got, want)
|
||
}
|
||
}
|
||
// The full-dims path keeps its usual total scaling.
|
||
fwdAll, err := FFTN(src, nil)
|
||
if err != nil {
|
||
t.Fatalf("FFTN all dims: %v", err)
|
||
}
|
||
backAll, err := IFFTN(fwdAll, nil)
|
||
if err != nil {
|
||
t.Fatalf("IFFTN all dims: %v", err)
|
||
}
|
||
for i := range 6 {
|
||
want, _ := core.FloatAt(src, i)
|
||
got, _ := core.ComplexAt(backAll, i)
|
||
if math.Abs(real(got)-want) > 1e-9 || math.Abs(imag(got)) > 1e-9 {
|
||
t.Fatalf("full round-trip [%d]: got %v, want %v", i, got, want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestFFT2DoesNotMutateComplexInput pins the immutability contract for
|
||
// the multi-dimension chain: FFT2 with an already complex input must
|
||
// not overwrite the caller's payload. An in-place first pass used to
|
||
// slip through here because the return values stayed correct; only the
|
||
// caller's array silently held the spectrum afterwards.
|
||
func TestFFT2DoesNotMutateComplexInput(t *testing.T) {
|
||
vals := []complex128{
|
||
1 + 2i, 3 - 1i,
|
||
-0.5 + 0.5i, 2 + 0i,
|
||
}
|
||
a, err := core.FromComplexes(vals, 2, 2)
|
||
if err != nil {
|
||
t.Fatalf("FromComplexes: %v", err)
|
||
}
|
||
got, err := FFT2(a)
|
||
if err != nil {
|
||
t.Fatalf("FFT2: %v", err)
|
||
}
|
||
for i, want := range vals {
|
||
if a.ComplexAt(i) != want {
|
||
t.Fatalf("FFT2 mutated input[%d] = %v, want %v", i, a.ComplexAt(i), want)
|
||
}
|
||
}
|
||
back, err := IFFT2(got)
|
||
if err != nil {
|
||
t.Fatalf("IFFT2: %v", err)
|
||
}
|
||
for i, want := range vals {
|
||
if back.ComplexAt(i) != want {
|
||
t.Fatalf("IFFT2 mutated gradient input or lost precision [%d] = %v, want %v", i, back.ComplexAt(i), want)
|
||
}
|
||
}
|
||
}
|