### Contrax Development Setup

Source: https://github.com/givani30/contrax/blob/master/README.md

Clone the Contrax repository, install development dependencies, set up pre-commit hooks, and run tests.

```bash
git clone https://github.com/givani30/Contrax.git
cd Contrax
uv sync --group dev
uv run pre-commit install
uv run pytest tests/ -q
```

--------------------------------

### Install Contrax

Source: https://github.com/givani30/contrax/blob/master/docs/getting-started.md

Use this command to install the library for local development.

```bash
uv sync
```

--------------------------------

### Setup Execution Model

Source: https://github.com/givani30/contrax/blob/master/docs/examples/lqr-optimal-execution.md

Initializes the Contrax execution model, which includes a DiscLTI system and an LQR solver. This is the basic setup for the execution problem.

```python
from jax import numpy as jnp

from contrax.execution import DiscLTI
from contrax.execution.lqr import solve_lqr

# Problem parameters
horizon = 20

# Build the execution model
# A is the state transition matrix, B is the control input matrix.
# For a single state x_k representing inventory, A = [[1.0]] and B = [[-1.0]].
A = jnp.array([[1.0]])
B = jnp.array([[-1.0]])

# The execution problem is a DiscLTI system.
# The state is remaining inventory, and the control is the change in inventory.
# The system dynamics are x_{k+1} = x_k + u_k, where u_k is the signed inventory change.
# The sell quantity is q_k^sell = -u_k.
model = DiscLTI(A, B)

```

--------------------------------

### Minimal LQR Example

Source: https://github.com/givani30/contrax/blob/master/docs/api/control.md

Demonstrates a minimal example of setting up a discrete-time system and performing an LQR design. Ensure JAX is configured for 64-bit precision.

```python
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)
```

--------------------------------

### Example Script Output

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/continuous-lqr.md

Expected console output from running the continuous LQR example script.

```text
Continuous LQR — double integrator
  K = [[1.         1.73205081]]
  closed-loop poles = [-0.8660254+0.5j -0.8660254-0.5j]
  all poles stable  = True
  residual norm     = 1.776357e-15
  x[0]  = [1. 0.]
  x[-1] = [-0.00023873  0.00033244]  (should be near zero)
  time horizon      = 10.000 s

Gradient smoke test (d/d(log q) of settling cost):
  grad = -3.848696  (finite: True)

All assertions passed.
```

--------------------------------

### Clone and Set Up Development Environment

Source: https://github.com/givani30/contrax/blob/master/CONTRIBUTING.md

Clone the Contrax repository and install development dependencies using `uv`.

```bash
git clone https://github.com/givani30/Contrax.git
cd Contrax
uv sync --group dev
```

--------------------------------

### Install Contrax Package

Source: https://github.com/givani30/contrax/blob/master/README.md

Install the Contrax library using pip. Requires Python 3.11+ and JAX 0.4+.

```bash
pip install contrax
```

--------------------------------

### Minimal Closed-Loop Simulation Example

Source: https://github.com/givani30/contrax/blob/master/docs/api/simulation.md

Demonstrates a basic closed-loop simulation using `cx.simulate`. This example sets up a discrete-time system, computes an LQR controller, and then simulates the system's response to an initial state using the computed controller. Ensure JAX 64-bit precision is enabled.

```python
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,
)
K = cx.lqr(sys, jnp.eye(2), jnp.array([[1.0]])).K
ts, xs, ys = cx.simulate(
    sys,
    jnp.array([1.0, 0.0]),
    lambda t, x: -K @ x,
    num_steps=80,
)
```

--------------------------------

### Setup and Model Configuration

Source: https://github.com/givani30/contrax/blob/master/docs/examples/foh-estimation.md

Initializes the system parameters and defines the continuous dynamics model for the Van der Pol oscillator.

```python
--8<-- "examples/foh_estimation.py:setup"
```

```python
--8<-- "examples/foh_estimation.py:model"
```

--------------------------------

### Setup Structured Nonlinear Model

Source: https://github.com/givani30/contrax/blob/master/docs/examples/structured-nonlinear-estimation.md

Defines the continuous structured model, builds the measurement map, and samples the discrete estimation model. Requires setup for the system dynamics.

```python
from contrax.systems import phs_system
from contrax.estimation import block_observation, sample_system

# Define the continuous structured model
system = phs_system(
    H=lambda q, p: 0.5 * (1.2 * q**2 + p**2),
    J=lambda q, p: [[0, 1], [-1, 0]],
    R=lambda q, p: [[0, 0], [0, 1]],
    u=lambda q, p: [[1], [0]],
)

# Build the position-only measurement map
# The state is (q, p), so we select the first component for measurement
observation_map = block_observation(system, indices=[0])

# Sample the continuous system to get a discrete estimation model
# dt is the time step, and noise_cov is the covariance of the process noise
est_model = sample_system(system, dt=0.01, noise_cov=0.01)

```

--------------------------------

### Install Pre-commit Hooks

Source: https://github.com/givani30/contrax/blob/master/CONTRIBUTING.md

Install pre-commit hooks to automatically run code style checks and formatting before commits.

```bash
uv run pre-commit install
```

--------------------------------

### Simulation Output Summary

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/linearize-lqr-simulate.md

Example output from running the simulation script.

```text
Linearize -> c2d -> lqr -> simulate
closed-loop poles = [0.96558903+0.04716561j 0.96558903-0.04716561j]
initial state norm = 0.250000
final state norm   = 0.026099
final state        = [-0.02007372  0.01667922]
time horizon       = 3.950 s
```

--------------------------------

### Minimal Kalman Filter Example

Source: https://github.com/givani30/contrax/blob/master/docs/api/estimation.md

Demonstrates the basic usage of the `kalman()` function for linear Gaussian filtering. Requires `jax.numpy` and `contrax`. Ensure system matrices and noise covariances are correctly defined.

```python
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,
)
result = cx.kalman(
    sys,
    Q_noise=1e-3 * jnp.eye(2),
    R_noise=1e-2 * jnp.eye(2),
    ys=jnp.zeros((20, 2)),
)
```

--------------------------------

### Setup for Nonlinear MHE

Source: https://github.com/givani30/contrax/blob/master/docs/examples/nonlinear-mhe.md

Sets up the nonlinear MHE problem, including system dynamics, parameterization, and initial conditions. This code is essential for defining the problem to be solved.

```python
from math import sin
from numpy import array, pi, cos

from proxddp import dynamics, constraints, meshcat_scene, quaternions
from proxddp.utils import positive_softplus

# system parameters
g = 9.81  # m/s^2
l = 0.5   # m

def pendulum_ode(x, u, params):
    # x = [theta, theta_dot]
    # u = [torque]
    theta, theta_dot = x
    torque = u[0]
    # state-dependent damping
    b_raw = params[0]
    b = positive_softplus(b_raw)
    return array([
        theta_dot,
        -g / l * sin(theta) - b * theta_dot + torque,
    ])


def pendulum_ode_with_param(x, u, params):
    # x = [theta, theta_dot, b_raw]
    # u = [torque]
    theta, theta_dot, b_raw = x
    torque = u[0]
    b = positive_softplus(b_raw)
    return array([
        theta_dot,
        -g / l * sin(theta) - b * theta_dot + torque,
        0.0,  # b_raw has identity dynamics (random walk)
    ])


def pendulum_measurement(x, params):
    # x = [theta, theta_dot, b_raw]
    theta, _, _ = x
    return array([theta])


def setup_pendulum_mhe(dt, T):
    # state: [theta, theta_dot, b_raw]
    # input: [torque]
    nx = 3
    nu = 1
    ndx = 3

    # dynamics
    ode = dynamics.ODEProblem(pendulum_ode_with_param, None, nx, nu)
    f = dynamics.IntegratorEuler(ode, dt)

    # measurement model
    # y = theta
    h = lambda x, params: pendulum_measurement(x, params)

    # initial state
    x0 = array([0.0, 0.0, 0.0])  # [theta, theta_dot, b_raw]

    # noise
    Q_noise = array([[1e-4, 0.0, 0.0], [0.0, 1e-4, 0.0], [0.0, 0.0, 1e-4]])
    R_noise = array([[1e-3]])

    # create MHE problem
    mhe = proxddp.MHEProblem(nx, nu, ndx)
    mhe.add_dynamics(f, Q_noise)
    mhe.add_measurement(h, R_noise, Q_noise)

    # create trajectory
    # initial guess for state and control
    # state: [theta, theta_dot, b_raw]
    # control: [torque]
    # b_raw is unconstrained, so we can initialize it to 0
    # b = softplus(0) = log(1+exp(0)) = log(2) = 0.693
    # we want b > 0, so b_raw=0 is a valid initialization
    x_init = array([0.0, 0.0, 0.0])
    u_init = array([0.0])
    mhe.init(x_init, u_init)

    return mhe

```

--------------------------------

### Example Script Output

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/kalman-filtering.md

Displays the summary output generated by the Kalman filtering and RTS smoothing script.

```text
Kalman filtering and RTS smoothing
final measurement        = 1.010000
mid filtered position    = 1.060639
mid smoothed position    = 1.003812
final filtered position  = 1.024387
final smoothed position  = 1.024387
final filtered velocity  = 0.055978
innovation norm          = 0.932813
final covariance trace   = 0.007035
```

--------------------------------

### Differentiable Closed-Loop Simulation with Contrax LQR

Source: https://github.com/givani30/contrax/blob/master/docs/index.md

This example demonstrates differentiating through a closed-loop simulation using Contrax's LQR controller. It sets up a discrete-time state-space system, defines a cost function based on LQR control and simulation, and then computes the cost and its gradients using JAX's `jit` and `value_and_grad`. This is useful for automatically tuning controller parameters by optimizing performance metrics.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
B = jnp.array([[0.0], [0.05]])
SYS = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)
X0 = jnp.array([1.0, 0.0])


def closed_loop_cost(log_q_diag, log_r):
    Q = jnp.diag(jnp.exp(log_q_diag))
    R = jnp.exp(log_r)[None, None]
    K = cx.lqr(SYS, Q, R).K
    _, xs, _ = cx.simulate(SYS, X0, lambda t, x: -K @ x, num_steps=80)
    return jnp.sum(xs**2)


objective_and_grad = jax.jit(jax.value_and_grad(closed_loop_cost, argnums=(0, 1)))
cost, grads = objective_and_grad(jnp.zeros(2), jnp.array(0.0))
```

--------------------------------

### Complete Working Code for Missing Measurements

Source: https://github.com/givani30/contrax/blob/master/docs/how-to/handle-missing-measurements.md

This snippet shows a complete, runnable example of using Contrax's `ekf_step` with `has_measurement=False`. It includes necessary imports, system model definitions, and the `one_step` function demonstrating the pattern. Ensure dummy measurements have the correct shape.

```python
import jax
import jax.numpy as jnp
import contrax as cx

Q = jnp.array([[1e-3]])
R = jnp.array([[1e-2]])


def f(x, u):
    return jnp.array([0.8 * x[0] + u[0]])


def h(x):
    return jnp.array([x[0] ** 2])


@jax.jit
def one_step(x, P, y, u, has_measurement):
    return cx.ekf_step(
        f,
        Q,
        R,
        y,
        u,
        x,
        P,
        observation=h,
        has_measurement=has_measurement,
    )


x = jnp.array([1.0])
P = jnp.array([[0.5]])
y = jnp.array([10.0])       # dummy value with the right shape
u = jnp.array([0.0])

x_next, P_next, innovation = one_step(x, P, y, u, jnp.array(False))
```

--------------------------------

### Setup Continuous Nonlinear System and FOH Bridge

Source: https://github.com/givani30/contrax/blob/master/docs/examples/continuous-nonlinear-estimation.md

Defines the continuous nonlinear system for a damped pendulum and sets up the estimation bridge using first-order hold (FOH) for input interpolation. This approach keeps the plant model continuous and nonlinear, allowing Contrax to handle the discrete sampling and estimation.

```python
from math import sin
from numpy import array, pi, cos

from contrax.system import NonlinearSystem
from contrax.estimation import sample_system


def setup():
    # pendulum dynamics
    def f(x, u):
        theta, dtheta = x
        return array([dtheta, -0.35 * dtheta - sin(theta) + u])

    # angle measurement
    def h(x, u):
        theta, _ = x
        return array([theta])

    # system parameters
    dt = 0.05  # sample time
    x0 = array([pi, 0.0])  # initial state
    u0 = 0.0  # initial input

    # continuous system
    system = NonlinearSystem(f, h, dt, x0, u0)

    # discrete estimator model
    # use foh for input interpolation
    estimator_model = sample_system(system, input_interpolation="foh")

    return system, estimator_model

```

--------------------------------

### Setup and Objective Function for Differentiable LQR

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/differentiable-lqr.md

Sets up the system and defines the closed-loop cost objective. This code is part of a larger JAX objective where controller weights are trainable parameters.

```python
import jax
import jax.numpy as jnp

from contrax.control import lqr
from contrax.systems import LinearSystem
from contrax.simulate import simulate

# System parameters
A = jnp.array([[1.0, 0.1], [0.0, 0.9]])
B = jnp.array([[0.0], [0.1]])

system = LinearSystem(A, B)

# Optimization parameters
# Q and R are parameterized by log-parameters to ensure positivity
def Q_from_log(log_diag_Q):
    return jnp.diag(jnp.exp(log_diag_Q))

def R_from_log(log_R):
    return jnp.array([[jnp.exp(log_R)]])

# Closed-loop cost objective
def closed_loop_cost(log_diag_Q, log_R):
    Q = Q_from_log(log_diag_Q)
    R = R_from_log(log_R)

    # LQR solve to get gain K
    K = lqr(system, Q, R)

    # Simulate closed-loop system
    T = 50  # simulation horizon
    x0 = jnp.zeros(2)
    u0 = None
    cl_state = simulate(system, K, x0, T, u0=u0)

    # Cost is sum of state and input costs
    state_cost = jnp.sum(cl_state.states**2, axis=1)
    input_cost = jnp.sum(cl_state.inputs**2, axis=1)
    return jnp.sum(state_cost) + 0.1 * jnp.sum(input_cost)
```

--------------------------------

### Define Pendulum Dynamics

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/linearize-lqr-simulate.md

Setup and dynamics definition for the pendulum model.

```python
--8<-- "examples/linearize_lqr_simulate.py:setup"

--8<-- "examples/linearize_lqr_simulate.py:dynamics"
```

--------------------------------

### Create and Analyze a State-Space System in Python

Source: https://github.com/givani30/contrax/blob/master/docs/api/analysis.md

Demonstrates creating a state-space system and computing its poles, DC gain, and controllability Gramian. Requires JAX 64-bit precision.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

sys = cx.ss(
    jnp.array([[0.0, 1.0], [-2.0, -0.4]]),
    jnp.array([[0.0], [1.0]]),
    jnp.array([[1.0, 0.0]]),
    jnp.zeros((1, 1)),
)

poles = cx.poles(sys)
dc = cx.dcgain(sys)
Wc = cx.ctrb_gramian(sys, t=2.0)
```

--------------------------------

### Simulate Discrete and Continuous Systems

Source: https://context7.com/givani30/contrax/llms.txt

Demonstrates simulating a discrete system with an LQR controller and a continuous system using Contrax.

```python
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,
)

K = cx.lqr(sys, jnp.eye(2), jnp.array([[1.0]])).K
x0 = jnp.array([1.0, 0.0])  # Initial state: position=1, velocity=0

# Discrete simulation
ts, xs, ys = cx.simulate(
    sys, x0,
    policy=lambda t, x: -K @ x,  # State feedback
    num_steps=100,               # 100 discrete steps
)
# ts.shape = (100,), xs.shape = (101, 2), ys.shape = (100, 2)

# Continuous system simulation (requires diffrax)
sys_c = cx.ss(
    jnp.array([[0.0, 1.0], [-1.0, -0.1]]),
    jnp.array([[0.0], [1.0]]),
    jnp.eye(2),
    jnp.zeros((2, 1)),
)
ts_c, xs_c, ys_c = cx.simulate(
    sys_c, x0,
    policy=lambda t, x: jnp.zeros(1),
    duration=10.0,  # Simulate for 10 seconds
    dt=0.01,        # Output sample spacing
)
```

--------------------------------

### JIT The Full Closed-Loop Path

Source: https://github.com/givani30/contrax/blob/master/docs/examples/jax-native-workflows.md

Compile the entire controller design and simulation pipeline into a single JAX-optimized function.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
B = jnp.array([[0.0], [0.05]])
sys = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)


@jax.jit
def run(q_scale):
    Q = jnp.diag(jnp.array([10.0, 1.0]) * q_scale)
    R = jnp.array([[1.0]])
    K = cx.lqr(sys, Q, R).K
    return cx.simulate(sys, jnp.array([1.0, 0.0]), lambda t, x: -K @ x, num_steps=60)


ts, xs, ys = run(jnp.array(1.0))
```

--------------------------------

### Define Continuous System

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/continuous-lqr.md

Sets up the double integrator state-space model using the Contrax library.

```python
--8<-- "examples/continuous_lqr.py:setup"

--8<-- "examples/continuous_lqr.py:system"
```

--------------------------------

### Complete Working Code for Batch Controller Design

Source: https://github.com/givani30/contrax/blob/master/docs/how-to/batch-controller-design.md

This code snippet demonstrates a complete workflow for batch controller design using Contrax. It includes defining the system dynamics, output function, design function, and then applying `jax.vmap` to compute controller gains across multiple operating points. Ensure JAX float64 is enabled for accurate calculations.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx


def pendulum(x, u):
    return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])

def output(x, u):
    return x

def design(x_eq, u_eq):
    sys_c = cx.linearize_ss(pendulum, x_eq, u_eq, output=output)
    sys_d = cx.c2d(sys_c, dt=0.05)
    return cx.lqr(sys_d, jnp.eye(2), jnp.ones((1, 1))).K

x_eqs = jnp.array([[0.0, 0.0], [0.1, 0.0], [-0.1, 0.0]])
u_eqs = jnp.zeros((3, 1))
Ks = jax.jit(jax.vmap(design))(x_eqs, u_eqs)
```

--------------------------------

### Nonlinear System Initialization and Sampling

Source: https://github.com/givani30/contrax/blob/master/docs/api/estimation.md

Initializes a continuous-time nonlinear system and then samples it to obtain a discrete-time representation suitable for EKF or UKF. The dt parameter specifies the sampling time.

```python
sys_c = cx.nonlinear_system(f_continuous, output=h, dt=None)
sys_d = cx.sample_system(sys_c, dt=0.1)
```

--------------------------------

### Design Batched Controllers with vmap

Source: https://context7.com/givani30/contrax/llms.txt

Uses JAX's vmap to parallelize the design of multiple controllers across different Q matrix configurations.

```python
import jax
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,
)

# Batch of Q matrices with different position weights
q_weights = jnp.array([1.0, 5.0, 10.0, 50.0])  # Different position penalties

def design_controller(q_weight):
    Q = jnp.diag(jnp.array([q_weight, 1.0]))
    R = jnp.array([[1.0]])
    return cx.lqr(sys, Q, R).K

# Vectorize over different Q configurations
batch_design = jax.vmap(design_controller)
K_batch = batch_design(q_weights)  # Shape (4, 1, 2)
```

--------------------------------

### Simulate Open-Loop Input with lsim

Source: https://context7.com/givani30/contrax/llms.txt

Simulates a discrete system given a specific input sequence.

```python
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,
)

# Input sequence: step input
T = 50
us = jnp.ones((T, 1))  # Unit step for 50 samples

x0 = jnp.zeros(2)  # Start at origin

ts, xs, ys = cx.lsim(sys, us, x0=x0)
# ts.shape = (50,), xs.shape = (51, 2), ys.shape = (50, 2)
```

--------------------------------

### Create Continuous-Time State-Space System (ss)

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a continuous-time LTI system using the `ss` function. Ensure JAX 64-bit precision is enabled for numerical stability.

```python
import jax
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

# Double integrator system (position + velocity)
A = jnp.array([[0.0, 1.0], [0.0, 0.0]])  # State matrix
B = jnp.array([[0.0], [1.0]])              # Input matrix (force -> acceleration)
C = jnp.eye(2)                             # Output: full state
D = jnp.zeros((2, 1))                      # No feedthrough

sys_continuous = cx.ss(A, B, C, D)
# sys_continuous.A, sys_continuous.B, sys_continuous.C, sys_continuous.D available
```

--------------------------------

### Differentiable Control Objective with JAX

Source: https://github.com/givani30/contrax/blob/master/README.md

Demonstrates integrating a controller design step (LQR) within a differentiable objective function using JAX. This allows for gradient-based optimization of control parameters.

```python
import jax
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
B = jnp.array([[0.0], [0.05]])
SYS = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)
X0 = jnp.array([1.0, 0.0])


def closed_loop_cost(log_q_diag, log_r):
    K = cx.lqr(SYS, jnp.diag(jnp.exp(log_q_diag)), jnp.exp(log_r)[None, None]).K
    _, xs, _ = cx.simulate(SYS, X0, lambda t, x: -K @ x, num_steps=80)
    return jnp.sum(xs**2) + 1e-2 * jnp.sum((xs[:-1] @ K.T) ** 2)


cost, (dq, dr) = jax.jit(jax.value_and_grad(closed_loop_cost, argnums=(0, 1)))(
    jnp.zeros(2), jnp.array(0.0)
)
```

--------------------------------

### Format Code with Ruff

Source: https://github.com/givani30/contrax/blob/master/CONTRIBUTING.md

Apply code formatting rules to the entire project using `ruff format`.

```bash
uv run ruff format .
```

--------------------------------

### cx.nonlinear_system - Create Nonlinear System Model

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a reusable nonlinear system model with dynamics and output functions.

```APIDOC
## cx.nonlinear_system

### Description
Constructs a reusable nonlinear system model with dynamics `dynamics(t, x, u)` and output `output(t, x, u)`.

### Parameters
- **dynamics** (callable) - Required - Dynamics function
- **output** (callable) - Required - Output function
- **dt** (float/None) - Optional - Sample time (None for continuous)
- **state_dim** (int) - Required - State dimension
- **input_dim** (int) - Required - Input dimension
- **output_dim** (int) - Required - Output dimension
```

--------------------------------

### Continuous LQR Workflow Overview

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/continuous-lqr.md

Conceptual representation of the continuous-time control pipeline.

```text
continuous state-space model -> care-backed lqr -> continuous simulate -> gradient smoke test
```

--------------------------------

### Run Unscented Kalman Filter (UKF)

Source: https://context7.com/givani30/contrax/llms.txt

Executes a UKF using sigma-point transforms. Supports tuning parameters like alpha, beta, and kappa for nonlinear systems.

```python
import jax.numpy as jnp
import contrax as cx

# Nonlinear dynamics
def f(x, u):
    return jnp.array([x[0] + 0.1 * x[1], x[1] + u[0]])

# Highly nonlinear observation
def h(x):
    return jnp.array([x[0] ** 2])  # Squared position

Q_noise = 1e-3 * jnp.eye(2)
R_noise = 1e-2 * jnp.eye(1)
ys = jnp.ones((10, 1))
us = jnp.zeros((10, 1))

result = cx.ukf(
    f, Q_noise, R_noise, ys, us,
    x0=jnp.ones(2),
    P0=jnp.eye(2),
    observation=h,
    alpha=1.0,   # Sigma-point spread
    beta=2.0,    # Prior parameter (2.0 for Gaussian)
    kappa=0.0,   # Secondary scaling
)

# UKFResult includes additional diagnostics
x_hat = result.x_hat
innovations = result.innovations
log_likelihood = result.log_likelihood_terms
```

--------------------------------

### Data Generation and EKF Filtering

Source: https://github.com/givani30/contrax/blob/master/docs/examples/foh-estimation.md

Simulates the system with FOH inputs and executes the EKF estimation loop.

```python
--8<-- "examples/foh_estimation.py:generate-data"
```

```python
--8<-- "examples/foh_estimation.py:filter"
```

--------------------------------

### Design Controllers Over Batches Of Operating Points

Source: https://github.com/givani30/contrax/blob/master/docs/examples/jax-native-workflows.md

Use jax.vmap to parallelize controller design across multiple operating points or conditions.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx


def pendulum(x, u):
    return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])


def sensor(x, u):
    return x


def design(x_eq, u_eq):
    sys_c = cx.linearize_ss(pendulum, x_eq, u_eq, output=sensor)
    sys_d = cx.c2d(sys_c, dt=0.05)
    return cx.lqr(sys_d, jnp.eye(2), jnp.ones((1, 1))).K


x_eqs = jnp.array([[0.0, 0.0], [0.1, 0.0], [-0.1, 0.0]])
u_eqs = jnp.zeros((3, 1))
batched_design = jax.jit(jax.vmap(design))
Ks = batched_design(x_eqs, u_eqs)
```

--------------------------------

### Print Summary of Results

Source: https://github.com/givani30/contrax/blob/master/docs/examples/lqr-optimal-execution.md

Prints a summary of the baseline controller, the tuned design parameters, and the results from batched execution. This provides a quick overview of the optimization outcomes.

```text
LQR optimal execution
baseline gain          = [[0.87695257]]
initial tuning loss    = 0.052402
final tuning loss      = 0.046743
tuned inventory risk   = 0.116998
tuned trading cost     = 0.156547
first sell quantities  = [8.76952567e-01 1.07906744e-01 1.32776339e-02 ...]
batched first sells    = [0.65586909 0.87695257 0.96291202]
```

--------------------------------

### Run Extended Kalman Filter (EKF)

Source: https://context7.com/givani30/contrax/llms.txt

Executes an EKF for nonlinear dynamics and observations. Requires defining nonlinear functions for dynamics and observations.

```python
import jax.numpy as jnp
import contrax as cx

# Nonlinear discrete dynamics
def f(x, u):
    return jnp.array([x[0] + 0.1 * x[1], x[1] + 0.1 * u[0]])

# Nonlinear observation (e.g., range measurement)
def h(x):
    return jnp.array([jnp.sqrt(x[0]**2 + x[1]**2)])

Q_noise = 1e-3 * jnp.eye(2)
R_noise = 1e-2 * jnp.eye(1)

# Simulated measurements
ys = jnp.array([[1.0], [1.1], [1.15], [1.2], [1.3]])
us = jnp.zeros((5, 1))

x0 = jnp.array([1.0, 0.0])
P0 = jnp.eye(2)

result = cx.ekf(
    f, Q_noise, R_noise, ys, us, x0, P0,
    observation=h,  # Required when f is a plain function
)

x_hat = result.x_hat  # Shape (5, 2)
```

--------------------------------

### Create Nonlinear System Model (nonlinear_system)

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a reusable nonlinear system model by defining `dynamics` and `output` functions. Specify `dt=None` for continuous-time systems or a float for discrete-time.

```python
import jax.numpy as jnp
import contrax as cx

# Pendulum dynamics: theta'' = -g/L * sin(theta) + u/mL^2
def pendulum_dynamics(t, x, u):
    theta, omega = x[0], x[1]
    g, L = 9.81, 1.0
    theta_dot = omega
    omega_dot = -(g / L) * jnp.sin(theta) + u[0]
    return jnp.array([theta_dot, omega_dot])

def pendulum_output(t, x, u):
    return x[:1]  # Only measure angle

sys_nonlinear = cx.nonlinear_system(
    dynamics=pendulum_dynamics,
    output=pendulum_output,
    dt=None,  # Continuous (set dt for discrete)
    state_dim=2,
    input_dim=1,
    output_dim=1,
)

# Use in simulation or estimation
x0 = jnp.array([0.1, 0.0])  # Initial state
u = jnp.array([0.0])        # Zero input
x_dot = sys_nonlinear.dynamics(0.0, x0, u)
y = sys_nonlinear.output(0.0, x0, u)
```

--------------------------------

### System Dynamics and Discretization

Source: https://context7.com/givani30/contrax/llms.txt

Tools for rolling out arbitrary dynamics and converting continuous systems to discrete-time models.

```APIDOC
## cx.rollout

### Description
Rolls out arbitrary discrete-time dynamics over a fixed input sequence.

### Parameters
- **dynamics** (Callable) - Required - Function defining state transition.
- **x0** (Array) - Required - Initial state.
- **us** (Array) - Required - Input sequence.
- **params** (Any) - Optional - Additional parameters for the dynamics function.

## cx.sample_system

### Description
Samples a continuous nonlinear system into a discrete transition model.

### Parameters
- **sys** (System) - Required - Continuous system model.
- **dt** (float) - Required - Timestep for discretization.
- **input_interpolation** (str) - Optional - Interpolation method (e.g., 'zoh').
```

--------------------------------

### Create Discrete-Time State-Space System (dss)

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a discrete-time LTI system using the `dss` function, specifying the state transition matrix, input matrix, output matrix, feedthrough matrix, and sampling time `dt`.

```python
import jax.numpy as jnp
import contrax as cx

# Discrete double integrator with 50ms sampling
A = jnp.array([[1.0, 0.05], [0.0, 1.0]])  # State transition
B = jnp.array([[0.0], [0.05]])             # Input matrix
C = jnp.eye(2)                             # Full state output
D = jnp.zeros((2, 1))                      # No feedthrough
dt = 0.05                                  # 50ms sample period

sys_discrete = cx.dss(A, B, C, D, dt=dt)
# Access: sys_discrete.A, sys_discrete.B, sys_discrete.C, sys_discrete.D, sys_discrete.dt
```

--------------------------------

### Rollout Discrete Trajectories

Source: https://context7.com/givani30/contrax/llms.txt

Executes arbitrary discrete-time dynamics over a fixed input sequence.

```python
import jax.numpy as jnp
import contrax as cx

# Custom nonlinear dynamics with parameters
def dynamics(x, u, gain):
    return gain * x + u

x0 = jnp.array([1.0])
us = jnp.zeros((20, 1))  # Zero input
gain = 0.9               # Decay parameter

xs = cx.rollout(dynamics, x0, us, params=gain)
# xs.shape = (21, 1) - includes initial state
```

--------------------------------

### Run Filtering, Smoothing, and Diagnostics

Source: https://github.com/givani30/contrax/blob/master/docs/examples/continuous-nonlinear-estimation.md

Simulates the continuous pendulum, generates noisy measurements, builds the discrete estimator model, and runs the Unscented Kalman Filter (UKF) and Smoother (UKS). It also includes diagnostics to assess estimation quality.

```python
from numpy import linspace, cos, pi, random

from contrax.estimation import ukf, uks, ukf_diagnostics


def run_example():
    system, estimator_model = setup()

    # simulation parameters
    T = 10.0  # simulation time
    dt = system.dt
    t = linspace(0, T, int(T / dt))
    random.seed(0)  # for reproducible noise

    # torque profile
    def u_profile(t):
        return 0.5 * cos(t)

    # simulate system and measurements
    x_true, u_true, y = system.simulate(t, u_profile, noise_std=0.1)

    # run UKF
    x_filtered, P_filtered, u_filtered, intermediates = ukf(estimator_model, y)

    # run UKS
    x_smoothed, P_smoothed, u_smoothed = uks(estimator_model, y, intermediates)

    # diagnostics
    nis, cond, loglik = ukf_diagnostics(estimator_model, y, intermediates)

    # print results
    print("Continuous nonlinear estimation")
    print(f"filtered theta rmse     = {sqrt(mean((x_filtered[:, 0] - x_true[:, 0]) ** 2)):.5f}")
    print(f"smoothed theta rmse     = {sqrt(mean((x_smoothed[:, 0] - x_true[:, 0]) ** 2)):.5f}")
    print(f"filtered rate rmse      = {sqrt(mean((x_filtered[:, 1] - x_true[:, 1]) ** 2)):.5f}")
    print(f"smoothed rate rmse      = {sqrt(mean((x_smoothed[:, 1] - x_true[:, 1]) ** 2)):.5f}")
    print(f"mean NIS                = {mean(nis):.5f}")
    print(f"max innovation cond     = {max(cond):.5f}")
    print(f"total log likelihood    = {sum(loglik):.5f}")

    return x_true, u_true, y, x_filtered, x_smoothed


if __name__ == "__main__":
    from numpy import sqrt, mean
    run_example()

```

--------------------------------

### Cast Quadratic Optimal Execution as LQR

Source: https://github.com/givani30/contrax/blob/master/docs/examples/jax-native-workflows.md

This snippet demonstrates how to model a simple execution problem as an LQR problem. It sets up a discrete-time system and defines risk and trading cost matrices to compute the LQR feedback gain. The entire process remains within JAX for differentiability.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0]])
B = jnp.array([[1.0]])
sys = cx.dss(A, B, jnp.array([[1.0]]), jnp.zeros((1, 1)), dt=1.0)

Q = jnp.array([[2.5]])   # inventory risk
R = jnp.array([[0.4]])   # temporary impact / trading cost
K = cx.lqr(sys, Q, R).K
```

--------------------------------

### cx.ss - Create Continuous-Time State-Space System

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a continuous-time LTI system with state-space representation: x_dot = Ax + Bu, y = Cx + Du.

```APIDOC
## cx.ss

### Description
Constructs a continuous-time LTI system with state-space representation: x_dot = Ax + Bu, y = Cx + Du.

### Parameters
- **A** (jnp.array) - Required - State matrix
- **B** (jnp.array) - Required - Input matrix
- **C** (jnp.array) - Required - Output matrix
- **D** (jnp.array) - Required - Feedthrough matrix
```

--------------------------------

### Turn Nonlinear Dynamics Into State-Space Models

Source: https://github.com/givani30/contrax/blob/master/docs/examples/jax-native-workflows.md

Linearize nonlinear plant functions around specific operating points to generate local state-space models.

```python
import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import contrax as cx


def pendulum(x, u):
    theta, theta_dot = x
    torque = u[0]
    return jnp.array([theta_dot, -jnp.sin(theta) + torque])


def sensor(x, u):
    return x


x_eq = jnp.array([0.1, 0.0])
u_eq = jnp.array([jnp.sin(0.1)])

sys_c = cx.linearize_ss(pendulum, x_eq, u_eq, output=sensor)
sys_d = cx.c2d(sys_c, dt=0.05)
```

--------------------------------

### Add Numerical Algorithm Reference

Source: https://github.com/givani30/contrax/blob/master/CONTRIBUTING.md

Include a reference comment for any non-trivial numerical algorithm added to the project.

```python
# Reference: Author (year), "Title", Journal.
```

--------------------------------

### cx.dss - Create Discrete-Time State-Space System

Source: https://context7.com/givani30/contrax/llms.txt

Constructs a discrete-time LTI system with state-space representation: x[k+1] = Ax[k] + Bu[k], y[k] = Cx[k] + Du[k].

```APIDOC
## cx.dss

### Description
Constructs a discrete-time LTI system with state-space representation: x[k+1] = Ax[k] + Bu[k], y[k] = Cx[k] + Du[k].

### Parameters
- **A** (jnp.array) - Required - State transition matrix
- **B** (jnp.array) - Required - Input matrix
- **C** (jnp.array) - Required - Output matrix
- **D** (jnp.array) - Required - Feedthrough matrix
- **dt** (float) - Required - Sample period
```

--------------------------------

### Define the Estimation Problem

Source: https://github.com/givani30/contrax/blob/master/docs/tutorials/kalman-filtering.md

Sets up a constant-velocity discrete system with position-only measurements.

```python
--8<-- "examples/kalman_filtering.py:setup"
```