feat: initial release

Assisted-by: GLM 5.3 Flash

signal/arma.go

@@ -0,0 +1,476 @@
+// 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
+}