Kalman Filter & State-Space Models

"""Kalman filter - predict and update, transcribed.

Two steps forever. Predict pushes the state forward and lets the uncertainty
grow; update pulls it toward the measurement by an amount set by which one is
trusted more. Matrix operations are written out by hand so the arithmetic is
visible rather than delegated.
"""


def matmul(A, B):
    rows, inner, cols = len(A), len(B), len(B[0])
    out = [[0.0] * cols for _ in range(rows)]
    for i in range(rows):
        for j in range(cols):
            total = 0.0
            for k in range(inner):
                total += A[i][k] * B[k][j]
            out[i][j] = total
    return out


def transpose(A):
    return [[A[i][j] for i in range(len(A))] for j in range(len(A[0]))]


def add(A, B):
    return [[A[i][j] + B[i][j] for j in range(len(A[0]))] for i in range(len(A))]


def invert(A):
    """Gauss-Jordan on the innovation covariance.

    In a real filter this matrix is the size of the MEASUREMENT, not the state,
    so it is usually 1x1 or 2x2 - which is why inverting it directly is
    acceptable here and inverting the state covariance would not be.
    """
    n = len(A)
    augmented = [row[:] + [1.0 if i == j else 0.0 for j in range(n)]
                 for i, row in enumerate(A)]

    for column in range(n):
        pivot = max(range(column, n), key=lambda r: abs(augmented[r][column]))
        if abs(augmented[pivot][column]) < 1e-12:
            raise ValueError("innovation covariance is singular")
        augmented[column], augmented[pivot] = augmented[pivot], augmented[column]

        divisor = augmented[column][column]
        for k in range(2 * n):
            augmented[column][k] /= divisor

        for row in range(n):
            if row == column:
                continue
            factor = augmented[row][column]
            for k in range(2 * n):
                augmented[row][k] -= factor * augmented[column][k]

    return [row[n:] for row in augmented]


def predict(x, P, F, Q):
    """x <- F x, P <- F P F^T + Q. Uncertainty GROWS: time passing makes the
    estimate worse, which is what Q encodes."""
    x_next = matmul(F, x)
    P_next = add(matmul(matmul(F, P), transpose(F)), Q)
    return x_next, P_next


def update(x, P, y, H, R):
    """The gain is prediction uncertainty over TOTAL uncertainty - the trust
    ratio, written as arithmetic."""
    n = len(x)

    # innovation: what was measured minus what was expected
    predicted = matmul(H, x)
    v = [[y[i][0] - predicted[i][0]] for i in range(len(y))]

    # S = H P H^T + R: how surprised the filter was entitled to be
    S = add(matmul(matmul(H, P), transpose(H)), R)
    K = matmul(matmul(P, transpose(H)), invert(S))

    x_next = add(x, matmul(K, v))

    # P <- (I - K H) P. Simple, and numerically the fragile form: rounding
    # can drive P non-symmetric and then indefinite, at which point the
    # filter reports impossible confidence.
    identity = [[1.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
    factor = [[identity[i][j] - matmul(K, H)[i][j] for j in range(n)] for i in range(n)]
    P_next = matmul(factor, P)

    return x_next, P_next


def run(observations, x0, P0, F, H, Q, R):
    x, P = x0, P0
    states = []
    for y in observations:
        x, P = predict(x, P, F, Q)
        x, P = update(x, P, y, H, R)
        states.append(x)
    return states

"""Kalman filter - typed, Joseph form, log-likelihood accumulated."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray
from scipy.linalg import cho_factor, cho_solve

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


@dataclass(frozen=True)
class LinearGaussianModel:
    """The model: dynamics, observation, and the two noise covariances.

    Q and R are the only real tuning here, and their RATIO is what decides
    whether the filter trusts its own prediction or the measurement. Almost
    every field failure of a Kalman filter traces back to this rather than to
    the algorithm.
    """

    F: Matrix       # (n_x, n_x) state transition
    H: Matrix       # (n_y, n_x) observation
    Q: Matrix       # (n_x, n_x) process noise
    R: Matrix       # (n_y, n_y) measurement noise

    def __post_init__(self) -> None:
        if self.F.shape[0] != self.F.shape[1]:
            raise ValueError(f"F must be square, got {self.F.shape}")
        if self.H.shape[1] != self.F.shape[0]:
            raise ValueError(f"H has {self.H.shape[1]} columns but the state is {self.F.shape[0]}")
        if not np.allclose(self.Q, self.Q.T) or not np.allclose(self.R, self.R.T):
            raise ValueError("noise covariances must be symmetric")


@dataclass
class FilterState:
    mean: Vector
    covariance: Matrix
    log_likelihood: float = 0.0


def predict(state: FilterState, model: LinearGaussianModel) -> FilterState:
    """Uncertainty GROWS: time passing makes the estimate worse."""
    mean = model.F @ state.mean
    covariance = model.F @ state.covariance @ model.F.T + model.Q
    return FilterState(mean, _symmetrize(covariance), state.log_likelihood)


def update(state: FilterState, observation: Vector, model: LinearGaussianModel) -> FilterState:
    """Joseph form update, and the log-likelihood contribution for free."""
    innovation = observation - model.H @ state.mean
    S = model.H @ state.covariance @ model.H.T + model.R

    # Cholesky rather than an explicit inverse: S is symmetric positive
    # definite by construction, and its log-determinant - needed for the
    # likelihood - falls out of the factor at no extra cost.
    factor = cho_factor(S, lower=True)
    gain = cho_solve(factor, model.H @ state.covariance).T

    mean = state.mean + gain @ innovation

    # JOSEPH FORM: (I - KH) P (I - KH)^T + K R K^T.
    #
    # Algebraically identical to (I - KH) P, and numerically far better: it is
    # a sum of two symmetric positive-semidefinite terms, so it stays positive
    # definite under round-off. The simple form does not, and a covariance
    # that loses positive definiteness makes the filter report impossible
    # confidence and diverge - with nothing raised.
    identity = np.eye(state.mean.size)
    closed = identity - gain @ model.H
    covariance = closed @ state.covariance @ closed.T + gain @ model.R @ gain.T

    # Prediction-error decomposition: the filter produces the one-step
    # predictive density as it runs, so the likelihood costs nothing extra.
    lower = factor[0]
    log_det = 2.0 * np.log(np.diag(lower)).sum()
    quadratic = float(innovation @ cho_solve(factor, innovation))
    contribution = -0.5 * (log_det + quadratic + innovation.size * np.log(2.0 * np.pi))

    return FilterState(mean, _symmetrize(covariance), state.log_likelihood + contribution)


def _symmetrize(covariance: Matrix) -> Matrix:
    """Round-off makes P asymmetric, and asymmetric becomes indefinite."""
    return 0.5 * (covariance + covariance.T)


def run(
    observations: Matrix,
    initial: FilterState,
    model: LinearGaussianModel,
) -> tuple[list[FilterState], float]:
    """Filter a whole series. A missing observation is a skipped update - the
    state simply propagates and the uncertainty grows, which is one of the
    framework's genuine advantages over a lag-based model."""
    if observations.ndim != 2 or observations.shape[1] != model.H.shape[0]:
        raise ValueError(f"expected (T, {model.H.shape[0]}) observations, got {observations.shape}")

    state = initial
    history: list[FilterState] = []

    for observation in observations:
        state = predict(state, model)
        if np.isfinite(observation).all():
            state = update(state, observation, model)
        history.append(state)

    return history, state.log_likelihood

"""Kalman filter - steady-state gain, and vectorized across many tracks."""

import numpy as np
from numpy.typing import NDArray
from scipy.linalg import solve_discrete_are

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


def steady_state_gain(F: Matrix, H: Matrix, Q: Matrix, R: Matrix) -> Matrix:
    """For a TIME-INVARIANT system the covariance recursion converges.

    P is a fixed point of the discrete algebraic Riccati equation, so after a
    transient the gain stops changing. Solving for that fixed point once
    removes the entire covariance propagation from the loop - the O(n_x^3)
    work per step disappears and each step becomes two matrix-vector products.

    The cost is that the transient is gone too: the filter no longer converges
    from an uncertain start, it simply begins at steady state. That is right
    for a long-running tracker and wrong for a filter that must recover from
    a bad initialization.
    """
    P = solve_discrete_are(F.T, H.T, Q, R)
    S = H @ P @ H.T + R
    return P @ H.T @ np.linalg.inv(S)


class BatchTracker:
    """Many independent filters, one interpreter iteration per timestep.

    A tracker runs hundreds of tracks with identical dynamics. Each is
    sequential in time, but they are completely INDEPENDENT of one another -
    so the loop over tracks disappears and only the loop over time remains,
    exactly as for a panel of time series.
    """

    def __init__(self, F: Matrix, H: Matrix, gain: Matrix, n_tracks: int) -> None:
        # One contiguous float64 block per operand: every step multiplies
        # against these, and a non-contiguous array forces a copy each time.
        self._F = np.ascontiguousarray(F, dtype=np.float64)
        self._H = np.ascontiguousarray(H, dtype=np.float64)
        self._K = np.ascontiguousarray(gain, dtype=np.float64)

        n_x = F.shape[0]
        # State for every track, allocated once and updated in place.
        self._state = np.zeros((n_tracks, n_x), dtype=np.float64)
        self._predicted = np.empty_like(self._state)
        self._innovation = np.empty((n_tracks, H.shape[0]), dtype=np.float64)

    def step(self, measurements: Matrix, valid: NDArray) -> Matrix:
        """One timestep for every track at once.

        measurements is (n_tracks, n_y); valid marks which tracks actually
        received one. A track with no measurement simply keeps its prediction,
        which is how occlusion is handled.
        """
        # Predict: state @ F.T is the batched form of F @ x for every track.
        np.matmul(self._state, self._F.T, out=self._predicted)

        # Innovation for every track, built in place.
        np.matmul(self._predicted, self._H.T, out=self._innovation)
        np.subtract(measurements, self._innovation, out=self._innovation)

        # Update only the tracks that were measured; the rest carry the
        # prediction forward untouched.
        np.copyto(self._state, self._predicted)
        if valid.any():
            self._state[valid] += self._innovation[valid] @ self._K.T

        return self._state

    def gate(self, measurements: Matrix, threshold: float, inverse_S: Matrix) -> NDArray:
        """Normalized innovation squared, for every track at once.

        Chi-squared distributed under a correct model, so the threshold comes
        from a distribution rather than a tuned multiple - which is what makes
        this gate principled and a distance cutoff not.
        """
        np.matmul(self._state, self._H.T, out=self._innovation)
        np.subtract(measurements, self._innovation, out=self._innovation)
        # v^T S^-1 v for every track, fused so the intermediate never exists.
        nis = np.einsum("ij,jk,ik->i", self._innovation, inverse_S, self._innovation)
        return nis < threshold