feat: initial release

Assisted-by: GLM 5.3 Flash

signal/kalman.go

@@ -0,0 +1,941 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package signal
+
+import (
+	"math"
+	"slices"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// State-space filtering: the Kalman family over the model
+// x_{t+1} = f(x_t) + w_t with w ~ N(0, Q), z_t = h(x_t) + v_t with
+// v ~ N(0, R). The linear filter runs the standard Riccati recursion
+// with a Joseph-form correction, the extended filter linearises f and
+// h at the current estimate (analytic Jacobians when supplied, the
+// house central-difference helper otherwise), and the unscented filter
+// carries the state distribution through the deterministic sigma-point
+// set. All three accumulate the exact Gaussian log-likelihood of the
+// innovation sequence, and every covariance they publish is mirrored
+// into its symmetric average, with positive definiteness enforced
+// where the mathematics demands it: a Cholesky factorisation that
+// cannot be taken names the step it failed at.
+
+// StateFunc maps one state vector to the image vector the model's
+// transition or observation applies: the f of x_{t+1} = f(x_t) + w_t,
+// the h of z_t = h(x_t) + v_t. The input array is private to the call
+// and may be kept until the call returns; the output must be a real
+// rank-1 array of one fixed length.
+type StateFunc func(x *core.Array) (*core.Array, error)
+
+// JacobianFunc returns the Jacobian of a StateFunc at x as an
+// (m × n) float array, row i holding the partials of output i. An
+// analytic Jacobian is always the better instrument; when the option
+// is left unset the extended filter builds one by central differences
+// at every step.
+type JacobianFunc func(x *core.Array) (*core.Array, error)
+
+// KalmanOptions carries the initial condition, the noise levels and
+// the filter-specific knobs. An unset (nil) array field takes its
+// documented default: a zero initial state, an identity initial
+// covariance, zero process noise, an identity measurement noise.
+// SigmaAlpha, SigmaBeta and SigmaKappa follow the same rule with the
+// defaults 0.001, 2 and 0, the standard scaled unscented choice for a
+// Gaussian prior. The extended and unscented filters cannot infer the
+// state dimension from their callbacks, so at least one of
+// InitialState and InitialCovariance must be set for them.
+//
+// TransitionJacobian and ObservationJacobian give the extended filter
+// the analytic partials it linearises with; a nil one is replaced by
+// central differences on the underlying map, two evaluations per
+// partial per step.
+type KalmanOptions struct {
+	InitialState        *core.Array
+	InitialCovariance   *core.Array
+	ProcessNoise        *core.Array
+	MeasurementNoise    *core.Array
+	TransitionJacobian  JacobianFunc
+	ObservationJacobian JacobianFunc
+	SigmaAlpha          float64
+	SigmaBeta           float64
+	SigmaKappa          float64
+}
+
+// KalmanResult holds one filtering pass over the measurement stack.
+type KalmanResult struct {
+	// States is the (n × d) stack of filtered means x̂_{t|t}, one row
+	// per measurement.
+	States *core.Array
+	// Covariances is the (n × d × d) stack of filtered covariances
+	// P_{t|t}, one symmetric positive-definite block per measurement.
+	Covariances *core.Array
+	// Innovations is the (n × m) stack of one-step prediction errors
+	// z_t − h(x̂_{t|t−1}).
+	Innovations *core.Array
+	// InnovationCovariances is the (n × m × m) stack of the innovation
+	// covariances S_t the likelihood reads.
+	InnovationCovariances *core.Array
+	// LogLikelihood is Σ_t log N(z_t; h(x̂_{t|t−1}), S_t), the exact
+	// Gaussian likelihood of the measurement sequence under the model
+	// and the quantity noise and parameter estimation maximises.
+	LogLikelihood float64
+}
+
+// kalmanStepOut carries one step's outputs from a filter's step
+// closure into the result stack.
+type kalmanStepOut struct {
+	mean          []float64
+	covariance    []float64
+	innovation    []float64
+	innovationCov []float64
+	logLikelihood float64
+}
+
+// kfSymmetryEps is the relative tolerance a covariance's mirror check
+// allows, the same reading the multivariate normal takes: matrices
+// assembled from products differ from their mirror by an ulp of
+// rounding, a genuinely asymmetric pair by far more.
+const kfSymmetryEps = 1e-12
+
+// kfFinite refuses the non-finite values a filter would otherwise
+// carry silently through every recursion.
+func kfFinite(name, what string, vals []float64) error {
+	for i, v := range vals {
+		if math.IsNaN(v) || math.IsInf(v, 0) {
+			return base.Errf("%s: %s holds the non-finite value %g at %d", name, what, v, i)
+		}
+	}
+	return nil
+}
+
+// kfVector reads a finite real rank-1 array into a float slice, the
+// payload itself for a contiguous float64 array,
+// with copies made only where the payload cannot serve, of exactly want entries when want is non-negative.
+func kfVector(name, what string, a *core.Array, want int) ([]float64, error) {
+	if a.NDim() != 1 {
+		return nil, base.Errf("%s: %s must be a vector, got shape %s", name, what, base.ShapeText(a.Shape()))
+	}
+	if a.Dtype() == core.Complex {
+		return nil, base.Errf("%s: %s must be real, got complex", name, what)
+	}
+	if want >= 0 && a.Len() != want {
+		return nil, base.Errf("%s: %s holds %d entries, want %d", name, what, a.Len(), want)
+	}
+	vals := widenFloats(a)
+	if err := kfFinite(name, what, vals); err != nil {
+		return nil, err
+	}
+	return vals, nil
+}
+
+// kfMatrix reads a finite real rank-2 array into a fresh row-major
+// float slice of exactly rows×cols, either side wildcarded at −1.
+func kfMatrix(name, what string, a *core.Array, rows, cols int) ([]float64, error) {
+	if a.NDim() != 2 {
+		return nil, base.Errf("%s: %s must be rank 2, got shape %s", name, what, base.ShapeText(a.Shape()))
+	}
+	if a.Dtype() == core.Complex {
+		return nil, base.Errf("%s: %s must be real, got complex", name, what)
+	}
+	shape := a.Shape()
+	if (rows >= 0 && shape[0] != rows) || (cols >= 0 && shape[1] != cols) {
+		return nil, base.Errf("%s: %s has shape %s, want %d×%d", name, what, base.ShapeText(shape), rows, cols)
+	}
+	vals := widenFloats(a)
+	if err := kfFinite(name, what, vals); err != nil {
+		return nil, err
+	}
+	return vals, nil
+}
+
+// kfSymmetric demands every mirror pair agree within a relative
+// tolerance, because only one triangle is ever read.
+func kfSymmetric(name, what string, a []float64, n int) error {
+	for i := range n {
+		for j := range i {
+			lo, hi := a[i*n+j], a[j*n+i]
+			if math.Abs(lo-hi) > kfSymmetryEps*math.Max(math.Abs(lo), math.Abs(hi)) {
+				return base.Errf("%s: %s is not symmetric at (%d, %d): %g against %g",
+					name, what, i+1, j+1, lo, hi)
+			}
+		}
+	}
+	return nil
+}
+
+// kfNoise reads one noise covariance: an unset array takes its
+// documented default (zeros for the process noise, whose only job is
+// to enter sums, an identity for the measurement noise, which is
+// factored every step and so must be positive definite), a set one
+// must be a finite symmetric matrix of the right size.
+func kfNoise(name, what string, a *core.Array, dim int, identityDefault bool) ([]float64, error) {
+	if a == nil {
+		if identityDefault {
+			return kfIdentity(dim), nil
+		}
+		return make([]float64, dim*dim), nil
+	}
+	vals, err := kfMatrix(name, what, a, dim, dim)
+	if err != nil {
+		return nil, err
+	}
+	if err := kfSymmetric(name, what, vals, dim); err != nil {
+		return nil, err
+	}
+	if identityDefault {
+		if _, err := kfCholesky(what, vals, dim); err != nil {
+			return nil, base.Errf("%s: %w", name, err)
+		}
+	}
+	return vals, nil
+}
+
+// kfInitialCondition reads the initial mean and covariance, zeros and
+// identity where the options leave them unset. The initial covariance
+// must be symmetric positive definite: the first predict would
+// otherwise hand the recursion a structure it cannot factor.
+func kfInitialCondition(name string, opts KalmanOptions, d int) (x0, p0 []float64, err error) {
+	if opts.InitialState == nil {
+		x0 = make([]float64, d)
+	} else if x0, err = kfVector(name, "the initial state", opts.InitialState, d); err != nil {
+		return nil, nil, err
+	}
+	if opts.InitialCovariance == nil {
+		p0 = kfIdentity(d)
+	} else {
+		if p0, err = kfMatrix(name, "the initial covariance", opts.InitialCovariance, d, d); err != nil {
+			return nil, nil, err
+		}
+		if err := kfSymmetric(name, "the initial covariance", p0, d); err != nil {
+			return nil, nil, err
+		}
+		if _, err := kfCholesky("the initial covariance", p0, d); err != nil {
+			return nil, nil, base.Errf("%s: %w", name, err)
+		}
+	}
+	return x0, p0, nil
+}
+
+// kfNonlinearInputs resolves the shared inputs of the extended and
+// unscented filters: the state dimension, which the callbacks cannot
+// carry, comes from the initial state or the initial covariance, at
+// least one of which must be set.
+func kfNonlinearInputs(name string, opts KalmanOptions, m int) (d int, x0, p0, q, r []float64, err error) {
+	switch {
+	case opts.InitialState != nil:
+		if x0, err = kfVector(name, "the initial state", opts.InitialState, -1); err != nil {
+			return 0, nil, nil, nil, nil, err
+		}
+		d = len(x0)
+	case opts.InitialCovariance != nil:
+		if opts.InitialCovariance.NDim() != 2 {
+			return 0, nil, nil, nil, nil, base.Errf("%s: the initial covariance must be rank 2, got shape %s",
+				name, base.ShapeText(opts.InitialCovariance.Shape()))
+		}
+		d = opts.InitialCovariance.Shape()[0]
+	default:
+		return 0, nil, nil, nil, nil, base.Errf("%s: the state dimension must come from an initial state or an initial covariance; neither is set", name)
+	}
+	if d < 1 {
+		return 0, nil, nil, nil, nil, base.Errf("%s: the state dimension must be at least 1, got %d", name, d)
+	}
+	if x0, p0, err = kfInitialCondition(name, opts, d); err != nil {
+		return 0, nil, nil, nil, nil, err
+	}
+	if q, err = kfNoise(name, "the process noise", opts.ProcessNoise, d, false); err != nil {
+		return 0, nil, nil, nil, nil, err
+	}
+	if r, err = kfNoise(name, "the measurement noise", opts.MeasurementNoise, m, true); err != nil {
+		return 0, nil, nil, nil, nil, err
+	}
+	return d, x0, p0, q, r, nil
+}
+
+// kfMeasurements reads the measurement stack: a rank-1 array holds n
+// scalar observations, a rank-2 array holds n rows of m channels. At
+// least one measurement is needed, every value must be finite.
+func kfMeasurements(name string, z *core.Array) (n, m int, rows []float64, err error) {
+	if z.NDim() != 1 && z.NDim() != 2 {
+		return 0, 0, nil, base.Errf("%s: the measurements must be rank 1 or rank 2, got shape %s",
+			name, base.ShapeText(z.Shape()))
+	}
+	if z.Dtype() == core.Complex {
+		return 0, 0, nil, base.Errf("%s: complex measurements are not supported", name)
+	}
+	if z.NDim() == 1 {
+		m = 1
+	} else {
+		m = z.Shape()[1]
+	}
+	if m < 1 {
+		return 0, 0, nil, base.Errf("%s: the measurement width must be at least 1, got %d", name, m)
+	}
+	n = z.Len() / m
+	if n < 1 {
+		return 0, 0, nil, base.Errf("%s: at least one measurement is needed", name)
+	}
+	rows = widenFloats(z)
+	if err := kfFinite(name, "the measurements", rows); err != nil {
+		return 0, 0, nil, err
+	}
+	return n, m, rows, nil
+}
+
+// kfIdentity returns the n×n identity, row-major.
+func kfIdentity(n int) []float64 {
+	a := make([]float64, n*n)
+	for i := range n {
+		a[i*n+i] = 1
+	}
+	return a
+}
+
+// kfTranspose returns the transpose of the row-major rows×cols
+// matrix.
+func kfTranspose(a []float64, rows, cols int) []float64 {
+	out := make([]float64, rows*cols)
+	for i := range rows {
+		for j := range cols {
+			out[j*rows+i] = a[i*cols+j]
+		}
+	}
+	return out
+}
+
+// kfMatVec returns A·x for the row-major rows×cols matrix A.
+func kfMatVec(a []float64, rows, cols int, x []float64) []float64 {
+	out := make([]float64, rows)
+	for i := range rows {
+		total := 0.0
+		row := a[i*cols : (i+1)*cols]
+		for j, v := range row {
+			total += v * x[j]
+		}
+		out[i] = total
+	}
+	return out
+}
+
+// kfMatMul returns A·B for the row-major A of size ra×ca and B of
+// size ca×cb.
+func kfMatMul(a []float64, ra, ca int, b []float64, cb int) []float64 {
+	out := make([]float64, ra*cb)
+	for i := range ra {
+		row := out[i*cb : (i+1)*cb]
+		arow := a[i*ca : (i+1)*ca]
+		for k, aik := range arow {
+			brow := b[k*cb : (k+1)*cb]
+			for j := range cb {
+				row[j] += aik * brow[j]
+			}
+		}
+	}
+	return out
+}
+
+// kfCholesky factors the symmetric positive-definite row-major n×n
+// matrix into the lower triangular L with A = L·Lᵀ, only the lower
+// mirror read. A non-positive pivot names its row.
+func kfCholesky(what string, a []float64, n int) ([]float64, error) {
+	l := make([]float64, n*n)
+	for i := range n {
+		for j := range i + 1 {
+			total := a[i*n+j]
+			for k := range j {
+				total -= l[i*n+k] * l[j*n+k]
+			}
+			if i == j {
+				if !(total > 0) {
+					return nil, base.Errf("%s is not positive definite at row %d", what, i+1)
+				}
+				l[i*n+i] = math.Sqrt(total)
+			} else {
+				l[i*n+j] = total / l[j*n+j]
+			}
+		}
+	}
+	return l, nil
+}
+
+// kfCholSolve solves L·Lᵀ·x = b through the forward and backward
+// substitutions.
+func kfCholSolve(l []float64, n int, b []float64) []float64 {
+	x := make([]float64, n)
+	for i := range n {
+		total := b[i]
+		for k := range i {
+			total -= l[i*n+k] * x[k]
+		}
+		x[i] = total / l[i*n+i]
+	}
+	for i := n - 1; i >= 0; i-- {
+		total := x[i]
+		for k := i + 1; k < n; k++ {
+			total -= l[k*n+i] * x[k]
+		}
+		x[i] = total / l[i*n+i]
+	}
+	return x
+}
+
+// kfCholSolveMatrix solves L·Lᵀ·X = B for the row-major B of size
+// n×cols, row by row.
+func kfCholSolveMatrix(l []float64, n int, b []float64, cols int) []float64 {
+	x := make([]float64, len(b))
+	copy(x, b)
+	for i := range n {
+		row := x[i*cols : (i+1)*cols]
+		for k := range i {
+			lk := l[i*n+k]
+			krow := x[k*cols : (k+1)*cols]
+			for j := range cols {
+				row[j] -= lk * krow[j]
+			}
+		}
+		di := l[i*n+i]
+		for j := range cols {
+			row[j] /= di
+		}
+	}
+	for i := n - 1; i >= 0; i-- {
+		row := x[i*cols : (i+1)*cols]
+		for k := i + 1; k < n; k++ {
+			lk := l[k*n+i]
+			krow := x[k*cols : (k+1)*cols]
+			for j := range cols {
+				row[j] -= lk * krow[j]
+			}
+		}
+		di := l[i*n+i]
+		for j := range cols {
+			row[j] /= di
+		}
+	}
+	return x
+}
+
+// kfLogDet returns the log determinant of the Cholesky factor: twice
+// the sum of the log diagonal.
+func kfLogDet(l []float64, n int) float64 {
+	total := 0.0
+	for i := range n {
+		total += math.Log(l[i*n+i])
+	}
+	return 2 * total
+}
+
+// kfSymmetrise replaces A by (A + Aᵀ)/2 in place: the rounding drift
+// of a recursion that touches a covariance only through symmetric
+// expressions cannot survive the mirror average.
+func kfSymmetrise(a []float64, n int) {
+	for i := range n {
+		for j := range i {
+			avg := (a[i*n+j] + a[j*n+i]) / 2
+			a[i*n+j] = avg
+			a[j*n+i] = avg
+		}
+	}
+}
+
+// kalmanRun walks the measurement stack through the step closure,
+// which owns one predict-and-correct cycle: it receives the current
+// mean and covariance (read-only; it must return fresh slices) and
+// the step's measurement row, and returns everything the result stack
+// records, the Gaussian log-likelihood contribution included.
+func kalmanRun(nMeas, m, d int, meas []float64, x0, p0 []float64,
+	step func(t int, zRow, mean, cov []float64) (kalmanStepOut, error),
+) (*KalmanResult, error) {
+	out := &KalmanResult{
+		States:                core.New(core.Float, nMeas, d),
+		Covariances:           core.New(core.Float, nMeas, d, d),
+		Innovations:           core.New(core.Float, nMeas, m),
+		InnovationCovariances: core.New(core.Float, nMeas, m, m),
+	}
+	mean := slices.Clone(x0)
+	cov := slices.Clone(p0)
+	total := 0.0
+	for t := range nMeas {
+		res, err := step(t, meas[t*m:(t+1)*m], mean, cov)
+		if err != nil {
+			return nil, err
+		}
+		copy(out.States.RawFloats()[t*d:(t+1)*d], res.mean)
+		copy(out.Covariances.RawFloats()[t*d*d:(t+1)*d*d], res.covariance)
+		copy(out.Innovations.RawFloats()[t*m:(t+1)*m], res.innovation)
+		copy(out.InnovationCovariances.RawFloats()[t*m*m:(t+1)*m*m], res.innovationCov)
+		total += res.logLikelihood
+		mean, cov = res.mean, res.covariance
+	}
+	out.LogLikelihood = total
+	return out, nil
+}
+
+// kfCorrect runs the shared correction of the linear and extended
+// filters: the innovation against the predicted observation, its
+// covariance S = H·P⁻·Hᵀ + R, the gain K = P⁻·Hᵀ·S⁻¹ through the
+// Cholesky solve, the state update and the Joseph-form covariance
+// (I−K·H)·P⁻·(I−K·H)ᵀ + K·R·Kᵀ, mirrored into its symmetric average.
+// The Joseph form keeps the filtered covariance symmetric positive
+// definite by construction, where the plain (I−K·H)·P⁻ recursion
+// preserves it only in exact arithmetic.
+func kfCorrect(name string, t int, zRow, predicted, xp, pp []float64, d, m int,
+	h, ht, r []float64,
+) (kalmanStepOut, error) {
+	innovation := make([]float64, m)
+	for i := range m {
+		innovation[i] = zRow[i] - predicted[i]
+	}
+	hp := kfMatMul(h, m, d, pp, d)
+	s := kfMatMul(hp, m, d, ht, m)
+	for i := range s {
+		s[i] += r[i]
+	}
+	ls, err := kfCholesky("the innovation covariance", s, m)
+	if err != nil {
+		return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t, err)
+	}
+	solved := kfCholSolve(ls, m, innovation)
+	quad := 0.0
+	for i := range m {
+		quad += innovation[i] * solved[i]
+	}
+	// K = P⁻·Hᵀ·S⁻¹ arrives through the transposed solve: S·X =
+	// H·P⁻ gives X = S⁻¹·H·P⁻ = Kᵀ and K = Xᵀ.
+	gain := kfTranspose(kfCholSolveMatrix(ls, m, hp, d), m, d)
+	mean := make([]float64, d)
+	copy(mean, xp)
+	for i := range d {
+		total := 0.0
+		grow := gain[i*m : (i+1)*m]
+		for j, y := range innovation {
+			total += grow[j] * y
+		}
+		mean[i] += total
+	}
+	kh := kfMatMul(gain, d, m, h, d)
+	a := make([]float64, d*d)
+	for i := range d {
+		for j := range d {
+			a[i*d+j] = -kh[i*d+j]
+		}
+		a[i*d+i]++
+	}
+	at := kfTranspose(a, d, d)
+	cov := kfMatMul(kfMatMul(a, d, d, pp, d), d, d, at, d)
+	krkt := kfMatMul(kfMatMul(gain, d, m, r, m), d, m, kfTranspose(gain, d, m), d)
+	for i := range cov {
+		cov[i] += krkt[i]
+	}
+	kfSymmetrise(cov, d)
+	logLik := -0.5 * (float64(m)*math.Log(2*math.Pi) + kfLogDet(ls, m) + quad)
+	return kalmanStepOut{mean: mean, covariance: cov, innovation: innovation,
+		innovationCov: s, logLikelihood: logLik}, nil
+}
+
+// KalmanFilter runs the linear Kalman filter over the measurement
+// stack z under the model x_{t+1} = F·x_t + w_t, z_t = H·x_t + v_t:
+// transition is the (d × d) state matrix F, observation the (m × d)
+// matrix H. Every step predicts with F and corrects with a
+// Joseph-form update, and the exact Gaussian log-likelihood of the
+// innovation sequence accumulates into the result.
+//
+// A rank-1 z holds n scalar observations, a rank-2 z holds n rows of
+// m channels. Nil option fields take their defaults (see
+// KalmanOptions). A non-square transition, a shape mismatch, an
+// asymmetric noise covariance, a singular measurement noise, a
+// non-finite value anywhere, and an innovation covariance that loses
+// positive definiteness (naming the step) are errors.
+func KalmanFilter(z *core.Array, transition, observation *core.Array, opts KalmanOptions) (*KalmanResult, error) {
+	const name = "KalmanFilter"
+	if transition == nil || observation == nil {
+		return nil, base.Errf("%s: the transition and observation matrices are required", name)
+	}
+	if transition.NDim() != 2 {
+		return nil, base.Errf("%s: the transition matrix must be rank 2, got shape %s",
+			name, base.ShapeText(transition.Shape()))
+	}
+	shape := transition.Shape()
+	if shape[0] != shape[1] {
+		return nil, base.Errf("%s: the transition matrix must be square, got %d×%d", name, shape[0], shape[1])
+	}
+	d := shape[0]
+	if d < 1 {
+		return nil, base.Errf("%s: the state dimension must be at least 1, got %d", name, d)
+	}
+	nMeas, m, meas, err := kfMeasurements(name, z)
+	if err != nil {
+		return nil, err
+	}
+	f, err := kfMatrix(name, "the transition matrix", transition, d, d)
+	if err != nil {
+		return nil, err
+	}
+	h, err := kfMatrix(name, "the observation matrix", observation, m, d)
+	if err != nil {
+		return nil, err
+	}
+	q, err := kfNoise(name, "the process noise", opts.ProcessNoise, d, false)
+	if err != nil {
+		return nil, err
+	}
+	r, err := kfNoise(name, "the measurement noise", opts.MeasurementNoise, m, true)
+	if err != nil {
+		return nil, err
+	}
+	x0, p0, err := kfInitialCondition(name, opts, d)
+	if err != nil {
+		return nil, err
+	}
+	ft := kfTranspose(f, d, d)
+	ht := kfTranspose(h, m, d)
+	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
+		// Predict: x⁻ = F·x, P⁻ = F·P·Fᵀ + Q.
+		xp := kfMatVec(f, d, d, x)
+		pp := kfMatMul(kfMatMul(f, d, d, p, d), d, d, ft, d)
+		for i := range pp {
+			pp[i] += q[i]
+		}
+		predicted := kfMatVec(h, m, d, xp)
+		return kfCorrect(name, t+1, zRow, predicted, xp, pp, d, m, h, ht, r)
+	})
+}
+
+// ExtendedKalmanFilter runs the extended Kalman filter over the
+// measurement stack z under the nonlinear model x_{t+1} = f(x_t) +
+// w_t, z_t = h(x_t) + v_t: every step predicts by propagating the
+// mean through f and the covariance through the linearisation F =
+// ∂f/∂x at x̂_{t|t}, then corrects through the Jacobian H = ∂h/∂x at
+// the predicted mean, with the same Joseph-form update and the same
+// exact log-likelihood the linear filter carries. The Jacobians come
+// from the options when supplied analytically and from central
+// differences otherwise (see JacobianFunc).
+//
+// The extended filter is the linear one applied to local linear
+// models: it inherits the Kalman recursions and with them the first
+// order's blindness to the curvature of f and h, so a strongly bent
+// observation map wants the unscented filter instead. The state
+// dimension must be fixed by an initial state or covariance (see
+// KalmanOptions); everything else follows the linear filter's
+// contract, with a failing callback or Jacobian reported with the
+// step it failed at.
+func ExtendedKalmanFilter(z *core.Array, transition, observation StateFunc, opts KalmanOptions) (*KalmanResult, error) {
+	const name = "ExtendedKalmanFilter"
+	if transition == nil || observation == nil {
+		return nil, base.Errf("%s: the transition and observation maps are required", name)
+	}
+	nMeas, m, meas, err := kfMeasurements(name, z)
+	if err != nil {
+		return nil, err
+	}
+	d, x0, p0, q, r, err := kfNonlinearInputs(name, opts, m)
+	if err != nil {
+		return nil, err
+	}
+	fjac := opts.TransitionJacobian
+	if fjac == nil {
+		fjac = func(x *core.Array) (*core.Array, error) {
+			return core.Jacobian(transition, x, core.JacobianOptions{})
+		}
+	}
+	hjac := opts.ObservationJacobian
+	if hjac == nil {
+		hjac = func(x *core.Array) (*core.Array, error) {
+			return core.Jacobian(observation, x, core.JacobianOptions{})
+		}
+	}
+	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
+		xArr, err := core.FromFloats(x, d)
+		if err != nil {
+			return kalmanStepOut{}, err
+		}
+		fOut, err := transition(xArr)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: the transition: %w", name, t+1, err)
+		}
+		xp, err := kfVector(name, "the transition output", fOut, d)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		fJacArr, err := fjac(xArr)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: the transition Jacobian: %w", name, t+1, err)
+		}
+		fMat, err := kfMatrix(name, "the transition Jacobian", fJacArr, d, d)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		// Predict through the linearisation at the filtered mean.
+		pp := kfMatMul(kfMatMul(fMat, d, d, p, d), d, d, kfTranspose(fMat, d, d), d)
+		for i := range pp {
+			pp[i] += q[i]
+		}
+		// Correct through the linearisation at the predicted mean, but
+		// innovate against the true observation map.
+		xpArr, err := core.FromFloats(xp, d)
+		if err != nil {
+			return kalmanStepOut{}, err
+		}
+		hOut, err := observation(xpArr)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: the observation: %w", name, t+1, err)
+		}
+		predicted, err := kfVector(name, "the observation output", hOut, m)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		hJacArr, err := hjac(xpArr)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: the observation Jacobian: %w", name, t+1, err)
+		}
+		hMat, err := kfMatrix(name, "the observation Jacobian", hJacArr, m, d)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		return kfCorrect(name, t+1, zRow, predicted, xp, pp, d, m, hMat, kfTranspose(hMat, m, d), r)
+	})
+}
+
+// UnscentedKalmanFilter runs the unscented Kalman filter over the
+// measurement stack z under the nonlinear model x_{t+1} = f(x_t) +
+// w_t, z_t = h(x_t) + v_t: the state distribution N(x̂, P) is carried
+// through f and h exactly to second order by the deterministic
+// sigma-point set x̂ ± sqrt(d+λ)·L[:, i], L the Cholesky factor of P,
+// whose weighted moments reconstruct the predicted mean and
+// covariance. The update re-draws the set from the predicted
+// distribution, so the state-observation cross-covariance carries the
+// process noise the innovation covariance does. λ = α²(d+κ) − d, with
+// d the state dimension the sigma-point set spans, comes from
+// KalmanOptions' SigmaAlpha, SigmaBeta and SigmaKappa; the
+// weights are the standard scaled set, with the covariance weight of
+// the centre point carrying the 1 − α² + β prior correction.
+//
+// The correction carries no Joseph form: the gain comes from the
+// explicit cross-covariance of state and observation rather than an
+// observation matrix, so the covariance leaves the update as
+// P⁻ − K·S·Kᵀ mirrored into its symmetric average, and positive
+// definiteness is enforced by the next predict's Cholesky
+// factorisation, which names the step when it fails. The log-likelihood
+// accumulates exactly as in the linear filter. On a linear model the
+// sigma transforms are exact and the filter degenerates to the
+// Kalman answer to rounding; the state dimension must be fixed by an
+// initial state or covariance (see KalmanOptions).
+func UnscentedKalmanFilter(z *core.Array, transition, observation StateFunc, opts KalmanOptions) (*KalmanResult, error) {
+	const name = "UnscentedKalmanFilter"
+	if transition == nil || observation == nil {
+		return nil, base.Errf("%s: the transition and observation maps are required", name)
+	}
+	nMeas, m, meas, err := kfMeasurements(name, z)
+	if err != nil {
+		return nil, err
+	}
+	d, x0, p0, q, r, err := kfNonlinearInputs(name, opts, m)
+	if err != nil {
+		return nil, err
+	}
+	alpha := opts.SigmaAlpha
+	if alpha == 0 {
+		alpha = 0.001
+	}
+	beta := opts.SigmaBeta
+	if beta == 0 {
+		beta = 2
+	}
+	kappa := opts.SigmaKappa
+	if math.IsNaN(alpha) || alpha < 0 {
+		return nil, base.Errf("%s: the sigma alpha must be unset (the default 0.001) or positive, got %g",
+			name, opts.SigmaAlpha)
+	}
+	if math.IsNaN(beta) || beta < 0 {
+		return nil, base.Errf("%s: the sigma beta must be unset (the default 2) or positive, got %g",
+			name, opts.SigmaBeta)
+	}
+	if math.IsNaN(kappa) {
+		return nil, base.Errf("%s: the sigma kappa must not be NaN, got %g", name, kappa)
+	}
+	scale := alpha * alpha * (float64(d) + kappa)
+	if scale <= 0 {
+		return nil, base.Errf("%s: the sigma spread vanishes: alpha %g and kappa %g leave no positive scale for %d states",
+			name, alpha, kappa, d)
+	}
+	points := 2*d + 1
+	lambda := scale - float64(d)
+	wm := make([]float64, points)
+	wc := make([]float64, points)
+	wm[0] = lambda / scale
+	wc[0] = wm[0] + 1 - alpha*alpha + beta
+	for i := 1; i < points; i++ {
+		wm[i] = 1 / (2 * scale)
+		wc[i] = wm[i]
+	}
+	spread := math.Sqrt(scale)
+	sig := make([]float64, points*d)
+	prop := make([]float64, points*d)
+	obs := make([]float64, points*m)
+	return kalmanRun(nMeas, m, d, meas, x0, p0, func(t int, zRow, x, p []float64) (kalmanStepOut, error) {
+		// The sigma set: the mean plus each Cholesky column of the
+		// covariance, pushed out by sqrt(d+λ) both ways.
+		l, err := kfCholesky("the state covariance", p, d)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		copy(sig[:d], x)
+		for i := range d {
+			plus := sig[(1+2*i)*d : (2+2*i)*d]
+			minus := sig[(2+2*i)*d : (3+2*i)*d]
+			copy(plus, x)
+			copy(minus, x)
+			for j := range d {
+				v := spread * l[j*d+i]
+				plus[j] += v
+				minus[j] -= v
+			}
+		}
+		// Predict: every sigma through f, then the weighted moments.
+		for s := range points {
+			sa, err := core.FromFloats(sig[s*d:(s+1)*d], d)
+			if err != nil {
+				return kalmanStepOut{}, err
+			}
+			fa, err := transition(sa)
+			if err != nil {
+				return kalmanStepOut{}, base.Errf("%s: at step %d: the transition: %w", name, t+1, err)
+			}
+			fv, err := kfVector(name, "the transition output", fa, d)
+			if err != nil {
+				return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+			}
+			copy(prop[s*d:(s+1)*d], fv)
+		}
+		xp := make([]float64, d)
+		for s := range points {
+			for j := range d {
+				xp[j] += wm[s] * prop[s*d+j]
+			}
+		}
+		pp := make([]float64, d*d)
+		for s := range points {
+			w := wc[s]
+			prows := prop[s*d : (s+1)*d]
+			for a := range d {
+				da := prows[a] - xp[a]
+				for b := range d {
+					pp[a*d+b] += w * da * (prows[b] - xp[b])
+				}
+			}
+		}
+		for i := range pp {
+			pp[i] += q[i]
+		}
+		kfSymmetrise(pp, d)
+		// The update re-draws the sigma set from the predicted
+		// distribution N(x⁻, P⁻), the process noise included: the
+		// state-observation cross-covariance must carry the same
+		// spread the innovation covariance does, or the gain loses the
+		// Q·Hᵀ term and a linear model stops degenerating to the
+		// Kalman answer.
+		lp2, err := kfCholesky("the predicted state covariance", pp, d)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		copy(sig[:d], xp)
+		for i := range d {
+			plus := sig[(1+2*i)*d : (2+2*i)*d]
+			minus := sig[(2+2*i)*d : (3+2*i)*d]
+			copy(plus, xp)
+			copy(minus, xp)
+			for j := range d {
+				v := spread * lp2[j*d+i]
+				plus[j] += v
+				minus[j] -= v
+			}
+		}
+		// Observation: the predicted sigmas through h, the weighted
+		// moments, and the state-observation cross-covariance.
+		for s := range points {
+			pa, err := core.FromFloats(sig[s*d:(s+1)*d], d)
+			if err != nil {
+				return kalmanStepOut{}, err
+			}
+			ha, err := observation(pa)
+			if err != nil {
+				return kalmanStepOut{}, base.Errf("%s: at step %d: the observation: %w", name, t+1, err)
+			}
+			hv, err := kfVector(name, "the observation output", ha, m)
+			if err != nil {
+				return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+			}
+			copy(obs[s*m:(s+1)*m], hv)
+		}
+		zp := make([]float64, m)
+		for s := range points {
+			for j := range m {
+				zp[j] += wm[s] * obs[s*m+j]
+			}
+		}
+		s := make([]float64, m*m)
+		for k := range points {
+			w := wc[k]
+			orow := obs[k*m : (k+1)*m]
+			for a := range m {
+				da := orow[a] - zp[a]
+				for b := range m {
+					s[a*m+b] += w * da * (orow[b] - zp[b])
+				}
+			}
+		}
+		for i := range s {
+			s[i] += r[i]
+		}
+		kfSymmetrise(s, m)
+		cross := make([]float64, d*m)
+		for k := range points {
+			w := wc[k]
+			prow := sig[k*d : (k+1)*d]
+			orow := obs[k*m : (k+1)*m]
+			for a := range d {
+				da := prow[a] - xp[a]
+				for b := range m {
+					cross[a*m+b] += w * da * (orow[b] - zp[b])
+				}
+			}
+		}
+		ls, err := kfCholesky("the innovation covariance", s, m)
+		if err != nil {
+			return kalmanStepOut{}, base.Errf("%s: at step %d: %w", name, t+1, err)
+		}
+		innovation := make([]float64, m)
+		for i := range m {
+			innovation[i] = zRow[i] - zp[i]
+		}
+		solved := kfCholSolve(ls, m, innovation)
+		quad := 0.0
+		for i := range m {
+			quad += innovation[i] * solved[i]
+		}
+		// K = P_xz·S⁻¹ through the transposed solve: S·X = P_xzᵀ
+		// gives X = S⁻¹·P_xzᵀ and K = Xᵀ.
+		gain := kfTranspose(kfCholSolveMatrix(ls, m, kfTranspose(cross, d, m), d), m, d)
+		mean := make([]float64, d)
+		copy(mean, xp)
+		for i := range d {
+			total := 0.0
+			grow := gain[i*m : (i+1)*m]
+			for j, y := range innovation {
+				total += grow[j] * y
+			}
+			mean[i] += total
+		}
+		kskt := kfMatMul(kfMatMul(gain, d, m, s, m), d, m, kfTranspose(gain, d, m), d)
+		covariance := make([]float64, d*d)
+		for i := range covariance {
+			covariance[i] = pp[i] - kskt[i]
+		}
+		kfSymmetrise(covariance, d)
+		logLik := -0.5 * (float64(m)*math.Log(2*math.Pi) + kfLogDet(ls, m) + quad)
+		return kalmanStepOut{mean: mean, covariance: covariance, innovation: innovation,
+			innovationCov: s, logLikelihood: logLik}, nil
+	})
+}