Files

428 lines
15 KiB
Go
Raw Permalink Normal View History

2026-09-03 10:00:00 +02:00
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT
package signal
import (
"math"
"testing"
"sourcedock.dev/petrbalvin/tensor/internal/core"
)
// arSeries generates a zero-mean AR(p) series: x_t = Σ φ_j·x_{t−j} +
// e_t with standard normal innovations, the first burn values of the
// burn-in discarded.
func arSeries(g *core.Generator, phi []float64, n, burn int) []float64 {
p := len(phi)
x := make([]float64, n+burn)
for i := p; i < len(x); i++ {
v := g.NormalUnit()
for j := range p {
v += phi[j] * x[i-1-j]
}
x[i] = v
}
return x[burn:]
}
// arma11Series generates a zero-mean ARMA(1,1) series: x_t =
// φ·x_{t−1} + e_t + θ·e_{t−1}, the burn-in discarded.
func arma11Series(g *core.Generator, phi, theta float64, n, burn int) []float64 {
x := make([]float64, n)
var xPrev, ePrev float64
for i := -burn; i < n; i++ {
e := g.NormalUnit()
v := phi*xPrev + e + theta*ePrev
xPrev, ePrev = v, e
if i >= 0 {
x[i] = v
}
}
return x
}
// TestEstimateAR2YuleWalker recovers a known AR(2) from a long
// generated series, and then demands the recovered coefficients solve
// the Yule-Walker equations on the empirical autocovariance the fit
// itself read: the Durbin-Levinson solve is exact, so the residuals
// sit at rounding level. The partial autocorrelation, produced by the
// same recursion answering a different question, must agree: lag one
// carries φ1, lag two the signature φ2/(1 − φ1²)... for an AR(2) the
// second partial is φ2 and the rest are noise.
func TestEstimateAR2YuleWalker(t *testing.T) {
const (
phi1 = 0.8
phi2 = -0.4
n = 60000
)
g := core.NewGenerator(7)
x := arSeries(g, []float64{phi1, phi2}, n, 500)
res, err := EstimateAR(mustFloats(t, x), 2)
if err != nil {
t.Fatalf("EstimateAR: %v", err)
}
if math.Abs(res.AR[0]-phi1) > 0.02 || math.Abs(res.AR[1]-phi2) > 0.02 {
t.Fatalf("AR(2) coefficients (%.4f, %.4f), want (%.2f, %.2f)", res.AR[0], res.AR[1], phi1, phi2)
}
if math.Abs(res.InnovationVariance-1) > 0.1 {
t.Fatalf("innovation variance %.4f, want 1", res.InnovationVariance)
}
if math.Abs(res.Mean) > 0.05 {
t.Fatalf("mean %.4f, want 0", res.Mean)
}
// The Yule-Walker equations on the empirical autocovariance.
r, _, _, err := armaAutocovariance("test", mustFloats(t, x), 2)
if err != nil {
t.Fatalf("armaAutocovariance: %v", err)
}
for k := range 2 {
got := res.AR[k]*r[0] + res.AR[1-k]*r[1] - r[k+1]
if math.Abs(got) > 1e-8*r[0] {
t.Fatalf("Yule-Walker residual at lag %d = %g, want rounding only", k+1, got)
}
}
// The PACF, the same recursion read for its reflection
// coefficients, must agree with the fit: lag two's partial is the
// order-two coefficient itself, lag one's is the lag-one
// autocorrelation φ1/(1 − φ2), both exact relations of the same
// autocovariances the Yule-Walker solve consumed.
pacf, err := PartialAutocorrelate(mustFloats(t, x), 3)
if err != nil {
t.Fatalf("PartialAutocorrelate: %v", err)
}
if math.Abs(pacf.FloatAt(1)-res.AR[1]) > 1e-9 {
t.Fatalf("PACF lag 2 = %.10f against the fitted φ2 %.10f", pacf.FloatAt(1), res.AR[1])
}
rho1 := res.AR[0] / (1 - res.AR[1])
if math.Abs(pacf.FloatAt(0)-rho1) > 1e-9 {
t.Fatalf("PACF lag 1 = %.10f against φ1/(1 − φ2) = %.10f", pacf.FloatAt(0), rho1)
}
if math.Abs(pacf.FloatAt(2)) > 0.02 {
t.Fatalf("PACF lag 3 = %.4f, an AR(2) cuts off after lag 2", pacf.FloatAt(2))
}
// The criteria carry the concentrated likelihood of the whole
// series with one parameter per coefficient plus the variance.
wantAIC := float64(n)*(math.Log(2*math.Pi)+math.Log(res.InnovationVariance)+1) + 2*3
if math.Abs(res.AIC-wantAIC) > 1e-6 {
t.Fatalf("AIC %.6f, want %.6f", res.AIC, wantAIC)
}
wantBIC := float64(n)*(math.Log(2*math.Pi)+math.Log(res.InnovationVariance)+1) + math.Log(float64(n))*3
if math.Abs(res.BIC-wantBIC) > 1e-6 {
t.Fatalf("BIC %.6f, want %.6f", res.BIC, wantBIC)
}
}
// TestEstimateARErrors pins the input gates and the stationary-data
// requirement.
func TestEstimateARErrors(t *testing.T) {
x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 6)
if _, err := EstimateAR(x, 0); err == nil {
t.Error("order zero accepted")
}
if _, err := EstimateAR(x, 5); err == nil {
t.Error("order at the sample count accepted")
}
rank2, _ := core.FromFloats([]float64{1, 2, 3, 4}, 2, 2)
if _, err := EstimateAR(rank2, 1); err == nil {
t.Error("rank-2 series accepted")
}
bad, _ := core.FromFloats([]float64{1, math.NaN(), 3, 4}, 4)
if _, err := EstimateAR(bad, 1); err == nil {
t.Error("non-finite series accepted")
}
constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2}, 6)
if _, err := EstimateAR(constant, 1); err == nil {
t.Error("constant series accepted")
}
}
// TestEstimateARMA11HannanRissanen recovers a known ARMA(1,1) from a
// long generated series. The honest caveat, documented with the
// function: Hannan-Rissen conditions on proxy residuals, so the
// estimates carry more sampling noise than a maximum-likelihood fit
// would and the tolerances here are an order looser than the AR(2)
// ones above, the MA side most of all.
func TestEstimateARMA11HannanRissanen(t *testing.T) {
const (
phi = 0.6
theta = 0.4
n = 80000
)
g := core.NewGenerator(11)
x := arma11Series(g, phi, theta, n, 500)
res, err := EstimateARMA(mustFloats(t, x), 1, 1, ARMAOptions{})
if err != nil {
t.Fatalf("EstimateARMA: %v", err)
}
if res.P != 1 || res.Q != 1 {
t.Fatalf("orders (%d, %d), want (1, 1)", res.P, res.Q)
}
if math.Abs(res.AR[0]-phi) > 0.05 {
t.Fatalf("AR coefficient %.4f, want %.2f within 0.05", res.AR[0], phi)
}
if math.Abs(res.MA[0]-theta) > 0.05 {
t.Fatalf("MA coefficient %.4f, want %.2f within 0.05", res.MA[0], theta)
}
if math.Abs(res.InnovationVariance-1) > 0.15 {
t.Fatalf("innovation variance %.4f, want 1 within 0.15", res.InnovationVariance)
}
if math.Abs(res.Mean) > 0.05 {
t.Fatalf("mean %.4f, want 0", res.Mean)
}
// The criteria carry the concentrated whole-series likelihood
// here too, not one windowed by the high-order AR proxy: a
// regression to the windowed count shifts AIC by dozens of points
// and breaks cross-model comparability, so the closed form pins
// the value on the ARMA path exactly as it does on the AR one.
wantLL := -float64(n) / 2 * (math.Log(2*math.Pi) + math.Log(res.InnovationVariance) + 1)
if math.Abs(res.LogLikelihood-wantLL) > 1e-6 {
t.Fatalf("log likelihood %.6f, want %.6f", res.LogLikelihood, wantLL)
}
wantAIC := -2*res.LogLikelihood + 2*3
if math.Abs(res.AIC-wantAIC) > 1e-6 {
t.Fatalf("AIC %.6f, want %.6f", res.AIC, wantAIC)
}
}
// TestEstimateARMAErrors pins the order and data gates, including the
// routing of a pure AR model to the Yule-Walker fit.
func TestEstimateARMAErrors(t *testing.T) {
x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 10)
if _, err := EstimateARMA(x, 0, 0, ARMAOptions{}); err == nil {
t.Error("orders (0, 0) accepted")
}
if _, err := EstimateARMA(x, -1, 1, ARMAOptions{}); err == nil {
t.Error("negative order accepted")
}
if _, err := EstimateARMA(x, 1, 1, ARMAOptions{HighAROrder: 1}); err == nil {
t.Error("proxy order below p+q+1 accepted")
}
if _, err := EstimateARMA(x, 1, 1, ARMAOptions{}); err == nil {
t.Error("series too short for the regression window accepted")
}
// A pure AR request is routed to the Yule-Walker fit: orders below
// its own gate are refused there.
if _, err := EstimateARMA(x, 0, 1, ARMAOptions{}); err == nil {
t.Error("degenerate ARMA with p = 0 accepted on a short series")
}
constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}, 20)
if _, err := EstimateARMA(constant, 1, 1, ARMAOptions{}); err == nil {
t.Error("constant series accepted")
}
}
// TestSelectARMAInformationCriteria generates an AR(2) and demands the
// grid search pick exactly (2, 0) by both criteria, then an ARMA(1,1)
// and demands both criteria pick exactly (1, 1). Both pins are
// deterministic on the seeded generator: the streams are bit-stable,
// so the choice is a regression pin, not a restatement of the
// criteria's large-sample guarantees. The seeds are chosen for wide
// decision margins (the closest competitor sits 1.999 AIC points
// behind on the AR(2) grid, 1.991 AIC and 11.3 BIC points on the
// ARMA(1,1) grid), because the AIC's documented willingness to
// overfit by one order makes exact selection a coin weighted by the
// sample: with margin 2 the decision cannot move on arithmetic-level
// perturbations.
func TestSelectARMAInformationCriteria(t *testing.T) {
g := core.NewGenerator(29)
x := arSeries(g, []float64{0.8, -0.4}, 60000, 500)
for _, criterion := range []string{"aic", "bic"} {
res, err := SelectARMA(mustFloats(t, x), 3, 1, ARMAOptions{Criterion: criterion})
if err != nil {
t.Fatalf("SelectARMA(%s): %v", criterion, err)
}
if res.P != 2 || res.Q != 0 {
t.Fatalf("%s picked ARMA(%d, %d), want ARMA(2, 0)", criterion, res.P, res.Q)
}
}
// Both criteria of the winning model are reported.
res, err := SelectARMA(mustFloats(t, x), 3, 1, ARMAOptions{Criterion: "bic"})
if err != nil {
t.Fatalf("SelectARMA: %v", err)
}
if !(res.BIC > res.AIC) {
t.Fatalf("BIC %.4f does not exceed AIC %.4f at n = 60000", res.BIC, res.AIC)
}
// The ARMA(1,1) series on a wider grid: both criteria pin (1, 1).
g2 := core.NewGenerator(4)
y := arma11Series(g2, 0.6, 0.4, 80000, 500)
for _, criterion := range []string{"aic", "bic"} {
best, err := SelectARMA(mustFloats(t, y), 2, 2, ARMAOptions{Criterion: criterion})
if err != nil {
t.Fatalf("SelectARMA(%s): %v", criterion, err)
}
if best.P != 1 || best.Q != 1 {
t.Fatalf("%s picked ARMA(%d, %d), want ARMA(1, 1)", criterion, best.P, best.Q)
}
}
}
// TestSelectARMAErrors pins the grid and criterion gates, including
// the reporting when no cell can be fitted.
func TestSelectARMAErrors(t *testing.T) {
x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 6)
if _, err := SelectARMA(x, 0, 0, ARMAOptions{}); err == nil {
t.Error("an empty grid accepted")
}
if _, err := SelectARMA(x, -1, 2, ARMAOptions{}); err == nil {
t.Error("a negative grid bound accepted")
}
if _, err := SelectARMA(x, 1, 1, ARMAOptions{Criterion: "sic"}); err == nil {
t.Error("an unknown criterion accepted")
}
constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}, 18)
if _, err := SelectARMA(constant, 2, 2, ARMAOptions{}); err == nil {
t.Error("a grid where nothing can be fitted reported no error")
}
}
// TestARMASpectrumMatchesPeriodogram evaluates the fitted AR(2) model's
// theoretical spectrum against Welch's estimate of a long series
// generated from that fitted model: the frequency grids line up bin
// for bin, and on the broad band between 0.03 and 0.47
// cycles-per-sample the two log-spectra must track each other well
// inside the averaging fluctuation of the estimate.
func TestARMASpectrumMatchesPeriodogram(t *testing.T) {
const n = 40000
g := core.NewGenerator(31)
x := arSeries(g, []float64{0.8, -0.4}, n, 500)
fitted, err := EstimateAR(mustFloats(t, x), 2)
if err != nil {
t.Fatalf("EstimateAR: %v", err)
}
// A long series from the fitted coefficients: the fitted model's
// own output, whose periodogram the theory must describe.
const n2 = 1 << 18
sigma := math.Sqrt(fitted.InnovationVariance)
y := make([]float64, n2+500)
for i := 2; i < len(y); i++ {
y[i] = fitted.AR[0]*y[i-1] + fitted.AR[1]*y[i-2] + sigma*g.NormalUnit()
}
const segment = 1024
freqs, psd, err := WelchPSD(mustFloats(t, y[500:]), 1, segment, segment/2, "hann")
if err != nil {
t.Fatalf("WelchPSD: %v", err)
}
bins := segment/2 + 1
tFreqs, tPsd, err := ARMASpectrum(fitted, bins)
if err != nil {
t.Fatalf("ARMASpectrum: %v", err)
}
for k := range bins {
if math.Abs(freqs.FloatAt(k)-tFreqs.FloatAt(k)) > 1e-12 {
t.Fatalf("frequency grids disagree at bin %d", k)
}
}
worst := 0.0
for k := range bins {
f := freqs.FloatAt(k)
if f < 0.03 || f > 0.47 {
continue
}
d := math.Abs(math.Log(psd.FloatAt(k)) - math.Log(tPsd.FloatAt(k)))
if d > worst {
worst = d
}
}
if worst > 0.2 {
t.Fatalf("log-spectral gap %.4f on the broad band, want under 0.2", worst)
}
// The resonance: cos ω₀ = φ1(φ2 − 1)/(4φ2) locates the AR(2)
// peak; both the Welch estimate and the theoretical spectrum must
// peak inside ±0.03 cycles-per-sample of it, and the theoretical
// argmax must sit within two bins of the formula's own optimum
// (the peak is broad, so bins are the honest resolution).
phi1, phi2 := fitted.AR[0], fitted.AR[1]
cosw0 := phi1 * (phi2 - 1) / (4 * phi2)
f0 := math.Acos(cosw0) / (2 * math.Pi)
if f0 < 0.02 || f0 > 0.48 {
t.Fatalf("the resonance formula left the band: %g", f0)
}
binOf := func(a *core.Array) int {
best, at := math.Inf(-1), -1
for k := range bins {
f := freqs.FloatAt(k)
if f < 0.02 || f > 0.48 {
continue
}
if a.FloatAt(k) > best {
best, at = a.FloatAt(k), k
}
}
return at
}
if d := math.Abs(float64(binOf(psd))/float64(segment) - f0); d > 0.03 {
t.Fatalf("the empirical peak sits %.4f from the resonance %g", d, f0)
}
if d := math.Abs(float64(binOf(tPsd))/float64(segment) - f0); d > 2.0/float64(segment) {
t.Fatalf("the theoretical peak sits %.4f from the resonance %g", d, f0)
}
// The theoretical values, checked against the formula evaluated
// independently on a fine grid: the ARMA(1,0) polynomial arithmetic
// is the whole function, so this pins it to rounding.
const fine = 4097
fFreqs, fPsd, err := ARMASpectrum(fitted, fine)
if err != nil {
t.Fatalf("ARMASpectrum: %v", err)
}
for k := range fine {
f := 0.5 * float64(k) / float64(fine-1)
w := 2 * math.Pi * f
re := 1 - phi1*math.Cos(w) - phi2*math.Cos(2*w)
im := phi1*math.Sin(w) + phi2*math.Sin(2*w)
want := fitted.InnovationVariance / (re*re + im*im)
if k > 0 && k < fine-1 {
want *= 2
}
if math.Abs(fPsd.FloatAt(k)-want) > 1e-10*want {
t.Fatalf("spectrum %.12g at bin %d, formula %.12g", fPsd.FloatAt(k), k, want)
}
if math.Abs(fFreqs.FloatAt(k)-f) > 1e-12 {
t.Fatalf("frequency %.12g at bin %d, want %.12g", fFreqs.FloatAt(k), k, f)
}
}
// The flat model answers the innovation variance at the edges and
// its double between them: the one-sided fold the Welch estimate
// of the same white noise prints.
flat := &ARMAResult{InnovationVariance: 2.5}
_, flatPsd, err := ARMASpectrum(flat, 16)
if err != nil {
t.Fatalf("ARMASpectrum: %v", err)
}
for k := range 16 {
want := 5.0
if k == 0 || k == 15 {
want = 2.5
}
if flatPsd.FloatAt(k) != want {
t.Fatalf("the flat model's spectrum is %g at bin %d, want %g", flatPsd.FloatAt(k), k, want)
}
}
}
// TestARMASpectrumErrors pins the gates, including the pole on the
// unit circle being named instead of answered with infinities.
func TestARMASpectrumErrors(t *testing.T) {
if _, _, err := ARMASpectrum(nil, 8); err == nil {
t.Error("a nil model accepted")
}
res := &ARMAResult{AR: []float64{0.5}, InnovationVariance: 1}
if _, _, err := ARMASpectrum(res, 1); err == nil {
t.Error("a single frequency accepted")
}
res.InnovationVariance = -1
if _, _, err := ARMASpectrum(res, 8); err == nil {
t.Error("a negative innovation variance accepted")
}
res.InnovationVariance = 1
res.AR = []float64{math.NaN()}
if _, _, err := ARMASpectrum(res, 8); err == nil {
t.Error("a non-finite coefficient accepted")
}
res.AR = []float64{1}
if _, _, err := ARMASpectrum(res, 8); err == nil {
t.Error("a pole on the unit circle accepted")
}
}