Analysis¶

def ctrb(sys: DiscLTI | ContLTI) -> Array:
    """Construct the controllability matrix of a state-space system.

    Forms `[B, AB, A^2 B, ..., A^{n-1} B]` for a continuous or discrete linear
    system.

    Use the singular values or a tolerance-aware rank of this matrix to assess
    controllability for a particular model and scaling.

    Args:
        sys: Continuous or discrete LTI system.

    Returns:
        Array: Controllability matrix. Shape: `(n, n * m)`.
    """
    A, B = sys.A, sys.B
    n = A.shape[0]

    def step(block, _):
        return A @ block, block

    _, blocks = jax.lax.scan(step, B, None, length=n)
    return jnp.transpose(blocks, (1, 0, 2)).reshape(n, n * B.shape[1])

def ctrb_gramian(
    sys: ContLTI | DiscLTI,
    *,
    t: float | Array | None = None,
    num_steps: int | None = None,
    solver: Lyapunov | None = None,
    throw: bool = False,
) -> LyapunovResult:
    """Compute a continuous or discrete controllability Gramian.

    With no horizon, solves the infinite-horizon equation for a strictly stable
    system. Set `t` for a continuous finite integral or `num_steps` for a
    discrete finite sum; finite horizons also apply to marginal or unstable
    systems. Requires float64 precision. A discrete horizon is a static Python
    integer, and each discrete sum term has unit weight.

    Args:
        sys: Continuous or discrete LTI realization.
        t: Nonnegative continuous-time integration horizon.
        num_steps: Nonnegative number of discrete Gramian terms.
        solver: Infinite-horizon Lyapunov configuration; defaults to stable doubling.
        throw: Raise on numerical failure or a nonstable infinite-horizon model.

    Returns:
        LyapunovResult: Gramian X, validated status and equation residual.
            Finite discrete num_steps reports the horizon term count.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> sys = cx.dss([[0.8]], [[1.0]], [[1.0]], [[0.0]], dt=1)
        >>> result = ctrb_gramian(sys, num_steps=20)
        >>> Wc = result.X
    """
    return _gramian(
        sys, observable=False, t=t, num_steps=num_steps, solver=solver, throw=throw
    )

def dcgain(sys: DiscLTI | ContLTI) -> Array:
    """Return the zero-frequency or steady-state gain of an LTI system.

    For continuous systems this evaluates `G(0)`. For discrete systems this
    evaluates `G(1)`.

    The corresponding resolvent must be nonsingular at zero frequency.

    Args:
        sys: Continuous or discrete LTI system.

    Returns:
        Array: DC gain matrix. Shape: `(p, m)`.
    """
    point = 0.0 if isinstance(sys, ContLTI) else 1.0
    return evalfr(sys, point)

def dlyap(
    A: Array, Q: Array, *, solver: Lyapunov | None = None, throw: bool = False
) -> LyapunovResult:
    """Solve the discrete Lyapunov equation AXA^T - X + Q = 0.

    The default uses squared Smith iteration, with residual validation and
    dense fallback bounded to eight states. Select `Lyapunov(method="dense")`
    for general uniquely solvable equations, including unstable systems.
    Uniqueness requires every eigenvalue product `lambda_i * lambda_j` to
    differ from one. Symmetric Q gives symmetric X.

    Requires float64 precision. JAX derivatives follow the implicit equation
    solve. Automatic dense fallback is limited to eight states; larger systems
    use the doubling method.

    Args:
        A: Real state transition matrix. Shape: `(n, n)`.
        Q: Real forcing matrix. Shape: `(n, n)`.
        solver: Algorithm, tolerances and iteration limit.
        throw: Raise on numerical failure, including under `jit`.

    Returns:
        LyapunovResult: Solution `X`, status, iteration count and residual norm.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> result = cx.dlyap(0.5 * jnp.eye(2), jnp.eye(2))
        >>> bool(result.successful)
        True
    """
    return lyapunov(A, Q, continuous=False, solver=solver, throw=throw)

def evalfr(sys: DiscLTI | ContLTI, point: complex | Array) -> Array:
    """Evaluate the transfer matrix at a complex frequency point.

    For continuous systems this computes `G(s) = C (sI - A)^-1 B + D`. For
    discrete systems it computes `G(z) = C (zI - A)^-1 B + D`.

    Use point evaluation for spot checks or as the basis of a frequency-response
    calculation. The resolvent must be nonsingular; evaluation at a realization
    pole may return nonfinite values even when a nonminimal realization has a
    removable cancellation. The matrices and point determine the promoted
    complex dtype.

    Args:
        sys: Continuous or discrete LTI system.
        point: Complex evaluation point `s` or `z`. Shape: scalar.

    Returns:
        Array: Transfer matrix evaluated at `point`. Shape: `(p, m)`.
    """
    if not isinstance(sys, (ContLTI, DiscLTI)):
        raise TypeError("evalfr requires a continuous or discrete LTI model")
    point = jnp.asarray(point, dtype=jnp.result_type(sys.A, point, 1j))
    if point.ndim != 0:
        raise ValueError("point must be scalar; use vmap or freqresp for a grid")
    n = sys.A.shape[0]
    identity = jnp.eye(n, dtype=point.dtype)
    A = sys.A.astype(point.dtype)
    B = sys.B.astype(point.dtype)
    C = sys.C.astype(point.dtype)
    D = sys.D.astype(point.dtype)
    resolvent_rhs = jnp.linalg.solve(point * identity - A, B)
    return C @ resolvent_rhs + D

def freqresp(sys: DiscLTI | ContLTI, omega: Array) -> Array:
    """Evaluate the frequency response over angular frequencies.

    Continuous systems are evaluated on the imaginary axis at `s = j omega`.
    Discrete systems are evaluated on the unit circle at
    `z = exp(j omega dt)`.

    Returns the complex transfer matrix at each frequency. An empty frequency
    vector returns shape `(0, p, m)`.

    Args:
        sys: Continuous or discrete LTI system.
        omega: Angular frequencies in rad/s. Shape: `(k,)`.

    Returns:
        Array: Complex frequency response. Shape: `(k, p, m)`.
    """
    omega = jnp.asarray(omega)
    if omega.ndim != 1 or jnp.iscomplexobj(omega):
        raise ValueError("omega must be a real vector of angular frequencies")
    omega = omega.astype(jnp.result_type(sys.A.real, omega, 1.0))
    if isinstance(sys, ContLTI):
        points = 1j * omega
    else:
        points = jnp.exp(1j * omega * sys.dt)
    return jax.vmap(lambda point: evalfr(sys, point))(points)

def lyap(
    A: Array, Q: Array, *, solver: Lyapunov | None = None, throw: bool = False
) -> LyapunovResult:
    """Solve the continuous Lyapunov equation AX + XA^T + Q = 0.

    The default uses squared Smith iteration after a Cayley transform, with
    residual validation and dense fallback bounded to eight states. Select
    `Lyapunov(method="dense")` for general uniquely solvable equations,
    including unstable systems. Uniqueness requires every eigenvalue sum
    `lambda_i + lambda_j` to be nonzero. Symmetric Q gives symmetric X.

    Requires float64 precision. JAX derivatives follow the implicit equation
    solve.
    Numerical failures return diagnostics unless `throw=True`.

    Args:
        A: Real state matrix. Shape: `(n, n)`.
        Q: Real forcing matrix. Shape: `(n, n)`.
        solver: Algorithm, tolerances and iteration limit.
        throw: Raise on numerical failure, including under `jit`.

    Returns:
        LyapunovResult: Solution `X`, status, iteration count and residual norm.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> result = cx.lyap(-jnp.eye(2), jnp.eye(2))
        >>> bool(result.successful)
        True
    """
    return lyapunov(A, Q, continuous=True, solver=solver, throw=throw)

def obsv(sys: DiscLTI | ContLTI) -> Array:
    """Construct the observability matrix of a state-space system.

    Forms `[C; CA; CA^2; ...; CA^{n-1}]` for a continuous or discrete linear
    system.

    Use the singular values or a tolerance-aware rank of this matrix to assess
    observability for a particular model and scaling.

    Args:
        sys: Continuous or discrete LTI system.

    Returns:
        Array: Observability matrix. Shape: `(n * p, n)`.
    """
    A, C = sys.A, sys.C
    n = A.shape[0]

    def step(block, _):
        return block @ A, block

    _, blocks = jax.lax.scan(step, C, None, length=n)
    return blocks.reshape(n * C.shape[0], n)

def obsv_gramian(
    sys: ContLTI | DiscLTI,
    *,
    t: float | Array | None = None,
    num_steps: int | None = None,
    solver: Lyapunov | None = None,
    throw: bool = False,
) -> LyapunovResult:
    """Compute a continuous or discrete observability Gramian.

    Uses the dual equation with forcing ``C.T @ C`` and transposed state
    dynamics. With no horizon, the system must be strictly stable. Set `t` for
    a finite continuous integral or `num_steps` for a finite discrete sum; zero
    horizons are valid. Requires float64 precision. A discrete horizon is a
    static Python integer, and each discrete sum term has unit weight.

    Args:
        sys: Continuous or discrete LTI realization.
        t: Nonnegative continuous-time integration horizon.
        num_steps: Nonnegative number of discrete Gramian terms.
        solver: Infinite-horizon Lyapunov configuration; defaults to stable doubling.
        throw: Raise on numerical failure or a nonstable infinite-horizon model.

    Returns:
        LyapunovResult: Gramian X, validated status and equation residual.
            Finite discrete num_steps reports the horizon term count.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> sys = cx.ss([[-1.0]], [[1.0]], [[1.0]], [[0.0]])
        >>> result = obsv_gramian(sys, t=2.0)
        >>> Wo = result.X
    """
    return _gramian(
        sys, observable=True, t=t, num_steps=num_steps, solver=solver, throw=throw
    )

def poles(sys: DiscLTI | ContLTI) -> Array:
    """Return the poles of a state-space system.

    The poles are the eigenvalues of the state matrix `A`. The returned dtype
    is generally complex, including for real systems with real poles.

    Args:
        sys: Continuous or discrete LTI system.

    Returns:
        Array: System poles. Shape: `(n,)`.
    """
    return jnp.linalg.eigvals(sys.A)

def quadratic_rollout_cost(
    sys: DiscLTI,
    K: Array,
    Q: Array,
    R: Array,
    x0: Array,
    *,
    num_steps: int,
    Q_terminal: Array | None = None,
) -> Array:
    """Evaluate a finite-horizon discrete LTI feedback cost.

    Computes ``sum(x_k.T @ Q @ x_k + u_k.T @ R @ u_k)`` for ``k=0,...,N-1``,
    with ``u_k = -K @ x_k``, and optionally adds
    ``x_N.T @ Q_terminal @ x_N``. Binary Gramian composition uses O(log N)
    matrix products and O(n**2) storage. `num_steps` is a static nonnegative
    Python integer; the system matrices, gain, costs, and initial state are JAX
    inputs. Long horizons with unstable closed-loop dynamics can overflow.

    Args:
        sys: Discrete LTI system with n states and m inputs.
        K: Feedback gain. Shape: `(m, n)`.
        Q: State cost. Shape: `(n, n)`.
        R: Input cost. Shape: `(m, m)`.
        x0: Initial state. Shape: `(n,)`.
        num_steps: Number of cost terms/transitions. Zero is valid.
        Q_terminal: Optional terminal-state cost. Shape: `(n, n)`.

    Returns:
        Array: Scalar quadratic cost for the single initial state.

    Raises:
        TypeError: For a continuous model or noninteger horizon.
        ValueError: For a negative horizon or inconsistent shapes.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> sys = cx.dss(jnp.eye(1), jnp.eye(1), jnp.eye(1), jnp.zeros((1, 1)), 1)
        >>> cx.quadratic_rollout_cost(sys, jnp.zeros((1, 1)), jnp.eye(1),
        ...                           jnp.eye(1), jnp.ones(1), num_steps=3)
        Array(3., dtype=float64)
    """
    if not isinstance(sys, DiscLTI):
        raise TypeError("quadratic_rollout_cost requires a discrete LTI system")
    if not isinstance(num_steps, int) or isinstance(num_steps, bool):
        raise TypeError("num_steps must be a static Python integer")
    if num_steps < 0:
        raise ValueError("num_steps must be nonnegative")
    K, Q, R, x0 = (jnp.asarray(x) for x in (K, Q, R, x0))
    n, m = sys.B.shape
    if K.shape != (m, n) or Q.shape != (n, n) or R.shape != (m, m) or x0.shape != (n,):
        raise ValueError("K, Q, R and x0 must have shapes (m,n), (n,n), (m,m), (n,)")
    terminal = jnp.zeros_like(Q) if Q_terminal is None else jnp.asarray(Q_terminal)
    if terminal.shape != (n, n):
        raise ValueError("Q_terminal must have shape (n, n)")
    A = sys.A - sys.B @ K
    W = Q + K.T @ R @ K
    phi, gramian = _compose_gramian(A, W, num_steps)
    gramian = gramian + phi.T @ terminal @ phi
    return x0 @ ((gramian + gramian.T) / 2) @ x0

def zeros(
    sys: ContLTI | DiscLTI, *, rtol: float | None = None, atol: float = 0.0
) -> Array:
    """Return finite invariant zeros of a real continuous or discrete LTI system.

    Computes rank drops of the Rosenbrock pencil relative to its normal rank
    for square or rectangular systems with any feedthrough rank. For minimal
    realizations, invariant zeros coincide with transmission zeros. Nonminimal
    realizations can also have decoupling zeros absent from a minimal transfer
    realization. The result contains finite zeros; infinite zeros are omitted.

    Orthogonal staircase reduction removes output constraints and infinite
    structure. A controllability staircase removes remaining right singular
    structure, then JAX eigensolvers solve the finite reduced pencil. Rank-
    dependent shapes and host-side decisions make this a design-time analysis
    routine; it does not support `jit`, `vmap`, or differentiation. Requires
    float64 mode. Rank decisions depend on scaling, so set tolerances for nearly
    rank-deficient realizations and compare transfer behavior after rescaling.

    Args:
        sys: Real continuous or discrete LTI realization.
        rtol: Relative numerical rank tolerance. Defaults to ten times machine
            epsilon times the larger Rosenbrock pencil dimension.
        atol: Absolute rank threshold. Nonnegative.

    Returns:
        Array: Finite invariant zeros including multiplicities. Complex shape `(k,)`.

    Raises:
        ValueError: For complex/nonfinite matrices or invalid rank tolerances.

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> import contrax as cx
        >>> import jax.numpy as jnp
        >>> sys = cx.ss(jnp.array([[-2.]]), jnp.ones((1, 1)),
        ...             jnp.array([[3.]]), jnp.array([[2.]]))
        >>> cx.zeros(sys)
        Array([-3.5+0.j], dtype=complex128)
    """
    require_x64("zeros")
    A, B, C, D = sys.A, sys.B, sys.C, sys.D
    if any(
        jnp.iscomplexobj(x) or not bool(jnp.all(jnp.isfinite(x))) for x in (A, B, C, D)
    ):
        raise ValueError("zeros requires finite real matrices")
    n, m, p = A.shape[0], B.shape[1], C.shape[0]
    rtol = (
        float(jnp.finfo(A.dtype).eps) * max(n + m, n + p, 1) * 10
        if rtol is None
        else rtol
    )
    if not math.isfinite(rtol) or not math.isfinite(atol) or rtol < 0 or atol < 0:
        raise ValueError("rank tolerances must be finite and nonnegative")
    compound = jnp.block([[A, B], [C, D]])
    scale = float(jnp.linalg.norm(compound))
    threshold = atol + rtol * scale

    # Reference: Emami-Naeini & Van Dooren (1982), "Computation of Zeros of
    # Linear Multivariable Systems", Automatica 18, 415-430; SLICOT AB08ND.
    # SVD supplies rank-revealing orthogonal coordinates. Eliminate the
    # nonsingular D block, then restrict x to ker(C_remaining). The discarded
    # state rows become the next algebraic output constraints. Each such step
    # removes at least one state (infinite structure), preserving finite rank
    # drops. Once constraints vanish, remove the right singular staircase.
    for _ in range(n + 1):
        if A.shape[0] == 0:
            return jnp.empty((0,), dtype=jnp.complex128)
        U, values, Vh, rank = _svd(D, threshold)
        B = B @ Vh.T
        C = U.T @ C
        if rank:
            A = A - (B[:, :rank] / values[:rank]) @ C[:rank]
        B = B[:, rank:]
        C = C[rank:]
        # Use an absolute floor from the original pencil to distinguish exact
        # structural zeros from roundoff created by orthogonal transformations.
        threshold = max(threshold, atol + rtol * float(jnp.linalg.norm(A)))
        _, _, Vh, rank_c = _svd(C, threshold)
        if rank_c == 0:
            finite = _uncontrollable_pencil(A, B, threshold)
            return jnp.linalg.eigvals(finite).astype(jnp.complex128)
        constrained, free = Vh[:rank_c].T, Vh[rank_c:].T
        A_free, B_free = free.T @ A @ free, free.T @ B
        C, D = constrained.T @ A @ free, constrained.T @ B
        A, B = A_free, B_free
    raise RuntimeError("invariant-zero staircase did not terminate")