Jaxonomy — hybrid dynamical systems on JAX

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?](#-why-jax) · [Install](#-installation) · [Quick Start](#-quick-start) · [Gallery](#-gallery) · [Examples](#-examples) · [Docs](https://py.jaxonomy.com)

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

¹ Simulink + Simscape + Stateflow + Control System Toolbox · ² Via Simulink Design Optimization, no end-to-end AD

---

⚡ Key Capabilities

---

📦 Installation

Requires Python 3.10+.

# Create and activate a virtual environment (recommended)

python -m venv .venv source .venv/bin/activate # Windows: .venv\Scripts\activate

Install

pip install jaxonomy # core pip install jaxonomy[safe] # + SciPy, Matplotlib, control, jaxopt pip install jaxonomy[nmpc] # + nonlinear MPC (requires IPOPT on PATH) pip install jaxonomy[all] # + everything

From source:

git clone https://github.com/machinavitalis/jaxonomy

cd jaxonomy pip install -e .

CLI runner:

jaxonomy_cli run --model path/to/model.json

---

🚀 Quick Start

A first simulation in a few lines — a custom block, built into a diagram, integrated through its ODE:

import jaxonomy as jx

import jax.numpy as jnp

Van der Pol oscillator as a custom block

class VanDerPol(jx.LeafSystem): def init(self, mu=1.0, kwargs): super().__init__(kwargs) self.declare_dynamic_parameter("mu", mu) self.declare_continuous_state( default_value=jnp.array([0.0, 2.0]), ode=self._ode ) self.declare_continuous_state_output(name="x")

def _ode(self, time, state, *inputs, **params): x, mu = state.continuous_state, params["mu"] return jnp.array([x[1], mu * (1 - x[0]**2) * x[1] - x[0]])

builder = jx.DiagramBuilder() vdp = builder.add(VanDerPol(mu=2.0, name="vdp")) diagram = builder.build()

results = jx.simulate( diagram, diagram.create_context(), (0.0, 20.0), options=jx.SimulatorOptions(buffer_length=4000), # room for adaptive steps recorded_signals={"x": vdp.output_ports[0]}, )

results.outputs["x"] → time-series of shape (T, 2)

---

📚 Documentation

• Online docs & tutorials: py.jaxonomy.com
• Local docs:

  pip install -r requirements.docs.txt

mkdocs serve # → http://127.0.0.1:8000

- **Example notebooks:** [`docs/examples/`](https://github.com/machinavitalis/jaxonomy/tree/main/docs/examples)

• Scope notes: PINNs & PDE surrogates — classical PDE PINNs are out of scope; physics-informed *dynamics* learning (UDE / Neural DAE / Neural ODE / SINDy) is core.

---

🤖 Driving Jaxonomy from an AI agent (MCP)

Jaxonomy ships an MCP server that exposes the engine as tools an AI agent can call directly — it can enumerate library blocks, build and validate a model, run a simulation, fit parameters to data, and linearize a system, then reason over the actual results. This is worth wiring up if you drive Jaxonomy from an agent (Claude Desktop/Code, Cursor, …); if you're writing Python by hand, the pip install above is all you need and you can skip this.

pip install jaxonomy[mcp]

Then register the server with your agent client. For Claude Desktop, add to claude_desktop_config.json:

{

"mcpServers": { "jaxonomy": { "command": "python", "args": ["-m", "jaxonomy.mcp.server"] } } }

Use the interpreter where jaxonomy[mcp] is installed (or the jaxonomy-mcp entry point). The server is listed in the MCP Registry as io.github.machinavitalis/jaxonomy; full tool reference, client configuration and limitations are at py.jaxonomy.com/mcp.

If you'd rather just point an agent at the documentation, start it on SKILL.md — when to use Jaxonomy, when to reach for something else, the core API, and the pitfalls that most often break a first script. The docs site also publishes llms.txt views: py.jaxonomy.com/llms.txt for an index of every page, and py.jaxonomy.com/llms-full.txt for the whole documentation set as one Markdown file.

---

📖 Examples

1 · Hybrid Dynamics: The Bouncing Ball

Jaxonomy is designed for hybrid systems — models where continuous physics interacts with instantaneous discrete resets. The bouncing ball is the canonical example: free-fall ODE interrupted by a collision event that reverses velocity.

Governing equations

The dynamics between bounces follow:

$$ \dot{x} = v, \qquad \dot{v} = -g $$

When the ball hits the ground ($x = 0$, $v < 0$) a zero-crossing event fires and the state resets:

$$ x^+ = 0, \qquad v^+ = -e \cdot v \quad (e \in [0, 1] \text{ — coefficient of restitution}) $$

import jaxonomy as jx

import jax.numpy as jnp

class BouncingBall(jx.LeafSystem): def init(self, g=9.81, e=0.8, *