Estimation¶

def ekf(
    model_or_f: DynamicsLike,
    Q_noise: Array,
    R_noise: Array,
    ys: Array,
    us: Array,
    x0: Array,
    P0: Array,
    *,
    observation: ObservationFn | None = None,
) -> KalmanResult:
    """Run an extended Kalman filter for nonlinear dynamics and observations.

    Uses automatic Jacobians of the transition and observation functions rather
    than requiring a hand-derived local linearization.

    This is the main nonlinear recursive estimation primitive in the current
    library. It works well with JAX-defined model functions, but the API is
    intentionally simple and does not yet expose multiple EKF variants.

    Args:
        model_or_f: Nonlinear system model or dynamics function.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        ys: Measurement sequence. Shape: `(T, p)`.
        us: Input sequence. Shape: `(T, m)`.
        x0: Initial mean. Shape: `(n,)`.
        P0: Initial covariance. Shape: `(n, n)`.
        observation: Observation function `h(x)` for plain dynamics callables.
            Omit this when passing a system model with its own output map.

    Returns:
        [KalmanResult][contrax.types.KalmanResult]: A bundle containing
            filtered means, covariances, and innovations.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> def f(x, u):
        ...     return jnp.array([x[0] + 0.1 * x[1], x[1] + u[0]])
        >>> def h(x):
        ...     return x[:1]
        >>> result = cx.ekf(
        ...     f,
        ...     Q_noise=1e-3 * jnp.eye(2),
        ...     R_noise=1e-2 * jnp.eye(1),
        ...     ys=jnp.zeros((10, 1)),
        ...     us=jnp.zeros((10, 1)),
        ...     x0=jnp.zeros(2),
        ...     P0=jnp.eye(2),
        ...     observation=h,
        ... )
    """
    if _is_system_model(model_or_f):
        if observation is not None:
            raise ValueError("ekf(sys, ...) does not accept observation=.")
        obs = model_or_f
    else:
        if observation is None:
            raise ValueError(
                "ekf(f, ..., observation=...) requires an observation function."
            )
        obs = observation

    dt = _system_dt(model_or_f)
    ts = jnp.arange(ys.shape[0], dtype=x0.dtype)
    if dt is not None:
        ts = ts * dt

    # Update-then-predict: x0/P0 is the prior on x_0.
    def update_predict(carry, inputs):
        x_prior, P_prior = carry
        y, u, t = inputs
        x_post, P_post, innov = ekf_update(
            obs,
            x_prior,
            P_prior,
            y,
            R_noise,
            u=u,
            t=t,
        )
        x_next, P_next = ekf_predict(model_or_f, x_post, P_post, u, Q_noise, t=t)
        return (x_next, P_next), (x_post, P_post, innov)

    _, (x_hats, Ps, innovations) = jax.lax.scan(update_predict, (x0, P0), (ys, us, ts))
    return KalmanResult(x_hat=x_hats, P=Ps, innovations=innovations)

def ekf_predict(
    model_or_f: DynamicsLike,
    x: Array,
    P: Array,
    u: Array,
    Q_noise: Array,
    *,
    t: Array | float = 0.0,
) -> tuple[Array, Array]:
    """Predict one extended Kalman filter step.

    Args:
        model_or_f: Nonlinear system model or dynamics function mapping
            `(x, u)` to the next state.
        x: Filtered mean at the previous time step. Shape: `(n,)`.
        P: Filtered covariance at the previous time step. Shape: `(n, n)`.
        u: Current input. Shape: `(m,)`.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        t: Current sample time.

    Returns:
        Tuple `(x_pred, P_pred)` with the predicted mean and covariance.
    """
    f = _coerce_dynamics(model_or_f)
    x_pred = f(x, u, t)
    F = jax.jacfwd(f, argnums=0)(x, u, t)
    P_pred = _symmetrize(F @ P @ F.T + Q_noise)
    return x_pred, P_pred

def ekf_step(
    model_or_f: DynamicsLike,
    x: Array,
    P: Array,
    u: Array,
    y: Array,
    Q_noise: Array,
    R_noise: Array,
    *,
    t: Array | float = 0.0,
    has_measurement: Array | bool = True,
    num_iter: int = 1,
    observation: ObservationFn | None = None,
) -> tuple[Array, Array, Array]:
    """Run one predict-update step of an extended Kalman filter.

    Args:
        model_or_f: Nonlinear system model or dynamics function.
        x: Filtered mean at the previous time step. Shape: `(n,)`.
        P: Filtered covariance at the previous time step. Shape: `(n, n)`.
        u: Current input. Shape: `(m,)`.
        y: Measurement vector. Shape: `(p,)`.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        has_measurement: Whether to apply the measurement update.
        num_iter: Number of iterated-EKF measurement linearization passes.
        observation: Observation function `h(x)` for plain dynamics callables.
            Omit this when passing a system model with its own output map.

    Returns:
        Tuple `(x, P, innovation)` after the step.
    """
    if _is_system_model(model_or_f):
        if observation is not None:
            raise ValueError("ekf_step(sys, ...) does not accept observation=.")
        obs = model_or_f
    else:
        if observation is None:
            raise ValueError(
                "ekf_step(f, ..., observation=...) requires an observation function."
            )
        obs = observation
    x_pred, P_pred = ekf_predict(model_or_f, x, P, u, Q_noise, t=t)
    return ekf_update(
        obs,
        x_pred,
        P_pred,
        y,
        R_noise,
        u=u,
        t=t,
        has_measurement=has_measurement,
        num_iter=num_iter,
    )

def ekf_update(
    model_or_h: ObservationLike,
    x_pred: Array,
    P_pred: Array,
    y: Array,
    R_noise: Array,
    *,
    u: Array | None = None,
    t: Array | float = 0.0,
    has_measurement: Array | bool = True,
    num_iter: int = 1,
) -> tuple[Array, Array, Array]:
    """Update one extended Kalman filter step.

    Args:
        model_or_h: Nonlinear system model or observation function.
        x_pred: Predicted mean. Shape: `(n,)`.
        P_pred: Predicted covariance. Shape: `(n, n)`.
        y: Measurement vector. Shape: `(p,)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        u: Current input used by the observation model. Shape: `(m,)`.
        t: Current sample time.
        has_measurement: Whether to apply the measurement update.
        num_iter: Number of iterated-EKF measurement linearization passes.
            `1` gives the ordinary EKF update.

    Returns:
        Tuple `(x, P, innovation)` with the filtered mean, filtered covariance,
        and final measurement residual. Missing-measurement steps return the
        prediction and a zero innovation.
    """
    if num_iter < 1:
        raise ValueError("num_iter must be at least 1")

    n = x_pred.shape[0]
    if u is None:
        u = jnp.zeros((0,), dtype=x_pred.dtype)
    h = _coerce_observation(
        model_or_h,
        None if _is_system_model(model_or_h) else model_or_h,
    )

    # Reference: Jazwinski (1970), "Stochastic Processes and Filtering
    # Theory"; the iterated EKF relinearizes the measurement model around the
    # current update iterate while keeping the prediction covariance fixed.
    def one_iteration(x_lin, _):
        H = jax.jacfwd(h, argnums=0)(x_lin, u, t)
        S = H @ P_pred @ H.T + R_noise
        K = jnp.linalg.solve(S.T, (P_pred @ H.T).T).T
        innov = y - h(x_lin, u, t) + H @ (x_lin - x_pred)
        x_next = x_pred + K @ innov
        return x_next, (x_next, H, K, innov)

    x_final, (_, Hs, Ks, innovations) = jax.lax.scan(
        one_iteration,
        x_pred,
        xs=None,
        length=num_iter,
    )
    H_final = Hs[-1]
    K_final = Ks[-1]
    innov_final = innovations[-1]
    P_final = _symmetrize((jnp.eye(n, dtype=P_pred.dtype) - K_final @ H_final) @ P_pred)

    def apply_update():
        return x_final, P_final, innov_final

    def skip_update():
        return x_pred, P_pred, jnp.zeros_like(y)

    return jax.lax.cond(
        _measurement_flag(has_measurement),
        apply_update,
        skip_update,
    )

def innovation_diagnostics(
    innovations: Array,
    innovation_covariances: Array,
) -> InnovationDiagnostics:
    """Summarize innovation magnitude and covariance conditioning.

    The returned NIS values are the per-step normalized innovation squared
    terms `v_k^T S_k^{-1} v_k`. They are useful for routine filter health
    inspection, but Contrax does not present them as a formal standalone
    consistency certificate.
    """

    def step_nis(innovation, covariance):
        solved = jnp.linalg.solve(covariance, innovation[:, None])
        return (innovation[None, :] @ solved).squeeze()

    nis = jax.vmap(step_nis)(innovations, innovation_covariances)
    innovation_norms = jnp.linalg.norm(innovations, axis=-1)
    innovation_cov_condition_numbers = jax.vmap(jnp.linalg.cond)(innovation_covariances)
    all_values = (
        jnp.isfinite(nis)
        & jnp.isfinite(innovation_norms)
        & jnp.isfinite(innovation_cov_condition_numbers)
    )
    return InnovationDiagnostics(
        nis=nis,
        innovation_norms=innovation_norms,
        innovation_cov_condition_numbers=innovation_cov_condition_numbers,
        mean_nis=jnp.mean(nis),
        max_nis=jnp.max(nis),
        mean_innovation_norm=jnp.mean(innovation_norms),
        max_innovation_norm=jnp.max(innovation_norms),
        max_innovation_cov_condition_number=jnp.max(innovation_cov_condition_numbers),
        nonfinite=jnp.logical_not(jnp.all(all_values)),
    )

def innovation_rms(result: KalmanResult | UKFResult) -> Array:
    """Return the RMS innovation magnitude for a filtered result."""
    return jnp.sqrt(jnp.mean(result.innovations**2))

def kalman(
    sys: DiscLTI,
    Q_noise: Array,  # (n, n) process noise covariance
    R_noise: Array,  # (p, p) measurement noise covariance
    ys: Array,  # (T, p) measurement sequence
    x0: Array | None = None,
    P0: Array | None = None,
) -> KalmanResult:
    """Run a linear discrete-time Kalman filter.

    Filters a measurement sequence for a
    [DiscLTI][contrax.systems.DiscLTI] system using the standard
    predict-update recursion expressed as a pure `lax.scan`.

    This is the baseline linear-Gaussian estimation path in Contrax. Innovation
    gains are computed with linear solves rather than explicit inverses, and
    the scan structure makes small differentiable covariance-tuning loops
    practical.

    Args:
        sys: Discrete-time linear system.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        ys: Measurement sequence. Shape: `(T, p)`.
        x0: Optional initial mean. Shape: `(n,)`. Defaults to zeros.
        P0: Optional initial covariance. Shape: `(n, n)`. Defaults to the
            identity matrix.

    Returns:
        [KalmanResult][contrax.types.KalmanResult]: A bundle containing
            filtered means `x_hat`, filtered covariances `P`, and measurement
            innovations `innovations`.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys = cx.dss(
        ...     jnp.array([[1.0, 0.1], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.1]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.1,
        ... )
        >>> result = cx.kalman(
        ...     sys,
        ...     Q_noise=1e-3 * jnp.eye(2),
        ...     R_noise=1e-2 * jnp.eye(2),
        ...     ys=jnp.zeros((20, 2)),
        ... )
    """
    n = sys.A.shape[0]
    if x0 is None:
        x0 = jnp.zeros(n, dtype=ys.dtype)
    if P0 is None:
        P0 = jnp.eye(n, dtype=ys.dtype)

    # Update-then-predict: x0/P0 is the prior on x_0.
    # Each scan step updates with y_k first (using the current prior), then
    # predicts forward to give the prior on x_{k+1}.
    def update_predict(carry, y):
        x_prior, P_prior = carry
        x_post, P_post, innov = kalman_update(sys, x_prior, P_prior, y, R_noise)
        x_next, P_next = kalman_predict(sys, x_post, P_post, Q_noise)
        return (x_next, P_next), (x_post, P_post, innov)

    _, (x_hats, Ps, innovations) = jax.lax.scan(update_predict, (x0, P0), ys)
    return KalmanResult(x_hat=x_hats, P=Ps, innovations=innovations)

def kalman_gain(
    sys: DiscLTI,
    Q_noise: Array,
    R_noise: Array,
) -> KalmanGainResult:
    """Design a steady-state Kalman filter gain for a discrete LTI system.

    This is the estimator dual of discrete LQR. It solves the DARE for
    `(A.T, C.T, Q_noise, R_noise)` and returns the measurement-update gain used
    by `kalman_update()`.

    Args:
        sys: Discrete-time linear system.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.

    Returns:
        [KalmanGainResult][contrax.types.KalmanGainResult]: A bundle with
            update gain `K`, steady-state predicted covariance `P`, and
            estimator error-dynamics poles.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys = cx.dss(
        ...     jnp.array([[1.0, 0.1], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.1]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.1,
        ... )
        >>> design = cx.kalman_gain(sys, 1e-3 * jnp.eye(2), 1e-2 * jnp.eye(2))
    """
    P_pred = dare(sys.A.T, sys.C.T, Q_noise, R_noise).S
    innovation_cov = sys.C @ P_pred @ sys.C.T + R_noise
    K = jnp.linalg.solve(innovation_cov.T, (P_pred @ sys.C.T).T).T
    poles = jnp.linalg.eigvals(sys.A - sys.A @ K @ sys.C)
    return KalmanGainResult(K=K, P=P_pred, poles=poles)

def kalman_predict(
    sys: DiscLTI,
    x: Array,
    P: Array,
    Q_noise: Array,
    u: Array | None = None,
) -> tuple[Array, Array]:
    """Predict one linear Kalman filter step.

    Args:
        sys: Discrete-time linear system.
        x: Filtered mean at the previous time step. Shape: `(n,)`.
        P: Filtered covariance at the previous time step. Shape: `(n, n)`.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        u: Optional input for the transition. Shape: `(m,)`. If omitted, the
            input term is zero.

    Returns:
        Tuple `(x_pred, P_pred)` with the predicted mean and covariance.
    """
    u = _zero_input(sys, x) if u is None else u
    x_pred = sys.A @ x + sys.B @ u
    P_pred = _symmetrize(sys.A @ P @ sys.A.T + Q_noise)
    return x_pred, P_pred

def kalman_step(
    sys: DiscLTI,
    x: Array,
    P: Array,
    y: Array,
    Q_noise: Array,
    R_noise: Array,
    u: Array | None = None,
    *,
    has_measurement: Array | bool = True,
) -> tuple[Array, Array, Array]:
    """Run one predict-update step of a linear Kalman filter.

    This is the online counterpart to `kalman()`. It is a pure array function,
    so service loops can call it directly and batch workflows can place it
    inside `jax.lax.scan`.

    Args:
        sys: Discrete-time linear system.
        x: Filtered mean at the previous time step. Shape: `(n,)`.
        P: Filtered covariance at the previous time step. Shape: `(n, n)`.
        y: Measurement vector. Shape: `(p,)`.
        Q_noise: Process noise covariance. Shape: `(n, n)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        u: Optional input. Shape: `(m,)`.
        has_measurement: Whether to apply the measurement update.

    Returns:
        Tuple `(x, P, innovation)` after the step.
    """
    x_pred, P_pred = kalman_predict(sys, x, P, Q_noise, u)
    return kalman_update(
        sys,
        x_pred,
        P_pred,
        y,
        R_noise,
        u,
        has_measurement=has_measurement,
    )

def kalman_update(
    sys: DiscLTI,
    x_pred: Array,
    P_pred: Array,
    y: Array,
    R_noise: Array,
    u: Array | None = None,
    *,
    has_measurement: Array | bool = True,
) -> tuple[Array, Array, Array]:
    """Update one linear Kalman filter step.

    `has_measurement=False` returns the predicted state and covariance without
    applying a measurement update. The flag may be a JAX scalar boolean, making
    this helper usable inside `jax.lax.scan` for streams with missing samples.

    Args:
        sys: Discrete-time linear system.
        x_pred: Predicted mean. Shape: `(n,)`.
        P_pred: Predicted covariance. Shape: `(n, n)`.
        y: Measurement vector. Shape: `(p,)`.
        R_noise: Measurement noise covariance. Shape: `(p, p)`.
        u: Optional input for direct-feedthrough measurement models. Shape:
            `(m,)`. If omitted, the feedthrough term is zero.
        has_measurement: Whether to apply the measurement update.

    Returns:
        Tuple `(x, P, innovation)` with the filtered mean, filtered covariance,
        and measurement residual. Missing-measurement steps return a zero
        innovation with the same shape as `y`.
    """
    u = _zero_input(sys, x_pred) if u is None else u
    n = x_pred.shape[0]
    y_pred = sys.C @ x_pred + sys.D @ u
    S = sys.C @ P_pred @ sys.C.T + R_noise
    K = jnp.linalg.solve(S.T, (P_pred @ sys.C.T).T).T
    innov = y - y_pred
    x_new = x_pred + K @ innov
    P_new = _symmetrize((jnp.eye(n, dtype=P_pred.dtype) - K @ sys.C) @ P_pred)

    def apply_update():
        return x_new, P_new, innov

    def skip_update():
        return x_pred, P_pred, jnp.zeros_like(y)

    return jax.lax.cond(
        _measurement_flag(has_measurement),
        apply_update,
        skip_update,
    )

def likelihood_diagnostics(log_likelihood_terms: Array) -> LikelihoodDiagnostics:
    """Summarize per-step innovation-form log-likelihood terms."""
    return LikelihoodDiagnostics(
        log_likelihood_terms=log_likelihood_terms,
        total_log_likelihood=jnp.sum(log_likelihood_terms),
        mean_log_likelihood=jnp.mean(log_likelihood_terms),
        min_log_likelihood=jnp.min(log_likelihood_terms),
        max_log_likelihood=jnp.max(log_likelihood_terms),
        nonfinite=jnp.logical_not(jnp.all(jnp.isfinite(log_likelihood_terms))),
    )

def mhe(
    f: Callable[..., Array],
    h: Callable[..., Array],
    xs_init: Array,
    us: Array,
    ys: Array,
    x_prior: Array,
    P_prior: Array,
    Q_noise: Array,
    R_noise: Array,
    params: object | None = None,
    extra_cost: Callable[..., Array] | None = None,
    solver: optx.AbstractMinimiser | None = None,
    max_steps: int = 256,
) -> MHEResult:
    """Solve a fixed-window moving-horizon-estimation problem.

    Minimises `mhe_objective` over the candidate state trajectory `xs_init`
    using an `optimistix` solver.  The default solver is
    `optx.LBFGS(rtol=1e-6, atol=1e-6)`, which is well-suited to the
    deterministic, smooth, fixed-horizon objectives that arise in MHE.

    The optimised terminal state `x_hat = xs[-1]` is the natural input to a
    downstream controller or observer.  For a linear-Gaussian model, the
    full trajectory `xs` matches the RTS smoother output within solver
    tolerance.

    Args:
        f: Transition function.  Called as `f(x, u)` when `params=None`,
            and `f(x, u, params)` otherwise.
        h: Observation function.  Called as `h(x)` when `params=None`,
            and `h(x, params)` otherwise.
        xs_init: Initial trajectory guess.  Shape: `(T, n)`.  Using the
            Kalman filtered means as a warm start is recommended.
        us: Input sequence.  Shape: `(T - 1, m)`.
        ys: Measurement sequence aligned with `xs`.  Shape: `(T, p)`.
        x_prior: Prior mean for the start of the window.  Shape: `(n,)`.
        P_prior: Arrival covariance.  Shape: `(n, n)`.
        Q_noise: Process noise covariance.  Shape: `(n, n)`.
        R_noise: Measurement noise covariance.  Shape: `(p, p)`.
        params: Optional parameters passed to `f`, `h`, and `extra_cost`.
        extra_cost: Optional callable adding additional smooth costs.
            Called as `extra_cost(xs, us, ys)` when `params=None`, and
            `extra_cost(xs, us, ys, params)` otherwise.
        solver: Optimistix minimiser.  Defaults to
            `optx.LBFGS(rtol=1e-6, atol=1e-6)`.
        max_steps: Maximum solver iterations.  Defaults to 256.

    Returns:
        [MHEResult][contrax.types.MHEResult]: Estimated trajectory, terminal
            state, final cost, and convergence flag.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys = cx.dss(
        ...     jnp.array([[0.8]]), jnp.zeros((1, 1)),
        ...     jnp.array([[1.0]]), jnp.zeros((1, 1)), dt=1.0,
        ... )
        >>> f = lambda x, u: sys.A @ x
        >>> h = lambda x: sys.C @ x
        >>> ys = jnp.array([[0.1], [0.4], [0.6], [0.5], [0.4]])
        >>> result = cx.mhe(
        ...     f, h,
        ...     xs_init=jnp.zeros((5, 1)),
        ...     us=jnp.zeros((4, 1)),
        ...     ys=ys,
        ...     x_prior=jnp.zeros(1),
        ...     P_prior=jnp.eye(1),
        ...     Q_noise=0.05 * jnp.eye(1),
        ...     R_noise=0.2 * jnp.eye(1),
        ... )
    """
    if solver is None:
        solver = optx.LBFGS(rtol=1e-6, atol=1e-6)

    def fn(xs, args):
        us_, ys_, x_prior_, P_prior_, Q_noise_, R_noise_ = args
        return mhe_objective(
            f,
            h,
            xs,
            us_,
            ys_,
            x_prior_,
            P_prior_,
            Q_noise_,
            R_noise_,
            params,
            extra_cost,
        )

    args = (us, ys, x_prior, P_prior, Q_noise, R_noise)
    sol = optx.minimise(
        fn, solver, xs_init, args=args, max_steps=max_steps, throw=False
    )
    xs_opt = sol.value
    return MHEResult(
        xs=xs_opt,
        x_hat=xs_opt[-1],
        final_cost=fn(xs_opt, args),
        solver_converged=sol.result == optx.RESULTS.successful,
    )