{
  "markdown": "<div align=\"center\">\n\n<img src=\"https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/hero_triptych.gif\" width=\"100%\" alt=\"Jaxonomy — compose, simulate, control\">\n\n# Jaxonomy\n\n**Differentiable simulation of hybrid dynamical systems — powered by JAX.**\n\n[![PyPI](https://img.shields.io/pypi/v/jaxonomy?color=blue&label=PyPI)](https://pypi.org/project/jaxonomy/)\n[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)](https://www.python.org/)\n[![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/machinavitalis/jaxonomy/blob/main/LICENSE.md)\n[![Docs](https://img.shields.io/badge/docs-py.jaxonomy.com-blue)](https://py.jaxonomy.com)\n\n*Block diagrams meet automatic differentiation. Model physical systems, close the loop with LQR/MPC/Kalman, and differentiate through everything.*\n\n**[Why JAX?](#-why-jax)** · **[Install](#-installation)** · **[Quick Start](#-quick-start)** · **[Gallery](#-gallery)** · **[Examples](#-examples)** · **[Docs](https://py.jaxonomy.com)**\n\n<sub>Every panel above is produced by <code>jaxonomy.simulate</code> on a model built from the library.</sub>\n\n</div>\n\n---\n\n## What is Jaxonomy?\n\nJaxonomy 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**.\n\n```python\nimport jax.numpy as jnp\nimport jaxonomy as jx\n\n# Double integrator: A, B, C, D define the plant; Q, R weight the LQR cost.\nA, B = jnp.array([[0., 1.], [0., 0.]]), jnp.array([[0.], [1.]])\nC, D = jnp.eye(2), jnp.zeros((2, 1))\nQ, R = jnp.eye(2), jnp.array([[1.]])\n\nbuilder = jx.DiagramBuilder()\n# Start displaced from the origin, so the regulator has something to do.\nplant      = builder.add(jx.library.LTISystem(A, B, C, D,\n                                              initialize_states=jnp.array([1.0, 0.0])))\ncontroller = builder.add(jx.library.LinearQuadraticRegulator(A, B, Q, R))\nbuilder.connect(plant.output_ports[0],      controller.input_ports[0])\nbuilder.connect(controller.output_ports[0], plant.input_ports[0])\n\ndiagram = builder.build()\nresults = jx.simulate(\n    diagram,\n    diagram.create_context(),\n    (0.0, 10.0),\n    # Nothing is stored unless you name it here: without recorded_signals,\n    # results.time and results.outputs come back as None.\n    recorded_signals={\"x\": plant.output_ports[0],\n                      \"u\": controller.output_ports[0]},\n)\n\nresults.outputs[\"x\"]    # (T, 2) — position and velocity, driven back to zero\n```\n\n---\n\n## 🔥 Why JAX?\n\nChoosing JAX as the compute backbone unlocks capabilities that are impractical with NumPy-based simulators:\n\n```\nTraditional simulator          Jaxonomy / JAX\n──────────────────────         ─────────────────────────────────────────\nsimulate(params)          →    jit(simulate)(params)          10–100× faster\nfor p in param_grid: …    →    vmap(simulate)(param_grid)     embarrassingly parallel\nfinite_diff_gradient(…)   →    grad(simulate)(params)         exact gradients, free\n```\n\n| Feature | SciPy / NumPy | Julia / DiffEq | Modelica | MathWorks¹ | **Jaxonomy** |\n|---|:---:|:---:|:---:|:---:|:---:|\n| Python-native | ✓ | ✗ | ✗ | ✗ | **✓** |\n| JIT / code generation | ✗ | ✓ | ✓ (C++) | ✓ (C/C++) | **✓** |\n| Full autodiff through ODE | ✗ | Partial | ✗ | Partial² | **✓** |\n| Hybrid events & zero-crossing | Partial | ✓ | ✓ | ✓ | **✓** |\n| Acausal / equation-based | ✗ | ✗ | ✓ | ✓ (Simscape) | **✓** |\n| Block-diagram composition | ✗ | Partial | Partial | ✓ (Simulink) | **✓** |\n| State-machine modeling | ✗ | ✗ | ✗ | ✓ (Stateflow) | **✓** |\n| LQR / MPC / Kalman built-in | ✗ | Partial | Via libs | ✓ (Toolboxes) | **✓** |\n| Neural ODE / SINDy | ✗ | ✓ | ✗ | ✗ | **✓** |\n| Reduced-order modeling (balred / POD-DEIM / DMD / Koopman) | ✗ | Partial | ✗ | ✓ (Toolboxes) | **✓** |\n| Batch / ensemble (vmap) | ✗ | ✗ | ✗ | ✗ | **✓** |\n| Open-source (MIT) | ✓ | ✓ | Partial | ✗ | **✓** |\n\n<sub>¹ Simulink + Simscape + Stateflow + Control System Toolbox &nbsp;·&nbsp; ² Via Simulink Design Optimization, no end-to-end AD</sub>\n\n---\n\n## ⚡ Key Capabilities\n\n| Capability | What it enables |\n|---|---|\n| ⚡ **JAX-native engine** | JIT-compile simulations, run ensembles with `vmap`, differentiate through ODE solvers including event handling |\n| 🔀 **Hybrid dynamics + state machines** | Continuous ODEs, periodic discrete updates, zero-crossing events, and `StateMachineBuilder`-authored finite state machines composed in one model. `jax.grad` flows through event times for hybrid trajectory optimisation. |\n| 🔌 **Acausal modeling** | Modelica-inspired multi-domain components (electrical, mechanical, thermal, fluid, battery) with Pantelides index reduction and a BDF mass-matrix DAE solver |\n| 🎯 **Control & estimation** | LQR (continuous, discrete, finite-horizon, LQG), linear MPC (native + OSQP), nonlinear MPC (shooting / transcription / Hermite-Simpson), Kalman / EKF / UKF / RLS / Luenberger, 2-DOF PID with classical tuning helpers |\n| 🧮 **Unit-aware wiring** | Optional `BusUnit` annotations on ports and signals; the diagram compiler catches dimensional mismatches at build time instead of as silent runtime bugs |\n| 🧠 **Data-driven modeling** | Neural ODEs, Universal Differential Equations, SINDy symbolic regression, neural-network blocks (`MLP` / `PyTorch` / `TensorFlow` / `ONNX`), differentiable lookup-table fitting, and statistical surrogates (Gaussian process, polynomial chaos, RBF) |\n| 📉 **Reduced-order modeling** | `jaxonomy.library.rom`: linear MOR (balanced truncation, `minreal`, modal / residualization), POD–Galerkin with DEIM hyper-reduction, and data-driven operator ROM (DMD / DMDc / ERA, Koopman / eDMD lifted-linear predictors). One `reduce(...)` front door; every reduced model is a differentiable, simulatable block |\n| 🎲 **Uncertainty & sensitivity** | First-class `jaxonomy.uq` workflow: Monte Carlo with parameter distributions, Latin Hypercube + quasi-Monte Carlo sampling, Sobol sensitivity decomposition, Morris screening |\n| 🤝 **FMI 2.0 / 3.0 interop** | Import FMI co-simulation FMUs (`ModelicaFMU`, mixed-type and array I/O) or model-exchange FMUs (`ModelicaFMUME`, integrated by Jaxonomy's own solver with FMI event indicators as zero-crossings); export a diagram as a binary `.fmu` via `build_fmu`. Verified against the Reference FMUs, OpenModelica, and the `fmusim` reference simulator in CI |\n| 🧩 **150+ library blocks** | Integrators, filters, state machines, look-up tables, coordinate transforms, container blocks, bus / mux family, stochastic sources, and more |\n\n---\n\n## 📦 Installation\n\nRequires **Python 3.10+**.\n\n```bash\n# Create and activate a virtual environment (recommended)\npython -m venv .venv\nsource .venv/bin/activate        # Windows: .venv\\Scripts\\activate\n\n# Install\npip install jaxonomy             # core\npip install jaxonomy[safe]       # + SciPy, Matplotlib, control, jaxopt\npip install jaxonomy[nmpc]       # + nonlinear MPC (requires IPOPT on PATH)\npip install jaxonomy[all]        # + everything\n```\n\n**From source:**\n\n```bash\ngit clone https://github.com/machinavitalis/jaxonomy\ncd jaxonomy\npip install -e .\n```\n\n**CLI runner:**\n\n```bash\njaxonomy_cli run --model path/to/model.json\n```\n\n---\n\n## 🚀 Quick Start\n\nA first simulation in a few lines — a custom block, built into a diagram, integrated through its ODE:\n\n```python\nimport jaxonomy as jx\nimport jax.numpy as jnp\n\n# Van der Pol oscillator as a custom block\nclass VanDerPol(jx.LeafSystem):\n    def __init__(self, mu=1.0, **kwargs):\n        super().__init__(**kwargs)\n        self.declare_dynamic_parameter(\"mu\", mu)\n        self.declare_continuous_state(\n            default_value=jnp.array([0.0, 2.0]), ode=self._ode\n        )\n        self.declare_continuous_state_output(name=\"x\")\n\n    def _ode(self, time, state, *inputs, **params):\n        x, mu = state.continuous_state, params[\"mu\"]\n        return jnp.array([x[1],  mu * (1 - x[0]**2) * x[1] - x[0]])\n\nbuilder = jx.DiagramBuilder()\nvdp = builder.add(VanDerPol(mu=2.0, name=\"vdp\"))\ndiagram = builder.build()\n\nresults = jx.simulate(\n    diagram, diagram.create_context(), (0.0, 20.0),\n    options=jx.SimulatorOptions(buffer_length=4000),  # room for adaptive steps\n    recorded_signals={\"x\": vdp.output_ports[0]},\n)\n# results.outputs[\"x\"] → time-series of shape (T, 2)\n```\n\n---\n\n## 📚 Documentation\n\n- **Online docs & tutorials:** [py.jaxonomy.com](https://py.jaxonomy.com)\n- **Local docs:**\n  ```bash\n  pip install -r requirements.docs.txt\n  mkdocs serve   # → http://127.0.0.1:8000\n  ```\n- **Example notebooks:** [`docs/examples/`](https://github.com/machinavitalis/jaxonomy/tree/main/docs/examples)\n- **Scope notes:** [PINNs & PDE surrogates](https://github.com/machinavitalis/jaxonomy/blob/main/docs/scope/pinn.md) — classical\n  PDE PINNs are out of scope; physics-informed *dynamics* learning (UDE /\n  Neural DAE / Neural ODE / SINDy) is core.\n\n---\n\n## 🤖 Driving Jaxonomy from an AI agent (MCP)\n\n<!--\n  The line below is the MCP Registry's package-ownership proof. The registry\n  reads it from this README as published to PyPI (pyproject sets\n  readme = \"README.md\") and requires it to match the \"name\" in server.json\n  exactly. Do not remove or reword it, and keep it on its own line.\n-->\n<!-- mcp-name: io.github.machinavitalis/jaxonomy -->\n\nJaxonomy ships an [MCP](https://modelcontextprotocol.io) server that exposes the\nengine as tools an AI agent can call directly — it can enumerate library blocks,\nbuild and validate a model, run a simulation, fit parameters to data, and\nlinearize a system, then reason over the actual results. This is worth wiring up\nif you drive Jaxonomy from an agent (Claude Desktop/Code, Cursor, …); if you're\nwriting Python by hand, the `pip install` above is all you need and you can skip\nthis.\n\n```bash\npip install jaxonomy[mcp]\n```\n\nThen register the server with your agent client. For Claude Desktop, add to\n`claude_desktop_config.json`:\n\n```json\n{\n  \"mcpServers\": {\n    \"jaxonomy\": {\n      \"command\": \"python\",\n      \"args\": [\"-m\", \"jaxonomy.mcp.server\"]\n    }\n  }\n}\n```\n\nUse the interpreter where `jaxonomy[mcp]` is installed (or the `jaxonomy-mcp`\nentry point). The server is listed in the\n[MCP Registry](https://registry.modelcontextprotocol.io) as\n`io.github.machinavitalis/jaxonomy`; full tool reference, client configuration\nand limitations are at [py.jaxonomy.com/mcp](https://py.jaxonomy.com/mcp/).\n\nIf you'd rather just point an agent at the documentation, start it on\n[`SKILL.md`](SKILL.md) — when to use Jaxonomy, when to reach for something else,\nthe core API, and the pitfalls that most often break a first script. The docs\nsite also publishes [llms.txt](https://llmstxt.org/) views:\n[py.jaxonomy.com/llms.txt](https://py.jaxonomy.com/llms.txt) for an index of\nevery page, and\n[py.jaxonomy.com/llms-full.txt](https://py.jaxonomy.com/llms-full.txt) for the\nwhole documentation set as one Markdown file.\n\n---\n\n## 📖 Examples\n\n### 1 · Hybrid Dynamics: The Bouncing Ball\n\n<div align=\"center\">\n\n<img src=\"https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/grid_bouncing_ball.gif\" width=\"45%\" alt=\"Bouncing ball — hybrid contact events\">\n\n</div>\n\nJaxonomy 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.\n\n#### Governing equations\n\nThe dynamics between bounces follow:\n\n$$\n\\dot{x} = v, \\qquad \\dot{v} = -g\n$$\n\nWhen the ball hits the ground ($x = 0$, $v < 0$) a zero-crossing event fires and the state resets:\n\n$$\nx^+ = 0, \\qquad v^+ = -e \\cdot v \\quad (e \\in [0, 1] \\text{ — coefficient of restitution})\n$$\n\n```python\nimport jaxonomy as jx\nimport jax.numpy as jnp\n\nclass BouncingBall(jx.LeafSystem):\n    def __init__(self, g=9.81, e=0.8, **kwargs):\n        super().__init__(**kwargs)\n        self.declare_dynamic_parameter(\"g\", g)\n        self.declare_dynamic_parameter(\"e\", e)\n        # [height, velocity]\n        self.declare_continuous_state(\n            default_value=jnp.array([5.0, 0.0]), ode=self._ode\n        )\n        self.declare_continuous_state_output(name=\"state\")\n\n        # Zero-crossing: fires as height crosses zero from above\n        self.declare_zero_crossing(\n            guard=self._hit_ground,\n            reset_map=self._bounce,\n            direction=\"positive_then_non_positive\",\n        )\n\n    def _ode(self, time, state, *inputs, **params):\n        v = state.continuous_state[1]\n        return jnp.array([v, -params[\"g\"]])\n\n    def _hit_ground(self, time, state, *inputs, **params):\n        return state.continuous_state[0]  # guard on height\n\n    def _bounce(self, time, state, *inputs, **params):\n        x, v = state.continuous_state\n        new_state = jnp.array([0.0, -params[\"e\"] * v])\n        return state.with_continuous_state(new_state)\n\nbuilder = jx.DiagramBuilder()\nball = builder.add(BouncingBall(g=9.81, e=0.85, name=\"ball\"))\ndiagram  = builder.build()\nresults  = jx.simulate(\n    diagram, diagram.create_context(), (0.0, 8.0),\n    recorded_signals={\"state\": ball.output_ports[0]},\n)\n```\n\nZero-crossing events are located with **40-step bisection** on the solver's dense output polynomial — giving sub-microsecond temporal accuracy without user-specified tolerances.\n\n---\n\n### 2 · Optimal Control: LQR Pendulum\n\n<div align=\"center\">\n\n![LQR Pendulum Block Diagram](https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/lqr_pendulum_block_diagram.png)\n\n</div>\n\nFor a linearized pendulum with state $x = [\\theta, \\dot\\theta]^\\top$:\n\n$$\n\\dot{x} = Ax + Bu, \\qquad\nA = \\begin{bmatrix} 0 & 1 \\\\ g/L & 0 \\end{bmatrix}, \\quad\nB = \\begin{bmatrix} 0 \\\\ 1/mL^2 \\end{bmatrix}\n$$\n\nThe **Linear Quadratic Regulator** minimizes infinite-horizon cost:\n\n$$\nJ = \\int_0^\\infty \\bigl( x^\\top Q\\, x + u^\\top R\\, u \\bigr)\\, dt\n$$\n\nby solving the algebraic Riccati equation $A^\\top P + PA - PBR^{-1}B^\\top P + Q = 0$ for the optimal gain $K = R^{-1}B^\\top P$, so $u^* = -Kx$.\n\n```python\nimport jaxonomy as jx\nfrom jaxonomy.library import LTISystem, LinearQuadraticRegulator\nimport jax.numpy as jnp\n\ng, L, m = 9.81, 1.0, 1.0\nA = jnp.array([[0, 1], [g/L, 0]])\nB = jnp.array([[0], [1 / (m * L**2)]])\nC, D = jnp.eye(2), jnp.zeros((2, 1))\n\nQ = jnp.diag(jnp.array([10.0, 1.0]))  # penalise angle more than rate\nR = jnp.array([[0.1]])                 # control effort cost\n\nbuilder    = jx.DiagramBuilder()\nplant      = builder.add(LTISystem(A, B, C, D, name=\"pendulum\"))\ncontroller = builder.add(LinearQuadraticRegulator(A, B, Q, R, name=\"lqr\"))\n\nbuilder.connect(plant.output_ports[0],      controller.input_ports[0])\nbuilder.connect(controller.output_ports[0], plant.input_ports[0])\n\ndiagram = builder.build()\ncontext = diagram.create_context()\n\n# Perturb the pendulum's initial angle by 15° (set one block's sub-state)\ncontext = context.with_subcontext(\n    plant.system_id,\n    context[plant.system_id].with_continuous_state(jnp.array([jnp.pi / 12, 0.0])),\n)\n\nresults = jx.simulate(\n    diagram, context, (0.0, 5.0),\n    recorded_signals={\"x\": plant.output_ports[0]},\n)\n```\n\n---\n\n### 3 · Differentiable Parameter Identification\n\n<div align=\"center\">\n\n![Differentiable parameter identification](https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/battery_optimization.gif)\n\n</div>\n\nJaxonomy can **differentiate through complete simulations** to fit model parameters to data — no finite-difference approximations, no hand-written adjoint code. Here we recover a spring–damper's stiffness and damping by differentiating the whole rollout and descending with [Optax](https://optax.readthedocs.io):\n\n```python\nimport jax\nimport jax.numpy as jnp\nimport optax\nimport jaxonomy as jx\n\nclass SpringDamper(jx.LeafSystem):\n    def __init__(self, k=1.0, c=0.3, **kwargs):\n        super().__init__(**kwargs)\n        self.declare_dynamic_parameter(\"k\", k)\n        self.declare_dynamic_parameter(\"c\", c)\n        self.declare_continuous_state(\n            default_value=jnp.array([1.0, 0.0]), ode=self._ode\n        )\n\n    def _ode(self, time, state, *inputs, **params):\n        x, v = state.continuous_state\n        return jnp.array([v, -params[\"k\"] * x - params[\"c\"] * v])\n\nsd   = SpringDamper(name=\"sd\")\nopts = jx.SimulatorOptions(enable_autodiff=True, max_major_steps=200)\n\ndef final_state(theta):                       # theta = {\"k\": ..., \"c\": ...}\n    ctx = sd.create_context()\n    ctx.parameters[\"k\"], ctx.parameters[\"c\"] = theta[\"k\"], theta[\"c\"]\n    res = jx.simulate(sd, ctx, (0.0, 1.5), options=opts)\n    return res.context.continuous_state       # differentiable final [x, v]\n\ntarget = jax.lax.stop_gradient(final_state({\"k\": 4.0, \"c\": 0.5}))  # \"measured\"\nloss   = lambda theta: jnp.sum((final_state(theta) - target) ** 2)\n\ntheta   = {\"k\": 1.0, \"c\": 0.1}\nopt     = optax.adam(2e-1)\nstate   = opt.init(theta)\ngrad_fn = jax.jit(jax.grad(loss))             # gradient through the ODE solver\nfor _ in range(300):\n    updates, state = opt.update(grad_fn(theta), state)\n    theta = optax.apply_updates(theta, updates)\n# theta → {\"k\": 4.00, \"c\": 0.50}\n```\n\n`jax.grad(loss)` flows back through every ODE step automatically, so the same recipe scales to battery ECMs (above), powertrains, or any parametric model — and `jaxonomy.optimization` wraps it in a higher-level `Optimizable` API when you want bounds, transforms, and constraints.\n\n---\n\n### 4 · Acausal Physical Modeling\n\n<div align=\"center\">\n\n![Acausal RC circuit schematic](https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/rc_circuit_acausal.png)\n\n</div>\n\nJaxonomy includes a Modelica-inspired **acausal modeling layer** for multi-domain physical systems. You describe component connections symbolically; the compiler automatically derives the governing DAE, reduces its index, and emits a `LeafSystem` ready to drop into any diagram.\n\n**RC circuit with initial conditions:**\n\n$$\nC\\,\\dot{V}_C = I, \\qquad V_C(0) = 0\\,\\text{V}, \\qquad V_{\\text{src}} = 1\\,\\text{V}\n$$\n\n```python\nimport jaxonomy as jx\nfrom jaxonomy.acausal import AcausalCompiler, AcausalDiagram, EqnEnv\nfrom jaxonomy.acausal import electrical as elec\n\nev = EqnEnv()\nad = AcausalDiagram()\n\nvs  = elec.VoltageSource(ev, name=\"vs\", v=1.0)\nr   = elec.Resistor(ev,      name=\"r\",  R=1.0)\nc   = elec.Capacitor(ev,     name=\"c\",  C=1.0,\n                              initial_voltage=0.0, initial_voltage_fixed=True)\ngnd = elec.Ground(ev,        name=\"gnd\")\n\nad.connect(vs, \"p\", r,  \"n\")\nad.connect(r,  \"p\", c,  \"p\")\nad.connect(c,  \"n\", vs, \"n\")\nad.connect(vs, \"n\", gnd, \"p\")\n\ncompiler = AcausalCompiler(ev, ad)\nrc_block = compiler()         # → LeafSystem, JIT-compiled ODE\n\nbuilder = jx.DiagramBuilder()\nrc = builder.add(rc_block)\ndiagram = builder.build()\nresults = jx.simulate(\n    diagram,\n    diagram.create_context(),\n    (0.0, 5.0),\n    recorded_signals={\"rc\": rc.output_ports[0]},\n)\n\nresults.outputs[\"rc\"]   # capacitor charges 0 → 0.993 V over 5 τ (R = C = 1)\n```\n\nThe same pipeline handles **multi-domain systems** — an electro-mechanical actuator connecting `electrical`, `rotational`, and `translational` domains compiles to a single optimized system.\n\nAvailable acausal domains: `electrical` · `rotational` · `translational` · `thermal` · `fluid`\n\n---\n\n### 5 · Differentiable Sensitivity & Batch Simulation\n\n<div align=\"center\">\n\n![Differentiable sensitivity and batch simulation of a spring–mass system](https://raw.githubusercontent.com/machinavitalis/jaxonomy/main/docs/examples/media/sensitivity_batch.png)\n\n</div>\n\nBecause `jax.grad`, `jax.jit`, and `jax.vmap` all compose with `simulate`, you get powerful workflows with minimal boilerplate:\n\n```python\nimport jax\nimport jax.numpy as jnp\nimport jaxonomy as jx\n\nclass SpringMass(jx.LeafSystem):\n    def __init__(self, mass=1.0, damping=0.1, stiffness=10.0, **kwargs):\n        super().__init__(**kwargs)\n        self.declare_dynamic_parameter(\"mass\", mass)\n        self.declare_dynamic_parameter(\"damping\", damping)\n        self.declare_dynamic_parameter(\"stiffness\", stiffness)\n        self.declare_continuous_state(\n            default_value=jnp.array([1.0, 0.0]), ode=self._ode\n        )\n        self.declare_continuous_state_output(name=\"x\")\n\n    def _ode(self, time, state, *inputs, **params):\n        x, v = state.continuous_state\n        a = -(params[\"stiffness\"] * x + params[\"damping\"] * v) / params[\"mass\"]\n        return jnp.array([v, a])\n\n# ── Sensitivity: ∂(final position)/∂(all parameters) in one reverse pass ─────\nplant = SpringMass(name=\"plant\")\ngrad_opts = jx.SimulatorOptions(enable_autodiff=True, max_major_steps=200)\n\ndef final_position(theta):\n    ctx = plant.create_context()\n    for name, value in theta.items():\n        ctx.parameters[name] = value\n    res = jx.simulate(plant, ctx, (0.0, 5.0), options=grad_opts)\n    return res.context.continuous_state[0]\n\ngrads = jax.grad(final_position)({\"mass\": 1.0, \"damping\": 0.1, \"stiffness\": 10.0})\n\n# ── Monte Carlo ensemble: 1000 trajectories over a mass sweep (vmap) ─────────\nbuilder = jx.DiagramBuilder()\nplant_b = builder.add(SpringMass(name=\"plant\"))\ndiagram = builder.build()\n\nresults = jx.simulate_batch(\n    diagram, t_span=(0.0, 5.0),\n    param_batches={\"plant.mass\": jnp.linspace(0.5, 2.0, 1000)},\n    options=jx.SimulatorOptions(math_backend=\"jax\", max_major_steps=200),\n    recorded_signals={\"x\": plant_b.output_ports[0]},\n)\n# results.outputs[\"x\"] shape: (1000, T, 2)\n```\n\n---\n\n## 🧩 Library Overview\n\nOver **150 built-in blocks** covering the full signal-processing and control toolkit. A representative slice — not exhaustive:\n\n| Category | Blocks |\n|---|---|\n| **Sources** | `Constant`, `Step`, `Ramp`, `Chirp`, `Pulse`, `Sawtooth`, `Clock`, `DataSource` |\n| **Stochastic** | `RandomNumber`, `UniformRandomNumber`, `WhiteNoise`, `BandLimitedNoise`, `PRBS`, `RandomSource` — all support `with_key` for independent noise streams under `vmap` |\n| **Arithmetic & nonlinearities** | `Adder`, `Gain`, `Product`, `Abs`, `Power`, `Trigonometric`, `Saturate` / `SoftSaturate`, `DeadZone`, `RateLimiter` / `SoftRateLimiter`, `Quantizer`, `Backlash` |\n| **Continuous dynamics** | `Integrator`, `Derivative`, `TransferFunction`, `LTISystem`, `PID` / `PIDContinuous` |\n| **Discrete dynamics** | `IntegratorDiscrete`, `UnitDelay`, `FilterDiscrete`, `LowPassDiscrete`, `LeadLag`, `Notch`, `PIDDiscrete`, `PIDController2DOF`, `Decimator`, `RateTransition` |\n| **Routing, buses, matrix** | `Mux` / `Demux`, `BusCreator` / `BusSelector` / `BusUpdate` (with `BusUnit` annotations), `Slice`, `Stack`, `Switch` / `MultiPortSwitch`, `IfThenElse`, `MatrixMultiplication`, `MatrixInversion`, `DotProduct`, `CrossProduct` |\n| **Logic & state machines** | `TruthTable`, `LogicalOperator`, `Comparator`, `Relay`, `EdgeDetection`, `StateMachine` (authored via `StateMachineBuilder` DSL) |\n| **Lookup tables** | `LookupTable1d` / `LookupTable2d` / `LookupTableND`, `Prelookup` / `InterpolationUsingPrelookup`, `TableSearch` — all differentiable and fittable from data via `fit_lookup_table_*` |\n| **Delays & containers** | `TransportDelay`, `VariableTransportDelay`, `EnabledSubsystem`, `TriggeredSubsystem`, `ForEach`, `Conditional` |\n| **Control** | `LinearQuadraticRegulator` / `DiscreteTimeLinearQuadraticRegulator` / `FiniteHorizonLinearQuadraticRegulator` / `LinearQuadraticGaussian`, `LinearDiscreteTimeMPC` (native + `LinearDiscreteTimeMPC_OSQP`), `DirectShootingNMPC` / `DirectTranscriptionNMPC` / `HermiteSimpsonNMPC` |\n| **Estimation** | `KalmanFilter`, `ExtendedKalmanFilter`, `UnscentedKalmanFilter`, `InfiniteHorizonKalmanFilter`, `RecursiveLeastSquares`, `AugmentedStateEKF`, `Luenberger` |\n| **Physics & coordinates** | `CoordinateRotation`, `RigidBody`, `BatteryCell` |\n| **ML / Data** | `MLP` (Equinox), `Sindy`, `PyTorch`, `TensorFlow`, `ONNX` / `ONNXJax` |\n| **ROM & surrogates** | `reduce(...)` → `ReducedOrderModel`; linear MOR (`balred` / `minreal` / `modal_truncation` / `residualize`), `galerkin_reduce` + `deim` (POD–DEIM), `dmd` / `dmdc` / `era`, `DMDForecaster` / `KoopmanPredictor`, and surrogate blocks `GaussianProcess` / `PolynomialChaos` / `RadialBasisSurrogate` |\n| **Interop** | `ModelicaFMU` / `ModelicaFMUME` (FMI 2.0 / 3.0 co-simulation and model-exchange import) + FMU export via `build_fmu`; `MuJoCo` / `MJX`; `Ros2Publisher` / `Ros2Subscriber`; `QuanserHAL`; `PyTwin` |\n| **Custom** | `CustomPythonBlock`, `CustomJaxBlock` for user-authored algorithms with persistent per-instance state |\n\nCustom blocks are first-class beyond the wrappers above: subclass `LeafSystem`, declare ports and states, and your block integrates with the full framework including JIT, autodiff, and event detection.\n\n---\n\n## 🏗️ System Architecture\n\nA Jaxonomy model is a **Diagram** — a directed graph of interconnected blocks. Blocks (`LeafSystem`) declare ports, continuous/discrete states, and event guards. The simulator orchestrates ODE integration, discrete updates, and event detection automatically.\n\n```mermaid\ngraph LR\n    subgraph Diagram[\"DiagramBuilder.build()\"]\n        direction LR\n        r[/\"r(t)\\nReference\"/] --> sum((\"Σ\"))\n        sum -->|\"e(t)\"| ctrl[\"Controller\\nLQR · MPC · PID\"]\n        ctrl -->|\"u(t)\"| plant[\"Plant\\nLTI · Nonlinear · Acausal\"]\n        plant -->|\"x(t)\"| obs[\"State Estimator\\nKalman · EKF · UKF\"]\n        obs --> sum\n        plant -->|\"y(t)\"| out[/\"y(t)\\nOutput\"/]\n    end\n\n    style Diagram fill:#f5f8ff,stroke:#4a6fa5,stroke-width:2px\n    style ctrl fill:#dbeafe,stroke:#2563eb\n    style plant fill:#dcfce7,stroke:#16a34a\n    style obs fill:#fef9c3,stroke:#ca8a04\n```\n\nThe JAX backend sits underneath every simulation, giving you a clean Python API backed by XLA compilation:\n\n```mermaid\ngraph TD\n    A[\"Python API\\nDiagram · LeafSystem · simulate()\"] --> B[\"JAX Backend\\nJIT · vmap · grad\"]\n    B --> C1[\"Dopri5\\nAdaptive RK45\\nnon-stiff ODEs\"]\n    B --> C2[\"BDF\\nImplicit solver\\nstiff systems\"]\n    B --> C3[\"Event Handler\\nZero-crossing\\nbisection\"]\n\n    style A fill:#e0f2fe,stroke:#0284c7\n    style B fill:#fdf4ff,stroke:#9333ea\n    style C1 fill:#f0fdf4,stroke:#16a34a\n    style C2 fill:#f0fdf4,stroke:#16a34a\n    style C3 fill:#f0fdf4,stroke:#16a34a\n```\n\n---\n\n## The Jaxonomy stack\n\nJaxonomy is the engine at the base of a three-package stack. Each package is its\nown repository, MIT-licensed, and depends only on the package(s) above it:\n\n- **Jaxonomy** — *this package.* The general-purpose, JAX-native simulation\n  engine for hybrid dynamical systems. Not robotics-specific; depends on nothing\n  else in the stack.\n- **[Jaxterity](https://github.com/machinavitalis/jaxterity)** — the robotics\n  layer built on top of Jaxonomy: URDF/MJCF import, MJX-backed articulated\n  dynamics, calibrated actuators/sensors, system identification, and whole-body\n  control. Imports Jaxonomy; never re-implements its primitives.\n- **[Jaxility](https://github.com/machinavitalis/jaxility)** — the deployment\n  artifact factory: compiles a calibrated robot to embedded C for Arm SoCs\n  (Cortex-A / Cortex-M) with a signed attestation manifest. Consumes Jaxterity.\n\nThe boundary rule: anything useful to controls engineers *outside* robotics\n(HVAC, battery, aerospace, energy) belongs in Jaxonomy; anything specific to\njoints, actuators, contacts, and kinematic chains belongs in Jaxterity;\nembedded codegen, targets, and attestation belong in Jaxility.\n\nWhat Jaxonomy does **not** yet do (or does only partially) is tracked in\n[`KNOWN_GAPS.md`](https://github.com/machinavitalis/jaxonomy/blob/main/KNOWN_GAPS.md) — the public inverse of the internal evidence\nledger in `CLAIMS.md`.\n\n## License\n\nReleased under the [MIT License](https://github.com/machinavitalis/jaxonomy/blob/main/LICENSE.md) from version 2.2.0 onward.\n\nDerived from the MIT-licensed open-source package **pycollimator** by **Collimator, Inc.**\n</content>\n</invoke>\n",
  "bytes": 27982,
  "sha": "6353708b2b6d51cc75850fa5e649eb78e4deb2fd047a38b822788ed295576558",
  "repo_slug": "machinavitalis/jaxonomy",
  "fonte": "repo",
  "truncated": false,
  "api": "https://agentalog.com/api/listings/mcp_io_github_machinavitalis_jaxonomy_39550151/readme"
}