Estimation¶

def ekf(
    model_or_f: DynamicsLike,
    ys: Array,
    *,
    Q: Array,
    R: Array,
    us: Array | None = None,
    x0: Array | None = None,
    P0: Array | None = None,
    observation: ObservationFn | None = None,
    has_measurement: Array | None = None,
    solver: EKF | None = None,
    throw: bool = False,
) -> KalmanResult:
    """Run an extended Kalman filter for discrete nonlinear dynamics and observations.

    Computes transition and observation Jacobians with automatic
    differentiation.

    Joseph covariance updates and one measurement linearization are the defaults.
    Solver configuration selects compact covariance form or iterated updates.
    Gradients flow through the filter recursion and its JAX-defined model
    functions.

    Args:
        model_or_f: Nonlinear system model or dynamics function.
        Q: Process noise covariance. Shape: `(n, n)`.
        R: Measurement noise covariance. Shape: `(p, p)`.
        ys: Measurement sequence. Shape: `(T, p)`.
        us: Optional controls, shape `(T, m)` or `(T, 2, m)` for sampled FOH.
            Defaults to zero controls; declare model.input_dim or supply us for
            controlled nonlinear models.
        x0: Initial prior mean, shape `(n,)`. Defaults to zeros.
        P0: Initial prior covariance, shape `(n, n)`. Defaults to identity.
        observation: Observation function `h(x)` for plain dynamics callables.
            Omit this when passing a system model with its own output map.
        has_measurement: Whole-observation availability flags. Shape: `(T,)`.
            Missing observations may contain NaNs and return zero innovations.
        solver: Joseph covariance form and iterated measurement-update settings.

        throw: Raise on numerical failure; otherwise inspect status.

    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(t, 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, jnp.zeros((10, 1)), Q=1e-3 * jnp.eye(2), R=1e-2 * jnp.eye(1),
        ...     us=jnp.zeros((10, 1)), 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

    ys, Q, R, us, x0, P0, flags = prepare_filter(
        model_or_f, ys, Q, R, us, x0, P0, has_measurement
    )
    dt = _system_dt(model_or_f)
    ts = jnp.arange(ys.shape[0], dtype=x0.dtype)
    if dt is not None:
        ts = ts * jnp.asarray(dt, dtype=x0.dtype)

    # Update-then-predict: x0/P0 is the prior on x_0.
    def update_predict(carry, inputs):
        x_prior, P_prior, code = carry
        y, u, t, available, k = inputs
        x_post, P_post, innov = ekf_update(
            obs,
            x_prior,
            P_prior,
            y,
            R=R,
            u=u,
            t=t,
            has_measurement=available,
            solver=solver,
        )

        def advance():
            prediction = _predict(model_or_f, x_post, P_post, Q, u, t)
            return (
                prediction.x,
                prediction.P,
                prediction.solver_status.astype(code.dtype),
            )

        x_next, P_next, next_code = jax.lax.cond(
            k + 1 < ys.shape[0],
            advance,
            lambda: (x_post, P_post, jnp.zeros_like(code)),
        )
        code = jnp.where(code != 0, code, next_code)
        return (x_next, P_next, code), (x_post, P_post, innov)

    (_, _, code), (x_hats, Ps, innovations) = jax.lax.scan(
        update_predict,
        (x0, P0, input_status(P0, Q, R)),
        (ys, us, ts, flags, jnp.arange(ys.shape[0])),
    )
    result = KalmanResult(
        xs=x_hats,
        Ps=Ps,
        innovations=innovations,
        has_measurement=flags,
        solver_status=code,
    )
    return checked(result, throw, "EKF filter failed; inspect status")

def ekf_predict(
    model_or_f: DynamicsLike,
    x: Array,
    P: Array,
    *,
    Q: Array,
    u: Array | None = None,
    t: Array | float = 0.0,
    throw: bool = False,
) -> FilterPredictionResult:
    """Predict one extended Kalman filter step.

    Args:
        model_or_f: A `NonlinearSystem` or a plain dynamics function with
            signature `(t, x, u) -> x_next`. Both signatures are now the same,
            so `sys.dynamics` can be passed directly when using a system model.
        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: Process noise covariance. Shape: `(n, n)`.
        t: Current sample time.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        FilterPredictionResult: Mean x, covariance P, status and successful.
    """
    x, P, Q = prepare_state(x, P, Q)
    u = zero_input(model_or_f, x.dtype) if u is None else u
    prediction = _predict(model_or_f, x, P, Q, u, t)
    valid = input_status(P, Q)
    code = jnp.where(valid != 0, valid, prediction.solver_status)
    result = FilterPredictionResult(prediction.x, prediction.P, code)
    return checked(result, throw, "EKF prediction failed; inspect status")

def ekf_step(
    model_or_f: DynamicsLike,
    x: Array,
    P: Array,
    y: Array,
    *,
    Q: Array,
    R: Array,
    u: Array | None = None,
    t: Array | float = 0.0,
    has_measurement: Array | bool = True,
    solver: EKF | None = None,
    observation: ObservationFn | None = None,
    u_measurement: Array | None = None,
    measurement_time: Array | float | None = None,
    throw: bool = False,
) -> FilterUpdateResult:
    """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,)`.
        t: Transition start time.
        y: Measurement vector. Shape: `(p,)`.
        Q: Process noise covariance. Shape: `(n, n)`.
        R: Measurement noise covariance. Shape: `(p, p)`.
        has_measurement: Whether to apply the measurement update.
        solver: Covariance update and number of iterated measurement passes.
        observation: Observation function `h(x)` for plain dynamics callables.
            Omit this when passing a system model with its own output map.
        u_measurement: Input at the measurement time, defaulting to u.
        measurement_time: Observation time. Defaults to t plus the model's
            sample time when available, otherwise t.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        FilterUpdateResult: Mean x, covariance P, innovation, status and successful.
    """
    u = zero_input(model_or_f, x.dtype) if u is None else u
    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
    prediction = ekf_predict(model_or_f, x, P, u=u, Q=Q, t=t)
    x_pred, P_pred = prediction
    if measurement_time is None:
        dt = _system_dt(model_or_f)
        measurement_time = t if dt is None else t + dt
    update = ekf_update(
        obs,
        x_pred,
        P_pred,
        y,
        R=R,
        u=u if u_measurement is None else u_measurement,
        t=measurement_time,
        has_measurement=has_measurement,
        solver=solver,
    )
    code = jnp.where(
        prediction.solver_status != 0, prediction.solver_status, update.solver_status
    )
    result = FilterUpdateResult(update.x, update.P, update.innovation, code)
    return checked(result, throw, "EKF step failed; inspect status")

def ekf_update(
    model_or_h: ObservationLike,
    x_pred: Array,
    P_pred: Array,
    y: Array,
    *,
    R: Array,
    u: Array | None = None,
    t: Array | float = 0.0,
    has_measurement: Array | bool = True,
    solver: EKF | None = None,
    throw: bool = False,
) -> FilterUpdateResult:
    """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: 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.
        solver: Joseph covariance form and number of iterated linearizations.
            `EKF(num_iter=1)` gives the ordinary EKF update.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        FilterUpdateResult: Mean x, covariance P, innovation and status; unpacks
        as (x, P, innovation), with the filtered mean, filtered covariance,
        and final measurement residual. Missing-measurement steps return the
        prediction and a zero innovation.
    """
    x_pred, P_pred, R = prepare_state(x_pred, P_pred, R, measurement=y)
    solver = EKF() if solver is None else solver
    if not isinstance(solver, EKF):
        raise TypeError("ekf_update solver must be contrax.solvers.EKF")

    if u is None:
        u = zero_input(model_or_h, 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=1)(t, x_lin, u)
        if H.shape != (y.shape[0], x_pred.shape[0]):
            raise ValueError("The observation map must return a vector matching y")
        S = H @ P_pred @ H.T + R
        K = jnp.linalg.solve(S.T, (P_pred @ H.T).T).T
        observed = jnp.where(_measurement_flag(has_measurement), y, jnp.zeros_like(y))
        innov = observed - h(t, x_lin, u) + H @ (x_lin - x_pred)
        x_next = x_pred + K @ innov
        return x_next, (x_next, H, K, innov)

    def apply_update():
        x_final, (_, Hs, Ks, innovations) = jax.lax.scan(
            one_iteration,
            x_pred,
            xs=None,
            length=solver.num_iter,
        )
        P_final = _covariance_update(
            P_pred,
            Hs[-1],
            Ks[-1],
            R,
            solver.covariance_update,
        )
        return x_final, P_final, innovations[-1]

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

    x, P, innovation = jax.lax.cond(
        _measurement_flag(has_measurement),
        apply_update,
        skip_update,
    )
    result = FilterUpdateResult(x, P, innovation, input_status(P_pred, R))
    return checked(result, throw, "EKF update failed; inspect status")

def innovation_diagnostics(
    innovations: Array,
    innovation_covariances: Array,
    *,
    has_measurement: Array | None = None,
) -> 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; interpret them alongside the model and noise assumptions.
    Optional whole-observation flags exclude missing entries from summaries and
    covariance solves. Missing per-step values are zero placeholders. With no
    observations, summaries are NaN and `nonfinite` remains false.
    """
    flags = _observations(innovations, has_measurement)
    innovations = jnp.where(flags[:, None], innovations, 0)
    innovation_covariances = jnp.where(
        flags[:, None, None],
        innovation_covariances,
        jnp.eye(innovations.shape[-1], dtype=innovations.dtype),
    )

    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 = jnp.where(
        flags, jax.vmap(jnp.linalg.cond)(innovation_covariances), 0
    )
    mean_nis, _, max_nis = _masked_summary(nis, flags)
    mean_norm, _, max_norm = _masked_summary(innovation_norms, flags)
    _, _, max_condition = _masked_summary(innovation_cov_condition_numbers, flags)
    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=mean_nis,
        max_nis=max_nis,
        mean_innovation_norm=mean_norm,
        max_innovation_norm=max_norm,
        max_innovation_cov_condition_number=max_condition,
        nonfinite=jnp.logical_not(jnp.all(all_values)),
    )

def innovation_rms(result: KalmanResult | UKFResult) -> Array:
    """Return observed innovation RMS; NaN if there are no observed components."""
    observed = jnp.where(result.has_measurement[:, None], result.innovations, 0)
    count = jnp.sum(result.has_measurement) * result.innovations.shape[-1]
    return jnp.where(
        count > 0, jnp.sqrt(jnp.sum(observed**2) / jnp.maximum(count, 1)), jnp.nan
    )

def kalman(
    sys: DiscLTI,
    ys: Array,  # (T, p) measurement sequence
    *,
    Q: Array,
    R: Array,
    x0: Array | None = None,
    P0: Array | None = None,
    us: Array | None = None,
    has_measurement: Array | None = None,
    solver: Kalman | SteadyStateKalman | InformationKalman | None = None,
    throw: bool = False,
) -> 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`.

    Innovation gains use Cholesky solves and posterior covariances use Joseph
    form by default. Innovation covariances must be positive definite. Gradients
    flow through the filter recursion, including covariance tuning. Inputs are
    promoted to a common floating dtype, including cached design data.

    Args:
        sys: Discrete-time linear system.
        Q: Process noise covariance. Shape: `(n, n)`.
        R: Measurement covariance `(p, p)`, or a positive variance vector `(p,)`
            with InformationKalman, which requires positive-definite priors.
        ys: Measurement sequence. Shape: `(T, p)`.
        x0: Optional initial mean. Shape: `(n,)`. Defaults to zeros.
        P0: Initial predicted covariance `(n, n)`. Defaults to identity for
            covariance/information filtering, or design.P with SteadyStateKalman.
            A provided cached-mode P0 must match the design's predicted covariance.
        us: Controls for feedthrough at k and transition k to k+1. Shape:
            `(T, m)`. Defaults to zeros; the final transition is not returned.
        has_measurement: Whole-observation availability flags. Shape: `(T,)`.
            Missing observations may contain NaNs and return zero innovations.
        solver: Kalman (Joseph default, compact optional), InformationKalman for
            diagonal R and state-size solves, or SteadyStateKalman(design).
            Cached mode requires a fully observed stream and a gain/covariance
            consistent with A, C, Q and R; mismatch reports invalid_problem.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        [KalmanResult][contrax.types.KalmanResult]: A bundle containing
            filtered means `xs`, filtered covariances `Ps`, 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, jnp.zeros((20, 2)), Q=1e-3 * jnp.eye(2), R=1e-2 * jnp.eye(2)
        ... )
    """
    if not isinstance(sys, DiscLTI):
        raise TypeError("kalman requires a discrete LTI system")
    solver = Kalman() if solver is None else solver
    if not isinstance(solver, (Kalman, SteadyStateKalman, InformationKalman)):
        raise TypeError("solver must be Kalman, SteadyStateKalman or InformationKalman")
    information = isinstance(solver, InformationKalman)
    extra = ()
    if isinstance(solver, SteadyStateKalman):
        extra = (solver.design.K, solver.design.P)
        P0 = solver.design.P if P0 is None else P0
    ys, Q, R, us, x0, P0, flags = prepare_filter(
        sys,
        ys,
        Q,
        R,
        us,
        x0,
        P0,
        has_measurement,
        diagonal_R=information,
        extra_arrays=extra,
    )
    if isinstance(solver, SteadyStateKalman):
        result = cached_filter(sys, ys, Q, R, us, x0, P0, flags, solver)
        return checked(result, throw, "Kalman filter failed; inspect status")

    # 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, inputs):
        x_prior, P_prior = carry
        y, u, available, k = inputs
        x_post, P_post, innov = kalman_update(
            sys, x_prior, P_prior, y, R=R, u=u, has_measurement=available, solver=solver
        )
        x_next, P_next = jax.lax.cond(
            k + 1 < ys.shape[0],
            lambda: tuple(kalman_predict(sys, x_post, P_post, Q=Q, u=u)),
            lambda: (x_post, P_post),
        )
        return (x_next, P_next), (x_post, P_post, innov)

    _, (x_hats, Ps, innovations) = jax.lax.scan(
        update_predict, (x0, P0), (ys, us, flags, jnp.arange(ys.shape[0]))
    )
    result = KalmanResult(
        xs=x_hats,
        Ps=Ps,
        innovations=innovations,
        has_measurement=flags,
        solver_status=jnp.where(
            input_status(P0, Q) != 0, 2, noise_status(R, information)
        ),
    )
    return checked(result, throw, "Kalman filter failed; inspect status")

def kalman_gain(
    sys: DiscLTI,
    Q: Array,
    R: Array,
    *,
    solver: DARE | None = None,
    throw: bool = False,
) -> KalmanGainResult:
    """Design a steady-state Kalman filter gain for a discrete LTI system.

    Solves the estimator-dual DARE for `(A.T, C.T, Q, R)` and returns the
    measurement-update gain for `kalman_update()`.

    Args:
        sys: Discrete-time linear system.
        Q: Process noise covariance. Shape: `(n, n)`.
        R: Measurement noise covariance. Shape: `(p, p)`.
        solver: Dual DARE solver configuration.
        throw: Raise on failed Riccati design.

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

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> 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, 0.001 * jnp.eye(2), 0.01 * jnp.eye(2))
    """
    if not isinstance(sys, DiscLTI):
        raise TypeError("kalman_gain requires a discrete LTI system")
    design = dare(sys.A.T, sys.C.T, Q, R, solver=solver, throw=throw)
    P_pred = design.S
    innovation_cov = sys.C @ P_pred @ sys.C.T + R
    K = jnp.linalg.solve(innovation_cov.T, (P_pred @ sys.C.T).T).T
    closed_loop = sys.A - sys.A @ K @ sys.C
    return KalmanGainResult(
        K=K,
        P=P_pred,
        closed_loop=closed_loop,
        status=design.status,
        residual_norm=design.residual_norm,
        num_steps=design.num_steps,
    )

def kalman_predict(
    sys: DiscLTI,
    x: Array,
    P: Array,
    *,
    Q: Array,
    u: Array | None = None,
    throw: bool = False,
) -> FilterPredictionResult:
    """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: Process noise covariance. Shape: `(n, n)`.
        u: Optional input for the transition. Shape: `(m,)`. If omitted, the
            input term is zero.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        FilterPredictionResult: Mean x, covariance P, status and successful.
    """
    x, P, Q = prepare_state(x, P, Q)
    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)
    result = FilterPredictionResult(x_pred, P_pred, input_status(P, Q))
    return checked(result, throw, "Kalman prediction failed; inspect status")

def kalman_step(
    sys: DiscLTI,
    x: Array,
    P: Array,
    y: Array,
    *,
    Q: Array,
    R: Array,
    u: Array | None = None,
    has_measurement: Array | bool = True,
    solver: Kalman | SteadyStateKalman | InformationKalman | None = None,
    u_measurement: Array | None = None,
    throw: bool = False,
) -> FilterUpdateResult:
    """Run one predict-update step of a linear Kalman filter.

    The step can be called directly from a service loop or placed inside
    `jax.lax.scan` for a batch workflow.

    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: Process noise covariance. Shape: `(n, n)`.
        R: Measurement covariance `(p, p)`, or variances `(p,)` in information mode.
        u: Optional input. Shape: `(m,)`.
        has_measurement: Whether to apply the measurement update.
        solver: Covariance, information or explicit cached-design configuration,
            shared with batch filtering.
        u_measurement: Input at the measurement time for direct feedthrough.
            Defaults to u; pass separately when consecutive controls differ.

        throw: Raise on numerical failure; otherwise inspect status.

    Returns:
        FilterUpdateResult: Mean x, covariance P, innovation, status and successful.
    """
    prediction = kalman_predict(sys, x, P, Q=Q, u=u)
    x_pred, P_pred = prediction
    update = kalman_update(
        sys,
        x_pred,
        P_pred,
        y,
        R=R,
        u=u if u_measurement is None else u_measurement,
        has_measurement=has_measurement,
        solver=solver,
    )
    code = jnp.where(
        prediction.solver_status != 0, prediction.solver_status, update.solver_status
    )
    result = FilterUpdateResult(update.x, update.P, update.innovation, code)
    return checked(result, throw, "Kalman step failed; inspect status")