428 lines
15 KiB
Go
428 lines
15 KiB
Go
// 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")
|
||
}
|
||
}
|