Systems¶

def block_matrix(
    row_block_sizes: Sequence[int],
    col_block_sizes: Sequence[int],
    blocks: dict[tuple[int, int], ArrayLike] | None = None,
    *,
    dtype: jnp.dtype | None = None,
) -> Array:
    """Assemble a dense matrix from block entries."""
    row_sizes = tuple(int(size) for size in row_block_sizes)
    col_sizes = tuple(int(size) for size in col_block_sizes)
    if not row_sizes or not col_sizes:
        raise ValueError("row_block_sizes and col_block_sizes must not be empty.")
    if any(size <= 0 for size in row_sizes + col_sizes):
        raise ValueError("Block sizes must contain only positive integers.")

    block_entries = {} if blocks is None else dict(blocks)

    # First pass: determine the fully promoted dtype across all user-supplied blocks.
    result_dtype = jnp.float32 if dtype is None else dtype
    for block_val in block_entries.values():
        result_dtype = jnp.result_type(result_dtype, jnp.asarray(block_val).dtype)

    # Second pass: build rows with a consistent dtype so zero-fill blocks match.
    rows = []
    for i, row_size in enumerate(row_sizes):
        row = []
        for j, col_size in enumerate(col_sizes):
            if (i, j) in block_entries:
                block = jnp.asarray(block_entries[(i, j)], dtype=result_dtype)
                if block.shape != (row_size, col_size):
                    raise ValueError(
                        f"Block {(i, j)} has shape {block.shape}, expected "
                        f"{(row_size, col_size)}."
                    )
            else:
                block = jnp.zeros((row_size, col_size), dtype=result_dtype)
            row.append(block)
        rows.append(row)

    return jnp.block(rows)

def block_observation(
    block_sizes: Sequence[int],
    block_indices: Sequence[int],
) -> Callable[[Array | float, Array, Array], Array]:
    """Build an output map that concatenates selected state blocks."""
    sizes = tuple(int(size) for size in block_sizes)
    indices = tuple(int(index) for index in block_indices)
    if not sizes:
        raise ValueError("block_sizes must not be empty.")
    if not indices:
        raise ValueError("block_indices must not be empty.")
    if any(index < 0 or index >= len(sizes) for index in indices):
        raise ValueError("block_indices must refer to valid state blocks.")

    def output(t: Array | float, x: Array, u: Array) -> Array:
        del t, u
        parts = partition_state(x, sizes)
        return jnp.concatenate([parts[index] for index in indices], axis=0)

    return output

def c2d(sys: ContLTI, dt: float, method: str = "zoh") -> DiscLTI:
    """Discretize a continuous-time state-space system.

    Converts a [ContLTI][contrax.systems.ContLTI] into a
    [DiscLTI][contrax.systems.DiscLTI] using either zero-order hold or the
    bilinear Tustin transform.

    The `zoh` path is the stronger current implementation and is differentiated
    through a custom VJP on the matrix exponential. `tustin` is available as a
    convenience discretization.

    Args:
        sys: Continuous-time system.
        dt: Sampling period. Shape: scalar.
        method: Discretization method. Supported values are `"zoh"` and
            `"tustin"`.

    Returns:
        [DiscLTI][contrax.systems.DiscLTI]: A discrete-time realization with
            the same output matrices and the requested sampling period.

    Raises:
        ValueError: If `method` is not one of the supported discretizations.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> sys_c = cx.ss(
        ...     jnp.array([[0.0, 1.0], [0.0, 0.0]]),
        ...     jnp.array([[0.0], [1.0]]),
        ...     jnp.eye(2),
        ...     jnp.zeros((2, 1)),
        ... )
        >>> sys_d = cx.c2d(sys_c, dt=0.05)
    """
    require_x64("c2d")
    if method == "zoh":
        return _c2d_zoh(sys, dt)
    elif method == "tustin":
        return _c2d_tustin(sys, dt)
    else:
        raise ValueError(f"Unknown discretization method: {method!r}")

def canonical_J(n_dofs: int, *, dtype: jnp.dtype | None = None) -> Array:
    """Construct the canonical skew-symmetric PHS structure matrix."""
    if n_dofs < 1:
        raise ValueError("n_dofs must be positive.")
    dtype = jnp.float32 if dtype is None else dtype
    eye = jnp.eye(n_dofs, dtype=dtype)
    zeros = jnp.zeros((n_dofs, n_dofs), dtype=dtype)
    return jnp.block([[zeros, eye], [-eye, zeros]])

def dss(
    A: ArrayLike,
    B: ArrayLike,
    C: ArrayLike,
    D: ArrayLike,
    dt: ArrayLike,
) -> DiscLTI:
    """Construct a discrete-time state-space system.

    Wraps discrete-time matrices in a
    [DiscLTI][contrax.systems.DiscLTI] pytree with MATLAB-familiar
    `ss(A, B, C, D, dt)` ergonomics.

    This is a lightweight constructor rather than a solver. The sampling period
    is stored as a plain array leaf instead of static metadata, which keeps
    batching over systems more natural in JAX.

    Args:
        A: State transition matrix. Shape: `(n, n)`.
        B: Input matrix. Shape: `(n, m)`.
        C: Output matrix. Shape: `(p, n)`.
        D: Feedthrough matrix. Shape: `(p, m)`.
        dt: Sampling period. Shape: scalar.

    Returns:
        [DiscLTI][contrax.systems.DiscLTI]: A discrete-time system with
            array-valued `A`, `B`, `C`, `D`, and `dt` fields.

    Raises:
        AssertionError: If the matrix shapes are inconsistent.

    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,
        ... )
    """
    A = jnp.asarray(A, dtype=float)
    B = jnp.asarray(B, dtype=float)
    C = jnp.asarray(C, dtype=float)
    D = jnp.asarray(D, dtype=float)
    dt = jnp.asarray(dt, dtype=float)
    n = A.shape[0]
    assert A.shape == (n, n)
    assert B.shape[0] == n
    assert C.shape[1] == n
    assert D.shape == (C.shape[0], B.shape[1])
    return DiscLTI(A=A, B=B, C=C, D=D, dt=dt)

def linearize(
    model_or_f: DynamicsLike,
    x0: Array,
    u0: Array,
    *,
    t: ArrayLike = 0.0,
) -> tuple[Array, Array]:
    """Linearize a dynamics function around an operating point.

    Computes the Jacobians `A = df/dx` and `B = df/du` at `(x0, u0)` using
    forward-mode automatic differentiation.

    This is one of the main JAX-native workflow bridges in Contrax. It works
    cleanly with `jit`, `vmap`, and differentiation as long as `f` is written
    as a pure JAX function without Python control flow over traced values.

    Args:
        model_or_f: A nonlinear system model or pure JAX dynamics function
            mapping `(x, u)` to `x_dot` or `x_next`.
        x0: State operating point. Shape: `(n,)`.
        u0: Input operating point. Shape: `(m,)`.
        t: Operating-point time for system-model linearization.

    Returns:
        tuple[Array, Array]: Jacobian pair `(A, B)`. Shapes: `(n, n)` and
        `(n, m)`.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> def dynamics(x, u):
        ...     return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])
        >>> A, B = cx.linearize(dynamics, jnp.zeros(2), jnp.zeros(1))
    """
    f = _coerce_dynamics(model_or_f)
    t = jnp.asarray(t, dtype=x0.dtype)
    A = jax.jacfwd(f, argnums=0)(x0, u0, t)
    B = jax.jacfwd(f, argnums=1)(x0, u0, t)
    return A, B

def linearize_ss(
    model_or_f: DynamicsLike,
    x0: Array,
    u0: Array,
    *,
    output: OutputFn | None = None,
    t: ArrayLike = 0.0,
) -> ContLTI:
    """Linearize dynamics and output maps into a continuous state-space model.

    Computes `A`, `B`, `C`, and `D` Jacobians at `(x0, u0)` and returns them as
    a [ContLTI][contrax.systems.ContLTI].

    This is the main bridge from nonlinear plant code into the linear control
    stack. Pair it with `c2d()` and `lqr()` when building local controllers
    around an operating point.

    Args:
        model_or_f: A nonlinear system model or pure JAX dynamics function
            mapping `(x, u)` to `x_dot`.
        x0: State operating point. Shape: `(n,)`.
        u0: Input operating point. Shape: `(m,)`.
        output: Output function mapping `(x, u)` to `y` when `model_or_f` is a
            plain dynamics callable. Omit this when passing a system model with
            its own output map.
        t: Operating-point time for system-model linearization.

    Returns:
        [ContLTI][contrax.systems.ContLTI]: A continuous-time state-space
            model with linearized `A`, `B`, `C`, and `D` matrices.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> def dynamics(x, u):
        ...     return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])
        >>> def output(x, u):
        ...     return x
        >>> sys_c = cx.linearize_ss(
        ...     dynamics,
        ...     jnp.zeros(2),
        ...     jnp.zeros(1),
        ...     output=output,
        ... )
    """
    if _is_system_model(model_or_f):
        if output is not None:
            raise ValueError("linearize_ss(sys, x0, u0) does not accept output=.")
        sys = model_or_f
        f = _coerce_dynamics(sys)
        h = _coerce_observation(sys)
    else:
        if output is None:
            raise ValueError(
                "linearize_ss(f, x0, u0, output=...) requires an output function."
            )
        f = _coerce_dynamics(model_or_f)

        def h(x: Array, u: Array, t_: Array) -> Array:
            del t_
            return output(x, u)

    t = jnp.asarray(t, dtype=x0.dtype)
    A = jax.jacfwd(f, argnums=0)(x0, u0, t)
    B = jax.jacfwd(f, argnums=1)(x0, u0, t)
    C = jax.jacfwd(h, argnums=0)(x0, u0, t)
    D = jax.jacfwd(h, argnums=1)(x0, u0, t)
    return ContLTI(A=A, B=B, C=C, D=D)

def nonlinear_system(
    dynamics: Callable,
    output: Callable | None = None,
    *,
    observation: Callable | None = None,
    dt: ArrayLike | None = None,
    state_dim: int | None = None,
    input_dim: int | None = None,
    output_dim: int | None = None,
) -> NonlinearSystem:
    """Construct a reusable nonlinear system model.

    Contrax uses [NonlinearSystem][contrax.systems.NonlinearSystem] as the
    shared contract for nonlinear simulation, linearization, and estimation.
    The intended callable shape is `dynamics(t, x, u)` and, when present,
    `output(t, x, u)`.

    Args:
        dynamics: State transition or vector-field callable with signature
            `(t, x, u)`.
        output: Optional output/measurement map with signature `(t, x, u)`.
            When omitted, the full state is used as the output.
        observation: Deprecated synonym for `output`. Pass only one of
            `output` or `observation`.
        dt: Optional discrete sample time. Omit for continuous models.
        state_dim: Optional static state dimension metadata.
        input_dim: Optional static input dimension metadata.
        output_dim: Optional static output dimension metadata.

    Returns:
        [NonlinearSystem][contrax.systems.NonlinearSystem]: A reusable
            nonlinear system model.
    """
    if output is not None and observation is not None:
        raise ValueError("Pass only one of output= or observation=.")
    observation = output if output is not None else observation
    dt_arr = None if dt is None else jnp.asarray(dt, dtype=float)
    return NonlinearSystem(
        dynamics=dynamics,
        observation=observation,
        dt=dt_arr,
        state_dim=state_dim,
        input_dim=input_dim,
        output_dim=output_dim,
    )

def parallel(
    sys1: DiscLTI | ContLTI,
    sys2: DiscLTI | ContLTI,
    *,
    sign: float = 1.0,
) -> DiscLTI | ContLTI:
    """Connect two systems in parallel and sum their outputs.

    Args:
        sys1: First system.
        sys2: Second system.
        sign: Output sign applied to `sys2`. Use `-1.0` for subtraction.

    Returns:
        [DiscLTI][contrax.systems.DiscLTI] |
            [ContLTI][contrax.systems.ContLTI]: Interconnected system.
    """
    _check_same_family(sys1, sys2)
    if sys1.B.shape[1] != sys2.B.shape[1]:
        raise ValueError("Parallel interconnection requires matching input dimensions.")
    if sys1.C.shape[0] != sys2.C.shape[0]:
        raise ValueError(
            "Parallel interconnection requires matching output dimensions."
        )

    A1, B1, C1, D1 = sys1.A, sys1.B, sys1.C, sys1.D
    A2, B2, C2, D2 = sys2.A, sys2.B, sys2.C, sys2.D
    n1, n2 = A1.shape[0], A2.shape[0]
    sign_arr = jnp.asarray(sign, dtype=A1.dtype)

    A = jnp.block(
        [
            [A1, jnp.zeros((n1, n2), dtype=A1.dtype)],
            [jnp.zeros((n2, n1), dtype=A2.dtype), A2],
        ]
    )
    B = jnp.vstack([B1, B2])
    C = jnp.hstack([C1, sign_arr * C2])
    D = D1 + sign_arr * D2

    if isinstance(sys1, DiscLTI):
        return DiscLTI(A=A, B=B, C=C, D=D, dt=sys1.dt)
    return ContLTI(A=A, B=B, C=C, D=D)

def partition_state(x: ArrayLike, block_sizes: Sequence[int]) -> tuple[Array, ...]:
    """Split a one-dimensional state vector into consecutive logical blocks."""
    x = jnp.asarray(x)
    sizes = tuple(int(size) for size in block_sizes)
    if x.ndim != 1:
        raise ValueError("partition_state() expects a one-dimensional state vector.")
    if not sizes or any(size <= 0 for size in sizes):
        raise ValueError("block_sizes must contain only positive integers.")
    total = sum(sizes)
    if x.shape[0] != total:
        raise ValueError(
            f"State length {x.shape[0]} does not match block_sizes sum {total}."
        )
    offsets = [0]
    for size in sizes:
        offsets.append(offsets[-1] + size)
    return tuple(x[offsets[i] : offsets[i + 1]] for i in range(len(sizes)))

def phs_diagnostics(
    sys: PHSSystem,
    x: ArrayLike,
    u: ArrayLike | None = None,
) -> PHSStructureDiagnostics:
    """Evaluate local structure and power-balance diagnostics for a PHS state."""
    x = jnp.asarray(x)
    u_arr = jnp.zeros((0,), dtype=x.dtype) if u is None else jnp.asarray(u)
    dim = x.shape[0]
    J = sys.J(x) if sys.J is not None else canonical_J(dim // 2, dtype=x.dtype)
    R = sys.R(x) if sys.R is not None else jnp.zeros((dim, dim), dtype=x.dtype)
    G = sys.G(x) if sys.G is not None else _zero_input_matrix(x, u_arr)
    grad_H = jax.grad(sys.H)(x)
    R_sym = symmetrize_matrix(R)
    # Reuse already-computed J, R, G rather than calling sys.dynamics() again.
    Gu = G @ u_arr
    xdot = (J - R) @ grad_H + Gu
    storage_rate = jnp.vdot(grad_H, xdot)
    dissipation_power = jnp.vdot(grad_H, R_sym @ grad_H)
    supplied_power = jnp.vdot(grad_H, Gu)
    # Standard PHS energy balance: dH/dt = supply - dissipation
    # => supplied_power - dissipation_power - storage_rate = 0
    return PHSStructureDiagnostics(
        skew_symmetry_error=jnp.max(jnp.abs(J + J.T)),
        dissipation_symmetry_error=jnp.max(jnp.abs(R - R.T)),
        min_dissipation_eigenvalue=jnp.min(jnp.linalg.eigvalsh(R_sym)),
        storage_rate=storage_rate,
        dissipation_power=dissipation_power,
        supplied_power=supplied_power,
        power_balance_residual=supplied_power - dissipation_power - storage_rate,
    )

def phs_system(
    H: Callable[[Array], Array],
    *,
    R: Callable[[Array], Array] | None = None,
    G: Callable[[Array], Array] | None = None,
    J: Callable[[Array], Array] | None = None,
    output: Callable | None = None,
    observation: Callable | None = None,
    dt: ArrayLike | None = None,
    state_dim: int | None = None,
    input_dim: int | None = None,
    output_dim: int | None = None,
) -> PHSSystem:
    """Construct a port-Hamiltonian system object.

    Args:
        H: Storage/Hamiltonian function.
        R: Optional dissipation map.
        G: Optional input map.
        J: Optional structure map.
        output: Optional output/measurement map with signature `(t, x, u)`.
        observation: Deprecated synonym for `output`. Pass only one of
            `output` or `observation`.
        dt: Optional discrete sample time.
        state_dim: Optional static state dimension metadata.
        input_dim: Optional static input dimension metadata.
        output_dim: Optional static output dimension metadata.

    Returns:
        [PHSSystem][contrax.systems.PHSSystem]: A reusable port-Hamiltonian
            model.
    """
    if output is not None and observation is not None:
        raise ValueError("Pass only one of output= or observation=.")
    observation = output if output is not None else observation
    dt_arr = None if dt is None else jnp.asarray(dt, dtype=float)
    return PHSSystem(
        H=H,
        R=R,
        G=G,
        J=J,
        observation=observation,
        dt=dt_arr,
        state_dim=state_dim,
        input_dim=input_dim,
        output_dim=output_dim,
    )