Jaxonomy — hybrid dynamical systems on JAX MCP server

Jaxonomy

Differentiable simulation of hybrid dynamical systems — powered by JAX.

Block diagrams meet automatic differentiation. Model physical systems, close the loop with LQR/MPC/Kalman, and differentiate through everything.

Why JAX? · Install · Quick Start · Gallery · Examples · Docs

Every panel above is produced by jaxonomy.simulate on a model built from the library.

What is Jaxonomy?

Jaxonomy is a Python framework for modeling, simulating, and optimizing hybrid dynamical systems — systems that combine continuous physics, discrete control laws, and event-driven logic in a single model. Every simulation runs on JAX, so it is JIT-compilable, batchable with vmap, and fully differentiable from end to end.

import jax.numpy as jnp
import jaxonomy as jx

# Double integrator: A, B, C, D define the plant; Q, R weight the LQR cost.
A, B = jnp.array([[0., 1.], [0., 0.]]), jnp.array([[0.], [1.]])
C, D = jnp.eye(2), jnp.zeros((2, 1))
Q, R = jnp.eye(2), jnp.array([[1.]])

builder = jx.DiagramBuilder()
# Start displaced from the origin, so the regulator has something to do.
plant      = builder.add(jx.library.LTISystem(A, B, C, D,
                                              initialize_states=jnp.array([1.0, 0.0])))
controller = builder.add(jx.library.LinearQuadraticRegulator(A, B, Q, R))
builder.connect(plant.output_ports[0],      controller.input_ports[0])
builder.connect(controller.output_ports[0], plant.input_ports[0])

diagram = builder.build()
results = jx.simulate(
    diagram,
    diagram.create_context(),
    (0.0, 10.0),
    # Nothing is stored unless you name it here: without recorded_signals,
    # results.time and results.outputs come back as None.
    recorded_signals={"x": plant.output_ports[0],
                      "u": controller.output_ports[0]},
)

results.outputs["x"]    # (T, 2) — position and velocity, driven back to zero

🔥 Why JAX?

Choosing JAX as the compute backbone unlocks capabilities that are impractical with NumPy-based simulators:

Traditional simulator          Jaxonomy / JAX
──────────────────────         ─────────────────────────────────────────
simulate(params)          →    jit(simulate)(params)          10–100× faster
for p in param_grid: …    →    vmap(simulate)(param_grid)     embarrassingly parallel
finite_diff_gradient(…)   →    grad(simulate)(params)         exact gradients, free

Shortened. The full README is on GitHub.