477 lines
17 KiB
Go
477 lines
17 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package signal
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Time-series model estimation: the autoregressive and
|
||
// ARMA staples. The pure AR fit solves the Yule-Walker equations by
|
||
// the Durbin-Levinson recursion over the biased autocovariance, the
|
||
// small solver this package's own PartialAutocorrelate runs in
|
||
// miniature (the recursion lives locally here; the domain boundary
|
||
// allows no linalg import, and the earlier recursion answers a
|
||
// different question). The mixed model is fitted by the
|
||
// Hannan-Rissanen innovations method: a high-order autoregression
|
||
// first, whose residuals stand in for the unobserved innovations,
|
||
// then one least-squares regression of the series on its own lags and
|
||
// the lagged residuals. Both report the concentrated Gaussian
|
||
// likelihood of their innovations, and with it the AIC and BIC the
|
||
// grid search in SelectARMA minimises.
|
||
|
||
// ARMAResult holds one fitted time-series model. The coefficients
|
||
// follow the convention x_t − μ = Σ_j φ_j·(x_{t−j} − μ) + e_t +
|
||
// Σ_k θ_k·e_{t−k}, innovations white with variance
|
||
// InnovationVariance.
|
||
type ARMAResult struct {
|
||
// AR holds φ_1..φ_P, empty only for a pure MA model.
|
||
AR []float64
|
||
// MA holds θ_1..θ_Q, empty only for a pure AR model.
|
||
MA []float64
|
||
// InnovationVariance is the estimated variance σ² of e_t.
|
||
InnovationVariance float64
|
||
// Mean is the sample mean μ the fit removed and reports back.
|
||
Mean float64
|
||
// LogLikelihood is the concentrated Gaussian likelihood of the
|
||
// innovations over the whole series: −n/2·(log 2πσ² + 1) at the
|
||
// fitted innovation variance, so the criteria of different models
|
||
// on one series sit on a common footing.
|
||
LogLikelihood float64
|
||
// AIC is −2·LogLikelihood + 2k with k = P + Q + 1, the variance
|
||
// counted as a parameter.
|
||
AIC float64
|
||
// BIC is −2·LogLikelihood + k·log n, the same k against the
|
||
// innovation count's logarithm.
|
||
BIC float64
|
||
// P and Q record the orders the result carries.
|
||
P int
|
||
Q int
|
||
}
|
||
|
||
// ARMAOptions tunes the Hannan-Rissanen estimation and the grid
|
||
// search. HighAROrder is the order of the proxy autoregression whose
|
||
// residuals supply the MA stage; the zero value implies the default
|
||
// max(16, p+q+8), which the data length bounds. Criterion names the
|
||
// information criterion SelectARMA minimises: "aic" or "bic", with
|
||
// "aic" implied by the zero value.
|
||
type ARMAOptions struct {
|
||
HighAROrder int
|
||
Criterion string
|
||
}
|
||
|
||
// armaLikelihood fills the concentrated likelihood and the criteria
|
||
// from the innovation variance, evaluated over the whole series: nEff
|
||
// is the same sample count for every model fitted to one series,
|
||
// which is what makes the criteria comparable across a grid, and k
|
||
// counts the coefficients plus the variance. The values a fit
|
||
// conditions on (an AR's first p samples, the Hannan-Rissanen
|
||
// window's proxy residuals) are estimation detail, not a data
|
||
// difference.
|
||
func armaLikelihood(res *ARMAResult, nEff int) {
|
||
k := res.P + res.Q + 1
|
||
logL := -0.5 * float64(nEff) * (math.Log(2*math.Pi) + math.Log(res.InnovationVariance) + 1)
|
||
res.LogLikelihood = logL
|
||
res.AIC = 2*float64(k) - 2*logL
|
||
res.BIC = math.Log(float64(nEff))*float64(k) - 2*logL
|
||
}
|
||
|
||
// armaAutocovariance returns the biased autocovariances r_0..r_maxLag
|
||
// of a real rank-1 series, the estimator the Yule-Walker theory is
|
||
// written for: the mean removed, every lag divided by the sample
|
||
// count. The series' values are read through the house autocorrelation
|
||
// and rescaled by the lag-zero power. The mean comes back too, because
|
||
// every model is fitted to the demeaned series and reported with it.
|
||
func armaAutocovariance(name string, x *core.Array, maxLag int) (r []float64, n int, mean float64, err error) {
|
||
if x.NDim() != 1 {
|
||
return nil, 0, 0, base.Errf("%s: needs a rank-1 series, got shape %s", name, base.ShapeText(x.Shape()))
|
||
}
|
||
if x.Dtype() == core.Complex {
|
||
return nil, 0, 0, base.Errf("%s: complex series are not supported", name)
|
||
}
|
||
n = x.Len()
|
||
if n < maxLag+2 {
|
||
return nil, 0, 0, base.Errf("%s: at least %d samples are needed for lag %d, got %d",
|
||
name, maxLag+2, maxLag, n)
|
||
}
|
||
vals := widenFloats(x)
|
||
if err := kfFinite(name, "the series", vals); err != nil {
|
||
return nil, 0, 0, err
|
||
}
|
||
for _, v := range vals {
|
||
mean += v
|
||
}
|
||
mean /= float64(n)
|
||
power := 0.0
|
||
for _, v := range vals {
|
||
d := v - mean
|
||
power += d * d
|
||
}
|
||
variance := power / float64(n)
|
||
acf, err := Autocorrelate(x, maxLag)
|
||
if err != nil {
|
||
return nil, 0, 0, err
|
||
}
|
||
r = make([]float64, maxLag+1)
|
||
for k := range r {
|
||
r[k] = acf.FloatAt(k) * variance
|
||
}
|
||
return r, n, mean, nil
|
||
}
|
||
|
||
// yuleWalker solves the Yule-Walker equations R·φ = r for the
|
||
// order-p autoregression's coefficients, R the Toeplitz matrix of the
|
||
// autocovariances r[0..p], through the Durbin-Levinson recursion: p²
|
||
// work, no matrix factored. It returns φ_1..φ_p and the innovation
|
||
// variance r₀·(1 − Σφ_j·ρ_j). A non-positive recursion denominator
|
||
// means the autocovariances do not belong to a stationary process;
|
||
// the same recursion, run for its reflection coefficients, is what
|
||
// PartialAutocorrelate reports.
|
||
func yuleWalker(name string, r []float64, p int) (phi []float64, sigma2 float64, err error) {
|
||
phi = make([]float64, p+1) // 1-based: phi[k] holds φ_k
|
||
prev := make([]float64, p+1)
|
||
for k := 1; k <= p; k++ {
|
||
num, den := r[k], r[0]
|
||
for j := 1; j < k; j++ {
|
||
num -= prev[j] * r[k-j]
|
||
den -= prev[j] * r[j]
|
||
}
|
||
if !(den > 0) || math.IsNaN(den) {
|
||
return nil, 0, base.Errf("%s: the Durbin-Levinson recursion broke down at order %d: the autocovariances do not describe a stationary process", name, k)
|
||
}
|
||
phi[k] = num / den
|
||
for j := 1; j < k; j++ {
|
||
phi[j] = prev[j] - phi[k]*prev[k-j]
|
||
}
|
||
copy(prev, phi)
|
||
}
|
||
out := make([]float64, p)
|
||
copy(out, phi[1:])
|
||
sigma2 = r[0]
|
||
for j := 1; j <= p; j++ {
|
||
sigma2 -= phi[j] * r[j]
|
||
}
|
||
if !(sigma2 > 0) || math.IsNaN(sigma2) {
|
||
return nil, 0, base.Errf("%s: the innovation variance came out %g: the autocovariances do not describe a stationary process", name, sigma2)
|
||
}
|
||
return out, sigma2, nil
|
||
}
|
||
|
||
// armaLstsq solves the small least-squares problem min ‖A·c − b‖²
|
||
// through the normal equations and Gaussian elimination with partial
|
||
// pivoting. The width stays at a model order's size, where the
|
||
// squandered conditioning costs nothing that matters; a pivot too
|
||
// small against the normal matrix's largest entry reports a
|
||
// degenerate regression rather than returning coefficients for it.
|
||
func armaLstsq(name string, design [][]float64, target []float64) ([]float64, error) {
|
||
width := len(design[0])
|
||
g := make([]float64, width*width)
|
||
rhs := make([]float64, width)
|
||
for i, row := range design {
|
||
for a := range width {
|
||
rhs[a] += row[a] * target[i]
|
||
for b := range width {
|
||
g[a*width+b] += row[a] * row[b]
|
||
}
|
||
}
|
||
}
|
||
big := 0.0
|
||
for _, v := range g {
|
||
big = max(big, math.Abs(v))
|
||
}
|
||
for a := range width {
|
||
piv := a
|
||
for c := a + 1; c < width; c++ {
|
||
if math.Abs(g[c*width+a]) > math.Abs(g[piv*width+a]) {
|
||
piv = c
|
||
}
|
||
}
|
||
if !(math.Abs(g[piv*width+a]) > 1e-12*big) {
|
||
return nil, base.Errf("%s: the regression design is degenerate: its normal equations have no pivot at column %d", name, a+1)
|
||
}
|
||
if piv != a {
|
||
for b := range width {
|
||
g[a*width+b], g[piv*width+b] = g[piv*width+b], g[a*width+b]
|
||
}
|
||
rhs[a], rhs[piv] = rhs[piv], rhs[a]
|
||
}
|
||
for c := a + 1; c < width; c++ {
|
||
f := g[c*width+a] / g[a*width+a]
|
||
for b := a; b < width; b++ {
|
||
g[c*width+b] -= f * g[a*width+b]
|
||
}
|
||
rhs[c] -= f * rhs[a]
|
||
}
|
||
}
|
||
coef := make([]float64, width)
|
||
for a := width - 1; a >= 0; a-- {
|
||
total := rhs[a]
|
||
for b := a + 1; b < width; b++ {
|
||
total -= g[a*width+b] * coef[b]
|
||
}
|
||
coef[a] = total / g[a*width+a]
|
||
}
|
||
return coef, nil
|
||
}
|
||
|
||
// EstimateAR fits the order-p autoregression x_t − μ = Σ_j
|
||
// φ_j·(x_{t−j} − μ) + e_t to the real rank-1 series x: the
|
||
// Yule-Walker equations over the biased autocovariance, solved by the
|
||
// Durbin-Levinson recursion. The result carries the coefficients, the
|
||
// innovation variance, the concentrated Gaussian likelihood of the
|
||
// whole series, and the AIC and BIC that compare models fitted to
|
||
// one series. An order below 1, a series too short for the lag, and
|
||
// non-finite or complex data are errors; autocovariances that do not
|
||
// describe a stationary process break the recursion and are reported.
|
||
func EstimateAR(x *core.Array, order int) (*ARMAResult, error) {
|
||
const name = "EstimateAR"
|
||
if order < 1 {
|
||
return nil, base.Errf("%s: the order must be at least 1, got %d", name, order)
|
||
}
|
||
r, n, mean, err := armaAutocovariance(name, x, order)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
phi, sigma2, err := yuleWalker(name, r, order)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
res := &ARMAResult{AR: phi, InnovationVariance: sigma2, Mean: mean, P: order}
|
||
armaLikelihood(res, n)
|
||
return res, nil
|
||
}
|
||
|
||
// EstimateARMA fits the ARMA(p, q) model x_t − μ = Σ_j φ_j·(x_{t−j} −
|
||
// μ) + e_t + Σ_k θ_k·e_{t−k} by the Hannan-Rissanen innovations
|
||
// method: an order-m autoregression first (m from ARMAOptions'
|
||
// HighAROrder, defaulting to max(16, p+q+8)), whose residuals stand
|
||
// in for the unobserved innovations, then one least-squares
|
||
// regression of the demeaned series on its own lags and the lagged
|
||
// residuals. The result carries the coefficients in the documented
|
||
// convention, the regression residuals' variance as the innovation
|
||
// variance and the likelihood and criteria computed from it over the
|
||
// whole series, so the criteria line up with the pure AR fit's and
|
||
// with the other cells of a SelectARMA grid.
|
||
//
|
||
// The honest caveat: the method conditions on proxy residuals, so its
|
||
// estimates carry more sampling noise than a maximum-likelihood fit
|
||
// would, the MA side most of all; tolerances set against it should be
|
||
// correspondingly looser. A pure AR model (q = 0) is routed to the
|
||
// Yule-Walker fit, which solves that problem exactly; p + q must be
|
||
// at least 1, the proxy order at least p+q+1, and the series long
|
||
// enough to leave a regression window worth fitting.
|
||
func EstimateARMA(x *core.Array, p, q int, opts ARMAOptions) (*ARMAResult, error) {
|
||
const name = "EstimateARMA"
|
||
if p < 0 || q < 0 {
|
||
return nil, base.Errf("%s: the orders must not be negative, got %d and %d", name, p, q)
|
||
}
|
||
if p+q < 1 {
|
||
return nil, base.Errf("%s: at least one coefficient is needed, got orders %d and %d", name, p, q)
|
||
}
|
||
if q == 0 {
|
||
return EstimateAR(x, p)
|
||
}
|
||
m := opts.HighAROrder
|
||
if m == 0 {
|
||
m = max(16, p+q+8)
|
||
}
|
||
if m < p+q+1 {
|
||
return nil, base.Errf("%s: the proxy AR order must be at least p+q+1 = %d, got %d", name, p+q+1, m)
|
||
}
|
||
r, n, mean, err := armaAutocovariance(name, x, m)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
rows := n - m - q
|
||
if rows < p+q+8 {
|
||
return nil, base.Errf("%s: %d samples leave only %d regression rows for %d coefficients; shorten the proxy or lengthen the series",
|
||
name, n, rows, p+q)
|
||
}
|
||
phiAR, _, err := yuleWalker(name, r[:m+1], m)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
demeaned := make([]float64, n)
|
||
for i := range n {
|
||
demeaned[i] = x.FloatAt(i) - mean
|
||
}
|
||
// The proxy AR's residuals, the innovations' stand-ins: e_t for
|
||
// t ≥ m.
|
||
resid := make([]float64, n-m)
|
||
for t := m; t < n; t++ {
|
||
v := demeaned[t]
|
||
for j := 1; j <= m; j++ {
|
||
v -= phiAR[j-1] * demeaned[t-j]
|
||
}
|
||
resid[t-m] = v
|
||
}
|
||
// The regression: x_t − μ on x_{t−1..t−p} − μ and e_{t−1..t−q},
|
||
// over the window t = m+q..n−1 the proxies cover.
|
||
T := n - m - q
|
||
design := make([][]float64, T)
|
||
target := make([]float64, T)
|
||
for i := range T {
|
||
t := m + q + i
|
||
row := make([]float64, 0, p+q)
|
||
for j := 1; j <= p; j++ {
|
||
row = append(row, demeaned[t-j])
|
||
}
|
||
for k := 1; k <= q; k++ {
|
||
row = append(row, resid[t-k-m])
|
||
}
|
||
design[i] = row
|
||
target[i] = demeaned[t]
|
||
}
|
||
coef, err := armaLstsq(name, design, target)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
rss := 0.0
|
||
for i := range T {
|
||
fitted := 0.0
|
||
for a, v := range design[i] {
|
||
fitted += v * coef[a]
|
||
}
|
||
d := target[i] - fitted
|
||
rss += d * d
|
||
}
|
||
sigma2 := rss / float64(T)
|
||
if !(sigma2 > 0) || math.IsNaN(sigma2) {
|
||
return nil, base.Errf("%s: the innovation variance came out %g: the fit carries no information", name, sigma2)
|
||
}
|
||
res := &ARMAResult{AR: coef[:p], MA: coef[p:], InnovationVariance: sigma2, Mean: mean, P: p, Q: q}
|
||
armaLikelihood(res, n)
|
||
return res, nil
|
||
}
|
||
|
||
// SelectARMA searches the order grid 0..maxAR × 0..maxMA for the ARMA
|
||
// model the chosen information criterion ranks best: every cell is
|
||
// fitted by EstimateARMA, the criterion (ARMAOptions' Criterion,
|
||
// "aic" by default) is computed from the innovation variance, and the
|
||
// smallest score wins. The (0, 0) cell carries no dynamics and is
|
||
// skipped; a cell whose fit fails is skipped too, and a grid where
|
||
// nothing can be fitted reports the last failure. The result carries
|
||
// both criteria of the winning model, so the runner-up is one more
|
||
// sweep away.
|
||
//
|
||
// The two criteria answer different questions: the AIC's linear
|
||
// parameter penalty makes it efficient but willing to overfit by an
|
||
// order with probability that does not vanish with the sample count,
|
||
// while the BIC's logarithmic penalty makes it consistent, choosing
|
||
// the true orders with probability tending to one. Where they
|
||
// disagree on long series, the BIC is the one to trust.
|
||
func SelectARMA(x *core.Array, maxAR, maxMA int, opts ARMAOptions) (*ARMAResult, error) {
|
||
const name = "SelectARMA"
|
||
if maxAR < 0 || maxMA < 0 {
|
||
return nil, base.Errf("%s: the grid bounds must not be negative, got %d and %d", name, maxAR, maxMA)
|
||
}
|
||
if maxAR+maxMA < 1 {
|
||
return nil, base.Errf("%s: the grid must hold at least one candidate, got %d×%d", name, maxAR+1, maxMA+1)
|
||
}
|
||
criterion := opts.Criterion
|
||
if criterion == "" {
|
||
criterion = "aic"
|
||
}
|
||
if criterion != "aic" && criterion != "bic" {
|
||
return nil, base.Errf("%s: the criterion must be %q or %q, got %q", name, "aic", "bic", opts.Criterion)
|
||
}
|
||
var (
|
||
best *ARMAResult
|
||
bestScore float64
|
||
lastErr error
|
||
)
|
||
for p := range maxAR + 1 {
|
||
for q := range maxMA + 1 {
|
||
if p == 0 && q == 0 {
|
||
continue
|
||
}
|
||
res, err := EstimateARMA(x, p, q, opts)
|
||
if err != nil {
|
||
lastErr = err
|
||
continue
|
||
}
|
||
score := res.AIC
|
||
if criterion == "bic" {
|
||
score = res.BIC
|
||
}
|
||
if best == nil || score < bestScore {
|
||
best, bestScore = res, score
|
||
}
|
||
}
|
||
}
|
||
if best == nil {
|
||
return nil, base.Errf("%s: no model on the grid could be fitted: %w", name, lastErr)
|
||
}
|
||
return best, nil
|
||
}
|
||
|
||
// ARMASpectrum evaluates the fitted model's theoretical power
|
||
// spectrum at the nFreq frequencies evenly spaced from 0 to the
|
||
// Nyquist frequency 0.5 inclusive, on the one-sided convention the
|
||
// house periodogram prints: the bins strictly between DC and Nyquist
|
||
// carry the folded double weight, DC and Nyquist the single one. The
|
||
// value is σ²·|Θ(e^{−2πif})|² / |Φ(e^{−2πif})|² with that weighting,
|
||
// so model and data line up bin for bin against WelchPSD at fs = 1:
|
||
// a white-noise sequence of variance σ² estimates σ² at the edges and
|
||
// 2σ² between them, and the fitted model's spectrum describes the
|
||
// data's own periodogram without a scale fudge. Fewer than two
|
||
// frequencies, a nil model, and a non-positive innovation variance or
|
||
// non-finite coefficient are errors; a pole of Φ on the unit circle
|
||
// names its frequency instead of answering infinities.
|
||
func ARMASpectrum(res *ARMAResult, nFreq int) (freqs, psd *core.Array, err error) {
|
||
const name = "ARMASpectrum"
|
||
if res == nil {
|
||
return nil, nil, base.Errf("%s: the model is nil", name)
|
||
}
|
||
if nFreq < 2 {
|
||
return nil, nil, base.Errf("%s: at least two frequencies are needed, got %d", name, nFreq)
|
||
}
|
||
if !(res.InnovationVariance > 0) || math.IsNaN(res.InnovationVariance) || math.IsInf(res.InnovationVariance, 0) {
|
||
return nil, nil, base.Errf("%s: the innovation variance must be positive and finite, got %g",
|
||
name, res.InnovationVariance)
|
||
}
|
||
vals := append(append([]float64{}, res.AR...), res.MA...)
|
||
if err := kfFinite(name, "the coefficients", vals); err != nil {
|
||
return nil, nil, err
|
||
}
|
||
freqs = core.New(core.Float, nFreq)
|
||
psd = core.New(core.Float, nFreq)
|
||
fRow := freqs.RawFloats()
|
||
pRow := psd.RawFloats()
|
||
for k := range nFreq {
|
||
f := 0.5 * float64(k) / float64(nFreq-1)
|
||
omega := 2 * math.Pi * f
|
||
// |Φ|² and |Θ|² by real arithmetic: the phase parts multiply
|
||
// out to squares of sine sums. Both loops read a sine and a
|
||
// cosine of the same argument; math.Sincos returns exactly the
|
||
// pair math.Sin and math.Cos produce (verified bit-for-bit), so
|
||
// the sums are unchanged while the trig work halves.
|
||
rePhi, imPhi := 1.0, 0.0
|
||
for j, phi := range res.AR {
|
||
s, c := math.Sincos(float64(j+1) * omega)
|
||
rePhi -= phi * c
|
||
imPhi += phi * s
|
||
}
|
||
reTheta, imTheta := 1.0, 0.0
|
||
for k2, theta := range res.MA {
|
||
s, c := math.Sincos(float64(k2+1) * omega)
|
||
reTheta += theta * c
|
||
imTheta -= theta * s
|
||
}
|
||
den := rePhi*rePhi + imPhi*imPhi
|
||
if den == 0 {
|
||
return nil, nil, base.Errf("%s: the model has a pole on the unit circle at frequency %g", name, f)
|
||
}
|
||
fRow[k] = f
|
||
pRow[k] = res.InnovationVariance * (reTheta*reTheta + imTheta*imTheta) / den
|
||
if k > 0 && k < nFreq-1 {
|
||
pRow[k] *= 2 // the one-sided fold between DC and Nyquist
|
||
}
|
||
}
|
||
return freqs, psd, nil
|
||
}
|