contrax 0.1.0__py3-none-any.whl

contrax/__init__.py

@@ -0,0 +1,197 @@
+from contrax.analysis import (
+    ctrb,
+    ctrb_gramian,
+    dcgain,
+    dlyap,
+    evalfr,
+    freqresp,
+    lyap,
+    obsv,
+    obsv_gramian,
+    poles,
+    zeros,
+)
+from contrax.control import (
+    augment_integrator,
+    care,
+    dare,
+    feedback,
+    lqi,
+    lqr,
+    place,
+    state_feedback,
+)
+from contrax.estimation import (
+    ekf,
+    ekf_predict,
+    ekf_step,
+    ekf_update,
+    innovation_diagnostics,
+    innovation_rms,
+    kalman,
+    kalman_gain,
+    kalman_predict,
+    kalman_step,
+    kalman_update,
+    likelihood_diagnostics,
+    mhe,
+    mhe_objective,
+    mhe_warm_start,
+    rts,
+    smoother_diagnostics,
+    soft_quadratic_penalty,
+    ukf,
+    ukf_diagnostics,
+    uks,
+)
+from contrax.parameterization import (
+    diagonal_spd,
+    lower_triangular,
+    positive_exp,
+    positive_softplus,
+    spd_from_cholesky_raw,
+)
+from contrax.simulation import (
+    as_ode_term,
+    foh_inputs,
+    impulse_response,
+    initial_response,
+    lsim,
+    rollout,
+    sample_system,
+    simulate,
+    step_response,
+)
+from contrax.systems import (
+    ContLTI,
+    DiscLTI,
+    NonlinearSystem,
+    PHSSystem,
+    block_matrix,
+    block_observation,
+    c2d,
+    canonical_J,
+    dss,
+    linearize,
+    linearize_ss,
+    nonlinear_system,
+    parallel,
+    partition_state,
+    phs_diagnostics,
+    phs_system,
+    phs_to_ss,
+    project_psd,
+    schedule_phs,
+    series,
+    ss,
+    symmetrize_matrix,
+)
+from contrax.types import (
+    InnovationDiagnostics,
+    KalmanGainResult,
+    KalmanResult,
+    LikelihoodDiagnostics,
+    LQRResult,
+    MHEResult,
+    PHSStructureDiagnostics,
+    RTSResult,
+    SmootherDiagnostics,
+    UKFResult,
+)
+
+__all__ = [
+    # system types
+    "ContLTI",
+    "DiscLTI",
+    "NonlinearSystem",
+    "PHSSystem",
+    "partition_state",
+    "block_observation",
+    "block_matrix",
+    "symmetrize_matrix",
+    "project_psd",
+    "phs_diagnostics",
+    # constructors
+    "ss",
+    "dss",
+    "c2d",
+    "nonlinear_system",
+    "phs_system",
+    "phs_to_ss",
+    "schedule_phs",
+    "canonical_J",
+    "linearize",
+    "linearize_ss",
+    "series",
+    "parallel",
+    # control
+    "lqr",
+    "lqi",
+    "dare",
+    "care",
+    "place",
+    "state_feedback",
+    "feedback",
+    "augment_integrator",
+    # simulation
+    "rollout",
+    "foh_inputs",
+    "sample_system",
+    "lsim",
+    "simulate",
+    "step_response",
+    "impulse_response",
+    "initial_response",
+    "as_ode_term",
+    # estimation
+    "kalman",
+    "kalman_gain",
+    "kalman_predict",
+    "kalman_update",
+    "kalman_step",
+    "innovation_diagnostics",
+    "innovation_rms",
+    "likelihood_diagnostics",
+    "ekf",
+    "ekf_predict",
+    "ekf_update",
+    "ekf_step",
+    "mhe_objective",
+    "mhe_warm_start",
+    "mhe",
+    "soft_quadratic_penalty",
+    "smoother_diagnostics",
+    "ukf",
+    "ukf_diagnostics",
+    "uks",
+    "rts",
+    # analysis
+    "ctrb",
+    "obsv",
+    "poles",
+    "zeros",
+    "evalfr",
+    "freqresp",
+    "dcgain",
+    "ctrb_gramian",
+    "obsv_gramian",
+    "lyap",
+    "dlyap",
+    # parameterization
+    "positive_exp",
+    "positive_softplus",
+    "lower_triangular",
+    "spd_from_cholesky_raw",
+    "diagonal_spd",
+    # result types
+    "LQRResult",
+    "InnovationDiagnostics",
+    "SmootherDiagnostics",
+    "KalmanGainResult",
+    "KalmanResult",
+    "LikelihoodDiagnostics",
+    "UKFResult",
+    "RTSResult",
+    "MHEResult",
+    "PHSStructureDiagnostics",
+]

contrax/_ekf.py

@@ -0,0 +1,269 @@
+"""contrax._ekf — Extended Kalman filter."""
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from contrax._kalman import _measurement_flag, _symmetrize
+from contrax.nonlinear import (
+    DynamicsLike,
+    ObservationFn,
+    ObservationLike,
+    _coerce_dynamics,
+    _coerce_observation,
+    _is_system_model,
+    _system_dt,
+)
+from contrax.types import KalmanResult
+
+
+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_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 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(
+    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)

contrax/_estimation_diagnostics.py

@@ -0,0 +1,155 @@
+"""contrax._estimation_diagnostics — estimation diagnostics and health checks."""
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from contrax.types import (
+    InnovationDiagnostics,
+    KalmanResult,
+    LikelihoodDiagnostics,
+    RTSResult,
+    SmootherDiagnostics,
+    UKFResult,
+)
+
+
+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 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 ukf_diagnostics(
+    result: UKFResult,
+) -> tuple[InnovationDiagnostics, LikelihoodDiagnostics]:
+    """Build the standard diagnostics pair for a UKF result."""
+    return (
+        innovation_diagnostics(result.innovations, result.innovation_covariances),
+        likelihood_diagnostics(result.log_likelihood_terms),
+    )
+
+
+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 smoother_diagnostics(
+    smoothed: RTSResult,
+    filtered: KalmanResult | UKFResult,
+) -> SmootherDiagnostics:
+    """Compute health diagnostics for an RTS or UKS smoother result.
+
+    A numerically healthy smoother satisfies ``P_smooth ≤ P_filtered`` in the
+    PSD sense at every step. This function checks that guarantee and surfaces
+    the worst-case violation so callers can detect breakdown before trusting
+    smoothed trajectories.
+
+    Args:
+        smoothed: Output of `rts()` or `uks()`.
+        filtered: The forward filter result used to produce ``smoothed``.
+            Accepts both `KalmanResult` and `UKFResult`.
+
+    Returns:
+        [SmootherDiagnostics][contrax.types.SmootherDiagnostics]: Diagnostic
+            bundle. ``min_covariance_reduction < 0`` is the primary health flag.
+
+    Examples:
+        >>> import jax.numpy as jnp
+        >>> import contrax as cx
+        >>> sys = cx.dss(
+        ...     jnp.array([[0.9]]), jnp.zeros((1, 1)),
+        ...     jnp.array([[1.0]]), jnp.zeros((1, 1)), dt=1.0,
+        ... )
+        >>> filtered = cx.kalman(sys, 1e-3 * jnp.eye(1), 1e-2 * jnp.eye(1),
+        ...                      jnp.zeros((10, 1)))
+        >>> smoothed = cx.rts(sys, filtered, Q_noise=1e-3 * jnp.eye(1))
+        >>> diag = cx.smoother_diagnostics(smoothed, filtered)
+    """
+    x_smooth = smoothed.x_smooth  # (T, n)
+    P_smooth = smoothed.P_smooth  # (T, n, n)
+    x_filt = filtered.x_hat  # (T, n)
+    P_filt = filtered.P  # (T, n, n)
+
+    def _sym(M: Array) -> Array:
+        return (M + M.T) / 2
+
+    # Min eigenvalue of (P_f - P_s): >= 0 means smoother reduced uncertainty.
+    cov_reduction = jax.vmap(
+        lambda Pf, Ps: jnp.min(jnp.linalg.eigvalsh(_sym(Pf) - _sym(Ps)))
+    )(P_filt, P_smooth)
+
+    state_corrections = jnp.linalg.norm(x_smooth - x_filt, axis=-1)
+
+    smoothed_min_eig = jax.vmap(lambda Ps: jnp.min(jnp.linalg.eigvalsh(_sym(Ps))))(
+        P_smooth
+    )
+
+    nonfinite = jnp.logical_not(
+        jnp.all(jnp.isfinite(x_smooth)) & jnp.all(jnp.isfinite(P_smooth))
+    )
+
+    return SmootherDiagnostics(
+        covariance_reduction=cov_reduction,
+        min_covariance_reduction=jnp.min(cov_reduction),
+        state_corrections=state_corrections,
+        max_state_correction=jnp.max(state_corrections),
+        smoothed_min_eigenvalue=smoothed_min_eig,
+        nonfinite=nonfinite,
+    )
+
+
+__all__ = [
+    "InnovationDiagnostics",
+    "LikelihoodDiagnostics",
+    "SmootherDiagnostics",
+    "innovation_diagnostics",
+    "innovation_rms",
+    "likelihood_diagnostics",
+    "smoother_diagnostics",
+    "ukf_diagnostics",
+]