"""Model-based (optimal) filtering: the Wiener and Kalman filters. Where a matched filter detects a *known* waveform in noise, the filters here estimate a *random* signal whose statistics we model. Two classics: * **Wiener filter** - for a wide-sense-stationary signal in stationary noise, the linear minimum-mean-square-error (MMSE) estimator. Its FIR form is the solution of the Wiener-Hopf normal equations ``R_x w = r_dx``; its non-causal frequency-domain form is ``H = S_dd / (S_dd + S_vv)``. * **Kalman filter** - the recursive MMSE estimator for a linear state-space model driven by Gaussian noise. It carries an estimate and its error covariance forward through a predict/update cycle. Everything here is plain NumPy and importable; the topic page imports these functions for its demos, and ``test_model_based.py`` checks them against brute-force least squares, the orthogonality principle, and Monte-Carlo covariance estimates. References: N. Wiener, *Extrapolation, Interpolation, and Smoothing of Stationary Time Series*, MIT Press, 1949. R. E. Kalman, "A new approach to linear filtering and prediction problems," *J. Basic Eng.* 82(1), 1960. """ from __future__ import annotations import warnings import numpy as np from scipy.linalg import solve_toeplitz # --- Wiener filter ---------------------------------------------------------- def wiener_hopf(rxx, rdx): """Solve the Wiener-Hopf normal equations ``R_x w = r_dx`` for the FIR taps. ``rxx`` is the autocorrelation sequence of the observed signal ``x`` for lags ``0 .. M-1`` (so ``rxx[0]`` is the signal power). The ``M x M`` matrix ``R_x`` it generates is symmetric Toeplitz. ``rdx`` is the cross-correlation ``r_dx[k] = E[d[n] x[n-k]]`` between the desired signal ``d`` and ``x`` for the same lags. Returns the length-``M`` optimal filter ``w`` that minimises ``E[(d[n] - sum_k w[k] x[n-k])^2]``. """ rxx = np.asarray(rxx, dtype=float) rdx = np.asarray(rdx, dtype=float) # R_x is symmetric Toeplitz with first row/column = rxx. return solve_toeplitz((rxx, rxx), rdx) def wiener_min_mse(rxx, rdx, sigma_d_sq): """Minimum MSE achieved by the FIR Wiener filter. ``J_min = sigma_d^2 - r_dx^T w = sigma_d^2 - r_dx^T R_x^{-1} r_dx`` where ``sigma_d_sq`` is the power ``E[d[n]^2]`` of the desired signal. This is the floor the orthogonality principle guarantees: among linear filters, no FIR filter of this length can do better. (A nonlinear estimator of the same inputs can, when the signal is non-Gaussian.) """ rdx = np.asarray(rdx, dtype=float) w = wiener_hopf(rxx, rdx) return float(sigma_d_sq - rdx @ w) def wiener_fir_design(x, d, num_taps): """Design an ``num_taps``-tap Wiener filter from data records of ``x``, ``d``. Estimates the autocorrelation of ``x`` and the cross-correlation of ``d`` with ``x`` from the supplied sequences (biased estimator, dividing by the record length), then solves the Wiener-Hopf equations. Returns ``(w, j_min)``: the taps and the predicted minimum MSE. This is the practical, data-driven path: you do not need the true spectra, only enough data to estimate the correlations. """ x = np.asarray(x, dtype=float) d = np.asarray(d, dtype=float) n = len(x) if len(d) != n: raise ValueError("x and d must have the same length") # rxx[k] = E[x[n] x[n-k]] (even, so the lag sign does not matter); # rdx[k] = E[d[n] x[n-k]] -> sum_n d[n] x[n-k] = dot(d[k:], x[:n-k]). # Getting rdx's lag sign right matters whenever d and x are asymmetrically # related (d a delayed or predicted function of x), the usual reason to use # a Wiener filter; for x = d + noise the cross-correlation is even and the # sign is invisible. rxx = np.array([np.dot(x[: n - k], x[k:]) / n for k in range(num_taps)]) rdx = np.array([np.dot(d[k:], x[: n - k]) / n for k in range(num_taps)]) w = wiener_hopf(rxx, rdx) sigma_d_sq = float(np.dot(d, d) / n) return w, float(sigma_d_sq - rdx @ w) def wiener_apply(w, x): """Apply an FIR Wiener filter (causal convolution, same length as ``x``).""" w = np.asarray(w, dtype=float) x = np.asarray(x, dtype=float) return np.convolve(x, w)[: len(x)] def wiener_freq(s_dd, s_vv): """Non-causal frequency-domain Wiener filter ``H = S_dd / (S_dd + S_vv)``. ``s_dd`` and ``s_vv`` are the power spectral densities of the desired signal and the (uncorrelated) noise, sampled on the same frequency grid. The optimal gain is between 0 and 1 at every frequency: it passes a band where the signal dominates and attenuates one where the noise does. Returns the real gain on that grid. """ s_dd = np.asarray(s_dd, dtype=float) s_vv = np.asarray(s_vv, dtype=float) return s_dd / (s_dd + s_vv) # --- Kalman filter ---------------------------------------------------------- def _as2d(M): M = np.atleast_2d(np.asarray(M, dtype=float)) return M def kalman_filter(z, F, H, Q, R, x0, P0, B=None, u=None): """Run the linear Kalman filter over a sequence of measurements. State-space model:: x_k = F x_{k-1} + B u_k + w_k, w_k ~ N(0, Q) z_k = H x_k + v_k, v_k ~ N(0, R) Parameters ---------- z : array (T, m) or (T,) Measurements, one row per time step (``m`` is the measurement dimension; a 1-D array is treated as ``m = 1``). F, H, Q, R : arrays State-transition ``(n, n)``, measurement ``(m, n)``, process-noise covariance ``(n, n)``, measurement-noise covariance ``(m, m)``. x0, P0 : arrays Initial state estimate ``(n,)`` and its error covariance ``(n, n)``. B, u : optional Control matrix ``(n, l)`` and control sequence ``(T, l)``. Omitted means no control input. Returns ------- xs : array (T, n) Filtered (a posteriori) state estimates ``x_hat_{k|k}``. Ps : array (T, n, n) A posteriori error covariances ``P_{k|k}``. """ F, H, Q, R = _as2d(F), _as2d(H), _as2d(Q), _as2d(R) x = np.asarray(x0, dtype=float).copy() P = _as2d(P0).copy() n = F.shape[0] I = np.eye(n) z = np.asarray(z, dtype=float) if z.ndim == 1: z = z[:, None] T = z.shape[0] if B is not None: B = _as2d(B) u = np.asarray(u, dtype=float) if u.ndim == 1: u = u[:, None] xs = np.empty((T, n)) Ps = np.empty((T, n, n)) for k in range(T): # Predict x = F @ x if B is not None: x = x + B @ u[k] P = F @ P @ F.T + Q # Update y = z[k] - H @ x # innovation S = H @ P @ H.T + R # innovation covariance K = P @ H.T @ np.linalg.inv(S) # optimal Kalman gain x = x + K @ y P = (I - K @ H) @ P xs[k] = x Ps[k] = P return xs, Ps def cv_model(dt, sigma_a, sigma_z): """Constant-velocity ("truck on rails") state-space model. State ``x = [position, velocity]``. Between steps an unknown acceleration ``a_k ~ N(0, sigma_a^2)`` perturbs the motion; only position is measured, with noise standard deviation ``sigma_z``. Returns ``(F, H, Q, R)`` with:: F = [[1, dt], [0, 1]] H = [[1, 0]] Q = sigma_a^2 * [[dt^4/4, dt^3/2], [dt^3/2, dt^2]] R = [[sigma_z^2]] ``Q = G G^T sigma_a^2`` with ``G = [dt^2/2, dt]^T`` is the standard discrete white-noise-acceleration process-noise covariance. """ F = np.array([[1.0, dt], [0.0, 1.0]]) H = np.array([[1.0, 0.0]]) G = np.array([dt ** 2 / 2.0, dt]) Q = np.outer(G, G) * sigma_a ** 2 R = np.array([[sigma_z ** 2]]) return F, H, Q, R def simulate_cv(F, H, Q, R, x0, steps, rng): """Simulate a constant-velocity truck: true states and noisy measurements. Draws process noise ``w_k ~ N(0, Q)`` and measurement noise ``v_k ~ N(0, R)`` and rolls the state-space model forward ``steps`` times from ``x0``. Returns ``(xs_true, zs)`` with shapes ``(steps, n)`` and ``(steps, m)``. ``rng`` is a ``numpy.random.Generator``. """ F, H, Q, R = _as2d(F), _as2d(H), _as2d(Q), _as2d(R) n = F.shape[0] m = H.shape[0] x = np.asarray(x0, dtype=float).copy() xs = np.empty((steps, n)) zs = np.empty((steps, m)) for k in range(steps): x = F @ x + rng.multivariate_normal(np.zeros(n), Q) z = H @ x + rng.multivariate_normal(np.zeros(m), R) xs[k] = x zs[k] = z return xs, zs def steady_state_gain(F, H, Q, R, max_iters=1000, tol=1e-12): """Steady-state Kalman gain by iterating the Riccati recursion to convergence. For a time-invariant model the error covariance ``P`` and gain ``K`` settle to constant values. This iterates the predict/update covariance recursion (no data needed) until ``P`` stops changing (its largest entry moves by less than ``tol``), or until ``max_iters`` is reached, then returns ``(K_inf, P_inf)``. For a model that converges very slowly (eigenvalues of ``F`` near the unit circle, or an extreme ``Q``/``R`` ratio) the iteration cap can bind before convergence, so check ``tol`` was actually met if you push the model that far. A microcontroller can precompute this gain offline and skip the per-sample covariance algebra entirely (the alpha-beta filter). """ F, H, Q, R = _as2d(F), _as2d(H), _as2d(Q), _as2d(R) n = F.shape[0] I = np.eye(n) P = np.eye(n) K = np.zeros((n, H.shape[0])) converged = False for _ in range(max_iters): P_prev = P P = F @ P @ F.T + Q S = H @ P @ H.T + R K = P @ H.T @ np.linalg.inv(S) P = (I - K @ H) @ P if np.max(np.abs(P - P_prev)) < tol: converged = True break if not converged: warnings.warn( f"steady_state_gain did not converge to tol={tol} within " f"max_iters={max_iters}; the returned gain may be unconverged " "(slow model: F eigenvalues near 1, or extreme Q/R).", RuntimeWarning, stacklevel=2, ) return K, P