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 blocks with positive static sizes and (row, column) index keys.

    Entries outside the declared grid or with incorrect shapes raise ValueError.
    Omitted blocks are zero; provided blocks determine the promoted dtype.
    """
    row_sizes = _block_sizes(row_block_sizes)
    col_sizes = _block_sizes(col_block_sizes)

    block_entries = {} if blocks is None else dict(blocks)
    for key in block_entries:
        if not isinstance(key, tuple) or len(key) != 2:
            raise ValueError("Block keys must be (row, column) index pairs")
        i, j = (dimension(index, "block index") for index in key)
        if i is None or j is None or i >= len(row_sizes) or j >= len(col_sizes):
            raise ValueError("Block index lies outside the declared block grid")

    # 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]:
    """Concatenate selected blocks, using positive static sizes and valid indices."""
    sizes = _block_sizes(block_sizes)
    indices = tuple(dimension(index, "block index") for index in block_indices)
    if not indices:
        raise ValueError("block_indices must not be empty.")
    if any(index is None 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 zero-order-hold path uses the JAX matrix exponential, including for
    derivatives. Requires float64 precision.

    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 requested sample time and the selected hold/frequency mapping.

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

    Examples:
        >>> import jax
        >>> jax.config.update("jax_enable_x64", True)
        >>> 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)
    """
    if not isinstance(sys, ContLTI):
        raise TypeError(
            "c2d requires continuous LTI; use sample_system for nonlinear models"
        )
    require_x64("c2d")
    # Keep weak scalar periods in the model precision, while promoting explicit
    # higher-precision array periods consistently with the computed matrices.
    dt, _ = _period(dt, jnp.result_type(sys.A.real, dt, 1.0))
    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 (2n, 2n) canonical structure for a positive static integer n."""
    n_dofs = dimension(n_dofs, "n_dofs", minimum=1)
    if n_dofs is None:
        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.

    Stores validated matrices and the sampling period in a `DiscLTI` pytree.
    The sample period is an array leaf, so it can vary across batched systems.
    Matrix dtypes are promoted together, preserving explicit float32 data.
    The sample time uses the resulting real matrix dtype.

    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: Finite positive sampling period. Shape: scalar.

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

    Raises:
        ValueError: If matrix shapes or the real, positive sample time are invalid.

    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, B, C, D = _lti_arrays(A, B, C, D)
    dt, invalid = _period(dt, A.real.dtype)
    result = DiscLTI(A=A, B=B, C=C, D=D, dt=dt)
    return eqx.error_if(result, invalid, "dt must be finite and positive")

def linearize(
    model_or_f: DynamicsLike,
    x0: Array,
    u0: Array,
    *,
    output: OutputFn | None = None,
    t: ArrayLike = 0.0,
) -> ContLTI | DiscLTI:
    """Linearize dynamics and output maps, preserving a model's time domain.

    Computes `A`, `B`, `C`, and `D` Jacobians at `(x0, u0)` using forward-mode
    automatic differentiation and returns them as a
    [ContLTI][contrax.systems.ContLTI] or [DiscLTI][contrax.systems.DiscLTI].
    A discrete model supplies its sample period; plain dynamics callables are
    interpreted as continuous vector fields. Supports jit, vmap and derivatives
    whenever the underlying callbacks do. Sampled models inherit their integration
    adjoint's transform support. State and input operating points must be vectors;
    use jacfwd directly for structured inputs such as FOH endpoint pairs.

    Use the local model for discretization and controller design. For example:

    ```python
    sys_lin = cx.linearize(nonlinear_sys, x_eq, u_eq)
    sys_d   = cx.c2d(sys_lin, dt=0.05)
    K       = cx.lqr(sys_d, Q, R).K
    cx.simulate(nonlinear_sys, x0, lambda t, x: -K @ x, ...)
    ```

    Args:
        model_or_f: A continuous/discrete LTI or nonlinear system, or a pure JAX
            vector field with signature `(t, x, u) → x_dot`.
        x0: State operating point. Shape: `(n,)`.
        u0: Input operating point. Shape: `(m,)`.
        output: Output function `(x, u) → y` for plain dynamics callables.
            When omitted, defaults to the full state (`C = I`, `D = 0`).
            Omit this when passing a system model with its own output map.
        t: Operating-point time. Relevant for time-varying system models.

    Returns:
        ContLTI | DiscLTI: Local perturbation model with Jacobians A, B, C, D.
            Discrete models retain their sample period. Constant operating-point
            offsets are not stored; states/inputs represent deviations from x0/u0.

    Examples:
        >>> import jax.numpy as jnp
        >>> import contrax as cx
        >>> def dynamics(t, x, u):
        ...     return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])
        >>> sys_lin = cx.linearize(dynamics, jnp.zeros(2), jnp.zeros(1))
    """
    x0, u0 = jnp.asarray(x0), jnp.asarray(u0)
    if x0.ndim != 1 or u0.ndim != 1:
        raise ValueError("x0 and u0 must be vectors; use jacfwd for structured inputs")
    dtype = jnp.result_type(x0, u0, 1.0)
    x0, u0 = x0.astype(dtype), u0.astype(dtype)
    t_arr = jnp.asarray(t, dtype=x0.real.dtype)
    if t_arr.ndim != 0:
        raise ValueError("t must be scalar")

    if isinstance(model_or_f, (ContLTI, DiscLTI)):
        if output is not None:
            raise ValueError("linearize(sys, x0, u0) does not accept output=.")
        if x0.shape != (model_or_f.A.shape[0],) or u0.shape != (model_or_f.B.shape[1],):
            raise ValueError("Operating-point dimensions must match the LTI model")
        return model_or_f

    if _is_system_model(model_or_f):
        if output is not None:
            raise ValueError("linearize(sys, x0, u0) does not accept output=.")
        f = model_or_f.dynamics
        h = model_or_f.output
    else:
        f = model_or_f
        if output is not None:

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

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

    A = jax.jacfwd(f, argnums=1)(t_arr, x0, u0)
    B = jax.jacfwd(f, argnums=2)(t_arr, x0, u0)
    C = jax.jacfwd(h, argnums=1)(t_arr, x0, u0)
    D = jax.jacfwd(h, argnums=2)(t_arr, x0, u0)
    dt = _system_dt(model_or_f) if _is_system_model(model_or_f) else None
    return ss(A, B, C, D) if dt is None else dss(A, B, C, D, dt)

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.

    The model supplies dynamics with signature `dynamics(t, x, u)` and an
    optional output map with signature `output(t, x, u)`. Use it for nonlinear
    simulation, linearization, and estimation.

    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 finite positive real scalar sample time. Omit for
            continuous models.
        state_dim: Optional nonnegative static integer state dimension.
        input_dim: Optional static input dimension. Enables omitted zero controls
            in recursive estimators; otherwise provide us for controlled models.
        output_dim: Optional static output dimension hint. Same caveat as
            `state_dim`.

    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, invalid = (None, False) if dt is None else sample_period(dt)
    result = NonlinearSystem(
        dynamics=dynamics,
        observation=observation,
        dt=dt_arr,
        state_dim=dimension(state_dim, "state_dim"),
        input_dim=dimension(input_dim, "input_dim"),
        output_dim=dimension(output_dim, "output_dim"),
    )
    return (
        result
        if dt is None
        else eqx.error_if(result, invalid, "dt must be finite and positive")
    )

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.
    """
    mismatch = _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=jnp.result_type(A1, A2, D1, D2, sign, 1.0))
    if sign_arr.ndim != 0:
        raise ValueError("sign must be scalar")

    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):
        result = DiscLTI(A=A, B=B, C=C, D=D, dt=sys1.dt)
        return eqx.error_if(
            result, mismatch, "Discrete interconnections require matching dt"
        )
    return ContLTI(A=A, B=B, C=C, D=D)

def partition_state(x: ArrayLike, block_sizes: Sequence[int]) -> tuple[Array, ...]:
    """Split a vector into blocks with positive static integer sizes summing to n."""
    x = jnp.asarray(x)
    sizes = _block_sizes(block_sizes)
    if x.ndim != 1:
        raise ValueError("partition_state() expects a one-dimensional state vector.")
    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)))