feat: initial release

Assisted-by: GLM 5.3 Flash

signal/correlate.go

@@ -0,0 +1,250 @@
+// 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"
+)
+
+// Correlation functions: the time-series staples on top of
+// the FFT. The autocorrelation and cross-correlation run as spectral
+// products (one transform pair whatever the lag count), the partial
+// autocorrelation follows by Durbin-Levinson recursion over the ACF,
+// the algorithm that defined the PACF in the first place.
+
+// Autocorrelate returns the sample autocorrelation of a rank-1 real
+// signal at lags 0..maxLag: the mean is removed first and every value
+// is normalised by the lag-zero sum (the biased estimator, the one
+// whose Durbin-Levinson and Bartlett-band theory is written for).
+// maxLag must lie in [0, n−1]; the result has maxLag+1 entries and
+// starts at 1, unless the signal is constant: its zero power makes
+// every lag 0, lag zero included.
+func Autocorrelate(x *core.Array, maxLag int) (*core.Array, error) {
+	const name = "Autocorrelate"
+	if x.NDim() != 1 {
+		return nil, base.Errf("%s: needs a rank-1 signal, got shape %s", name, base.ShapeText(x.Shape()))
+	}
+	if x.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex signals are not supported", name)
+	}
+	n := x.Len()
+	if n == 0 {
+		return nil, base.Errf("%s: an empty signal has no correlation", name)
+	}
+	if maxLag < 0 || maxLag >= n {
+		return nil, base.Errf("%s: maxLag must lie in [0, %d], got %d", name, n-1, maxLag)
+	}
+	// Demean. The payload walk reads the same values FloatAt would
+	// (widenFloats widens exactly), in the same order, so the mean
+	// and every padded sample keep their bits.
+	xF := x.RawFloats()
+	if xF == nil || x.Strided() {
+		xF = widenFloats(x)
+	}
+	// A non-finite sample would drive the normaliser NaN and publish
+	// NaN correlations with no error, so it is refused up front (the
+	// guard SolvePoissonPeriodic applies to its source).
+	for i := range n {
+		if v := xF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
+			return nil, base.Errf("%s: x holds the non-finite value %g at %d", name, v, i)
+		}
+	}
+	mean := 0.0
+	for i := range n {
+		mean += xF[i]
+	}
+	mean /= float64(n)
+	// The correlation runs as one transform pair over the zero-padded
+	// signal, in place on a private buffer: the values that reach the
+	// butterflies and the arithmetic that follows are the FFT/IFFT
+	// path's unchanged, minus the copies into and out of arrays that
+	// nobody else ever read.
+	padded := make([]complex128, 2*n)
+	for i := range n {
+		padded[i] = complex(xF[i]-mean, 0)
+	}
+	transform(padded, -1)
+	// |F|²: multiply by the conjugate, then back. The zero-padded
+	// transform makes the circular correlation linear.
+	for i := range padded {
+		padded[i] = padded[i] * conj128(padded[i])
+	}
+	transform(padded, +1)
+	// The inverse scaling, applied with IFFT's own operation order.
+	m := float64(2 * n)
+	for i := range padded {
+		padded[i] /= complex(m, 0)
+	}
+	den := real(padded[0])
+	out := make([]float64, maxLag+1)
+	for k := range maxLag + 1 {
+		if den == 0 {
+			out[k] = 0
+			continue
+		}
+		out[k] = real(padded[k]) / den
+	}
+	return floatsFromArrayMust(out, []int{maxLag + 1}), nil
+}
+
+// conj128 conjugates one complex number (the local spelling keeps the
+// dependency surface to the packages already imported).
+func conj128(z complex128) complex128 { return complex(real(z), -imag(z)) }
+
+// CrossCorrelate returns the raw cross-correlation of two equal-length
+// real signals at every lag: the result has 2n−1 entries, index i
+// carrying lag i−(n−1), value r(lag) = Σ_t x[t+lag]·y[t] over the
+// overlapping t. No mean removal, no normalisation: correlation as
+// the inner product of shifted copies, the definition convolution and
+// matched filtering share.
+func CrossCorrelate(x, y *core.Array) (*core.Array, error) {
+	const name = "CrossCorrelate"
+	if x.NDim() != 1 || y.NDim() != 1 {
+		return nil, base.Errf("%s: needs two rank-1 signals, got %s and %s",
+			name, base.ShapeText(x.Shape()), base.ShapeText(y.Shape()))
+	}
+	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex signals are not supported", name)
+	}
+	n := x.Len()
+	if n == 0 || y.Len() != n {
+		return nil, base.Errf("%s: the signals must be equal-length non-empty, got %d and %d",
+			name, n, y.Len())
+	}
+	// Zero-pad both signals through the raw payloads where possible;
+	// the widened reads equal the FloatAt values bit for bit.
+	xF := x.RawFloats()
+	if xF == nil || x.Strided() {
+		xF = widenFloats(x)
+	}
+	yF := y.RawFloats()
+	if yF == nil || y.Strided() {
+		yF = widenFloats(y)
+	}
+	// A non-finite sample would publish NaN correlations with no
+	// error, so both signals are refused up front.
+	for i := range n {
+		if v := xF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
+			return nil, base.Errf("%s: x holds the non-finite value %g at %d", name, v, i)
+		}
+		if v := yF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
+			return nil, base.Errf("%s: y holds the non-finite value %g at %d", name, v, i)
+		}
+	}
+	// The spectra multiply in place on private buffers: the values
+	// reaching the transforms and the per-element products are the
+	// FFT/mulSpectraConj/IFFT path's unchanged, minus the copies into
+	// and out of arrays nobody else ever read.
+	px := make([]complex128, 2*n)
+	py := make([]complex128, 2*n)
+	for i := range n {
+		px[i] = complex(xF[i], 0)
+		py[i] = complex(yF[i], 0)
+	}
+	transform(px, -1)
+	transform(py, -1)
+	// Fx·conj(Fy) transforms back to Σ_t x[t+lag]·y[t], the documented
+	// orientation; the index mapping lands the positive lags in the
+	// first half.
+	// The product runs in mulSpectraConj, the array-based helper the
+	// FFT path always used: the compiler's FMA fusion of the complex
+	// multiply depends on the compiled shape of its surroundings, and
+	// the helper's shape is the one whose product bits this result is
+	// pinned to. The wrappers take the buffers without copying; the
+	// product array is the one allocation of the old path this keeps.
+	prod, perr := mulSpectraConj(complexFromArrayMust(px, []int{2 * n}),
+		complexFromArrayMust(py, []int{2 * n}))
+	if perr != nil {
+		return nil, base.Errf("%s: %w", name, perr)
+	}
+	pc := prod.RawComplexes()
+	transform(pc, +1)
+	m := float64(2 * n)
+	for i := range pc {
+		pc[i] /= complex(m, 0)
+	}
+	// Lag 0 sits at index 0 of the transformed product, lag k at k and
+	// negative lag −k at 2n−k; the output centres them on index n−1.
+	out := make([]float64, 2*n-1)
+	for k := range 2*n - 1 {
+		lag := k - (n - 1)
+		if lag >= 0 {
+			out[k] = real(pc[lag])
+		} else {
+			out[k] = real(pc[2*n+lag])
+		}
+	}
+	return floatsFromArrayMust(out, []int{2*n - 1}), nil
+}
+
+// mulSpectraConj multiplies the first spectrum by the conjugate of the
+// second, element-wise. Both spectra come from the FFT and stream
+// their complex payloads directly; the accessor loop stays as the
+// fallback for anything else. The function is kept byte-for-byte as
+// the FFT path always defined it: its compiled FMA pattern is what the
+// correlation results are pinned to.
+func mulSpectraConj(a, b *core.Array) (*core.Array, error) {
+	as := a.RawComplexes()
+	bs := b.RawComplexes()
+	if as != nil && bs != nil && !a.Strided() && !b.Strided() {
+		out := make([]complex128, len(as))
+		for i := range out {
+			out[i] = as[i] * conj128(bs[i])
+		}
+		return core.ComplexFromArray(out, a.Shape()...)
+	}
+	out := make([]complex128, a.Len())
+	for i := range out {
+		out[i] = a.ComplexAt(i) * conj128(b.ComplexAt(i))
+	}
+	return core.ComplexFromArray(out, a.Shape()...)
+}
+
+// PartialAutocorrelate returns the partial autocorrelation of a rank-1
+// real signal at lags 1..maxLag, by the Durbin-Levinson recursion over
+// the biased autocorrelation: pacf[k] is the last coefficient of the
+// order-k AR fit, the "correlation with the intermediate lags
+// regressed out" that box-jenkins identification reads. maxLag must
+// lie in [1, n−2] for the recursion to stay meaningful.
+func PartialAutocorrelate(x *core.Array, maxLag int) (*core.Array, error) {
+	const name = "PartialAutocorrelate"
+	n := x.Len()
+	if maxLag < 1 || maxLag > n-2 {
+		return nil, base.Errf("%s: maxLag must lie in [1, %d], got %d", name, n-2, maxLag)
+	}
+	acf, err := Autocorrelate(x, maxLag)
+	if err != nil {
+		return nil, err
+	}
+	// r[k] = acf at lag k (1-based recursion indexing).
+	r := make([]float64, maxLag+1)
+	r[0] = 1
+	for k := range maxLag {
+		r[k+1] = acf.FloatAt(k + 1)
+	}
+	phi := make([]float64, maxLag+1) // current order's coefficients
+	phiPrev := make([]float64, maxLag+1)
+	out := make([]float64, maxLag)
+	for k := 1; k <= maxLag; k++ {
+		num := r[k]
+		den := 1.0
+		for j := 1; j < k; j++ {
+			num -= phiPrev[j] * r[k-j]
+			den -= phiPrev[j] * r[j]
+		}
+		if den == 0 {
+			return nil, base.Errf("%s: the recursion broke down at lag %d", name, k)
+		}
+		phi[k] = num / den
+		for j := 1; j < k; j++ {
+			phi[j] = phiPrev[j] - phi[k]*phiPrev[k-j]
+		}
+		copy(phiPrev, phi)
+		out[k-1] = phi[k]
+	}
+	return floatsFromArrayMust(out, []int{maxLag}), nil
+}