Control¶

def augment_integrator(
    sys: DiscLTI | ContLTI,
    C_integral: Array | None = None,
    D_integral: Array | None = None,
    *,
    sign: float = 1.0,
    dt_scale: ArrayLike | None = None,
) -> DiscLTI | ContLTI:
    """Augment an LTI system with integral-of-output states.

    Returns an augmented system for LQI design. Pass it to `lqr()` with a state
    cost that includes the added integral states.

    For continuous systems, the added state evolves as
    `z_dot = sign * (C_integral x + D_integral u)`. For discrete systems, the
    added state evolves as
    `z[k+1] = z[k] + dt_scale * sign * (C_integral x[k] + D_integral u[k])`.
    When `dt_scale` is omitted for a discrete system, `sys.dt` is used.

    Args:
        sys: Continuous or discrete LTI system to augment.
        C_integral: Output map to integrate. Shape: `(r, n)`. Defaults to
            `sys.C`.
        D_integral: Feedthrough map for the integrated output. Shape: `(r, m)`.
            Defaults to `sys.D`.
        sign: Sign applied to the integrated output. Use `-1.0` for the common
            `reference - output` convention when the reference is handled
            outside the design model.
        dt_scale: Discrete integration scale. Defaults to `sys.dt` for
            discrete systems and is ignored for continuous systems.

    Returns:
        [DiscLTI][contrax.systems.DiscLTI] |
            [ContLTI][contrax.systems.ContLTI]: System with augmented state
            `[x; z]`.
    """
    C_int = sys.C if C_integral is None else jnp.asarray(C_integral, dtype=sys.A.dtype)
    D_int = sys.D if D_integral is None else jnp.asarray(D_integral, dtype=sys.A.dtype)
    scale = jnp.asarray(sign, dtype=sys.A.dtype)

    n = sys.A.shape[0]
    r = C_int.shape[0]
    zeros_xz = jnp.zeros((n, r), dtype=sys.A.dtype)
    eye_z = jnp.eye(r, dtype=sys.A.dtype)
    C_aug = jnp.hstack([sys.C, jnp.zeros((sys.C.shape[0], r), dtype=sys.A.dtype)])

    if isinstance(sys, ContLTI):
        A_aug = jnp.block(
            [
                [sys.A, zeros_xz],
                [scale * C_int, jnp.zeros((r, r), dtype=sys.A.dtype)],
            ]
        )
        B_aug = jnp.vstack([sys.B, scale * D_int])
        return ContLTI(A=A_aug, B=B_aug, C=C_aug, D=sys.D)

    if isinstance(sys, DiscLTI):
        dt = sys.dt if dt_scale is None else jnp.asarray(dt_scale, dtype=sys.A.dtype)
        row_scale = scale * dt
        A_aug = jnp.block([[sys.A, zeros_xz], [row_scale * C_int, eye_z]])
        B_aug = jnp.vstack([sys.B, row_scale * D_int])
        return DiscLTI(A=A_aug, B=B_aug, C=C_aug, D=sys.D, dt=sys.dt)

    raise TypeError(f"Expected DiscLTI or ContLTI, got {type(sys)}")

def care(
    A: Array,
    B: Array,
    Q: Array,
    R: Array,
    *,
    solver: CARE | None = None,
    throw: bool = False,
) -> LQRResult:
    """Solve the continuous algebraic Riccati equation and design a stabilizing gain.

    Solves $A^T S + S A - S B R^{-1} B^T S + Q = 0$ and computes
    $K = R^{-1} B^T S$.

    The baseline solver starts from a Hamiltonian stable-subspace candidate and
    applies residual-decreasing Newton corrections. Poor separation of the
    stable subspace can make the candidate inaccurate; residual and closed-loop
    stability checks report such cases as failures. The stabilizing solution
    requires a stabilizable pair (A, B) and detectable pair (Q**1/2, A).

    Uses real matrices with symmetric positive-semidefinite Q and
    positive-definite R, and requires float64 precision. JVPs and VJPs use
    implicit differentiation of the Riccati equation.

    Args:
        A: State matrix. Shape: `(n, n)`.
        B: Input matrix. Shape: `(n, m)`.
        Q: State cost matrix. Shape: `(n, n)`.
        R: Input cost matrix. Shape: `(m, m)`.
        solver: Hamiltonian candidate refinement and residual tolerances.
        throw: Raise on failed validation instead of returning failure diagnostics.

    Returns:
        LQRResult: Gain K, Riccati solution S, poles, status and diagnostics.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> result = cx.care(-jnp.eye(2), jnp.eye(2), jnp.eye(2), jnp.eye(2))
        >>> bool(result.successful)
        True
    """
    require_x64("care")
    solver = CARE() if solver is None else solver
    if not isinstance(solver, CARE):
        raise TypeError("care solver must be contrax.solvers.CARE")
    A, B, Q, R = _validate_matrices(A, B, Q, R)
    S, steps = _care_with_info(A, B, Q, R, solver.refinement_steps)
    return _result(A, B, Q, R, S, steps, solver, continuous=True, throw=throw)

def dare(
    A: Array,
    B: Array,
    Q: Array,
    R: Array,
    *,
    solver: DARE | None = None,
    throw: bool = False,
) -> LQRResult:
    """Solve the discrete algebraic Riccati equation and design a stabilizing gain.

    Solves $S = A^T S A - A^T S B (R + B^T S B)^{-1} B^T S A + Q$
    and computes $K = (R + B^T S B)^{-1} B^T S A$.

    The solver applies structured doubling with scale-aware stopping and
    validates the equation residual. It is the mature discrete LTI design path;
    check the returned status and residual for numerical failures.
    The stabilizing solution requires a stabilizable pair (A, B) and detectable
    pair (Q**1/2, A). Inputs are real matrices with symmetric
    positive-semidefinite Q and positive-definite R. Requires float64 precision.

    Gradients use implicit differentiation of the Riccati equation through a
    residual-validated Lyapunov solve.

    Args:
        A: State transition matrix. Shape: `(n, n)`.
        B: Input matrix. Shape: `(n, m)`.
        Q: State cost matrix. Shape: `(n, n)`.
        R: Input cost matrix. Shape: `(m, m)`.
        solver: Structured doubling tolerances and iteration limit.
        throw: Raise on failed validation instead of returning failure diagnostics.

    Returns:
        LQRResult: Gain K, Riccati solution S, poles, status and diagnostics.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> result = cx.dare(jnp.eye(2), jnp.eye(2), jnp.eye(2), jnp.eye(2))
        >>> bool(result.successful)
        True
    """
    require_x64("dare")
    solver = DARE() if solver is None else solver
    if not isinstance(solver, DARE):
        raise TypeError("dare solver must be contrax.solvers.DARE")
    A, B, Q, R = _validate_matrices(A, B, Q, R)
    S, steps = _dare_solve_with_info(
        A, B, Q, R, solver.max_steps, solver.rtol, solver.atol
    )
    return _result(A, B, Q, R, S, steps, solver, continuous=False, throw=throw)

def feedback(sys: DiscLTI | ContLTI, K: Array) -> DiscLTI | ContLTI:
    """Apply state feedback to an LTI system; alias for `state_feedback()`."""
    return state_feedback(sys, K)

def lqi(
    sys: DiscLTI | ContLTI,
    Q: Array,
    R: Array,
    C_integral: Array | None = None,
    D_integral: Array | None = None,
    *,
    solver: DARE | CARE | None = None,
    throw: bool = False,
    sign: float = 1.0,
    dt_scale: ArrayLike | None = None,
) -> LQRResult:
    """Design an LQI gain by augmenting integral states and calling `lqr()`.

    Augments the system with integral states and designs a gain for the
    resulting state `[x; z]`, where `z` contains the integrated outputs.

    Args:
        sys: Continuous or discrete LTI system.
        Q: Augmented state cost matrix. Shape: `(n + r, n + r)`.
        R: Input cost matrix. Shape: `(m, m)`.
        C_integral: Output map to integrate. Shape: `(r, n)`. Defaults to
            `sys.C`.
        D_integral: Feedthrough map for the integrated output. Shape: `(r, m)`.
            Defaults to `sys.D`.
        solver: DARE for discrete systems or CARE for continuous systems.
        throw: Raise on numerical failure.
        sign: Sign applied to the integrated output.
        dt_scale: Discrete integration scale. Defaults to `sys.dt` for
            discrete systems and is ignored for continuous systems.

    Returns:
        [LQRResult][contrax.types.LQRResult]: LQR result for the
            integrator-augmented design model.
    """
    augmented = augment_integrator(
        sys,
        C_integral=C_integral,
        D_integral=D_integral,
        sign=sign,
        dt_scale=dt_scale,
    )
    return lqr(augmented, Q, R, solver=solver, throw=throw)

def lqr(
    sys: DiscLTI | ContLTI,
    Q: Array,
    R: Array,
    *,
    solver: DARE | CARE | None = None,
    throw: bool = False,
) -> LQRResult:
    """Solve the infinite-horizon linear quadratic regulator problem.

    Dispatches to `dare()` for discrete systems and `care()` for continuous
    systems, returning a unified [LQRResult][contrax.types.LQRResult].

    Both routes validate the equation residual and closed-loop stability and
    support implicit JVP/VJP differentiation. Type dispatch occurs at trace time;
    the system's time domain must stay fixed within a compiled workflow.

    Args:
        sys: [DiscLTI][contrax.systems.DiscLTI] or
            [ContLTI][contrax.systems.ContLTI] system.
        Q: State cost matrix. Shape: `(n, n)`.
        R: Input cost matrix. Shape: `(m, m)`.
        solver: DARE for discrete systems or CARE for continuous systems.
        throw: Raise on numerical failure; otherwise return failure diagnostics.

    Returns:
        [LQRResult][contrax.types.LQRResult]: A bundle containing the optimal
            gain, Riccati solution, and closed-loop poles.

    Raises:
        TypeError: If `sys` is not a
            [DiscLTI][contrax.systems.DiscLTI] or
            [ContLTI][contrax.systems.ContLTI].

    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.05], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.05]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.05,
        ... )
        >>> result = cx.lqr(sys, jnp.eye(2), jnp.array([[1.0]]))
    """
    if isinstance(sys, DiscLTI):
        return dare(sys.A, sys.B, Q, R, solver=solver, throw=throw)
    elif isinstance(sys, ContLTI):
        return care(sys.A, sys.B, Q, R, solver=solver, throw=throw)
    else:
        raise TypeError(f"Expected DiscLTI or ContLTI, got {type(sys)}")

def place(
    sys: DiscLTI | ContLTI,
    poles: ArrayLike,
    *,
    solver: PolePlacement | None = None,
    throw: bool = False,
) -> PlaceResult:
    """Assign real/conjugate-paired poles and validate the achieved feedback.

    Uses native JAX KNV0 (real poles) or YT (real or complex pairs) linear
    algebra with design-time rank decisions and convergence control. Input
    subspaces use SVD so dependent and redundant input columns are supported.
    Small SISO failures may try bounded Ackermann fallback; every candidate
    must meet assignment and conditioning tolerances.

    Rank decisions and host-side matching make this a design-time routine. It
    requires float64 precision, promotes real model matrices to float64, and
    cannot be jitted, vmapped, or differentiated.

    Args:
        sys: Continuous or discrete real LTI model with `n > 0` and `m > 0`.
        poles: Finite requested poles. Shape: `(n,)`; values are real or
            occur in conjugate pairs.
        solver: PolePlacement method, stopping/acceptance tolerances and limits.
            Multiplicity beyond `rank(B)` requires the bounded SISO fallback;
            KNV0 and YT updates do not support that configuration.
        throw: Raise on failed assignment, convergence or conditioning checks.

    Returns:
        PlaceResult: Pole-placement gain and validation diagnostics. A failed
        result may contain a finite candidate for inspection; check
        `successful` before using it.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> sys = cx.ss([[0., 1.], [0., 0.]], [[0.], [1.]], [[1., 0.]], [[0.]])
        >>> result = cx.place(sys, [-1., -2.])
        >>> bool(result.successful)
        True
    """
    require_x64("place")
    solver = PolePlacement() if solver is None else solver
    if not isinstance(solver, PolePlacement):
        raise TypeError("place solver must be contrax.solvers.PolePlacement")
    if not isinstance(sys, (ContLTI, DiscLTI)):
        raise TypeError("place requires a continuous or discrete LTI model")
    if jnp.iscomplexobj(sys.A) or jnp.iscomplexobj(sys.B):
        raise ValueError("place requires real model matrices")
    try:
        np.asarray(sys.A)
        np.asarray(sys.B)
        requested = np.asarray(poles, dtype=complex)
    except jax.errors.TracerArrayConversionError as error:
        raise TypeError(
            "place is design-time; call it outside JAX transforms"
        ) from error
    A, B = jnp.asarray(sys.A, dtype=jnp.float64), jnp.asarray(sys.B, dtype=jnp.float64)
    if A.ndim != 2 or A.shape[0] == 0 or A.shape[1] != A.shape[0]:
        raise ValueError("A must be a nonempty square matrix")
    n = A.shape[0]
    if B.ndim != 2 or B.shape[0] != n or B.shape[1] == 0:
        raise ValueError("B must have shape (n, m), m>0")
    if requested.ndim != 1 or requested.shape[0] != n:
        raise ValueError("place requires exactly n requested poles")
    desired = _order_complex_poles(requested)
    complex_poles = bool(jnp.any(jnp.imag(desired) != 0))
    method = (
        ("YT" if complex_poles else "KNV0")
        if solver.method == "auto"
        else solver.method
    )
    if method == "KNV0" and complex_poles:
        raise ValueError("KNV0 does not support complex poles; select YT or auto")
    cc = float("inf")

    def result(
        K, actual, error, ec, status, steps=0, improvement=float("inf"), fallback=False
    ):
        answer = PlaceResult(
            K=K,
            poles=actual,
            status=jnp.asarray(status, dtype=jnp.int32),
            assignment_error=jnp.asarray(error),
            controllability_condition=jnp.asarray(cc),
            eigenvector_condition=jnp.asarray(ec),
            num_steps=jnp.asarray(steps, dtype=jnp.int32),
            relative_improvement=jnp.asarray(improvement),
            used_fallback=jnp.asarray(fallback),
            method="ackermann" if fallback else method,
        )
        if throw and status != SolverStatus.successful:
            raise RuntimeError(
                "Pole assignment failed; inspect assignment_error and conditioning"
            )
        return answer

    def failed(status):
        return result(
            jnp.full((B.shape[1], n), jnp.nan),
            jnp.full(n, jnp.nan, dtype=jnp.complex128),
            float("inf"),
            float("inf"),
            status,
        )

    if not bool(jnp.all(jnp.isfinite(A)) & jnp.all(jnp.isfinite(B))):
        return failed(SolverStatus.invalid_problem)
    controllability = _ctrb_matrix(A, B)
    if not bool(jnp.all(jnp.isfinite(controllability))):
        return failed(SolverStatus.numerical_failure)
    singular = jnp.linalg.svd(controllability, compute_uv=False)
    if not bool(jnp.all(jnp.isfinite(singular))):
        return failed(SolverStatus.numerical_failure)
    cc = float(jnp.where(singular[-1] > 0, singular[0] / singular[-1], jnp.inf))
    rank = int(
        jnp.sum(
            singular > singular[0] * max(controllability.shape) * jnp.finfo(A.dtype).eps
        )
    )
    if rank < n:
        return failed(SolverStatus.invalid_problem)
    if cc > solver.condition_limit:
        return failed(SolverStatus.numerical_failure)
    rankB = int(jnp.linalg.matrix_rank(B))
    repeated = (
        max(int(np.sum(np.asarray(desired) == pole)) for pole in np.asarray(desired))
        > rankB
    )
    allow_fallback = B.shape[1] == 1 and n <= solver.fallback_limit
    if repeated and not allow_fallback:
        raise ValueError("Requested multiplicity exceeds rank(B) for KNV0/YT")

    fallback = repeated
    if fallback:
        K, steps, improvement = _place_ackermann(A, B, desired), 0, 0.0
    else:
        K, steps, improvement = _place_robust(
            A, B, desired, method=method, rtol=solver.rtol, maxiter=solver.max_steps
        )
    actual, error, ec, good = _assess(A, B, K, desired, solver)
    if not good and allow_fallback and not fallback:
        alternate = _place_ackermann(A, B, desired)
        ap, ae, ac, accepted = _assess(A, B, alternate, desired, solver)
        if accepted or (
            not bool(jnp.all(jnp.isfinite(K)))
            and bool(jnp.all(jnp.isfinite(alternate)))
        ):
            K, actual, error, ec, good = alternate, ap, ae, ac, accepted
            steps, improvement, fallback = 0, 0.0, True
    status = (
        (
            SolverStatus.successful
            if improvement < solver.rtol
            else SolverStatus.max_steps_reached
        )
        if good
        else SolverStatus.numerical_failure
    )
    return result(K, actual, error, ec, status, steps, improvement, fallback)

def state_feedback(sys: DiscLTI | ContLTI, K: Array) -> DiscLTI | ContLTI:
    """Construct the closed-loop system under state feedback.

    Applies u = v - Kx, retaining v as the external input. The new matrices are
    A_cl = A - BK and C_cl = C - DK, with B, D and discrete dt unchanged.
    Nonzero direct feedthrough therefore contributes to the closed-loop output.

    Accepts either [DiscLTI][contrax.systems.DiscLTI] or
    [ContLTI][contrax.systems.ContLTI]. For external input `v`, the closed-loop
    matrices are `A_cl = A - B K` and `C_cl = C - D K`; `B`, `D`, and the
    discrete sample time are unchanged.

    Args:
        sys: Continuous or discrete-time system.
        K: State-feedback gain. Shape: `(m, n)`.

    Returns:
        [DiscLTI][contrax.systems.DiscLTI] |
            [ContLTI][contrax.systems.ContLTI]: A new system representing the
            closed-loop dynamics.

    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.05], [0.0, 1.0]]),
        ...     jnp.array([[0.0], [0.05]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ...     dt=0.05,
        ... )
        >>> result = cx.lqr(sys, jnp.eye(2), jnp.array([[1.0]]))
        >>> closed_loop = cx.state_feedback(sys, result.K)
    """
    if not isinstance(sys, (DiscLTI, ContLTI)):
        raise TypeError(f"Expected DiscLTI or ContLTI, got {type(sys)}")
    K = jnp.asarray(K)
    if K.shape != (sys.B.shape[1], sys.A.shape[0]):
        raise ValueError("K must have shape (number of inputs, number of states)")
    if isinstance(sys, DiscLTI):
        return dss(
            A=sys.A - sys.B @ K,
            B=sys.B,
            C=sys.C - sys.D @ K,
            D=sys.D,
            dt=sys.dt,
        )
    elif isinstance(sys, ContLTI):
        return ss(
            A=sys.A - sys.B @ K,
            B=sys.B,
            C=sys.C - sys.D @ K,
            D=sys.D,
        )
    else:
        raise TypeError(f"Expected DiscLTI or ContLTI, got {type(sys)}")