# jaxonomy > Jaxonomy — block-diagram simulation for hybrid dynamical systems (JAX, NumPy API, tutorials, examples, and API reference). Jaxonomy is a JAX-native Python engine for simulating hybrid dynamical systems built as block diagrams — wired blocks, continuous and discrete state, zero-crossing events, and acausal multi-physics networks (electrical, mechanical, thermal, hydraulic). Every simulation is JIT-compilable, vmap-batchable, and differentiable end to end, so calibration, trajectory optimization and controller tuning are ordinary gradient-based optimizations over the simulation itself. It ships 150+ library blocks including LQR, MPC, PID and Kalman/EKF/UKF, plus FMI 2.0 co-simulation and a reduced-order-modeling suite. MIT licensed. # Start here # Jaxonomy documentation **Jaxonomy** is a Python package for simulating **hybrid dynamical systems** described as **block diagrams**: wired blocks (integrators, gains, custom subsystems, acausal networks, and more), continuous and discrete states, and event/zero-crossing logic. The runtime is built around **JAX**, so you get JIT-friendly execution and **automatic differentiation** where the model allows it, while keeping a NumPy-style API for numerics. The library runs **entirely locally** and can serialise models to Collimator-format JSON. There is no hosted cloud service — see the [About](https://py.jaxonomy.com/about/index.md) page for the project's scope. The source lives at [github.com/machinavitalis/jaxonomy](https://github.com/machinavitalis/jaxonomy). ______________________________________________________________________ ## Install ``` pip install jaxonomy ``` Use a virtual environment when possible. Platform notes, optional extras (`[safe]`, `[nmpc]`, `[all]`), and **development installs from a git clone** are covered in the **[installation guide](https://py.jaxonomy.com/installation/index.md)**. ______________________________________________________________________ ## Where to go next | Goal | Link | | --------------------------------------- | -------------------------------------------------------------------------------------------- | | First simulation walkthrough | [Tutorials → Getting started](https://py.jaxonomy.com/tutorials/01-getting-started/index.md) | | Shorter topical guides | [Tutorials index](https://py.jaxonomy.com/tutorials/index.md) | | Applied notebooks (control, MPC, ML, …) | [Examples](https://py.jaxonomy.com/examples/index.md) | | `DiagramBuilder`, `LeafSystem`, ports | [Framework](https://py.jaxonomy.com/framework/index.md) | | Built-in blocks | [Block library](https://py.jaxonomy.com/library/index.md) | | `simulate`, solvers, options | [Simulation](https://py.jaxonomy.com/simulation/index.md) | | Training / optimization helpers | [Optimization](https://py.jaxonomy.com/optimization/index.md) | ______________________________________________________________________ ## Minimal pattern 1. Add blocks with `DiagramBuilder`, `connect` outputs to inputs, then `build()`. 1. Call `jaxonomy.simulate(diagram, context, t_span, ...)`, where `t_span` is a `(start, stop)` tuple (see [Simulation](https://py.jaxonomy.com/simulation/index.md) for `SimulatorOptions` and results handling). 1. Pass `recorded_signals={name: port}` for anything you want back — without it, `results.time` and `results.outputs` are `None`. The [Getting started](https://py.jaxonomy.com/tutorials/01-getting-started/index.md) tutorial builds a simple mass–spring–damper-style diagram step by step. ______________________________________________________________________ ## Using Jaxonomy from an AI agent If you are pointing an AI coding agent at Jaxonomy, start it on the **[agent guide](https://py.jaxonomy.com/agents/index.md)** — when to use the library, when to reach for something else, the core API, and the pitfalls that most often break a first script. Machine-readable views of this site, following the [llms.txt](https://llmstxt.org/) convention: | File | Contents | | --------------------------------------------------------- | ----------------------------------------------------------------- | | [`/llms.txt`](https://py.jaxonomy.com/llms.txt) | Index of the documentation, with a short description of each page | | [`/llms-full.txt`](https://py.jaxonomy.com/llms-full.txt) | The full documentation as a single Markdown file | Jaxonomy also ships an [MCP server](https://github.com/machinavitalis/jaxonomy/blob/main/jaxonomy/mcp/README.md) that exposes the engine as tools an agent can call directly. ______________________________________________________________________ ## Build this site locally ``` pip install -r requirements.docs.txt mkdocs serve ``` Source for this page: `docs/index.md`. ______________________________________________________________________ ## License and attribution This project is released under the [MIT License](https://mit-license.org/). See the `LICENSE.md` file in the [repository](https://github.com/machinavitalis/jaxonomy) for the full text. **Provenance:** This library is derived from the MIT-licensed open-source Python package **pycollimator**, developed by **Collimator, Inc.** # Jaxonomy Skill You are using Jaxonomy, a JAX-native engine for modeling and simulating hybrid dynamical systems by block-diagram composition. This file is your operating manual for *using* the library. Read it before suggesting any Jaxonomy code. (If you are *modifying* Jaxonomy itself, read `AGENTS.md` and the `AGENTS/` docs instead.) ## What Jaxonomy does Jaxonomy composes continuous physics, discrete control, and event-driven logic into a single model, and runs it on JAX — JIT-compilable, `vmap`-batchable, and **differentiable end to end**. Its thesis: modeling and simulation is most useful when it is optimization-ready, so every output, objective, and constraint is differentiable w.r.t. parameters, initial conditions, and network weights by default. It is the engine at the base of a larger stack (robotics and embedded deployment layer on top of it). For the full architecture, design philosophy, and invariants, read `AGENTS/CONTEXT.md` — don't restate it here. ## When to use Jaxonomy Use it when the user is: - Building a block-diagram model of a dynamical system (continuous, discrete, or hybrid) and simulating it. - Closing a control loop — LQR, MPC, PID, or a Kalman/EKF/UKF estimator around a plant. - Differentiating *through* a simulation: parameter calibration against data, trajectory optimization, controller tuning, neural ODE / SINDy, digital-twin updates. - Acausal / multi-physics modeling (electrical, thermal, fluid, mechanical) via `jaxonomy.acausal`. - Uncertainty quantification (Monte Carlo, Sobol, LHS, qMC) via `jaxonomy.uq`. - Batch/ensemble simulation with `vmap`, or GPU/TPU-accelerated runs. ## When NOT to use Jaxonomy - **Just integrating an ODE.** Use Diffrax directly. Jaxonomy adds block-diagram composition, hybrid dynamics, events, and state machines on top of a solver — if none of that is needed, it's overhead. - **Robotics with joints/actuators/contacts/kinematic chains.** Use Jaxterity, the robotics layer built on top of Jaxonomy. Don't re-implement URDF import, articulated dynamics, or WBC here. - **Embedded codegen / cross-compilation to silicon.** That's a downstream deployment concern, out of scope for Jaxonomy. - **A hosted/cloud simulation platform, web UI, or collaborative editor.** Not what this library is; see `AGENTS/CONTEXT.md` ("What Jaxonomy is NOT"). ## Core API surface, in order of how often agents will use it ### Build a diagram, create a context, simulate ``` import jaxonomy as jx builder = jx.DiagramBuilder() plant = builder.add(jx.library.LTISystem(A, B, C, D)) 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() context = diagram.create_context() # holds state + parameters results = jx.simulate( diagram, context, (0.0, 5.0), # t_span, a (start, stop) tuple recorded_signals={"x": plant.output_ports[0]}, ) print(results.time.shape, results.outputs["x"].shape) ``` - `DiagramBuilder.add(block)` returns a handle whose `.input_ports[i]` / `.output_ports[i]` you wire with `builder.connect(src_out, dst_in)`. - The `Context` is an immutable carrier of state and parameters — thread it through; don't mutate it in place. - `jx.simulate(system, context, t_span, ...)` is the entry point; the whole call is differentiable and `jit`/`vmap`-friendly. The time span is a `(start, stop)` tuple, passed positionally or as `t_span=` — there is no `stop_time`, `start_time`, or `end_time` keyword. - **Nothing is recorded unless you ask for it.** Without `recorded_signals=`, the run still succeeds but `results.time` and `results.outputs` are both `None`. This is the most common reason a first Jaxonomy script appears to produce nothing. ### Library blocks Standard blocks live under `jx.library`, split by category (sources, math_ops, logic, routing, dynamics, nonlinearities, tables); `jx.library.primitives` re-exports them for back-compat. Controls/estimation blocks (LQR, MPC, PID, Kalman/EKF/UKF) are library blocks too — prefer them over hand-rolling. ### Analysis and optimization - `linearize(...)` → a `LinearizedSystem`; analytical helpers (`bode_data`, `nyquist_data`, `step_response`, `frequency_response`, …) return plain dicts of arrays (no matplotlib inside Jaxonomy — you plot). - `influence_graph(system, context)` → an `InfluenceGraph`: the model's dependency structure with autodiff Jacobians on every edge. Use it to answer "what actually drives this signal" quantitatively rather than structurally — `.slice(target, threshold)` (vs `.structural_slice(target)` for the boolean over-approximation), `.attribute(target, source)` (signed per-path chain rule), `.dominant_paths`, `.dead_edges`, `.bottlenecks`. Read the module docstring before reading a weight: weights are dimensionless elasticities by default, edges into a continuous state are scaled by `tau` (which makes a path's product its gain at ω = 1/`tau`), and anything non-differentiable is labelled `local_gradient=False` rather than silently zeroed. `analysis.influence_subgraph(graph, focus, budget_tokens=…)` serializes a bounded, citable neighbourhood when you need to reason about one part of a model too large to read whole. - Because `simulate` is differentiable, parameter estimation, trajectory optimization, and controller tuning are just gradient-based optimizations over the simulation — use JAX autodiff / the provided tuning helpers. For exact signatures, see the docs at py.jaxonomy.com and the `examples/` notebooks — prefer those over guessing an API. ## Key gotchas - **Backend-neutrality: `npa` vs `jnp` vs `np`.** Inside models, numeric code goes through the `npa` abstraction, not `jnp` directly — it's a load-bearing invariant. If you write library-style code, follow the pattern in `AGENTS/PATTERNS.md`; as a *user* composing existing blocks you rarely touch it, but don't assume raw `jnp` everywhere. - **State is NamedTuple-shaped, not mutable objects.** Read state out of the results / context; don't try to assign into it. - **Differentiability is the default and the point.** If a construct would break gradients (Python-side branching on traced values, in-place mutation), reach for the block/pattern that preserves them. ## Common pitfalls & idioms Hard-won usage tips (harvested from prior consumer sessions): - **Parameter sweeps re-JIT per iteration.** `diagram.with_parameter("p", float(v))` in a Python loop keys the trace cache on the *value*, recompiling every step. Wrap the scalar as `jnp.asarray(v)`, or use `simulate_batch` / `jax.vmap` from the start. `scipy.optimize.minimize_scalar` inherits the same blowup. - **Update sub-system parameters by dot-path on the *outer* diagram:** `outer.with_parameters({"inner.gain.gain": jnp.array(5.0)})` — don't reach into `outer["inner"]` (you'll get a stale outer diagram). The dot-path form also composes under `vmap`. - **`declare_periodic_update` needs an explicit `offset=0.0`.** Omitting it constructs fine but crashes at the first step deep in the scheduler. - **Traced-mode guards:** gate host-side checks with `isinstance(x, jax.core.Tracer)`, *not* `jax.core.is_concrete` — the latter is `True` under `jax.grad`'s tracer and silently breaks gradients. - **`LookupTable1d.interp_1d` takes `method=`**, not `mode=` (the rest of the library uses `mode=` on `Quantizer`/`Saturate`, so the typo is natural). - **Symmetric saturation shorthand:** `Saturate(limit=L)` instead of `upper_limit=+L, lower_limit=-L`. - **`simulate` specifics:** `context` is a required positional arg (no context-less shortcut), the time span is one `(start, stop)` tuple rather than two scalars, and `options` must be a `SimulatorOptions`, not a dict. Results come back empty — `res.time is None` — unless you pass `recorded_signals={name: port}`; with it, read `res.time` and `res.outputs[name]`. - **Long / stiff / multi-rate runs:** bump `SimulatorOptions(buffer_length=...)` (the recorder silently truncates to the *tail* when its ring buffer fills) and use `ode_solver_method="bdf"` for stiff fast+slow coupling (the explicit default collapses its step size). - **Stateful feedback controller as a LeafSystem:** declare integrator state with `declare_continuous_state(..., ode=cb, requires_inputs=True)` and the command with `declare_output_port(..., requires_inputs=True, prerequisites_of_calc=[DependencyTicket.xc])`. Wiring controller↔plant is *not* an algebraic loop if the plant exposes state via `declare_continuous_state_output` (a state output, not feedthrough). ## When to escalate to the human - A model that won't `jit` or `vmap` cleanly and the cause isn't an obvious Python-control-flow-over-traced-values mistake. - Numerical divergence or stiffness that looks like a solver/tolerance choice rather than a modeling bug. - Any request that implies robotics-specific structure — redirect to Jaxterity rather than bolting kinematics into a Jaxonomy model. ## Where to find more - `AGENTS/CONTEXT.md` — architecture, design philosophy, what it is and is NOT. - `AGENTS/PATTERNS.md` — coding conventions (`npa`, `LeafSystem`, state). - `CLAIMS.md` / `KNOWN_GAPS.md` — what Jaxonomy actually does vs. what it does not (read both before relying on a surface). - `README.md` + `examples/` — quick start and end-to-end notebooks. - Docs: py.jaxonomy.com. # Installation ## Requirements - **Python** 3.10 or newer - A **64-bit** environment is assumed for JAX wheels on most platforms ## Recommended: virtual environment ``` python -m venv .venv source .venv/bin/activate # Linux / macOS # .venv\Scripts\activate # Windows cmd pip install --upgrade pip pip install jaxonomy ``` ## Platform notes - **Windows:** for double precision in JAX you may need\ `set JAX_ENABLE_X64=True` (cmd) or `$env:JAX_ENABLE_X64="True"` (PowerShell) before importing JAX. - **macOS:** some optional builds (e.g. pieces of the NMPC stack, and other compiled dependencies on Apple Silicon) need **cmake** (`brew install cmake`). ## Optional dependency groups Install extras in brackets: | Command | Typical use | | ---------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `pip install jaxonomy[safe]` | Scientific / ML extras, **no** NMPC/IPOPT. Note this is a **large, multi-GB** install: it pulls in PyTorch, TensorFlow, pandas, SymPy, python-control, PySINDy, pyTwin, Matplotlib, OpenCV, evosax, and nlopt. | | `pip install jaxonomy[nmpc]` | Nonlinear MPC blocks (cyipopt + OSQP) — requires **IPOPT** on the system. | | `pip install jaxonomy[all]` | Everything: the `[safe]` and `[nmpc]` sets **plus** MuJoCo / MJX / Brax. Largest install. | ### Nonlinear MPC (IPOPT) NMPC blocks expect **IPOPT** to be available on the machine: - **Ubuntu:** `sudo apt install coinor-libipopt-dev` - **macOS:** `brew install ipopt` (and `brew install cmake` if builds fail) Then: ``` pip install jaxonomy[nmpc] # or pip install jaxonomy[all] ``` If `pip install jaxonomy` or `pip install jaxonomy[all]` fails in the resolver (for example conflicting NumPy pins between transitive dependencies), fix or relax constraints in your environment, or install the project **from a clone** with `pip install -e . --no-deps` after you have compatible JAX/NumPy/etc. already installed. ## Development install (git clone) From the repository root: ``` pip install -e . ``` That makes `import jaxonomy` work from any working directory for that interpreter. **Jupyter / VS Code notebooks:** pick a kernel that uses the **same** Python where you ran `pip install -e .`, or add the repo root to `PYTHONPATH` / use: ``` import sys sys.path.insert(0, "/absolute/path/to/repo") ``` **Building the documentation site** (MkDocs): ``` pip install -r requirements.docs.txt mkdocs serve # preview at http://127.0.0.1:8000 by default ``` ______________________________________________________________________ License The `jaxonomy` package is released under the [MIT](https://mit-license.org/) license. # MCP server Jaxonomy ships a [Model Context Protocol](https://modelcontextprotocol.io) server that exposes the engine as tools an AI agent can call directly. Instead of writing Jaxonomy code and running it, the agent enumerates the block library, builds and validates a model, runs the simulation, and reads the actual numbers back. This page is the reference for that server. If you are writing Python by hand, you do not need any of it — `pip install jaxonomy` is enough. If you want an agent to *write* Jaxonomy for you rather than *drive* it, point it at [Using Jaxonomy from an AI agent](https://py.jaxonomy.com/agents/index.md) instead; the two are complementary. The server is registered in the [MCP Registry](https://registry.modelcontextprotocol.io) as `io.github.machinavitalis/jaxonomy`. ## Install The server lives behind an optional extra, so it is not installed by default: ``` pip install jaxonomy[mcp] ``` ## Configure a client The server speaks stdio. Point your client at the `jaxonomy-mcp` entry point, or equivalently at `python -m jaxonomy.mcp.server`. **Claude Code:** ``` claude mcp add jaxonomy -- jaxonomy-mcp ``` **Claude Desktop** — in `claude_desktop_config.json`: ``` { "mcpServers": { "jaxonomy": { "command": "jaxonomy-mcp" } } } ``` **Without installing first**, `uvx` can fetch the package and the extra in one step: ``` uvx --from 'jaxonomy[mcp]' jaxonomy-mcp ``` That is convenient for a one-off trial, but `uvx` builds a throwaway environment, so it pulls JAX and its dependencies on every cold start. For regular use, install into a real environment and point the client at that interpreter. Whichever form you use, the interpreter running the server must be the one where `jaxonomy[mcp]` is installed. A client that launches a bare `python` may pick up a different environment; give it an absolute path to the interpreter if the server fails to start. ## Tools The server exposes seven tools. Models are passed as JSON strings in Jaxonomy's model format; `list_blocks` is the usual starting point because it tells the agent what it has to work with. | Tool | What it does | | -------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `list_blocks` | Catalogue of available library block types, with descriptions and key parameters. | | `validate_model` | Checks a model JSON for structural and validation problems; returns `valid`, `errors`, `warnings`. | | `explain_model` | Plain-English description of a model's blocks, parameters, and signal flow. | | `run_simulation` | Runs a simulation over `[t_start, t_stop]`, recording named signals (e.g. `integrator.out_0`). Selectable `jax` or `numpy` backend. | | `fit_parameters` | Fits chosen parameters to measured data supplied as CSV, via finite-difference gradients and Adam, with optional bounds. | | `linearize_model` | Linearizes around an operating point; returns `A`, `B`, `C`, `D` and the eigenvalues. | | `influence_subgraph` | Serializes what actually drives a chosen signal — the dependency structure weighted by autodiff Jacobians, expanded strongest-edge-first under a token budget. | `influence_subgraph` exists for models too large to hand to an agent whole: it answers "what drives this signal, and by how much" while keeping the response inside a token budget, so what gets dropped is what mattered least. ## Limitations - `fit_parameters` uses **finite-difference** gradients, not Jaxonomy's end-to-end autodiff. It is a convenience path for an agent holding a CSV, not the recommended way to calibrate a model — for that, write the `jax.grad` loop directly (see [Using Jaxonomy from an AI agent](https://py.jaxonomy.com/agents/index.md)). - Models cross the boundary as JSON, so anything requiring a custom Python `LeafSystem` cannot be expressed through these tools. Custom blocks are a code-writing task. - The server is stdio-only; there is no hosted or HTTP transport. # API reference # Framework ## `jaxonomy.framework` ### `BlockInitializationError` Bases: `JaxonomyError` A generic error to be thrown when a block fails at init time, but the full exceptions are known to cause issues, eg. with ray serialization. ### `BlockParameterError` Bases: `StaticError` Block parameters are missing or have invalid values. ### `BlockRuntimeError` Bases: `JaxonomyError` A generic error to be thrown when a block fails at runtime, but the full exceptions are known to cause issues, eg. with ray serialization. ### `BusUnit` Compound unit carrying one :class:`Unit` per named bus field. Attached to the output port of a :class:`BusCreator` (and the matching input port of a :class:`BusSelector`) so that the connect-time consistency check can verify each field's unit individually. Attributes: | Name | Type | Description | | -------- | -------------------- | --------------------------------------------------------------------------------------------------- | | `fields` | `Mapping[str, Unit]` | Mapping from bus field name to its :class:Unit. Stored as a plain dict (insertion order preserved). | #### `field_unit(name)` Return the :class:`Unit` for `name`, or `None` if absent. Used by :class:`BusSelector` to look up its output-port unit when wired downstream of a unit-tagged bus. ### `ContextBase` Context object containing state, parameters, etc for a system. NOTE: Type hints in ContextBase indicate the union between what would be returned by a LeafContext and a DiagramContext. See type hints of the subclasses for the specific argument and return types. Attributes: | Name | Type | Description | | ---------------- | ------------ | ------------------------------------------------------------------------------------------------------------- | | `owning_system` | `SystemBase` | The owning system of the context. | | `time` | `Scalar` | The time associated with the context. Will be None unless the context is the root context. | | `is_initialized` | `bool` | Flag indicating if the context is initialized. This should only be set by the ContextFactory during creation. | #### `__getitem__(key)` Get the subcontext associated with the given system ID. For leaf contexts, this will return `self`, but the method is provided so that there is a consistent interface for working with either an individual LeafSystem or tree-structured Diagram. For nested diagrams, intermediate diagrams do not have associated contexts, so indexing will fail. #### `with_continuous_state(value)` Create a copy of this context, replacing the continuous state. #### `with_discrete_state(value)` Create a copy of this context, replacing the discrete state. #### `with_mode(value)` Create a copy of this context, replacing the mode. #### `with_new_state()` Create a copy of this context, replacing the state with a new state. #### `with_parameter(name, value)` Create a copy of this context, replacing the specified parameter. #### `with_parameters(new_parameters)` Create a copy of this context, replacing only the specified parameters. #### `with_state(state)` Create a copy of this context, replacing the entire state. #### `with_subcontext(key, ctx)` Create a copy of this context, replacing the specified subcontext. #### `with_time(value)` Create a copy of this context, replacing time with the given value. This should only be called on the root context, since it is expected that all subcontexts will have a time value of None to avoid any conflicts. #### `with_updated_parameters()` Create a copy of this context, updating all parameters to their current values. ### `DependencyTicket` Singleton class for managing unique dependency tickets. ### `Diagram` Bases: `SystemBase` Composite block-diagram representation of a dynamical system. A Diagram is a collection of Systems connected together to form a larger hybrid dynamical system. Diagrams can be nested to any depth, creating a tree-structured block diagram. NOTE: The Diagram class is not intended to be constructed directly. Instead, use the `DiagramBuilder` to construct a Diagram, which will pass the appropriate information to this constructor. #### `continuous_substep_vector` T-133: per-leaf multirate substep vectors, in `leaf_systems` order (the same ordering `mass_matrix` relies on for alignment with the flattened continuous state). #### `has_dirty_static_parameters` Check if any static parameters have been modified. #### `check_no_algebraic_loops()` Check for algebraic loops in the diagram. This is a more or less direct port of the Drake method DiagramBuilder::ThrowIfAlgebraicLoopExists. Some comments are verbatim explanations of the algorithm implemented there. #### `check_types(context, error_collector=None)` Perform any system-specific static analysis. #### `declare_dynamic_parameter(name, parameter)` Declare a parameter for this system. Parameters: | Name | Type | Description | Default | | ----------- | ----------- | -------------------------- | ---------- | | `name` | `str` | The name of the parameter. | *required* | | `parameter` | `Parameter` | The parameter object. | *required* | #### `eval_subsystem_input_port(context, port_locator)` Evaluate the input port for a child of this system given the root context. Parameters: | Name | Type | Description | Default | | -------------- | ------------------ | -------------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | root context for this system | *required* | | `port_locator` | `InputPortLocator` | tuple of (system, port_index) identifying the input port to evaluate | *required* | Returns: | Name | Type | Description | | ------- | ------- | -------------------------------------------------- | | `Array` | `Array` | Value returned from evaluating the subsystem port. | Raises: | Type | Description | | ------------------------ | ---------------------------------- | | `InputNotConnectedError` | if the input port is not connected | #### `eval_subsystem_output_port(context, port_locator)` "Evaluate the output port for a child of this system given the root context. Parameters: | Name | Type | Description | Default | | -------------- | ------------------- | --------------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | root context for this system | *required* | | `port_locator` | `OutputPortLocator` | tuple of (system, port_index) identifying the output port to evaluate | *required* | Returns: | Name | Type | Description | | ------- | ------- | -------------------------------------------------- | | `Array` | `Array` | Value returned from evaluating the subsystem port. | #### `export_input(locator, port_name)` Export a subsystem input port as a diagram-level input. This should typically only be called during construction by DiagramBuilder. The standard workflow will be to call export_input on the *builder* object, which will automatically call this method on the Diagram once created. Parameters: | Name | Type | Description | Default | | ----------- | ------------------ | ------------------------------------------------------------------ | ---------- | | `locator` | `InputPortLocator` | tuple of (system, port_index) identifying the input port to export | *required* | | `port_name` | `str` | name of the new exported input port | *required* | Returns: | Name | Type | Description | | ----- | ----- | ---------------------------------------------------------------- | | `int` | `int` | index of the exported input port in the diagram input_ports list | #### `export_output(locator, port_name)` Export a subsystem output port as a diagram-level output. This should typically only be called during construction by DiagramBuilder. The standard workflow will be to call export_input on the *builder* object, which will automatically call this method on the Diagram once created. Parameters: | Name | Type | Description | Default | | ----------- | ------------------- | ------------------------------------------------------------------- | ---------- | | `locator` | `OutputPortLocator` | tuple of (system, port_index) identifying the output port to export | *required* | | `port_name` | `str` | name of the new exported output port | *required* | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------------------------------ | | `int` | `int` | index of the exported output port in the diagram output_ports list | #### `get_parameter(path)` Get a parameter by dot-separated path (child blocks and nested diagrams). For a path `"block.param"`, `block` must be a direct child name of this diagram; the remainder is resolved on that child (recursively for nested diagrams). A single segment refers to this diagram's own parameters (same as :meth:`SystemBase.get_parameter`). Examples: `diagram.get_parameter("gain.gain")` for a child named `gain` with parameter `gain`. Raises: | Type | Description | | ---------- | ------------------------------------------------- | | `KeyError` | If a segment does not match a child or parameter. | #### `initialize_static_data(context)` Perform any system-specific static analysis. #### `list_parameters(prefix='')` Flatten parameters under this diagram with dot-notation keys. Includes this diagram's own parameters (if any), then each child's parameters prefixed by `child_name.`. Nested diagrams recurse. Parameters: | Name | Type | Description | Default | | -------- | ----- | --------------------------------------------------------------------------------------------- | ------- | | `prefix` | `str` | Internal use: prepend to every key (non-empty when called recursively from a parent diagram). | `''` | #### `post_simulation_finalize()` Perform any post-simulation cleanup for this system. #### `print_schedule(*, format='text', file=None, ensure_initialized=True)` Print the inferred sample-time schedule of this diagram. Renders one entry per rate group — period (for discrete blocks), the leaves that fire at that rate, any detected rate mismatches, and the deterministic execution order. Pure inspection helper; does not mutate the diagram or its contexts. T-105-followup-print-schedule-pre-context: by default, `print_schedule()` lazily calls :meth:`create_context()` first so that discrete blocks whose periodic events are registered in their :meth:`initialize` hook (e.g. :class:`PIDDiscrete`, :class:`Decimator`, :class:`UnitDelay`, :class:`ZeroOrderHold`) are bucketed into the correct rate group. Pre-fix those blocks showed up as `constant` when `print_schedule()` ran before any context had been created, because the periodic event hadn't yet been declared. The lazy `create_context()` call is idempotent and cheap on already-initialised diagrams; if it fails (e.g. the diagram has missing connections), the schedule is rendered anyway with a one-line warning explaining the rate-group output may be incomplete. Pass `ensure_initialized=False` to opt out and use the pre-fix behaviour (useful for debugging the initialisation path itself). Parameters: | Name | Type | Description | Default | | -------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------- | | `format` | `str` | "text" (default), "markdown", or "json". | `'text'` | | `file` | | Destination file-like object. Defaults to sys.stdout when None. | `None` | | `ensure_initialized` | `bool` | If True (default), call :meth:create_context once before rendering so all discrete blocks have registered their periodic events. Set to False to render against the current (possibly pre-init) state. | `True` | Examples: ``` >>> diagram.print_schedule() rate groups: discrete(period=0.001, offset=0.0): ctrl_inner discrete(period=0.01, offset=0.0): ctrl_outer continuous: plant execution order: ... ``` ``` >>> with open("schedule.md", "w") as f: ... diagram.print_schedule(format="markdown", file=f) ``` See also :func:`jaxonomy.simulation.rate_groups.rate_summary` — the underlying string formatter, useful when you want the result as a string for embedding in a manifest or PR body. :func:`jaxonomy.simulation.rate_groups.rate_summary_dot` — the DOT-format companion for graphviz visualization. #### `with_parameters(updates)` Return a new diagram with parameters replaced (dot-notation paths). Grouping is by top-level block name; nested paths are forwarded recursively. The original diagram is unchanged. Parameters: | Name | Type | Description | Default | | --------- | ---------------- | ---------------------------------------------------------------------------------------------------- | ---------- | | `updates` | `dict[str, Any]` | Map from dot paths to new values, e.g. {"motor.R": jnp.array(2.3), "controller.Kp": jnp.array(1.5)}. | *required* | Returns: | Name | Type | Description | | ----- | --------- | ----------------------- | | `New` | `Diagram` | class:Diagram instance. | Raises: | Type | Description | | ----------- | -------------------------------------- | | `KeyError` | Unknown block or parameter. | | `TypeError` | Attempt to replace a static parameter. | ### `DiagramBuilder` Class for constructing block diagram systems. The `DiagramBuilder` class is responsible for building a diagram by adding systems, connecting ports, and exporting inputs and outputs. It keeps track of the registered systems, input and output ports, and the connection map between input and output ports of the child systems. #### `__init__(*, validate_rates_at_connect=None, unit_conversion='auto', auto_insert_rate_transitions=False)` Construct a DiagramBuilder. Parameters: | Name | Type | Description | Default | | ------------------------------ | ------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------- | | `validate_rates_at_connect` | \`str | bool | None\` | | `unit_conversion` | `str` | T-104 followup — controls behaviour when two connected ports share base-dimensions but differ only by a scalar scale (e.g. meter vs kilometer): * "auto" (default) silently inserts the conversion factor on the destination input port; * "warn" inserts the factor and emits a :class:UserWarning; * "error" refuses the connection (preserves the Phase-1 strict-equal behaviour). Genuine dimensional mismatches (e.g. meter vs second) always raise regardless of mode. | `'auto'` | | `auto_insert_rate_transitions` | `bool` | T-105-followup-phase3 — when True, :meth:connect automatically synthesises a :func:jaxonomy.library.RateTransition block (a ZeroOrderHold for slow→fast, a Decimator for fast→slow) between any two adjacent leaves whose inferred discrete sample times differ. The rewritten wiring is src → rate_transition → dst and an informational log line documents the insertion. Composes with validate_rates_at_connect: when both are enabled, the warning still fires and the transition still gets inserted. Default False keeps the legacy code path byte-equivalent (the strict mode that surfaces rate mismatches rather than silently inserting transitions). | `False` | #### `add(*systems)` ``` add(system: SystemBase) -> SystemBase ``` ``` add(system: SystemBase, *systems: SystemBase) -> List[SystemBase] ``` Add one or more systems to the diagram. Parameters: | Name | Type | Description | Default | | --------------------- | ---- | -------------------------------- | ---------- | | `*systems SystemBase` | | System(s) to add to the diagram. | *required* | Returns: | Type | Description | | ------------------ | ------------ | | \`List[SystemBase] | SystemBase\` | Raises: | Type | Description | | -------------- | -------------------------------------- | | `BuilderError` | If the diagram has already been built. | | `BuilderError` | If the system is already registered. | | `BuilderError` | If the system name is not unique. | #### `build(name='root', ui_id=None, parameters=None)` Builds a Diagram system with the specified name and system ID. Parameters: | Name | Type | Description | Default | | ------------ | ---------------------- | -------------------------------------------------------------- | -------- | | `name` | `str` | The name of the diagram. Defaults to "root". | `'root'` | | `ui_id` | `str` | The unique identifier for the diagram. | `None` | | `parameters` | `dict[str, Parameter]` | A dictionary of dynamic parameters to declare for the diagram. | `None` | Returns: | Name | Type | Description | | --------- | --------- | ------------------------------ | | `Diagram` | `Diagram` | The newly constructed diagram. | Raises: | Type | Description | | ------------------------ | ------------------------------------------------ | | `EmptyDiagramError` | If no systems are registered in the diagram. | | `BuilderError` | If the diagram has already been built. | | `AlgebraicLoopError` | If an algebraic loop is detected in the diagram. | | `DisconnectedInputError` | If an input port is not connected. | #### `connect(src, dest)` Connect an output port to an input port. The input port and output port must both belong to systems that have already been added to the diagram. The input port must not already be connected to another output port. Parameters: | Name | Type | Description | Default | | ------ | ------------ | --------------------------- | ---------- | | `src` | `OutputPort` | The output port to connect. | *required* | | `dest` | `InputPort` | The input port to connect. | *required* | Raises: | Type | Description | | -------------- | ------------------------------------------------ | | `BuilderError` | If the diagram has already been built. | | `BuilderError` | If the source system is not registered. | | `BuilderError` | If the destination system is not registered. | | `BuilderError` | If the input port is already connected. | | `BuilderError` | If src is an InputPort or dest is an OutputPort. | #### `export_input(port, name=None)` Export an input port of a child system as a diagram-level input. The input port must belong to a system that has already been added to the diagram. The input port must not already be connected to another output port. Parameters: | Name | Type | Description | Default | | ------ | ----------- | ------------------------------------------------------------------------------------------------ | ---------- | | `port` | `InputPort` | The input port to export. | *required* | | `name` | `str` | The name to assign to the exported input port. If not provided, a unique name will be generated. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------------------------------ | | `int` | `int` | The index (in the to-be-built diagram) of the exported input port. | Raises: | Type | Description | | -------------- | --------------------------------------- | | `BuilderError` | If the diagram has already been built. | | `BuilderError` | If the system is not registered. | | `BuilderError` | If the input port is already connected. | | `BuilderError` | If the input port name is not unique. | #### `export_output(port, name=None)` Export an output port of a child system as a diagram-level output. The output port must belong to a system that has already been added to the diagram. Parameters: | Name | Type | Description | Default | | ------ | ------------ | ------------------------------------------------------------------------------------------------- | ---------- | | `port` | `OutputPort` | The output port to export. | *required* | | `name` | `str` | The name to assign to the exported output port. If not provided, a unique name will be generated. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------------------------------- | | `int` | `int` | The index (in the to-be-built diagram) of the exported output port. | Raises: | Type | Description | | -------------- | -------------------------------------- | | `BuilderError` | If the diagram has already been built. | | `BuilderError` | If the system is not registered. | | `BuilderError` | If the output port name is not unique. | ### `DiagramContext` Bases: `ContextBase` #### `with_parameters(new_parameters)` Create a copy of this context, replacing only the specified parameters. ### `DiscreteUpdateEvent` Bases: `Event` Event representing a discrete update in a hybrid system. ### `DtypeMismatchError` Bases: `StaticError` Block parameters or input/outputs have mismatched dtypes. ### `EnabledMode` Allowed string values for `EnabledSubsystem.mode`. ### `EnabledStateMode` Allowed string values for `EnabledSubsystem.state_mode`. Controls how the *continuous state* (declared via `state_dynamics`) evolves while the enable signal is false: - `HOLD` (default): freeze the state at its current value (`xdot = 0` while disabled). Resumes integration on re-enable. - `RESET`: snap the state back to `initial_state` on every disable→enable transition (so each enable window starts from the configured initial value). While disabled, the state is held. - `FREE`: the state evolves according to `state_dynamics` regardless of enable. Only the *output* is masked per `mode=`. ### `EnabledSubsystem` Bases: `LeafSystem` Container block: run a submodel only while an enable signal is true. This is the subsystem-framing wrapper around the existing :class:`jaxonomy.library.Conditional` primitive (T-009). It exists as a separate class so that: - The block-diagram-vocabulary name `EnabledSubsystem` is discoverable next to the rest of the container family. - We can later extend the `mode="hold"` path with subsystem-state semantics (per-block discrete-state binding) without disturbing the lighter `Conditional` primitive. Parameters: | Name | Type | Description | Default | | ---------------- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output (single output per phase 1). Must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of submodel inputs (does NOT include the enable port). Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel inputs. | `1` | | `mode` | `Literal['reset', 'passthrough', 'hold']` | One of "reset" / "passthrough" / "hold". reset: when disabled, output = initial_value. passthrough: when disabled, output = first user input (input port 1). Submodel and passthrough output must broadcast-compatibly. hold: when disabled, output holds the most recent snapshot taken at hold_period. Requires a positive hold_period. | `RESET` | | `initial_value` | | Output value when disabled in reset mode, and the seed for the held discrete state in hold mode. Used to infer output shape/dtype. | `0.0` | | `hold_period` | \`float | None\` | Sample period (seconds) for the held snapshot in hold mode. Required iff mode == "hold". | | `state_mode` | `Literal['hold', 'reset', 'free']` | One of "hold" / "reset" / "free". Controls the continuous-state behaviour while disabled (independent of mode= which gates only the output). See :class:EnabledStateMode. Default "hold". Only has an effect when state_dynamics is provided; for the stateless submodel default this kwarg is validated but otherwise a no-op (so the default-off path is byte-equivalent to phase 1). | `HOLD` | | `state_dynamics` | \`Callable | None\` | Optional callable f(t, x, \*user_inputs) -> xdot defining a continuous state for the EnabledSubsystem itself. When provided, the block declares a continuous state seeded by initial_state (or initial_value if initial_state is None) and applies state_mode semantics around it. When omitted, the block has no continuous state and behaves exactly as in T-120 phase 1. | | `initial_state` | | Initial value of the continuous state. Required when state_dynamics is provided. | `None` | | `name` | | Optional block name. | *required* | ### `ErrorCollector` Tool used to collect errors related to users model specification. Errors related to user model specification are identified during model static analysis, e.g. context creation, type checking, etc. An instance of this tool can be created, and then passed down a tree of function calls to collect errors found any where in the tree. Locally in the tree it can be determined whether it is ok to continue or not. This tool enables collecting errors up until the point when continuation is no longer possible. Note: this latter behavior, where sometimes there is early exit desired, and all other "pipeline" operations are "nullified", might better be implemented using pymonad:Either class. #### `add_error(error)` Add an error to the collection. #### `context(parent=None)` A context manager convenience to use when tracing errors. Use as: ``` with ErrorCollector.trace(error_context) as ec: ... ``` If the parent context is None, then exceptions will pass through without being collected. Else, exceptions will be collected in the parent context. ### `EventCollection` A collection of events owned by a system. Users should not need to interact with these objects directly. They are intended to be used internally by the simulation framework for handling events in hybrid system simulation. These contain callback functions that update the context in various ways when the event is triggered. There will be different "collections" for each trigger type in simulation (e.g. periodic vs zero-crossing). Within the collections, events are broken out by function (e.g. discrete vs unrestricted updates). There are separate implementations for leaf and diagram systems, where the DiagramCEventCollection preserves the tree structure of the underlying Diagram. However, the interface in both cases is the same and is identical to the interface defined by EventCollection. ### `ForLoop` Bases: `LeafSystem` Container block: run `body_fn` `n_iter` times per major step. `ForLoop` wraps :func:`jax.lax.fori_loop`. The block declares a single input port carrying the *initial carry value* and a single output port returning the carry after `n_iter` iterations. Parameters: | Name | Type | Description | Default | | --------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `body_fn` | `Callable` | Callable (i: int, carry) -> carry. Must be JAX-traceable. The carry pytree must have a fixed structure and shape across iterations (this is a lax.fori_loop requirement, not a Jaxonomy choice). | *required* | | `n_iter` | `int` | Number of iterations. Must be a non-negative Python int (static); a runtime-traced n_iter would force lax.while_loop semantics and is not supported here — use :class:WhileLoop for that case. | *required* | | `name` | | Optional block name. | *required* | Differentiability :func:`jax.grad` flows through `body_fn`'s parameters and through the initial carry. The loop count `n_iter` is static and not differentiable. Example A body that accumulates `i` into the carry over 10 iterations yields `carry_initial + (0+1+...+9) = carry + 45`. Notes - `body_fn` must close over any constants it needs; the `i`-th iteration receives only `(i, carry)`. - Per T-005, default float dtype is float64 unless the active precision policy says otherwise; `ForLoop` does not cast. ### `IntegerTime` Class for managing conversion between decimal and integer time. #### `as_decimal(time)` Convert an integer time to a floating-point time. #### `from_decimal(time)` Convert a floating-point time to an integer time. ### `JaxonomyError` Bases: `Exception` Base class for all custom jaxonomy errors. #### `__init__(message=None, *, system=None, system_id=None, name_path=None, ui_id_path=None, port_index=None, port_name=None, port_direction=None, parameter_name=None, loop=None)` Create a new JaxonomyError. Only `message` is a positional argument, all others are keyword arguments. Parameters: | Name | Type | Description | Default | | ---------------- | --------------------------- | --------------------------------------------------------------------------------------------- | ------- | | `message` | | A custom error message, defaults to the error class name. | `None` | | `system` | `SystemBase` | The system that the error occurred in, if available. | `None` | | `system_id` | `Hashable` | The id of the system that the error occurred in, use if system can't be passed. | `None` | | `name_path` | `list[str]` | The name path of the block that the error occurred in, use if system can't be passed. | `None` | | `ui_id_path` | `list[str]` | The ui_id (uuid) path of the block that the error occurred in, use if system can't be passed. | `None` | | `port_index` | `int` | The index of the port that the error occurred at. | `None` | | `port_name` | `str` | The name of the port that the error occurred at. | `None` | | `port_direction` | `str` | The direction of the port that the error occurred at. | `None` | | `parameter_name` | `str` | The name of the parameter that the error occurred at. | `None` | | `loop` | `list[DirectedPortLocator]` | A list of I/O ports where the error occurred (eg. AlgebraicLoopError). | `None` | #### `caused_by(exc_type)` Check if this error is or was caused by another error type. For instance, if a JaxonomyError is raised because of a TypeError, this method will return True when called with TypeError as exc_type. Parameters: | Name | Type | Description | Default | | ---------- | ------ | -------------------------------------------------- | ---------- | | `exc_type` | `type` | The type of exception to check for (eg. TypeError) | *required* | Returns: | Name | Type | Description | | ------ | ---- | --------------------------------------------------------------- | | `bool` | | True if the error is or was caused by the given exception type. | ### `LeafContext` Bases: `ContextBase` #### `__getitem__(key)` Dummy indexing for compatibility with DiagramContexts, returning self. #### `with_parameters(new_parameters)` Create a copy of this context, replacing only the specified parameters. #### `with_subcontext(key, ctx)` Dummy replacement for compatibility with DiagramContexts, returning ctx. ### `LeafState` Container for state information for a leaf system. Attributes: | Name | Type | Description | | ------------------ | --------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `name` | `str` | Name of the leaf system that owns this state. | | `continuous_state` | `LeafStateComponent` | Continuous state of the system, i.e. the component of state that evolves in continuous time. If the system has no continuous state, this will be None. | | `discrete_state` | `LeafStateComponent` | Discrete state of the system, i.e. one or more components of state that do not change continuously with ime (not necessarily discrete-valued). If the system has no discrete state, this will be None. | | `mode` | `int` | An integer value indicating the current "mode", "stage", or discrete-valued state component of the system. Used for finite state machines or multi-stage hybrid systems. If the system has no mode, this will be None. | | `cache` | `tuple[LeafStateComponent]` | The current values of sample-and-hold outputs from the system. In a pure discrete system these would not be state components (just results of feedthrough computations), but in a hybrid or multirate system they act as discrete state from the perspective of continuous or asynchronous discrete components of the system. Hence, they are stored in the state, but are maintained separately from the normal internal state of the system. | Notes (1) This class is immutable. To modify a LeafState, use the `with_*` methods. (2) The type annotations for state components are LeafStateComponent, which is a union of array, tuple, and named tuple. The most common case is arrays, but this allows for more flexibility in defining state components, e.g. a second-order system can define a named tuple of generalized coordinates and velocities rather than concatenating into a single array. #### `with_cached_value(index, value)` Create a copy of this LeafState with the specified cache value replaced. #### `with_continuous_state(value)` Create a copy of this LeafState with the continuous state replaced. #### `with_discrete_state(value)` Create a copy of this LeafState with the discrete state replaced. #### `with_mode(value)` Create a copy of this LeafState with the mode replaced. ### `LeafSystem` Bases: `SystemBase` Basic building block for dynamical systems. A LeafSystem is a minimal component of a system model in jaxonomy, containing no subsystems. Inputs, outputs, state, parameters, updates, etc. can be added to the block using the various `declare_*` methods. The built-in blocks in jaxonomy.library are all subclasses of LeafSystem, as are any custom blocks defined by the user. #### `continuous_state_default` The declared default continuous-state value (read-only). This is the `default_value` passed to `declare_continuous_state` (or the array inferred from `shape` / `dtype`), i.e. the value that seeds `context.continuous_state` before any user override. Returns `None` when the block has no continuous state. Exposed as a documented accessor so callers don't have to reach into the private `_default_continuous_state` attribute (T-C2-followup). #### `continuous_substep_vector` T-133: per-entry multirate substep factors for this block. Returns pytree-structured int vectors aligned with the flattened continuous state (same leaves-concatenation ordering the ODE solvers use for `mass_matrix`), or `None` when the block has no continuous state. Every entry carries the block-level factor declared via `declare_continuous_state(substeps=N)` (default 1). #### `has_multirate_substeps` True when this block declared `substeps > 1` (T-133). #### `configure_continuous_state(callback_idx, shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None)` Configure a continuous state component for the system. The `ode` callback computes the time derivative of the continuous state based on the current time, state, and any additional inputs. If `ode` is not provided, a default zero vector of the same size as the continuous state is used. If provided, the `ode` callback should have the signature `ode(time, state, *inputs, **params) -> xcdot`. Parameters: | Name | Type | Description | Default | | ----------------------- | ------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `callback_idx` | `int` | The index of the callback in the system's callback list. | *required* | | `shape` | `ShapeLike` | The shape of the continuous state vector. Defaults to None. | `None` | | `default_value` | `Array` | The initial value of the continuous state vector. Defaults to None. | `None` | | `dtype` | `DTypeLike` | The data type of the continuous state vector. Defaults to None. | `None` | | `ode` | `Callable` | The callback for computing the time derivative of the continuous state. Should have the signature: ode(time, state, \*inputs, \*\*parameters) -> xcdot. Defaults to None. | `None` | | `mass_matrix` | `Array` | The mass matrix for the continuous state. Defaults to None. If provided, must be a square matrix with the same shape as the continuous state. Using a mass matrix different from the identity in any LeafSystem will require the use of a compatible continuous-time solver (currently only BDF is supported). Currently mass matrices are also only supported for scalar- or vector-valued continuous states ( i.e. no matrices or other PyTree-structured states). | `None` | | `as_array` | `bool` | If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification. | `True` | | `requires_inputs` | `bool` | If True, indicates that the ODE computation requires inputs. | `True` | | `prerequisites_of_calc` | `List[DependencyTicket]` | The dependency tickets for the ODE computation. Defaults to None, in which case the assumption is a dependency on either (time, continuous state) if requires_inputs is False, otherwise (time, continuous state, inputs. | `None` | Raises: | Type | Description | | ---------------- | ---------------------------------------------------------------------------------------------------------------- | | `AssertionError` | If neither shape nor default_value is provided, or if the mass matrix is inconsistent with the continuous state. | Notes (1) Only one of `shape` and `default_value` should be provided. If `default_value` is provided, it will be used as the initial value of the continuous state. If `shape` is provided, the initial value will be a zero vector of the given shape and specified dtype. #### `configure_output_port(port_index, callback, period=None, offset=0.0, prerequisites_of_calc=None, default_value=None, requires_inputs=None)` Configure an output port in the LeafSystem. See `declare_output_port` for a description of the arguments. Parameters: | Name | Type | Description | Default | | ------------ | ----- | ------------------------------------------ | ---------- | | `port_index` | `int` | The index of the output port to configure. | *required* | Returns: | Type | Description | | ---- | ----------- | | | None | #### `configure_periodic_update(event_index, callback, period, offset, enable_tracing=None)` Configure an existing periodic update event. The event will be triggered at regular intervals defined by the period and offset parameters. The callback should have the signature `callback(time, state, *inputs, **params) -> xd_plus`, where `xd_plus` is the updated value of the discrete state. This callback should be written to compute the "plus" value of the discrete state component given the "minus" values of all state components and inputs. Parameters: | Name | Type | Description | Default | | ---------------- | ---------- | ------------------------------------------------------------------- | ---------- | | `event_index` | `int` | The index of the event to configure. | *required* | | `callback` | `Callable` | The callback function defining the update. | *required* | | `period` | `Scalar` | The period at which the update event occurs. | *required* | | `offset` | `Scalar` | The offset at which the first occurrence of the event is triggered. | *required* | | `enable_tracing` | `bool` | If True, enable tracing for this event. Defaults to None. | `None` | #### `declare_cache(callback, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=True)` Declare a stored computation for the system. This method accepts a callback function with the block-level signature `callback(time, state, *inputs, **parameters) -> value` and wraps it to have the signature `callback(context) -> value` This callback can optionally be used to define a periodic update event that refreshes the cached value. Other calculations (e.g. sample-and-hold output ports) can then depend on the cached value. Parameters: | Name | Type | Description | Default | | ----------------------- | ------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `callback` | `Callable` | The callback function defining the cached computation. | *required* | | `period` | `float` | If not None, the callback function will be used to define a periodic update event that refreshes the value. Defaults to None. | `None` | | `offset` | `float` | The offset of the periodic update event. Defaults to 0.0. Will be ignored unless period is not None. | `0.0` | | `name` | `str` | The name of the cached value. Defaults to None. | `None` | | `default_value` | `Array` | The default value of the result, if known. Defaults to None. | `None` | | `requires_inputs` | `bool` | If True, the callback will eval input ports to gather input values. This will add a bit to compile time, so setting to False where possible is recommended. Defaults to True. | `True` | | `prerequisites_of_calc` | `List[DependencyTicket]` | The dependency tickets for the computation. Defaults to None, in which case the default is to assume dependency on either (inputs) if requires_inputs is True, or (nothing) otherwise. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------------------------------------------------------------------------------------------- | | `int` | `int` | The index of the callback in system.callbacks. The cache index can recovered from system.callbacks[callback_index].cache_index. | #### `declare_continuous_state(shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None, substeps=1, project=None)` Declare a continuous state component for the system. The continuous state value is read inside callbacks as `state.continuous_state` (the `state` argument of the `ode` / output callbacks). **Unpack contract** (T-C3-followup): the shape of `state.continuous_state` mirrors exactly what you passed as `default_value` (or the zeros array implied by `shape` / `dtype`): - A scalar default (`jnp.array(0.0)`) gives a scalar `state.continuous_state` — read it directly, do **not** index. - A vector default (`jnp.zeros(3)`) gives a length-3 array — index / unpack as `x, y, z = state.continuous_state` or `state.continuous_state[i]`. - A PyTree default (tuple / NamedTuple / dict) gives back the same PyTree structure; your `ode` must return `xcdot` with the identical structure. The `ode` callback's return value must match the `default_value` structure element-for-element, since it is added to the state during integration. A common error is declaring a scalar state but returning `jnp.array([xdot])` (shape `(1,)`) from the ode — keep both scalar or both vector. Multirate substepping (T-133): `substeps=N` declares that this block's continuous dynamics have a fast time constant needing `N` inner integration steps per outer solver step (e.g. a motor's electrical winding inside a 1 kHz control loop). Honored by the fixed-step `rk4` solver (`SimulatorOptions(ode_solver_method= "rk4")`): the block's states advance with `N` RK4 substeps of `h/N` while the rest of the diagram takes one step of `h`, with first-order (zero-order-hold) coupling at the boundary — each side sees the other's start-of-step values, matching the semantics of a hand-rolled JIT-safe substep loop. Adaptive solvers (`dopri5`/`bdf`) ignore the declaration — they control stiffness through global step adaptation. `N` must be a static Python `int >= 1`; the default 1 is byte-equivalent to the pre-T-133 behavior. Reverse-mode autodiff (`enable_autodiff=True`) is supported — the substep loop has a static trip count and the checkpointed adjoint substeps the costates alongside their primals. Gradient accuracy carries the scheme's first-order coupling error: the adjoint converges to the true sensitivity linearly in the outer step `h` (exact FD agreement is only recovered as `h` is refined), and for dynamics *unstable at the outer step* the adjoint's reverse-time primal re-integration further limits accuracy. Reduce the outer step when gradients through the coupling interface need to be tight. Declared state projection (T-132): `project=fn` declares that this block's continuous state lives on a manifold and supplies the retraction back onto it — e.g. unit-quaternion renormalization for an attitude state (`nq=4` integrated componentwise drifts off the unit sphere under any one-step integrator). `fn(x) -> x` receives the state in its declared structure, must be shape-preserving and jit-safe, and is applied by the simulator **at the end of every major step** (composing with, and independent of, the T-003a DAE projection). Within-step drift is bounded by the step size; the recorded trajectory and all values other blocks see at major-step boundaries are on the manifold. Differentiable: the projection participates in reverse-mode AD as ordinary traced ops. #### `declare_continuous_state_output(name=None)` Declare a continuous state output port in the system. This method creates a new block-level output port which returns the full continuous state of the system. Parameters: | Name | Type | Description | Default | | ------ | ----- | ------------------------------------------------------------------ | ------- | | `name` | `str` | The name of the output port. Defaults to None (autogenerate name). | `None` | Returns: | Name | Type | Description | | ----- | ----- | --------------------------------- | | `int` | `int` | The index of the new output port. | #### `declare_discrete_state(shape=None, default_value=None, dtype=None, as_array=True, name=None)` Declare a discrete state component for the system. The discrete state is a component of the system's state that can be updated at specific events, such as zero-crossings or periodic updates. .. note:: Currently only **one** discrete state component is supported per `LeafSystem`. If `declare_discrete_state` is called more than once, the second call will silently overwrite the first. To store several independent values, pack them into a single array and split inside your update callback. Parameters: | Name | Type | Description | Default | | --------------- | ----------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `shape` | `ShapeLike` | The shape of the discrete state. Defaults to None. | `None` | | `default_value` | `Array` | The initial value of the discrete state. Defaults to None. | `None` | | `dtype` | `DTypeLike` | The data type of the discrete state. Defaults to None. | `None` | | `as_array` | `bool` | If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification. | `True` | | `name` | `str` | Readability label for the discrete state (parity with declare_continuous_state_output(name=...)). Stored as self.discrete_state_name for diagnostics/debugging; it does not change runtime behaviour, and the state is still read as state.discrete_state. | `None` | Raises: | Type | Description | | ---------------- | -------------------------------------------------------------------- | | `AssertionError` | If as_array is True and neither shape nor default_value is provided. | Notes (1) Only one of `shape` and `default_value` should be provided. If `default_value` is provided, it will be used as the initial value of the continuous state. If `shape` is provided, the initial value will be a zero vector of the given shape and specified dtype. (2) Use `declare_periodic_update` to declare an update event that modifies the discrete state at a recurring interval. #### `declare_mode_output(name=None)` Declare a mode output port in the system. This method creates a new block-level output port which returns the component of the system's state corresponding to the discrete "mode" or "stage". Parameters: | Name | Type | Description | Default | | ------ | ----- | ---------------------------------------------- | ------- | | `name` | `str` | The name of the output port. Defaults to None. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------- | | `int` | `int` | The index of the declared mode output port. | #### `declare_output_port(callback=None, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=None, units=None)` Declare an output port in the LeafSystem. This method accepts a callback function with the block-level signature `callback(time, state, *inputs, **parameters) -> value` and wraps it to the signature expected by SystemBase.declare_output_port: `callback(context) -> value` The callback is passed **positionally** (there is no `eval=` or `calc=` keyword). A computed (non-state) output looks like:: ``` class Thermometer(LeafSystem): def __init__(self): super().__init__() self.declare_input_port(name="q_in") self.declare_continuous_state(default_value=jnp.zeros(2), ode=self._ode) # Output reads the continuous state -> declare the xc # prerequisite; it does not read the input -> say so. self.declare_output_port( self._temperature, name="T", requires_inputs=False, prerequisites_of_calc=[DependencyTicket.xc], ) def _temperature(self, time, state, *inputs, **parameters): return state.continuous_state[0] ``` A feedthrough output that transforms its inputs instead declares `requires_inputs=True` (the default) and reads them positionally: `def _out(self, time, state, u, v): return u + v`. Mistakes in the callback signature surface at trace time, not at declaration. Parameters: | Name | Type | Description | Default | | ----------------------- | ------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `callback` | `Callable` | The callback function defining the output port. | `None` | | `period` | `float` | If not None, the port will act as a "sample-and-hold", with the callback function used to define a periodic update event that refreshes the value that will be returned by the port. Typically this should match the update period of some associated update event in the system. Defaults to None. | `None` | | `offset` | `float` | The offset of the periodic update event. Defaults to 0.0. Will be ignored unless period is not None. | `0.0` | | `name` | `str` | The name of the output port. Defaults to None. | `None` | | `default_value` | `Array` | The default value of the output port, if known. Defaults to None. | `None` | | `requires_inputs` | \`bool | list[int] | None\` | | `prerequisites_of_calc` | `List[DependencyTicket]` | The dependency tickets for the output port computation. Defaults to None, in which case the assumption is a dependency on either (nothing) if requires_inputs is False otherwise (inputs). | `None` | Returns: | Name | Type | Description | | ----- | ----- | -------------------------------------- | | `int` | `int` | The index of the declared output port. | #### `declare_zero_crossing(guard, reset_map=None, start_mode=None, end_mode=None, direction='crosses_zero', terminal=False, name=None, enable_tracing=None, zeno_tolerance=None, grad_guard=None)` Declare an event triggered by a zero-crossing of a guard function. Optionally, the system can also transition between discrete modes If `start_mode` and `end_mode` are specified, the system will transition from `start_mode` to `end_mode` when the event is triggered according to `guard`. This event will be active conditionally on `state.mode == start_mode` and when triggered will result in applying the reset map. In addition, the mode will be updated to `end_mode`. If `start_mode` and `end_mode` are not specified, the event will always be active and will not result in a mode transition. The guard function should have the signature `guard(time, state, *inputs, **parameters) -> float` and the reset map should have the signature of an unrestricted update `reset_map(time, state, *inputs, **parameters) -> state` Parameters: | Name | Type | Description | Default | | ---------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------- | | `guard` | `Callable` | The guard function which triggers updates on zero crossing. | *required* | | `reset_map` | `Callable` | The reset map which is applied when the event is triggered. If None (default), no reset is applied. | `None` | | `start_mode` | `int` | The mode or stage of the system in which the guard will be actively monitored. If None (default), the event will always be active. | `None` | | `end_mode` | `int` | The mode or stage of the system to which the system will transition when the event is triggered. If start_mode is None, this is ignored. Otherwise it must be specified, though it can be the same as start_mode. | `None` | | `direction` | `str` | The direction of the zero crossing. Options are "crosses_zero" (default), "positive_then_non_positive", "negative_then_non_negative", and "edge_detection". All except edge detection operate on continuous signals; edge detection operates on boolean signals and looks for a jump from False to True or vice versa. | `'crosses_zero'` | | `terminal` | `bool` | If True, the event will halt simulation if and when the zero-crossing occurs. If this event is triggered the reset map will still be applied as usual prior to termination. Defaults to False. | `False` | | `name` | `str` | The name of the event. Defaults to None. | `None` | | `enable_tracing` | `bool` | If True, enable tracing for this event. Defaults to None. | `None` | Notes By default the system state does not have a "mode" component, so in order to declare "state transitions" with non-null start and end modes, the user must first call `declare_default_mode` to set the default mode to be some integer (initial condition for the system). #### `initialize(**parameters)` Hook for initializing a system. Called during context creation. If the parameters are instances of Parameter, they will be resolved. If implemented, the function signature should contain all the declared parameters. This function should not be called directly. It will be called implicitly after **init** with the resolved parameters. #### `reset_default_values(**dynamic_parameters)` This function is used to reset default values for continuous/discrete states, ports and mode based on dynamic parameters. It is called in `create_state()` and used to reset states in ensemble sims and optimization with the context method `with_new_state()`. Note that dtypes and shapes can't be changed after initialization because the diagram may already have been jax-compiled. Only values may change. #### `with_parameter(name, value)` Return a copy of this system with one dynamic parameter replaced. The returned system is a new instance. The original is unchanged. Parameters: | Name | Type | Description | Default | | ------- | ----- | ---------------------------------------------------------- | ---------- | | `name` | `str` | Parameter name (must exist as a dynamic parameter). | *required* | | `value` | | New value (typically a JAX array for jax.grad / jax.vmap). | *required* | Raises: | Type | Description | | ----------- | ----------------------------------- | | `KeyError` | If name is not a dynamic parameter. | | `TypeError` | If name is a static parameter. | #### `wrap_callback(callback, collect_inputs=True)` Wrap an update function to unpack local variables and block inputs. The callback should have the signature `callback(time, state, *inputs, **params) -> result` and will be wrapped to have the signature `callback(context) -> result`, as expected by the event handling logic. This is used internally for declaration methods like `declare_periodic_update` so that users can write more intuitive block-level update functions without worrying about the "context", and have them automatically wrapped to have the right interface. It can also be called directly by users to wrap their own update functions, for example to create a callback function for `declare_output_port`. The context and state are strictly immutable, so the callback should not attempt to change any values in the context or state. Even in cases where it is impossible to *enforce* this (e.g. a state component is a list, which is always mutable in Python), the callback should be careful to avoid direct modification of the context or state, which may lead to unexpected behavior or JAX tracer errors. Parameters: | Name | Type | Description | Default | | ---------------- | ---------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `callback` | `Callable` | The (pure) function to be wrapped. See above for expected signature. | *required* | | `collect_inputs` | `bool` | If True, the callback will eval input ports to gather input values. Normally this should be True, but it can be set to False if the return value depends only on the state but not inputs, for instance. This helps reduce the number of expressions that need to be JIT compiled. Can also be specified as a list of integer port indices. Default is True (collect all inputs). | `True` | Returns: | Name | Type | Description | | ---------- | ---------- | ----------------------------------------------------------------- | | `Callable` | `Callable` | The wrapped function, with signature callback(context) -> result. | ### `Parameter` #### `__deepcopy__(memo)` Copy fields and re-run post-init so :class:`ParameterCache` bookkeeping matches. #### `unwrap(value)` Get the underlying value of raw arrays and Parameter objects alike. #### `value_as_api_param(allow_param_name=True, allow_string_literal=True)` Returns an API-compatible expression[1] that defines this parameter What we return depends on the caller's context, since it depends on whether we are serializing for a model, submodel or block parameter. The boolean is the value of 'is_string' (means "string literal" or "do not call eval"). [1] The returned string can be serialized to JSON, but it is not an already escaped JSON string! Parameters: | Name | Type | Description | Default | | ---------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `allow_param_name` | | Set to false for (sub)model parameters. Optional. If true, and the value is defined by a name, just the name will be returned. | `True` | | `allow_string_literal` | | Set to false for (sub)model parameters. Optional. If true, and the value is a string, then the string will be returned and 'is_string' will be returned as True. | `True` | ### `ParameterCache` Global parameter value cache used by all :class:`Parameter` instances. Thread safety All public methods are protected by a class-level reentrant lock (`threading.RLock`). Using an `RLock` rather than a plain `Lock` is necessary because `__compute__` may call `param.get()` recursively (for compound parameter expressions), which would deadlock under a non-reentrant lock held by the outer `get()` call. Concurrent simulations in separate threads sharing the same `Parameter` objects are serialised correctly. However, mutating a parameter from one thread while another thread is actively simulating with it is not recommended — the lock ensures the state remains consistent, but the simulation semantics of mid-run mutation are undefined. #### `epoch()` Monotonic version of all parameter values; changes on any mutation. #### `print_dependents(param, indent=0)` Prints the dependents tree of a parameter ### `RuntimeVariantSubsystem` Bases: `LeafSystem` Switch between pre-built submodel choices via a discrete selector input. This is the runtime counterpart to `select_variant` / `Variant`. Unlike the build-time selector (which never instantiates the unselected branches), `RuntimeVariantSubsystem` builds *every* choice and routes the selected branch's output through. The selector is a normal input port, so it can be driven by any discrete control signal in the diagram and the active branch follows at simulate time. This is the runtime-controlled variant pattern, as opposed to the label-mode build-time variant. Implementation: the block stacks all branches' outputs along a new leading axis and picks out the selected slice with integer indexing. This is the same mechanism used by `MultiPortSwitch` (T-118), reused here at the framework level so it does not pull a library dependency. ##### Contract — "all branches integrated each step" Because the underlying `stack` traces every branch, every choice's submodel runs on every step and sees the same input trajectory. The consequences: - Pure (memoryless) branches behave exactly as you'd expect: only the selected branch's output is exposed; gradients w.r.t. the active branch's parameters are non-zero, and gradients w.r.t. the others are zero (matching `MultiPortSwitch`'s data-input semantics). - The selector is non-differentiable (`round` + `clip` zero out its gradient), as expected for a control signal. - If a branch holds internal discrete state (e.g. a hold latch) the caller is responsible for supplying that state. `RuntimeVariantSubsystem` itself is stateless; if you need stateful sub-Diagrams, hoist the state out, or build the runtime switch by composing `MultiPortSwitch` (T-118) with N pre-built sub-diagrams in a parent `DiagramBuilder`. All branches must return outputs that are broadcast-compatible (the stack op requires a common shape/dtype after broadcasting). Parameters: | Name | Type | Description | Default | | ---------------- | ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `choices` | | Either a sequence of submodel callables [f0, f1, ..., f\_{N-1}] or a mapping {int: callable} (integer keys must be 0..N-1). Each callable has signature f(\*inputs) -> output and must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of user inputs forwarded to every branch. Input port 0 is always the selector; ports 1..n_inputs are the user inputs. Defaults to 1. | `1` | | `default_choice` | `int` | Index of the choice used as the default. Stored for introspection / documentation; the runtime selector value still controls which branch is exposed each step. Defaults to 0. | `0` | | `name` | | Optional block name. | *required* | Input ports (0) selector — scalar integer-valued signal in `[0, N-1]`. Floating values are rounded and clipped. (1..n_inputs) user inputs forwarded to every branch. Output ports (0) The selected branch's output. Raises: | Type | Description | | -------------- | ---------------------------------------------------------------------------------------- | | `VariantError` | If choices is empty / non-callable / has bad keys, or if default_choice is out of range. | #### `default_choice` Default choice index (documentary; runtime selector still rules). #### `n_choices` Number of variant choices held by this block. ### `ShapeMismatchError` Bases: `StaticError` Block parameters or input/outputs have mismatched shapes. ### `StaticError` Bases: `JaxonomyError` Wraps a Python exception to record the offending block id. The original exception is found in the '**cause**' field. See jaxonomy.framework.context_factory.\_check_types for use. This is called 'static' (as opposed to say 'runtime') meaning this is for wrapping errors detected prior to running a simulation. ### `SystemBase` Basic building block for simulation in jaxonomy. NOTE: Type hints in SystemBase indicate the union between what would be returned by a LeafSystem and a Diagram. See type hints of the subclasses for the specific argument and return types. #### `dependency_graph` Retrieve (or create if necessary) the dependency graph for this system. #### `has_dirty_static_parameters` Check if any static parameters have been modified. #### `has_feedthrough_side_effects` Check if the system includes any feedthrough calls to `io_callback`. #### `has_mass_matrix` Returns True if any component of the system has a nontrivial mass matrix. #### `has_ode_side_effects` Check if the ODE RHS for the system includes any calls to `io_callback`. #### `mass_matrix` Mass matrix for this system. Returns PyTree-structured data where each leaf is an (n, n) array. This is used for implicit integration methods (currently only BDF). #### `name_path` Get the human-readable path to this system. None if some names are not set. #### `name_path_str` Get the human-readable path to this system as a string. #### `ports` Dictionary of all ports in this system, indexed by name #### `root` Get the root system of the current system. #### `sorted_callbacks` Sort and return the callbacks for this system. #### `ui_id_path` Get the uuid node path to this system. None if some IDs are not set. #### `__deepcopy__(memo)` Deep-copy while keeping partially constructed copies hashable. Subsystems reference themselves via callbacks; the default deepcopy order can call :meth:`__hash__` (via dict/set operations) before `system_id` exists on the copy. Assign a new `system_id` immediately after memo registration. `_dependency_graph` is *not* copied — it is reset to `None` on the copy and rebuilt by the next `create_context`. It is a derived cache whose `DependencyTracker.prerequisites` form a linked chain spanning the whole signal path, so deep-copying it recurses once per block: a serial diagram of a few hundred blocks overflowed the interpreter stack with `RecursionError` on any copy of a *warm* system (one that had already built a context). Skipping it makes deepcopy nesting constant in block count instead of linear. #### `check_types(context, error_collector=None)` Perform any system-specific static analysis. #### `collect_inputs(context, port_indices=None)` Collect all current inputs for this system. Parameters: | Name | Type | Description | Default | | -------------- | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | root context for this system | *required* | | `port_indices` | `List[int]` | list of input port indices to collect. If None (default), will return values from all ports. Otherwise will return a list of length(num_input_ports), where the values are None for ports not in the list. | `None` | Returns: | Type | Description | | ------------- | ----------------------------------------------- | | `List[Array]` | List\[Array\]: list of all current input values | #### `configure_output_port(port_index, callback, prerequisites_of_calc=None, default_value=None, event=None, cache_index=None)` Configure an output port of the system. See `declare_output_port` for a description of the arguments. Parameters: | Name | Type | Description | Default | | ------------ | ----- | ------------------------------------- | ---------- | | `port_index` | `int` | index of the output port to configure | *required* | Returns: | Type | Description | | ---- | ----------- | | | None | #### `context_factory()` Factory object for creating contexts for this system. Should not be called directly - use `system.create_context` instead. #### `create_context(**kwargs)` Create a new context for this system. The context will contain all variable information used in simulation/analysis/optimization, such as state and parameters. Returns: | Name | Type | Description | | ------------- | ------------- | --------------------------- | | `ContextBase` | `ContextBase` | new context for this system | #### `create_dependency_graph()` Create a dependency graph for this system. #### `declare_dynamic_parameter(name, default_value=None, shape=None, dtype=None, as_array=True)` Declare a numeric parameter for the system. Parameters are declared in the system and accessed through the context to maintain separation of data ownership. This method creates an entry in the system's dynamic_parameters, recording the name, default value, and dependency ticket for later reference. The default value will be used to initialize the context, so it will also serve as the initial value unless explicitly overridden. In the simplest cases, parameters could be stored as attributes of the LeafSystem, but declaring them has the advantage of moving the values to the context, allowing them to be traced by JAX rather than stored as static data. This means they can be differentiated, vmapped, or otherwise modified without re-compiling the simulation. Parameters: | Name | Type | Description | Default | | --------------- | ------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `name` | `str` | The name of the parameter. | *required* | | `default_value` | `Union[Array, Parameter]` | The default value of the parameter. Parameters are used primarily internally for serialization and should not normally need to be used directly when implementing LeafSystems. Defaults to None. | `None` | | `shape` | `ShapeLike` | The shape of the parameter. Defaults to None. | `None` | | `dtype` | `DTypeLike` | The data type of the parameter. Defaults to None. | `None` | | `as_array` | `bool` | If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification. | `True` | Raises: | Type | Description | | ---------------- | --------------------------------------------------------- | | `AssertionError` | If the parameter with the given name is already declared. | Notes (1) Only one of `shape` and `default_value` should be provided. If `default_value` is provided, it will be used as the initial value of the continuous state. If `shape` is provided, the initial value will be a zero vector of the given shape and specified dtype. #### `declare_input_port(name=None, prerequisites_of_calc=None, units=None)` Add an input port to the system. Returns the corresponding index into the system input_port_indices list Note that this is different from the callbacks index - typically it will make more sense to retrieve via system.input_ports[port_index], but Parameters: | Name | Type | Description | Default | | ----------------------- | ------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `name` | `str` | name of the new port. Defaults to None, which will use the default naming scheme for the system (e.g. "u_0") | `None` | | `prerequisites_of_calc` | `List[DependencyTicket]` | list of dependencies for the callback function. Defaults to None. | `None` | | `units` | `Unit` | physical unit of the signal carried on this port (T-104 phase 1). Default None is treated as dimensionless — existing diagrams that never declare a unit continue to connect to anything. | `None` | Returns: | Name | Type | Description | | ----- | ----- | --------------------------------------------------- | | `int` | `int` | port index of the newly created port in input_ports | #### `declare_output_port(callback, name=None, prerequisites_of_calc=None, default_value=None, event=None, cache_index=None, units=None)` Add an output port to the system. This output port could represent any function of the context available to the system, so a callback function is required. This function should have the form `callback(context: ContextBase) -> Array` SystemBase implementations have some specific convenience wrappers, e.g.: `LeafSystem.declare_continuous_state_output` `Diagram.export_output` Common cases are: - Feedthrough blocks: gather inputs and return some function of the inputs (e.g. a gain) - Stateful blocks: use LeafSystem.declare\_(...)\_state_output_port to return the value of a particular state - Diagrams: create and export a diagram-level port to the parent system using the callback function associated with the system-level port Returns the corresponding index into the system output_port_indices list Note that this is different from the callbacks index - typically it will make more sense to retrieve via system.output_ports[port_index]. Parameters: | Name | Type | Description | Default | | ----------------------- | ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `callback` | `Callable` | computes the value of the output port given the root context. | *required* | | `name` | `str` | name of the new port. Defaults to None, which will use the default naming scheme for the system (e.g. "y_0") | `None` | | `prerequisites_of_calc` | `List[DependencyTicket]` | list of dependencies for the callback function. Defaults to None, which will use the default dependencies for the system (all sources). This may conservatively flag the system as having algebraic loops, so it is better to be specific here when possible. This is done automatically in the wrapper functions like LeafSystem.declare\_(...)\_output_port | `None` | | `default_value` | `Array` | A default array-like value used to seed the context and perform type inference, when this is known up front. Defaults to None, which will use information propagation through the graph along with type promotion to determine an appropriate value. | `None` | | `event` | `DiscreteUpdateEvent` | A discrete update event associated with this output port that will periodically refresh the value that will be returned by the callback function. This makes the port act as a sample-and-hold rather than a direct function evaluation. | `None` | | `cache_index` | `int` | Index into the cache state component corresponding to the output port result, if the output port is of periodically-updated sample-and-hold type. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------ | | `int` | `int` | port index of the newly created port | #### `declare_static_parameter(name, value)` Declare a single static parameter for the system. This is a convenience function for declaring a single static parameter. Parameters: | Name | Type | Description | Default | | ------- | ------------------------- | ---------------------- | ---------- | | `name` | `str` | name of the parameter | *required* | | `value` | `Union[Array, Parameter]` | value of the parameter | *required* | #### `declare_static_parameters(**params)` Declare a set of static parameters for the system. These parameters are not JAX-traceable and therefore can't be optimized. Examples of static parameters include booleans, strings, parameters used in shapes, etc. The args should be a dict of name-value pairs, where the values are either strings, bool, arrays, or Parameters. Typical usage: ``` class MyBlock(LeafSystem): def __init__(self, param1=True, param2=1.0): super().__init__() self.declare_static_parameters(param1=param1, param2=param2) ``` #### `dependency_graph_factory()` Factory object for creating dependency graphs for this system. Should not be called directly - use `system.create_dependency_graph` instead. #### `determine_active_guards(context)` Determine active guards for zero-crossing events. This method is responsible for evaluating and determining which zero-crossing events are active based on the current system mode and other conditions. This can be overridden to flag active/inactive guards on a block-specific basis, for instance in a StateMachine-type block. By default all guards are marked active at this point unless the zero-crossing event was declared with a non-default `start_mode`, in which case the guard is activated conditionally on the current mode. For example, in a system with finite state transitions, where a transition from mode A to mode B is triggered by a guard function g_AB and the inverse transition is triggered by a guard function g_BA, this function would activate g_AB if the system is in mode A and g_BA if the system is in mode B. The other guard function would be inactive. If the zero-crossing event is not associated with a starting mode, it is considered to be always active. Parameters: | Name | Type | Description | Default | | --------- | ------------- | ------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | The root context containing the overall state and parameters. | *required* | Returns: | Name | Type | Description | | ----------------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------------- | | `EventCollection` | `EventCollection` | A collection of zero-crossing events with active/inactive status updated based on the current system mode and other conditions. | #### `eval_input(context, port_index=0)` Get the input for a given port. This works by evaluating the callback function associated with the port, which will "pull" the upstream output port values. Parameters: | Name | Type | Description | Default | | ------------ | ------------- | ------------------------------------------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | root context for this system | *required* | | `port_index` | `int` | index into self.input_ports, for example the value returned by declare_input_port. Defaults to 0. | `0` | Returns: | Name | Type | Description | | ------- | ------- | -------------------- | | `Array` | `Array` | current input values | #### `eval_time_derivatives(context)` Evaluate the continuous time derivatives for this system. Given the *root* context, evaluate the continuous time derivatives, which must have the same PyTree structure as the continuous state. In principle, this can be overridden by custom implementations, but in general it is preferable to declare continuous states for LeafSystems using `declare_continuous_state`, which accepts a callback function that will be used to compute the derivatives. For Diagrams, the time derivatives are computed automatically using the callback functions for all child systems with continuous state. Parameters: | Name | Type | Description | Default | | --------- | ------------- | --------------------------- | ---------- | | `context` | `ContextBase` | root context of this system | *required* | Returns: | Name | Type | Description | | ---------------- | ---------------- | ------------------------------------------------------------------------------------------- | | `StateComponent` | `StateComponent` | Continuous time derivatives for this system, or None if the system has no continuous state. | #### `eval_zero_crossing_updates(context, events)` Evaluate reset maps associated with zero-crossing events. Parameters: | Name | Type | Description | Default | | --------- | ----------------- | -------------------------------------------------------------------------------------------------------- | ---------- | | `context` | `ContextBase` | The context for the system, containing the current state and parameters. | *required* | | `events` | `EventCollection` | The collection of events to be evaluated (for example zero-crossing or periodic events for this system). | *required* | Returns: | Name | Type | Description | | ------- | ------- | -------------------------------------------- | | `State` | `State` | The complete state with all updates applied. | Notes (1) Following the Drake definition, "unrestricted" updates are allowed to modify any component of the state: continuous, discrete, or mode. These updates are evaluated in the order in which they were declared, so it is *possible* (but should be strictly avoided) for multiple events to modify the same state component at the same time. Each update computes its results given the *current* state of the system (the "minus" values) and returns the *updated* state (the "plus" values). The update functions cannot access any information about the "plus" values of its own state or the state of any other block. This could change in the future but for now it ensures consistency with Drake's discrete semantices: More specifically, since all unrestricted updates can modify the entire state, any time there are multiple unrestricted updates, the resulting states are ALWAYS in conflict. For example, suppose a system has two unrestricted updates, `event1` and `event2`. At time t_n, `event1` is active and `event2` is inactive. First, `event1` is evaluated, and the state is updated. Then `event2` is evaluated, but the state is not updated. Which one is valid? Obviously, the `event1` return is valid, but how do we communicate this to JAX? The situation is more complicated if both `event1` and `event2` happen to be active. In this case the states have to be "merged" somehow. In the worst case, these two will modify the same components of the state in different ways. The implementation updates the state in a local copy of the context (since both are immutable). This allows multiple unrestricted updates, but leaves open the possibility of multiple active updates modifying the state in conflicting ways. This should strictly be avoided by the implementer of the LeafSystem. If it is at all unclear how to do this, it may be better to split the system into multiple blocks to be safe. (2) The events are evaluated conditionally on being marked "active" (indicating that their guard function triggered), so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns. #### `get_feedthrough()` Determine pairs of direct feedthrough ports for this system. By default, the algorithm relies on the dependency tracking system to determine feedthrough, but this can be overridden by implementing this method directly in a subclass, for instance if the automatic dependency tracking is too conservative in determining feedthrough. Returns: | Type | Description | | ----------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `List[Tuple[int, int]]` | List\[Tuple[int, int]\]: A list of tuples (u, v) indicating that output port v has a direct dependency on input port u, resulting in a feedthrough path in the system. The indices u and v correspond to the indices of the input and output ports in the system's input and output port lists. | #### `get_input_port(name)` Retrieve a specific input port by name. #### `get_output_port(name)` Retrieve a specific output port by name. #### `get_parameter(name)` Get a parameter value by name. Checks dynamic parameters first, then static parameters. Values are returned as concrete array-like / Python scalars via :meth:`Parameter.unwrap`. Parameters: | Name | Type | Description | Default | | ------ | ----- | -------------------------------------------------- | ---------- | | `name` | `str` | Parameter name on this system (not a dotted path). | *required* | Raises: | Type | Description | | ---------- | -------------------------------------------------------- | | `KeyError` | If name is not found; the message lists available names. | #### `handle_discrete_update(events, context, *, topological_order=False)` Compute and apply active discrete updates. Given the *root* context, evaluate the discrete updates, which must have the same PyTree structure as the discrete states of this system. This should be a pure function, so that it does not modify any aspect of the context in-place (even though it is difficult to strictly prevent this in Python). This will evaluate the set of events that result from declaring state or output update events on systems using `LeafSystem.declare_periodic_update` and `LeafSystem.declare_output_port` with an associated periodic update rate. This is intended for internal use by the simulator and should not normally need to be invoked directly by users. Events are evaluated conditionally on being marked "active", so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns. For a discrete system updating at a particular rate, the update rule for a particular block is: ``` x[n+1] = f(t[n], x[n], u[n]) y[n] = g(t[n], x[n], u[n]) ``` Additionally, the value y[n] is held constant until the next update from the point of view of other continuous-time or asynchronous discrete-time blocks. Because each output `y[n]` may in general depend on the input `u[n]` evaluated *at the same time*, the composite discrete update function represents a system of equations. However, since algebraic loops are prohibited, the events can be ordered and executed sequentially to ensure that the updates are applied in the correct order. This is implemented in `SystemBase.sorted_callbacks`. Multirate systems work in the same way, except that the events are evaluated conditionally on whether the current time corresponds to an update time for each event. Parameters: | Name | Type | Description | Default | | --------- | ----------------- | ------------------------------------ | ---------- | | `events` | `EventCollection` | collection of discrete update events | *required* | | `context` | `ContextBase` | root context for this system | *required* | Returns: | Name | Type | Description | | ------------- | ------------- | --------------------------------------------------------------------- | | `ContextBase` | `ContextBase` | updated context with all active updates applied to the discrete state | #### `handle_zero_crossings(events, context)` Compute and apply active zero-crossing events. This is intended for internal use by the simulator and should not normally need to be invoked directly by users. Events are evaluated conditionally on being marked "active", so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns. Parameters: | Name | Type | Description | Default | | --------- | ----------------- | ---------------------------------- | ---------- | | `events` | `EventCollection` | collection of zero-crossing events | *required* | | `context` | `ContextBase` | root context for this system | *required* | Returns: | Name | Type | Description | | ------------- | ------------- | ------------------------------------------------------------ | | `ContextBase` | `ContextBase` | updated context with all active zero-crossing events applied | #### `initialize_static_data(context)` Initialize any context data that has to be done after context creation. Use this to define custom auxiliary data or type inference that doesn't get traced by JAX. See the `ZeroOrderHold` implementation for an example. Since this is only applied during context initialization, it is allowed to modify the context directly (or the system itself). Typically this should not be called outside of the ContextFactory. Parameters: | Name | Type | Description | Default | | --------- | ------------- | ---------------------------------------------- | ---------- | | `context` | `ContextBase` | partially initialized context for this system. | *required* | #### `list_parameters()` Return all parameters as a flat `{name: value}` mapping. Dynamic parameters override static ones when names collide. Values are unwrapped the same way as :meth:`get_parameter`. #### `post_simulation_finalize()` Finalize the system after simulation has completed. This is only intended for special blocks that need to clean up resources and close files. #### `pprint(output=print, fancy=True)` Pretty-print the system and its hierarchy. #### `recompute_port_cache(context)` Recompute all output ports and return them in a dictionary. ### `SystemCallback` A function associated with a system that has has specified dependencies. This can include port update rules, discrete update functions, the right-hand-side of an ODE, etc. Storing these functions as SystemCallbacks allows the system, or a Diagram containing the system, to track dependencies across the system or diagram. Attributes: | Name | Type | Description | | ----------------------- | ------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `system` | `SystemBase` | The system that owns this callback. | | `ticket` | `int` | The dependency ticket associated with this callback. See DependencyTicket for built-in tickets. If None, a new ticket will be generated. | | `name` | `str` | A short description of this callback function. | | `prerequisites_of_calc` | `List[DependencyTicket]` | Direct prerequisites of the computation, used for dependency tracking. These might be built-in tickets or tickets associated with other SystemCallbacks. | | `default_value` | `Array` | A dummy value of the same shape/dtype as the result, if known. If None, any type checking will rely on propagating upstream information via the callback. | | `callback_index` | `int` | The index of this function in the system's list of associated callbacks. | | `event` | `Event` | Optionally, the callback function may be associated with an event. If so, the associated trackers can be used to sort event execution order in addition to the regular callback execution order. For example, if an OutputPort is of sample-and-hold type, then this will be the event that periodically updates the output value. Default is None. | #### `calc(root_context)` Unconditionally evaluate the callback function. This does not check the cache status, but will always recompute the value. Typically `eval` should be preferred to `calc` to take advantage of caching where possible. Parameters: | Name | Type | Description | Default | | -------------- | ------------- | ----------------------------------------- | ---------- | | `root_context` | `ContextBase` | The root context used for the evaluation. | *required* | Returns: | Type | Description | | ------- | --------------------------------------------------------------- | | `Array` | The calculated value from the callback, expected to be a Array. | #### `eval(root_context)` Evaluate the callback function and return the calculated value. Within a single top-level call, repeated evaluations of the same callback against the same context are memoized (see `_eval_memo` above) — this keeps eager evaluation of diagrams with fan-out / reconvergence linear in graph size instead of exponential in composition depth. Nothing is cached across top-level calls. Parameters: | Name | Type | Description | Default | | -------------- | ------------- | ----------------------------------------- | ---------- | | `root_context` | `ContextBase` | The root context used for the evaluation. | *required* | Returns: | Type | Description | | ------- | --------------------------------------------------------------- | | `Array` | The calculated value from the callback, expected to be a Array. | ### `TriggerEdge` Allowed string values for `TriggeredSubsystem.edge`. ### `TriggeredSubsystem` Bases: `LeafSystem` Container block: latch the submodel output on edge transitions (the child still RUNS every step — only the *output* is gated). Important: this does **not** skip execution of the submodel on non-triggered steps. The submodel is evaluated on every step so its inputs participate in the JAX trace; the trigger only controls whether a fresh result is *latched* into the held output. If you need to actually skip computation between triggers, gate it yourself with `jax.lax.cond` at the application level. Phase-1 implementation runs the submodel on every step (so the inputs participate in the trace) but only *latches* a new output on an edge transition of the trigger signal. Between transitions the output holds the most recently latched value. The trigger signal is sampled at `sample_period`. Edges are detected by comparing the current trigger sample against the previously-stored sample held in discrete state. This is *not* the eventual zero-crossing-driven `TriggeredSubsystem` described in the T-120 architecture notes (that requires hooking into the continuous-time event detector); but it is functionally correct for any sample-rate use case and matches the behaviour documented in the test fixtures. Parameters: | Name | Type | Description | Default | | --------------- | ---------------------------------------- | --------------------------------------------------------------------------------------------------- | ---------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output taking the non-trigger user inputs. Must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of user inputs (NOT counting the trigger). | `1` | | `edge` | `Literal['rising', 'falling', 'either']` | "rising" (low→high), "falling" (high→low) or "either". | `RISING` | | `sample_period` | `float` | Period (seconds) at which the trigger signal is sampled and the latch is updated. Must be positive. | `0.0` | | `initial_value` | | Latched output value before any edge has been detected. Defines output shape/dtype. | `0.0` | | `name` | | Optional block name. | *required* | Limitations (phase 1): - Trigger detection runs on the periodic sample grid, not on continuous-time zero crossings. Trigger pulses shorter than `sample_period` may be missed. - The latch is a single discrete state; the submodel must produce a single output array. - The submodel runs on every output evaluation; only the *output* is gated. Users who need to skip computation on non-triggered steps should use `jax.lax.cond` at the application level. ### `Unit` Immutable SI dimensional value. Units are compared by their dimension exponents and scale factor. The optional `name` is informational (used in error messages) and is not part of equality. #### `from_dict(data)` Construct a :class:`Unit` from a dict produced by :meth:`to_dict`. Missing keys take their dataclass defaults so the empty dict `{}` round-trips to `Unit()`. #### `from_json(json_str)` Inverse of :meth:`to_json`. Raises: | Type | Description | | ------------ | --------------------------------- | | `ValueError` | If json_str is not a JSON object. | #### `same_dimension_as(other)` True if exponents match (ignoring scale). Phase 1 doesn't use this for the connect check (which is strict-equal), but it's part of the public surface so Phase 2 can layer scalar-conversion warnings on top. #### `summary()` Return a human-readable one-line summary of this Unit. Designed for `print()` / display contexts where `repr(unit)` is too terse. Includes the dimension exponents (with SI labels), scale, offset, currency exponents, and the `physical_quantity` tag when set. #### `to_dict()` Return a JSON-friendly dict representation of this Unit. Round-trips losslessly via :meth:`from_dict`. Keys are stable across versions; new optional fields are always added with defaults so older serialised forms continue to load. The default value for any field is omitted from the output for compactness — every legacy `Unit()` instance serialises to `{}`. #### `to_json(*, indent=None)` Serialise :meth:`to_dict` via :func:`json.dumps`. Parameters: | Name | Type | Description | Default | | -------- | ----- | ----------- | -------------------------------------------------------------------------------------------- | | `indent` | \`int | None\` | Optional JSON indent (default None for compact form; pass an int for pretty-printed output). | ### `UnitMismatchError` Bases: `StaticError` Raised at diagram build time when two connected ports have incompatible units. Attributes are populated through :class:`StaticError` so the regular `ErrorCollector` / system-locator machinery still works. ### `Variant` A frozen description of N variant choices for build-time selection. Parameters: | Name | Type | Description | Default | | --------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `choices` | `Mapping[str, Callable[[], SystemBase]]` | Mapping from choice name to a zero-argument builder callable. Each callable, when invoked, must return a SystemBase (typically a fully-built Diagram). Unselected callables are never invoked. | *required* | | `default` | `str` | Name of the choice to use when select_variant is called without an explicit name. Required (no implicit "first choice") so that adding a new variant later doesn't silently change the default. Must be a key in choices. | *required* | | `name` | `Optional[str]` | Optional human-readable label for diagnostics / logging. Does not affect resolution. | `None` | Raises: | Type | Description | | -------------- | ------------------------------------------------------------------------------ | | `VariantError` | If choices is empty, default is not in choices, or any choice is not callable. | #### `choice_names` Stable tuple of available choice names (for introspection / CLI). ### `VariantError` Bases: `ValueError` Raised when a variant configuration is invalid or a selection is bad. ### `WhileLoop` Bases: `LeafSystem` Container block: run `body_fn` until `cond_fn` is False. `WhileLoop` wraps :func:`jax.lax.while_loop` with a built-in iteration counter that caps execution at `max_iter` to guarantee termination under jit. The block declares an input port for the *initial carry value* (port 0) plus `n_inputs` additional ports for upstream signals that the loop body / condition can consume. A single output port returns the carry after the loop exits (either because `cond_fn` returned False, or because `max_iter` was hit). Parameters: | Name | Type | Description | Default | | ---------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `body_fn` | `Callable` | Callable. Either carry -> carry (legacy) or (carry, \*inputs) -> carry when n_inputs > 0. Must be JAX-traceable. The signature is detected via :func:inspect.signature; functions that accept more than one positional argument (or \*args) receive all upstream input values. Callables that only need a subset should either accept the rest as throwaway positional args, or use \*args and index into it. | *required* | | `cond_fn` | `Callable` | Callable. Either carry -> bool (legacy) or (carry, \*inputs) -> bool. Loop continues while this returns True. Signature detection mirrors body_fn; the same all-inputs-or-none rule applies. | *required* | | `max_iter` | `int` | Positive integer cap on iterations. Required to keep traces bounded under jit. Defaults to 1000. | `1000` | | `n_inputs` | `int` | Number of additional upstream input ports (default 0). When n_inputs > 0 the block exposes ports u_0..u\_{n_inputs-1} after the carry_init port. The current values of these inputs are passed to cond_fn / body_fn (if they accept them) on every iteration — so the condition can compare the carry against a live upstream signal (e.g. "iterate until the input exceeds a threshold"). | `0` | | `name` | | Optional block name. | *required* | Differentiability `jax.grad` flows through the carry as long as `body_fn` and `cond_fn` are pure. The number of iterations is data-dependent and not differentiable; `lax.while_loop` is itself non-differentiable in reverse mode (use `jax.jvp` for forward mode, or refactor with :func:`jax.lax.scan` if you need a reverse-mode-friendly bounded loop). Notes - On hitting `max_iter` the loop exits silently. Users who want a runtime warning should `jax.debug.callback` from `body_fn` or test the post-loop carry. - The carry pytree structure must be invariant across iterations (a `lax.while_loop` requirement). - The condition is re-evaluated against the *current* upstream input values inside the loop trace — the inputs are captured once at output-evaluation time and held constant for the duration of the loop (the diagram doesn't re-tick during a single major step). ### `ZeroCrossingEvent` Bases: `Event` An event that triggers when a specified "guard" function crosses zero. The event is triggered when the guard function crosses zero in the specified direction. In addition to the guard callback, the event also has a "reset map" which is called when the event is triggered. The reset map may update any state component in the system. The event can also be defined as "terminal", which means that the simulation will terminate when the event is triggered. (TODO: Does the reset map still happen?) The "direction" of the zero-crossing is one of the following: - "none": Never trigger the event (can be useful for debugging) - "positive_then_non_positive": Trigger when the guard goes from positive to non-positive - "negative_then_non_negative": Trigger when the guard goes from negative to non-negative - "crosses_zero": Trigger when the guard crosses zero in either direction - "edge_detection": Trigger when the guard changes value Notes This class should typically not need to be used directly by users. Instead, declare the guard function and reset map on a LeafSystem using the `declare_zero_crossing` method. The event will then be auto-generated for simulation. #### `handle(context)` Conditionally compute the result of the zero crossing callback If the zero crossing is marked "inactive" via its event data attribute, the passthrough callback will be called instead of the update callback. Otherwise, the update callback will be called. The return types of both callbacks must match, but the specific type will depend on the kind of event. #### `should_trigger()` Determine if the event should trigger based on the stored guard values. ### `ZeroCrossingTriggeredSubsystem` Bases: `LeafSystem` Container block: latch the submodel output at zero-crossings. Like :class:`TriggeredSubsystem`, but uses the framework's continuous zero-crossing detector rather than a periodic sample grid. The submodel fires *exactly* when the trigger signal crosses zero in the configured direction — this gives sub-sample-period precision for the latched event time, which is the property normally expected from a triggered subsystem driven by a continuous signal. Wiring matches :class:`TriggeredSubsystem`: - Input port 0 is the trigger signal (a continuous scalar; the block monitors its sign). - Input ports 1..n_inputs are the submodel inputs. - The single output port returns the most recently latched submodel output (initialized to `initial_value`). Parameters: | Name | Type | Description | Default | | --------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------- | ---------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output taking the non-trigger user inputs. Must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of user inputs (NOT counting the trigger). | `1` | | `edge` | `Literal['rising', 'falling', 'either']` | "rising" (low→high zero crossing of the trigger signal), "falling" (high→low), or "either". | `RISING` | | `initial_value` | | Latched output value before the first crossing fires. Also defines the output shape/dtype. | `0.0` | | `name` | | Optional block name. | *required* | Differentiability `jax.grad` flows through the submodel inputs along the path through the latch (so when the latched value depends on a differentiable input, the gradient propagates). The trigger signal itself is consumed by the zero-crossing event detector; the gradient through the discontinuity at the firing instant is zero by design (the latched value is constant between crossings). Notes - The framework localizes the zero crossing to within the integrator's tolerance, so the latched output reflects the submodel inputs *at the crossing instant*, not at the next periodic sample. Compare with the phase-1 :class:`TriggeredSubsystem`, which can only resolve the edge to the nearest `sample_period`. - The latch is a single discrete-state component; the submodel must produce a single output array of fixed shape. - This is a leaf block (no nested mode machinery), so the `"hold between crossings"` semantics fall out naturally: the output port simply returns `state.discrete_state`. ### `ForEach(submodel, n, n_inputs=1, in_axes=None, name=None)` Container block: evaluate a submodel `n` times in parallel. `ForEach` is a block-diagram-vocabulary alias for the existing :class:`jaxonomy.library.ReplicatedFunction` (T-010). It exists so that users familiar with the `ForEach` block name can find it without paying a duplication tax: the implementation is exactly :class:`ReplicatedFunction` under the hood. Parameters: | Name | Type | Description | Default | | ---------- | ---------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output. Must be JAX-traceable. | *required* | | `n` | `int` | Number of replicas (the iteration count). | *required* | | `n_inputs` | `int` | Number of input ports the block declares. | `1` | | `in_axes` | | As in :func:jax.vmap / ReplicatedFunction: a length-n_inputs tuple of 0 (input is batched along the leading axis) or None (input is broadcast). Default is all-batched. | `None` | | `name` | \`str | None\` | Optional block name. | Returns: | Type | Description | | ---- | ------------------------------------------------------------ | | | A configured :class:ReplicatedFunction instance, ready to be | | | wired into a :class:DiagramBuilder. | ### `apply_variant_config(diagram, **overrides)` Return a copy of `diagram` with named variants reconfigured. Walks the diagram tree, finds every subsystem that was produced by :func:`select_variant` from a named `Variant`, and -- for each `override_name=choice` keyword -- replaces matching subsystems with a freshly-built copy from `select_variant(variant, name=choice)`. All other diagram structure (non-variant blocks, connections, exported ports) is preserved. The original diagram is not modified. Example:: ``` builder = DiagramBuilder() ctrl = select_variant(controller_variant, name="pid") # default plant = select_variant(plant_variant, name="lti") builder.add(ctrl) builder.add(plant) ... diagram = builder.build() # Reconfigure post-build: runtime_a = apply_variant_config(diagram, controller="pid", plant="lti") runtime_b = apply_variant_config(diagram, controller="lqr", plant="lti") ``` Parameters: | Name | Type | Description | Default | | ------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `diagram` | | A built Diagram (typically the output of DiagramBuilder.build). | *required* | | `**overrides` | | Map from a variant's name (the name= kwarg passed to :class:Variant) to the choice name to activate. Variants whose name is not mentioned in overrides keep their currently-active choice. | `{}` | Returns: | Type | Description | | ---- | ------------------------------------------------------------ | | | A new Diagram with the requested variant choices resolved. | | | If overrides is empty, returns a structurally identical deep | | | copy of diagram (the same default-off semantics as | | | meth:Diagram.with_parameters with no updates). | Raises: | Type | Description | | -------------- | --------------------------------------------------------------------------------------------------------------------------- | | `VariantError` | If an override name does not match any variant in the diagram, or if the requested choice is not in that variant's choices. | ### `are_units_compatible(src, dst)` Return True if a connection from `src` to `dst` should be allowed under the Phase-1 rules: - Either side being `None` (unset) is always OK. - If both sides are :class:`BusUnit`, compatible iff every shared field's :class:`Unit` is pair-wise compatible AND the field sets match. A `BusUnit` on one side and `None` on the other is always OK (default-off byte-equivalence with the no-units bus). - Otherwise, both sides must be plain :class:`Unit`; either being :data:`dimensionless` is OK, else units must be equal. Scalar conversion (Phase 2) is layered on top by :func:`assert_units_compatible_with_scale` — see there. ### `assert_unit_compatible(src, dst, *, src_label='source port', dst_label='destination port')` Raise :class:`UnitMismatchError` if the two units are not Phase-1-compatible. See :func:`are_units_compatible`. The labels are interpolated into the message so the caller (typically :meth:`DiagramBuilder.connect`) can name both ports. ### `clear_fx_rates()` Empty the FX rate table. Tests use this to keep their state isolated; production code should rarely need to call it. ### `convert_currency(value, from_unit, to_unit)` Convert a numeric `value` carried in `from_unit` to the equivalent value in `to_unit` using the current FX rate table. Self-conversion (same currency on both sides) is a no-op and returns the value unchanged. Cross-currency conversion looks up the rate via :func:`get_fx_rate` and multiplies; a missing rate raises :class:`KeyError`. Parameters: | Name | Type | Description | Default | | ----------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | | `value` | | Numeric value (Python scalar, NumPy array, JAX array). The helper only uses \*, so it composes transparently through jit / vmap / grad. | *required* | | `from_unit` | \`'Unit | str'\` | Source currency, either a :class:Unit (such as :data:usd) or a string code ("USD"). | | `to_unit` | \`'Unit | str'\` | Destination currency, ditto. | Returns: | Type | Description | | ---- | ---------------------------------------------------------- | | | value * rate where rate = get_fx_rate(from_unit, to_unit). | Raises: | Type | Description | | ------------------- | -------------------------------------------------------------------------------------------------- | | `UnitMismatchError` | if either argument carries non-currency dimensions (e.g. seconds), so the conversion is undefined. | | `KeyError` | if the relevant FX rate has not been registered. | Example: ``` >>> set_fx_rate("USD", "EUR", 0.92) >>> convert_currency(100.0, usd, eur) 92.0 ``` ### `derived_unit(name, symbol=None, components=None)` Define a new derived :class:`Unit` from existing components. This is a convenience constructor for users who want to spell a composite unit once (with a friendly name) rather than recomposing its base components at every port declaration site. The returned :class:`Unit` has the same `(dims, scale, offset)` as `components` — it is therefore equal (under `Unit.__eq__`) to any other unit with matching dimensions and scale — but carries a custom `name` for friendlier error messages and pprint output. Parameters: | Name | Type | Description | Default | | ------------ | ------- | -------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `name` | `str` | Long-form descriptive name (e.g. "my_torque"). Used only when symbol is omitted. | *required* | | `symbol` | \`str | None\` | Short-form printable symbol (e.g. "τ"). When provided, it overrides name as the Unit's display label. | | `components` | \`'Unit | None'\` | A :class:Unit expression describing the dimensions of the new unit (e.g. meter * newton). Must not be None and must have offset == 0 — affine units cannot be re-aliased this way. | Returns: | Type | Description | | -------- | ----------------------------------------------- | | `'Unit'` | A fresh :class:Unit with components.dims / | | `'Unit'` | components.scale / components.offset and a name | | `'Unit'` | set to symbol (when provided) or name. | Raises: | Type | Description | | ------------------- | --------------------------------------------------------------------- | | `TypeError` | if components is not a :class:Unit. | | `UnitMismatchError` | if components has a non-zero offset (affine units cannot be aliased). | Example: ``` >>> from jaxonomy.framework.units import ( ... derived_unit, meter, newton, ... ) >>> torque = derived_unit("torque", "N·m", meter * newton) >>> torque.dims == (1, 2, -2, 0, 0, 0, 0) True >>> torque == meter * newton True ``` ### `flatten_diagram(diagram)` Flatten a nested Diagram into a single-depth Diagram. All intermediate sub-Diagrams are dissolved. The resulting Diagram has: - nodes: all LeafSystem instances from the original tree - connection_map: remapped to only reference leaf-to-leaf connections - exported inputs/outputs: preserved (still reference the same leaf ports) Parameters: | Name | Type | Description | Default | | --------- | --------- | ----------------------------------------- | ---------- | | `diagram` | `Diagram` | The (possibly nested) Diagram to flatten. | *required* | Returns: | Type | Description | | --------- | ------------------------------------------------------------------ | | `Diagram` | A new single-depth Diagram with all original LeafSystems as direct | | `Diagram` | children and all connections resolved to the leaf level. | ### `get_active_variant(diagram, variant_name)` Return the currently-selected choice for the named variant. Parameters: | Name | Type | Description | Default | | -------------- | ----- | ---------------------------------------------------------------------------------------------- | ---------- | | `diagram` | | A built Diagram. | *required* | | `variant_name` | `str` | The human-readable label of the variant to look up (the name= kwarg passed to :class:Variant). | *required* | Returns: | Type | Description | | ---- | ------------------------------------------------------------ | | | The name of the active choice (a string in | | | Variant.choice_names), or None if no variant with the | | | given name is found in the diagram. The None sentinel lets | | | CLI / introspection code treat "no such variant" as a soft | | | miss; use :func:get_variant_choices if you want a hard error | | | for unknown names. | ### `get_fx_rate(from_currency, to_currency)` Return the previously-set FX rate from `from_currency` to `to_currency`. Self-rates are always `1.0` even when unset. Raises: | Type | Description | | ------------------- | ------------------------------------------------------------------------ | | `KeyError` | if no rate has been set for the requested pair AND the two codes differ. | | `UnitMismatchError` | if either argument is not a pure currency. | ### `get_variant_choices(diagram, variant_name)` Return the choice names of the named variant in `diagram`. Parameters: | Name | Type | Description | Default | | -------------- | ----- | ---------------------------------------------------------------------------------------------- | ---------- | | `diagram` | | A built Diagram. | *required* | | `variant_name` | `str` | The human-readable label of the variant to look up (the name= kwarg passed to :class:Variant). | *required* | Returns: | Type | Description | | ------- | -------------------------------------------------- | | `tuple` | A tuple of choice names (Variant.choice_names), in | | `tuple` | insertion order. | Raises: | Type | Description | | -------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `VariantError` | If no variant with the given name is found in the diagram. Anonymous Variants (built without name=) are never matched and so cannot be queried via this helper. | ### `leaf_connections(diagram)` Resolve the Diagram hierarchy's wiring to leaf-to-leaf connections. Same resolution :func:`flatten_diagram` performs, but returned as a plain list of `(input_locator, output_locator)` pairs without rebuilding a Diagram. Use this when you need the flattened *topology* while continuing to evaluate against the original hierarchy and its context — :func:`flatten_diagram` re-parents the leaf systems into a new Diagram, which invalidates contexts created from the original. A `LeafSystem` (or a Diagram with no connections) yields an empty list. ### `list_variants(diagram)` List every variant point found in a (possibly nested) diagram. Walks `diagram` recursively and returns a metadata triple for every subsystem that was produced by :func:`select_variant`. Each triple has the shape `(name, choice_names, active_choice)`: - `name` is the variant's human-readable label (`Variant.name`). `None` for anonymous Variants. - `choice_names` is the stable tuple of available choice names (`Variant.choice_names`). - `active_choice` is the name of the choice currently bound at this point in the diagram. Iteration order follows the diagram's tree-traversal order (parent before children, siblings in registration order). If the same `Variant` instance is reused at multiple points in the diagram, each occurrence yields its own entry — callers that want a deduped view should collapse on `name`. Parameters: | Name | Type | Description | Default | | --------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `diagram` | | A built Diagram (typically the output of DiagramBuilder.build) or any SystemBase. Passing a non-Diagram subsystem returns [] (no children to walk; a tagged-leaf-as-root case is not produced by the current API surface, but the helper degrades gracefully). | *required* | Returns: | Type | Description | | ------------- | ----------------------------------------------------------- | | `list[tuple]` | A list of (name, choice_names, active_choice) tuples. Empty | | `list[tuple]` | if diagram contains no variant points (default-off path). | ### `next_dependency_ticket()` Create a new unique dependency ticket using the next available value. ### `parameters(static=None, dynamic=None)` Decorator to apply to a system class to declare static or dynamic parameters. ### `select_variant(variant, name=None)` Resolve a `Variant` at build time and return the active sub-system. Only the chosen builder is invoked; the others are never called. This matches the "active variant only" code-generation behavior familiar from established block-diagram tools -- nothing about the unselected branches enters the JIT trace, the parameter pytree, or the diagram's registered-systems list. Parameters: | Name | Type | Description | Default | | --------- | --------------- | ----------------------------------------------------------------- | ---------- | | `variant` | `Variant` | The Variant to resolve. | *required* | | `name` | `Optional[str]` | Name of the choice to activate. If None, variant.default is used. | `None` | Returns: | Type | Description | | ------------ | ---------------------------------------------- | | `SystemBase` | The SystemBase returned by the chosen builder. | Raises: | Type | Description | | -------------- | ---------------------------------------------------------------------------------------------------------- | | `VariantError` | If name is not one of variant.choices, or if the chosen builder returns something that isn't a SystemBase. | ### `set_fx_rate(from_currency, to_currency, rate)` Record an FX rate so that one unit of `from_currency` equals `rate` units of `to_currency`. Both directions are written: setting USD→EUR at 0.92 simultaneously sets EUR→USD at `1.0 / 0.92` so round-trips are exact under the floating-point reciprocal. A zero or non-finite `rate` is rejected (FX rates must be positive finite numbers). Parameters: | Name | Type | Description | Default | | --------------- | ------- | --------------------------------------------------- | ----------------------------------------------------------------------------------- | | `from_currency` | \`'Unit | str'\` | Source currency, either a :class:Unit (such as :data:usd) or a string code ("USD"). | | `to_currency` | \`'Unit | str'\` | Destination currency, ditto. | | `rate` | `float` | Strictly positive multiplicative conversion factor. | *required* | Raises: | Type | Description | | ------------------- | ------------------------------------------ | | `ValueError` | if rate is non-positive or non-finite. | | `UnitMismatchError` | if either argument is not a pure currency. | Example: ``` >>> set_fx_rate("USD", "EUR", 0.92) >>> get_fx_rate(usd, eur) 0.92 >>> # Self-rate is always 1.0 and need not be set explicitly. >>> get_fx_rate(usd, usd) 1.0 ``` ### `submodel_function(system, output_ports=None, input_ports=None, auto_seed=True)` Wrap `system`'s ports as a pure function of (context, \*inputs). Parameters: | Name | Type | Description | Default | | -------------- | ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------- | | `system` | `'SystemBase'` | The LeafSystem or Diagram to wrap. | *required* | | `output_ports` | \`'Sequence[OutputPort] | None'\` | Output ports whose values to return. Defaults to all of system.output_ports. | | `input_ports` | \`'Sequence[InputPort] | None'\` | Input ports that the closure will feed. Defaults to all of system.input_ports. Inputs not listed here are assumed already connected or fixed. | | `auto_seed` | `bool` | If True (default), any input port in input_ports that is not already fixed or connected is pre-fixed to a zero placeholder so create_context succeeds on systems with dangling exported inputs. Set to False if you have seeded placeholders yourself. | `True` | Returns: | Type | Description | | ---------- | ------------------------------------------------------------- | | `Callable` | f(context, \*inputs) -> outputs. When a single output port | | `Callable` | is selected the return is a scalar / array; otherwise a tuple | | `Callable` | in output_ports declaration order. | Example:: ``` bld = jaxonomy.DiagramBuilder() plant = bld.add(MyPlant()) bld.export_input(plant.input_ports[0], name="u") bld.export_output(plant.output_ports[0], name="y") diagram = bld.build() f = jaxonomy.submodel_function(diagram) ctx = diagram.create_context() # auto-seeded placeholders y = f(ctx, u) dy_du = jax.grad(lambda u: f(ctx, u))(u0) y_batch = jax.vmap(f, in_axes=(None, 0))(ctx, u_batch) ``` Performance envelope (T-008, follow-up finding 2026-05-16): Each call invokes the diagram's full evaluation machinery — port-fix context managers, dependency-tracked output evaluation, cache invalidation. That overhead is fine for **one-shot rollouts**, **batched evaluation** (where the cost amortises across the batch via `jax.vmap`), and **gradient computation via `jax.grad`** (the closure is traced once, then the compiled XLA program runs at native speed). ``` It is **not** fine for tight Python-side loops that call ``f`` thousands of times per simulated second — typical MPC inner loops where the prediction model is re-evaluated at every sample of a ``jax.lax.scan``-style rollout. There the per-call Python overhead dominates and the wall-clock blows up by 100× or more relative to closing over the underlying primitive directly (e.g. ``interp_2d``, ``lookup_table_nd``, or a hand-rolled JAX function). The canonical workaround in that case is to skip ``submodel_function`` entirely for the inner loop and call the primitive directly inside the scan body. See ``docs/examples/engine_map_fitting_to_mpc.ipynb`` for an example of the hand-rolled-scan pattern. Rule of thumb: if the closure will be invoked from a Python-level loop more than ~100 times per simulation, profile first. ``jax.jit(f)`` + ``jax.vmap`` over the entire batch usually beats a Python loop by orders of magnitude. ``` ### `variant_subsystem(choices, name=None, default=None)` Build a resolver closure for a one-shot variant point. Convenience wrapper around `Variant` + `select_variant` for the common case:: ``` controller = variant_subsystem( choices={ "pid": lambda: build_pid(), "lqr": lambda: build_lqr(), }, default="pid", ) # Later, at "configure" time: active = controller(name="lqr") # returns the lqr Diagram active = controller() # returns the pid Diagram (default) ``` Parameters: | Name | Type | Description | Default | | --------- | ---------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `choices` | `Mapping[str, Callable[[], SystemBase]]` | See Variant.choices. | *required* | | `name` | `Optional[str]` | See Variant.name. | `None` | | `default` | `Optional[str]` | Choice name to use when the returned closure is called without an argument. If None, the first key of choices is used (insertion order, which is guaranteed in Python 3.7+). | `None` | Returns: | Type | Description | | --------------------------- | ------------------------------------------------------------- | | `Callable[..., SystemBase]` | A closure select(name=None) -> SystemBase that, on each call, | | `Callable[..., SystemBase]` | resolves to a freshly-built sub-system for the named choice. | # Block library **Notes on imported neural models** (`ONNX` / `ONNXJax` / `PyTorch` / `TensorFlow`): - **x64 at import.** `import jaxonomy` enables JAX 64-bit mode (`jax_enable_x64`) process-wide, so float32 artifacts see float64 inputs and silently compute in different arithmetic than they were trained and validated in. Cast explicitly at block boundaries: `cast_outputs_to_dtype="float32"` on the block, `x.astype(jnp.float32)` on upstream signals. - **Discrete-time policies need a `ZeroOrderHold`.** A sample-and-hold controller exported from a discrete-time training loop (torch/NEUROMANCER-style) is otherwise re-evaluated at every ODE solver stage, silently destroying step-for-step parity with the exporting framework. Follow the policy block with `ZeroOrderHold(dt=ts)` and pin the step grid with `SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts)` — with that, closed-loop parity is ~4e-8 over 400 steps on the two-tank benchmark. ## `jaxonomy.library` ### `Saturate = _inject_saturate_limit_kwarg(Saturate)` Clip the input signal to a specified range. Given an input signal `u` and upper and lower limits `ulim` and `llim`, the output signal is: ``` y = max(llim, min(ulim, u)) ``` where `max` and `min` are the element-wise maximum and minimum functions. This is equivalent to `y = clip(u, llim, ulim)`. Optionally, the block can also be configured with "dynamic" limits, which will add input ports for time-varying upper and lower limits. Input ports (0) The input signal. (1) The upper limit, if dynamic limits are enabled. (2) The lower limit, if dynamic limits are enabled. (Will be indexed as 1 if dynamic upper limits are not enabled.) Output ports (0) The clipped output signal. Parameters: | Name | Type | Description | Default | | ---------------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `upper_limit` | | The upper limit of the input signal. Default is np.inf. | *required* | | `enable_dynamic_upper_limit` | | If True, then the upper limit can be set by an external signal. Default is False. | *required* | | `lower_limit` | | The lower limit of the input signal. Default is -np.inf. | *required* | | `enable_dynamic_lower_limit` | | If True, then the lower limit can be set by an external signal. Default is False. | *required* | | `limit` | | T-115-followup-saturate-symmetric-kwarg: shorthand for the symmetric case. Saturate(limit=L) expands to Saturate(upper_limit=+L, lower_limit=-L). Mutually exclusive with explicit upper_limit / lower_limit and with the dynamic-limit flags. L must be a positive finite scalar. | *required* | Events The block will trigger an event when the input signal crosses either the upper or lower limit. For example, if the block is configured with static upper and lower limits and the input signal crosses the upper limit, then a zero-crossing event will be triggered. T-115-followup-mode-flag The `mode` kwarg unifies the smooth (differentiable) variant previously exposed as :class:`SoftSaturate`. `mode="hard"` (default) is byte-equivalent to the legacy behavior, including zero-crossing event declaration. `mode="smooth"` dispatches to :func:`soft_saturate` and does *not* declare zero-crossing events (the smooth output has no discontinuity for the solver to catch). The smooth path requires finite `upper_limit` / `lower_limit` and a `sharpness > 0` (defaults to `10.0`). ### `Abs` Bases: `FeedthroughBlock` Output the absolute value of the input signal. Input ports None Output ports (0) The absolute value of the input signal. Events An event is triggered when the output changes from positive to negative or vice versa. ### `Adder` Bases: `ReduceBlock` Computes the sum/difference of the input. The add/subtract operation can be switched by setting the `operators` parameter. For example, a 3-input block specified as `Adder(3, operators="+-+")` would add the first and third inputs and subtract the second input. Input ports (0..n_in-1) The input signals to add/subtract. Output ports (0) The sum/difference of the input signals. ### `Arithmetic` Bases: `ReduceBlock` Performs addition, subtraction, multiplication, and division on the input. The arithmetic operation is determined by setting the `operators` parameter. For example, a 4-input block specified as `Arithmetic(4, operators="+-*/")` would: - Add the first input, - Subtract the second input, - Multiply the third input, - Divide by the fourth input. Input ports (0..n_in-1) The input signals for the specified arithmetic operations. Output ports (0) The result of the specified arithmetic operations on the input signals. ### `AugmentedStateEKF` Bases: `LeafSystem` Extended Kalman Filter with augmented state for **joint** state and parameter estimation. The block estimates both the plant state *x* and unknown parameters *θ* online by augmenting the state vector: .. code-block:: text ``` z = [x; θ] (shape: nx + n_params) ``` with augmented dynamics and observation: .. code-block:: text ``` z[n+1] = f_aug(z[n], u[n]) + noise = [f(x[n], u[n], θ[n]); θ[n]] + [G_x w_x; w_θ] y[n] = h(x[n], u[n], θ[n]) + v[n] E(w_x) = E(w_θ) = E(v) = 0 Cov(w_x) = Q_x, Cov(w_θ) = Q_θ, Cov(v) = R ``` Parameters *θ* follow a **random-walk** model (`θ[n+1] = θ[n] + w_θ`). Setting `Q_theta` small makes parameters quasi-constant; increasing it allows tracking of slowly time-varying parameters. All Jacobians are computed automatically via `jax.jacfwd`. ``` +--------------------+ --- u[n] ------>| |----> x_hat[n] | AugmentedStateEKF | --- y[n] ------>| |----> theta_hat[n] +--------------------+ ``` Input ports (0) u : control vector at timestep n, shape `(nu,)` or scalar (1) y : measurement vector at timestep n, shape `(ny,)` or scalar Output ports (0) x_hat : state estimate, shape `(nx,)` (1) theta_hat : parameter estimate, shape `(n_params,)` Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Sampling period. | *required* | | `nx` | | int Dimension of the plant state x. | *required* | | `n_params` | | int Dimension of the parameter vector θ. | *required* | | `forward` | | Callable Discrete-time state transition: f(x, u, theta) -> x_next. Must be JAX-traceable. | *required* | | `observation` | | Callable Observation function: h(x, u, theta) -> y. Must be JAX-traceable. | *required* | | `G_x_func` | | Callable Process-noise input matrix for states: G_x(t) -> (nx, nw) array. Pass lambda t: jnp.eye(nx) for isotropic process noise. | *required* | | `Q_x_func` | | Callable Process-noise covariance for states: Q_x(t, x, u, theta) -> (nw, nw). | *required* | | `Q_theta` | | array_like Constant parameter diffusion covariance matrix (n_params, n_params). Small values → slow/no parameter drift. | *required* | | `R_func` | | Callable Measurement noise covariance: R(t) -> (ny, ny). | *required* | | `x_hat_0` | | array_like Initial state estimate, shape (nx,). | *required* | | `P_hat_0_x` | | array_like Initial state covariance, shape (nx, nx). | *required* | | `theta_hat_0` | | array_like Initial parameter estimate, shape (n_params,). | *required* | | `P_hat_0_theta` | | array_like Initial parameter covariance, shape (n_params, n_params). | *required* | Example:: ``` import jax.numpy as jnp from jaxonomy import library # Simple first-order system: x[n+1] = a*x[n] + b*u[n] # where 'a' (decay rate) is unknown and must be estimated. def forward(x, u, theta): a = theta[0] return jnp.array([a * x[0] + u[0]]) def observation(x, u, theta): return jnp.array([x[0]]) aekf = library.AugmentedStateEKF( dt=0.1, nx=1, n_params=1, forward=forward, observation=observation, G_x_func=lambda t: jnp.eye(1), Q_x_func=lambda t, x, u, th: jnp.array([[0.01]]), Q_theta=jnp.array([[1e-4]]), R_func=lambda t: jnp.array([[0.1]]), x_hat_0=jnp.zeros(1), P_hat_0_x=jnp.eye(1), theta_hat_0=jnp.zeros(1), P_hat_0_theta=jnp.eye(1), ) ``` #### `DiscreteStateType` Bases: `NamedTuple` Internal filter state (both minus=predicted and plus=corrected estimates). #### `initialize(dt, nx, n_params, forward, observation, G_x_func, Q_x_func, Q_theta, R_func, x_hat_0, P_hat_0_x, theta_hat_0, P_hat_0_theta)` Called at context-creation time to store resolved parameters. ### `AveragedInverter` Bases: `LeafSystem` Averaged voltage-source inverter with the bus-voltage amplitude limit. Switching-free (averaged) model: the commanded rotor-frame voltage passes through unchanged while its magnitude is realizable, and is scaled down onto the voltage circle — preserving its angle — when it is not:: ``` v_lim = V_dc / sqrt(3) (SVPWM) or V_dc / 2 (SPWM) v_out = v_cmd * min(1, v_lim / |v_cmd|) ``` Input ports (0) v_dq_cmd: commanded `[v_d, v_q]` (V). Output ports (0) v_dq: realizable `[v_d, v_q]` after the amplitude limit. Parameters: | Name | Type | Description | Default | | ------------ | ------- | --------------------------------------------------------------- | --------- | | `V_dc` | `float` | DC bus voltage (V). | `48.0` | | `modulation` | `str` | "svpwm" (default, limit V_dc/sqrt(3)) or "spwm" (limit V_dc/2). | `'svpwm'` | ### `Backlash` Bases: `LeafSystem` Hysteretic nonlinearity modelling mechanical slack / backlash. A `Backlash(width)` block has a single discrete state `last_output` that tracks the most recent output value. At each sample tick the update rule is:: ``` delta = u - last_output if delta > width/2: new_output = u - width/2 elif delta < -width/2: new_output = u + width/2 else: new_output = last_output ``` Equivalently, the output "follows" the input only after the input has moved by more than `width/2` from the last output; within that band the output sticks. This is the standard model for gearbox slack / actuator hysteresis: the input must take up the slack before the output moves. Differentiability: the per-step update is expressed via `npa.where` on `delta`; the output is a continuous (piecewise-linear) function of both `u` and `width`, so `jax.grad` w.r.t. `width` is finite. The non-smooth "knee" at `|delta| = width/2` has subgradient `0` (inside the band) or `-sign(delta)/2` (outside) w.r.t. `width` -- both finite, as required. Input ports (0) The driving input signal. Output ports (0) The hysteretic output signal. Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ------------------------------------------------------------------------------------- | ------- | | `width` | | Positive scalar; total hysteresis band width. width=0 recovers y = u (no hysteresis). | `1.0` | | `dt` | | Periodic update sample time. Required: this is a discrete block. | `0.01` | | `initial_output` | | Initial value of the output / discrete state. Default 0.0. | `0.0` | Notes For an exact match against a canonical discrete-time "Backlash" block the discrete sample time must match. For a *continuous* hysteresis approximation, choose `dt` much smaller than the dominant input timescale; the block then tracks the input modulo the `width/2` slack with single-step latency. ### `BatteryCell` Bases: `LeafSystem` Dynamic electro-checmical Li-ion cell model. Based on [Tremblay and Dessaint (2009)](https://doi.org/10.3390/wevj3020289). By using appropriate parameters, the cell model can be used to model a battery pack with the assumption that the cells of the pack behave as a single unit. Parameters E0, K, A, below are abstract parameters used in the model presented in the reference paper. As described in the reference paper, these parameters can be extracted from typical cell manufacturer datasheets; see section 3. Section 3 also provides a table of example values for these parameters. Input ports (0) The current (A) flowing through the cell. Positive is discharge. Output ports (0) The voltage across the cell terminals (V) (1) The state of charge of the cell (normalized between 0 and 1) Parameters: | Name | Type | Description | Default | | ------------- | ------- | ------------------------------------------------------------------------------------ | --------- | | `E0` | `float` | described as "battery constant voltage (V)" by the reference paper. | `3.366` | | `K` | `float` | described as "polarization constant (V/Ah)" by the reference paper. | `0.0076` | | `Q` | `float` | battery capacity in Ah | `2.3` | | `R` | `float` | internal resistance (Ohms) | `0.01` | | `A` | `float` | described as "exponential zone amplitude (V)" by the reference paper. | `0.26422` | | `B` | `float` | described as "exponential zone time constant inverse (1/Ah)" by the reference paper. | `26.5487` | | `initial_SOC` | `float` | initial state of charge, normalized between 0 and 1. | `1.0` | ### `BusCreator` Bases: `LeafSystem` Pack `n = len(field_names)` signals into a single named-bus output. The output is a `collections.namedtuple` (named `"Bus"`) whose fields are exactly `field_names` in declaration order. NamedTuples are first-class JAX pytrees, so the bus signal flows through `jax.jit`, `vmap`, and `grad` without any extra registration. Pair with :class:`BusSelector` to pull individual fields back out downstream. Use this when signals share a logical group identity (e.g. `("position", "velocity", "acceleration")` for a vehicle state bus) and you would rather refer to them by name than by positional `Mux`/`Demux` index. Parameters: | Name | Type | Description | Default | | -------------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `field_names` | | Tuple/list of strings — one name per input port, in declaration order. Must be unique, valid Python identifiers (NamedTuple constraint). | *required* | | `field_units` | | Optional mapping from field name to :class:Unit (T-117-followup-bus-units). When supplied, each input port is declared with the corresponding units= and the output port carries a :class:BusUnit so downstream :class:BusSelector blocks can recover per-field units at connect time. When None (the default), input/output ports carry no unit metadata — byte-equivalent to the T-117-fu-bus-namedtuple shipping behaviour. | `None` | | `field_shapes` | | Optional mapping from field name to a JAX-style shape tuple (T-117-followup-bus-array). When supplied, the named field carries an array of the declared shape rather than a scalar — useful for grouping e.g. an 8-element thermocouple readout under a single "sensors" slot without manually muxing. Fields not listed default to scalar shape (). Default None is byte-equivalent to the T-117-fu-bus-namedtuple all-scalar behaviour. | `None` | Input ports `(0..n-1)` — one port per field name; values are packed into the corresponding slot of the output NamedTuple. Output ports `(0)` — the bus, a NamedTuple with fields `field_names`. #### `bus_type` The underlying NamedTuple class for this bus. #### `bus_unit` The compound :class:`BusUnit` for this bus, or `None` if `field_units` was not supplied at construction time. #### `field_names` The tuple of field names in declaration / port order. #### `field_shapes` Mapping from field name to declared array shape tuple (T-117-followup-bus-array). Fields default to scalar `()` when `field_shapes` is omitted at construction time. ### `BusMerge` Bases: `LeafSystem` Merge two bus signals by union of fields (LeafSystem wrapper). Wraps :func:`merge_buses` as a block for use inside a Diagram. The two upstream bus signals are read from input ports 0 and 1; the merged bus is produced on output port 0. The merged-bus schema is fixed at construction time from `bus_spec_a` and `bus_spec_b` so the output port can be declared with a known NamedTuple type (and so the framework's pytree handling sees a consistent type across context-build and trace time, matching the T-117-fu-bus-namedtuple design for :class:`BusCreator`). Parameters: | Name | Type | Description | Default | | -------------- | ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `bus_spec_a` | | Schema of the first bus. Accepts a :class:BusCreator instance, a NamedTuple class, or a tuple/list of field-name strings. | *required* | | `bus_spec_b` | | Schema of the second bus. Same forms as bus_spec_a. | *required* | | `on_collision` | `str` | Policy for fields that appear in both schemas: "error" (default) — raise :class:ValueError at construction time. "prefer_a" — read the colliding leaf from input port 0. "prefer_b" — read the colliding leaf from input port 1. | `'error'` | Input ports `(0)` — bus_a (NamedTuple-shaped). `(1)` — bus_b (NamedTuple-shaped). Output ports `(0)` — the merged bus, a NamedTuple whose fields are the union of the two input schemas. #### `bus_type` The underlying NamedTuple class for the merged bus. #### `collisions` The tuple of field names that collided between the two input schemas. Empty unless `on_collision` is `"prefer_a"` or `"prefer_b"`. #### `field_names` The tuple of merged field names, in output / declaration order. #### `on_collision` The collision-resolution policy in effect for this block. ### `BusPassthrough` Bases: `LeafSystem` Identity copy of a bus signal — single input port, single output port. The output is the input value, returned as-is. Pass-through semantics: NamedTuple-shaped buses (as produced by :class:`BusCreator` / :class:`BusMerge`) flow through unchanged, scalar/array signals likewise. This block exists to give Diagrams an explicit "junction" node for rewiring, debugging, or scheduler-boundary purposes; it has no parameters and does no computation beyond forwarding. Differentiable: `jax.grad` flows from the output back to the input leaf-by-leaf (identity has Jacobian = identity), and the block is JIT-traceable since the underlying op is a no-op closure return. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `bus_unit` | | Optional :class:BusUnit describing the per-field units of the bus signal. When supplied, both input and output ports are tagged with this BusUnit so connect- time unit checks see a matched pair. None (the default) preserves the unit-less behaviour byte-for-byte. | `None` | Input ports `(0)` — the bus (or any) signal to forward. Output ports `(0)` — the same value, returned as-is. #### `bus_unit` The :class:`BusUnit` propagated through this passthrough, or `None` if no unit metadata was supplied at construction. ### `BusSelector` Bases: `LeafSystem` Pull one named field out of a bus signal. The bus input is expected to be a NamedTuple-shaped value (typically produced by :class:`BusCreator`). The selected field is read with plain `getattr`, so this block is the inverse of `BusCreator` when wired correctly: `BusSelector("a")(BusCreator(["a","b","c"]) (a, b, c)) == a`. Both `getattr` and the NamedTuple constructor are transparent to JAX autodiff, so gradients flow cleanly from the selector output back to the upstream input that filled the corresponding bus slot. T-117-followup-bus-dot-path: `field_name` may contain dots to descend into nested bus signals — e.g. `BusSelector("chassis.suspension.spring_force")` extracts the leaf in one block instead of three cascaded selectors. The path is resolved via :func:`operator.attrgetter`, so each segment must name a valid NamedTuple field at its level. Each segment is validated as a Python identifier at construction time. Parameters: | Name | Type | Description | Default | | ------------ | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `field_name` | | Name (or dot-separated path) of the bus field to extract. Raises AttributeError at execution time if the upstream bus does not have a field at any segment of the path. | *required* | | `bus_unit` | | Optional :class:BusUnit describing the per-field units of the upstream bus (T-117-followup-bus-units). When supplied, the selector's input port is tagged with this BusUnit (so the connect-time check verifies that the upstream :class:BusCreator produced a compatible bus) and its output port is tagged with bus_unit.fields[field_name] so further downstream blocks see the right scalar unit. Default None preserves the unit-less behaviour byte-for-byte. When field_name contains dots, the leaf unit cannot be propagated because :class:BusUnit is flat (one Unit per top-level field); the input bus is still tagged but the output unit is None. Pass a single-segment field_name if you need leaf-unit propagation. | `None` | Input ports `(0)` — the bus signal (NamedTuple-shaped). Output ports `(0)` — the value of `bus.`, optionally sliced at `slice_idx`. #### `bus_unit` The :class:`BusUnit` describing the upstream bus, or `None` if no unit metadata was supplied. #### `field_name` The name of the bus field this block selects. #### `slice_idx` The integer index into an array-valued bus field, or `None` if the selector returns the field value as-is (T-117-followup-bus-array). ### `BusUnit` Compound unit carrying one :class:`Unit` per named bus field. Attached to the output port of a :class:`BusCreator` (and the matching input port of a :class:`BusSelector`) so that the connect-time consistency check can verify each field's unit individually. Attributes: | Name | Type | Description | | -------- | -------------------- | --------------------------------------------------------------------------------------------------- | | `fields` | `Mapping[str, Unit]` | Mapping from bus field name to its :class:Unit. Stored as a plain dict (insertion order preserved). | #### `field_unit(name)` Return the :class:`Unit` for `name`, or `None` if absent. Used by :class:`BusSelector` to look up its output-port unit when wired downstream of a unit-tagged bus. ### `BusUpdate` Bases: `LeafSystem` Replace one field of a bus signal with a new value. Two-input / one-output LeafSystem. Input port 0 carries the upstream bus (a NamedTuple-shaped value, typically produced by :class:`BusCreator`). Input port 1 carries the new value for the field named `field_name`. The output port is a fresh bus identical to the input except that the `field_name` slot is replaced by the new value. Field order is preserved exactly; all other fields are forwarded unchanged. Use this instead of the BusSelector-modify-BusCreator triplet when you only need to edit one field of a wide bus -- the block makes the intent explicit and avoids manually wiring N-1 passthrough edges. Differentiable: the underlying op is NamedTuple construction over `getattr` lookups on the input bus plus the `new_value` leaf, all of which are transparent to `jax.grad` and `jax.jit` (same as :class:`BusCreator` and :func:`merge_buses`). Parameters: | Name | Type | Description | Default | | ------------ | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `bus_spec` | | Schema of the bus. Accepts a :class:BusCreator instance, a NamedTuple class, or a tuple/list of field-name strings (same forms as :class:BusMerge). Used both to validate field_name at construction time and to type the output port's NamedTuple class so the framework's pytree handling sees a consistent type across context-build and trace time. | *required* | | `field_name` | | The name of the field to replace. Must be one of the names declared in bus_spec; otherwise a clear :class:ValueError is raised at construction time. | *required* | | `bus_unit` | | Optional :class:BusUnit describing the per-field units of the upstream bus (T-117-followup-bus-update-units- prop). When supplied, the block's bus_in and bus_out ports advertise this BusUnit so the connect-time check verifies that the upstream :class:BusCreator produced a compatible bus and that downstream consumers see the right schema. The new_value input port is independently tagged with the per-field unit bus_unit.fields[field_name] so the connect-time check enforces unit compatibility on the replacement value. Composes with T-104 Phase 2 behaviour on Sum / Product / Integrator. Default None preserves the unit-less behaviour byte-for-byte. | `None` | Input ports `(0)` — `bus_in`, the upstream bus signal (NamedTuple-shaped). `(1)` — `new_value`, the value to put into the `field_name` slot of the output bus. Output ports `(0)` — `bus_out`, a NamedTuple of the same shape as `bus_in` with `field_name` replaced by `new_value`. #### `bus_type` The underlying NamedTuple class for the output bus. #### `bus_unit` The :class:`BusUnit` propagated through this update, or `None` if no unit metadata was supplied at construction (T-117-followup-bus-update-units-prop). #### `field_name` The name of the field this block replaces on each tick. #### `field_names` The tuple of bus field names, in declaration / output order. ### `Chirp` Bases: `SourceBlock` Produces a linear chirp signal — matches :func:`scipy.signal.chirp`. The output signal is `cos(2π·f(t)·t + phi)` with the linearly swept frequency `f(t) = f0 + (f1 − f0)·t/(2·stop_time)`. At `t=0` the instantaneous frequency is `f0` Hz; at `t=stop_time` it is `f1` Hz. See https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.chirp.html Parameters: | Name | Type | Description | Default | | ----------- | ------- | ----------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `f0` | `float` | Frequency (Hz) at time t=0. | *required* | | `f1` | `float` | Frequency (Hz) at time t=stop_time. | *required* | | `stop_time` | `float` | Time to end the signal (seconds). | *required* | | `phi` | `float` | Phase offset (radians). | `0.0` | | `units` | \`str | None\` | T-122-followup-chirp-hz-convention — frequency-unit convention for f0 / f1. Defaults to "hz" (matches the docstring and :func:scipy.signal.chirp). Legacy diagrams that depended on the pre-2026-05 behaviour (where f0 / f1 were silently interpreted in rad/s) can opt into the old semantics by passing units="rad/s"; doing so emits a :class:DeprecationWarning because the legacy path will be removed in a future release. | Input ports None Output ports (0) The chirp signal. ### `Clarke` Bases: `FeedthroughBlock` Clarke transform: three-phase `[a, b, c]` -> stationary `[alpha, beta]`. Amplitude-invariant scaling: a balanced sinusoidal three-phase set of peak amplitude `M` maps to an alpha-beta vector of magnitude `M`. The zero-sequence component is discarded. ### `Clock` Bases: `SourceBlock` Source block returning simulation time. Input ports None Output ports (0) The simulation time. Parameters: | Name | Type | Description | Default | | ------- | ---- | ----------------------------------------------------------------------------------------------------------------------------- | ------- | | `dtype` | | The data type of the output signal. The default is "None", which will default to the current default floating point precision | `None` | ### `Comparator` Bases: `LeafSystem` Compare two signals using typical relational operators. When using == and != operators, the block uses tolerances to determine if the expression is true or false. Parameters: | Name | Type | Description | Default | | ---------- | ---- | --------------------------------------------------- | ------- | | `operator` | | one of ("==", "!=", ">=", ">", ">=", "\<") | `None` | | `atol` | | the absolute tolerance value used with "==" or "!=" | `1e-05` | | `rtol` | | the relative tolerance value used with "==" or "!=" | `1e-08` | Input Ports (0) The left side operand (1) The right side operand Output Ports (0) The result of the comparison (boolean signal) Events An event is triggered when the output changes from true to false or vice versa. ### `Conditional` Bases: `LeafSystem` Container block that enables/disables a submodel. Parameters: | Name | Type | Description | Default | | --------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `submodel` | `Callable` | Callable taking \*inputs and returning a single output array. For a subdiagram, use jaxonomy.submodel_function to build a compatible callable, then wrap with a context-capturing lambda: lambda \*u: f(context, \*u). | *required* | | `n_inputs` | `int` | Number of non-enable inputs the submodel takes. Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel's inputs in order. | `1` | | `when_disabled` | `str` | "reset", "hold", or "passthrough". | `RESET` | | `initial_value` | | Output value when disabled (reset or hold mode's initial state). Also used to infer output shape/dtype when the submodel has not been evaluated yet. | `0.0` | | `name` | | Optional block name. | *required* | ### `Constant` Bases: `LeafSystem` A source block that emits a constant value. Parameters: | Name | Type | Description | Default | | ------- | --------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `value` | | The constant value of the block. | *required* | | `dtype` | `optional, T-038a-followup-other-blocks` | If set, the constant value is cast to this dtype on output. See LookupTable1d for the per-block dtype contract. | `None` | | `units` | `(optional, T - 104 - followup - units - on - source - blocks)` | If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams). | `None` | Input ports None Output ports (0) The constant value. ### `ContinuousTimeInfiniteHorizonKalmanFilter` Bases: `LeafSystem` Continuous-time Infinite Horizon Kalman Filter for the following system: ``` dot_x = A x + B u + G w y = C x + D u + v E(w) = E(v) = 0 E(ww') = Q E(vv') = R E(wv') = N = 0 ``` Input ports (0) u : continuous-time control vector (1) y : continuous-time measurement vector Output ports (1) x_hat : continuous-time state vector estimate Parameters: | Name | Type | Description | Default | | --------- | ---- | ------------------------------------------- | ---------- | | `A` | | ndarray State transition matrix | *required* | | `B` | | ndarray Input matrix | *required* | | `C` | | ndarray Output matrix | *required* | | `D` | | ndarray Feedthrough matrix | *required* | | `G` | | ndarray Process noise matrix | *required* | | `Q` | | ndarray Process noise covariance matrix | *required* | | `R` | | ndarray Measurement noise covariance matrix | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | #### `for_continuous_plant(plant, x_eq, u_eq, Q, R, G=None, x_hat_bar_0=None, name=None)` Obtain a continuous-time Infinite Horizon Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq) The input plant contains the deterministic forms of the forward and observation operators: ``` dx/dt = f(x,u) y = g(x,u) ``` Note: Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator. A plant with disturbances of the following form is then considered following form: ``` dx/dt = f(x,u) + G w y = g(x,u) + v ``` where: ``` `w` represents the process noise, `v` represents the measurement noise, ``` and ``` E(w) = E(v) = 0 E(ww') = Q E(vv') = R E(wv') = N = 0 ``` This plant with disturbances is linearized (only `f` and `q`) around the equilibrium point to obtain: ``` d/dt (x_bar) = A x_bar + B u_bar + G w --- (C1) y_bar = C x_bar + D u_bar + v --- (C2) ``` where, ``` x_bar = x - x_eq u_bar = u - u_eq y_bar = y - y_bar y_eq = g(x_eq, u_eq) ``` A continuous-time Kalman Filter estimator for the system of equations (C1) and (C2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar` states. The returned system will have Input ports (0) u_bar : continuous-time control vector relative to equilibrium point (1) y_bar : continuous-time measurement vector relative to equilibrium point Output ports (1) x_hat_bar : continuous-time state vector estimate relative to equilibrium point Parameters: | Name | Type | Description | Default | | ------------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. | *required* | | `x_eq` | | ndarray Equilibrium state vector for discretization | *required* | | `u_eq` | | ndarray Equilibrium control vector for discretization | *required* | | `Q` | | ndarray Process noise covariance matrix. | *required* | | `R` | | ndarray Measurement noise covariance matrix. | *required* | | `G` | | ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w. | `None` | | `x_hat_bar_0` | | ndarray Initial state estimate relative to equilibrium point. If None, an identity matrix is assumed. | `None` | ### `CoordinateRotation` Bases: `LeafSystem` Computes the rotation of a 3D vector between coordinate systems. Given sufficient information to construct a rotation matrix `C_AB` from orthogonal coordinate system `B` to orthogonal coordinate system `A`, along with an input vector `x_B` expressed in `B`-axes, this block will compute the matrix-vector product `x_A = C_AB @ x_B`. Note that depending on the type of rotation representation, this matrix may not be explicitly computed. The types of rotations supported are Quaternion, Euler Angles, and Direction Cosine Matrix (DCM). By default, the rotations have the following convention: - **Quaternion:** The rotation is represented by a 4-component quaternion `q`. The rotation is carried out by the product `p_A = q⁻¹ * p_B * q`, where `q⁻¹` is the quaternion inverse of `q`, `*` is the quaternion product, and `p_A` and `p_B` are the quaternion extensions of the vectors `x_A` and `x_B`, i.e. `p_A = [0, x_A]` and `p_B = [0, x_B]`. - **Roll-Pitch-Yaw (Euler Angles):** The rotation is represented by the set of Euler angles ϕ (roll), θ (pitch), and ψ (yaw), in the "1-2-3" convention for intrinsic rotations. The resulting rotation matrix `C_AB(ϕ, θ, ψ)` is the same as the product of the three single-axis rotation matrices `C_AB = Cz(ψ) * Cy(θ) * Cx(ϕ)`. For example, if `B` represents a fixed "world" frame with axes `xyz` and `A` is a body-fixed frame with axes `XYZ`, then `C_AB` represents a rotation from the world frame to the body frame, in the following sequence: 1. Right-hand rotation about the world frame `x`-axis by `ϕ` (roll), resulting in the intermediate frame `x'y'z'` with `x' = x`. 1. Right-hand rotation about the intermediate frame `y'`-axis by `θ` (pitch), resulting in the intermediate frame `x''y''z''` with `y'' = y'`. 1. Right-hand rotation about the intermediate frame `z''`-axis by `ψ` (yaw), resulting in the body frame `XYZ` with `z = z''`. - **Direction Cosine Matrix:** The rotation is directly represented as a 3x3 matrix `C_AB`. The rotation is carried out by the matrix-vector product `x_A = C_AB @ x_B`. Input ports (0): The input vector `x_B` expressed in the `B`-axes. (1): (if `enable_external_rotation_definition=True`) The rotation representation (quaternion, Euler angles, or cosine matrix) that defines the rotation from `B` to `A` (or `A` to `B` if `inverse=True`). Output ports (0): The output vector `x_A` expressed in the `A`-axes. Parameters: | Name | Type | Description | Default | | ------------------------------------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `rotation_type` | `str` | The type of rotation representation to use. Must be one of ("quaternion", "roll_pitch_yaw", "dcm"). | *required* | | `enable_external_rotation_definition` | | If True, the block will have one input port for the rotation representation (quaternion, Euler angles, or cosine matrix). Otherwise the rotation must be provided as a block parameter. | `True` | | `inverse` | | If True, the block will compute the inverse transformation, i.e. if the matrix representation of the rotation is C_AB from frame B to frame A, the block will compute the inverse transformation C_BA = C_AB⁻¹ = C_AB.T | `False` | | `quaternion` | `Array` | The quaternion representation of the rotation if enable_external_rotation_definition=False. | `None` | | `roll_pitch_yaw` | `Array` | The Euler angles representation of the rotation if enable_external_rotation_definition=False. | `None` | | `direction_cosine_matrix` | `Array` | The direction cosine matrix representation of the rotation if enable_external_rotation_definition=False. | `None` | ### `CoordinateRotationConversion` Bases: `LeafSystem` Converts between different representations of rotations. See CoordinateRotation block documentation for descriptions of the different rotation representations supported. This block supports conversion between quaternion, roll-pitch-yaw (Euler angles), and direction cosine matrix (DCM). Note that conversions are reversible in terms of the abstract rotation, although creating a quaternion from a direction cosine matrix (and therefore creating a quaternion from roll-pitch-yaw sequence) results in an arbitrary sign assignment. Input ports (0): The input rotation representation. Output ports (1): The output rotation representation. Parameters: | Name | Type | Description | Default | | ----------------- | ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `conversion_type` | `str` | The type of rotation conversion to perform. Must be one of ("quaternion_to_euler", "quaternion_to_dcm", "euler_to_quaternion", "euler_to_dcm", "dcm_to_quaternion", "dcm_to_euler") | *required* | ### `Counter` Bases: `LeafSystem` Discrete counter that increments on rising edges of its trigger input. The block samples a boolean/binary trigger signal every `dt` seconds. On each rising edge (`prev_trigger == False` and `current_trigger == True`) the internal count advances by `increment`. When `max_count` is set, the counter either *saturates* (clamps at `max_count`) or *wraps* to `0` after the increment that hits / exceeds `max_count`, according to `reset_on_max`. Input ports (0) Trigger signal — boolean / binary-valued. Output ports (0) Current count (integer, stored as `int32`). Parameters: | Name | Type | Description | Default | | --------------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `initial_count` | | Starting count value at t = 0. Default 0. | `0` | | `dt` | | Sample period (seconds) of the discrete update. | *required* | | `increment` | | Amount to add to the count on each rising edge. Default 1. | `1` | | `max_count` | | Optional cap on the count. If None the counter is unbounded. Default None. | `None` | | `reset_on_max` | | When max_count is set, controls behaviour at saturation. True wraps the count back to 0 once it reaches / exceeds max_count. False clamps the count at max_count. Default False. | `False` | Notes Edge detection is the same simple "previous-sample-was-False, current-sample-is-True" rule used by :class:`EdgeDetection` — adequate for boolean triggers driven at the block's own sample rate. For sub-sample-period precision use `ZeroCrossingTriggeredSubsystem` together with this block. The output is integer-typed and therefore non-differentiable in the strict sense; gradient-flow tests should not expect gradients with respect to the count itself. ### `CrossProduct` Bases: `ReduceBlock` Compute the cross product between the inputs. See NumPy docs for details: https://numpy.org/doc/stable/reference/generated/numpy.cross.html Input ports (0) The first input vector. (1) The second input vector. Output ports (0) The cross product of the inputs. ### `CustomJaxBlock` Bases: `LeafSystem` JAX implementation of the PythonScript block. A few important notes and changes/limitations to this JAX implementation: - For this block all code must be written using the JAX-supported subset of Python: * Numerical operations should use `jax.numpy = jnp` instead of `numpy = np` * Standard control flow is not supported (if/else, for, while, etc.). Instead use `lax.cond`, `lax.fori_loop`, `lax.while_loop`, etc. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#structured-control-flow-primitives Where possible, NumPy-style operations like `jnp.where` or `jnp.select` should be preferred to lax control flow primitives. * Functions must be pure and arrays treated as immutable. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates Provided these assumptions hold, the code can be JIT compiled, differentiated, run on GPU, etc. - Variable scoping: the `init_code` and `step_code` are executed in the same scope, so variables declared in the `init_code` will be available in the `step_code` and can be modified in that scope. Internally, everything declared in `init_code` is treated as a single state-like cache entry. However, variables declared in the `step_code` will NOT persist between evaluations. Users should think of `step_code` as a normal Python function where locally declared variables will disappear on leaving the scope. - Persistent variables (outputs and anything declared in `init_code`) must have static shapes and dtypes. This means that you cannot declare `x = 0.0` in `init_code` and then later assign `x = jnp.zeros(4)` in `step_code`. These changes mean that many older PythonScript blocks may not be backwards compatible. Input ports Variable number of input ports, one for each input variable declared in `inputs`. The order of the input ports is the same as the order of the input variables. Output ports Variable number of output ports, one for each output variable declared in `outputs`. The order of the output ports is the same as the order of the output variables. Parameters: | Name | Type | Description | Default | | --------------------- | ------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------ | | `dt` | `float` | The discrete time step of the block, or None if the block is in agnostic time mode. | `None` | | `init_script` | `str` | A string containing Python code that will be executed once when the block is initialized. This code can be used to declare persistent variables that will be available in the step_code. | `''` | | `user_statements` | `str` | A string containing Python code that will be executed once per time step (or per output port evaluation, in agnostic mode). This code can use the persistent variables declared in init_script and the block inputs. | `''` | | `finalize_script` | `str` | A string containing Python code that will be executed once when the simulation completes (or is otherwise torn down). This code can use the persistent variables declared in init_script. Supported for :class:CustomPythonBlock only; raises :class:PythonScriptError for :class:CustomJaxBlock. | `''` | | `accelerate_with_jax` | `bool` | If True, the block will be JIT compiled. If False, the block will be executed in pure Python. This parameter exists for compatibility with UI options; when creating pure Python blocks from code (e.g. for testing), explicitly create the CustomPythonBlock class. | `True` | | `time_mode` | `str` | One of "discrete" or "agnostic". If "discrete", the block step code will be evaluated at peridodic intervals specified by "dt". If "agnostic", the block step code will be evaluated once per output port evaluation, and the block will not have a discrete time step. | `'discrete'` | | `inputs` | `List[str]` | A list of input variable names. The order of the input ports is the same as the order of the input variables. | `None` | | `outputs` | `Mapping[str, Tuple[DTypeLike, ShapeLike]]` | A dictionary mapping output variable names to a tuple of dtype and shape. The order of the output ports is the same as the order of the output variables. | `None` | | `static_parameters` | `Mapping[str, Array]` | A dictionary mapping parameter names to values. Parameters are treated as immutable and cannot be modified in the step code. Static parameters can't be used in ensemble simulations or optimization workflows. | `None` | | `dynamic_parameters` | `Mapping[str, Array]` | A dictionary mapping parameter names to values. Parameters are treated as immutable and cannot be modified in the step code. Dynamic parameters can be arrays or scalars, but must have static shapes and dtypes in order to support JIT compilation. | `None` | #### `check_types(context, error_collector=None)` Test-compile the init and step code to check for errors. ### `CustomPythonBlock` Bases: `CustomJaxBlock` Container for arbitrary user-defined Python code. Implemented to support legacy PythonScript blocks. Not traceable (no JIT compilation or autodiff). The internal implementation and behavior of this block differs vastly from the JAX-compatible block as this block stores state directly within the Python instance. Objects and modules can be kept as discrete state. Note that in "agnostic" mode, the step code will be evaluated *once per output port evaluation*. Because locally defined environment variables (in the init script) are preserved between evaluations, any mutation of these variables will be preserved. This can lead to unexpected behavior and should be avoided. Stateful behavior should be implemented using discrete state variables instead. Warning: The finalize_script parameter is currently accepted but not executed. This is a known limitation. Do not rely on finalize_script for cleanup operations. #### `exec_finalize()` Execute the finalize_script using the current persistent environment. Called once at the end of the simulation via :meth:`post_simulation_finalize`. The script runs in the same environment that was maintained throughout the simulation, so all variables declared in `init_script` (and updated by `user_statements`) are available. Has no effect if `finalize_script` is empty. #### `post_simulation_finalize()` Run `finalize_script` and then call the base-class hook. ### `DMDForecaster` Bases: `LeafSystem` Discrete-time predictor for a fitted (reduced) linear operator. Propagates `x[k+1] = A x[k] (+ B u[k])` and outputs `y[k] = C x[k]` (`C` defaults to the identity, so the state itself is the output). This is the online counterpart of :func:`dmd` / :func:`dmdc`: fit `A` (and `B`) from snapshots offline, then drop the operator into a Jaxonomy diagram as a jax-traceable discrete block that runs inside :func:`jaxonomy.simulate`. An input port (and use of `B`) is created only when `B` is provided. Input ports (0) u\[k\]: control input, present iff `B` is given. Output ports (0) y[k] = C x[k]. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------------------------------------------- | ---------- | | `A` | | State operator (n, n) — a dynamic parameter. | *required* | | `B` | | Optional input operator (n, m) — a dynamic parameter when given. | `None` | | `C` | | Optional output operator (p, n) — a dynamic parameter when given; defaults to identity. | `None` | | `dt` | | Sampling period of the discrete update. | `1.0` | | `initial_state` | | Initial state x[0] of size n (default: zeros). | `None` | ### `DMDResult` Exact-DMD spectral decomposition (Tu et al. 2014). Attributes: | Name | Type | Description | | ------------- | ----- | ------------------------------------------------------------------------------------------------------------------------------------- | | `modes` | `Any` | DMD modes Φ (columns), shape (n, r), generally complex. | | `eigenvalues` | `Any` | Discrete-time DMD eigenvalues λ, shape (r,). The growth/decay and oscillation of the identified linear dynamics; a mode is stable iff | | `amplitudes` | `Any` | Mode amplitudes b fitting the first snapshot, shape (r,). | | `A_tilde` | `Any` | Reduced r×r operator in the POD-projected coordinates. | ### `DMDcResult` DMD-with-control operators (Proctor, Brunton & Kutz 2016). Attributes: | Name | Type | Description | | ------------- | ----- | ------------------------------------------------------------ | | `A` | `Any` | Full n×n state operator. | | `B` | `Any` | Full n×m input operator. | | `A_tilde` | `Any` | Reduced r×r state operator (POD-projected). | | `B_tilde` | `Any` | Reduced r×m input operator. | | `basis` | `Any` | POD basis Û (columns), shape (n, r), mapping reduced ↔ full. | | `eigenvalues` | `Any` | Eigenvalues of A_tilde, shape (r,). | ### `DataSource` Bases: `SourceBlock` Produces outputs from an imported data file (.csv, .npy, .npz). CSV files are read with pandas when installed; otherwise NumPy is used. Parameters: | Name | Type | Description | Default | | -------------- | --------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `file_name` | `str` | Path to .csv, .npy, or .npz. | *required* | | `column` | `Optional[str]` | Optional. When set, selects the signal column(s) by name or index string and overrides data_columns for CSV loading. When None, data_columns is used (default index "1" is the second column, i.e. first column is often time at index 0). | `None` | | `time_column` | `str` | For CSV with time_samples_as_column=True, column name (e.g. "t") or index string (e.g. "0"). If the name is missing but the file has a header row, the first column is used as time. | `'0'` | | `data_columns` | `str` | Column index, name, slice (e.g. 3:8), or list string for CSV. | `'1'` | ### `DeadZone` Bases: `FeedthroughBlock` Generates zero output within a specified range. Applies the following function: ``` [ input, input < -half_range output = | 0, -half_range <= input <= half_range [ input input > half_range ``` Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `half_range` | | The range of the dead zone. Must be > 0. | `1.0` | | `mode` | | "hard" (default, byte-equivalent to legacy behavior) or "smooth". Smooth mode replaces the discontinuous gate with a sigmoid-blended kernel so gradients flow through the dead zone region; the output then has no discontinuity at | x | | `sharpness` | | Positive scalar; default 10.0. Only used in smooth mode. Larger values give a tighter approximation to the hard dead zone (with smaller gradients inside the band). | `10.0` | | `output_shifted` | | False (default, byte-equivalent to legacy behavior) or True. When True, the hard-mode output outside the band is shifted by half_range * sign(input) so the output is continuous across the band boundary (slope 1 outside, value 0 at the boundary). The False form keeps the legacy "Coulomb friction" semantics where the output jumps at the band boundary. The smooth mode already produces a continuous output and is unaffected by this flag. | `False` | Input ports (0) The input signal. Output ports (0) The input signal modified by the dead zone. Events An event is triggered when the signal enters or exits the dead zone in either direction (hard mode only). T-115-followup-deadzone-backlash The `mode` kwarg unifies a smooth (differentiable) variant. `mode="hard"` (default) is byte-equivalent to the legacy behavior, including zero-crossing event declaration. `mode="smooth"` dispatches to a sigmoid-blended formula (see :func:`soft_dead_zone`) and does *not* declare zero-crossing events. T-115-followup-deadzone-bilinear The `output_shifted` kwarg toggles between the legacy Coulomb-style hard dead-zone (default, output jumps at the band boundary) and the shifted-output variant (continuous across the band boundary). Default `False` keeps the block byte-equivalent to phase 1. ### `Decimator` Bases: `LeafSystem` Fast-to-slow rate transition: subsample-and-hold at `output_dt`. Implements a discrete-time decimator that samples its input on a periodic clock at `output_dt` (the slow rate) and holds the value until the next slow tick. The input is assumed to be running at `input_dt` (the fast rate); the block is agnostic to the actual upstream sampling, but the rate-mismatch detector uses this declared pair to recognise the bridge. Difference equation, with input `u` and output `y`:: ``` x[k+1] = u[k * (output_dt / input_dt)] y(t) = x[k], t in (t_k, t_k + output_dt) ``` Equivalent to a "Rate Transition (fast to slow)" block in its default ZOH-with-decimation mode. Input ports (0) The fast-rate input signal. Output ports (0) The slow-rate held signal. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `input_dt` | | Sample period of the upstream (fast) source. Used for documentation / the rate-mismatch detector; the block does not actually read the upstream clock. | *required* | | `output_dt` | | Sample period of this block's output (slow rate). Must satisfy output_dt > input_dt for "fast to slow" semantics; the constructor warns if not. | *required* | | `initial_state` | | Initial output value held until the first slow tick fires. Default 0.0. | `0.0` | | `mode` | | How to combine the input samples within each output_dt window before emitting at the slow tick. One of: "pick_last" (default) — emit the most recent input sample at the slow tick (the standard "Rate Transition fast → slow" default). Byte-equivalent to T-123 phase 1. "mean" — emit the arithmetic mean of every input sample observed during the window. Standard anti-aliasing decimation for continuous signals; differentiable through the input (linear). "peak" — emit the input sample with the largest absolute value over the window. Preserves peak excursions for envelope tracking / detection. The selector itself is non-differentiable but gradients flow through the selected sample's value (a np.where branch picks the max- | u | Notes (`mode="mean"` / `mode="peak"`): The block declares a second periodic update at `input_dt` that accumulates samples into a running buffer. At each slow tick the emit-and-reset update fires first (declared before the accumulator in `__init__`), reads the buffer, computes the window result, and zeroes the buffer for the next window. At simultaneous slow+fast ticks the emit therefore sees the full window from the previous interval; the fast tick at the same `t` then starts the next window with the current input as its first sample. ### `Demultiplexer` Bases: `LeafSystem` Split a vector signal into its components. Input ports (0) The vector signal to split. Output ports (0..n_out-1) The components of the input signal. ### `Demux` Bases: `LeafSystem` Unstack a single array input into `n_outputs` separate signals. This is the standard `Demux` block and the inverse of :class:`Mux`: given a 1-D input `[a, b, c]` it produces three scalar outputs `a`, `b`, `c`; given a 2-D input of shape `(n_outputs, k)` it produces `n_outputs` outputs of shape `(k,)`. Internally this is index-based slicing along axis 0, which is fully differentiable through every output port (each output picks one slice of the input vector). Input ports (0) The vector or array signal to split. Its leading axis must have length `n_outputs`. Output ports (0..n_outputs-1) The components of the input signal. ### `Derivative` Bases: `LTISystem` Causal estimate of the derivative of a signal in continuous time. This is implemented as a state-space system with matrices (A, B, C, D), which are then used to create a (first-order) LTISystem. Note that this only supports single-input, single-output derivative blocks. The derivative is implemented as a filter with a filter coefficient of `N`, which is used to construct the following proper transfer function: ``` H(s) = Ns / (s + N) ``` As N -> ∞, the transfer function approaches a pure differentiator. However, this system becomes increasingly stiff and difficult to integrate, so it is recommended to select a value of N based on the time scales of the system. From the transfer function, `scipy.signal.tf2ss` is used to convert to state-space form and create an LTISystem. Input ports (0) u: Input (scalar) Output ports (0) y: Output (scalar), estimating the time derivative du/dt ### `DerivativeDiscrete` Bases: `LeafSystem` Discrete approximation to the derivative of the input signal w.r.t. time.' By default the block uses a simple backward difference approximation: ``` y[k] = (u[k] - u[k-1]) / dt ``` However, the block can also be configured to use a recursive filter for a better approximation. In this case the filter coefficients are determined by the `filter_type` and `filter_coefficient` parameters. The filter is a pair of two-element arrays `a` and `b` and the filter equation is: ``` a0*y[k] + a1*y[k-1] = b0*u[k] + b1*u[k-1] ``` Denoting the `filter_coefficient` parameter by `N`, the following filters are available: - "none": The default, a simple finite difference approximation. - "forward": A filtered forward Euler discretization. The filter is: `a = [1, (N*dt - 1)]` and `b = [N, -N]`. - "backward": A filtered backward Euler discretization. The filter is: `a = [(1 + N*dt), -1]` and `b = [N, -N]`. - "bilinear": A filtered bilinear transform discretization. The filter is: `a = [(2 + N*dt), (-2 + N*dt)]` and `b = [2*N, -2*N]`. Input ports (0) The input signal. Output ports (0) The approximate derivative of the input signal. Parameters: | Name | Type | Description | Default | | -------------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | The time step of the discrete approximation. | *required* | | `filter_type` | | One of "none", "forward", "backward", or "bilinear". This determines the type of filter used to approximate the derivative. The default is "none", corresponding to a simple backward difference approximation. | `'none'` | | `filter_coefficient` | | The coefficient in the filter (N in the equations above). This is only used if filter_type is not "none". The default is 1.0. | `1.0` | #### `initialize_static_data(context)` Infer the size and dtype of the internal states ### `DirectShootingNMPC` Bases: `NonlinearMPCIpopt` Implementation of nonlinear MPC with a direct shooting transcription and IPOPT as the NLP solver. Input ports (0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC. Output ports (1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC. Parameters: | Name | Type | Description | Default | | ------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | LeafSystem or Diagram The plant to be controlled. | *required* | | `Q` | | Array State weighting matrix in the cost function. | *required* | | `QN` | | Array Terminal state weighting matrix in the cost function. | *required* | | `R` | | Array Control input weighting matrix in the cost function. | *required* | | `N` | | int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now. | *required* | | `nh` | | int Number of minor steps to take within an RK4 major step. | *required* | | `dt` | | float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons. | *required* | | `lb_u` | | Array Lower bound on the control input vector. | `None` | | `ub_u` | | Array Upper bound on the control input vector. | `None` | | `u_optvars_0` | | Array Initial guess for the control vector optimization variables in the NLP. | `None` | ### `DirectTranscriptionNMPC` Bases: `NonlinearMPCIpopt` Implementation of nonlinear MPC with direct transcription and IPOPT as the NLP solver. Input ports (0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC. Output ports (1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC. Parameters: | Name | Type | Description | Default | | ------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | LeafSystem or Diagram The plant to be controlled. | *required* | | `Q` | | Array State weighting matrix in the cost function. | *required* | | `QN` | | Array Terminal state weighting matrix in the cost function. | *required* | | `R` | | Array Control input weighting matrix in the cost function. | *required* | | `N` | | int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now. | *required* | | `nh` | | int Number of minor steps to take within an RK4 major step. | *required* | | `dt` | | float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons. | *required* | | `lb_x` | | Array Lower bound on the state vector. | `None` | | `ub_x` | | Array Upper bound on the state vector. | `None` | | `lb_u` | | Array Lower bound on the control input vector. | `None` | | `ub_u` | | Array Upper bound on the control input vector. | `None` | | `x_optvars_0` | | Array Initial guess for the state vector optimization variables in the NLP. | `None` | | `u_optvars_0` | | Array Initial guess for the control vector optimization variables in the NLP. | `None` | ### `DiscreteClock` Bases: `LeafSystem` Source block that produces the time sampled at a fixed rate. The block maintains the most recently sampled time as a discrete state, provided to the output port during the following interval. Graphically, a discrete clock sampled at 100 Hz would have the following time series: ``` x(t) ●━ | ┆ .03 | ●━━━━○ | ┆ .02 | ●━━━━○ | ┆ .01 | ●━━━━○ | ┆ 0 ●━━━━○----+----+----+-- t 0 .01 .02 .03 .04 ``` The recorded states are the closed circles, which should be interpreted at index `n` as the value seen by all other blocks on the interval `(t[n], t[n+1])`. Input ports None Output ports (0) The sampled time. Parameters: | Name | Type | Description | Default | | ------------ | ---- | ------------------------------------------------------------- | ---------- | | `dt` | | The sampling period of the clock. | *required* | | `start_time` | | The simulation time at which the clock starts. Defaults to 0. | `0` | ### `DiscreteInitializer` Bases: `LeafSystem` Discrete Initializer. Outputs True for first discrete step, then outputs False there after. Or, outputs False for first discrete step, then outputs True there after. Practical for cases where it is necessary to have some signal fed initially by some initialization, but then after from else in the model. Input ports None Output ports (0) The dot product of the inputs. ### `DiscreteTimeLinearQuadraticRegulator` Bases: `LeafSystem` Linear Quadratic Regulator (LQR) for a discrete-time system: x[k+1] = A x[k] + B u[k]. Computes the optimal control input: u[k] = -K x[k], where u minimises the cost function over \[0, ∞)\]: J = ∑(x[k].T Q x[k] + u[k].T R u[k]). Input ports (0) x\[k\]: state vector of the system. Output ports (0) u\[k\]: optimal control vector. Parameters: | Name | Type | Description | Default | | ---- | ---- | ------------------------------------ | ---------- | | `A` | | Array State matrix of the system. | *required* | | `B` | | Array Input matrix of the system. | *required* | | `Q` | | Array State cost matrix. | *required* | | `R` | | Array Input cost matrix. | *required* | | `dt` | | float Sampling period of the system. | *required* | ### `DotProduct` Bases: `ReduceBlock` Compute the dot product between the inputs. This block dispatches to `jax.numpy.dot`, so the semantics, broadcasting rules, etc. are the same. See the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.dot.html Input ports (0) The first input vector. (1) The second input vector. Output ports (0) The dot product of the inputs. ### `EDMDResult` Fitted Koopman model (Williams, Kevrekidis & Rowley 2015). Attributes: | Name | Type | Description | | ------------ | ---------- | ----------------------------------------------------------------- | | `K` | `Any` | Koopman operator on lifted observables, shape (L, L). | | `B` | `Any` | Input operator on the lifted space, shape (L, m) (eDMDc) or None. | | `C` | `Any` | De-lift matrix mapping lifted → physical state, shape (n, L). | | `dictionary` | `Callable` | The observable dictionary g used for lifting. | ### `ERAResult` Minimal state-space realization from Markov parameters (Juang & Pappa 1985). Attributes: | Name | Type | Description | | ----------------- | ----- | ------------------------------------------------------------- | | `A` | `Any` | Realized r×r state matrix. | | `B` | `Any` | Realized r×n_inputs input matrix. | | `C` | `Any` | Realized n_outputs×r output matrix. | | `D` | `Any` | Feedthrough n_outputs×n_inputs (the zeroth Markov parameter). | | `singular_values` | `Any` | Hankel singular values (from the block-Hankel SVD). | ### `EdgeDetection` Bases: `LeafSystem` Output is true only when the input signal changes in a specified way. The block updates at a discrete rate, checking the boolean- or binary-valued input signal for changes. Available edge detection modes are: - "rising": Output is true when the input changes from False (0) to True (1). - "falling": Output is true when the input changes from True (1) to False (0). - "either": Output is true when the input changes in either direction Input ports (0) The input signal. Must be boolean or binary-valued. Output ports (0) The edge detection output signal. Boolean-valued. Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ------------------------------------------------------------------------------------------------------ | ---------- | | `dt` | | The sampling period of the block. | *required* | | `edge_detection` | | One of "rising", "falling", or "either". Determines the type of edge detection performed by the block. | *required* | | `initial_state` | | The initial value of the output signal. | `False` | ### `EnabledMode` Allowed string values for `EnabledSubsystem.mode`. ### `EnabledStateMode` Allowed string values for `EnabledSubsystem.state_mode`. Controls how the *continuous state* (declared via `state_dynamics`) evolves while the enable signal is false: - `HOLD` (default): freeze the state at its current value (`xdot = 0` while disabled). Resumes integration on re-enable. - `RESET`: snap the state back to `initial_state` on every disable→enable transition (so each enable window starts from the configured initial value). While disabled, the state is held. - `FREE`: the state evolves according to `state_dynamics` regardless of enable. Only the *output* is masked per `mode=`. ### `EnabledSubsystem` Bases: `LeafSystem` Container block: run a submodel only while an enable signal is true. This is the subsystem-framing wrapper around the existing :class:`jaxonomy.library.Conditional` primitive (T-009). It exists as a separate class so that: - The block-diagram-vocabulary name `EnabledSubsystem` is discoverable next to the rest of the container family. - We can later extend the `mode="hold"` path with subsystem-state semantics (per-block discrete-state binding) without disturbing the lighter `Conditional` primitive. Parameters: | Name | Type | Description | Default | | ---------------- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output (single output per phase 1). Must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of submodel inputs (does NOT include the enable port). Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel inputs. | `1` | | `mode` | `Literal['reset', 'passthrough', 'hold']` | One of "reset" / "passthrough" / "hold". reset: when disabled, output = initial_value. passthrough: when disabled, output = first user input (input port 1). Submodel and passthrough output must broadcast-compatibly. hold: when disabled, output holds the most recent snapshot taken at hold_period. Requires a positive hold_period. | `RESET` | | `initial_value` | | Output value when disabled in reset mode, and the seed for the held discrete state in hold mode. Used to infer output shape/dtype. | `0.0` | | `hold_period` | \`float | None\` | Sample period (seconds) for the held snapshot in hold mode. Required iff mode == "hold". | | `state_mode` | `Literal['hold', 'reset', 'free']` | One of "hold" / "reset" / "free". Controls the continuous-state behaviour while disabled (independent of mode= which gates only the output). See :class:EnabledStateMode. Default "hold". Only has an effect when state_dynamics is provided; for the stateless submodel default this kwarg is validated but otherwise a no-op (so the default-off path is byte-equivalent to phase 1). | `HOLD` | | `state_dynamics` | \`Callable | None\` | Optional callable f(t, x, \*user_inputs) -> xdot defining a continuous state for the EnabledSubsystem itself. When provided, the block declares a continuous state seeded by initial_state (or initial_value if initial_state is None) and applies state_mode semantics around it. When omitted, the block has no continuous state and behaves exactly as in T-120 phase 1. | | `initial_state` | | Initial value of the continuous state. Required when state_dynamics is provided. | `None` | | `name` | | Optional block name. | *required* | ### `Exponent` Bases: `FeedthroughBlock` Compute the exponential of the input signal. Input ports (0) The input signal. Output ports (0) The exponential of the input signal. Parameters: | Name | Type | Description | Default | | ------ | ---- | --------------------------------------------------------------------- | ---------- | | `base` | | One of "exp" or "2". Determines the base of the exponential function. | *required* | ### `ExtendedKalmanFilter` Bases: `KalmanFilterBase` Extended Kalman Filter (EKF) for the following system: ``` ``` x[n+1] = f(x[n], u[n]) + G(t[n]) w[n] y[n] = g(x[n], u[n]) + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q(t[n], x[n], u[n]) E(v[n]v'[n] = R(t[n]) E(w[n]v'[n] = N(t[n]) = 0 ``` ``` `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`, while R`and`G`are discrete-time functions of time`t[n]`.`Q`is a discrete-time function of`t[n], x[n], u[n]\`. This last aspect is included for zero-order-hold discretization of a continuous-time system Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ------------- | ---- | -------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `forward` | | Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations. | *required* | | `observation` | | Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations. | *required* | | `G_func` | | Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations. | *required* | | `Q_func` | | Callable A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents Q in the above equations. | *required* | | `R_func` | | Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations. | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate | *required* | #### `for_continuous_plant(plant, dt, G_func, Q_func, R_func, x_hat_0, P_hat_0, discretization_method='euler', discretized_noise=False, name=None, ui_id=None)` Extended Kalman Filter system for a continuous-time plant. The input plant contains the deterministic forms of the forward and observation operators: ``` dx/dt = f(x,u) y = g(x,u) ``` Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator; (ii) the user may pass a plant with disturbances (not recommended) as the input plant. In this case, the forward and observation evaluations will be corrupted by noise. A plant with disturbances of the following form is then considered: ``` dx/dt = f(x,u) + G(t) w -- (C1) y = g(x,u) + v -- (C2) ``` where: ``` `w` represents the process noise, `v` represents the measurement noise, ``` and ``` E(w) = E(v) = 0 E(ww') = Q(t) E(vv') = R(t) E(wv') = N(t) = 0 ``` This plant is discretized to obtain the following form: ``` x[n+1] = fd(x[n], u[n]) + Gd w[n] -- (D1) y[n] = gd(x[n], u[n]) + v[n] -- (D2) E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Qd E(v[n]v'[n] = Rd E(w[n]v'[n] = Nd = 0 ``` The above discretization is performed either via the `euler` or the `zoh` method, and an Extended Kalman Filter estimator for the system of equations (D1) and (D2) is returned. Note: If `discretized_noise` is True, then it is assumed that the user is directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to an Identity matrix. The returned system will have: Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. | *required* | | `dt` | | float Time step for the discretization. | *required* | | `G_func` | | Callable A function with signature G(t) -> G that represents G in the continuous-time equations (C1) and (C2). | *required* | | `Q_func` | | Callable A function with signature Q(t) -> Q that represents Q in the continuous-time equations (C1) and (C2). | *required* | | `R_func` | | Callable A function with signature R(t) -> R that represents R in the continuous-time equations (C1) and (C2). | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate. If None, an Identity matrix is assumed. | *required* | | `discretization_method` | | str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler". | `'euler'` | | `discretized_noise` | | bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G_func, Q_func, and R_func provide Gd(t), Qd(t), and Rd(t), respectively. | `False` | #### `from_operators(dt, forward, observation, G_func, Q_func, R_func, x_hat_0, P_hat_0, name=None, ui_id=None)` Extended Kalman Filter (UKF) for the following system: ``` x[n+1] = f(x[n], u[n]) + G(t[n]) w[n] y[n] = g(x[n], u[n]) + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q(t[n], x[n], u[n]) E(v[n]v'[n] = R(t[n]) E(w[n]v'[n] = N(t[n]) = 0 ``` `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`, while `Q` and `R` and `G` are discrete-time functions of time `t[n]`. Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ------------- | ---- | ---------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `forward` | | Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations. | *required* | | `observation` | | Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations. | *required* | | `G_func` | | Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations. | *required* | | `Q_func` | | Callable A function with signature Q(t[n]) -> Q[n] that represents Q in the above equations. | *required* | | `R_func` | | Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations. | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate | *required* | ### `FeedthroughBlock` Bases: `LeafSystem` Simple feedthrough blocks with a function of a single input ### `FilterDiscrete` Bases: `LeafSystem` Finite Impulse Response (FIR) filter. Similar to https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.lfilter.html Note: does not implement the IIR filter. Input ports (0) The input signal. Output ports (0) The filtered signal. Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ----------------------------- | ---------- | | `b_coefficients` | | Array of filter coefficients. | *required* | ### `FiniteHorizonLinearQuadraticRegulator` Bases: `LeafSystem` Finite Horizon Linear Quadratic Regulator (LQR) for a continuous-time system. Solves the Riccati Differential Equation (RDE) to compute the optimal control for the following finitie horizon cost function over \[t0, tf\]: Minimise cost J: ``` J = [x(tf) - xd(tf)].T Qf [x(tf) - xd(tf)] + ∫[(x(t) - xd(t)].T Q [(x(t) - xd(t)] dt + ∫[(u(t) - ud(t)].T R [(u(t) - ud(t)] dt + 2 ∫[(x(t) - xd(t)].T N [(u(t) - ud(t)] dt ``` subject to the constraints: dx(t)/dt - dx0(t)/dt = A [x(t)-x0(t)] + B [u(t)-u0(t)] - c(t), where, x(t) is the state vector, u(t) is the control vector, xd(t) is the desired state vector, ud(t) is the desired control vector, x0(t) is the nominal state vector, u0(t) is the nominal control vector, Q, R, and N are the state, input, and cross cost matrices, Qf is the final state cost matrix, and A, B, and c are computed from linearisation of the plant `df/dx = f(x, u)` around the nominal trajectory (x0(t), u0(t)). ``` A = df/dx(x0(t), u0(t), t) B = df/du(x0(t), u0(t), t) c = f(x0(t), u0(t), t) - dx0(t)/dt ``` The optimal control `u` obtained by the solution of the above problem is output. See Section 8.5.1 of https://underactuated.csail.mit.edu/lqr.html#finite_horizon Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `t0` | | float Initial time of the finite horizon. | *required* | | `tf` | | float Final time of the finite horizon. | *required* | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. The plant to be controlled. This represents df/dx = f(x, u). | *required* | | `Qf` | | Array Final state cost matrix. | *required* | | `func_Q` | | Callable A function that returns the state cost matrix Q at time t: func_Q(t)->Q | *required* | | `func_R` | | Callable A function that returns the input cost matrix R at time t: func_R(t)->R | *required* | | `func_N` | | Callable A function that returns the cross cost matrix N at time t. func_N(t)->N | *required* | | `func_x_0` | | Callable A function that returns the nominal state vector x0 at time t. func_x_0(t)->x0 | *required* | | `func_u_0` | | Callable A function that returns the nominal control vector u0 at time t. func_u_0(t)->u0 | *required* | | `func_x_d` | | Callable A function that returns the desired state vector xd at time t. func_x_d(t)->xd. If None, assumed to be the same as the nominal trajectory. | *required* | | `func_u_d` | | Callable A function that returns the desired control vector ud at time t. func_u_d(t)->ud. If None, assumed to be the same as the nominal trajectory. | *required* | ### `FrequencyResponse` Result of a frequency-response evaluation. Attributes: | Name | Type | Description | | ------------ | ----- | ----------------------------------------------------------------- | | `omegas` | `Any` | Angular frequency vector ω in rad/s, shape (K,). | | `response` | `Any` | Complex frequency response G(jω), shape (K, n_outputs, n_inputs). | | `magnitudes` | `Any` | | | `phases` | `Any` | Phase arg G(jω) in radians, shape (K, n_outputs, n_inputs). | ### `GPModel` Fitted Gaussian-process regressor (kriging). Stores the training inputs, the pre-solved weight vector `alpha = K^-1 y`, and the Cholesky factor of the (noisy) covariance matrix for variance prediction. See Rasmussen & Williams 2006, Algorithm 2.1. #### `predict(Xstar)` Posterior `(mean, variance)` at query points `Xstar`. jax-traceable. `Xstar` may be 1-D (single point / single feature) or 2-D `(m, d)`. Returns arrays of shape `(m,)`. ### `Gain` Bases: `FeedthroughBlock` Multiply the input signal by a constant value. Input ports (0) The input signal. Output ports (0) The input signal multiplied by the gain: `y = gain * u`. Parameters: | Name | Type | Description | Default | | ------- | ---------------------------------------- | ----------------------------------------------------------------------------------------------------- | ---------- | | `gain` | | The value to scale the input signal by. | *required* | | `dtype` | `optional, T-038a-followup-other-blocks` | If set, the block's output is cast to this dtype. See LookupTable1d for the per-block dtype contract. | `None` | ### `GaussianProcess` Bases: `LeafSystem` Gaussian-process (kriging) surrogate as a feedthrough block. Input port 0 is the feature vector `u`; output port 0 is the predictive mean and output port 1 the predictive variance. Training data and the Cholesky factor are stored statically on the block; the weight vector `alpha` and the kernel hyperparameters are dynamic parameters so the surrogate is differentiable in them (Rasmussen & Williams 2006). ### `HermiteSimpsonNMPC` Bases: `NonlinearMPCIpopt` Implementation of nonlinear MPC with Hermite-Simpson collocation and IPOPT as the NLP solver. Input ports (0) x_0 : current state vector. (1) x_ref : reference state trajectory for the nonlinear MPC. (2) u_ref : reference input trajectory for the nonlinear MPC. Output ports (1) u_opt : the optimal control input to be applied at the current time step as determined by the nonlinear MPC. Parameters: | Name | Type | Description | Default | | ---------------------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | LeafSystem or Diagram The plant to be controlled. | *required* | | `Q` | | Array State weighting matrix in the cost function. | *required* | | `QN` | | Array Terminal state weighting matrix in the cost function. | *required* | | `R` | | Array Control input weighting matrix in the cost function. | *required* | | `N` | | int The prediction horizon, an integer specifying the number of steps to predict. Note: prediction and control horizons are identical for now. | *required* | | `dt` | | float: Major time step, a scalar indicating the increment in time for each step in the prediction and control horizons. | *required* | | `lb_x` | | Array Lower bound on the state vector. | `None` | | `ub_x` | | Array Upper bound on the state vector. | `None` | | `lb_u` | | Array Lower bound on the control input vector. | `None` | | `ub_u` | | Array Upper bound on the control input vector. | `None` | | `include_terminal_x_as_constraint` | | bool If True, the terminal state is included as a constraint in the NLP. | `False` | | `include_terminal_u_as_constraint` | | bool If True, the terminal control input is included as a constraint in the NLP. | `False` | | `x_optvars_0` | | Array Initial guess for the state vector optimization variables in the NLP. | `None` | | `u_optvars_0` | | Array Initial guess for the control vector optimization variables in the NLP. | `None` | ### `IOPort` Bases: `FeedthroughBlock` Simple class for organizing input/output ports for groups/submodels. Since these are treated as standalone blocks in the UI rather than specific input/output ports exported to the parent model, it is more straightforward to represent them that way here as well. This class represents a simple one-input, one-output feedthrough block where the feedthrough function is an identity. The input (resp. output) port can then be exported to the parent model to create an Inport (resp. Outport). ### `IfThenElse` Bases: `LeafSystem` Applies a conditional expression to the input signals. Given inputs `pred`, `true_val`, and `false_val`, the block computes: ``` y = true_val if pred else false_val ``` The true and false values may be any arrays, but must have the same shape and dtype. Input ports (0) The boolean predicate. (1) The true value. (2) The false value. Output ports (0) The result of the conditional expression. Shape and dtype will match the true and false values. Events An event is triggered when the output changes from true to false or vice versa. ### `InfiniteHorizonKalmanFilter` Bases: `KalmanFilterBase` Infinite Horizon Kalman Filter for the following system: ``` x[n+1] = A x[n] + B u[n] + G w[n] y[n] = C x[n] + D u[n] + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q E(v[n]v'[n]) = R E(w[n]v'[n]) = N = 0 ``` Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | --------- | ---- | ----------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `A` | | ndarray State transition matrix | *required* | | `B` | | ndarray Input matrix | *required* | | `C` | | ndarray Output matrix. If None, full state output is assumed. | `None` | | `D` | | ndarray Feedthrough matrix. If None, no feedthrough is assumed. | `None` | | `G` | | ndarray Process noise matrix. If None, G=B is assumed. | `None` | | `Q` | | ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and A is assumed. | `None` | | `R` | | ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with C and A is assumed. | `None` | | `x_hat_0` | | ndarray Initial state estimate. If None, an array of zeros is assumed. | `None` | #### `for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_bar_0=None, discretization_method='zoh', discretized_noise=False, name=None, ui_id=None)` Obtain an Infinite Horizon Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq) The input plant contains the deterministic forms of the forward and observation operators: ``` dx/dt = f(x,u) y = g(x,u) ``` Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator. (ii) the user may pass a plant with disturbances as the input plant. However, computation of `y_eq` will be fraught with disturbances. A plant with disturbances of the following form is then considered following form: ``` dx/dt = f(x,u) + G w --- (C1) y = g(x,u) + v --- (C2) ``` where: ``` `w` represents the process noise, `v` represents the measurement noise, ``` and ``` E(w) = E(v) = 0 E(ww') = Q E(vv') = R E(wv') = N = 0 ``` This plant with disturbances is linearized (only `f` and `g`) around the equilibrium point to obtain: ``` d/dt (x_bar) = A x_bar + B u_bar + G w y_bar = C x_bar + D u_bar + v ``` where, ``` x_bar = x - x_eq u_bar = u - u_eq y_bar = y - y_bar y_eq = g(x_eq, u_eq) ``` The linearized plant is then discretized via `euler` or `zoh` method to obtain: ``` x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n] --- (L1) y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n] --- (L2) E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Qd E(v[n]v'[n]) = Rd E(w[n]v'[n]) = Nd = 0 ``` Note: If `discretized_noise` is True, then it is assumed that the user is providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to Identity matrix. An Infinite Horizon Kalman Filter estimator for the system of equations (L1) and (L2) is returned. This filter is in the `x_bar`, `u_bar`, and `y_bar` states. This returned system will have Input ports (0) u_bar[n] : control vector at timestep n, relative to equilibrium (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium Output ports (1) x_hat_bar[n] : state vector estimate at timestep n, relative to equilibrium Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. | *required* | | `x_eq` | | ndarray Equilibrium state vector for discretization | *required* | | `u_eq` | | ndarray Equilibrium control vector for discretization | *required* | | `dt` | | float Time step for the discretization. | *required* | | `Q` | | ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and and linearized system's A is assumed. | `None` | | `R` | | ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with linearized system's C and A is assumed. | `None` | | `G` | | ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w. | `None` | | `x_hat_bar_0` | | ndarray Initial state estimate relative to equilibrium. If None, an identity matrix is assumed. | `None` | | `discretization_method` | | str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler". | `'zoh'` | | `discretized_noise` | | bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G, Q, and R are assumed to be Gd, Qd, and Rd, respectively. | `False` | #### `global_filter_for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None)` See docs for `for_continuous_plant`, which returns the local infinite horizon Kalman Filter. This method additionally converts the local Kalman Filter to a global estimator. See docs for `make_global_estimator_from_local` for details. ### `Integrator` Bases: `LeafSystem` Integrate the input signal in time. The Integrator block is the main primitive for building continuous-time models. It is a first-order integrator, implementing the following linear time-invariant ordinary differential equation for input values `u` and output values `y`: ``` ẋ = u y = x ``` where `x` is the state of the integrator. The integrator is initialized with the value of the `initial_state` parameter. Options Reset: the integrator can be configured to reset its state on an input trigger. The reset value can be either the initial state of the integrator or an external value provided by an input port. Limits: the integrator can be configured such that the output and state are constrained by upper and lower limits. Hold: the integrator can be configured to hold integration based on an input trigger. The Integrator block is also designed to detect "Zeno" behavior, where the reset events happen asymptotically closer together. This is a pathological case that can cause numerical issues in the simulation and should typically be avoided by introducing some physically realistic hysteresis into the model. However, in the event that Zeno behavior is unavoidable, the integrator will enter a "Zeno" state where the output is held constant until the trigger changes value to False. See the "bouncing ball" demo for a Zeno example. Input ports (0) The input signal. Must match the shape and dtype of the initial continuous state. (1) The reset trigger. Optional, only if `enable_reset` is True. (2) The reset value. Optional, only if `enable_external_reset` is True. (3) The hold trigger. Optional, only if 'enable_hold' is True. Output ports (0) The continuous state of the integrator. Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `initial_state` | | The initial value of the integrator state. Can be any array, or even a nested structure of arrays, but the data type should be floating-point. | *required* | | `enable_reset` | | If True, the integrator will reset its state to the initial value when the reset trigger is True. Adds an additional input port for the reset trigger. This signal should be boolean- or binary-valued. | `False` | | `enable_external_reset` | | If True, the integrator will reset its state to the value provided by the reset value input port when the reset trigger is True. Otherwise, the integrator will reset to the initial value. Adds an additional input port for the reset value. This signal should match the shape and dtype of the initial continuous state. | `False` | | `enable_limits` | | If True, the integrator will constrain its state and output to within the upper and lower limits. Either limit may be disbale by setting its value to None. | `False` | | `enable_hold` | | If True, the integrator will hold integration when the hold trigger is True. | `False` | | `reset_on_enter_zeno` | | If True, the integrator will reset its state to the initial value when the integrator enters the Zeno state. This option is ignored unless enable_reset is True. | `False` | | `zeno_tolerance` | | The tolerance used to determine if the integrator is in the Zeno state. If the time between events is less than this tolerance, then the integrator is in the Zeno state. This option is ignored unless enable_reset is True. | `1e-06` | Events An event is triggered when the "reset" port changes. An event is triggered when the state hit one of the limits. An event is triggered when the "hold" port changes. Another guard is conditionally active when the integrator is in the Zeno state, and is triggered when the "reset" port changes from True to False. This event is used to exit the Zeno state and resume normal integration. ### `IntegratorDiscrete` Bases: `LeafSystem` Discrete first-order integrator. This block is a discrete-time approximation to the behavior of the Integrator block. It implements the following linear time-invariant difference equation for input values `u` and output values `y`: ``` x[k+1] = x[k] + dt * u[k] y[k] = x[k] ``` where `x` is the state of the integrator. The integrator is initialized with the value of the `initial_state` parameter. Options Reset: the integrator can be configured to reset its state on an input trigger. The reset value can be either the initial state of the integrator or an external value provided by an input port. Limits: the integrator can be configured such that the output and state are constrained by upper and lower limits. Hold: the integrator can be configured to hold integration based on an input trigger. Unlike the continuous-time integrator, the discrete integrator does not detect Zeno behavior, since this is not a concern in discrete-time systems. Input ports (0) The input signal. Must match the shape and dtype of the initial state. (1) The reset trigger. Optional, only if `enable_reset` is True. (2) The reset value. Optional, only if `enable_external_reset` is True. (3) The hold trigger. Optional, only if 'enable_hold' is True. Output ports (0) The current state of the integrator. Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `initial_state` | | The initial value of the integrator state. Can be any array, or even a nested structure of arrays, but the data type should be floating-point. | *required* | | `enable_reset` | | If True, the integrator will reset its state to the initial value when the reset trigger is True. Adds an additional input port for the reset trigger. This signal should be boolean- or binary-valued. | `False` | | `enable_external_reset` | | If True, the integrator will reset its state to the value provided by the reset value input port when the reset trigger is True. Otherwise, the integrator will reset to the initial value. Adds an additional input port for the reset value. This signal should match the shape and dtype of the initial continuous state. | `False` | | `enable_limits` | | If True, the integrator will constrain its state and output to within the upper and lower limits. Either limit may be disbale by setting its value to None. | `False` | | `enable_hold` | | If True, the integrator will hold integration when the hold trigger is True. | `False` | ### `InterpolationUsingPrelookup` Bases: `LeafSystem` Interpolate a static `output_array` using a precomputed (index, fraction) tuple from :class:`Prelookup`. This is the downstream half of the standard `Prelookup`/`InterpolationUsingPrelookup` pair. Plug as many of these as you like into a single :class:`Prelookup`'s output port -- each one interpolates its OWN `output_array` against the shared (index, fraction) signal, avoiding the redundant binary searches you would do with N independent :class:`LookupTable1d` blocks. Input ports `(0)` -- the (index, fraction) NamedTuple produced by an upstream :class:`Prelookup` block. Output ports `(0)` -- the linearly-interpolated value `(1 - alpha) * output_array[i] + alpha * output_array[i + 1]`. Parameters: | Name | Type | Description | Default | | --------------- | ----------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `output_array` | | 1-D table values, shape (N,) where N == len(prelookup input_array). output_array[i] corresponds to the i-th breakpoint in the upstream :class:Prelookup's grid. | *required* | | `dtype` | `optional` | If set (e.g. jnp.float32), the table values are cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d. | `None` | | `extrapolation` | `(optional, T - 114 - fu - prelookup - extrap)` | Out-of-range policy. Must match the upstream :class:Prelookup's extrapolation kwarg (the math is all done in the producer's alpha computation; this block just consumes alpha). Exists on this side of the API for symmetry, IDE discoverability and so model JSON / user code make intent explicit. One of "clip" (default, byte-equivalent to T-114-fu-prelookup), "linear", or "nan". | `'clip'` | Notes Differentiable through `output_array` (every time-step the interpolation is a convex combination of two of its entries; the gradient w.r.t. the picked entries is exact). The fraction-side gradient flows through the upstream :class:`Prelookup`'s `alpha` computation; the discrete `index` is piecewise constant (expected -- same as every `searchsorted`-based block in the library). Today only linear interpolation is supported -- PCHIP/Akima downstream interpolation is a deeper followup (`T-114-followup-prelookup-cubic`). #### `extrapolation` The declared OOB policy (must match the upstream Prelookup). #### `output_array` The 1-D table being interpolated. ### `InverseClarke` Bases: `FeedthroughBlock` Inverse Clarke transform: `[alpha, beta]` -> three-phase `[a, b, c]`. Amplitude-invariant, zero-sequence-free: `a + b + c = 0` by construction, and `InverseClarke(Clarke(abc)) == abc` for any zero-sequence-free input. ### `InversePark` Bases: `_AngleTransform` Inverse Park transform: rotor `[d, q]` -> stationary `[alpha, beta]`. Input ports (0) dq: rotor-frame vector. (1) theta: electrical angle `theta_e` (rad). Output ports (0) alpha_beta: `[cos*d - sin*q, sin*d + cos*q]`. ### `KalmanFilter` Bases: `KalmanFilterBase` Kalman Filter for the following system: ``` x[n+1] = A x[n] + B u[n] + G w[n] y[n] = C x[n] + D u[n] + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q E(v[n]v'[n] = R E(w[n]v'[n] = N = 0 ``` Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | --------- | ---- | ----------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `A` | | ndarray State transition matrix | *required* | | `B` | | ndarray Input matrix | *required* | | `C` | | ndarray Output matrix. If None, full state output is assumed. | `None` | | `D` | | ndarray Feedthrough matrix. If None, no feedthrough is assumed. | `None` | | `G` | | ndarray Process noise matrix. If None, G=B is assumed. | `None` | | `Q` | | ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and A is assumed. | `None` | | `R` | | ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with C and A is assumed. | `None` | | `x_hat_0` | | ndarray Initial state estimate. If None, an array of zeros is assumed. | `None` | | `P_hat_0` | | ndarray Initial state covariance matrix estimate. If None, Identity matrix of size identical to A is assumed. | `None` | #### `for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_bar_0=None, P_hat_bar_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None)` Obtain a Kalman Filter system for a continuous-time plant after linearization at equilibrium point (x_eq, u_eq) The input plant contains the deterministic forms of the forward and observation operators: ``` dx/dt = f(x,u) y = g(x,u) ``` Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator. A plant with disturbances of the following form is then considered following form: ``` dx/dt = f(x,u) + G w --- (C1) y = g(x,u) + v --- (C2) ``` where: ``` `w` represents the process noise, `v` represents the measurement noise, ``` and ``` E(w) = E(v) = 0 E(ww') = Q E(vv') = R E(wv') = N = 0 ``` This plant with disturbances is linearized (only `f` and `g`) around the equilibrium point to obtain: ``` d/dt (x_bar) = A x_bar + B u_bar + G w y_bar = C x_bar + D u_bar + v ``` where, ``` x_bar = x - x_eq u_bar = u - u_eq y_bar = y - y_bar y_eq = g(x_eq, u_eq) ``` The linearized plant is then discretized via `euler` or `zoh` method to obtain: ``` x_bar[n] = Ad x_bar[n] + Bd u_bar[n] + Gd w[n] --- (L1) y_bar[n] = Cd x_bar[n] + Dd u_bar[n] + v[n] --- (L2) E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Qd E(v[n]v'[n]) = Rd E(w[n]v'[n]) = Nd = 0 ``` Note: If `discretized_noise` is True, then it is assumed that the user is providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to Identity matrix. A Kalman Filter estimator for the system of equations (L1) and (L2) is created and returned. This filter is in the `x_bar`, `u_bar`, and `y_bar` states. This returned system will have Input ports (0) u_bar[n] : control vector at timestep n, relative to equilibrium (1) y_bar[n] : measurement vector at timestep n, relative to equilibrium Output ports (1) x_hat_bar[n] : state vector estimate at timestep n, relative to equilibrium Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. | *required* | | `x_eq` | | ndarray Equilibrium state vector for discretization | *required* | | `u_eq` | | ndarray Equilibrium control vector for discretization | *required* | | `dt` | | float Time step for the discretization. | *required* | | `Q` | | ndarray Process noise covariance matrix. If None, Identity matrix of size compatible with G and and linearized system's A is assumed. | `None` | | `R` | | ndarray Measurement noise covariance matrix. If None, Identity matrix of size compatible with linearized system's C and A is assumed. | `None` | | `G` | | ndarray Process noise matrix. If None, G=B is assumed making disrurbances additive to control vector u, i.e. u_disturbed = u_orig + w. | `None` | | `x_hat_bar_0` | | ndarray Initial state estimate, relative to equilirium. If None, an identity matrix is assumed. | `None` | | `P_hat_bar_0` | | ndarray Initial covariance matrix estimate for state, relative to equilibrium. If None, an Identity matrix is assumed. | `None` | | `discretization_method` | | str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler". | `'euler'` | | `discretized_noise` | | bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G, Q, and R are assumed to be Gd, Qd, and Rd, respectively. | `False` | #### `global_filter_for_continuous_plant(plant, x_eq, u_eq, dt, Q=None, R=None, G=None, x_hat_0=None, P_hat_0=None, discretization_method='euler', discretized_noise=False, name=None, ui_id=None)` See docs for `for_continuous_plant`, which returns the local Kalman Filter. This method additionally converts the local Kalman Filter to a global estimator. See docs for `make_global_estimator_from_local` for details. ### `KoopmanPredictor` Bases: `LeafSystem` Discrete-time Koopman predictor for a nonlinear system. Each step: lift the stored physical state `z = g(x)`, advance the lifted *linear* dynamics `z[k+1] = K z[k] (+ B u[k])`, then de-lift `x[k+1] = C z[k+1]`. The physical state is the block output. Because the lifted model is *linear* in `z`, `(K, B)` is a discrete linear model you can use for linear MPC / LQR-style control — design in lifted coordinates and de-lift with `C` (the Koopman→linear-MPC framing; Korda & Mezić 2018). Note a lifted model generally has no state that both satisfies `z = g(x)` and a hard terminal-equality constraint, so a *terminal-cost* MPC is the right fit; the current :class:`~jaxonomy.library.mpc.LinearDiscreteTimeMPC` block (hard terminal equality, continuous-time model input) does not compose directly — see the `rom_dmdc_koopman_mpc` example, which uses a compact terminal-cost MPC. An input port (and use of `B`) is created only when `B` is provided. Input ports (0) u\[k\]: control input, present iff `B` is given. Output ports (0) x\[k\]: de-lifted physical state. Parameters: | Name | Type | Description | Default | | --------------- | ---- | ----------------------------------------------------------------------- | ---------- | | `K` | | Koopman operator (L, L) — a dynamic parameter. | *required* | | `C` | | De-lift matrix (n, L) — a dynamic parameter. | *required* | | `dictionary` | | Observable dictionary g used for lifting (identity first). | *required* | | `B` | | Optional lifted input operator (L, m) — a dynamic parameter when given. | `None` | | `dt` | | Sampling period of the discrete update. | `1.0` | | `initial_state` | | Initial physical state x[0] of size n. | `None` | ### `LTISystem` Bases: `LTISystemBase` Continuous-time linear time-invariant system. Implements the following system of ODEs: ``` ẋ = Ax + Bu y = Cx + Du ``` Input ports (0) u: Input vector of size m Output ports (0) y: Output vector of size p. Note that this is feedthrough from the input port if and only if D is nonzero. Parameters: | Name | Type | Description | Default | | ------------------- | ---- | ------------------------------------------- | ---------- | | `A` | | State matrix of size n x n | *required* | | `B` | | Input matrix of size n x m | *required* | | `C` | | Output matrix of size p x n | *required* | | `D` | | Feedthrough matrix of size p x m | *required* | | `initialize_states` | | Initial state vector of size n (default: 0) | `None` | #### `ss` State-space representation of the system. ### `LTISystemDiscrete` Bases: `LTISystemBase` Discrete-time linear time-invariant system. Implements the following system of ODEs: ``` x[k+1] = A x[k] + B u[k] y[k] = C x[k] + D u[k] ``` Input ports (0) u\[k\]: Input vector of size m Output ports (0) y\[k\]: Output vector of size p. Note that this is feedthrough from the input port if and only if D is nonzero. Parameters: | Name | Type | Description | Default | | ------------------- | ---- | ------------------------------------------- | ---------- | | `A` | | State matrix of size n x n | *required* | | `B` | | Input matrix of size n x m | *required* | | `C` | | Output matrix of size p x n | *required* | | `D` | | Feedthrough matrix of size p x m | *required* | | `dt` | | Sampling period | *required* | | `initialize_states` | | Initial state vector of size n (default: 0) | `None` | #### `ss` State-space representation of the system. ### `LeadLag` Bases: `LeafSystem` Discrete first-order lead-lag compensator. Discretisation of the continuous compensator `G(s) = K * (1 + T_lead * s) / (1 + T_lag * s)` via the Tustin (bilinear) transform with `s = (2/dt)*(z-1)/(z+1)`. The resulting difference equation is:: ``` c = 2 / dt den = 1 + T_lag * c b0 = K * (1 + T_lead * c) / den b1 = K * (1 - T_lead * c) / den a1 = (1 - T_lag * c) / den y[k] = b0 * x[k] + b1 * x[k-1] - a1 * y[k-1] ``` With `T_lead = T_lag` the s-domain pole and zero cancel and the block reduces to a pure gain `K` (used as an identity check in the corpus). Input ports (0) The input signal `x`. Output ports (0) The compensated signal `y`. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------- | ---------- | | `dt` | | Sampling period of the block (s). | *required* | | `K` | | Compensator gain. Differentiable. | `1.0` | | `T_lead` | | Lead time constant (s). Differentiable. | `1.0` | | `T_lag` | | Lag time constant (s). Must be > 0. Differentiable. | `1.0` | | `initial_state` | | Initial value of y[-1]. Default 0.0. | `0.0` | Notes Differentiability: `K`, `T_lead` and `T_lag` flow into the biquad coefficients via smooth arithmetic, so `jax.grad` is finite through them. The block is feedthrough on its input port (`b0 != 0` whenever `K != 0`), so `y[k]` depends on `x[k]` directly. ### `LinearDiscreteTimeMPC` Bases: `LeafSystem` Model predictive control for a linear discrete-time system. Solves a constrained quadratic program at each time step using OSQP via `jax.pure_callback`, making it compatible with JAX's JIT compiler. Notes This block is *feedthrough*: the QP solver runs every time the output port is evaluated. Pair with a zero-order hold so the solver is invoked only once per MPC step. Data-driven discrete-time operators (:func:`~jaxonomy.library.rom.dmdc`, :func:`~jaxonomy.library.rom.edmd`) compose directly: wrap the fitted `A`, `B` as `LinearizedSystem(A, B, C, D, {}, dt=dt)` (`C`/`D` are unused by the MPC) and pass the result — no manual "un-discretization" needed. Parameters: | Name | Type | Description | Default | | ------------ | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `lin_sys` | | Prediction model. A continuous-time model (LTISystem block, or LinearizedSystem with dt=None) is discretized internally by forward Euler at dt. An already-discrete model (LinearizedSystem with dt set — e.g. from discretize(...) or a dmdc/edmd fit — or an LTISystemDiscrete block) has its A, B used as-is; its dt must equal the MPC dt. | *required* | | `Q` | | State cost matrix (n×n). | *required* | | `R` | | Input cost matrix (m×m). | *required* | | `N` | | Prediction horizon (number of steps). | *required* | | `dt` | | MPC sampling period (also the Euler discretization step for a continuous-time lin_sys). | *required* | | `x_ref` | | Terminal state reference (length-n array). | *required* | | `lbu` | | Lower bound on control input (scalar or length-m array). | `-inf` | | `ubu` | | Upper bound on control input (scalar or length-m array). | `inf` | | `warm_start` | | Whether to warm-start the OSQP solver between solves. | `False` | ### `LinearDiscreteTimeMPC_OSQP` Bases: `LinearDiscreteTimeMPC` Deprecated alias for :class:`LinearDiscreteTimeMPC`. Both classes now use OSQP via `jax.pure_callback`. Use :class:`LinearDiscreteTimeMPC` directly. ### `LinearQuadraticGaussian` Bases: `LeafSystem` Continuous-time infinite-horizon LQG controller (separation principle). The observer uses the algebraic Riccati solution for the Kalman gain `L` given `(A, G, C, Qn, Rn)`; the regulator uses the algebraic Riccati solution for the feedback gain `K` given `(A, B, Qc, Rc)`. Both use the `control` library's `lqe` / `lqr` helpers. Input / output shapes follow the plant: `y` is `(ny,)`, `u` is `(nu,)`, and the observer's internal state is `(nx,)`. ### `LinearQuadraticRegulator` Bases: `FeedthroughBlock` Linear Quadratic Regulator (LQR) for a continuous-time system: dx/dt = A x + B u. Computes the optimal control input: u = -K x, where u minimises the cost function over \[0, ∞)\]: J = ∫(x.T Q x + u.T R u) dt. Input ports (0) x: state vector of the system. Output ports (0) u: optimal control vector. Parameters: | Name | Type | Description | Default | | ---- | ---- | --------------------------------- | ---------- | | `A` | | Array State matrix of the system. | *required* | | `B` | | Array Input matrix of the system. | *required* | | `Q` | | Array State cost matrix. | *required* | | `R` | | Array Input cost matrix. | *required* | ### `LinearizedSystem` State-space linearization result. For a continuous-time linsys (`dt is None`): ``` dx/dt = Ax + Bu, y = Cx + Du ``` For a discrete-time linsys (`dt` is a positive float — produced by :func:`jaxonomy.library.linearization_workflow.discretize`): ``` x[k+1] = Ax[k] + Bu[k], y[k] = Cx[k] + Du[k] ``` Attributes: | Name | Type | Description | | ----------------- | ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | | `A` | `Any` | State matrix (n_states, n_states) | | `B` | `Any` | Input matrix (n_states, n_inputs) | | `C` | `Any` | Output matrix (n_outputs, n_states) | | `D` | `Any` | Feedthrough matrix (n_outputs, n_inputs) | | `operating_point` | `dict` | dict with state and input values used for linearization | | `dt` | `Optional[float]` | Sampling period in seconds when the linsys is discrete-time; None for continuous-time. Default None so existing constructor calls remain byte-equivalent. | #### `is_discrete` True if this LinearizedSystem carries a sampling period. #### `is_stable` True if the system is stable. For continuous-time linsys (`dt is None`) the criterion is `Re(eig(A)) < 0`; for a discrete-time linsys (`dt is not None`) it is `|eig(A)| < 1`. #### `eigenvalues()` Compute eigenvalues of A matrix. Returns complex array of shape (n_states,). #### `to_lti()` Convert to Jaxonomy LTISystem block. #### `to_scipy_lti()` Convert to :class:`scipy.signal.StateSpace` for frequency-domain analysis. Returns a `scipy.signal.StateSpace` object, which supports MIMO systems (multiple inputs and/or multiple outputs) as well as SISO systems. Use this for Bode plots, Nyquist diagrams, step/impulse responses, etc. Note `scipy.signal.lti` is SISO-only and is **not** used here. `scipy.signal.StateSpace` (a subclass of `lti`) is the correct target for state-space (A, B, C, D) representations with any number of I/Os. Requires scipy to be installed. ### `Logarithm` Bases: `FeedthroughBlock` Compute the logarithm of the input signal. This block dispatches to `jax.numpy.log`, `jax.numpy.log2`, or `jax.numpy.log10`, so the semantics, broadcasting rules, etc. are the same. See the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.log.html https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.log2.html https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.log10.html Input ports (0) The input signal. Output ports (0) The logarithm of the input signal. Parameters: | Name | Type | Description | Default | | ------ | ---- | ----------------------------------------------------------------------------------------------- | ----------- | | `base` | | One of "natural", "2", or "10". Determines the base of the logarithm. The default is "natural". | `'natural'` | ### `LogicalOperator` Bases: `LeafSystem` Apply a boolean function elementwise to the input signals. This block implements the following boolean functions - "or": same as np.logical_or - "and": same as np.logical_and - "not": same as np.logical_not - "nor": equivalent to np.logical_not(np.logical_or(in_0,in_1)) - "nand": equivalent to np.logical_not(np.logical_and(in_0,in_1)) - "xor": same as np.logical_xor Input ports (0,1) The input signals. If numeric, they are interpreted as boolean types (so 0 is False and any other value is True). Output ports (0) The result of the logical operation, a boolean-valued signal. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ---------------------------------------------------------------------------------- | ---------- | | `function` | | The boolean function to apply. One of "or", "and", "not", "nor", "nand", or "xor". | *required* | Events An event is triggered when the output changes from True to False or vice versa. ### `LogicalReduce` Bases: `FeedthroughBlock` Apply a boolean reduce function to the elements of the input signal. This block implements the following boolean functions - "any": Output is True if any input element is True. - "all": Output is True if all input elements are True. Input ports (0) The input signal. If numeric, they are interpreted as boolean types (so 0 is False and any other value is True). Output ports (0) The result of the logical operation, a boolean-valued signal. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------- | ---------- | | `function` | | The boolean function to apply. One of "any", "all". | *required* | | `axis` | | Axis or axes along which a logical OR/AND reduction is performed. | `None` | Events An event is triggered when the output changes from True to False or vice versa. ### `LookupTable1d` Bases: `FeedthroughBlock` Interpolate the input signal into a static lookup table. If a function `y = f(x)` is sampled at a set of points `(x_i, y_i)`, then this block will interpolate the input signal `x` to compute the output signal `y`. The behavior is modeled after `scipy.interpolate.interp1d` but is implemented in JAX. Available interpolation modes are: - "linear": Linear interpolation using `jax.interp`. - "pchip": Monotone cubic Hermite (Hyman/Fritsch-Carlson). T-106 phase 1 — smooth gradients everywhere, monotone on monotone data. - "akima": Akima 1970 cubic spline (T-114 phase 2). Smoother than PCHIP on non-monotone data, less prone to overshoot than natural cubic splines. Matches `scipy.interpolate.Akima1DInterpolator`. - "cubic": Natural cubic spline (T-114-followup-natural-cubic- spline). C^2-continuous, second derivative zero at the boundaries — the smoothest possible C^2 interpolant. Requires at least 4 breakpoints. Matches `scipy.interpolate.CubicSpline(bc_type='natural')`. - "nearest": Nearest-neighbor interpolation. - "flat": Flat interpolation. Input ports (0) The input signal, which is used as the interpolation coordinate. Output ports (0) The interpolated output signal. Parameters: | Name | Type | Description | Default | | --------------- | ------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `input_array` | | The array of input values at which the output values are provided. | *required* | | `output_array` | | The array of output values. | *required* | | `interpolation` | | One of "linear", "pchip", "nearest", or "flat". Determines the type of interpolation performed by the block. | *required* | | `extrapolation` | `optional, T-114 phase 1` | One of "clip" (default — standard lookup-table behaviour, holds the boundary value), "linear" (extends the boundary slope past the breakpoints), or "nan" (returns NaN outside \[input_array[0], input_array[-1]\]). Stored outside @parameters, so this kwarg is not round-tripped through model JSON — pre-existing models reload byte-equivalently. | `'clip'` | | `dtype` | `optional, T-038a` | If set (e.g. jnp.float32), the block's input_array and output_array are cast to this dtype on construction, so the interpolation arithmetic — and therefore the output signal — runs at this precision regardless of the global x64 setting. The default (None) preserves the pre-T-038a behavior: the arrays are stored verbatim and float64 is used under the default x64-enabled install. T-038a-followup-mixed-precision-cascade: when dtype is None and a :func:jaxonomy.precision_policy context manager is active, the block falls back to the context's dtype. Explicit dtype= always wins (explicit-over-implicit). Best-effort: this enforces dtype on the block's internal arrays and on the output of jnp.interp / jnp.argmin lookups, but downstream operations (e.g. connecting to a default-dtype block) are subject to JAX's standard promotion rules — the result of the wider arithmetic may be promoted to float64. | `None` | Notes Currently restricted to 1D input and output data. This may be expanded to support multi-dimensional output arrays in the future. #### `fit_from_data(xp, x_data, y_data, *, weights=None, smoothness=0.0, **block_kwargs)` Build a `LookupTable1d` whose output values are fitted by least squares to `(x_data, y_data)` at the fixed grid `xp`. Ergonomic wrapper around :func:`jaxonomy.library.fit_lookup_table_1d` so the fitting entry point is discoverable from the block class itself. Parameters: | Name | Type | Description | Default | | ---------------- | ------- | ------------------------------------------------------------------------------------------------------------- | ---------- | | `xp` | | Fixed grid of breakpoints (1-D, strictly increasing). | *required* | | `x_data` | | Measured input cloud, shape (K,). | *required* | | `y_data` | | Measured output cloud, shape (K,). | *required* | | `weights` | | Optional per-sample weights for weighted least squares. None = OLS. | `None` | | `smoothness` | `float` | Non-negative discrete first-difference penalty. Use small values (1e-3 .. 1.0) on noisy / sparse data. | `0.0` | | `**block_kwargs` | | Forwarded to :func:jaxonomy.library.fit_lookup_table_1d (e.g. interpolation=, extrapolation=, name=, dtype=). | `{}` | Returns: | Type | Description | | ---- | ------------------------------------------------ | | | A LookupTable1d instance with input_array=xp and | | | output_array set to the LS-fit table values. | ### `LookupTable2d` Bases: `LeafSystem` Interpolate the input signals into a static lookup table. The behavior is modeled on `scipy.interpolate.interp2d` but is implemented in JAX. `"linear"` (bilinear, default) and `"bicubic"` (Catmull-Rom) interpolation are supported. The input arrays must be 1D and the output array must be 2D. Input ports (0) The first input signal, used as the first interpolation coordinate. (1) The second input signal, used as the second interpolation coordinate. Output ports (0) The interpolated output signal. Parameters: | Name | Type | Description | Default | | -------------------- | ------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `input_x_array` | | The array of input values at which the output values are provided, corresponding to the first input signal. Must be 1D | *required* | | `input_y_array` | | The array of input values at which the output values are provided, corresponding to the second input signal. Must be 1D | *required* | | `output_table_array` | | The array of output values. Must be 2D with shape (m, n), where m = len(input_x_array) and n = len(input_y_array). | *required* | | `interpolation` | | "linear" (default, standard bilinear behaviour) or "bicubic" (T-114-followup-2d-bicubic — Catmull-Rom cubic-convolution kernel; C^1-continuous, exact at grid corners, smoother than bilinear off-grid). "bicubic" requires at least 4 breakpoints per axis. | `'linear'` | | `extrapolation` | `optional, T-114 phase 2` | One of "clip" (default — standard lookup-table behaviour, holds the boundary value), "linear" (bilinear extension past the grid via edge-slope continuation), or "nan" (returns NaN outside the grid). The default "clip" matches the legacy npa.interp2d behaviour byte-equivalently. Stored outside @parameters so this kwarg is not round-tripped through model JSON; pre-existing models reload byte- equivalently. | `'clip'` | #### `fit_from_data(xp, yp, x_data, y_data, z_data, *, weights=None, smoothness=0.0, **block_kwargs)` Build a `LookupTable2d` whose table values are fitted by bilinear least squares to `(x_data, y_data, z_data)` at the fixed grid `(xp, yp)`. Ergonomic classmethod mirror of :func:`jaxonomy.library.fit_lookup_table_2d` so the 2-D fitting entry point is discoverable from the block class itself. Parameters: | Name | Type | Description | Default | | ------------------------ | ------- | ------------------------------------------------------------------------------------------------------------- | ---------- | | `xp` | | Fixed grid along the first axis (1-D, strictly increasing). | *required* | | `yp` | | Fixed grid along the second axis (1-D, strictly increasing). | *required* | | `x_data, y_data, z_data` | | Measurement cloud, all shape (K,). | *required* | | `weights` | | Optional per-sample weights for weighted least squares. None = OLS. | `None` | | `smoothness` | `float` | Non-negative 5-point Laplacian penalty on the fitted table. 0.0 (default) is pure data-fit. | `0.0` | | `**block_kwargs` | | Forwarded to :func:jaxonomy.library.fit_lookup_table_2d (e.g. interpolation=, extrapolation=, name=, dtype=). | `{}` | Returns: | Type | Description | | ---- | --------------------------------------------------- | | | A LookupTable2d instance with input_x_array=xp, | | | input_y_array=yp, and output_table_array set to the | | | LS-fit table of shape (len(xp), len(yp)). | ### `LookupTableND` Bases: `LeafSystem` Interpolate the input signal into a static N-D lookup table. Generalises :class:`LookupTable1d` and :class:`LookupTable2d` to an arbitrary number of axes. The block takes a single input port whose value is a length-`N` query vector `[q_1, ..., q_N]` and returns the multilinearly interpolated table value at that point. Implementation: delegates to :func:`jaxonomy.library.lookup_table.interp_nd`, which performs `N` successive 1-D linear interpolations along each axis (no `jnp.interpn` exists today — see the deeper-followup note in that function's docstring). Input ports `(0)` — the query vector, shape `(N,)` where `N` is the number of grid axes. Output ports `(0)` — the multilinearly interpolated table value. Parameters: | Name | Type | Description | Default | | --------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `grid_axes` | | Tuple of N strictly-increasing 1-D breakpoint arrays. The i-th array has length B_i and corresponds to axis i of output_array. | *required* | | `output_array` | | Sample values, shape (B_1, B_2, ..., B_N). output_array[i_1, ..., i_N] = f(grid_axes[0][i_1], ..., grid_axes[N-1][i_N]). | *required* | | `interpolation` | | Currently only "linear" (multilinear). Reserved for future N-D smooth methods (filed under T-114-followup-phase4-nd-cubic). | `'linear'` | | `extrapolation` | `optional` | One of "clip" (default — clips each coordinate to its axis range, the standard lookup-table default), "linear" (multilinear continuation past the grid), or "nan" (returns NaN whenever any coordinate is outside its axis range). | `'clip'` | | `dtype` | `optional` | If set (e.g. jnp.float32), the block's grid arrays and output array are cast to this dtype on construction. Mirrors the per-block dtype contract of LookupTable1d. | `None` | Notes Differentiable through both the query vector and the table values (modulo the discrete bucket-index `searchsorted`, whose gradient is piecewise constant — within a cell the gradient is exact). ### `LowPassDiscrete` Bases: `LeafSystem` Discrete first-order (single-pole RC) low-pass filter. Implements the difference equation:: ``` tau = 1 / (2*pi*cutoff_hz) alpha = dt / (dt + tau) y[k] = alpha * x[k] + (1 - alpha) * y[k-1] ``` The continuous-time analogue is `H(s) = 1 / (1 + tau*s)`, with -3 dB crossover at `f = cutoff_hz` in the small-`dt` limit. Input ports (0) The input signal `x`. Output ports (0) The filtered signal `y`. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------- | ---------- | | `dt` | | Sampling period of the block (s). | *required* | | `cutoff_hz` | | Design cutoff frequency (Hz). Differentiable. | `1.0` | | `initial_state` | | Initial value of y[-1]. Default 0.0. | `0.0` | Notes Differentiability: `cutoff_hz` enters the recursive update via smooth arithmetic (`alpha = dt/(dt + 1/(2*pi*cutoff_hz))`), so `jax.grad` is finite through it. ### `Luenberger` Bases: `LeafSystem` Discrete-time Luenberger observer with user-supplied gain `L`. State-update equation: .. code-block:: text ``` x_hat[k+1] = A·x_hat[k] + B·u[k] + L·(y[k] - C·x_hat[k] - D·u[k]) ``` Parameters: | Name | Type | Description | Default | | ------------ | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | Discrete sample period (seconds). Must match the plant's sample period (or, for continuous plants, the discretisation period chosen for the design). | *required* | | `A, B, C, D` | | Plant state-space matrices (typically the output of jaxonomy.discretize(linearize(plant, op), dt)). D defaults to a zero matrix of compatible shape. | *required* | | `L` | | Observer gain matrix, shape (n_states, n_outputs). The caller computes this — e.g. via scipy.signal.place for pole placement, or via a steady-state Kalman gain. | *required* | | `x_hat_0` | | Initial state estimate. Defaults to zeros. | `None` | Input ports (0) `u`: control input vector, shape `(n_inputs,)`. (1) `y`: noisy measurement vector, shape `(n_outputs,)`. Output ports (0) `x_hat`: state estimate vector, shape `(n_states,)`. Notes This block is the simpler half of the Kalman pair — the design cost (computing `L`) is paid offline, leaving only the cheap runtime update. If you want online Riccati-based gain updates instead, use :class:`KalmanFilter`. If you have a continuous plant and want the steady-state infinite-horizon Kalman gain, use :class:`InfiniteHorizonKalmanFilter`. ### `MJX` Bases: `MuJoCoBase` A system that wraps a MuJoCo model and provides a continuous-time ODE LeafSystem. Currently only supports a single body system. Input ports (0) The control input vector `control`. Output ports (0) The generalized position coordinates `qpos`. (1) The generalized velocity coordinates `qvel`. (2) The actuator coordinates `act`. (3) The sensor data `sensor_data` (if enabled). (4) The video output `video` as RGB frames of shape (H,W,3) (if enabled). (5) A fake output port, present only if `vHIL=True` and outputs Array(0.0). (6+) Custom output ports, defined with user-specified python scripts. Parameters: | Name | Type | Description | Default | | --------------------- | ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------ | | `file_name` | `str` | The path to the MuJoCo XML model file. | *required* | | `dt` | `float` | If None, jaxonomy's internal solver will be used and this block can be considered as a continuous block. If set, the model will be run in a discrete mode with the specified timestep, using MJX's solver, more like Co-Simulation. In that case, it might be favorable to set use_mjx=False. | `None` | | `key_frame_0` | \`int | str\` | The keyframe to initialize the model from. | | `qpos_0` | `Array` | The initial generalized position coordinates. | `None` | | `qvel_0` | `Array` | The initial generalized velocity coordinates. | `None` | | `act_0` | `Array` | The initial actuator coordinates. | `None` | | `enable_sensor_data` | `bool` | Whether to output the sensor data to an optional port named 'sensor_data'. | `False` | | `enable_video_output` | `bool` | Whether to output the rendered video frames to an optional port named 'video'. | `False` | | `video_size` | `tuple[int, int]` | The size of the video output frames as a (H,W) tuple. | `None` | | `enable_mocap_pos` | `bool` | Whether to enable the mocap_pos input port for motion capture tracking. | `False` | | `vHIL` | `bool` | Whether to run in virtual hardware-in-the-loop mode. | `False` | | `vHIL_dt` | `float` | The timestep for the virtual hardware-in-the-loop mode. | `0.01` | Notes: (i) `_model` and `_data` refer to MuJoCo's `mjModel` and `mjData` objects respectively. `model` and `data` are the corresponding MJX objects. (ii) While `sensordata` output is supported as a pure callback to MuJoCo since MJX has not yet implemented this aspect. This can be expensive. #### `normalize_qpos_quat(qpos)` Normalize the quaternion components of the generalized position coordinates. ### `MLP` Bases: `FeedthroughBlock` A feedforward neural network block representing an Equinox multi-layer perceptron (MLP). The output `y` of the MLP is computed as ``` y = MLP(x, theta) ``` where `theta` are the parameters of the MLP, and `x` is the input to the MLP. This block is differentialble w.r.t. the MLP parameters `theta`. Note that `theta`, does not include the hyperparameters representing the architecture of the MLP. Input ports (0) The input to the MLP. Output ports (0) The output of the MLP. Parameters: | Name | Type | Description | Default | | ---------------------- | ------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------ | | `in_size` | `int` | The dimension of the input to the MLP. | `None` | | `out_size` | `int` | The dimension of the output of the MLP. | `None` | | `width_size` | `int` | The width of every hidden layers of the MLP. | `None` | | `depth` | `int` | The depth of the MLP. This represents the number of hidden layers, including the output layer. | `None` | | `seed` | `int` | The seed for the random number generator for initialization of the MLP parameters (weights and biases of every layer). If None, a random 32-bit seed will be generated. | `None` | | `activation_str` | `str` | The activation function to use after each internal layer of the MLP. Possible values are "relu", "sigmoid", "tanh", "elu", "swish", "gelu", "leaky_relu", "rbf", and "identity". Default is "relu". | `'relu'` | | `final_activation_str` | `str` | The activation function to use for the output layer of the MLP. Same choices as activation_str. Default is "identity". | `'identity'` | | `use_bias` | `bool` | Whether to add a bias to the internal layers of the MLP. Default is True. | `True` | | `use_final_bias` | `bool` | Wheter to add a bias to the output layer of the MLP. Default is True. | `True` | | `file_name` | `str` | Optional file name containing the serialized parameters of the MLP. If provided, the parameters are loaded from the file, and set as the parameters of the MLP. Default is None. | `None` | #### `mlp` The underlying Equinox `eqx.nn.MLP` object. Built lazily in :meth:`initialize`, which runs the first time `create_context()` is called on a diagram containing this block. Accessing it before then raises a clear error pointing at `create_context()` (T-B4-followup-mlp-pre-context). #### `__init__(in_size=None, out_size=None, width_size=None, depth=None, seed=None, activation_str='relu', final_activation_str='identity', use_bias=True, use_final_bias=True, file_name=None, **kwargs)` see https://docs.kidger.site/equinox/examples/serialisation/ for rationale of implementation here. We can't serialize the activation function, so we serialize a string representing a selection for activation function amongst a finite set of options. #### `serialize(file_name, mlp_params=None)` Serialize only the parameters of the MLP. Note that the hyperparameters representing the architecture of the MLP are not serialized. This is because of the following use-cases imagined: (i) The user may train the Equinox MLP outside of Jaxonomy. In this case, it seems unnecessary to force the user to serialize the hyperparameters of the MLP in the strict form chosen by Jaxonomy. It would seem much easier for the user to just input these hyperparameters when creating the MLP block in Jaxonomy UI, and upload the naturally produced serialized parameters file by Equinox. (ii) The user may want to train the Equinox MLP within Jaxonomy in a notebook, and then use the block within Colimator UI. In this case, while serialization of the hyperparameters of the MLP would be a litte more convenient compared to manually inputting the hyperparameters in the UI, it seems like a small convenience relative to disadvantages of (i). Ideally the user should be able to use the API to push the learnt parameters. (iii) When we support training in the UI, the hyperparameters are naturally serialzed with `declare_configuraton_parameters`, and thus, in this case too, only serializatio of the MLP parameters is necessary. The choice of an optional `mlp_params` is to enable training of the models in a notebook and easily seralizing them for use in the UI. ### `MaskedDelayBuffer` Bases: `LeafSystem` Delay buffer where the delay length can be set at runtime (up to max_steps). Like ShiftRegister, but allows the delay to be specified as an input signal. Uses masking (not dynamic indexing) for JAX compatibility. Parameters: | Name | Type | Description | Default | | -------------- | ------- | ------------------------------- | ---------- | | `max_steps` | `int` | Maximum possible delay. STATIC. | *required* | | `signal_shape` | `tuple` | Shape of each signal frame. | `()` | | `dt` | `float` | Discrete update interval. | `0.01` | Ports Input[0] "u": signal to delay Input[1] "delay_steps": integer scalar, 0 < delay_steps \<= max_steps Output[0] "y": delayed signal ### `MatrixConcatenation` Bases: `ReduceBlock` Concatenate two matrices along a given axis. Dispatches to `jax.numpy.concatenate`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.concatenate.html Parameters: | Name | Type | Description | Default | | ------ | ---- | ------------------------------------------------------------------------------------------------------ | ------- | | `axis` | | The axis along which the matrices are concatenated. 0 for vertical and 1 for horizontal. Default is 0. | `0` | Input ports (0, 1) The input matrices `A` and `B` Output ports (0) The concatenation input matrices: e.g. `[A,B]`. ### `MatrixInversion` Bases: `FeedthroughBlock` Compute the matrix inverse of the input signal. Dispatches to `jax.numpy.inv`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.linalg.inv.html Input ports (0) The input matrix. Output ports (0) The inverse of the input matrix. ### `MatrixMultiplication` Bases: `ReduceBlock` Compute the matrix product of the input signals. Dispatches to `jax.numpy.matmul`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.matmul.html Input ports (0, 1) The input matrices `A` and `B` Output ports (0) The matrix product of the input matrices: `A @ B`. ### `MatrixTransposition` Bases: `FeedthroughBlock` Compute the matrix transpose of the input signal. Dispatches to `jax.numpy.transpose`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.transpose.html Input ports (0) The input matrix. Output ports (0) The transpose of the input matrix. ### `MinMax` Bases: `ReduceBlock` Return the extremum of the input signals. Input ports (0..n_in-1) The input signals. Output ports (0) The minimum or maximum of the input signals. Parameters: | Name | Type | Description | Default | | ---------- | ---- | -------------------------------------------------------------------------------------------------------- | ---------- | | `operator` | | One of "min" or "max". Determines whether the block returns the minimum or maximum of the input signals. | *required* | Events An event is triggered when the extreme input signal changes. For example, if the block is configured as a "max" block with two inputs and the second signal becomes greater than the first, a zero-crossing event will be triggered. ### `ModelicaFMU` Bases: `LeafSystem` #### `__init__(file_name, dt, name=None, input_names=None, output_names=None, parameters=None, start_time=0.0, first_step_at_zero=False, **kwargs)` Load and execute an FMU for Co-Simulation. .. warning:: **One instance per FMU per process** for FMUs built with pythonfmu (e.g. via :func:`jaxonomy.library.build_fmu`). The embedded-Python wrapper holds a process-wide `Py_Initialize` singleton, so instantiating the same `.fmu` dylib twice in one Python process fails. For multi-start or batched co-simulation, isolate each instance in its own process (`multiprocessing` / `concurrent.futures.ProcessPoolExecutor` with the *spawn* start method, or a `subprocess` running a small driver script) and aggregate results afterwards. This is an upstream pythonfmu limitation, not a Jaxonomy one; FMUs from other exporters (OpenModelica, Dymola, Reference-FMUs) do not carry it. Parameters: | Name | Type | Description | Default | | -------------------- | ----------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `file_name` | `str` | path to FMU file | *required* | | `dt` | `float` | stepsize for FMU simulation | *required* | | `name` | `str` | name of block | `None` | | `input_names` | `list[str]` | if set, only expose these inputs | `None` | | `output_names` | `list[str]` | if set, only expose these outputs | `None` | | `parameters` | `dict` | dictionary of parameter overrides | `None` | | `start_time` | `float` | FMU experiment start time. | `0.0` | | `first_step_at_zero` | `bool` | Default False. The first FMU step fires at t=dt (Modelica clocked-block convention: the periodic update is offset by one sample), so the block's outputs at t=0 reflect the FMU's initial state (post-setupExperiment, pre-step) and only switch to "post-first-step" values from t=dt onward. This introduces a one-sample phase lag versus an FMU exported with offset=0 semantics, which weakens FMU round-trip byte-equivalence. Pass first_step_at_zero=True to fire the first step at t=0 so the outputs leave their initial value as soon as the simulation begins. Surfaced as the FMU offset asymmetry in a follow-up finding. | `False` | | `kwargs` | | ignored | `{}` | ### `ModelicaFMUME` Bases: `LeafSystem` Import an FMI **model-exchange** FMU as a continuous-time block. Where :class:`ModelicaFMU` imports a co-simulation FMU (the FMU owns a solver; this block samples it on a fixed communication grid), a model-exchange FMU exposes only its right-hand side and leaves integration to the importer. Jaxonomy's solvers take that over, which buys two things over the co-simulation path: - **No communication-step error.** A co-simulation import holds its inputs constant across each step, so agreement with the exporting tool is first-order in `dt`. Here the FMU's derivatives are evaluated inside the adaptive solver and accuracy is set by `rtol`/`atol` instead. - **Events resolved by the host.** Each FMI event indicator becomes a jaxonomy zero-crossing, so the solver localizes the crossing and then runs the FMU's event iteration, rather than stepping over it. This is also the interface many tools export: OpenModelica emits model exchange, and its own importer accepts nothing else. Parameters: | Name | Type | Description | Default | | -------------- | ----------- | -------------------------------------------------- | ---------- | | `file_name` | `str` | path to the FMU file. | *required* | | `name` | `str` | name of the block. | `None` | | `input_names` | `list[str]` | if set, expose only these inputs. | `None` | | `output_names` | `list[str]` | if set, expose only these outputs. | `None` | | `parameters` | `dict` | parameter overrides applied during initialization. | `None` | | `start_time` | `float` | FMI experiment start time. | `0.0` | .. warning:: An FMU instance is a stateful C object reached through `io_callback`, so this block is **not** `vmap`-safe and is not differentiable — the same boundary caveats as :class:`ModelicaFMU`. Every derivative evaluation re-sends time, state, and inputs before reading, which keeps the callback a pure function of `(t, x, u)` and therefore safe under the rejected steps and stage evaluations of an adaptive solver. .. note:: `fmi2CompletedIntegratorStep` is not called: jaxonomy's solvers expose no accepted-step hook to drive it from. FMUs that rely on it to flag step events (rather than on event indicators) may miss those events. .. important:: Every callback here passes `ordered=True`. An FMU is a stateful C object and the event reset *mutates* it (the FMI event iteration advances discrete state), so with the default unordered `io_callback` XLA is free to schedule that mutation against the derivative, guard, and output reads in any order. The result was nondeterministic — identical inputs produced a correct bounce on one run and a ball falling through the floor on the next. Reads alone would tolerate reordering, since each re-sends `(t, x, u)` before reading; the mutation is what makes ordering load-bearing. Do not drop it. ### `MuJoCo` Bases: `MuJoCoBase` MuJoCo implementation without MJX. Refer to MJX for the main docs. Unlike the MJX variant of the block, this version uses the solver provided by mujoco itself and the physics are fully handled by mujoco. This behaves like a Co-Simulation environment. This variant may be used to speed up compilation times or in situations where full JAX is not available or practical. ### `MultiPortSwitch` Bases: `LeafSystem` Route one of N data signals based on an integer selector input. Inputs are `(selector, data_0, data_1, ..., data_{n-1})` and the output is `data_{clip(round(selector), 0, n-1)}`. Implementation strategy: stack the data inputs along a new leading axis and pick out the selected slice with integer indexing. This is fully differentiable through the *selected* data input (zero gradient on the others), which is the standard documented semantics for this block. The `selector` is rounded and clipped to `[0, n-1]` so floating-point inputs are tolerated; the selector itself is non-differentiable (`round`/`clip` zero out the gradient). All data inputs must share the same shape and dtype (the stack requires it). Mixed shapes are rejected by `npa.stack` at trace time. Parameters: | Name | Type | Description | Default | | --------------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `n_data_inputs` | | number of data input ports. Must be >= 1. | *required* | | `choice_names` | | optional tuple of unique non-empty string labels, one per data input, used for self-documenting diagrams. When supplied its length MUST equal n_data_inputs and entries must be distinct. Resolve a friendly name to its integer selector via index_of(name) when wiring the diagram — see "Named choices" below. Default None (integer-only, byte-equivalent to phase 1). | `None` | Input ports (0) selector — scalar integer-valued signal in `[0, n-1]`. Floating values are rounded and clipped. (1..n_data_inputs) data inputs. Output ports (0) The data input at index `selector`. Notes The original T-118 spec includes `indexing="one-based"` and `mode="smooth"` (softmax-blend across data inputs). Both are deferred — see T-118-followup-modes. Zero-based indexing is the only mode supported in phase 1, matching Python conventions. Named choices (T-118-followup-multi-port-string-keys, 2026-05-13): `choice_names=("low", "medium", "high")` lets a diagram document which port means what. The selector port itself still expects an integer at runtime — JAX cannot trace strings, so this is the build-time-only interpretation the followup spec calls out. Look up the integer for a name on the Python side: ``` .. code-block:: python mps = MultiPortSwitch(3, choice_names=("low", "med", "high")) sel = library.Constant(mps.index_of("med")) # → 1 Passing a string to ``index_of`` returns the matching integer; passing an int returns it unchanged after a range check, so callers can mix the two without branching. Unknown strings and out-of-range ints raise ``BlockParameterError`` at construction (build) time, not at trace time. The runtime _compute_output path is unchanged when ``choice_names`` is ``None`` — the default-off byte-equivalence guarantee. ``` #### `choice_names` Tuple of channel labels, or `None` if unlabeled. #### `index_of(selector)` Resolve a string or int selector to its integer index. `selector` may be a string (looked up in `choice_names`) or any object convertible via `int()`. Strings only resolve when `choice_names` was supplied at construction. Out-of-range ints and unknown strings raise `BlockParameterError` at *build time*; runtime selectors flowing through the input port are still clipped silently by `_compute_output` (no change to the runtime path). ### `Multiplexer` Bases: `ReduceBlock` Stack the input signals into a single output signal. Dispatches to `jax.numpy.hstack`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.hstack.html Input ports (0..n_in-1) The input signals. Output ports (0) The stacked output signal. ### `Mux` Bases: `ReduceBlock` Stack `n_inputs` homogeneous signals into a single output signal. This is the standard `Mux` block. It dispatches to `npa.stack` along axis 0, so: - `Mux(3)([1.0, 2.0, 3.0]) -> array([1.0, 2.0, 3.0])` (shape `(3,)`). - `Mux(2)([(1.0, 2.0), (3.0, 4.0)]) -> array([[1.0, 2.0], [3.0, 4.0]])` (shape `(2, 2)`). All inputs must be the same shape and dtype; this matches the conventional `Mux` semantics and `npa.stack`'s broadcasting rules. For the older flatten-by-concatenation behavior (`hstack`), use :class:`Multiplexer` instead. Input ports (0..n_inputs-1) The input signals (must share shape and dtype). Output ports (0) The stacked output signal, with one extra leading axis. ### `Notch` Bases: `LeafSystem` Discrete biquad band-stop ("notch") filter. Implements the textbook biquad notch:: ``` omega0 = 2*pi * frequency_hz * dt r = 1 - pi * bandwidth_hz * dt (pole radius) rho2 = r*r + depth * (1 - r*r) (zero radius**2) b0, b1, b2 = 1, -2*sqrt(rho2)*cos(omega0), rho2 a1, a2 = -2*r*cos(omega0), r*r y[k] = b0*x[k] + b1*x[k-1] + b2*x[k-2] - a1*y[k-1] - a2*y[k-2] ``` With `depth = 1` the zeros sit on the unit circle and the notch is infinitely deep. With `depth = 0` the numerator collapses to the denominator and the block becomes a unit-gain pass-through (useful as an identity check). Input ports (0) The input signal `x`. Output ports (0) The filtered signal `y`. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | Sampling period of the block (s). | *required* | | `frequency_hz` | | Notch centre frequency (Hz). Differentiable. | `1.0` | | `bandwidth_hz` | | Approximate -3 dB bandwidth of the notch (Hz). Differentiable. Must satisfy pi * bandwidth_hz * dt < 1 so the pole radius r remains in (0, 1). | `0.1` | | `depth` | | Notch depth in \[0, 1). Default 0.99. Larger ⇒ deeper notch; depth = 1 puts the zeros on the unit circle for an infinitely deep notch (allowed but numerically borderline). | `0.99` | | `initial_state` | | Initial value of y[-1] (and, implicitly, y[-2], x[-1], x[-2]). Default 0.0. | `0.0` | Notes Differentiability: `frequency_hz`, `bandwidth_hz`, and `depth` enter the biquad coefficients via smooth arithmetic (`cos`, `sqrt`), so `jax.grad` is finite through them. The block is feedthrough on its input port (`b0 = 1`), so `y[k]` depends on `x[k]` directly. ### `ONNX` Bases: `LeafSystem` ONNX inference block. Parameters: | Name | Type | Description | Default | | ----------------------- | ----- | ---------------------------------------------------------------------------------------------------------------------------------------- | --------------------------- | | `file_name` | `str` | Path to the .onnx model file. | *required* | | `num_inputs` | `int` | Number of input tensors the model expects. | `1` | | `num_outputs` | `int` | Number of output tensors the model produces. | `1` | | `cast_outputs_to_dtype` | | Optional jnp dtype name ("float32", "float64", etc.) to cast every output to. If None, the output dtype matches the model's output spec. | `None` | | `providers` | | onnxruntime execution providers; defaults to CPU. Pass e.g. ("CUDAExecutionProvider", "CPUExecutionProvider") on a GPU host. | `('CPUExecutionProvider',)` | | `name` | | Optional block name. | *required* | Notes Differentiability is **best-effort** — `jax.pure_callback` does not define a VJP, so reverse-mode autodiff through the block raises. For end-to-end gradients, look at the T-023a follow-up on a JAX-traceable conversion (`onnx2jax` or similar). .. note:: **float32 artifacts under jaxonomy's global x64.** `import jaxonomy` enables `jax_enable_x64` process-wide, so a float32 ONNX model receives float64 inputs unless you cast at the block boundary — a silent arithmetic change relative to the framework the model was exported and validated in. One-line idiom: pass `cast_outputs_to_dtype="float32"` and feed the block `x.astype(jnp.float32)` inputs. ### `ONNXJax` Bases: `LeafSystem` JAX-traceable ONNX inference (T-023a). Parameters: | Name | Type | Description | Default | | ----------------------- | ----- | ------------------------------------------------------------------------------ | ---------- | | `file_name` | `str` | Path to the .onnx model file. | *required* | | `num_inputs` | `int` | Number of input tensors the model expects. | `1` | | `num_outputs` | `int` | Number of output tensors the model produces. | `1` | | `cast_outputs_to_dtype` | | Optional jnp dtype name to cast every output to ("float32" / "float64" / ...). | `None` | | `name` | | Optional block name. | *required* | Differentiability: end-to-end via `jaxonnxruntime`'s JAX primitive implementations. Op coverage failure shows up at initialize() time as a clear `RuntimeError` from `jaxonnxruntime`. For models that use ops outside `jaxonnxruntime`'s coverage, fall back to :class:`ONNX` — same constructor signature, runs via `onnxruntime` host callback (no autodiff). .. note:: **float32 artifacts under jaxonomy's global x64.** `import jaxonomy` enables `jax_enable_x64` process-wide, so a float32 model here computes against float64 inputs unless you cast at the block boundary. One-line idiom: pass `cast_outputs_to_dtype="float32"` and feed the block `x.astype(jnp.float32)` inputs. .. note:: **Imported discrete-time policies need a ZeroOrderHold.** A sample-and-hold controller exported from a discrete-time training loop (torch / NEUROMANCER-style: compute :math:`u_k` once per sample, hold for `ts`) is re-evaluated at every ODE solver stage (e.g. all four RK4 stages) when wired directly into a continuous plant — continuous-feedback semantics. Both loops "work", but step-for-step parity with the exporting framework is silently destroyed. Follow the block with `ZeroOrderHold(dt=ts)` and pin the step grid with `SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts)`; with that, closed-loop parity is ~4e-8 over 400 steps on the two-tank benchmark. ### `Offset` Bases: `FeedthroughBlock` Add a constant offset or bias to the input signal. Given an input signal `u` and offset value `b`, this will return `y = u + b`. Input ports (0) The input signal. Output ports (0) The input signal plus the offset. Parameters: | Name | Type | Description | Default | | -------- | ---- | ----------------------------------------------- | ---------- | | `offset` | | The constant offset to add to the input signal. | *required* | ### `OperatingPoint` Result of :func:`findop`. Attributes: | Name | Type | Description | | --------------- | ------- | -------------------------------------------------------------------------------- | | `x` | `Any` | Equilibrium continuous state. | | `u` | `Any` | Input value held fixed during the search (taken from base_context at call time). | | `residual_norm` | `float` | Final ‖ẋ(x\*, u)‖\_∞ after Newton iterations. | | `converged` | `bool` | True if residual_norm met tol within max_iter steps. | | `iterations` | `int` | Number of Newton iterations actually executed. | ### `PCEModel` Fitted polynomial-chaos expansion `y = sum_k c_k Psi_k(xi)`. The orthonormal basis (Askey scheme) gives closed-form statistics: the mean is the constant coefficient, the variance is the sum of squared non-constant coefficients, and Sobol indices follow from partitioning that sum by which inputs each basis term depends on (Xiu & Karniadakis 2002; Sudret, *Reliab. Eng. Syst. Saf.* 93(7):964--979, 2008). #### `mean()` Analytic mean = constant-term coefficient. #### `predict(Xstar)` Surrogate response at `Xstar` (jax-traceable). #### `sobol_indices()` Main-effect (first-order) and total Sobol indices per input. Returns a dict `{"first_order": (dim,), "total": (dim,)}`. #### `variance()` Analytic variance = sum of squared non-constant coefficients. ### `PID` Bases: `LTISystem` Continuous-time PID controller. The PID controller is implemented as a state-space system with matrices (A, B, C, D), which are then used to create a (second-order) LTISystem. Note that this only supports single-input, single-output PID controllers. The PID controller implements the following control law: ``` u = kp * e + ki * ∫e + kd * ė ``` where e is the error signal, and ∫e and ė are the integral and derivative of the error signal, respectively. With a filter coefficient of `n` (to make the transfer function proper), the state-space form of the system is: ``` A = [[0, 1], [0, -n]] B = [[0], [1]] C = [[ki * n, (kp * n + ki) - (kp + kd * n) * n]] D = [[kp + kd * n]] ``` Since D is nonzero, the block is feedthrough. Input ports (0) e: Error signal (scalar) Output ports (0) u: Control signal (scalar) Parameters: | Name | Type | Description | Default | | --------------- | ---- | ----------------------------------------------- | ---------- | | `kp` | | Proportional gain | *required* | | `ki` | | Integral gain | *required* | | `kd` | | Derivative gain | *required* | | `n` | | Derivative filter coefficient | *required* | | `initial_state` | | Initial state of the integral term (default: 0) | `0.0` | ### `PIDController2DOF` Bases: `LeafSystem` Two-degree-of-freedom discrete-time PID controller. Implements the standard 2-DOF PID control law:: ``` u = Kp * (b*r - y) + Ki * integral(r - y) + Kd * d/dt(c*r - y) ``` where `r` is the setpoint, `y` is the measurement, and `b` and `c` are setpoint weights in `[0, 1]` for the proportional and derivative paths respectively. With `b = c = 1` the block is numerically equivalent to the existing :class:`PIDDiscrete` block on the error signal `e = r - y`; with `b = c = 0` it becomes an "I-PD" controller (only the integral term reacts to setpoint changes). The integral term uses a forward-Euler approximation:: ``` e_int[k+1] = e_int[k] + (r[k] - y[k]) * dt ``` and the derivative term is computed exactly as for :class:`DerivativeDiscrete` / :class:`PIDDiscrete`, including the optional first-order filter (`filter_type`, `filter_coefficient`). Input ports (0) Setpoint signal `r`. (1) Measurement signal `y`. Dynamic-port appendices, declared in this deterministic order when the corresponding `*_dynamic` flag is True: (a) `b` if `b_dynamic=True` (T-127-followup-external-weights) (b) `c` if `c_dynamic=True` (c) `kp` if `kp_dynamic=True` (T-127-followup-gain-scheduling) (d) `ki` if `ki_dynamic=True` (e) `kd` if `kd_dynamic=True` (f) `kff` if `kff_dynamic=True` (g) `u_ext` if `tracking_enabled=True` (T-127-followup-tracking-mode) (h) `mode_flag` if `tracking_enabled_dynamic=True` (T-127-followup-bumpless-mode-switch). Scalar input cast to a {0, 1} gate that selects whether the tracking-pull branch runs this tick (0 = OFF / AUTO, non-zero = ON / MANUAL / TRACKING). Requires `tracking_enabled=True` so the `u_ext` port exists. Port indices skip any flag set to False, so e.g. with only `kp_dynamic=True` the `kp` port is at index 2; with `b_dynamic=True` + `kp_dynamic=True` the `b` port is at index 2 and `kp` is at index 3. The instance attributes `self.b_index` / `self.c_index` / `self.kp_index` / `self.ki_index` / `self.kd_index` / `self.kff_index` / `self.u_ext_index` / `self.mode_flag_index` expose the resolved positions. Output ports (0) The control signal `u` computed by the 2-DOF PID law. Parameters: | Name | Type | Description | Default | | -------------------------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------- | | `kp` | | Proportional gain (scalar). | `1.0` | | `ki` | | Integral gain (scalar). | `1.0` | | `kd` | | Derivative gain (scalar). | `1.0` | | `b` | | Setpoint weight for the proportional term, in [0, 1]. Default 1.0 (matches 1-DOF PID). Ignored when b_dynamic=True (the runtime port value is used instead). | `1.0` | | `c` | | Setpoint weight for the derivative term, in [0, 1]. Default 1.0 (matches 1-DOF PID). Ignored when c_dynamic=True. Recommendation for real-world controllers: prefer c = 0 (a.k.a. "derivative on measurement only"). With c = 1 (textbook 2-DOF / 1-DOF PID), a step change in the setpoint produces a one-tick spike of magnitude Kd / dt in d/dt(c\*r - y) — "derivative kick" — that propagates into a brief, often saturation-inducing control transient. Setting c = 0 routes the derivative through the measurement only (-d/dt(y)), so setpoint steps no longer kick the derivative term and integral / proportional action alone drive the response. See the :meth:with_derivative_on_measurement factory and the derivative_on_measurement_only=True convenience kwarg below for the standard recipe. | `1.0` | | `derivative_on_measurement_only` | | T-127-followup-derivative-on-measurement. Convenience flag equivalent to c=0 + c_dynamic=False. When True the derivative term sees only the measurement (-d/dt(y)) and is immune to setpoint-step kick; the user-supplied c / c_dynamic kwargs are rejected with a ValueError to keep the contract unambiguous. Default False (c=1) preserves byte-equivalence with phase 1. See :meth:with_derivative_on_measurement for the matching factory function. | *required* | | `b_dynamic` | | If True, the proportional setpoint weight b is read from an additional input port (index 2) rather than from the static b parameter. The static b value is then ignored at runtime; the user MUST connect a signal to the new port. Mirrors the enable_dynamic\_\* pattern used by :class:Saturate and :class:RateLimiter. Default False (byte-equivalent to phase 1). | `False` | | `c_dynamic` | | If True, the derivative setpoint weight c is read from an additional input port instead of the static c parameter. The new port lives at index 2 (when only c_dynamic is set) or index 3 (when both b_dynamic and c_dynamic are set). Default False. | `False` | | `dt` | | Sampling period of the block. | *required* | | `initial_state` | | Initial value of the integral. Default 0.0. | `0.0` | | `filter_type` | | One of "none", "forward", "backward", or "bilinear" — derivative-filter mode. Default "none". | `'none'` | | `filter_coefficient` | | Filter coefficient N for the derivative filter (the conventional "filter coefficient" PID-tuning parameter). Default 1.0. | `1.0` | | `output_min` | | Lower saturation limit on the control output (T-127-followup- anti-windup). None (default) disables the lower clip. When either output_min or output_max is set the saturated control value is published on port (0); the unsaturated value also feeds the anti-windup correction. | `None` | | `output_max` | | Upper saturation limit on the control output. None (default) disables the upper clip. | `None` | | `anti_windup_method` | | One of "none", "back_calc", or "clamping" (T-127-followup-anti-windup). Default "none". "none" — no anti-windup; integrator update unchanged. "back_calc" — back-calculation: subtract (u_unsat - u_sat) / anti_windup_gain * dt from the integrator each tick. Smooth, fully differentiable. "clamping" — integrator-tracking: only update the integral when the controller is not pushing further into saturation (u_unsat == u_sat OR the error sign points away from the saturated direction). Anti-windup is a no-op unless output_min or output_max is also set. | `'none'` | | `anti_windup_gain` | | Tracking time constant Tt used by "back_calc". Smaller values pull the integrator back faster. Default 1.0. Differentiable through jax.grad. | `1.0` | | `integrator_method` | | One of "forward_euler" (default), "backward_euler", or "trapezoidal" — selects the discretisation used to advance the integral term (T-127-followup-discrete-integrator- derivative). Backward-Euler is more stable for stiff loops; trapezoidal is more accurate. Defaults to byte-equivalence with phase 1. Non-default values add an extra e_i_prev delay cell to the discrete state. | `'forward_euler'` | | `derivative_method` | | One of "forward_diff" (default), "backward_diff", or "centered_diff" — selects the unfiltered derivative kernel. "backward_diff" uses past samples only (typical for real-time control, introduces a one-tick delay relative to "forward_diff"). "centered_diff" is less noisy but adds one extra delay cell (e_d_prev_prev). Only applies when filter_type='none' — combining a non-default derivative_method with a recursive filter raises a ValueError at construction. | `'forward_diff'` | | `kff` | | Feedforward gain on the setpoint (T-127-followup-feedforward). Adds kff * r to the PID output before saturation / anti-windup, so the feedforward term participates in the output_min / output_max clip and the anti-windup comparison between unsaturated and saturated control values. For a plant whose steady-state transfer function from u to y is G(0), choosing kff = 1/G(0) makes the controller track step changes in r without integrator action — fastest possible step response (pair with PID for disturbance rejection). Differentiable through jax.grad. Default 0.0 → byte-equivalent to phase 1 (no feedforward). | `0.0` | | `kp_dynamic` | | T-127-followup-gain-scheduling. If True, the proportional gain kp is read from a runtime input port instead of the static kp parameter. The new port is appended after r, y, and any active b / c ports (see "Input ports" above). The static kp value is then ignored at runtime; the user MUST connect a signal to the new port. Default False (byte-equivalent to phase 1). | `False` | | `ki_dynamic` | | Same as kp_dynamic but for the integral gain ki. Default False. | `False` | | `kd_dynamic` | | Same as kp_dynamic but for the derivative gain kd. Default False. | `False` | | `kff_dynamic` | | Same as kp_dynamic but for the feedforward gain kff. Default False. | `False` | | `error_deadband` | | T-127-followup-deadband-error. Non-negative scalar; when positive, the raw error signal e_raw (the unweighted r - y for the integral path and the weighted b\*r - y / c\*r - y for the P / D paths) is gated through a deadband before it feeds each PID term. In hard mode the gate is e = e_raw for | e_raw | | `error_deadband_mode` | | "hard" (default) or "smooth". Hard mode uses an npa.where gate; smooth mode uses :func:soft_dead_zone for a sigmoid-blended kernel with finite gradient through the band. Only relevant when error_deadband > 0. | `'hard'` | | `error_deadband_sharpness` | | Positive scalar controlling the steepness of the smooth deadband transition (passed straight to :func:soft_dead_zone). Default 10.0. Ignored when error_deadband_mode='hard'. | `10.0` | | `tracking_enabled` | | T-127-followup-tracking-mode. If True, declares an extra input port u_ext (appended after every other dynamic port) and folds a tracking-error term into the integrator update:: e_track = u_ext - u_unsat I[k+1] += (e_track / tracking_gain) * dt This is the standard "tracking mode" / "manual mode" mechanism: while another controller (or an operator) drives u_ext, the PID's integrator is pulled toward the value that would produce u_ext so the handoff back to PID- driven control is bumpless. Implementation-wise it is back- calculation on the external signal — the same kernel as anti_windup_method="back_calc" but using u_ext instead of u_sat as the target. Both mechanisms compose: their corrections sum into the integrator each tick. Default False (byte-equivalent to phase 1). | `False` | | `tracking_gain` | | Tracking time constant Tt for the tracking-mode back- calculation kernel (T-127-followup-tracking-mode). Smaller values pull the integrator toward u_ext faster. Differentiable through jax.grad. Default 1.0. | `1.0` | | `integrate_tracking_error` | | T-127-followup-i-on-error-only. Selects whether the tracking-error term (u_ext - u_unsat)/Tt * dt is folded into the integrator each tick. When True (default) the integrator update is:: I[k+1] = I[k] + Ki\*(r-y)\*dt + (u_ext - u_unsat)/Tt * dt preserving the T-127-followup-tracking-mode kernel exactly. When False the integrator only accumulates the regulation error:: I[k+1] = I[k] + Ki\*(r-y)\*dt and u_ext does NOT pull the integrator at all (the tracking signal still flows through any parallel path the user has wired up, e.g. a feedforward addition outside the block). Only meaningful when tracking_enabled=True; the flag is silently irrelevant otherwise but still round-trips through :meth:to_dict / :meth:from_dict. Default True is byte-equivalent to T-127-followup-tracking-mode. | `True` | | `tracking_enabled_dynamic` | | T-127-followup-bumpless-mode-switch. When True, promotes the tracking-mode flag from a static construction- time choice to a runtime SCALAR INPUT port appended after u_ext. The port value is treated as a boolean (0 = OFF / AUTO, non-zero = ON / MANUAL/TRACKING) — multiplying the per-tick tracking-pull correction by that gate. This lets a model toggle between PID-driven control and external override mid-simulation while keeping the integrator loaded with the value that would produce u_ext (so handoffs in either direction remain bumpless). Requires tracking_enabled=True; tracking_enabled=False plus tracking_enabled_dynamic=True raises ValueError at construction (without the u_ext port the runtime gate has nothing to multiply). Composes with integrate_tracking_error (the gate multiplies the correction term that flag exposes). Default False is byte-equivalent to T-127-followup-tracking-mode. | `False` | Gain-scheduling recipe Each of `kp_dynamic`, `ki_dynamic`, `kd_dynamic`, and `kff_dynamic` is the natural plug for a lookup-table-driven gain. The standard wiring uses one :class:`LookupTable1d` (or :class:`LookupTable2d` for two scheduling variables) per scheduled gain and the same scheduling-variable signal source for all of them:: ``` import jaxonomy from jaxonomy.library import ( Constant, LookupTable1d, PIDController2DOF, ) builder = jaxonomy.DiagramBuilder() r = builder.add(Constant(1.0, name="r")) y = builder.add(Constant(0.0, name="y")) # Scheduling variable -- e.g. engine speed, Mach number, # tank level. Replace with whatever source you have. sched = builder.add(Constant(0.5, name="sched")) # Schedule kp as a function of the scheduling variable. kp_tbl = builder.add( LookupTable1d( input_array=[0.0, 0.5, 1.0], output_array=[1.0, 2.0, 4.0], interpolation="linear", name="kp_schedule", ) ) pid = builder.add( PIDController2DOF( dt=0.01, kp_dynamic=True, name="pid" ) ) builder.connect(r.output_ports[0], pid.input_ports[0]) builder.connect(y.output_ports[0], pid.input_ports[1]) builder.connect(sched.output_ports[0], kp_tbl.input_ports[0]) # kp port is at index 2 when only kp_dynamic is set. builder.connect(kp_tbl.output_ports[0], pid.input_ports[2]) ``` Because every dynamic port is a regular signal port, gradients flow through the lookup-table parameters (breakpoints / values) AND through the scheduling-variable signal — the standard T-114 guarantee. Multiple gains can be scheduled simultaneously; flagging `ki_dynamic` and `kd_dynamic` simply adds two more ports for the integral / derivative tables. #### `cohen_coon(K, tau, theta, dt, mode='PID', **kwargs)` Construct a PID tuned by the Cohen-Coon rule for a FOPDT plant. For a first-order-plus-dead-time plant `G(s) = K * exp(-theta*s) / (tau*s + 1)` (process gain `K`, time constant `tau`, dead time `theta`), the Cohen-Coon (1953) formulas are:: ``` r = theta / tau P: Kp = (1/K) * (1/r) * (1 + r/3) PI: Kp = (1/K) * (1/r) * (9/10 + r/12) Ti = theta * (30 + 3*r) / (9 + 20*r) PID: Kp = (1/K) * (1/r) * (4/3 + r/4) Ti = theta * (32 + 6*r) / (13 + 8*r) Td = theta * 4 / (11 + 2*r) ``` with `Ki = Kp / Ti` and `Kd = Kp * Td`. Cohen-Coon is more aggressive than Ziegler-Nichols on plants where `theta / tau` is large (dead-time-dominated processes); it is widely used in chemical-process control. Parameters: | Name | Type | Description | Default | | ---------- | ---- | -------------------------------------------------------------------------- | ---------- | | `K` | | Process (steady-state) gain. Must be non-zero. | *required* | | `tau` | | First-order time constant in seconds. Must be positive. | *required* | | `theta` | | Dead time in seconds. Must be positive. | *required* | | `dt` | | Sampling period for the discrete PID. | *required* | | `mode` | | One of "P", "PI", "PID" (default "PID"). Selects which gains are non-zero. | `'PID'` | | `**kwargs` | | Forwarded to :class:PIDController2DOF. | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ------------------------------------------------ | | `A` | | class:PIDController2DOF whose (Kp, Ki, Kd) match | | | | the Cohen-Coon formulas for the requested mode. | Raises: | Type | Description | | ------------ | -------------------------------------------------------------------------------------------- | | `ValueError` | If mode is not one of "P", "PI", "PID", or if K is zero, or if tau / theta are non-positive. | #### `from_dict(data, **block_kwargs)` Reconstruct a :class:`PIDController2DOF` from a config dict. Extra keyword arguments (`name=`, `system_id=`, ...) are forwarded to the constructor so a deserialized block can pick up a fresh name in its target diagram. Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `data` | | dict produced by :meth:to_dict (or any equivalent mapping with the same key set). dt is the only required field; every other field falls back to the constructor default when absent. | *required* | | `**block_kwargs` | | forwarded to PIDController2DOF.__init__ (typically name=...). | `{}` | Raises: | Type | Description | | ------------ | ------------------------------------ | | `ValueError` | if data lacks the required dt field. | Returns: | Type | Description | | ---- | -------------------------------------------------- | | | A new :class:PIDController2DOF whose configuration | | | matches data. | #### `standard(kp, ki, kd, dt, **kwargs)` Construct a textbook PID with `b = c = 1`. Convenience factory equivalent to `PIDController2DOF(dt, kp, ki, kd)`: the proportional and derivative paths both see the setpoint with weight 1, so the block reduces to a 1-DOF PID on the error signal `e = r - y`. Useful as the explicit counterpart to :meth:`with_derivative_on_measurement` — callers self-document which 2-DOF configuration they want. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `kp` | | Proportional gain. | *required* | | `ki` | | Integral gain. | *required* | | `kd` | | Derivative gain. | *required* | | `dt` | | Sampling period. | *required* | | `**kwargs` | | Forwarded to :class:PIDController2DOF. Setting b or c here is allowed but discouraged (use the main constructor for non-standard weights). | `{}` | Returns: | Name | Type | Description | | ---- | ---- | --------------------------------------- | | `A` | | class:PIDController2DOF with b = c = 1. | #### `to_dict()` Return a JSON-serializable dict describing this controller. The dict captures every construction-time field that controls the block's behavior — all `@parameters`-registered fields plus the mode strings and `*_dynamic` port-topology flags that live outside `@parameters`. Round-tripping through :meth:`from_dict` (optionally via `json.dumps` / `json.loads`) produces a block with identical step-response behavior on a fixed input. Note Diagram wiring (which signal feeds which input port) is not part of the block config and must be re-established by the caller after :meth:`from_dict`. This is especially relevant when any `*_dynamic` flag is True — the reconstructed block still declares the runtime port, but it has no upstream connection until the caller wires it up. Returns: | Type | Description | | ---- | ----------------------------------------------------- | | | dict mapping each of :attr:\_CONFIG_STATIC_FIELDS and | | | attr:\_CONFIG_DYNAMIC_FIELDS to a JSON primitive. | #### `tyreus_luyben(Ku, Tu, dt, **kwargs)` Construct a PI controller tuned by the Tyreus-Luyben rule. A gentler Ziegler-Nichols alternative that trades response speed for robustness. Tyreus & Luyben (1992) recommend the PI form for most chemical-process applications because the derivative term tends to amplify measurement noise:: ``` Kp = Ku / 3.2, Ti = 2.2 * Tu → Ki = Kp / Ti ``` and `Kd = 0` (no derivative action). Compared with Z-N, Tyreus-Luyben gives roughly 1/3 the proportional gain and a ~4x longer integral time, producing a much less aggressive loop with substantially better robustness to model error. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------------------------- | ---------- | | `Ku` | | Ultimate gain (proportional-only gain at sustained oscillation). Must be positive. | *required* | | `Tu` | | Ultimate period (period of the sustained oscillation in seconds). Must be positive. | *required* | | `dt` | | Sampling period for the discrete PID. | *required* | | `**kwargs` | | Forwarded to :class:PIDController2DOF. | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ------------------------------------------------- | | `A` | | class:PIDController2DOF configured as a PI | | | | controller (Kd = 0) with the Tyreus-Luyben gains. | Raises: | Type | Description | | ------------ | ---------------------------- | | `ValueError` | If Ku / Tu are non-positive. | #### `with_derivative_on_measurement(kp, ki, kd, dt, **kwargs)` Construct a PID with derivative-on-measurement-only (`b=1, c=0`). Convenience factory for the standard "no derivative kick" recipe used by most real-world controllers. With `c = 0` the derivative term sees only the measurement (`-d/dt(y)`), so a step change in the setpoint does NOT inject a `Kd / dt` spike through the derivative path — only integral and proportional action drive the transient. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `kp` | | Proportional gain. | *required* | | `ki` | | Integral gain. | *required* | | `kd` | | Derivative gain. | *required* | | `dt` | | Sampling period. | *required* | | `**kwargs` | | Forwarded to :class:PIDController2DOF. Passing c or c_dynamic here raises ValueError via the constructor's derivative_on_measurement_only contract — by construction this factory pins c=0. | `{}` | Returns: | Name | Type | Description | | ---- | ---- | --------------------------------------------- | | `A` | | class:PIDController2DOF with b = 1 and c = 0. | #### `ziegler_nichols(Ku, Tu, dt, mode='PID', **kwargs)` Construct a PID tuned by the Ziegler-Nichols ultimate-cycle rule. Given the ultimate gain `Ku` (the proportional-only gain at which the closed loop just sustains oscillation) and the corresponding ultimate period `Tu`, the Z-N table maps to controller gains as:: ``` P: Kp = 0.5 * Ku, Ki = 0, Kd = 0 PI: Kp = 0.45 * Ku, Ti = Tu / 1.2 → Ki = 0.54*Ku/Tu, Kd = 0 PID: Kp = 0.6 * Ku, Ti = Tu / 2.0 → Ki = 1.2 *Ku/Tu, Td = Tu / 8.0 → Kd = 0.075*Ku*Tu ``` The coefficients are the canonical Ziegler & Nichols (1942) values; see e.g. Astrom & Hagglund, *PID Controllers: Theory, Design, and Tuning* (1995), Table 4.1. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------------------------- | ---------- | | `Ku` | | Ultimate gain (proportional-only gain at sustained oscillation). Must be positive. | *required* | | `Tu` | | Ultimate period (period of the sustained oscillation in seconds). Must be positive. | *required* | | `dt` | | Sampling period for the discrete PID. | *required* | | `mode` | | One of "P", "PI", "PID" (default "PID"). Selects which gains are non-zero. | `'PID'` | | `**kwargs` | | Forwarded to :class:PIDController2DOF. | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ------------------------------------------------ | | `A` | | class:PIDController2DOF whose (Kp, Ki, Kd) match | | | | the Z-N table for the requested mode. | Raises: | Type | Description | | ------------ | ----------------------------------------------------------------------- | | `ValueError` | If mode is not one of "P", "PI", "PID", or if Ku / Tu are non-positive. | ### `PIDDiscrete` Bases: `LeafSystem` Discrete-time PID controller. This block implements a discrete-time PID controller with a first-order approximation to the integrated error and an optional derivative filter. The integrated error term is computed as: ``` e_int[k+1] = e_int[k] + e[k] * dt ``` where `e` is the error signal and `dt` is the sampling period. The derivative term is computed in the same way as for the DerivativeDiscrete block, including filter options described there. With the running error integral `e_int` and current estimate of the time derivative of the error `e_dot`, the output is: ``` u[k] = kp * e[k] + ki * e_int[k] + kd * e_dot[k] ``` Input ports (0) The error signal. Output ports (0) The control signal computed by the PID algorithm. Parameters: | Name | Type | Description | Default | | ------------------------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `kp` | | The proportional gain (scalar) | `1.0` | | `ki` | | The integral gain (scalar) | `1.0` | | `kd` | | The derivative gain (scalar) | `1.0` | | `dt` | | The sampling period of the block. | *required* | | `initial_state` | | The initial value of the running error integral. Default is 0. | `0.0` | | `enable_external_initial_state` | | Source for the value used for the integrator initial state. True=from inport, False=from the initial_state parameter. | `False` | | `filter_type` | | One of "none", "forward", "backward", or "bilinear". Determines the type of filter used to estimate the derivative of the error signal. Default is "none". See DerivativeDiscrete documentation for details. | `'none'` | | `filter_coefficient` | | The filter coefficient for the derivative filter. Default is 1.0. See DerivativeDiscrete documentation for details. | `1.0` | #### `initialize_static_data(context)` Set the initial state from the input port, if specified via config ### `PMSM` Bases: `LeafSystem` Interior permanent-magnet synchronous machine in the rotor (dq) frame. Electrical dynamics (amplitude-invariant dq, electrical speed `w_e = pole_pairs * w_m`):: ``` Ld * di_d/dt = v_d - R*i_d + w_e*Lq*i_q Lq * di_q/dt = v_q - R*i_q - w_e*(Ld*i_d + lambda_m) ``` Electromagnetic torque (magnet + reluctance):: ``` Te = 1.5 * pole_pairs * (lambda_m*i_q + (Ld - Lq)*i_d*i_q) ``` Mechanical dynamics:: ``` J * dw_m/dt = Te - B*w_m - T_load dtheta_m/dt = w_m ``` A surface PMSM or BLDC (sinusoidal back-EMF approximation) is the special case `Ld == Lq`. Input ports (0) v_dq: rotor-frame stator voltage `[v_d, v_q]` (V). (1) T_load: load torque (N\*m) — only when `enable_load_port=True`; otherwise the `T_load` parameter is used. Output ports (0) state: `[i_d, i_q, w_m, theta_m]` (A, A, rad/s, rad). (1) torque: electromagnetic torque `Te` (N\*m). Parameters: | Name | Type | Description | Default | | ------------------ | ------- | ------------------------------------------------------------------------------------------------------------------------------- | -------- | | `R` | `float` | Stator phase resistance (ohm). | `0.45` | | `Ld` | `float` | d-axis inductance (H). | `0.0032` | | `Lq` | `float` | q-axis inductance (H). Lq > Ld models an interior PMSM. | `0.0058` | | `lambda_m` | `float` | Permanent-magnet flux linkage (Wb). | `0.0533` | | `pole_pairs` | `float` | Number of pole pairs (electrical/mechanical speed ratio). | `4.0` | | `J` | `float` | Rotor inertia (kg\*m^2). | `0.0012` | | `B` | `float` | Viscous friction coefficient (Nms). | `8e-05` | | `T_load` | `float` | Constant load torque (N\*m); ignored when the load port is on. | `0.0` | | `initial_state` | | Initial [i_d, i_q, w_m, theta_m] (default zeros). | `None` | | `locked` | `bool` | Clamp the rotor (dw_m = dtheta_m = 0) — a standstill bench test that reduces each axis to an RL circuit with time constant L/R. | `False` | | `enable_load_port` | `bool` | Expose input port (1) for the load torque. | `False` | The electrical angle for the Park transforms is `theta_e = pole_pairs * theta_m`. ### `PRBS` Bases: `LeafSystem` Pseudo-Random Binary Sequence (PRBS) source. Emits `+amplitude` or `-amplitude` at each `sample_time` tick, drawn from a fair Bernoulli(0.5) under `jax.random.bernoulli` and remapped via `2*b - 1`. Useful as a broad-band excitation signal for system identification. The `amplitude` parameter is differentiable (scaling the binary selector); the `+1 / -1` selector itself is wrapped in `lax.stop_gradient`. Input ports None. Output ports (0) The most recent `±amplitude` sample. Parameters: | Name | Type | Description | Default | | ------------- | ------- | ---------------------------------------------------------------------------------------- | ---------- | | `sample_time` | `float` | Period (s) at which a fresh bit is drawn. | *required* | | `amplitude` | `float` | Magnitude of the binary output (differentiable). | `1.0` | | `seed` | `int` | Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random. | `None` | Notes Phase 1 uses Bernoulli(0.5) sampling rather than a true maximal-length LFSR. Period-faithful PRBS-N (n_bits register size) is deferred — see `T-122-followup-lfsr`. Per-vmap-batch independence: pass `fold_in_batch_index=True` (T-122-followup-vmap-fold-in) to derive a per-replica independent PRNG stream via `jax.lax.axis_index("batch")` inside `simulate_batch(use_vmap=True)` / `simulate_distributed`. Default `False` preserves bit-identical phase 1 behaviour. ### `PRBSLFSR` Bases: `LeafSystem` True maximal-length PRBS-N source built on a binary LFSR. Emits `+amplitude` or `-amplitude` at each `sample_time` tick, drawn from a Linear-Feedback Shift Register (LFSR) of length `register_length` configured with the standard primitive feedback polynomial. The output sequence has period exactly `2^N - 1` and a flat power spectrum below `1/(2N)` of the sample rate, making it the canonical "white" excitation for system identification. Reproducibility: same `seed` -> bit-identical sequence. Distinct seeds traverse the same cyclic orbit at different starting phases. Differentiability: `amplitude` is a dynamic parameter and flows through gradients linearly; the binary selector itself is wrapped in `lax.stop_gradient` (the LSB extraction is non-differentiable). Input ports None. Output ports (0) The most recent `±amplitude` sample. Parameters: | Name | Type | Description | Default | | ----------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `sample_time` | `float` | Period (s) at which the LFSR advances by one step. | *required* | | `amplitude` | `float` | Magnitude of the binary output (differentiable). | `1.0` | | `register_length` | `int` | One of {7, 9, 11, 15, 17, 23, 31}. Determines the period (2^N - 1) and the tap polynomial. | `15` | | `seed` | `int` | Non-zero integer seeding the LFSR register state. 0 is silently promoted to 1 since the all-zero register is a fixed point of any LFSR (would emit a constant zero). | `1` | Notes Per-vmap-batch independence: pass `fold_in_batch_index=True` (T-122-followup-vmap-fold-in) to derive a per-replica independent starting phase on the same maximal-length cycle. Inside `simulate_batch(use_vmap=True)` / `simulate_distributed` (which wrap their vmap with `axis_name="batch"`), the LFSR register is XOR-perturbed by a per-replica non-zero salt derived from `jax.lax.axis_index("batch")` on the very first update step (tracked via a `phase_advanced` flag in the discrete state). The salt is masked to the low N bits of the register and promoted from 0 to 1 to avoid the all-zero fixed point. Default `False` preserves bit-identical behaviour with the original LFSR follow-up. ### `Park` Bases: `_AngleTransform` Park transform: stationary `[alpha, beta]` -> rotor `[d, q]`. Input ports (0) alpha_beta: stationary-frame vector. (1) theta: electrical angle `theta_e` (rad); for :class:`PMSM`, `theta_e = pole_pairs * theta_m`. Output ports (0) dq: `[d, q] = [ cos*alpha + sin*beta, -sin*alpha + cos*beta ]`. ### `PolynomialChaos` Bases: `LeafSystem` Polynomial-chaos surrogate `y = sum_k c_k Psi_k(u)` as a feedthrough block. Input port 0 is the feature vector `u`; the coefficients are a dynamic parameter (Xiu & Karniadakis 2002). ### `Power` Bases: `FeedthroughBlock` Raise the input signal to a constant power. Dispatches to `jax.numpy.power`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.power.html For input signal `u` with exponent `p`, the output will be `y = u ** p`. Input ports (0) The input signal. Output ports (0) The input signal raised to the power of the exponent. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ------------------------------------------------- | ---------- | | `exponent` | | The exponent to which the input signal is raised. | *required* | ### `Prelookup` Bases: `LeafSystem` Compute the (bucket_index, fraction) pair for a query against a precomputed grid. This is the upstream half of the standard `Prelookup`/`InterpolationUsingPrelookup` pair. Pair with one or more :class:`InterpolationUsingPrelookup` blocks downstream -- each can interpolate a DIFFERENT output table that shares the same grid axis without re-running the bucket search. The marketing wedge: when N downstream tables share one query axis (e.g. 10 lookup maps on a common engine-RPM input), this saves `N - 1` binary searches per evaluation. Input ports `(0)` -- the query coordinate (scalar). Output ports `(0)` -- a NamedTuple with fields `(index, fraction)` ready to plug into one or more :class:`InterpolationUsingPrelookup` blocks. Parameters: | Name | Type | Description | Default | | --------------- | ----------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `input_array` | | 1-D, strictly-increasing grid of breakpoints. Stored verbatim for the bucket search. | *required* | | `dtype` | `optional` | If set (e.g. jnp.float32), the grid array is cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d. | `None` | | `extrapolation` | `(optional, T - 114 - fu - prelookup - extrap)` | Out-of-range policy for the alpha blend weight. One of "clip" (default; alpha clamped to [0, 1] so OOB queries map to the nearest endpoint), "linear" (alpha left raw -- the downstream blend extends the boundary slope linearly), or "nan" (alpha set to NaN on OOB queries -- the downstream blend propagates NaN). All three modes match the corresponding :class:LookupTable1d extrapolation policies byte-for-byte. Any paired :class:InterpolationUsingPrelookup should declare the SAME mode for API agreement (the math is owned by Prelookup's alpha computation). | `'clip'` | Notes Differentiable through the query coordinate via `fraction` (the discrete `index` is piecewise-constant). #### `extrapolation` The OOB policy applied to `alpha` (`"clip"`/`"linear"`/`"nan"`). #### `input_array` The breakpoint array used for the bucket search. ### `PrelookupInverse` Bases: `LeafSystem` Compute the (bucket_index, fraction) pair for an INVERSE-direction lookup against a strictly-monotonic value array. Forward :class:`Prelookup` answers: given `x`, find `(i, alpha)` s.t. `xp[i] + alpha * (xp[i+1] - xp[i]) ≈ x`. Inverse :class:`PrelookupInverse` answers: given `y`, find `(i, alpha)` s.t. `yp[i] + alpha * (yp[i+1] - yp[i]) ≈ y`. The output is the same NamedTuple-typed :class:`_PrelookupResult` produced by :class:`Prelookup`, so it plugs straight into one or more :class:`InterpolationUsingPrelookup` blocks connected to OTHER tables -- typically the inverse table that maps `i` back to the recovered `x` (for example the breakpoints of the forward table). Marketing wedge: implicit equations `y = f(x)` where `f` is a monotonic 1-D table -- gain scheduling, sensor calibration, etc. Input ports `(0)` -- the query coordinate in the OUTPUT space (`y`). Output ports `(0)` -- a :class:`_PrelookupResult` NamedTuple `(index, fraction)` ready to plug into an :class:`InterpolationUsingPrelookup` block. Parameters: | Name | Type | Description | Default | | --------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `output_array` | | 1-D, strictly-monotonic (increasing OR decreasing) array of values to invert. Stored verbatim for the bucket search; decreasing tables are handled by reversing the search direction. | *required* | | `dtype` | `optional` | If set, the value array is cast to this dtype on construction. Mirrors the :class:Prelookup / :class:LookupTable1d dtype contract. | `None` | | `extrapolation` | `optional` | Only "clip" is supported in this followup -- queries outside the monotone range collapse to the nearest endpoint. "linear"/"nan" are filed as T-114-followup-prelookup-inverse-extrap because the OOB definition on a value-axis interacts with the direction-flipping logic non-trivially. | `'clip'` | Notes Differentiable through the query coordinate and through `output_array`. Non-monotonic `output_array` raises `ValueError` at construction time. #### `direction` `"increasing"` or `"decreasing"` -- monotonicity sense. #### `extrapolation` The OOB policy (always `"clip"` in this followup). #### `output_array` The 1-D monotonic value array being inverted. ### `Product` Bases: `ReduceBlock` Compute the product and/or quotient of the input signals. The block will multiply or divide the input signals, depending on the specified operators. For example, if the block has three inputs `u1`, `u2`, and `u3` and is configured with operators="\*\*/", then the output signal will be `y = u1 * u2 / u3`. By default, the block will multiply all of the input signals. Input ports (0..n_in-1) The input signals. Output ports (0) The product and/or quotient of the input signals. Parameters: | Name | Type | Description | Default | | ------------------------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `n_in` | | The number of input ports. | *required* | | `operators` | | A string of length n_in specifying the operators to apply to each of the input signals. Each character in the string must be either "" or "/". The default is "". | `None` | | `denominator_limit` | | Currently unsupported | `None` | | `divide_by_zero_behavior` | | Currently unsupported | `None` | ### `ProductOfElements` Bases: `FeedthroughBlock` Compute the product of the elements of the input signal. Dispatches to `jax.numpy.prod`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.prod.html Input ports (0) The input signal. Output ports (0) The product of the elements of the input signal. ### `Pulse` Bases: `SourceBlock` A periodic pulse signal. Given amplitude `a`, pulse width `w`, and period `p`, the output signal is: ``` y(t) = a if t % p < w else 0 ``` where `%` is the modulo operator. Input ports None Output ports (0) The pulse signal. Parameters: | Name | Type | Description | Default | | ------------- | ---- | ------------------------------------------------------------ | ------- | | `amplitude` | | The amplitude of the pulse signal. | `1.0` | | `pulse_width` | | The fraction of the period during which the pulse is "high". | `0.5` | | `period` | | The period of the pulse signal. | `1.0` | | `phase_delay` | | Currently unsupported. | `0.0` | ### `PyTorch` Bases: `LeafSystem` Block to perform inference with a pre-trained PyTorch model saved as TorchScript. The input to the block should be of compatible type and shape expected by the TorchScript. For example, if the TorchScript model expects a `torch.float32` tensor of shape `(3, 224, 224)`, the input to the block should be a `jax.numpy` array of shape (3, 224, 224) of dtype `jnp.float32`. For output types, if no casting is specified through the `cast_outputs_to_dtype` parameter, the output of the block will have the same dtype as the TorchScript model output, but expressed as `jax.numpy` types. For example. if the TorchScript model outputs a `torch.float32` tensor, the output of the block will be a `jax.numpy` array of dtype `jnp.float32`. If casting is specified through `cast_outputs_to_dtype` parameter, all the outputs, of the block will be casted to this specific `jax.numpy` dtype. .. note:: **float32 models under jaxonomy's global x64.** `import jaxonomy` enables JAX 64-bit mode (`jax_enable_x64`) for the whole process, so upstream signals are float64 by default. A TorchScript model traced in `torch.float32` therefore receives float64 inputs (a dtype error or a silent arithmetic change) unless you cast at the block boundary. One-line idiom: pass `cast_outputs_to_dtype="float32"` and feed the block `x.astype(jnp.float32)` inputs. Input ports (i) The ith input to the model. Output ports (j) The jth output of the model. Parameters: | Name | Type | Description | Default | | ------------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `file_name` | `str` | Path to the model Torchscript .pt file. | *required* | | `num_inputs` | `int` | The number of inputs to the model. Only required for TorchScript models. | `1` | | `num_outputs` | `int` | The number of outputs of the model. | `1` | | `cast_outputs_to_dtype` | `str` | The dtype to cast all the outputs of the block to. Must correspond to a jax.numpy datatype. For example, "float32", "float64", "int32", "int64". | `None` | | `add_batch_dim_to_inputs` | `bool` | Whether to add a new first dimension to the inputs before evaluating the TorchScript or TensorFlow model. This is useful when the model expects a batch dimension. | `False` | #### `initialize_static_data(context)` Infer the output shapes and dtypes of the ML model. ### `QuadraticCost` Bases: `ReduceBlock` LQR-type quadratic cost function for a state and input. Computes the cost as x'Qx + u'Ru, where Q and R are the cost matrices. In order to compute a running cost, combine this with an `Integrator` or `IntegratorDiscrete` block. ### `QuanserHAL` Bases: `LeafSystem` Hardware Abstraction Layer for Quanser hardware. This block provides an interface to virtual or physical Quanser hardware. It requires that the Quanser hardware or QLabs simulator be properly configured and that the Quanser python library is available on the system path. See the Quanser documentation for more information. To use an idealized model of the Qube Servo hardware, see the `jaxonomy.library.QubeServoModel` block, which may be run without hardware or in the cloud-based simulation UI. Input ports (0) Control signal to the motor in volts Output ports (0) The observed rotor and pendulum angles in radians Parameters: | Name | Type | Description | Default | | ---------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------- | -------------- | | `dt` | | The time step of the simulation. | *required* | | `version` | | The version of the Qube hardware (2 or 3). By default, version 2 is used with hardware=False, or version 3 is used with hardware=True. | `2` | | `hardware` | | If True, connect to the physical hardware. If False, connect to the QLabs simulator. | `False` | | `name` | | The name of the system in the Jaxonomy model. | `'QuanserHAL'` | ### `Quantizer` Bases: `FeedthroughBlock` Discritize the input signal into a set of discrete values. Given an input signal `u` and a resolution `interval`, this block quantizes the input signal onto the integer multiples of `interval`. The output signal is `y = interval * f(u / interval)` where `f` is selected by `mode`: - `"round"` (default): round-half-to-even (banker's rounding, IEEE-754 default). Byte-equivalent with the phase-1 implementation (which used `npa.round` unconditionally). - `"floor"`: round toward -inf (truncation in many DSP impls). - `"ceil"`: round toward +inf. - `"trunc"`: round toward zero (chops the fractional part regardless of sign). Quantization is non-differentiable: the output is piecewise-constant with measure-zero jumps. The block wraps the rounded result in :func:`jax.lax.stop_gradient` (JAX backend only) so JAX always sees a zero gradient through the block. This both matches the underlying mathematical reality and prevents spurious gradient leakage if any backend ever provides a smoothed surrogate for `round`/`floor`/ `ceil`/`trunc`. Under the numpy backend the helper is the identity, preserving dtype/value byte-equivalence. Input ports (0) The continuous input signal. In most cases, should be scaled to the range `[0, interval]`. Output ports (0) The quantized output signal, on the same scale as the input signal. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------------------- | ---------- | | `interval` | | The quantization step size — output values are integer multiples of interval. | *required* | | `mode` | | One of "round", "floor", "ceil", "trunc". Default "round". | `'round'` | ### `QubeServoModel` Bases: `LeafSystem` Plant model for the Quanser Qube Servo Furuta Pendulum. The Quanser Qube Servo is a pendulum controlled by a rotary arm. The rotary arm is actuated by a DC motor. The pendulum is free to rotate about the rotary arm. The state of the system is given by the rotor angle (theta), pendulum angle (alpha), rotor angular velocity, and pendulum angular velocity. The input to the system is the voltage applied to the motor, which is converted to torque by a simple linear model. Input ports (0) The motor voltage signal Output ports (0) If `full_state_output` is False, the rotor angle and pendulum angle. Otherwise, will return the entire continuous state vector. Parameters: | Name | Type | Description | Default | | ------------------- | ---- | -------------------------------------------------------------------------------------------- | ---------------------- | | `x0` | | Initial state of the system [theta, alpha, theta_dot, alpha_dot] | `[0.0, 0.0, 0.0, 0.0]` | | `Rm` | | Motor resistance (Ohms) | `8.4` | | `km` | | Back-emf constant (V-s/rad) | `0.042` | | `mr` | | Rotary arm mass (kg) | `0.095` | | `Lr` | | Rotor arm length (m) | `0.085` | | `br` | | Rotor arm damping coefficient (N-m-s/rad) | `0.0005` | | `mp` | | Pendulum mass (kg) | `0.024` | | `Lp` | | Pendulum arm length (m) | `0.129` | | `bp` | | Pendulum damping coefficient (N-m-s/rad) | `2.5e-05` | | `g` | | Gravitational constant (m/s^2) | `9.81` | | `kr` | | Feedback control to send the rotor back to zero | `0.0` | | `full_state_output` | | If True, output the full state vector. Otherwise, only output the rotor and pendulum angles. | `False` | ### `RBFModel` Fitted radial-basis-function interpolant with optional polynomial tail. `s(x) = sum_i w_i phi(||x - c_i||) + sum_k c_k p_k(x)` (Hardy 1971; Wendland 2005). #### `predict(Xstar)` Interpolant value at `Xstar` (jax-traceable). ### `RadialBasisSurrogate` Bases: `LeafSystem` RBF surrogate `y = sum_i w_i phi(||u - c_i||) (+ poly tail)` as a feedthrough block. Input port 0 is the feature vector `u`; the RBF weights (and polynomial-tail coefficients) are dynamic parameters (Hardy 1971). ### `Ramp` Bases: `SourceBlock` Output a linear ramp signal in time. Given a slope `m`, a start value `y0`, and a start time `t0`, the output signal is: ``` y(t) = m * (t - t0) + y0 if t >= t0 else y0 ``` where `t` is the current simulation time. Input ports None Output ports (0) The ramp signal. Parameters: | Name | Type | Description | Default | | ------------- | --------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `start_value` | | The value of the output signal at the start time. | `0.0` | | `slope` | | The slope of the ramp signal. | `1.0` | | `start_time` | | The time at which the ramp signal begins. | `1.0` | | `units` | `(optional, T - 104 - followup - units - on - source - blocks)` | If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams). | `None` | ### `RandomNumber` Bases: `LeafSystem` Discrete-time random number generator. Generates independent, identically distributed random numbers at each time step. Dispatches to `jax.random` for the actual random number generation. Supported distributions include "ball", "cauchy", "choice", "dirichlet", "exponential", "gamma", "lognormal", "maxwell", "normal", "orthogonal", "poisson", "randint", "truncated_normal", and "uniform". See https://jax.readthedocs.io/en/latest/jax.random.html#random-samplers for a full list of available distributions and associated parameters. Although the JAX random number generator is a deterministic function of the key, this block maintains the key as part of the discrete state, making it a stateful RNG. The block can be seeded for reproducibility by passing an integer seed; if None, a random seed will be generated using numpy.random. Note that this block should typically not be used as a source of randomness for continuous-time systems, as it generates a discrete-time signal. For continuous systems, use a continuous-time noise source, such as `WhiteNoise`. Input ports None Output ports (0) The most recently generated random number. Parameters: | Name | Type | Description | Default | | ------------------------- | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | `float` | The rate at which random numbers are generated. | *required* | | `distribution` | `str` | The name of the random distribution to sample from. | `'normal'` | | `seed` | `int` | An integer seed for the random number generator. If None, a random 32-bit seed will be generated. | `None` | | `dtype` | `DTypeLike` | data type of the random number. If None, the default data type for the specified distribution will be used. Not all distributions support all data types; check the JAX documentation for details. | `None` | | `distribution_parameters` | | A dictionary of additional parameters to pass to the distribution function. | `{}` | #### `with_key(key, **kwargs)` Construct RandomNumber with an explicit JAX PRNGKey. Use this when you need independent noise streams in batched (jax.vmap) simulations. Example keys = jax.random.split(jax.random.PRNGKey(0), 16) ##### Each diagram gets a different key diagrams = \[ build_diagram_with( RandomNumber.with_key(keys[i], ...) ) for i in range(16) \] ##### OR with with_parameters (preferred): diagram.with_parameters({"noise.key": keys[i]}) Parameters: | Name | Type | Description | Default | | ---------- | ------------- | ---------------------------------------------- | ---------- | | `key` | `'jax.Array'` | JAX PRNGKey array (shape (2,) for default RNG) | *required* | | `**kwargs` | | other constructor arguments | `{}` | ### `RandomSource` Bases: `LeafSystem` Multi-distribution discrete-time random source. Unified rebuild of the single-distribution `UniformRandomNumber` pattern from T-122 phase 1: one block, four distributions, selected at construction by a string flag plus a `params` dict. Supported distributions:: ``` distribution="uniform" params={"low": ..., "high": ...} distribution="normal" params={"mean": ..., "std": ...} distribution="lognormal" params={"mu": ..., "sigma": ...} distribution="triangular" params={"low": ..., "peak": ..., "high": ...} distribution="exponential" params={"rate": ...} distribution="poisson" params={"rate": ...} # integer-typed output distribution="bernoulli" params={"p": ...} # integer-typed 0/1 output distribution="beta" params={"alpha": ..., "beta": ...} distribution="gamma" params={"shape": ..., "scale": ...} distribution="weibull" params={"shape": ..., "scale": ...} distribution="pareto" params={"scale": ..., "alpha": ...} ``` `"exponential"` is differentiable through `rate` via the standard inverse-CDF reparameterisation `x = -log(1-u) / rate`; `"poisson"` is the discrete count distribution and is *not* differentiable through `rate` w.r.t. its samples (the per-sample grad is zero by construction — the sampler is wrapped in `stop_gradient`). See T-122-followup-poisson. Same seed -> bit-identical sequence (determinism contract). Under `simulate_batch(use_vmap=True)` / `simulate_distributed`, pass `fold_in_batch_index=True` (T-122-followup-vmap-fold-in) to derive a per-replica independent stream from the same master seed via `jax.lax.axis_index("batch")`. Default `False` preserves bit-identical behaviour with the original distributions follow-up. All named `params` flow through smooth, differentiable reparameterisations of an underlying `Uniform[0,1)` or `N(0,1)` draw; the random draw itself is wrapped in `lax.stop_gradient` so gradients of downstream losses flow cleanly through the distribution parameters but never attempt to differentiate the PRNG key. Input ports None. Output ports (0) The most recent sample. Parameters: | Name | Type | Description | Default | | -------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------- | ----------- | | `sample_time` | `float` | Period (s) at which a fresh sample is drawn. | *required* | | `distribution` | `str` | One of "uniform", "normal", "lognormal", "triangular", "exponential", "poisson". | `'uniform'` | | `params` | `dict` | Dict of distribution parameters (see above for keys). Each value is registered as a dynamic parameter and is differentiable / vmap-mappable. | `None` | | `seed` | `int` | Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random. | `None` | | `shape` | | Output shape. Default () (scalar). | `()` | Notes Honest-fallback note: the spec contemplated a `lax.switch` over distributions to share one block class. Because each distribution has different `params` keys (and `triangular` has three), a runtime switch would require padding/aligning the param tuples per distribution, which is clunky and defeats the whole point of named params. Instead we dispatch at *Python time* on the static `distribution` flag -- different distributions trace into different compute graphs, which is exactly the JAX-idiomatic path for static-flag polymorphism. ### `RateLimiter` Bases: `LeafSystem` Limit the time derivative of the block output. Given an input signal `u` computes the derivative of the output signal as: ``` y_rate = (u(t) - y(Tprev))/(t - Tprev) ``` Where Tprev is the last time the block was called for output update. When y_rate is greater than the upper_limit, the output is: ``` y(t) = (t - Tprev)*upper_limit + y(Tprev) ``` When y_rate is less than the lower_limit, the output is: ``` y(t) = (t - Tprev)*lower_limit + y(Tprev) ``` If the lower_limit is greater than the upper_limit, and both are being violated, the upper_limit takes precedence. Optionally, the block can also be configured with "dynamic" limits, which will add input ports for time-varying upper and lower limits. Presently, the block is constrainted to periodic updates. Input ports (0) The input signal. (1) The upper limit, if dynamic limits are enabled. (2) The lower limit, if dynamic limits are enabled. (Will be indexed as 1 if dynamic upper limits are not enabled.) Output ports (0) The rate limited output signal. Parameters: | Name | Type | Description | Default | | ---------------------------- | ---- | --------------------------------------------------------------------------------- | ------- | | `upper_limit` | | The upper limit of the input signal. Default is np.inf. | `inf` | | `enable_dynamic_upper_limit` | | If True, then the upper limit can be set by an external signal. Default is False. | `False` | | `lower_limit` | | The lower limit of the input signal. Default is -np.inf. | `-inf` | | `enable_dynamic_lower_limit` | | If True, then the lower limit can be set by an external signal. Default is False. | `False` | T-115-followup-mode-flag The `mode` kwarg unifies the smooth (differentiable) variant previously exposed as :class:`SoftRateLimiter`. `mode="hard"` (default) is byte-equivalent to the legacy behavior. `mode="smooth"` replaces the inner per-step delta clip with :func:`soft_saturate` so gradients flow through active rate limiting. Smooth mode requires finite (static) `upper_limit` / `lower_limit` and `sharpness > 0` (defaults to `10.0`). #### `initialize_static_data(context)` Infer the size and dtype of the internal states ### `Reciprocal` Bases: `FeedthroughBlock` Compute the reciprocal of the input signal. Input ports (0) The input signal. Output ports (0) The reciprocal of the input signal: `y = 1 / u`. ### `RecursiveLeastSquares` Bases: `LeafSystem` Recursive Least Squares (RLS) estimator for online parameter identification in linear-in-parameters models: ``` ``y[k] = φ[k]ᵀ θ + noise`` ``` where - `y[k]` is a scalar (or vector) measurement at timestep k, - `φ[k]` is a regressor vector of size `n_params`, - `θ` is the unknown parameter vector to be estimated. The RLS update equations with forgetting factor λ are: .. code-block:: text ``` e[k] = y[k] − φ[k]ᵀ θ̂[k−1] (prediction error) K[k] = P[k−1] φ[k] / (λ + φ[k]ᵀ P[k−1] φ[k]) (Kalman gain) θ̂[k] = θ̂[k−1] + K[k] e[k] (parameter update) P[k] = (P[k−1] − K[k] φ[k]ᵀ P[k−1]) / λ (covariance update) ``` A forgetting factor `λ < 1` down-weights older measurements, making the estimator track slowly time-varying parameters. `λ = 1` (default) is the classic batch RLS equivalent. The block is fully JAX-traceable and compatible with JIT/autodiff. ``` +--------------------+ --- phi[k] ---->| |----> theta_hat[k] | Recursive Least |----> P[k] --- y[k] ------>| Squares |----> prediction_error[k] +--------------------+ ``` Input ports (0) phi : regressor vector at timestep k, shape `(n_params,)` (1) y : scalar measurement at timestep k Output ports (0) theta_hat : parameter estimate, shape `(n_params,)` (1) P : parameter covariance matrix, shape `(n_params, n_params)` (2) prediction_error : scalar prediction residual `e = y − φᵀ θ̂` Parameters: | Name | Type | Description | Default | | ------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Sampling period. | *required* | | `n_params` | | int Number of parameters to estimate (dimension of θ). | *required* | | `theta_0` | | array_like, optional Initial parameter estimate, shape (n_params,). Defaults to the zero vector. | *required* | | `P_0` | | array_like, optional Initial covariance matrix, shape (n_params, n_params). Defaults to 1e4 * I, which encodes high initial uncertainty. | *required* | | `forgetting_factor` | | float, optional Forgetting factor λ ∈ (0, 1\]. Default 1.0 (no forgetting). | *required* | Example:: ``` import numpy as np import jaxonomy from jaxonomy import library, DiagramBuilder, SimulatorOptions # True parameters: y = 2*phi_0 + 3*phi_1 TRUE_THETA = np.array([2.0, 3.0]) DT = 0.1 rls = library.RecursiveLeastSquares( dt=DT, n_params=2, forgetting_factor=1.0, ) ``` #### `DiscreteStateType` Bases: `NamedTuple` Internal state: current parameter estimate and covariance. #### `initialize(dt, n_params, theta_0=None, P_0=None, forgetting_factor=1.0)` Called at context-creation time to store resolved parameters. ### `ReducedOrderModel` A reduced model plus its provenance. Attributes: | Name | Type | Description | | --------------- | --------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `system` | `Any` | The reduced, simulatable Jaxonomy block — an LTISystem / LTISystemDiscrete for linear MOR, or a discrete-time predictor LeafSystem for data-driven methods. Drop it straight into a diagram or jaxonomy.simulate. | | `method` | `str` | The reduction method that produced it. | | `full_order` | `Optional[int]` | State dimension of the source model (when known). | | `reduced_order` | `Optional[int]` | State dimension of system. | | `info` | `dict` | Method-specific extras — e.g. error_bound and hsv for balanced truncation, eigenvalues / basis for DMD, and the raw result object from the underlying routine. | #### `to_block()` Return the reduced Jaxonomy block (alias for `.system`). ### `ReferenceSubdiagram` Registry for reusable diagram templates ("reference subdiagrams"). A reference subdiagram is a parameterized diagram factory. It is registered once via :meth:`register` and can then be instantiated multiple times with different parameter values via :meth:`create_diagram`. Example:: ``` def my_submodel(instance_name, parameters): builder = DiagramBuilder() gain = parameters["gain"].get() ... return builder.build(instance_name) ref_id = ReferenceSubdiagram.register( my_submodel, default_parameters=[Parameter("gain", 1.0)], ) diagram = ReferenceSubdiagram.create_diagram(ref_id, "my_instance") ``` #### `create_diagram(ref_id, instance_name, *args, instance_parameters=None, **kwargs)` Create a diagram based on the given reference ID and parameters. Note that for submodels we evaluate all parameters, there is no "pure" string parameters. Parameters: | Name | Type | Description | Default | | --------------------- | ---------------- | ------------------------------------------------------------------------------------------------------------- | ---------- | | `ref_id` | `str` | The reference ID of the diagram. | *required* | | `instance_name` | `str` | Name for this specific instance. | *required* | | `*args` | | Variable length arguments passed to the constructor. | `()` | | `instance_parameters` | `dict[str, Any]` | Per-instance parameter overrides. Keys must match names declared at registration time. Example: {"gain": 3.0} | `None` | | `**kwargs` | | Keyword arguments passed to the constructor. | `{}` | Returns: | Name | Type | Description | | --------- | --------- | -------------------- | | `Diagram` | `Diagram` | The created diagram. | Raises: | Type | Description | | ------------ | --------------------------------------------------------------------- | | `ValueError` | If the reference subdiagram with the given ref_id is not found. | | `ValueError` | If an instance_parameter key does not match any registered parameter. | #### `get_default_parameters(ref_id)` Return the default parameters for the given reference subdiagram. #### `get_parameter_definitions(ref_id)` Return the default parameters for the given reference subdiagram. .. deprecated:: Use :meth:`get_default_parameters` instead. #### `register(constructor, default_parameters=None, ref_id=None, parameter_definitions=None)` Register a diagram constructor as a reusable reference subdiagram. Parameters: | Name | Type | Description | Default | | ----------------------- | ----------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------- | | `constructor` | `ReferenceSubdiagramProtocol` | A callable that builds a :class:Diagram given instance_name and parameters. | *required* | | `default_parameters` | `list[Parameter]` | Default :class:Parameter values for this subdiagram. Instances can override individual parameters at creation time via :meth:create_diagram. | `None` | | `ref_id` | \`str | None\` | Optional stable identifier. A UUID is generated if omitted. | | `parameter_definitions` | `list[Parameter]` | Deprecated – use default_parameters. | `None` | Returns: | Name | Type | Description | | ----- | ----- | ------------------------------------------------------ | | `str` | `str` | The ref_id that can be passed to :meth:create_diagram. | ### `Relay` Bases: `LeafSystem` Simple state machine implementing hysteresis behavior. The input-output map is as follows: ``` output | on_value | -------<------<--------------------- | | | | ⌄ ^ | | | off_value |----------|-------->----->-----| | |---------------------------------------------- input | off_threshold | on_threshold ``` Note that the "time mode" behavior of this block will follow the input signal. That is, if the input signal varies continuously in time, then the zero-crossing event from OFF->ON or vice versa will be localized in time. On the other hand, if the input signal varies only as a result of periodic updates to the discrete state, the relay will only change state at those instants. If the input signal is continuous, the block can be "forced" to this discrete-time periodic behavior by adding a ZeroOrderHold block before the input. The exception to this is the case where there are no blocks in the system containing either discrete or continuous state. In this case the state changes will only be localized to the resolution of the major step. Input ports (0) The input signal. Output ports (0) The relay output signal, which is equal to either the on_value or the off_value, depending on the internal state of the relay. Parameters: | Name | Type | Description | Default | | --------------- | ---- | ------------------------------------------------------------------------------------------------------------------------ | ---------- | | `on_threshold` | | When input rises above this value, the internal state transitions to ON. | *required* | | `off_threshold` | | When input falls below this value, the internal state transitions to OFF. | *required* | | `on_value` | | Value of the output signal when state is ON. | *required* | | `off_value` | | Value of the output signal when state is OFF | *required* | | `initial_state` | | If equal to on_value, the block will be initialized in the ON state. Otherwise, it will be initialized to the OFF state. | *required* | Events There are two zero-crossing events: one to transition from OFF->ON and one for the opposite transition from ON->OFF. ### `ReplicatedFunction` Bases: `LeafSystem` Container block: evaluate a submodel N times in parallel via vmap. Parameters: | Name | Type | Description | Default | | ---------- | --------------- | -------------------------------------------------------------------------------------------------------- | ---------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output. Must be JAX-traceable so vmap can transform it. | *required* | | `n` | `int` | Number of replicas. | *required* | | `n_inputs` | `int` | Number of input ports the block should declare (and the number of positional inputs the submodel takes). | `1` | | `in_axes` | \`Sequence\[int | None\] | None\` | | `name` | | Optional block name. | *required* | ### `RigidBody` Bases: `LeafSystem` Implements dynamics of a single three-dimensional body. The block models both translational and rotational degrees of freedom, for a total of 6 degrees of freedom. With second-order equations, the block has 12 state variables, 6 for the position/orientation and 6 for the velocities/rates. Currently only a roll-pitch-yaw (Euler angle) representation is supported for the orientation. The full 12-dof state vector is `x = [p_i, Φ, vᵇ, ωᵇ]`, where `pⁱ` is the position in an inertial "world" frame `i`, `Φ` is the (roll, pitch, and yaw) Euler angle sequence defining the rotation from the inertial "world" frame to the body frame, `vᵇ` is the translational velocity with respect to body-fixed axes `b`, and `ωᵇ` is the angular velocity about the body-fixed axes. The mass and inertia properties of the block can independently be defined statically as parameters, or dynamically as inputs to the block. Input ports (0) force_vector: 3D force vector, defined in the *body-fixed* coordinate frame. For example, if gravity is acting on the body, the gravity vector should be pre-rotated using CoordinateRotation. (1) torque_vector: 3D torque vector, be defined in the *body-fixed* coordinate frame. (2) inertia: If `enable_external_inertia_matrix=True`, this input provides the time-varying body-fixed inertia matrix. Output ports (0): The position in the inertial "world" frame `pⁱ`. (1): The orientation of the body, represented as a roll-pitch-yaw Euler angle sequence. (2): The translational velocity with respect to body-fixed axes `vᵇ`. (3): The angular velocity about the body-fixed axes `ωᵇ`. (4): (if `enable_output_state_derivatives=True`) The time derivatives of the position variables in the world frame `ṗⁱ`. Not generally equal to the state `vᵇ`, defining time derivatives in the body frame. (5): (if `enable_output_state_derivatives=True`) The "Euler rates" `Φ̇`, which are the time derivatives of the Euler angles. Not generally equal to the angular velocity `ωᵇ`. (6): (if `enable_output_state_derivatives=True`) The body-fixed acceleration vector `aᵇ`. (7): (if `enable_output_state_derivatives=True`) The angular acceleration in body-fixed axes `ω̇ᵇ`. Parameters: | Name | Type | Description | Default | | --------------------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `initial_position` | `Array` | The initial position in the inertial frame. | *required* | | `initial_orientation` | `Array` | The initial orientation of the body, represented as a roll-pitch-yaw Euler angle sequence. | *required* | | `initial_velocity` | `Array` | The initial translational velocity with respect to body-fixed axes. | *required* | | `initial_angular_velocity` | `Array` | The initial angular velocity about the body-fixed axes. | *required* | | `enable_external_mass` | `bool` | If True, the block will have one input port for the mass. Otherwise the mass must be provided as a block parameter. | `False` | | `mass` | `float` | The constant value for the body mass when enable_external_mass=False. If None, will default to 1.0. | `1.0` | | `enable_external_inertia_matrix` | `bool` | If True, the block will have one input port for a (3x3) inertia matrix. Otherwise the inertia matrix must be provided as a block parameter. | `False` | | `inertia_matrix` | | The constant value for the body inertia matrix when enable_external_inertia_matrix=False. If None, will default to the 3x3 identity matrix. | `eye(3)` | | `enable_output_state_derivatives` | `bool` | If True, the block will output the time derivatives of the state variables. | `False` | | `gravity_vector` | `Array` | The constant gravitational acceleration vector acting on the body, defined in the inertial frame. If None, will default to the zero vector. | `zeros(3)` | Notes Assumes that the inertia matrix is computed at the center of mass. Assumes that the mass and inertia matrix are quasi-steady. This means that if one or both is specified as "dynamic" inputs their time derivative is neglected in the dynamics. For instance, for pure translation (`w_b=0`) the approximation to Newton's law is `F_net = (d/dt)(m * v) ≈ m * (dv/dt)`. ### `Ros2Publisher` Bases: `LeafSystem` Ros2Publisher block can emit signals to a ROS2 topic, based on input signal data. #### `__init__(dt, topic, msg_type, fields, **kwargs)` Publish messages to a ROS2 topic. Parameters: | Name | Type | Description | Default | | ---------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | `float` | Period of the system, in both sim and real (ros2) time. | *required* | | `topic` | `str` | ROS2 topic to publish to. Eg. /turtle1/cmd_vel. | *required* | | `msg_type` | `type` | ROS2 message type, e.g. Twist from geometry_msgs.msg. Unlike the corresponding UI parameter, this must be a Python type object. | *required* | | `fields` | `dict[str, type]` | Ordered dictionary of default values to extract from the received message. The keys are the full attribute path (with dots) to the value in the message, and the values are the default values. This is used to create the output ports with valid data types. Use Python or Numpy data types, not JAX. For instance, for a `geometry_msgs.msg.Twist` message, the `fields` could be `{"linear.x": float, "angular.z": float}`. | *required* | ### `Ros2Subscriber` Bases: `LeafSystem` Ros2Subscriber block listens to messages over a ROS2 topic and outputs them as signals in jaxonomy. #### `__init__(dt, topic, msg_type, fields, read_before_start=True, **kwargs)` Subscribe to a ROS2 topic and extract message values to output ports. Parameters: | Name | Type | Description | Default | | ------------------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | Period of the system, in both sim and real (ros2) time. | *required* | | `topic` | `str` | ROS2 topic to subscribe to. Eg. /turtle1/pose. | *required* | | `msg_type` | `type` | ROS2 message type, e.g. Pose from turtlesim.msg. Unlike the corresponding UI parameter, this must be a Python type object. | *required* | | `fields` | `dict[str, type]` | Ordered dictionary of default values to extract from the received message. The keys are the full attribute path (with dots) to the value in the message, and the values are the default values. This is used to create the output ports with valid data types. Use Python or Numpy data types, not JAX. For instance, for a `geometry_msgs.msg.Twist` message, the `fields` could be `{"linear.x": float, "angular.z": float}`. | *required* | | `read_before_start` | | If True, the subscriber will read the first message before the simulation starts. Otherwise, the initial outputs will be 0. | `True` | ### `Sawtooth` Bases: `SourceBlock` Produces a modulated linear sawtooth signal. The signal is similar to: https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.sawtooth.html Given amplitude `a`, period `p`, and phase delay `phi`, the output signal is: ``` y(t) = a * ((t - phi) % p) ``` where `%` is the modulo operator. Input ports None Output ports (0) The sawtooth signal. ### `ScalarBroadcast` Bases: `FeedthroughBlock` Broadcast a scalar to a vector or matrix. Given a scalar input `u` and dimensions `m` and `n`, this block will return a vector or matrix of shape `(m, n)` with all elements equal to `u`. Input ports (0) The scalar input signal. Output ports (0) The broadcasted output signal. Parameters: | Name | Type | Description | Default | | ---- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `m` | | The number of rows in the output matrix. If m is None, then the output will be a vector with shape (n,). To get a row vector of size (1,n), set m=1 expliclty. | *required* | | `n` | | The number of columns in the output matrix. If n is None, then the output will be a vector with shape (m,). To get a column vector of size (m,1), set n=1 expliclty. | *required* | ### `ShiftRegister` Bases: `LeafSystem` Fixed-length shift register delay line. Delays an input signal by exactly n_steps discrete timesteps. Output at time t is the input value from n_steps timesteps ago. Parameters: | Name | Type | Description | Default | | --------------- | -------------- | --------------------------------------------------------------------------------------------------- | ---------- | | `n_steps` | `int` | Number of steps to delay. STATIC — set at construction, cannot be changed at runtime. Must be >= 1. | *required* | | `signal_shape` | `tuple` | Shape of each signal frame. Use () for scalar, (3,) for 3-vector, etc. | `()` | | `initial_value` | `array - like` | Value to fill the buffer with before any input has been received. Default: zeros. | `None` | | `dt` | `float` | Discrete update interval in seconds. | `0.01` | Ports Input[0] "u": signal to delay, shape=signal_shape Output[0] "y": delayed signal, shape=signal_shape ### `SignalDatatypeConversion` Bases: `FeedthroughBlock` Convert the input signal to a different data type. Input ports: (0) The input signal. Output ports: (0) The input signal converted to the specified data type. Parameters: dtype: The data type to which the input signal is converted. Must be a valid NumPy data type, e.g. "float32", "int64", etc. ### `SimulationResultsSource` Bases: `LeafSystem` Replays one recorded trajectory from :class:`~jaxonomy.simulation.types.SimulationResults`. Output port `y` is the signal value at the current simulation time, using linear interpolation or zero-order hold. Values clamp to the first/last sample outside the recorded time range (`jnp.interp` semantics for linear mode). ### `Sindy` Bases: `LeafSystem` This class implements System Identification (SINDy) algorithm with or without control inputs for contiuous-time and discrete-time systems. The learned continuous-time dynamical system model will be of the form: ``` dx/dt = f(x, u) ``` where `x` is the state vector and `u` is the optional control input vector. The block will output the full state vector `x` of the system. The learned discrete-time dynamical system model will be of the form: ``` x_{k+1} = f(x_k, u_k) ``` where `x_k` is the state vector at time step `k` and `u_k` is the optional control vector. The block will update the output to `x_k` at an interval provided by the parameter `discrete_time_update_interval`. Input ports (0) u: control vector for the system. This port is only available if the Sindy model is trained with control inputs, i.e. `control_input_columns` is not `None` during training. Output ports (0) x: full state output of the system. Parameters: | Name | Type | Description | Default | | ------------------------------- | ------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------------- | | `file_name` | `str` | Path to the CSV file containing training data. | `None` | | `header_as_first_row` | `bool` | If True, the first row of the CSV file is treated as the header. | `False` | | `state_columns` | \`int | str | list[int] | | `control_input_columns` | \`int | str | list[int] | | `dt` | `float` | Fixed value of dt if rows of the CSV file represent equidistant time steps. | `None` | | `time_column` | `(str, int)` | Column name (str) for column index (int) for time data t. If time_column is provided, then fixed dt above will be ignored. If neither dt nor time_column is provided, then the SINDy model will use a fixed detault time step of dt=1. | `None` | | `state_derivatives_columns` | \`int | str | list[int] | | `discrete_time` | `bool` | If True, the SINDy model will be trained for discrete-time systems. In this case, the dynamical system is treated as a map. Rather than predicting derivatives, the right hand side functions step the system forward by one time step. If False, dynamical system is assumed to be a flow (right-hand side functions predict continuous time derivatives). See documentation for pysindy. | `False` | | `differentiation_method` | `str` | Method to use for differentiating the state data x to obtain state derivatives dot_x = dx/dt. Available options are: 'centered difference' (default) | `'centered difference'` | | `threshold` | `float` | Threshold for the Sequentially thresholded least squares (STLSQ) algorithm used for training SINDy model. | `0.1` | | `alpha` | `float` | Regularization strength for the STLSQ algorithm. | `0.05` | | `max_iter` | `int` | Maximum number of iterations for the STLSQ algorithm. | `20` | | `normalize_columns` | `bool` | If True, normalize the columns of the data matrix before regression. | `False` | | `poly_order` | `int` | Degree of polynomial features. Set to None to omit this library. | `2` | | `fourier_n_frequencies` | `int` | Number of Fourier frequencies. Set to None to omit this library. | `None` | | `custom_basis_functions` | `list of functions` | A list of custom basis functions to use for training the SINDy model. For example to include f(x) = 1/x and g(x) = exp(-x), provide [lambda x: 1.0/(x.0 + 1e-06), lamda x: jnp.exp(-x)] Currently only supported for jaxonomy interface. Calls from UI and pretrained model loading does not support custom basis functions. | `None` | | `pretrained` | `bool` | If True, use a pretrained model specified by the pretrained_file_path argument. | `False` | | `pretrained_file_path` | `str` | Path to the pretrained model file. | `None` | | `initial_state` | `ndarray` | Initial state of the system for propagating the continuous-time or discrete-time system forward duiring simulation. | `None` | | `discrete_time_update_interval` | `float` | Interval at which the discrete-time model should be updated. Default is 1.0. | `1.0` | | `equations` | `list of strings` | (For internal UI use only) The identified system equations. | `None` | | `base_feature_names` | `list of strings` | (For internal UI use only) Features x_i and u_i. | `None` | | `feature_names` | `list of strings` | (For internal UI use only) Composed features with basis libraries. | `None` | | `coefficients` | `ndarray` | (For internal UI use only) Coefficients of the identified model. | `None` | | `has_control_input` | `bool` | (For internal UI use only) If True, the model was trained with control. For standard training from CSV file, this is inferred from the parameter control_input_columns. | `True` | #### `serialize(filename)` Save the relevant class attributes post training so that model state can be restored #### `serialize_trained_pysindy_model(model, filename)` Serialize a PySindy model trained outside of Jaxonomy. The saved file can be used as a pretrained model in Jaxonomy. ### `Sine` Bases: `SourceBlock` Generates a sinusoidal signal. Given amplitude `a`, frequency `f`, phase `phi`, and bias `b`, the output signal is: ``` y(t) = a * sin(f * t + phi) + b ``` Input ports None Output ports (0) The sinusoidal signal. Parameters: | Name | Type | Description | Default | | ----------- | --------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `amplitude` | | The amplitude of the sinusoidal signal. | `1.0` | | `frequency` | | The frequency of the sinusoidal signal. | `1.0` | | `phase` | | The phase of the sinusoidal signal. | `0.0` | | `bias` | | The bias of the sinusoidal signal. | `0.0` | | `units` | `(optional, T - 104 - followup - units - on - source - blocks)` | If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams). | `None` | ### `Slice` Bases: `FeedthroughBlock` Slice the input signal using Python indexing rules. Input ports (0) The input signal. Output ports (0) The sliced output signal. Parameters: | Name | Type | Description | Default | | -------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `slice_` | | The slice operator to apply to the input signal. Must be specified as a string input, e.g. the output u[1:3] would be created with the block Slice("1:3"). | *required* | Notes Currently only up to 3-dimensional slices are supported. ### `SnapshotData` Container for a column-wise snapshot matrix. Attributes: | Name | Type | Description | | -------- | ------------------- | ------------------------------------------------------------------ | | `X` | `ndarray` | State/output snapshots, shape (n_features, n_samples). | | `time` | `Optional[ndarray]` | Optional sample times, shape (n_samples,). | | `inputs` | `Optional[ndarray]` | Optional input snapshots U, shape (n_inputs, n_samples). | | `Xdot` | `Optional[ndarray]` | Optional time-derivative snapshots, shape (n_features, n_samples). | ### `SoftRateLimiter` Bases: `LeafSystem` Smooth (differentiable) rate limiter. Drop-in differentiable variant of :class:`RateLimiter`. Identical discrete update semantics, except the inner hard `clip` on the desired step `(u - y_prev)` is replaced by a smooth saturation so gradients flow through the limiter even when it is actively limiting. The smooth clip is implemented via :func:`soft_saturate` and recovers the hard rate limiter as `sharpness -> inf`. Parameters mirror :class:`RateLimiter` plus: sharpness: Scalar > 0 controlling how sharply the smooth saturation transitions at the rate limits. Larger `sharpness` -> closer to the hard rate limiter. Default `10.0`. See :class:`RateLimiter` for the (non-smoothed) reference behavior. ### `SoftSaturate` Bases: `FeedthroughBlock` Smooth (differentiable) saturation block. Drop-in differentiable variant of :class:`Saturate` that uses :func:`soft_saturate` instead of `npa.clip`. The original hard :class:`Saturate` block is unchanged. Why a separate block: the standard :class:`Saturate` block returns `npa.clip(u, lo, hi)`, whose gradient is exactly zero outside the bounds. That kills gradient signal in any optimization that drives the input past the limits. `SoftSaturate` keeps gradient flow alive so e.g. trajectory optimization through actuator limits actually converges. See :func:`soft_saturate` for the formula. Unlike :class:`Saturate`, this block does *not* declare zero-crossing events (it has no discontinuity to catch). Parameters: | Name | Type | Description | Default | | ------------- | ---- | ------------------------------------------------------------------------------------------------------ | ------- | | `upper_limit` | | Upper limit; default 1.0. Must be finite. | `1.0` | | `lower_limit` | | Lower limit; default 0.0. Must be finite and strictly less than upper_limit. | `0.0` | | `sharpness` | | Smoothing knob, > 0; default 10.0. As sharpness -> inf this approaches the hard :class:Saturate block. | `10.0` | Input ports (0) The input signal. Output ports (0) The smoothly-saturated output signal. ### `SourceBlock` Bases: `LeafSystem` Simple blocks with a single time-dependent output #### `__init__(func, **kwargs)` Create a source block with a time-dependent output. Parameters: | Name | Type | Description | Default | | ------ | ---------- | ----------------------------------------------------------------------------------------------------------------------- | ---------- | | `func` | `Callable` | A function of time and parameters that returns a single value. Signature should be func(time, \*\*parameters) -> Array. | *required* | ### `SquareRoot` Bases: `FeedthroughBlock` Compute the square root of the input signal. Dispatches to `jax.numpy.sqrt`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.sqrt.html Input ports (0) The input signal. Output ports (0) The square root of the input signal. ### `Stack` Bases: `ReduceBlock` Stack the input signals into a single output signal along a new axis. Dispatches to `jax.numpy.stack`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.stack.html Input ports (0..n_in-1) The input signals. Output ports (0) The stacked output signal. Parameters: | Name | Type | Description | Default | | ------ | ---- | ----------------------------------------------------------------- | ------- | | `axis` | | The axis along which the input signals are stacked. Default is 0. | `0` | ### `StateMachine` Bases: `LeafSystem` Finite State Machine similar to Mealy Machine. https://en.wikipedia.org/wiki/Mealy_machine The state machine can be executed either periodically or by zero_crossings. Each state as 0 or more exit transitions. These are prioritized such that when 2 exits are simultaneously valid, the higher priority is executed. It is not allowed for a state to have more than one exit transition with no guard. Guardless exits only make sense in the periodic case. Each transitions may have 0 or more actions. Each action is a python statement that modifies the value of an output. When a transitions is executed (i.e. it's guard evaluates to true), its actions are then processed. If 'time' is needed for guards or actions, pass 'time' in from clock block. Whether executed periodically or by zero_crossings, the states are constant between transitions executions. In the zero_crossing case, all guards for transitions exiting the current state are continuously checked, and if any 'triggers', then the earlist point in time that any guard becomes true is determined, the actions of the earliest (and highest priority if multiple trigger simultaneously) guard are executed at that time, and the simulation continues afterwards. Input ports User specified. Output ports User specified. Parameters: | Name | Type | Description | Default | | --------------------- | ------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `dt` | | Either Float or None. When not None, state machine is executed periodically. When None, the transitions are monitored by zero_crossing events. | `None` | | `accelerate_with_jax` | `bool` | Bool. When True, the actions and guards are JIT-compiled with JAX. Default is False. | `False` | ### `Step` Bases: `SourceBlock` A step signal. Given start value `y0`, end value `y1`, and step time `t0`, the output signal is: ``` y(t) = y0 if t < t0 else y1 ``` Input ports None Output ports (0) The step signal. Parameters: | Name | Type | Description | Default | | ------------- | --------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `start_value` | | The value of the output signal before the step time. | `0.0` | | `end_value` | | The value of the output signal after the step time. | `1.0` | | `step_time` | | The time at which the step occurs. | `1.0` | | `units` | `(optional, T - 104 - followup - units - on - source - blocks)` | If set, the output port advertises this :class:Unit. The connect-time consistency check (T-104) then enforces downstream ports declare a compatible unit. Default None keeps the legacy "no-units" behaviour (byte-equivalent to pre-T-104 diagrams). | `None` | ### `Stop` Bases: `LeafSystem` Stop the simulation early as soon as the input signal becomes True. If the input signal changes as a result of a discrete update, the simulation will terminate the major step early (before advancing continuous time). Input ports (0): the boolean- or binary-valued termination signal Output ports None ### `SumOfElements` Bases: `FeedthroughBlock` Compute the sum of the elements of the input signal. Dispatches to `jax.numpy.sum`, so see the JAX docs for details: https://jax.readthedocs.io/en/latest/\_autosummary/jax.numpy.sum.html Input ports (0) The input signal. Output ports (0) The sum of the elements of the input signal. ### `Switch` Bases: `LeafSystem` Route one of two data signals based on a thresholded control signal. Three inputs `(data_a, control, data_b)` and one output: .. code-block:: python ``` y = data_a if criteria(control, threshold) else data_b ``` Default `mode="where"` is implemented via `npa.where`, so JAX gradients flow through *both* data branches simultaneously (the selector branch is treated as non-differentiable, which is the only well-defined choice for a hard threshold). `mode="smooth"` replaces the hard `where` with a sigmoid blend .. code-block:: python ``` alpha = sigmoid(sharpness * sign * (control - threshold)) y = alpha * data_a + (1 - alpha) * data_b ``` where `sign` is +1 for the `>=`/`>` criteria and -1 for the `<=`/`<` criteria, so the smooth output approaches the hard answer in the strict-active region as `sharpness -> inf`. This mode lets gradients flow through the *threshold itself*, which the hard `where` zeroes out — the killer feature for trajectory optimization where the threshold is a tunable parameter. `data_a` and `data_b` must be broadcast-compatible (same as `npa.where`'s requirements). The block does not enforce that `control` is a scalar — element-wise selection is supported when `control` and the data inputs broadcast together. Parameters: | Name | Type | Description | Default | | ----------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | | `threshold` | | scalar threshold against which control is compared. | `0.0` | | `criteria` | | one of ">=", ">", "\<=", "\<", "==", "!=". Default ">=". The equality criteria ("=="/"!=") are not supported in mode="smooth" (no sigmoid approximation makes sense for them). | `'>='` | | `mode` | | one of "where" (default), "smooth", or "hard". "where" is byte-equivalent to T-118 phase 1 and the right pick for simulation. "smooth" is the right pick for gradient-based optimization through the threshold. "hard" dispatches to jax.lax.cond so only the active branch is evaluated — useful when one branch is much more expensive than the other or when branches have incompatible side effects. Switch(mode='hard') is incompatible with vmap; use mode='where' or mode='smooth' for batched use (e.g. under simulate_batch). On older JAX (\<0.4) a batched predicate raises TracerBoolConversionError; on modern JAX the cond is silently rewritten to evaluate both branches with a select, which is numerically correct but defeats the entire point of picking mode='hard' over mode='where'. | `'where'` | | `sharpness` | | positive scalar controlling sigmoid steepness in mode="smooth". Default 10.0. Larger values give a tighter approximation to the hard switch but smaller (and faster vanishing) gradients in the strict-active region. Ignored when mode="where". | `10.0` | Input ports (0) data_a — output when criteria(control, threshold) is True. (1) control — the selector signal compared to `threshold`. (2) data_b — output when criteria(control, threshold) is False. Output ports (0) The selected data signal, with shape determined by broadcasting between data_a and data_b. ### `TableSearch` Bases: `LeafSystem` Search a monotonic table for the bucket containing a query value. Given a strictly-increasing 1-D grid `xp` of length `n` and a scalar query `x`, returns the bucket index `i` (as a float) such that `xp[i] <= x < xp[i+1]`. Out-of-range queries clamp to the nearest endpoint: `x < xp[0]` returns `0`; `x >= xp[-1]` returns `n - 1`. The standard "Direct Lookup" pattern. Different from :class:`Prelookup` in that the output is just the bucket index -- no fractional `alpha` is computed. Useful for binning, threshold detection, and inverse-table indexing. Input ports `(0)` -- scalar query coordinate `x`. Output ports `(0)` -- scalar bucket index, returned as a float (so it composes with the float-defaulting numeric pipeline). The output is wrapped in `jax.lax.stop_gradient` -- gradient is zero almost everywhere by construction (step function), so we make that non-differentiability explicit to avoid spurious grad-flow surprises. Parameters: | Name | Type | Description | Default | | ------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `xp` | | 1-D, strictly-monotonically-increasing grid of breakpoints (length >= 2). Stored verbatim for the bucket search. | *required* | | `mode` | | "binary" (default) uses jnp.searchsorted (O(log n)); "linear" uses a linear scan via a sum over a comparison mask (O(n)). Both modes return byte-identical results on valid (strictly-monotonic) grids. | `'binary'` | | `dtype` | `optional` | If set (e.g. jnp.float32), the grid array is cast to this dtype on construction. Mirrors the per-block dtype contract of :class:LookupTable1d / :class:Prelookup. | `None` | Notes Index is wrapped in `jax.lax.stop_gradient` -- the gradient through the query coordinate is zero, by construction. Callers who need a differentiable index-like quantity should use :class:`Prelookup` (which exposes the fractional `alpha`) or :class:`LookupTable1d` directly. #### `mode` Search mode (`"binary"` or `"linear"`). #### `xp` The 1-D strictly-increasing breakpoint array. ### `TensorFlow` Bases: `LeafSystem` Block to perform inference with a pre-trained TensorFlow SavedModel. The input to the block should be of compatible type and shape expected by the TensorFlow model. For example, if the TensorFlow SavedModel model expects a `tf.float32` tensor of shape `(3, 224, 224)`, the input to the block should be a `jax.numpy` array of shape (3, 224, 224) of dtype `jnp.float32`. For output types, if no casting is specified through the `cast_outputs_to_dtype` parameter, the output of the block will have the same dtype as the TensorFlow model output, but expressed as `jax.numpy` types. For example. if the TensorFlow model outputs a `tf.float32` tensor, the output of the block will be a `jax.numpy` array of dtype `jnp.float32`. If casting is specified through `cast_outputs_to_dtype` parameter, all the outputs, of the block will be casted to this specific `jax.numpy` dtype. .. note:: **float32 models under jaxonomy's global x64.** `import jaxonomy` enables JAX 64-bit mode (`jax_enable_x64`) for the whole process, so upstream signals are float64 by default. A SavedModel with `tf.float32` signatures therefore receives float64 inputs (a dtype error or a silent arithmetic change) unless you cast at the block boundary. One-line idiom: pass `cast_outputs_to_dtype="float32"` and feed the block `x.astype(jnp.float32)` inputs. Input ports (i) The ith input to the model. Output ports (j) The jth output of the model. Parameters: | Name | Type | Description | Default | | ------------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `file_name` | `str` | Path to the model file. This should be a .zip containing the SavedModel. | *required* | | `cast_outputs_to_dtype` | `str` | The dtype to cast all the outputs of the block to. Must correspond to a jax.numpy datatype. For example, "float32", "float64", "int32", "int64". | `None` | | `add_batch_dim_to_inputs` | `bool` | Whether to add a new first dimension to the inputs before evaluating the TorchScript or TensorFlow model. This is useful when the model expects a batch dimension. | `False` | #### `initialize_static_data(context)` Infer the output shapes and dtypes of the ML model. ### `TransferFunction` Bases: `LTISystem` Continuous-time LTI system specified as a transfer function. The transfer function is converted to state-space form using `scipy.signal.tf2ss`. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.tf2ss.html The resulting system will be in canonical controller form with matrices (A, B, C, D), which are then used to create an LTISystem. Note that this only supports single-input, single-output systems. Input ports (0) u: Input vector (scalar) Output ports (0) y: Output vector (scalar). Note that this is feedthrough from the input port iff D is nonzero. Parameters: | Name | Type | Description | Default | | ----- | ---- | -------------------------------------------------------------- | ---------- | | `num` | | Numerator polynomial coefficients, in descending powers of s | *required* | | `den` | | Denominator polynomial coefficients, in descending powers of s | *required* | ### `TransferFunctionDiscrete` Bases: `LTISystemDiscrete` Implements a Discrete Time Transfer Function. https://en.wikipedia.org/wiki/Z-transform#Transfer_function The resulting system will be in canonical controller form with matrices (A, B, C, D), which are then used to create an LTISystem. Note that this only supports single-input, single-output systems. Input ports (0) u\[k\]: Input vector (scalar) Output ports (0) y\[k\]: Output vector (scalar). Note that this is feedthrough from the input port if and only if D is nonzero. Parameters: | Name | Type | Description | Default | | ------------------- | ---- | -------------------------------------------------------------- | ---------- | | `dt` | | Sampling period of the discrete system. | *required* | | `num` | | Numerator polynomial coefficients, in descending powers of z | *required* | | `den` | | Denominator polynomial coefficients, in descending powers of z | *required* | | `initialize_states` | | Initial state vector (default: 0) | `None` | ### `TransportDelay` Bases: `LeafSystem` Continuous-time fixed transport (pure) delay. Implements `y(t) = u(t - delay_seconds)` for `t >= delay_seconds`; for `t < delay_seconds` the output is `initial_output` (the standard "Initial output" semantics). The block samples its input on a periodic clock with period `dt` and stores the most recent `history_length` `(time, value)` pairs in a discrete-state ring buffer. Output evaluation at any continuous time `t` is a linear interpolation over the buffered `(time, value)` samples at `t - delay_seconds`. The delay is differentiable via the input signal (gradient flows through `npa.interp` over the values buffer). Differentiability w.r.t. the delay value itself is well-defined wherever the buffer interpolant is differentiable; the linear interpolant has a kink at sample boundaries — pass `method="pchip"` on :class:`VariableTransportDelay` for a C¹-smooth alternative. The buffer is sized statically as `history_length` samples. To cover a delay of `delay_seconds` at sample period `dt`, you need at least `ceil(delay_seconds / dt) + 1` slots; we recommend a small safety margin. `history_length` defaults to `max(8, ceil(delay_seconds / dt) + 4)` which is sufficient for the default constant delay. Input ports (0) The input signal `u(t)`. Scalar or array. Output ports (0) The delayed signal `y(t) = u(t - delay_seconds)` (or `initial_output` while `t < delay_seconds`). Parameters: | Name | Type | Description | Default | | ---------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `dt` | | Sampling period for the history buffer. Smaller dt ⇒ finer interpolation but a larger ring buffer to cover the same physical delay. | *required* | | `delay_seconds` | | Fixed delay τ in seconds. Dynamic parameter (may be tuned via with_parameters); see notes on differentiability above. | *required* | | `initial_output` | | Output value while t < delay_seconds. Default is 0.0. | `0.0` | | `history_length` | | Number of (time, value) pairs stored. Static (compile-time) — required for vmap/JIT-safe buffer sizing. If None, defaults to max(8, ceil(delay_seconds / dt) + 4). | `None` | Notes - For arbitrary array-shaped signals the interpolation is applied elementwise via `jax.vmap` over the trailing axes. - Buffer overflow (delay larger than `history_length * dt`) is not raised; `npa.interp` clamps to the boundary, which means the oldest stored sample is repeated. This is a documented T-107 follow-up; for now, size `history_length` generously. - `VariableTransportDelay` (signal-driven τ) is the natural phase-2 extension; it reuses the same ring-buffer machinery with the delay sourced from an input port. ### `TriggerEdge` Allowed string values for `TriggeredSubsystem.edge`. ### `TriggeredSubsystem` Bases: `LeafSystem` Container block: latch the submodel output on edge transitions (the child still RUNS every step — only the *output* is gated). Important: this does **not** skip execution of the submodel on non-triggered steps. The submodel is evaluated on every step so its inputs participate in the JAX trace; the trigger only controls whether a fresh result is *latched* into the held output. If you need to actually skip computation between triggers, gate it yourself with `jax.lax.cond` at the application level. Phase-1 implementation runs the submodel on every step (so the inputs participate in the trace) but only *latches* a new output on an edge transition of the trigger signal. Between transitions the output holds the most recently latched value. The trigger signal is sampled at `sample_period`. Edges are detected by comparing the current trigger sample against the previously-stored sample held in discrete state. This is *not* the eventual zero-crossing-driven `TriggeredSubsystem` described in the T-120 architecture notes (that requires hooking into the continuous-time event detector); but it is functionally correct for any sample-rate use case and matches the behaviour documented in the test fixtures. Parameters: | Name | Type | Description | Default | | --------------- | ---------------------------------------- | --------------------------------------------------------------------------------------------------- | ---------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output taking the non-trigger user inputs. Must be JAX-traceable. | *required* | | `n_inputs` | `int` | Number of user inputs (NOT counting the trigger). | `1` | | `edge` | `Literal['rising', 'falling', 'either']` | "rising" (low→high), "falling" (high→low) or "either". | `RISING` | | `sample_period` | `float` | Period (seconds) at which the trigger signal is sampled and the latch is updated. Must be positive. | `0.0` | | `initial_value` | | Latched output value before any edge has been detected. Defines output shape/dtype. | `0.0` | | `name` | | Optional block name. | *required* | Limitations (phase 1): - Trigger detection runs on the periodic sample grid, not on continuous-time zero crossings. Trigger pulses shorter than `sample_period` may be missed. - The latch is a single discrete state; the submodel must produce a single output array. - The submodel runs on every output evaluation; only the *output* is gated. Users who need to skip computation on non-triggered steps should use `jax.lax.cond` at the application level. ### `Trigonometric` Bases: `FeedthroughBlock` Apply a trigonometric function to the input signal. Available functions are sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh Dispatches to `jax.numpy.sin`, `jax.numpy.cos`, etc, so see the JAX docs for details. Input ports (0) The input signal. Output ports (0) The trigonometric function applied to the input signal. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `function` | | The trigonometric function to apply to the input signal. Must be one of "sin", "cos", "tan", "asin", "acos", "atan", "sinh", "cosh", "tanh", "asinh", "acosh", "atanh". | *required* | ### `TruthTable` Bases: `LeafSystem` Evaluate a fixed truth table over boolean-castable inputs. Given a list of `(input_pattern, output)` rows, this block compares its inputs against each pattern and emits the output of the first matching row (or `default_output` if none match). Patterns are tuples of `bool` values or the string `"X"` as a wildcard. Example — a 2-input AND gate: .. code-block:: python ``` tt = TruthTable( rows=[ ((True, True), 1.0), ((True, False), 0.0), ((False, True), 0.0), ((False, False), 0.0), ], n_inputs=2, default_output=0.0, ) ``` Wildcard example — ignore the first input: .. code-block:: python ``` tt = TruthTable( rows=[(("X", True), 1.0), (("X", False), 0.0)], n_inputs=2, default_output=0.0, ) ``` Callable output example — row output depends on raw input values (T-119-followup-numeric-output): .. code-block:: python ``` tt = TruthTable( rows=[ ((True, True), lambda a, b: a + b), ((True, False), lambda a, b: a - b), ((False, "X"), 0.0), # constant fallback row ], n_inputs=2, default_output=0.0, ) ``` Parameters: | Name | Type | Description | Default | | ---------------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `rows` | | list of (pattern, output) tuples. pattern is a length-n_inputs tuple whose entries are bool (matched literally) or the string "X" (wildcard, matches anything). output may be a scalar/array (constant for that row) or a callable f(\*inputs) -> scalar_or_vector invoked with the RAW inputs (every row callable is evaluated at every step under JAX's branchless where semantics; the result is selected only when the row matches). All row outputs (and the callable return values) must broadcast against default_output. | *required* | | `n_inputs` | | number of input ports. | *required* | | `default_output` | | value emitted when no row matches. May be a scalar/array (constant fallback) or a callable f(\*inputs) -> scalar_or_vector invoked with the RAW inputs (T-119-followup-default-callable). For a callable default, the output shape/dtype is determined at trace time by the callable's return value; for a constant default, it is determined statically from the value. | *required* | Input ports (0..n_inputs-1) Boolean-castable scalars. Non-boolean inputs are coerced to bool before pattern matching (any non-zero is True). Output ports (0) The output of the first matching row, or `default_output`. Notes Earlier rows take precedence: if multiple patterns would match the same input combination, the one listed first in `rows` wins. The static-completeness/ambiguity checker is deferred (see `T-119-followup-completeness-checker`); JSON serialization of the rows table is deferred (see `T-119-followup-serialization`). #### `builder(n_inputs, default_output, input_names=None, **block_kwargs)` Construct a fluent builder for this truth table. See :class:`TruthTableBuilder` for usage. Equivalent to `TruthTableBuilder(n_inputs, default_output, input_names, **block_kwargs)`. #### `from_csv(path, **block_kwargs)` Load a TruthTable from a CSV file. The CSV must have a header row whose last column(s) are named `output` (single scalar output) or any sequence of columns whose names start with `output` (e.g. `output_x, output_y`) which are stacked into a 1-D vector output per row. All columns preceding the first `output*` column are treated as input columns, in order. Input cells accept `T`/`True`/`1` (True), `F`/`False`/ `0` (False), and `X`/`-`/`*` or an empty cell (wildcard). Matching is case-insensitive and whitespace is stripped. Output cells must parse as `float`. Example CSV:: ``` in1,in2,in3,output T,T,T,1.0 T,T,F,0.5 T,F,X,0.25 F,X,X,0.0 ``` Parameters: | Name | Type | Description | Default | | ---------------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `path` | | filesystem path (str or os.PathLike) to a readable CSV file. | *required* | | `**block_kwargs` | | forwarded to TruthTable.__init__ — typically name / system_id. May also include default_output to override the zero-default that this loader picks (a scalar 0.0 for single-output CSVs, a zeros vector of the right shape for multi-output CSVs). | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ------------------------------------------- | | `A` | | class:TruthTable whose rows mirror the CSV. | Raises: | Type | Description | | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------ | | `ValueError` | if the file is empty, has no header, has no output column, or any row has the wrong number of cells / an unparseable input or output cell. | #### `from_dict(data, **block_kwargs)` Reconstruct a TruthTable from the dict produced by :meth:`to_dict`. Extra keyword arguments (`name=`, `system_id=`, ...) are forwarded to the underlying `TruthTable` constructor, so a deserialized block can pick up a fresh name in its target diagram. Parameters: | Name | Type | Description | Default | | ---------------- | ---- | -------------------------------------------------------- | ---------- | | `data` | | dict with the keys documented on :meth:to_dict. | *required* | | `**block_kwargs` | | forwarded to TruthTable.__init__ (e.g. name, system_id). | `{}` | Returns: | Type | Description | | ---- | ----------------------------------------- | | | A new :class:TruthTable whose rows and | | | default_output match the serialized form. | #### `to_csv(path, **csv_kwargs)` Write this TruthTable to a CSV file (inverse of `from_csv`). Emits a header row of `in1,in2,...,output` (single-output) or `in1,...,output_0,output_1,...` (vector output), followed by one data row per `rows` entry. Input cells are written as `T` / `F` / `X`; output cells are written as `float(...)`. Parameters: | Name | Type | Description | Default | | -------------- | ---- | --------------------------------------------------------------------- | ---------- | | `path` | | filesystem path (str or os.PathLike) for the CSV file to (over)write. | *required* | | `**csv_kwargs` | | forwarded to csv.writer (e.g. delimiter, quoting). | `{}` | Returns: | Type | Description | | ---- | ---------------------------------- | | | path (so the caller can chain | | | TruthTable.from_csv(t.to_csv(p))). | Raises: | Type | Description | | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `ValueError` | if any row's output is a callable (T-119-followup-numeric-output) or default_output is a callable (T-119-followup-default-callable) — callables are not representable in CSV. | #### `to_dict()` Return a JSON-serializable dict describing this TruthTable. The dict round-trips through :meth:`from_dict` to a TruthTable with identical behaviour for every input combination. Pattern wildcards (`"X"`) and vector outputs are preserved. Returns: | Type | Description | | ---- | ---------------------------------------------- | | | dict with keys: | | | n_inputs (int) | | | default_output (float | | | rows (list of {"pattern": str, "output": ...}) | #### `validate(strict_completeness=False, strict_disjointness=False)` Static analysis of the truth-table rows. Enumerates all `2**n_inputs` boolean input vectors and checks: - **Completeness** — each vector is matched by at least one row's pattern (with `"X"` as wildcard). Vectors that no row matches are reported as `missing_patterns`; without coverage they silently hit `default_output` at runtime. - **Disjointness** — no two rows match the same input vector. Overlaps are reported as `(earlier_idx, later_idx)` pairs. Jaxonomy's runtime resolves overlaps by earlier-row-wins, so this is informational unless `strict_disjointness=True`. Parameters: | Name | Type | Description | Default | | --------------------- | ---- | ------------------------------------------------------------------------------------------------------------ | ------- | | `strict_completeness` | | if True, raise :class:BlockParameterError when any input combination is uncovered. Default False. | `False` | | `strict_disjointness` | | if True, raise :class:BlockParameterError when any two rows match the same input combination. Default False. | `False` | Returns: | Type | Description | | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------ | | | dict with keys: | | | covered_combinations (int) — number of distinct input vectors matched by at least one row. | | | total_combinations (int) — 2 \*\* n_inputs. | | | missing_patterns (list\[tuple[bool, ...]\]) — input vectors not matched by any row. | | | overlapping_pairs (list\[tuple[int, int]\]) — sorted (i, j) row-index pairs (i < j) where row i and row j both match at least one common input vector. | Notes For `n_inputs > 10` (i.e. > 1024 enumerated combinations) the check emits a :class:`UserWarning` since cost grows as `2 ** n_inputs * len(rows)`. ### `TruthTableBuilder` Fluent builder for :class:`TruthTable` rows by named-input keywords. Example — a 2-input AND gate: .. code-block:: python ``` tt = ( TruthTable.builder(n_inputs=2, default_output=0.0) .row(in1=True, in2=True, output=1.0) .row(in1=True, in2=False, output=0.0) .row(in1=False, in2="X", output=0.0) .build() ) ``` Inputs omitted from a `.row(...)` call default to the wildcard `"X"`, so partial decision tables are concise. Custom input names may be supplied via the `input_names=` constructor argument; the default names are `in1, in2, ..., inN`. #### `build()` Materialize the accumulated rows into a :class:`TruthTable`. #### `row(output, **input_assignments)` Append a row, named by input keyword. `output` is the value emitted when the row matches. Each keyword in `input_assignments` must be one of the configured input names; omitted inputs default to the wildcard `"X"`. Returns `self` for fluent chaining. ### `UniformRandomNumber` Bases: `LeafSystem` Discrete-time uniform random number generator. Emits a fresh Uniform[low, high] sample every `sample_time` seconds, using `jax.random.uniform` with a key carried in the block's discrete state. Reproducible: same `seed` and same diagram → bit-identical sequence. The sample is computed as `low + (high - low) * u` where `u ~ Uniform[0, 1)`, so gradients of downstream losses flow cleanly through `low` and `high` via the reparameterization trick. The `u` draw is wrapped in `lax.stop_gradient` so JAX never tries to differentiate the random sequence w.r.t. the key. Input ports None. Output ports (0) The most recent uniform sample. Parameters: | Name | Type | Description | Default | | ------------- | ------- | ---------------------------------------------------------------------------------------- | ---------- | | `sample_time` | `float` | Period (s) at which a fresh sample is drawn. | *required* | | `low` | `float` | Lower bound of the uniform interval (differentiable). | `0.0` | | `high` | `float` | Upper bound of the uniform interval (differentiable). | `1.0` | | `seed` | `int` | Integer seed for the PRNG key. If None, a 32-bit random seed is drawn from numpy.random. | `None` | | `shape` | | Output shape. Default () (scalar). | `()` | Notes Per-vmap-batch independence: pass `fold_in_batch_index=True` (T-122-followup-vmap-fold-in) to derive a per-replica independent PRNG stream via `jax.lax.axis_index("batch")` inside `simulate_batch(use_vmap=True)` / `simulate_distributed`. Outside any vmap context the kwarg is a no-op (the unbound-axis `NameError` is caught gracefully and the plain seed-derived key is used). The default `False` preserves bit-identical behaviour with T-122 phase 1. ### `UnitDelay` Bases: `LeafSystem` Hold and delay the input signal by one time step. This block implements a "unit delay" with the following difference equation for internal state `x`, input signal `u`, and output signal `y`: ``` x[k+1] = u[k] y[k] = x[k] ``` Or, in a hybrid context, the discrete update advances the internal state from the "pre" or "minus" value x⁻ to the "post" or "plus" value x⁺ at time `tₖ = t0 + k * dt`. According to the discrete update rules, this calculation happens using the input values computed during the update step (i.e. by computing upstream outputs before evaluating the inputs to this block). That is, the update rule can be written `x⁺(tₖ) = f(tₖ, x⁻(tₖ), u(tₖ))`. The values of `u` are not distinguished as "pre" or "post" because there is only one value at the update time. In the difference equation notation, x⁺(tₖ) ≡ x[k+1]`,`x⁻(tₖ) ≡ x[k], and u(tₖ) ≡ u[k]. The hybrid update rule is then: ``` x⁺(tₖ) = u(tₖ) y(t) = x⁻(tₖ), between tₖ⁺ and (tₖ+dt)⁻ ``` The output signal "seen" by all other blocks on the time interval (tₖ, tₖ+dt) is then the value of the input signal u(tₖ) at the previous update. Therefore, all downstream discrete-time blocks updating at the same time tₖ will still see the value of x⁻(tₖ), the value of the internal state prior to the update. Input ports (0) The input signal. Output ports (0) The input signal delayed by one time step Parameters: | Name | Type | Description | Default | | --------------- | ---- | ----------------------------------------------- | ---------- | | `dt` | | The time step of the discrete update. | *required* | | `initial_state` | | The initial state of the block. Default is 0.0. | *required* | Note For a *multi-step* / fixed transport latency, do not chain N `UnitDelay` blocks — use a single :class:`TransportDelay` (`delay_seconds = N * dt`), which buffers the history in one block and is differentiable through the signal. `UnitDelay` is the exact one-sample `z⁻¹` primitive; :class:`TransportDelay` is the parameterized N-sample delay line. ### `UnscentedKalmanFilter` Bases: `KalmanFilterBase` Unscented Kalman Filter (UKF) for the following system: ``` ``` x[n+1] = f(x[n], u[n]) + G(t[n]) w[n] y[n] = g(x[n], u[n]) + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q(t[n], x[n], u[n]) E(v[n]v'[n] = R(t[n]) E(w[n]v'[n] = N(t[n]) = 0 ``` ``` `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`, while R`and`G`are discrete-time functions of time`t[n]`.`Q`is a discrete-time function of`t[n], x[n], u[n]\`. This last aspect is included for zero-order-hold discretization of a continuous-time system Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ------------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `forward` | | Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations. | *required* | | `observation` | | Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations. | *required* | | `G_func` | | Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations. | *required* | | `Q_func` | | Callable A function with signature Q(t[n], x[n], u[n]) -> Q[n] that represents Q in the above equations. | *required* | | `R_func` | | Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations. | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate | *required* | | `alpha` | | float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04\<= alpha \<= 1.0). Default is 1.0. | `1.0` | | `beta` | | float Scaling constant to include prior information about the distribution of the state. Default is 0.0. | `0.0` | | `kappa` | | float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0. | `0.0` | #### `for_continuous_plant(plant, dt, G_func, Q_func, R_func, x_hat_0, P_hat_0, discretization_method='euler', discretized_noise=False, alpha=1.0, beta=0.0, kappa=0.0, name=None, ui_id=None)` Unscented Kalman Filter system for a continuous-time plant. The input plant contains the deterministic forms of the forward and observation operators: ``` dx/dt = f(x,u) y = g(x,u) ``` Note: (i) Only plants with one vector-valued input and one vector-valued output are currently supported. Furthermore, the plant LeafSystem/Diagram should have only one vector-valued integrator; (ii) the user may pass a plant with disturbances (not recommended) as the input plant. In this case, the forward and observation evaluations will be corrupted by noise. A plant with disturbances of the following form is then considered: ``` dx/dt = f(x,u) + G(t) w -- (C1) y = g(x,u) + v -- (C2) ``` where: ``` `w` represents the process noise, `v` represents the measurement noise, ``` and ``` E(w) = E(v) = 0 E(ww') = Q(t) E(vv') = R(t) E(wv') = N(t) = 0 ``` This plant is discretized to obtain the following form: ``` x[n+1] = fd(x[n], u[n]) + Gd w[n] -- (D1) y[n] = gd(x[n], u[n]) + v[n] -- (D2) E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Qd E(v[n]v'[n] = Rd E(w[n]v'[n] = Nd = 0 ``` The above discretization is performed either via the `euler` or the `zoh` method, and an Unscented Kalman Filter estimator for the system of equations (D1) and (D2) is returned. Note: If `discretized_noise` is True, then it is assumed that the user is directly providing Gd, Qd and Rd. If False, then Qd and Rd are computed from continuous-time Q, R, and G, and Gd is set to an Identity matrix. The returned system will have: Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ----------------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `plant` | | a Plant object which can be a LeafSystem or a Diagram. | *required* | | `dt` | | float Time step for the discretization. | *required* | | `G_func` | | Callable A function with signature G(t) -> G that represents G in the continuous-time equations (C1) and (C2). | *required* | | `Q_func` | | Callable A function with signature Q(t) -> Q that represents Q in the continuous-time equations (C1) and (C2). | *required* | | `R_func` | | Callable A function with signature R(t) -> R that represents R in the continuous-time equations (C1) and (C2). | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate. If None, an Identity matrix is assumed. | *required* | | `discretization_method` | | str ("euler" or "zoh") Method to discretize the continuous-time plant. Default is "euler". | `'euler'` | | `discretized_noise` | | bool Whether the user is directly providing Gd, Qd and Rd. Default is False. If True, G_func, Q_func, and R_func provide Gd(t), Qd(t), and Rd(t), respectively. | `False` | | `alpha` | | float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04\<= alpha \<= 1.0). Default is 1.0. | `1.0` | | `beta` | | float Scaling constant to include prior information about the distribution of the state. Default is 0.0. | `0.0` | | `kappa` | | float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0. | `0.0` | #### `from_operators(dt, forward, observation, G_func, Q_func, R_func, x_hat_0, P_hat_0, alpha=1.0, beta=0.0, kappa=0.0, name=None, ui_id=None)` Unscented Kalman Filter (UKF) for the following system: ``` x[n+1] = f(x[n], u[n]) + G(t[n]) w[n] y[n] = g(x[n], u[n]) + v[n] E(w[n]) = E(v[n]) = 0 E(w[n]w'[n]) = Q(t[n], x[n], u[n]) E(v[n]v'[n] = R(t[n]) E(w[n]v'[n] = N(t[n]) = 0 ``` `f` and `g` are discrete-time functions of state `x[n]` and control `u[n]`, while `Q` and `R` and `G` are discrete-time functions of time `t[n]`. Input ports (0) u[n] : control vector at timestep n (1) y[n] : measurement vector at timestep n Output ports (1) x_hat[n] : state vector estimate at timestep n Parameters: | Name | Type | Description | Default | | ------------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | float Time step of the discrete-time system | *required* | | `forward` | | Callable A function with signature f(x[n], u[n]) -> x[n+1] that represents f in the above equations. | *required* | | `observation` | | Callable A function with signature g(x[n], u[n]) -> y[n] that represents g in the above equations. | *required* | | `G_func` | | Callable A function with signature G(t[n]) -> G[n] that represents G in the above equations. | *required* | | `Q_func` | | Callable A function with signature Q(t[n]) -> Q[n] that represents Q in the above equations. | *required* | | `R_func` | | Callable A function with signature R(t[n]) -> R[n] that represents R in the above equations. | *required* | | `x_hat_0` | | ndarray Initial state estimate | *required* | | `P_hat_0` | | ndarray Initial state covariance matrix estimate | *required* | | `alpha` | | float Sigma point spread to control the amount of nonlinearities taken into account. Usually set to a value (1e-04\<= alpha \<= 1.0). Default is 1.0. | `1.0` | | `beta` | | float Scaling constant to include prior information about the distribution of the state. Default is 0.0. | `0.0` | | `kappa` | | float Relatively non-critical parameter to control the kurtosis of sigma point distribution. Default is 0.0. | `0.0` | ### `VariableTransportDelay` Bases: `LeafSystem` Continuous-time variable transport (pure) delay. Implements `y(t) = u(t - tau(t))` where the delay `tau` is supplied as a runtime input signal (second input port) rather than as a static parameter. This is the T-107-followup-variable-tau extension to the fixed-delay :class:`TransportDelay` block (T-107 phase 1). Mechanism: identical to :class:`TransportDelay` — a periodic clock at period `dt` writes the most recent `history_length` `(time, u)` pairs into a discrete-state ring buffer, and the (continuous-time) output port performs a linear interpolation over the buffer at `t - clip(tau, 0, max_delay_seconds)`. The clip guards the interpolation against transient out-of-range delay values from upstream blocks; out-of-band `tau` is clamped (not raised) so that the block remains differentiable everywhere. Differentiability: - w.r.t. the data input `u`: via `npa.interp` over `values`, same as :class:`TransportDelay`. - w.r.t. the delay input `tau`: under `method="linear"` (default, phase 3) via `npa.interp`'s gradient w.r.t. its query coordinate — the standard linear-interp Jacobian, well defined except at sample boundaries where the gradient has a jump discontinuity. Pass `method="pchip"` (T-107 phase 4) to route through the T-106 backend's monotone cubic Hermite interpolant instead: smooth (C^1) gradient w.r.t. tau across every sample boundary, at the cost of one extra slope-array compute per output evaluation. The buffer is sized statically from `max_delay_seconds`: at sample period `dt` you need at least `ceil(max_delay_seconds / dt) + 1` slots; `history_length` defaults to `max(8, ceil(max_delay_seconds / dt) + 4)`. Input ports (0) The input signal `u(t)`. Scalar or array. (1) The delay signal `tau(t)` in seconds. Runtime scalar in `[0, max_delay_seconds]`; values outside that range are silently clamped. Output ports (0) The delayed signal `y(t) = u(t - tau(t))` (or `initial_output` while `t < tau(t)`). Parameters: | Name | Type | Description | Default | | ------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `dt` | | Sampling period for the history buffer. Smaller dt ⇒ finer interpolation but a larger ring buffer to cover the same physical delay. | *required* | | `max_delay_seconds` | | Upper bound on the runtime delay value. Used to size the ring buffer and to clip out-of-range tau inputs. Static (compile-time). | *required* | | `initial_output` | | Output value while t < tau(t). Default is 0.0. | `0.0` | | `history_length` | | Number of (time, value) pairs stored. Static (compile-time) — required for vmap/JIT-safe buffer sizing. If None, defaults to max(8, ceil(max_delay_seconds / dt) + 4). | `None` | Notes - Default-off / non-touched-block path is byte-equivalent: the existing :class:`TransportDelay` is untouched. - Buffer overflow (`tau > max_delay_seconds`) is clamped to `max_delay_seconds` rather than raised; this keeps the block differentiable but means the user is responsible for choosing a sufficiently large `max_delay_seconds`. - The variable-tau interpolation runs once per output evaluation (continuous-time semantics). For workloads where the delay changes only at major-step granularity, sampling `tau` at the periodic update would be cheaper — deferred until profiling demands it. - For arbitrary array-shaped data signals the interpolation is applied elementwise via a static loop over the trailing axes (mirrors :class:`TransportDelay`). ### `VideoSink` Bases: `LeafSystem` Records RGB frames to a video file. Parameters: | Name | Type | Description | Default | | ----------- | ------- | ---------------------------------------------- | ---------- | | `dt` | `float` | Interval at which to record frames. | *required* | | `file_name` | `str` | Name of the video file to write to (optional). | *required* | ### `VideoSource` Bases: `LeafSystem` Reads frames from a video file. Parameters: | Name | Type | Description | Default | | ----------- | ----- | ---------------------------------------------------------------------- | ---------- | | `file_name` | `str` | Name of the video file to read from. | *required* | | `no_repeat` | | Whether to stop at the end of the video or loop back to the beginning. | `False` | ### `WhenDisabled` Allowed string values for the `when_disabled` kwarg. ### `WhiteNoise` Bases: `LeafSystem` Continuous-time white noise generator. Generates a band-limited white noise signal using a sinc-interpolated random number generator. The output signal is a continuous-time signal, but the underlying random number generator is discrete-time. As a result, the signal is not truly white, but is band-limited by the sample rate. The resulting signal has the following approximate power spectral density: ``` S(f) = A * fs if |f| < fs else 0, ``` where `A` is the noise power and `fs = 1/dt` is the sample rate. See Ch. 10.4 in Baraniuk, "Signal Processing and Modeling" for details: https://shorturl.at/floRZ The output signal will have variance `A`, zero mean, and will decorrelate at the sample rate. Input ports None Output ports (0) The band-limited white noise signal with variance `noise_power`, zero mean, and correlation time `dt`. Parameters: | Name | Type | Description | Default | | ------------------ | ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `correlation_time` | | The correlation time of the output signal and the inverse of the bandwidth. It is the sample frequency of the underlying random number generator. | *required* | | `noise_power` | `float` | The variance of the white noise signal. Also scales the amplitude of the power spectral density. | `1.0` | | `num_samples` | `int` | The number of samples to use for sinc interpolation. More samples will result in a more accurate approximation of the ideal power spectrum, but will also increase the computational cost. The default of 10 is sufficient for most applications. | `10` | | `seed` | `int` | An integer seed for the random number generator. If None, a random 32-bit seed will be generated. | `None` | | `dtype` | `DTypeLike` | data type of the random number. If None, defaults to float. | `None` | | `shape` | `ShapeLike` | The shape of the output signal. If empty, the output will be a scalar. | `()` | #### `with_key(key, **kwargs)` Construct WhiteNoise with an explicit JAX PRNGKey. Use this when you need independent noise streams in batched (jax.vmap) simulations. Example keys = jax.random.split(jax.random.PRNGKey(0), 16) ##### Each diagram gets a different key diagrams = \[ build_diagram_with( WhiteNoise.with_key(keys[i], ...) ) for i in range(16) \] ##### OR with with_parameters (preferred): diagram.with_parameters({"noise.key": keys[i]}) Parameters: | Name | Type | Description | Default | | ---------- | ------------- | ---------------------------------------------- | ---------- | | `key` | `'jax.Array'` | JAX PRNGKey array (shape (2,) for default RNG) | *required* | | `**kwargs` | | other constructor arguments | `{}` | ### `ZeroOrderHold` Bases: `LeafSystem` Implements a "zero-order hold" A/D conversion. https://en.wikipedia.org/wiki/Zero-order_hold The block implements a "zero-order hold" with the following difference equation for input signal `u` and output signal `y`: ``` y[k] = u[k] ``` The block does not maintain an internal state, but simply holds the value of the input signal at the previous update time. As a result, the block is "feedthrough" from its inputs to outputs and cannot be used to break an algebraic loop. The data type of this hold value is inferred from upstream blocks. Input ports (0) The input signal. Output ports (0) The "hold" value of the input signal. If the input signal is continuous, then the output will be the value of the input signal at the previous update time. If the input signal is discrete and synchonous with the block, the output will be the value of the input signal at the current time (i.e. identical to the input signal). Parameters: | Name | Type | Description | Default | | ---- | ---- | ------------------------------------- | ---------- | | `dt` | | The time step of the discrete update. | *required* | ### `ForEach(submodel, n, n_inputs=1, in_axes=None, name=None)` Container block: evaluate a submodel `n` times in parallel. `ForEach` is a block-diagram-vocabulary alias for the existing :class:`jaxonomy.library.ReplicatedFunction` (T-010). It exists so that users familiar with the `ForEach` block name can find it without paying a duplication tax: the implementation is exactly :class:`ReplicatedFunction` under the hood. Parameters: | Name | Type | Description | Default | | ---------- | ---------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | | `submodel` | `Callable` | Callable f(\*inputs) -> output. Must be JAX-traceable. | *required* | | `n` | `int` | Number of replicas (the iteration count). | *required* | | `n_inputs` | `int` | Number of input ports the block declares. | `1` | | `in_axes` | | As in :func:jax.vmap / ReplicatedFunction: a length-n_inputs tuple of 0 (input is batched along the leading axis) or None (input is broadcast). Default is all-batched. | `None` | | `name` | \`str | None\` | Optional block name. | Returns: | Type | Description | | ---- | ------------------------------------------------------------ | | | A configured :class:ReplicatedFunction instance, ready to be | | | wired into a :class:DiagramBuilder. | ### `RateTransition(input_dt, output_dt, initial_state=0.0, *, name=None, dtype=None, **kwargs)` Auto-pick the right rate-bridging block based on `input_dt` vs `output_dt`. - `input_dt > output_dt` (slow source → fast destination): :class:`ZeroOrderHold` at `output_dt` (the fast rate). The held value is whatever the upstream slow block last produced; the ZOH re-samples on every fast tick. - `input_dt < output_dt` (fast source → slow destination): :class:`Decimator` at `output_dt` (the slow rate). - `input_dt == output_dt` (same rate): :class:`UnitDelay` at `input_dt` — a one-step delay so adjacent same-rate blocks can still break feedthrough loops. Both ZOH and Decimator paths are tagged with the `_jaxonomy_rate_transition` marker so :func:`jaxonomy.simulation.rate_groups.detect_rate_mismatches` silences the rate-mismatch warning across the connection. The same-rate (`UnitDelay`) path does not need the marker because it cannot itself be a rate mismatch. Parameters: | Name | Type | Description | Default | | --------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `input_dt` | | Sample period of the upstream block. | *required* | | `output_dt` | | Sample period of the downstream block. | *required* | | `initial_state` | | Initial output value (only meaningful for the same-rate UnitDelay and the fast→slow Decimator branches; ZeroOrderHold ignores it in Phase 1). | `0.0` | | `name` | | Optional block name. | `None` | | `dtype` | | Optional per-block dtype (forwarded to the underlying block). See T-038a-followup-other-blocks. | `None` | | `**kwargs` | | Forwarded to the underlying block constructor. | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ------------------------------------------- | | `A` | | class:LeafSystem instance: ZeroOrderHold, | | | | class:Decimator, or :class:\`UnitDelay\`\`. | ### `balanced_realization(sys)` Internally-balanced realization of `sys`. Returns `(balanced_system, hsv)` where `balanced_system` is an equivalent :class:`LinearizedSystem` whose controllability and observability Gramians are equal and diagonal, with the Hankel singular values `hsv` on the diagonal (Moore 1981; square-root algorithm of Laub, Heath, Paige & Ward 1987). Requires a stable `sys`. A non-minimal or stiff system (Gramians only numerically semidefinite) is handled — see :func:`_psd_sqrt` — but its ~zero Hankel-value states are ill-defined in the *full* balanced form; use :func:`balanced_truncation` or :func:`minimal_realization` to remove them. ### `balanced_truncation(sys, order=None, tol=None)` Balanced truncation (Moore 1981). Balances `sys` and keeps the states associated with the largest Hankel singular values. Order selection: - `order` given — keep exactly that many states. - `tol` given (and `order` is `None`) — keep the fewest states whose retained "energy" `Σσ_kept² / Σσ²` is at least `1 - tol`; i.e. `tol` is the fraction of Gramian energy allowed to be discarded. - neither given — no truncation (returns the balanced realization). The returned :class:`LinearizedSystem` additionally exposes: - `.hsv` — the full Hankel-singular-value spectrum, - `.reduced_order` — the retained state count `r`, - `.error_bound` — the a priori :math:`H_\infty` error bound :math:`\lVert G - G_r\rVert_\infty \le 2\sum_{i>r}\sigma_i` (Glover 1984 / Enns 1984). ### `bode_data(linsys, omegas)` Return matplotlib-ready Bode arrays for `linsys`. Handles MIMO systems out of the box: when the underlying :func:`frequency_response` returns shape `(K, p, m)` with `p > 1` or `m > 1`, the returned `magnitude_db` / `phase_deg` arrays keep the same `(K, p, m)` shape — one Bode pair per `(output, input)` channel pair. The phase is unwrapped along the frequency axis (`axis=0`) independently for each channel, which is the standard convention for MIMO Bode plots. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | ------------------------------------------------ | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem (p outputs, m inputs). | *required* | | `omegas` | | 1-D array-like of angular frequencies ω (rad/s). | *required* | Returns: | Type | Description | | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | | | Dictionary with keys: "omega" — angular frequencies (rad/s), shape (K,); "freq_hz" — ω / (2π) for log-Hz plotting, shape (K,); "magnitude_db" — 20 log₁₀ | ### `collect_snapshots(results, signals=None)` Assemble a snapshot matrix from a jaxonomy `SimulationResults`. Selected recorded signals (`results.outputs`) are stacked column-wise into `X` of shape `(n_features, n_samples)` where `n_features` is the total width of the selected signals and `n_samples == len(results.time)`. Parameters: | Name | Type | Description | Default | | --------- | ------------------------- | --------------------------------------------------------------------------------------------------------------------------- | ---------- | | `results` | | A SimulationResults with .time and an .outputs dict mapping signal name -> array of shape (n_samples,) or (n_samples, dim). | *required* | | `signals` | `Optional[Sequence[str]]` | Names to include (in order). None selects every output. | `None` | Returns: | Name | Type | Description | | ---- | -------------- | --------------------------------------------- | | `A` | `SnapshotData` | class:SnapshotData with X and time populated. | ### `controllability_gramian(A, B, dt=None)` Controllability Gramian :math:`W_c`. Continuous time (`dt is None`) solves the Lyapunov equation .. math:: A W_c + W_c A^\\mathsf{T} = -B B^\\mathsf{T} Discrete time (`dt` given) solves the Stein equation .. math:: A W_c A^\\mathsf{T} - W_c + B B^\\mathsf{T} = 0 Both require `A` stable (continuous: `Re(eig) < 0`; discrete: `|eig| < 1`) for a positive-semidefinite solution. ### `deim(nonlinear_snapshots, rank=None, energy=None)` Greedy DEIM point selection (Chaturantabut & Sorensen 2010). Takes an SVD basis `U` of the nonlinear-term snapshots and greedily selects `m` interpolation indices, then forms the oblique DEIM projector `U (Pᵀ U)⁻¹` (`P` selects the chosen rows). Parameters: | Name | Type | Description | Default | | --------------------- | ----------------- | --------------------------------------------------------------- | ---------- | | `nonlinear_snapshots` | | Snapshots of the nonlinear term, shape (n_features, n_samples). | *required* | | `rank` | `Optional[int]` | Number of DEIM modes/points m to keep. | `None` | | `energy` | `Optional[float]` | Cumulative-energy threshold used when rank is None. | `None` | Returns: | Type | Description | | --------- | --------------------------------------------------------- | | `ndarray` | (indices, projector) — indices are m distinct row indices | | `ndarray` | (np.ndarray of int), projector has shape (n_features, m). | ### `deim_galerkin_reduce(linear_rhs_fn, nonlinear_fn, basis, deim_result, x_ref=None, input_size=0, name=None)` Build a DEIM hyper-reduced POD-Galerkin ROM. The full-order dynamics are split as `ẋ = f_lin(t, x, u) + g(x)` with a (affine-)linear part `f_lin` and an elementwise nonlinearity `g`. The linear operator is reduced offline to dense `r×r` / `r×m` operators, and the nonlinearity is approximated by DEIM so it is evaluated only at the selected points (Chaturantabut & Sorensen 2010). Parameters: | Name | Type | Description | Default | | --------------- | ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `linear_rhs_fn` | `Callable` | Affine-linear part. Called linear_rhs_fn(t, x, u) if input_size > 0 else linear_rhs_fn(t, x); jax-traceable, returns (n_features,). Probed offline at t=0 to extract its reduced operators, so it must be affine in (x, u). | *required* | | `nonlinear_fn` | `Callable` | Elementwise nonlinearity g; called with a (m,) vector of states at the DEIM points and returns (m,). | *required* | | `basis` | | POD trial basis Φ, shape (n_features, r). | *required* | | `deim_result` | `Tuple[ndarray, ndarray]` | The (indices, projector) pair from :func:deim. | *required* | | `x_ref` | | Reference/offset state (default zeros). | `None` | | `input_size` | `int` | Width of the single input port; 0 for autonomous. | `0` | | `name` | `Optional[str]` | Optional block name. | `None` | Returns: | Type | Description | | ------------------ | ------------------------------------------------------------ | | `_DEIMGalerkinROM` | A jaxonomy LeafSystem with r reduced continuous states whose | | `_DEIMGalerkinROM` | per-step cost is independent of the full dimension n. | ### `discretize(linsys, dt, *, method='zoh', base_context=None, input_port=None, output_port=None)` Discretize a continuous-time linear system (T-109 phase 4). Two call patterns, dispatched on the type of `linsys`: 1. `discretize(linsys: LinearizedSystem, dt, *, method)` — the LTI-level path (shipped first as the T-109 phase-4 sub-piece). Wraps the matrix-level helpers in :mod:`jaxonomy.library.state_estimators.utils`. 1. `discretize(system: SystemBase, dt, *, method, base_context, input_port, output_port)` — the **diagram-level lift** (T-109 phase 4 completion). Linearizes `system` about `base_context` (via :func:`linearize`) then routes the result through path 1. Equivalent to `discretize(linearize(system, base_context, ...), dt, method=method)`; provided so controller-design workflows can write `ddiagram = jaxonomy.discretize(diagram, dt, base_context=ctx)` in one call. Converts `dx/dt = Ax + Bu` into `x[k+1] = A_d x[k] + B_d u[k]` while keeping `C`, `D`, and the operating point untouched (the output map is unaffected by discretization). The returned :class:`LinearizedSystem` carries `dt` so downstream consumers (e.g. :meth:`LinearizedSystem.is_stable`) interpret it as discrete-time. Parameters: | Name | Type | Description | Default | | -------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | | Either a continuous-time :class:LinearizedSystem (dt must be None) or a :class:SystemBase / :class:Diagram. In the diagram case, base_context is required. | *required* | | `dt` | `float` | Sampling period in seconds. Must be positive. | *required* | | `method` | `str` | Discretization rule. "zoh" (default) — exact zero-order-hold, read off the augmented matrix exponential expm(\[[A, B], [0, 0]\]·dt) = \[[A_d, B_d], [0, I]\]. Exact for singular A too, so integrator dynamics need no special-casing. "euler" — first-order forward-Euler: A_d = I + A·dt, B_d = B·dt. Cheap and JAX-clean but biased; use "zoh" unless you specifically need the Euler shape for hardware-in-the-loop parity. | `'zoh'` | | `base_context` | | Required when linsys is a SystemBase / Diagram; ignored when it's already a LinearizedSystem. The operating point about which to linearize. | `None` | | `input_port` | | Optional input port for :func:linearize (diagram path only). Defaults to the diagram's single input. | `None` | | `output_port` | | Optional output port for :func:linearize (diagram path only). Defaults to the diagram's single output. | `None` | Returns: | Type | Description | | ------------------ | -------------------------------------------------------- | | `LinearizedSystem` | A new :class:LinearizedSystem with discrete matrices and | | `LinearizedSystem` | dt set. The output map (C, D) and | | `LinearizedSystem` | operating_point are forwarded unchanged. | Raises: | Type | Description | | ------------ | ----------------------------------------------------------------------------------------------------------------------------- | | `ValueError` | If dt \<= 0, method is not "zoh" or "euler", linsys is already discrete, or the diagram path is invoked without base_context. | Notes Differentiable through `A`, `B`, `C`, `D`, and `dt` via the JAX-traceable matrix exponential. The diagram path is differentiable through whatever :func:`linearize` is itself differentiable through. See also :func:`linearize` — the continuous-time linearization step. :func:`jaxonomy.library.state_estimators.utils.discretize_forward_zoh` and :func:`discretize_forward_euler` — the matrix-level primitives the LTI path wraps. ### `dmdc(X, Xp, U, rank=None, B_known=None)` Dynamic Mode Decomposition with control (Proctor, Brunton & Kutz 2016). Fits `x[k+1] ≈ A x[k] + B u[k]` from snapshot pairs and control inputs. Two cases are handled: - **Unknown `B`** (default): regress on the augmented snapshot `Ω = [X; U]` so `[A B] = Xp Ω⁺`. - **Known `B`** (pass `B_known`): subtract the known control effect first, `A = (Xp − B U) X⁺`. Parameters: | Name | Type | Description | Default | | --------- | ---- | -------------------------------------------------------------- | ---------- | | `X` | | State snapshots x[k], shape (n, k). | *required* | | `Xp` | | Advanced snapshots x[k+1], shape (n, k). | *required* | | `U` | | Control inputs u[k], shape (m, k). | *required* | | `rank` | | Optional POD rank r for the reduced operators (defaults full). | `None` | | `B_known` | | Optional known input matrix (n, m) for the known-B case. | `None` | Returns: | Type | Description | | ---- | ------------------------------------------------------------- | | | class:DMDcResult with full A, B and reduced A_tilde, B_tilde. | ### `edmd(X, Xp, dictionary, U=None)` Extended DMD — approximate the Koopman operator on lifted snapshots. Lifts the snapshot pair through `dictionary` and least-squares fits the lifted linear dynamics `z[k+1] ≈ K z[k] (+ B u[k])`. Parameters: | Name | Type | Description | Default | | ------------ | ---- | ------------------------------------------------------------ | ---------- | | `X` | | State snapshots x[k], shape (n, k). | *required* | | `Xp` | | Advanced snapshots x[k+1], shape (n, k). | *required* | | `dictionary` | | Callable g(x) -> lifted vector (identity observables first). | *required* | | `U` | | Optional control inputs (m, k) for eDMDc. | `None` | Returns: | Type | Description | | ---- | ---------------------------------------------------------------- | | | class:EDMDResult with the Koopman operator K, the input operator | | | B (or None), and the de-lift matrix C. | ### `era(markov, n_inputs, n_outputs, num_rows=None, num_cols=None, rank=None)` Eigensystem Realization Algorithm (Juang & Pappa 1985). Builds a minimal discrete-time state-space realization `(A, B, C, D)` from a sequence of impulse-response Markov parameters `Y_0 = D`, `Y_1 = C B`, `Y_2 = C A B` ... Parameters: | Name | Type | Description | Default | | ----------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `markov` | | Markov parameters. Either an array of shape (L+1, n_outputs, n_inputs) or a length-L+1 sequence of such blocks; SISO impulse responses may be passed as a 1-D array. | *required* | | `n_inputs` | | Number of inputs m. | *required* | | `n_outputs` | | Number of outputs p. | *required* | | `num_rows` | | Block rows α of the Hankel matrix (default ~half the data). | `None` | | `num_cols` | | Block cols β of the Hankel matrix (default ~half the data). | `None` | | `rank` | | Optional model order r (SVD truncation of the Hankel matrix). | `None` | Returns: | Type | Description | | ---- | --------------------------------------------------------- | | | class:ERAResult with the realized (A, B, C, D) and Hankel | | | singular values. | ### `estimate_frequency_response(diagram, ctx, t_span, input_port, output_port, freq_grid, *, options=None, recorded_signals_extra=None, window=True, coherence_floor=1e-12, n_segments=8, segment_overlap=0.5)` Empirically estimate the SISO transfer function of `diagram`. Drives `diagram` with whatever signal is already wired to `input_port` (typically a :class:`jaxonomy.library.Chirp`, :class:`PRBS`, or :class:`BandLimitedNoise` source connected upstream of `input_port`) and records the input/output trajectories. The empirical transfer function is computed as the cross-spectral ratio `G(f) = Sxy(f) / Sxx(f)` where `Sxx` and `Sxy` are the (Hann- windowed) auto- and cross-spectral densities of input/output. Results are interpolated onto the user-supplied `freq_grid` (Hz). This is the practical alternative to analytic :func:`linearize` / :func:`frequency_response` when: - the system contains hard nonlinearities (lookup tables, saturation, contact dynamics) that make symbolic linearization fragile, or - you want an empirical sanity-check against the linearized model. Parameters: | Name | Type | Description | Default | | ------------------------ | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `diagram` | | A built diagram (typically with a chirp or PRBS source wired to the block-under-test's input port). | *required* | | `ctx` | | Initial simulation context. | *required* | | `t_span` | | (t0, tf) simulation horizon. Make this comfortably longer than the slowest period of interest in freq_grid. | *required* | | `input_port` | | OutputPort whose recorded trajectory provides the excitation samples u(t). This is typically the upstream source's output port (the same signal that drives the block-under-test). | *required* | | `output_port` | | OutputPort whose recorded trajectory provides the measured response y(t). | *required* | | `freq_grid` | | 1-D array of frequencies (Hz) at which the empirical response should be evaluated. Frequencies outside the simulation's resolved band [1/T, fs/2] are clamped — the caller should keep freq_grid inside that band. | *required* | | `options` | | Optional :class:SimulatorOptions. recorded_signals is overridden internally; everything else (rtol, atol, solver, etc.) is honoured. | `None` | | `recorded_signals_extra` | | Optional dict[str, OutputPort] of additional signals to record alongside the input/output (useful for debugging / plotting). Not used by the estimator itself. | `None` | | `window` | `bool` | If True (default) apply a Hann window before the FFT to suppress spectral leakage. If False (rectangular) the transfer-function ratio is more sensitive to leakage but faithful to the raw FFT. | `True` | | `coherence_floor` | `float` | Minimum | U(f) | | `n_segments` | `int` | Number of overlapping segments to average (Welch's method). More segments → less variance, lower frequency resolution. n_segments=1 falls back to a single-window FFT estimate. Default 8 is a reasonable trade-off for most chirp/PRBS excitations. | `8` | | `segment_overlap` | `float` | Fractional overlap between consecutive segments (Welch's method), in \[0, 1). Default 0.5 (50%). | `0.5` | Returns: | Type | Description | | ------------------- | ----------------------------------------------------------- | | `FrequencyResponse` | class:FrequencyResponse with omegas = 2π·freq_grid, complex | | `FrequencyResponse` | response of shape (K, 1, 1), and corresponding | | `FrequencyResponse` | magnitudes and phases. Drop-in compatible with | | `FrequencyResponse` | func:bode_data. | Notes - The implementation is intentionally pure-NumPy on the post-simulation arrays; it does not need to be JAX-traceable (callers can JIT downstream code that consumes the returned `response` array). - For best results pick an excitation that covers the band of interest densely: a linear :class:`Chirp` from `f0 ≪ freq_min` to `f1 ≳ freq_max` over a horizon of several seconds, or a :class:`PRBS` with sample time `≪ 1/(2·freq_max)`. - The returned `response` is a NumPy complex array (consumers calling :func:`bode_data` will see `jnp.asarray` promotion); this is fine because :class:`FrequencyResponse` fields are typed `Any`. ### `findop(system, base_context, *, initial_guess=None, input_port=None, tol=1e-08, max_iter=50, damping=1e-10, axis_mask=None, residual_fn=None, residual_scaling=None, scaling_eps=1e-08)` Find a continuous-state operating point `x*` such that `ẋ(x*, u₀) ≈ 0`. Performs damped Newton iteration on the residual `r(x) = ẋ(x, u₀)` where `u₀` is read from `base_context` (and held fixed for the duration of the search). The Jacobian is computed with `jax.jacrev` and the linear update is solved with `jnp.linalg.solve` plus a small Levenberg regularisation so singular Jacobians degrade to a least-squares step instead of NaN. Parameters: | Name | Type | Description | Default | | ------------------ | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------- | | `system` | | The system whose equilibrium is sought. | *required* | | `base_context` | | A context that supplies the initial state, parameter values, and (via input_port.eval) the held-fixed input. | *required* | | `initial_guess` | | Optional initial state. Defaults to base_context.continuous_state. | `None` | | `input_port` | | Input port to read u₀ from. Defaults to system.input_ports[0] (errors if system has multiple inputs and none is specified). | `None` | | `tol` | `float` | Stop when max( | scaled residual | | `max_iter` | `int` | Hard cap on Newton iterations. | `50` | | `damping` | `float` | Tikhonov damping added to JᵀJ for ill-conditioned solves. | `1e-10` | | `axis_mask` | | Optional selector for which state components the Newton iteration drives to zero. Either a boolean array (length = number of flat state components, True = solve this component) or a sequence of integer indices. Components not selected are held at initial_guess and excluded from both the residual and the unknowns — use this to trim systems with passive states whose equilibrium derivative is intrinsically nonzero (a cornering vehicle's heading ψ̇ = r ≠ 0, a free integrator), which would otherwise dominate the full-state Newton step and prevent convergence. None (default) solves the full state. | `None` | | `residual_fn` | | Optional (x) -> residual_vector overriding the default ẋ(x, u₀) — e.g. to add a custom equilibrium condition or drop terms. Receives the full state x; its output is masked / scaled like the default residual. | `None` | | `residual_scaling` | | Optional per-component residual weighting to put disparate units on a common footing (cf. MATLAB findop's XScaling / YScaling) — e.g. a chassis residual in m/s² ~10 alongside a wheel residual in rad/s² ~100. One of: None (no scaling), "auto" (component i scaled by 1/max( | rᵢ(x₀) | | `scaling_eps` | `float` | Floor for the "auto" scaling denominator. | `1e-08` | Returns: | Type | Description | | ---------------- | ------------------------------------------------------------------- | | `OperatingPoint` | class:OperatingPoint carrying the equilibrium state and convergence | | `OperatingPoint` | metadata. x always has the shape of initial_guess (held | | `OperatingPoint` | components carry their initial values when axis_mask is used). | Notes The returned `x` is a JAX array, so the residual function used here is differentiable: `jax.grad(lambda x0: jnp.sum(residual(x0)**2))` works. Composing :func:`findop` itself under `jax.grad` requires an implicit-differentiation wrapper which is deferred to a follow-up. **Robust fallback.** Newton operating-point search can stall on stiff, strongly-coupled, or badly-scaled systems even with `axis_mask` and `residual_scaling`. The most robust equilibrium finder is simply to *integrate to steady state*: `simulate` the system from a reasonable initial condition over a horizon long relative to its slowest mode and take the final state (optionally asserting `max(|ẋ|)` is small there). Use that when `findop` reports `converged=False`. ### `fit_gp(X, y, kernel='rbf', length_scale=1.0, signal_var=1.0, noise=1e-08, optimize=False, n_restarts=0, lr=0.05, n_steps=200, matern_nu=2.5)` Fit a Gaussian-process (kriging) surrogate. Parameters: | Name | Type | Description | Default | | --------------------------------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `X` | | training inputs, shape (n,) or (n, d). | *required* | | `y` | | training targets, shape (n,). | *required* | | `kernel` | | "rbf" / "squared_exponential" or "matern" / "matern32" / "matern52". | `'rbf'` | | `length_scale, signal_var, noise` | | kernel hyperparameters (initial values when optimize=True). | *required* | | `optimize` | | if True, maximize the marginal log-likelihood over (length_scale, signal_var, noise) by gradient ascent in log-space (Rasmussen & Williams 2006, Eq. 5.9). | `False` | Returns: | Name | Type | Description | | ---- | ---- | -------------- | | `A` | | class:GPModel. | ### `fit_lookup_table_1d(xp, x_data, y_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)` Fit a 1-D lookup table to data and return a `LookupTable1d` block. Parameters: | Name | Type | Description | Default | | ---------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------ | | `xp` | | Fixed grid of breakpoints (1-D, strictly increasing). | *required* | | `x_data` | | Measured input cloud, shape (K,). | *required* | | `y_data` | | Measured output cloud, shape (K,). | *required* | | `interpolation` | `str` | Interpolation rule for the runtime block ("linear" / "pchip" / "nearest" / "flat"). The fit itself is always linear-LS — see the module docstring of :mod:jaxonomy.library.lookup_table for why. | `'linear'` | | `extrapolation` | `str` | Out-of-range policy for the runtime block; see :class:jaxonomy.library.LookupTable1d. | `'clip'` | | `weights` | | Optional per-sample weights for weighted least- squares. None = OLS. | `None` | | `smoothness` | `float` | Non-negative discrete first-difference penalty. Use small values (1e-3 .. 1.0) on noisy / sparse data. | `0.0` | | `name` | \`str | None\` | Optional block name, forwarded to LookupTable1d. | | `**block_kwargs` | | Additional kwargs forwarded to the LookupTable1d constructor (e.g. dtype=). | `{}` | Returns: | Type | Description | | ---- | ------------------------------------------------ | | | A LookupTable1d instance with input_array=xp and | | | output_array set to the LS-fit table values. | ### `fit_lookup_table_2d(xp, yp, x_data, y_data, z_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)` Fit a 2-D lookup table to data and return a `LookupTable2d` block. Parameters: | Name | Type | Description | Default | | ------------------------ | ------- | -------------------------------------------------------------------------------------------------------------------- | -------------------- | | `xp` | | Fixed grid of breakpoints along the first axis (1-D, strictly increasing). | *required* | | `yp` | | Fixed grid of breakpoints along the second axis (1-D, strictly increasing). | *required* | | `x_data, y_data, z_data` | | Measurement cloud, all shape (K,). | *required* | | `interpolation` | `str` | Interpolation rule for the runtime block (currently only "linear" / bilinear). The fit itself is always bilinear-LS. | `'linear'` | | `extrapolation` | `str` | Out-of-range policy for the runtime block. | `'clip'` | | `weights` | | Optional per-sample weights for weighted least-squares. | `None` | | `smoothness` | `float` | Non-negative 5-point-Laplacian smoothness penalty. | `0.0` | | `name` | \`str | None\` | Optional block name. | | `**block_kwargs` | | Additional kwargs forwarded to the LookupTable2d constructor (e.g. dtype=). | `{}` | Returns: | Type | Description | | ---- | --------------------------------------------------- | | | A LookupTable2d instance with input_x_array=xp, | | | input_y_array=yp, and output_table_array set to the | | | LS-fit table values of shape (len(xp), len(yp)). | ### `fit_lookup_table_nd(grid_axes, x_data, y_data, *, interpolation='linear', extrapolation='clip', weights=None, smoothness=0.0, name=None, **block_kwargs)` Fit an N-D lookup table to data and return a `LookupTableND` block. The public N-D counterpart to :func:`fit_lookup_table_1d` and :func:`fit_lookup_table_2d`. Returns a fully-built block whose `output_array` is the LS-fit table. Parameters: | Name | Type | Description | Default | | ---------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------- | -------------------- | | `grid_axes` | | Tuple of N strictly-increasing 1-D breakpoint arrays. | *required* | | `x_data, y_data` | | Measurement cloud — x_data shape (K, N), y_data shape (K,). | *required* | | `interpolation` | `str` | Interpolation rule for the runtime block. Only "linear" (multilinear) is supported today; the fit itself is always multilinear-LS. | `'linear'` | | `extrapolation` | `str` | Out-of-range policy for the runtime block; see :class:jaxonomy.library.LookupTableND. | `'clip'` | | `weights` | | Optional per-sample weights for weighted least-squares. | `None` | | `smoothness` | `float` | Non-negative coefficient on the N-D Laplacian smoothness penalty. | `0.0` | | `name` | \`str | None\` | Optional block name. | | `**block_kwargs` | | Additional kwargs forwarded to the LookupTableND constructor (e.g. dtype=). | `{}` | Returns: | Name | Type | Description | | ---- | ---- | -------------------------------------------------------- | | `A` | | class:LookupTableND instance with the supplied grid_axes | | | | and output_array set to the LS-fit table of shape | | | | (len(grid_axes[0]), ..., len(grid_axes[N-1])). | ### `fit_pce(X, y, distributions, order)` Fit a polynomial-chaos expansion by least-squares regression. Parameters: | Name | Type | Description | Default | | --------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `X` | | training inputs, shape (n,) or (n, d). | *required* | | `y` | | training targets, shape (n,). | *required* | | `distributions` | `Sequence` | per-dimension germ, e.g. [("normal", mu, sigma), ("uniform", a, b)]. Hermite basis for normal, Legendre for uniform (Wiener--Askey scheme, Xiu & Karniadakis 2002). | *required* | | `order` | `int` | total-degree truncation. | *required* | Returns: | Name | Type | Description | | ---- | ---- | --------------- | | `A` | | class:PCEModel. | ### `fit_rbf(X, y, kernel='multiquadric', epsilon=1.0, smoothing=0.0, poly_degree=None)` Fit a radial-basis-function surrogate. Parameters: | Name | Type | Description | Default | | ------------- | ---- | ---------------------------------------------------------------------------------------------------------------------- | ---------------- | | `X` | | training inputs, shape (n,) or (n, d). | *required* | | `y` | | training targets, shape (n,). | *required* | | `kernel` | | "multiquadric", "inverse_multiquadric", "gaussian", or "thin_plate_spline". | `'multiquadric'` | | `epsilon` | | shape parameter (ignored by the thin-plate spline). | `1.0` | | `smoothing` | | ridge regularization added to the kernel diagonal; 0 gives exact interpolation. | `0.0` | | `poly_degree` | | if set, augment with a total-degree polynomial tail and solve the bordered saddle-point system (Wendland 2005, Ch. 8). | `None` | Returns: | Name | Type | Description | | ---- | ---- | --------------- | | `An` | | class:RBFModel. | ### `fit_table_1d_with_grid(n_grid_points, x_data, y_data, x_lo=None, x_hi=None, init_xp=None, *, smoothness=0.0, optimizer='gd', max_iter=200, learning_rate=0.001, auto_normalize=True)` Jointly optimise the grid `xp` AND the table values `yp`. This is the T-124-followup-grid-optimization deliverable. Phase 1's :func:`fit_table_1d` fits `yp` at a fixed user-supplied `xp`; here we ALSO move the breakpoints to better resolve regions where the data has strong features (sharp peaks, kinks). The math: for any candidate grid `xp`, the inner problem is still a linear least-squares solve for `yp` (closed form). The outer loop minimises the resulting data residual w.r.t. `xp`, with monotonicity enforced via a smooth `cumsum(softplus(deltas))` parametrisation rather than projection. Parameters: | Name | Type | Description | Default | | ---------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------- | | `n_grid_points` | `int` | Number of breakpoints to place (must be ≥ 2). | *required* | | `x_data` | | Measured input cloud, shape (K,). | *required* | | `y_data` | | Measured output cloud, shape (K,). | *required* | | `x_lo` | \`float | None\` | Lower endpoint of the grid. None (default) = min(x_data). The endpoint is pinned — the optimiser only moves the interior breakpoints. | | `x_hi` | \`float | None\` | Upper endpoint of the grid. None (default) = max(x_data). | | `init_xp` | | Optional initial grid (1-D, strictly increasing, spanning [x_lo, x_hi]). None (default) starts from a uniform grid. | `None` | | `smoothness` | `float` | Forwarded to the inner LS solve as a discrete first-difference penalty on yp. 0.0 (default) is pure data-residual. | `0.0` | | `optimizer` | `str` | "gd" (default) — hand-rolled fixed-step gradient descent on the unconstrained deltas. Differentiable end-to-end, jit-friendly, reliable. "lbfgs" — delegates the outer loop to :func:jax.scipy.optimize.minimize (BFGS). Faster on well-conditioned problems but does NOT support differentiation through itself (jax.grad of the joint fit w.r.t. y_data will fail with this option — use "gd" if you need the gradient). See T-124-followup-grid-optimization-lbfgs for the proper differentiable L-BFGS implementation. | `'gd'` | | `max_iter` | `int` | Outer-loop iteration budget. For optimizer="gd" each iter is one gradient step; for "lbfgs" it is the BFGS maxiter. | `200` | | `learning_rate` | `float` | Step size for optimizer="gd". Default is 1e-3 — the residual landscape in deltas-space has steep cliffs near sharp data features and aggressive step sizes overshoot. Ignored by "lbfgs". When auto_normalize=True (the default), the learning rate is applied in the normalised [-1, +1] x-space and [-1, +1] y-space rather than in the user's natural units, so a single sensible default works across orders-of-magnitude data scales. | `0.001` | | `auto_normalize` | `bool` | When True (default, T-124-followup-grid-fit-auto-normalize), the optimiser internally rescales x_data and y_data to roughly [-1, +1] (zero-mean, unit-half-range affine transform) so the learning_rate=1e-3 default works on wide-but-smooth features (e.g. an engine-map slice with rpm ∈ [80, 650] and torque ~250 N·m) without blowing up to NaN. The optimised grid and table are transformed back to the user's natural units on return — results are byte-equivalent to the pre-normalisation path on data that was already centred near unit scale. Set to False to disable (e.g. for byte-equivalent reproduction of pre-normalisation runs, or when the learning_rate was tuned in natural units). | `True` | Returns: | Type | Description | | ---- | ------------------------------------------------------ | | | (xp_opt, yp_opt) — the optimised grid (shape | | | (n_grid_points,)) and the corresponding optimal table | | | values (shape (n_grid_points,)). Both differentiable | | | through y_data (and through x_data modulo the discrete | | | bucket index in the design matrix) when | | | optimizer="gd". | Honest fallback note: this ships the gradient-descent path as the primary solver (rather than full L-BFGS) per the task spec — it's slower but more robust on the inner-outer formulation. The proper differentiable L-BFGS via implicit-function-theorem unrolling is filed as `T-124-followup-grid-optimization-lbfgs`. ### `fit_table_2d(xp, yp, x_data, y_data, z_data, weights=None, smoothness=0.0, rcond=None)` Fit a 2-D lookup table `zp` at the fixed grid `(xp, yp)` to `(x_data, y_data, z_data)`. Solves the bilinear least-squares problem ``` min_zp Σ_k w_k * (z_data[k] - bilinear_interp(x_data[k], y_data[k]; xp, yp, zp))² + smoothness * Σ_{i,j} (zp[i,j] - mean(neighbours))² ``` via :func:`jnp.linalg.lstsq` on the bilinear design matrix. Linear bilinear only — the 2-D analogue of the `fit_table_1d` linear-only restriction. See `T-124-followup-2d-pchip-fit` for non-linear extensions. Parameters: | Name | Type | Description | Default | | ------------------------ | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------ | | `xp` | | 1-D, strictly increasing grid along the first axis (Nx). | *required* | | `yp` | | 1-D, strictly increasing grid along the second axis (Ny). | *required* | | `x_data, y_data, z_data` | | Measurement cloud, all shape (K,). | *required* | | `weights` | | Optional per-sample weights, shape (K,). None means uniform weighting. | `None` | | `smoothness` | `float` | Non-negative coefficient for the 5-point Laplacian penalty. 0.0 (default) is pure data-fit; small values (1e-3 .. 1.0) regularise on noisy / sparse data. | `0.0` | | `rcond` | \`float | None\` | Forwarded to :func:jnp.linalg.lstsq. | Returns: | Type | Description | | ---- | ---------------------------------------------------- | | | zp of shape (len(xp), len(yp)) — the optimal table | | | values. Differentiable through z_data (and through | | | x_data / y_data modulo the discrete bucket indices). | ### `fit_table_nd(grid_axes, x_data, y_data, *, weights=None, smoothness=0.0, rcond=None)` Fit an N-D lookup table at fixed grid breakpoints. Solves the multilinear least-squares problem ``` min_zp Σ_k w_k * (y_data[k] - multilinear_interp(x_data[k]; grid_axes, zp))² + smoothness * Σ_cells (zp[cell] - mean(in-bounds neighbours))² ``` via :func:`jnp.linalg.lstsq` on the multilinear design matrix. Generalises :func:`fit_table_2d` (and :func:`fit_table_1d` for `N=1`) to an arbitrary number of grid axes. Parameters: | Name | Type | Description | Default | | ------------ | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------ | | `grid_axes` | | Tuple of N strictly-increasing 1-D breakpoint arrays. Axis d has length B_d. | *required* | | `x_data` | | Query points, shape (K, N). Column d is the d-th coordinate. | *required* | | `y_data` | | Sample values at the query points, shape (K,). | *required* | | `weights` | | Optional per-sample weights, shape (K,). None means uniform weighting. | `None` | | `smoothness` | `float` | Non-negative coefficient on the N-D-Laplacian penalty. 0.0 (default) is a pure data fit; small values (1e-3 .. 1.0) regularise on noisy / sparse data. The Laplacian is the canonical smoother on a regular grid — far better than diagonal Tikhonov, especially for cells without nearby measurements. | `0.0` | | `rcond` | \`float | None\` | Forwarded to :func:jnp.linalg.lstsq. | Returns: | Type | Description | | ---- | ------------------------------------------------------------ | | | zp of shape (B_1, ..., B_N) — the optimal table values. | | | Layout matches :class:LookupTableND.output_array exactly, so | | | the result can be passed straight through. | Memory note: builds a dense `(K + prod(B_i), prod(B_i))` design matrix. For `N=5` with `B_i = 10` that's `10^5` columns — fine on CPU up to a few thousand measurements. For larger tables or higher-D problems, switch to a sparse solver (filed under `T-104-followup-fit-table-nd-sparse`). ### `frequency_response(linsys, omegas)` Compute the frequency response of a linearized state-space system. For a continuous-time LTI `ẋ = Ax + Bu, y = Cx + Du` the transfer function evaluated at `s = jω` is the `(p, m)` transfer-function matrix `G(s) = C (sI − A)⁻¹ B + D`. This helper vectorises that evaluation across an `omegas` array and naturally handles MIMO systems (`m > 1` inputs and/or `p > 1` outputs) — the returned array shape is always `(K, p, m)`. SISO is the special case `p = m = 1` which produces shape `(K, 1, 1)`. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | ---------------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem (typically produced by :func:linearize). A is (n, n), B is (n, m), C is (p, n), D is (p, m). | *required* | | `omegas` | | 1-D array-like of angular frequencies ω (rad/s). | *required* | Returns: | Type | Description | | ------------------- | --------------------------------------------------------- | | `FrequencyResponse` | class:FrequencyResponse with omegas (shape (K,)), complex | | `FrequencyResponse` | response (shape (K, p, m)), and corresponding magnitudes | | `FrequencyResponse` | and phases (radians). For MIMO response[k, i, j] is the | | `FrequencyResponse` | transfer function from input j to output i evaluated at | | `FrequencyResponse` | ω = omegas[k]. | Notes The implementation is fully JAX-traceable and differentiable through `A, B, C, D` and `omegas` so it composes with `jax.grad` and `jax.vmap`. `jnp.linalg.solve` solves the matrix RHS `B` in one shot per frequency so the per-omega cost is one LU decomposition plus `m` triangular back-substitutions — substantially cheaper than looping over input channels. ### `galerkin_reduce(rhs_fn, basis, x_ref=None, output_fn=None, input_size=0, test_basis=None, name=None)` Project a full-order RHS onto a reduced basis (POD-Galerkin / LSPG). Parameters: | Name | Type | Description | Default | | ------------ | -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `rhs_fn` | `Callable` | Full-order dynamics. Called as rhs_fn(t, x_full, u) when input_size > 0 else rhs_fn(t, x_full); must be jax-traceable and return dx_full of shape (n_features,). | *required* | | `basis` | | Trial basis Φ, shape (n_features, r). | *required* | | `x_ref` | | Reference/offset state added on reconstruction (default zeros). | `None` | | `output_fn` | `Optional[Callable]` | Optional map applied to the reconstructed full state for the output port. | `None` | | `input_size` | `int` | Width of the single input port; 0 for an autonomous block (no input port). | `0` | | `test_basis` | | Optional test basis Ψ (shape (n_features, r)) for a Petrov-Galerkin/LSPG projection W = (Ψᵀ Φ)⁻¹ Ψᵀ. When None, Galerkin W = Φᵀ. | `None` | | `name` | `Optional[str]` | Optional block name. | `None` | Returns: | Type | Description | | ---------------- | ------------------------------------------------------- | | `_ProjectionROM` | A jaxonomy LeafSystem with r reduced continuous states. | ### `hankel_singular_values(sys)` Hankel singular values, sorted descending. :math:`\sigma_i = \sqrt{\lambda_i(W_c W_o)}` for the controllability and observability Gramians of `sys` (Moore 1981). ### `identity_dictionary()` Trivial dictionary `g(x) = x`. eDMD with this dictionary reduces to plain (linear) DMD — a useful baseline. ### `impulse_response(linsys, t_grid)` Impulse response of an LTI system, continuous- or discrete-time. **Continuous** (`linsys.dt is None`): for zero initial state the (finite part of the) impulse response is .. code-block:: text ``` y(t) = C · expm(A·t) · B for t > 0 ``` The Dirac component `D · δ(t)` is omitted from the returned samples since it is not representable on a numeric grid; consumers that need it can add `D` to the `t = 0` sample explicitly. **Discrete** (`linsys.dt` set): the response to the unit pulse `u[0] = 1` (which, unlike the Dirac, *is* representable): `y[0] = D`, `y[k] = C·A^{k-1}·B` for `k ≥ 1`, evaluated by the exact recurrence and sampled on the integer grid `k = t/dt`. Off-grid times raise `ValueError` (same policy as :func:`step_response`); negative times return zero. This matches the `scipy.signal.dimpulse` convention. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | ----------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem (either time base). | *required* | | `t_grid` | | Scalar or 1-D array of evaluation times. Discrete: must be integer multiples of dt and concrete (not JAX-traced). | *required* | Returns: | Type | Description | | ---- | ----------------------------------------------------- | | | Array of shape (K, p, m) for vector t_grid, or (p, m) | | | for scalar t_grid. K = len(t_grid), p = n_outputs, | | | m = n_inputs. | Notes Fully differentiable through `A, B, C, D`. ### `linearize(system, base_context, name=None, output_index=None, input_port=None, output_port=None)` Linearize the system about an operating point specified by the base context. Note: Deprecated return type. Previously returned LTISystem directly, now returns a LinearizedSystem object. Use `.to_lti()` on the result if you need an LTISystem block. ### `linearize_to_lti(system, base_context, input_port=None, output_port=None, name=None)` Linearize `system` at `base_context` and return an `LTISystem`. Parameters: | Name | Type | Description | Default | | -------------- | --------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `'SystemBase'` | A LeafSystem or Diagram to linearize. If it has multiple input or output ports, input_port and output_port must be supplied explicitly. | *required* | | `base_context` | `'ContextBase'` | The operating-point context. State, inputs, and parameters read from this context define the point about which the linearization is performed. | *required* | | `input_port` | | Input port to linearize against. Required when system has more than one input port. | `None` | | `output_port` | | Output port to linearize against. Required when system has more than one output port. | `None` | | `name` | `Optional[str]` | Optional name for the returned LTISystem block. | `None` | Returns: | Type | Description | | ------------- | --------------------------------------------------------- | | `'LTISystem'` | An LTISystem block with the derived (A, B, C, D) | | `'LTISystem'` | matrices. Drop this into a DiagramBuilder wherever the | | `'LTISystem'` | original subdiagram would go; downstream blocks should be | | `'LTISystem'` | wired to lti.output_ports[0] and upstream blocks to | | `'LTISystem'` | lti.input_ports[0]. | ### `merge_buses(bus_a, bus_b, *, on_collision='error')` Merge two NamedTuple-shaped bus signals by union of fields. The merged bus is a fresh NamedTuple (named `"MergedBus"`) whose fields are `bus_a._fields` followed by the fields of `bus_b._fields` not already in `bus_a` (de-duplicated while preserving declaration order). The result is a JAX-pytree-friendly value identical in shape to what :class:`BusCreator` would produce for the merged schema. Differentiability: gradients flow from each merged-bus leaf back to whichever input bus contributed the leaf — the underlying op is NamedTuple construction over `getattr` lookups, both of which are transparent to `jax.grad` / `jax.jit`. Parameters: | Name | Type | Description | Default | | -------------- | ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `bus_a` | | First bus signal. Must be a NamedTuple-shaped value (isinstance(bus_a, tuple) and hasattr(bus_a, "\_fields")). | *required* | | `bus_b` | | Second bus signal. Same contract as bus_a. | *required* | | `on_collision` | `str` | Policy for fields that appear in both inputs. "error" (default) — raise :class:ValueError. "prefer_a" — keep the value from bus_a. "prefer_b" — keep the value from bus_b. The merged-bus schema (field order) is independent of the policy: collisions only change which leaf value lands in the colliding slot. | `'error'` | Returns: | Type | Description | | ---- | --------------------------------------------------- | | | A NamedTuple instance whose fields are the union of | | | bus_a.\_fields and bus_b.\_fields. | Raises: | Type | Description | | ------------ | --------------------------------------------------------------------------------------------------------- | | `TypeError` | If either input is not a NamedTuple-shaped value. | | `ValueError` | If on_collision is not one of the supported policies, or if collisions exist and on_collision == "error". | ### `minimal_realization(sys, tol=1e-08)` Minimal realization of `sys` (Kalman decomposition). Removes uncontrollable and unobservable modes by projecting onto the controllable subspace (range of the controllability matrix) and then onto the observable subspace (range of the observability matrix transposed). Ranks are decided from singular values with the relative threshold `tol`. The input/output transfer function is preserved. ### `modal_truncation(sys, order=None, keep=None)` Modal truncation. Transforms `sys` to a real block-diagonal modal realization and keeps the dominant (slowest) modes, discarding the rest. Complex-conjugate pairs are always kept or dropped together, so the reduced model stays real; the retained poles are exactly the retained eigenvalues. - `order` — target number of retained states (a straddling conjugate pair may push the actual count to `order + 1`). - `keep` — explicit iterable of modal-state indices to retain (expanded to whole blocks). - neither — no truncation (returns the modal-form equivalent). Because coupling to the discarded modes is dropped outright, the DC gain generally shifts; use :func:`residualize` to preserve it. ### `model_description_xml(diagram, *, model_name, guid=None, description='Exported by jaxonomy.library.fmu_export', generation_tool='jaxonomy')` Build the FMI 2.0 modelDescription XML as a string. Parameters: | Name | Type | Description | Default | | ----------------- | ----------- | ------------------------------------------------------------------------------------------------------- | ------------------------------------------- | | `diagram` | `'Diagram'` | A :class:~jaxonomy.framework.diagram.Diagram whose input and output ports define the FMU's I/O surface. | *required* | | `model_name` | `str` | Human-readable model name. Also used as the modelIdentifier (with non-identifier characters stripped). | *required* | | `guid` | \`str | None\` | Optional FMU GUID; auto-generated if None. | | `description` | `str` | Free-form description string. | `'Exported by jaxonomy.library.fmu_export'` | | `generation_tool` | `str` | Stored in the FMU metadata. | `'jaxonomy'` | Returns: | Type | Description | | ----- | ------------------------------------------------ | | `str` | UTF-8 XML string ending with a trailing newline. | ### `nyquist_data(linsys, omegas)` Return Nyquist-contour arrays for `linsys`. The Nyquist plot traces `G(jω)` through the complex plane. This helper returns the real and imaginary parts of `G(jω)` over the supplied positive angular frequencies and additionally the reflected negative-frequency arrays (since `G(-jω) = conj(G(jω))` for a real-coefficient LTI, the reflection is given exactly by `(Re, -Im)`). Consumers can concatenate the negative and positive arrays to obtain the full closed contour used for encirclement counting; for stability margins computed only from the positive sweep, `real` and `imag` are sufficient. For MIMO systems the returned `real` / `imag` arrays preserve the `(K, p, m)` channel structure from :func:`frequency_response`; for SISO they are squeezed to `(K,)`, matching the convention of :func:`bode_data`. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem. | *required* | | `omegas` | | 1-D array-like of positive angular frequencies ω (rad/s). Negative or zero entries are not rejected — G(0) is the DC gain and is returned unchanged, but duplicating with the mirror is not meaningful at ω = 0. | *required* | Returns: | Type | Description | | ------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `dict` | Dictionary with keys: "omega" — positive angular frequencies (rad/s), shape (K,); "real" — Re G(jω); shape (K,) for SISO, (K, p, m) for MIMO; "imag" — Im G(jω); same shape as real; "real_neg" — Re G(-jω) = Re G(jω); same shape as real (mirror; provided for convenience when plotting the full closed contour); "imag_neg" — Im G(-jω) = -Im G(jω); same shape as imag. | Notes Fully differentiable through `A, B, C, D` and `omegas` via the underlying :func:`frequency_response`. ### `observability_gramian(A, C, dt=None)` Observability Gramian :math:`W_o`. Continuous time (`dt is None`) solves .. math:: A^\\mathsf{T} W_o + W_o A = -C^\\mathsf{T} C Discrete time (`dt` given) solves .. math:: A^\\mathsf{T} W_o A - W_o + C^\\mathsf{T} C = 0 ### `pod_basis(X, rank=None, energy=None)` Proper-orthogonal-decomposition basis of a snapshot matrix. Computes the (host-side) thin SVD `X = U Σ Vᵀ` and truncates to `r` left singular vectors, which are the energetically optimal orthonormal modes (Sirovich 1987). Parameters: | Name | Type | Description | Default | | -------- | ----------------- | ------------------------------------------------------------------------------------------------------------ | ---------- | | `X` | | Snapshot matrix, shape (n_features, n_samples). | *required* | | `rank` | `Optional[int]` | Explicit number of modes to keep. | `None` | | `energy` | `Optional[float]` | Cumulative-energy threshold in (0, 1\] (used when rank is None), e.g. 0.99 keeps 99% of the snapshot energy. | `None` | Returns: | Type | Description | | --------- | ----------------------------------------------------------------- | | `ndarray` | (Phi, sigma, r) where Phi has shape (n_features, r) with | | `ndarray` | orthonormal columns, sigma is the full singular-value vector, and | | `int` | r is the retained rank. | ### `pole_zero_map(linsys)` Compute poles and zeros of a :class:`LinearizedSystem`. Poles are the eigenvalues of `A`. Zeros are the (transmission) zeros of the SISO transfer function `G(s) = C (sI − A)⁻¹ B + D`, computed as the finite generalised eigenvalues of the Rosenbrock system pencil .. code-block:: text ``` P(s) = [[ sI − A, −B ], [ C , D ]] ``` by solving the generalised eigenproblem `λ E v = M v` with .. code-block:: text ``` E = [[ I, 0 ], M = [[ A, B ], [ 0, 0 ]] [ C, D ]] ``` Finite eigenvalues (those with non-zero `E`-weight) of this pencil are the invariant zeros of the system; for SISO they coincide with the numerator roots of the transfer function. The high-frequency gain is reported as `D[0, 0]` (the asymptotic value of `G(s) → D` for `|s| → ∞`); for a strictly-proper system (`D = 0`) the leading-coefficient gain is harder to define unambiguously without polynomial fitting and is left to a deeper follow-up. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | --------------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem. Phase 1 ships SISO support only — for MIMO systems the first input/output channel is used. | *required* | Returns: | Type | Description | | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `dict` | Dictionary with keys: "poles" — complex 1-D array of eigenvalues of A, "zeros" — complex 1-D array of (finite) invariant zeros, "gain" — feedthrough scalar D[0, 0]. | Notes Pole computation is differentiable through `A` (via `jnp.linalg.eigvals`). Zero computation uses :func:`scipy.linalg.eig` on the generalised problem and is therefore not currently traceable by JAX — call it eagerly, outside `jit`. A differentiable variant is a deeper follow-up. ### `polynomial_dictionary(degree, include_constant=True)` Monomial dictionary up to `degree`. Layout: `[x_1..x_n, (1), (degree-2..degree monomials)]` — identity first so the state is recoverable. The constant term is included by default (it makes affine dynamics representable in the lifted space). ### `projection_error(X, basis)` Relative projection error `‖X − ΦΦᵀX‖ / ‖X‖` of `X` onto `basis`. `basis` (`Φ`) is assumed to have orthonormal columns. ### `rbf_dictionary(centers, epsilon=1.0)` Gaussian radial-basis dictionary. Lifts to `[x, exp(-epsilon ||x - c_j||²) for each center c_j]` — identity observables first, followed by one RBF feature per center. Parameters: | Name | Type | Description | Default | | --------- | ------- | --------------------------------------- | ---------- | | `centers` | | Array (n_centers, n) of RBF centers. | *required* | | `epsilon` | `float` | Shape parameter of the Gaussian kernel. | `1.0` | ### `reduce(target, method='balred', *, order=None, tol=None, dt=1.0, **kwargs)` Reduce `target` by `method` and return a :class:`ReducedOrderModel`. Parameters: | Name | Type | Description | Default | | ---------- | ---- | ---------------------------------------------------------------------------------------------------------------- | ---------- | | `target` | | An LTI model (linear MOR) or snapshot data (data-driven). | *required* | | `method` | | See the module docstring for the supported names. | `'balred'` | | `order` | | Target reduced order, where the method takes one. | `None` | | `tol` | | Energy/tolerance selector for balanced truncation / minreal. | `None` | | `dt` | | Sampling period for the data-driven predictor blocks. | `1.0` | | `**kwargs` | | Forwarded to the underlying routine (e.g. keep= for modal methods, dictionary= and U= for eDMD, initial_state=). | `{}` | Returns: | Name | Type | Description | | ---- | ---- | ----------------------------------------------------------- | | `A` | | class:ReducedOrderModel whose .system is ready to simulate. | ### `relative_error(x_true, x_approx)` Relative L2 (Frobenius) error `‖x_true − x_approx‖ / ‖x_true‖`. Works for a single trajectory column or a full snapshot matrix. ### `residualize(sys, order=None, keep=None)` Singular-perturbation (residualization) reduction. Like :func:`modal_truncation`, but instead of deleting the fast modes it sets their derivative (continuous) or their increment (discrete) to zero and solves for their quasi-steady value, folding it back into the retained model. This matches the DC gain of the discarded modes (Kokotović, Khalil & O'Reilly 1986). Partitioning the modal realization into retained `(1)` and discarded `(2)` states, continuous time gives .. math:: ``` A_r &= A_{11} - A_{12} A_{22}^{-1} A_{21}, & B_r &= B_1 - A_{12} A_{22}^{-1} B_2, \\ C_r &= C_1 - C_2 A_{22}^{-1} A_{21}, & D_r &= D - C_2 A_{22}^{-1} B_2, ``` and discrete time replaces `A_{22}^{-1}` by `-(I - A_{22})^{-1}`. Selection arguments match :func:`modal_truncation`. ### `retained_energy(singular_values, r)` Fraction of total energy captured by the first `r` POD modes. Energy is measured in squared singular values, `Σ_{i 1` outside the band, so `y -> u` for `|u| >> half_range`. - `gate -> 0` inside the band, so `y -> 0` for `|u| << half_range`. - Continuous everywhere; finite gradient even inside the band (this is the whole reason it exists). - As `sharpness -> inf` the function converges to the hard gate. Parameters: | Name | Type | Description | Default | | ------------ | ---- | ------------------------------------------------------------------------------------------------ | ---------- | | `u` | | Input array. | *required* | | `half_range` | | Positive scalar; band half-width. | *required* | | `sharpness` | | Positive scalar; default 10.0. Larger values give a tighter approximation to the hard dead zone. | `10.0` | Returns: | Type | Description | | ---- | -------------------------------------- | | | Smoothly gated array, same shape as u. | ### `soft_saturate(u, lower, upper, sharpness=10.0)` Smooth (differentiable) saturation between `lower` and `upper`. Approximates `npa.clip(u, lower, upper)` with a tanh-based smooth clamp (per the T-115 spec): :: ``` mid = (lower + upper) / 2 span = upper - lower y = mid + (span / 2) * tanh(sharpness * (u - mid) / span) ``` With `sharpness = 10.0` and a unit span this gives ~99% saturation by `|u - mid| = span`, which is the design default. Properties - As `sharpness -> inf` the function converges to a hard `clip`. - Strictly monotonically increasing in `u` (analytically; in finite precision the slope underflows to zero very far from `mid` because tanh saturates exponentially). - Has a positive derivative across and inside the bound region -- gradients flow through saturation, which is the whole reason this exists. - Requires *finite* `lower` / `upper`; for unbounded sides use the hard :class:`Saturate` block. Parameters: | Name | Type | Description | Default | | ----------- | ---- | -------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `u` | | Input array. | *required* | | `lower` | | Lower limit (scalar or broadcastable array). | *required* | | `upper` | | Upper limit (scalar or broadcastable array). | *required* | | `sharpness` | | Positive scalar; default 10.0. Larger values give a tighter approximation to clip but smaller (and faster vanishing) gradients outside the bounds. | `10.0` | Returns: | Type | Description | | ---- | ------------------------------------------ | | | Smoothly saturated array, same shape as u. | ### `step_response(linsys, t_grid)` Step response of an LTI system, continuous- or discrete-time. **Continuous** (`linsys.dt is None`): for zero initial state and unit step input `u(t) = 1` (for `t ≥ 0`) the closed-form response is .. code-block:: text ``` y(t) = C · ∫₀ᵗ expm(A·s) ds · B + D·1 ``` The matrix integral is computed via the augmented-matrix expm trick (see :func:`_augmented_step_block`) so the routine works correctly for non-invertible `A` (e.g. integrators). When `A` is invertible the same value equals `C·A⁻¹·(expm(A·t) − I)·B + D`. **Discrete** (`linsys.dt` set, e.g. from :func:`discretize`): the exact recurrence `x[k+1] = A x[k] + B`, `y[k] = C x[k] + D` from zero initial state, sampled on the integer grid `k = t/dt`. Every entry of `t_grid` must lie on the sampling grid (within 1e-6·dt) — off-grid times raise `ValueError` rather than silently interpolating. Negative times return the causal pre-input response (zero). For a `"zoh"` discretization the discrete samples equal the continuous step response at `t = k·dt` exactly. Parameters: | Name | Type | Description | Default | | -------- | ------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `linsys` | `LinearizedSystem` | A :class:LinearizedSystem (either time base). | *required* | | `t_grid` | | Scalar or 1-D array of evaluation times. Continuous: negative times are evaluated formally (the closed-form result is still well defined; physically the step starts at t = 0). Discrete: must be integer multiples of dt (see above), and must be concrete (not JAX-traced). | *required* | Returns: | Type | Description | | ---- | -------------------------------------------------------- | | | Array of shape (K, p, m) — step response at each t for | | | each (output, input) pair, where K = len(t_grid), | | | p = n_outputs, m = n_inputs. If t_grid is a scalar | | | the returned shape is (p, m). For SISO systems with | | | scalar t_grid the result squeezes naturally to a scalar. | Notes Fully differentiable through `A, B, C, D` — via :func:`jax.scipy.linalg.expm` (continuous) or `lax.scan` (discrete). For very large continuous state dimensions (`n > 50`) the augmented expm may be slow — the honest fall- back is to simulate the diagram with a :class:`Step` source (deferred to a deeper follow-up). ### `with_observer(plant, observer, *, plant_u_port=0, plant_y_port=0, name='plant_with_observer')` Build a new diagram with `observer` attached to `plant`. Wires the plant's control input through to both the plant and the observer; wires the plant's measurement output through to the observer; exports the observer's `x_hat` estimate as a top-level output of the augmented diagram. The plant's other output ports are not re-exported automatically — call sites that need them can wire them up in a parent diagram. Parameters: | Name | Type | Description | Default | | -------------- | ----- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------- | | `plant` | | A built :class:Diagram (or :class:SystemBase) representing the open-loop plant. Must expose at least one input port (for control u) and one output port (for measurement y). | *required* | | `observer` | | A built observer block (typically :class:jaxonomy.library.Luenberger or any block that accepts (u, y) inputs and produces x_hat as its first output port). | *required* | | `plant_u_port` | `int` | Index of the plant input port carrying the control signal u. Defaults to 0. | `0` | | `plant_y_port` | `int` | Index of the plant output port carrying the measurement y. Defaults to 0. | `0` | | `name` | `str` | Name for the resulting wrapper diagram. | `'plant_with_observer'` | Returns: | Type | Description | | ---- | ------------------------------------------------------------------------------------------------------ | | | A new :class:Diagram containing both plant and | | | observer wired as described, with: | | | a single exported input port u (the control signal, fed to both the plant and the observer's u input), | | | a single exported output port x_hat (the observer's state estimate). | Notes The result is a "passive" instrumentation pattern — the observer reads `(u, y)` and emits an estimate, but the estimate is not fed back into the plant. To close a loop around the estimate (state-feedback control with observed state), build a controller subdiagram and wire its output back to `plant.u` in a parent diagram. ### `write_model_description(diagram, path, *, model_name=None, guid=None, description=None)` Write a Jaxonomy diagram's FMI 2.0 modelDescription.xml to disk. Parameters: | Name | Type | Description | Default | | ------------- | ----------- | ------------------ | ------------------------------- | | `diagram` | `'Diagram'` | Diagram to export. | *required* | | `path` | `str` | Output file path. | *required* | | `model_name` | \`str | None\` | Defaults to diagram.name. | | `guid` | \`str | None\` | Optional GUID. | | `description` | \`str | None\` | Optional free-form description. | Returns: | Type | Description | | ----- | -------------------------------------------------- | | `str` | The same path argument (for chaining convenience). | # Simulation ## `jaxonomy.simulation` ### `BatchSimulationResults` Results from :func:`simulate_batch`. Attributes: | Name | Type | Description | | ------------ | -------------------- | -------------------------------------------------------------------------------------------------------------------------------------- | | `time` | `Any` | Time vector of shape (T,) taken from the first batch run. Later runs are linearly interpolated onto this grid so all batch rows align. | | `outputs` | `dict[str, Any]` | Mapping signal_name -> array with shape (N, T, ...) where N is batch size. | | `used_vmap` | `bool` | True if the vectorised vmap path was used. | | `provenance` | \`ProvenanceManifest | None\` | #### `mean(signal)` Mean trajectory across the batch (axis 0). #### `percentile(signal, p)` `p`-th percentile across the batch at each time index; `p` in `[0, 100]`. #### `std(signal)` Standard deviation across the batch (axis 0). #### `to_simulation_results(idx)` Slice one batch index into a :class:`SimulationResults` (no final context). ### `Decay` Bases: `LeafSystem` xdot = -k * x; output = x. ### `FastRestartSimulator` Stateful single-simulation runner that reuses one JIT-compiled kernel. The simulator is built lazily on the first :meth:`run` so that the `recorded_signals` set passed to the constructor (which selects the set of recorded ports baked into the kernel) is locked in before compilation. Subsequent :meth:`run` calls reuse the same compiled XLA program — only the parameter pytree changes. Parameters: | Name | Type | Description | Default | | ------------------ | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `system` | | A :class:~jaxonomy.framework.diagram.Diagram (or any :class:~jaxonomy.framework.system_base.SystemBase). The structural shape (block topology, port shapes/dtypes, parameter pytree shape) must remain constant across :meth:run calls — only parameter values may change. | *required* | | `t_span` | `tuple[float, float]` | (t_start, t_stop). Locked in on construction since it affects the auto-estimated max_major_steps and the recorder buffer length. A run(t_span=...) override is deferred to a follow-up. | *required* | | `options` | \`SimulatorOptions | None\` | :class:SimulatorOptions. math_backend="jax" and enable_tracing=True (the defaults) are required to get any warm-start benefit; the JIT cache is what makes subsequent calls fast. | | `recorded_signals` | \`dict[str, OutputPort] | None\` | Mapping signal_name -> OutputPort (same convention as :func:simulate). Required so the recorder buffer shape is fixed at construction. | The context-manager protocol is supported but not strictly required; use `with FastRestartSimulator(...) as sim: ...` for symmetry with other resource-holding APIs (the `__exit__` clears the JIT cache reference, freeing the compiled kernel). #### `close()` Drop references to the compiled kernel and base context. Subsequent :meth:`run` calls will rebuild and recompile. The underlying JAX persistent cache (T-017) still holds the compiled XLA program, so the second build remains fast. The per-diagram-identity kernel cache used by :meth:`run_with_diagram` is also cleared. #### `reset(diagram=None)` Clear the cached compiled kernel; optionally rebind to a new diagram. When `diagram` is `None` (default) this is equivalent to :meth:`close` — the next :meth:`run` rebuilds the simulator and recompiles the kernel. The JAX persistent cache typically makes this a fast operation if the diagram structure is unchanged. When `diagram` is provided, the simulator rebinds to it. This is the "swap subsystem variant" path: a parameter sweep where the *structure* (block topology, port shapes, parameter pytree layout) varies between runs. The next :meth:`run` will perform a full recompile against the new diagram. Parameters: | Name | Type | Description | Default | | --------- | --------- | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `diagram` | \`Diagram | None\` | Optional new diagram (or any :class:~jaxonomy.framework.system_base.SystemBase) to bind the simulator to. None means "keep the current one — just drop the cached kernel". | #### `run(parameters=None, initial_state=None)` Run one simulation, optionally patching parameters / initial state first. The first call builds the simulator and JIT-compiles the kernel. Subsequent calls reuse the same compiled program; only the parameter and initial-state values change. Parameters: | Name | Type | Description | Default | | --------------- | ---------------- | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `parameters` | \`dict[str, Any] | None\` | Optional dot-path mapping {"block.param": value, ...} — same convention as :func:simulate_batch's param_batches (without the leading batch axis). Values must have the same shape / dtype as the parameters they replace; otherwise the JIT cache will miss and a recompile will occur (with a UserWarning). Pass None (the default) to run with the base context unchanged. | | `initial_state` | \`Any | None\` | Optional override for the simulator's continuous state at t = t_span[0]. For a :class:~jaxonomy.framework.context.LeafContext-rooted system, pass a single array. For a multi-block :class:~jaxonomy.framework.diagram.Diagram with more than one continuous-state block, pass a sequence of arrays in the order returned by ctx.continuous_state (one per continuous-state subcontext). As a convenience, a single array is auto-wrapped into a single-element list when the diagram has exactly one continuous-state block. Shape/dtype must match the diagram's default continuous state — otherwise the JIT cache will miss and a recompile will occur (with a UserWarning). | Returns: | Name | Type | Description | | ---- | ------------------- | -------------------------------------------- | | `A` | `SimulationResults` | class:SimulationResults populated with time, | | | `SimulationResults` | outputs, and (when | | | `SimulationResults` | options.return_context=True) context. | #### `run_batch(parameters_batch, initial_states_batch=None)` Run `N` simulations differing only by parameters, vmap'd over the cached kernel. Counterpart to :meth:`run` for batched parameter sweeps. Equivalent to :func:`simulate_batch` with `use_vmap=True` but reuses the warm-cached kernel built by the most recent :meth:`run` call (or builds it lazily on first use). Calling :meth:`run` first to warm the cache and then :meth:`run_batch` for the sweep is the typical UX pattern; both code paths share one JIT compile. Parameters: | Name | Type | Description | Default | | ---------------------- | ---------------- | -------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `parameters_batch` | `dict[str, Any]` | {path: (N, ...) array} — same convention as :func:simulate_batch's param_batches. Every value must have the same leading batch size N. | *required* | | `initial_states_batch` | \`Any | None\` | Optional batched initial-state override. Either a single array of shape (N, ...) (auto-wrapped for diagrams with a single continuous-state block) or a list/tuple of (N, ...) arrays (one per continuous-state block, matching ctx.continuous_state ordering). Default None reuses the diagram's default IC for every batch element. | Returns: | Name | Type | Description | | ---- | -------------------------- | --------------------------------- | | `A` | `'BatchSimulationResults'` | class:BatchSimulationResults with | | | `'BatchSimulationResults'` | outputs[name].shape[0] == N. | Notes - The kernel is vmap'd over all leaves of the patched context, not just the explicitly-batched parameter paths. Unpatched leaves are broadcast to shape `(N, ...)` so the vmap'd kernel sees a uniformly batched pytree. - The cached kernel was JIT'd against a *scalar* context signature. `jax.vmap` traces against the batched signature, so the very first call to :meth:`run_batch` incurs one extra trace (still cheap; XLA caches the inner compiled program). Subsequent :meth:`run_batch` calls with the same `N` and the same parameter pytree shape reuse the vmap-cached kernel. #### `run_with_diagram(diagram, parameters=None, initial_state=None, recorded_signals=None)` Run one simulation against `diagram`, caching its kernel by identity. Use this when you hold a *pool* of structurally-different diagrams (e.g. controller variants) and want to rapidly switch between them. The compiled kernel for each distinct Diagram instance is built on first use and reused thereafter — no recompile on cache hit. Compared to :meth:`reset` + :meth:`run` (which drops and rebuilds the kernel on every swap), :meth:`run_with_diagram` keeps one compiled kernel *per Diagram identity* alive, so toggling back and forth between N diagrams in a loop costs N compiles total, not one per call. The user's :meth:`run` and :meth:`reset` APIs are unaffected; this is a purely-additive surface. Parameters: | Name | Type | Description | Default | | ------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `diagram` | `Diagram` | Diagram (or any :class:~jaxonomy.framework.system_base.SystemBase) to simulate. The cache is keyed on id(diagram) so a strong reference to the diagram is held internally for the lifetime of this :class:FastRestartSimulator. Limitation: calling diagram.with_config(...) produces a new Diagram object — the returned object's id differs from the original, so a with_config-rewritten derivative will miss the cache. A cache_key= user-supplied identifier is a natural future extension. | *required* | | `parameters` | \`dict[str, Any] | None\` | Optional dot-path mapping {"block.param": value, ...} — same convention as :meth:run. | | `initial_state` | \`Any | None\` | Optional override for the simulator's continuous state at t = t_span[0] — same convention as :meth:run. | | `recorded_signals` | \`dict[str, OutputPort] | None\` | Optional {name: OutputPort} mapping. Required on the first call for a given diagram (the ports are diagram-specific and locked into the kernel buffer shape at compile time). On warm cache hits the originally-cached mapping is reused; passing a different recorded_signals for the same cached diagram has no effect. If omitted on the first call, falls back to self.recorded_signals — which is typically only valid for the diagram passed to the constructor. | Returns: | Name | Type | Description | | ---- | ------------------- | -------------------------------------------- | | `A` | `SimulationResults` | class:SimulationResults populated with time, | | | `SimulationResults` | outputs, and (when | | | `SimulationResults` | options.return_context=True) context. | ### `LazyResults` A deferred-evaluation wrapper around a :class:`SimulationResults`. Construct via :meth:`SimulationResults.lazy` rather than directly. #### `align_to(name)` Resample every signal to `name`'s native cadence (defers). Convenience wrapper over :meth:`resample` that targets the per-signal time vector for `name`. Useful when one signal is the natural reference clock (e.g. a 1 Hz sensor) and you want every other recorded signal aligned to its ticks before materialising. The returned chain inherits the active backend (eager / polars / duckdb) and routes through the same `resample` translator — i.e. the polars backend uses the asof-join + linear-interp plan from T-015a-followup-resample-pushdown. Raises: | Type | Description | | ---------- | --------------------------------- | | `KeyError` | if name is not a recorded signal. | #### `cadence_of(name)` Classify `name`'s recording cadence. Returns one of: - `"continuous"` — sampled every major step (`time_for(name).shape == self._time.shape` and the value array is full-length). - `"periodic"` — sampled on a fixed schedule (Mode A path: both per-signal times AND outputs are shorter than the global vector and have matching length). - `"event-driven"` — Mode B value-diff dedup populated per-signal times but the output array remained at the global cadence (the recording pipeline could not pin a fixed period to the source `OutputPort`). - `"default"` — no per-signal cadence info available; the signal shares the global :attr:`_time` vector. This is a structural classification derived from the recorded- array shapes — it does not re-invoke the static `ResultsRecorder.classify_signal_cadence` (which requires live `OutputPort` references that aren't carried on :class:`SimulationResults`). The four buckets nevertheless line up 1-to-1 with the four cadence kinds the recording pipeline produces (continuous / periodic / event-driven / default), so a downstream consumer can plan I/O without reaching back into the simulator. Raises: | Type | Description | | ---------- | --------------------------------- | | `KeyError` | if name is not a recorded signal. | #### `collect()` Materialise the chain. Returns `{"time": t, **signals}`. #### `explain()` Render the deferred operation chain as a human-readable string. #### `from_parquet(path, backend='polars')` Load a parquet file written by :meth:`to_parquet`. ###### Parameters path Path to a parquet file produced by :meth:`to_parquet` (or any parquet file with a `time` column). backend `"polars"` (default; T-015a) returns a :class:`LazyResults` with the polars backend pre-enabled. `"duckdb"` (T-015a-followup-resample-pushdown-duckdb) opens the file via DuckDB's `read_parquet(...)` against a fresh in-memory connection — the out-of-core entry point for SQL-style queries. In both cases vector-valued signals stored as `name__i` columns are re-collapsed into `(T, k)` numpy arrays for compatibility with the eager-numpy fallback path. #### `resample(t_new, *, method='linear')` Interpolate every signal onto `t_new` (defers). T-108 phase 2 wires the optional `method=` kwarg through to the T-106 backend (:func:`jaxonomy.library.lookup_table.interp_1d`), so callers can pick the smoother interpolation rules without leaving the lazy pipeline: - `"linear"` (default) — uses the existing fast paths (`np.interp` eager, native polars asof-join + linear-interp). - `"pchip"` — monotone cubic Hermite; smooth gradients, no overshoot near monotonic data. - `"akima"` — Akima 1970 cubic spline; less overshoot than the natural cubic on non-monotone data. - `"cubic"` — natural cubic spline (C^2 continuous, second derivative zero at boundaries). - `"nearest"` / `"flat"` — zero-gradient piecewise constant. For any non-linear method, the polars / DuckDB lazy paths fall back to materialising the upstream chain first and then routing each signal through `interp_1d` per-channel — non-linear interpolation is not expressible as a single polars expression. `method="linear"` keeps the native-polars / native-DuckDB pushdown so large lazy plans stay out-of-core. Polars backend (T-015a-followup-resample-pushdown): for `method="linear"` only, translated natively via two `join_asof` calls (backward + forward) plus a linear-interp expression — no Python `map_batches` callback. Target times must lie within the source range; non-monotonic `t_new` is supported (sorted internally, then re-permuted on output). #### `select(*signals)` Project to a subset of signals (defers). #### `signal(name)` Return `(time, value)` for `name` at its NATIVE cadence. Eager (non-lazy) accessor: bypasses the deferred op chain and reads directly from the underlying recorded arrays. Returns the per-signal timestamp vector populated by `T-013` / `T-013a` (Mode A or Mode B) when available, else falls back to the global :attr:`_time` vector — matching the semantics of :meth:`SimulationResults.time_for`. For Mode B "default"-classified signals (per-signal times are deduplicated but `outputs` stays at full length), the value array is back-projected onto the deduplicated times via `searchsorted` so the returned `(time, value)` pair has consistent shape — same trick used by :meth:`SimulationResults.align`. Raises: | Type | Description | | ---------- | --------------------------------- | | `KeyError` | if name is not a recorded signal. | #### `to_hdf5(path, key='results', chunk_size=10000)` Stream-write the materialised result to an HDF5 file (T-108-followup-streaming-export). Layout: a top-level `time` dataset and an `outputs/` group holding one dataset per signal (vector-valued signals are exploded into `outputs/__` to mirror the parquet column convention). Each dataset is created with `maxshape=(None, ...)` and extended chunk-by-chunk so the file never has to hold the full frame in memory at once. ###### Parameters path Destination `.h5` file path. Overwritten if it exists. key Currently unused — reserved for forward compatibility with multi-result HDF5 files; the layout described above is relative to the file root and not under `key`. chunk_size Rows written per extend. Tune for memory / I/O trade-off; defaults to 10 000 rows. ###### Notes Optional dep: requires `h5py` (`pip install h5py`). Raises :class:`ImportError` if not available. #### `to_numpy()` Alias for :meth:`collect`. #### `to_pandas()` Materialise to a `pandas.DataFrame` (requires pandas). Vector-valued signals are exploded into `name__0`, `name__1` columns. #### `to_parquet(path, batch_size=None)` Write the materialised result to `path` as Parquet. With the polars backend (T-015a) and `batch_size=None`, writes via `LazyFrame.sink_parquet` for true streaming output that never materialises the whole frame in memory. With `batch_size=N`, partitions the output into multiple files `path.0.parquet` / `path.1.parquet` / ... each holding at most `N` rows. With the DuckDB backend (T-015a-followup-resample-pushdown-duckdb) and `batch_size=None`, writes via DuckDB's native `COPY (sql) TO 'path' (FORMAT PARQUET)` — genuinely streaming (DuckDB never materialises the whole result in Python memory). `batch_size=N` partitions on the Python side just like polars. Without an opt-in backend, uses pandas (`pyarrow`) and falls back to polars when pandas is unavailable. #### `to_polars()` Materialise to a `polars.DataFrame` (requires polars). #### `to_zarr(path, chunk_size=10000)` Stream-write the materialised result to a zarr store (T-108-followup-streaming-export). Layout mirrors :meth:`to_hdf5`: a `time` array at the group root and one array per signal under `outputs/` (vector-valued signals exploded as `outputs/__`). Each array is created with `shape=(0,)` and resized in place per chunk. ###### Parameters path Destination directory (a zarr v3 store). Created if absent; overwritten otherwise. chunk_size Rows written per extend. Also used as the underlying zarr chunk dimension so I/O alignment matches the write cadence. ###### Notes Optional dep: requires `zarr` (`pip install zarr`). Raises :class:`ImportError` if not available. #### `where(mask)` Boolean-mask filter on rows (defers). `mask` may be: - a boolean numpy array of length `len(time)`; - a callable `f(t, outputs) -> bool array`; - a string expression that uses `t` and any signal name as free variables (e.g. `"t > 5"`, `"x > 0 & t < 1.5"`). #### `with_duckdb_backend(connection=None)` Opt in to the DuckDB SQL execution path (T-015a-followup-...-duckdb). ###### Parameters connection An existing :class:`duckdb.DuckDBPyConnection`, or `None` (default) to allocate a fresh in-memory connection. Pass an explicit connection to control persistence, extension loading, or thread count. Returns a copy of this :class:`LazyResults` whose terminal materialisers run a single SQL query against an in-memory DuckDB table built from the recorded `(time, outputs)` arrays. Vector-valued signals are exposed as `name__i` columns (matching the polars backend convention). Per-op fallback: `with_signal`, callable `where` predicates, and `resample` are not generally SQL-able and emit :class:`RuntimeWarning` at materialise time, falling back to the eager-numpy path for that op (the chain re-enters DuckDB afterwards). `select` and `where` with a string predicate translate cleanly. #### `with_polars_backend()` Opt in to the polars LazyFrame execution path (T-015a). Returns a copy of this :class:`LazyResults` whose terminal materialisers (`to_polars`/`to_pandas`/`to_parquet`/ `to_numpy`/`collect`) build a `polars.LazyFrame` plan rather than evaluating ops eagerly on numpy arrays. Falls back to eager-numpy on a per-op basis (with :class:`RuntimeWarning`) for ops that polars cannot express natively — currently only callable `where` predicates. `resample` is native polars (asof-join + linear-interp expression; T-015a-followup-resample-pushdown). `with_signal` is executed via collect-and-re-lazy. #### `with_signal(name, fn)` Derive a new signal `name` from existing ones (defers). `fn` receives `(t, outputs)` and returns an array shaped like `time`. ### `ManifestMismatch` Bases: `AssertionError` Raised by :func:`verify_manifest` when two manifests differ. Inherits from :class:`AssertionError` so it composes with `pytest` and standard assertion-style verification flows without callers needing to import the exception explicitly. The exception instance carries a `differences` attribute holding the same `list[tuple[str, Any, Any]]` that :func:`compare_manifests` returns, so programmatic consumers can introspect the drift instead of parsing the message. ### `ODESolverOptions` Options for the ODE solver. See documentation for `simulate` for details on these options. ### `ProvenanceManifest` Reproducibility snapshot for one `simulate(...)` call (T-110). Phase 1 captures library versions, the resolved precision policy, a deterministic system fingerprint, and the relevant :class:`SimulatorOptions` field values. An ISO-8601 UTC timestamp is included so the manifest is self-describing; `git_head` is populated when `simulate` is called from inside a git checkout. The `config_hash` field (T-110-followup-config-hash) is a deterministic SHA-256 of the relevant configuration — same options - same system + same jaxonomy/jax versions yield the same hash across runs and across git commits (timestamp and git HEAD are deliberately excluded). The dataclass is frozen so a recorded manifest can't be silently mutated downstream. #### `from_dict(data)` Construct a :class:`ProvenanceManifest` from a dict produced by :meth:`to_dict` (round-trip helper for serialisation tests). #### `to_dict()` Return a JSON-friendly dict representation of the manifest. #### `to_json(*, indent=None)` Serialise :meth:`to_dict` via `json.dumps`. ### `ResultsWithProvenance` Pair a results object with its :class:`ProvenanceManifest`. Attribute access is forwarded to the underlying `results` instance, so `wrapped.outputs[name]` works exactly like `results.outputs[name]`. `wrapped.results` and `wrapped.provenance` give explicit access to either side. The wrapper is frozen so the pairing can't be silently mutated. Construction does not copy or wrap the underlying results — the wrapper holds a reference, nothing else. ### `SimulationError` Bases: `JaxonomyError` Raised when a simulator entry point fails at trace or run time. Attributes: | Name | Type | Description | | ------- | ---- | -------------------------------------------------------------------------------------------------------- | | `cause` | | The original exception. Accessible as __cause__ too. | | `block` | | Name of the block that appeared innermost in the traceback, or None if no block context was recoverable. | | `port` | | Name of the port if the failure was inside a port callback (best-effort), else None. | ### `SimulationResults` Bases: `NamedTuple` Data structure for the results of a simulation. Attributes: | Name | Type | Description | | ------------ | ------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `context` | `ContextBase` | The output context of the simulation, containing final states, times, etc. May be None if return_context=False was passed to simulate. | | `outputs` | `dict[str, Array]` | A dictionary of the outputs of the simulation, keyed by the name provided to recorded_signals in simulate. May be None if recorded_signals is not provided to simulate. | | `time` | `Array` | The time vector of the simulation. | | `parameters` | `dict[str, Any]` | The parameters used in the simulation, used in ensemble simulations to identify different runs. | #### `align(time_vector, signals=None)` Resample recorded signals onto a common time vector (T-013). Useful when per-signal timestamps have been captured at different native rates and a rectangular timeline is required for plotting or further processing. Parameters: | Name | Type | Description | Default | | ------------- | ---- | ------------------------------------------------------------------------------- | ---------- | | `time_vector` | | 1-D array of times to sample at. | *required* | | `signals` | | Optional iterable of signal names to include. Defaults to all recorded signals. | `None` | Returns: | Type | Description | | ---- | ------------------------------------------------------- | | | A new :class:SimulationResults where every requested | | | signal has been linearly interpolated onto time_vector. | | | per_signal_times is reset to None because all signals | | | now share the same timeline. | #### `lazy()` Return a :class:`LazyResults` wrapper for fluent / deferred queries. See :mod:`jaxonomy.simulation.lazy_results` for the full API. #### `query(t, signal=None)` Interpolate recorded signal(s) at time `t` (T-012, T-012a). Default path uses a linear interpolant over the recorded time/value arrays — fast, consistent across solvers, sufficient for the common post-hoc-sampling workflow. When the simulation was run with `SimulatorOptions(record_solver_states=True)` the `solver_states` field is populated and `query` switches to a PCHIP cubic-Hermite interpolant built from the same recorded samples (T-012a partial). PCHIP is shape-preserving — no overshoot at zero-order-hold plateaus — and gives ~3 orders of magnitude better accuracy than linear on smooth (continuous) signals. Discrete (zero-order-hold) signals are detected by constant-plateau runs and fall back to step interpolation rather than smoothing through the steps. The ODE solver's *native* dense interpolant (Dopri5's 5th-order polynomial, BDF's polynomial predictor) — which would give sub-ULP accuracy — remains a follow-up since it requires plumbing per-major-step solver state through the recording pipeline. Parameters: | Name | Type | Description | Default | | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `t` | | Scalar time, or 1-D array of times. | *required* | | `signal` | `Optional[str]` | Optional signal name. If provided, return only that signal's interpolated value. If None, return a dict of all recorded signals. | `None` | Returns: | Type | Description | | ---- | --------------------------------------------------------------------------------------- | | | If signal is provided: the interpolated array (scalar when t is scalar, 1-D otherwise). | | | Otherwise: dict[str, Array] matching self.outputs. | Raises: | Type | Description | | ------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------- | | `ValueError` | if t falls outside \[time[0], time[-1]\], or if recorded_signals was not supplied to simulate (self.outputs is None), or if signal is not in self.outputs. | #### `time_for(signal)` Return the time vector associated with `signal`. Falls back to `self.time` when `per_signal_times` is None or does not contain `signal` — matching the legacy behaviour where all recorded signals share one timeline. ### `Simulator` Class for orchestrating simulations of hybrid dynamical systems. See the `simulate` function for more details. #### `__init__(system, ode_solver=None, options=None)` Initialize the simulator. Parameters: | Name | Type | Description | Default | | ------------ | ------------------ | ---------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `SystemBase` | The hybrid dynamical system to simulate. | *required* | | `ode_solver` | `ODESolverBase` | The ODE solver to use for integrating the continuous-time component of the system. If not provided, a default solver will be used. | `None` | | `options` | `SimulatorOptions` | Options for the simulation process. See simulate for details. | `None` | #### `compile(tf, context)` Warm up / pre-compile the simulation advance_to method on the device. #### `initialize(context)` Perform initial setup for the simulation. #### `while_loop(cond_fun, body_fun, val)` Structured control flow primitive for a while loop. Dispatches to a bounded while loop when • `enable_autodiff=True` (required for reverse-mode AD), or • the caller explicitly set `max_major_steps` in SimulatorOptions (acts as a hard simulation budget, e.g. for Zeno protection). Otherwise the standard unbounded `lax.while_loop` (JAX backend) or a pure-Python loop (NumPy backend) is used. ### `SimulatorOptions` Options for the hybrid simulator. See documentation for `simulate` for details on these options. This also contains all configuration for the ODE solver as a subset of options so that multiple options classes don't need to be created separately. ### `algebraic_row_mask(system)` Boolean mask: True for rows of M that are identically zero. Returns `None` if the system has no mass matrix (purely ODE form). Rows with `M[i, :] == 0` correspond to algebraic constraints in the semi-explicit form `M·ẋ = f`. ### `attach_provenance_to_batch(results, system, options)` Attach a :class:`ProvenanceManifest` to `results` in place and return it. Standalone helper for the rare case where a user has a :class:`BatchSimulationResults` produced without `record_provenance=True` and now wants reproducibility metadata attached (for example, after-the-fact archival). In the normal flow, :func:`simulate_batch` and :func:`simulate_distributed` already wire up the manifest when `options.record_provenance=True`; this helper is just the explicit, opt-in escape hatch. Parameters: | Name | Type | Description | Default | | --------- | ------------------------ | ------------------------------------------ | ---------------------------------------------------------------------------- | | `results` | `BatchSimulationResults` | A :class:BatchSimulationResults to mutate. | *required* | | `system` | \`Diagram | None\` | The diagram that was simulated (passed through to :func:compute_provenance). | | `options` | \`SimulatorOptions | None\` | The :class:SimulatorOptions used for the batch run. | Returns: | Type | Description | | ------------------------ | -------------------------------------------------- | | `BatchSimulationResults` | The same results instance, with results.provenance | | `BatchSimulationResults` | populated. | ### `bundle_results(results)` Wrap `results` in a :class:`ResultsWithProvenance` if applicable. When `results.provenance` is populated (a non-None manifest), the return value is a :class:`ResultsWithProvenance` carrying both the original results object and its provenance. When `results` has no `provenance` attribute or that attribute is `None`, the original `results` object is returned unchanged — so callers can sprinkle `bundle_results(...)` in front of every simulate call without breaking byte-equivalent default-off paths. This helper is purely ergonomic. The legacy `results.provenance` field is left in place; nothing about the underlying results object is mutated. ### `compare_manifests(actual, expected, *, ignore_fields=None)` Diff two manifests field-by-field. Parameters: | Name | Type | Description | Default | | --------------- | -------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `actual` | `ProvenanceManifest` | the manifest produced by the run being checked. | *required* | | `expected` | `ProvenanceManifest` | the reference manifest (e.g. loaded from a published release-tag artifact via :func:load_manifest). | *required* | | `ignore_fields` | `Optional[set[str]]` | top-level field names whose drift is acceptable. Defaults to {"timestamp"} since the timestamp is always different and never load-bearing for reproducibility. Pass ignore_fields=set() to compare every field including the timestamp. | `None` | Returns: | Type | Description | | ---------------------------- | ------------------------------------------------------------- | | `list[tuple[str, Any, Any]]` | A flat list of (dotted_path, actual_value, expected_value) | | `list[tuple[str, Any, Any]]` | triples — one per differing leaf. An empty list means the two | | `list[tuple[str, Any, Any]]` | manifests agree on every compared field. | ### `compute_constraint_residual(system, context)` Return the residual of the algebraic constraints at the given context. For a semi-explicit DAE `M·ẋ = f(t, x, p)`, rows of `M` that are zero enforce `f_a(t, x, p) = 0`. This function returns the concatenated `f_a` vector — ideally near zero on a converged solver step, and any growth over simulation time indicates constraint drift. Returns `None` for systems without a mass matrix (no constraints to satisfy; `M` is identity). ### `compute_provenance(system, options=None, *, include_git=True, timestamp=None)` Build a :class:`ProvenanceManifest` for `system` + `options`. All capture happens in plain Python — no JAX tracing — so the function is safe to call before or after a JIT'd simulation kernel. Parameters: | Name | Type | Description | Default | | ------------- | ------------------------------ | ---------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `Optional['SystemBase']` | the system being simulated (may be None for tests or pre-built recordings). | *required* | | `options` | `Optional['SimulatorOptions']` | the active :class:SimulatorOptions; None records an empty options dict. | `None` | | `include_git` | `bool` | when False, skip the git-HEAD lookup (useful when the caller knows it isn't in a git checkout or wants a faster path). | `True` | | `timestamp` | `Optional[str]` | optional override (ISO-8601 string). Defaults to the current UTC time. Override is useful for deterministic tests. | `None` | Returns: | Type | Description | | -------------------- | -------------------------------------- | | `ProvenanceManifest` | A populated :class:ProvenanceManifest. | ### `constraint_residual_norm(system, context)` Max-abs residual of the algebraic constraints, or `None` for pure ODE. `||f_a||_∞` is the natural comparison quantity for a drift threshold: a single violated constraint should trigger the warning even if the average residual is tiny. ### `estimate_max_major_steps(system, tspan, max_major_step_length=None, safety_factor=2)` Heuristic for estimating the required number of major steps. This is used to bound the number of iterations in the while loop in the `simulate` function when automatic differentiation is enabled. The number of major steps is determined by the smallest discrete period in the system and the length of the simulation interval. The number of major steps is bounded by the length of the simulation interval divided by the smallest discrete period, with a safety factor applied. The safety factor accounts for unscheduled major steps that may be triggered by zero-crossing events. This function assumes static time variables, so cannot be called from within traced (JAX-transformed) functions. This is typically the case when the beginning or end time of the simulation is a variable that will be differentiated. In this case `estimate_max_major_steps` should be called statically ahead of time to determine a reasonable bound for `max_major_steps`. Parameters: | Name | Type | Description | Default | | ----------------------- | --------------------- | ------------------------------------------------------------------------------------------------------------------------------------ | ---------- | | `system` | `SystemBase` | The system to simulate. | *required* | | `tspan` | `tuple[float, float]` | The time interval to simulate over. | *required* | | `max_major_step_length` | `float` | The maximum length of a major step. If provided, this will be used to bound the number of major steps. Otherwise it will be ignored. | `None` | | `safety_factor` | `int` | The safety factor to apply to the number of major steps. Defaults to 2. | `2` | ### `event_time_gradient(guard_fn, ode_rhs_fn, t_event, state_at_event_fn, params, *, eps=1e-30)` Compute `∂t_event/∂params` via the implicit-function theorem. Parameters: | Name | Type | Description | Default | | ------------------- | -------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `guard_fn` | `Callable[[float, Any, Any], ndarray]` | (t, state, params) -> scalar — the zero-crossing guard. Must be JAX-traceable in all three arguments. | *required* | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt — the continuous RHS evaluated at the event boundary. Same PyTree structure as the state. | *required* | | `t_event` | `ndarray` | Scalar time at which the guard fires. | *required* | | `state_at_event_fn` | \`Callable\[[Any], Any\] | Any\` | Either * a callable params -> state that reconstructs the recorded event state from the parameters (so JAX can propagate the trajectory sensitivity ∂x_e/∂p), or * a constant PyTree of state values (no implicit dependence on params). The callable form is the general case; the constant form is equivalent to passing lambda p: and is useful when the user only wants the explicit ∂g/∂p contribution. | | `params` | `Any` | Parameter PyTree to differentiate with respect to. May be a scalar, ndarray, or any nested container. | *required* | | `eps` | `float` | Floor used to clip the denominator (∂g/∂x · ẋ + ∂g/∂t) away from zero before division — keeps jax.grad finite at grazing crossings. Sign-preserving. | `1e-30` | Returns: | Type | Description | | ----- | --------------------------------------------------------- | | `Any` | The PyTree of ∂t_event/∂params with the same structure as | | `Any` | params. | ### `event_time_jacobian(guard_fn, ode_rhs_fn, t_event, state_at_event_fn, params, *, eps=1e-30)` Vector-valued convenience wrapper of :func:`event_time_gradient`. Identical semantics, but returns a flat ndarray so that the result composes cleanly with downstream linear-algebra (Sobol sampling, Fisher information, etc.). `params` should be a 1-D array. For a 1-D `params` array of length `n_p`, returns shape `(n_p,)`. ### `event_times_gradient(results, params, guards, ode_rhs_fn, state_at_event_fn, *, event_indices=None, eps=1e-30)` Batch event-time gradients across all firings recorded by `simulate(..., options=SimulatorOptions(record_event_times=True))`. For each recorded event in `results.event_times`, applies the implicit-function theorem (T-125 phase 1) to every firing instant and returns the per-firing gradient PyTrees keyed by event index. Parameters: | Name | Type | Description | Default | | ------------------- | ---------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `results` | `Any` | A :class:SimulationResults whose event_times is populated (i.e., the simulation was run with SimulatorOptions(record_event_times=True)). Calling this helper on a results whose event_times is None raises ValueError with the remediation hint — the default-off path is preserved by simply not invoking this helper. | *required* | | `params` | `Any` | Parameter PyTree to differentiate with respect to. Same semantics as :func:event_time_gradient. | *required* | | `guards` | `Any` | Either a single guard callable (t, state, params) -> scalar applied to every recorded event, or a mapping {event_index: guard_fn} providing a distinct guard per event slot. The latter form is intended for multi-event diagrams where each event index has its own zero-crossing function. | *required* | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt — same as in :func:event_time_gradient. Reused across all firings. | *required* | | `state_at_event_fn` | `Callable[[float, Any], Any]` | (t_e, params) -> state — reconstructs the trajectory state at firing time t_e parametrized by params. This is the simpler state_fn form noted in the T-125-followup-multi-event task spec: callers express per-event-class behaviour via the t_e argument rather than per-event callables. The deeper per-event-class form is a deferred followup. | *required* | | `event_indices` | `Any` | Optional iterable of event indices to compute gradients for. When None (default), every event index present in results.event_times is processed. Indices not present in results.event_times raise KeyError. | `None` | | `eps` | `float` | Forwarded to :func:event_time_gradient — denominator floor for grazing crossings. | `1e-30` | Returns: | Type | Description | | ------ | ----------------------------------------------------------- | | `dict` | {event_index: stacked_gradient} — for each event index, | | `dict` | the per-firing gradients stacked along a leading axis (so a | | `dict` | gradient that is itself a PyTree leaf of shape S becomes | | `dict` | an array of shape (n_firings,) + S; PyTree containers are | | `dict` | preserved by mapping the stack over leaves). Events that | | `dict` | fired zero times yield an empty leading axis. | ### `load_manifest(path)` Load a :class:`ProvenanceManifest` from a JSON file written by :meth:`ProvenanceManifest.save`. Parameters: | Name | Type | Description | Default | | ------ | ---- | ------------------------------------------------------------ | ---------- | | `path` | | filesystem path (str or pathlib.Path) of the saved manifest. | *required* | Returns: | Type | Description | | -------------------- | -------------------------------------------- | | `ProvenanceManifest` | The reconstructed :class:ProvenanceManifest. | Raises: | Type | Description | | ------------------- | ------------------------------ | | `FileNotFoundError` | if path does not exist. | | `JSONDecodeError` | if the file is not valid JSON. | ### `multi_event_time_gradient(guard_fn, ode_rhs_fn, reset_map_fn, initial_state, event_times, params, *, t0=0.0, eps=1e-30, rtol=1e-10, atol=1e-12, return_state_sensitivity=False)` Saltation gradient `dt_e/dp` for *every* firing along a hybrid trajectory, propagating the forward sensitivity through reset maps. Unlike :func:`event_time_gradient` — which needs the caller to supply a closed-form `state_at_event_fn` for the trajectory sensitivity, and so only gets the first firing right — this helper reconstructs `∂x_e/∂p` itself by integrating the variational equation along each recorded arc and applying the saltation jump at each event. It is the correct path for multi-bounce / repeated-event problems where each firing re-initialises the arc from the previous reset map. Parameters: | Name | Type | Description | Default | | -------------------------- | ---------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `guard_fn` | \`Callable\[[float, Any, Any], ndarray\] | Any\` | (t, state, params) -> scalar zero-crossing guard, or a sequence of such callables aligned with event_times (one per firing) for heterogeneous events. | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt continuous RHS, shared across all arcs. Must be JAX-traceable. | *required* | | `reset_map_fn` | \`Callable\[[float, Any, Any], Any\] | Any\` | (t_e, state_minus, params) -> state_plus reset map applied at each firing, or a sequence aligned with event_times. Use the identity map (lambda t, x, p: x) for events that only observe a crossing without resetting state. | | `initial_state` | \`Callable\[[Any], Any\] | Any\` | either a callable params -> x0 (so the seed sensitivity S(t0) = ∂x0/∂p is captured) or a constant state PyTree (seed sensitivity is then zero). | | `event_times` | `Any` | ordered sequence / array of recorded firing instants [t_1, ..., t_n] (strictly increasing, all > t0). These are the recorded primal event times — e.g. from results.event_times. | *required* | | `params` | `Any` | parameter PyTree to differentiate with respect to. | *required* | | `t0` | `float` | trajectory start time (default 0.0). | `0.0` | | `eps` | `float` | sign-preserving floor on the implicit-function denominator (∂g/∂x · ẋ + ∂g/∂t) — guards grazing crossings. | `1e-30` | | `rtol` | `float` | relative tolerance for the augmented (state + sensitivity) arc integration. | `1e-10` | | `atol` | `float` | absolute tolerance for the augmented arc integration. | `1e-12` | | `return_state_sensitivity` | `bool` | when True, also return the list of pre-event forward sensitivities S⁻(t_e) (flat (n_x, n_p) arrays) for inspection / debugging. | `False` | Returns: | Type | Description | | ----- | --------------------------------------------------------------- | | `Any` | The per-firing dt_e/dp stacked along a leading axis of length | | `Any` | n and shaped like params (a scalar parameter yields shape | | `Any` | (n,); a PyTree parameter yields the same PyTree with each leaf | | `Any` | carrying a leading firing axis). If return_state_sensitivity is | | `Any` | True, returns (grads, [S_minus_1, ..., S_minus_n]). | Notes Fully JAX-traceable (the arc integration uses `jax.experimental.ode.odeint`). Default-off and purely additive: the simulator path is untouched and callers who don't import this helper pay zero cost. ### `scalar_cost_simulate(system, context_fn, t_span, params, cost_fn=None, *, options=None, return_grad=False)` Reverse-mode differentiable scalar cost from a simulation (T-A1). Resolves the most common autodiff friction in jaxonomy: you cannot record a trajectory and reduce it to a cost under `jax.grad`, because `enable_autodiff=True` forbids `save_time_series=True` (recording is not `vmap`/AD-safe). The supported pattern is to **accumulate the cost inside the diagram** — e.g. add an `Integrator` whose input is the running cost `L(t, x, u)` — and read the final accumulated value off the context at `t_span[1]`. This helper packages that pattern so the canonical `cost = f(params)` / `grad = jax.grad(f)(params)` workflow works out of the box. It is the reverse-mode counterpart to :func:`simulate_jacfwd`: use this for a *scalar* objective (optimisation / tuning), and `simulate_jacfwd` for a Jacobian when `n_params` is small relative to the output size. Parameters: | Name | Type | Description | Default | | ------------- | ------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `SystemBase` | the Diagram / LeafSystem to simulate. | *required* | | `context_fn` | `Callable[[Any], ContextBase]` | context_fn(params) -> Context building the initial context with params applied (e.g. via diagram.with_parameters(...).create_context() or by setting ctx.parameters). Differentiation flows through this. | *required* | | `t_span` | `tuple[float, float]` | (t0, tf) simulation interval. | *required* | | `params` | `Any` | parameter pytree — the differentiation argument. | *required* | | `cost_fn` | `Callable[[ContextBase], Any]` | cost_fn(final_context) -> scalar reducing the final context to the objective (typically reading the accumulated-cost state slot, e.g. lambda ctx: ctx[acc.system_id].continuous_state[0]). Defaults to sum(final continuous_state) with a note that you almost always want to supply your own. | `None` | | `options` | `SimulatorOptions` | SimulatorOptions. enable_autodiff is forced True and recorded_signals is cleared (recording is incompatible with AD). Set max_major_steps for systems with many events or when differentiating w.r.t. tf. | `None` | | `return_grad` | `bool` | when True, return (value, grad) via jax.value_and_grad; otherwise return just the scalar value (compose your own jax.grad / jax.value_and_grad over a lambda p: scalar_cost_simulate(...) closure). | `False` | Returns: | Type | Description | | ---- | -------------------------------------------------------- | | | The scalar cost, or (value, grad) when return_grad=True. | Example > > > #### `acc` is an Integrator accumulating the running cost inside the diagram > > > > > > def make_ctx(theta): ... return diagram.with_parameters({"ctrl.kp": theta}).create_context() cost = lambda ctx: ctx[acc.system_id].continuous_state[0] f = lambda th: scalar_cost_simulate(diagram, make_ctx, (0., 5.), th, cost) J = jax.grad(f)(jnp.array(1.0)) # doctest: +SKIP val, grad = scalar_cost_simulate(diagram, make_ctx, (0., 5.), ... jnp.array(1.0), cost, return_grad=True) ### `simulate(system, context, t_span=None, options=None, results_options=None, recorded_signals=None, postprocess=True, flatten=False, *, tspan=None)` Simulate the hybrid dynamical system defined by `system`. The parameters and initial state are defined by `context`. The simulation time runs from `tspan[0]` to `tspan[1]`. The simulation is "hybrid" in the sense that it handles dynamical systems with both discrete and continuous components. The continuous components are integrated using an ODE solver, while discrete components are updated periodically as specified by the individual system components. The continuous and discrete states can also be modified by "zero-crossing" events, which trigger when scalar-valued guard functions cross zero in a specified direction. The simulation is thus broken into "major" steps, which consist of the following, in order: (1) Perform any periodic updates to the discrete state. (2) Check if the discrete update triggered any zero-crossing events and handle associated reset maps if necessary. (3) Advance the continuous state using an ODE solver until the next discrete update or zero-crossing, localizing the zero-crossing with a bisection search. (4) Store the results data. (5) If the ODE solver terminated due to a zero-crossing, handle the reset map. The steps taken by the ODE solver are "minor" steps in this simulation. The behavior of the ODE solver and the hybrid simulation in general can be controlled by configuring `SimulatorOptions`. Available settings are as follows: SimulatorOptions enable_tracing (bool): Allow JAX tracing for JIT compilation max_major_step_length (float): Maximum length of a major step max_major_steps (int): The maximum number of major steps to take in the simulation. This is necessary for automatic differentiation - otherwise the "while" loop is non-differentiable. With the default value of None, a heuristic is used to determine the maximum number of steps based on the periodic update events and time interval. rtol (float): Relative tolerance for the ODE solver. Default is 1e-6. atol (float): Absolute tolerance for the ODE solver. Default is 1e-8. min_minor_step_size (float): Minimum step size for the ODE solver. max_minor_step_size (float): Maximum step size for the ODE solver. ode_solver_method (str): The DE solver to use. Default is "auto", which will use the Dopri5/Jax if JAX tracing is enabled, otherwise the SciPy Dopri5 solver. save_time_series (bool): This option determines whether the simulator saves any data. If the simulation is initiated from `simulate` this will be set automatically depending on whether `recorded_signals` is provided. Hence, this should not need to be manually configured. recorded_signals (dict[str, OutputPort]): Dictionary of ports or other cache sources for which the time series should be recorded. Note that if the simulation is initiated from `simulate` and `recorded_signals` is provided as a kwarg to `simulate`, anything set here will be overridden. Hence, this should not need to be manually configured. return_context (bool): If the context is not needed for anything, opting to not return it can speed up compilation times. For instance, typical simulation calls from the UI don't use the context for anything, so model_interface.py will set `return_context=False` for performance. postprocess (bool): If using buffered results recording (i.e. with JAX numerical backend), this determines whether to automatically trim the buffer after the simulation is complete. This is the default behavior, which will serve unless the full call to `simulate` needs to be traced (e.g. with `grad` or `vmap`). The return value is a `SimulationResults` object, which is a named tuple containing all recorded signals as well as the final context (if `options.return_context` is `True`). Signals can be recorded by providing a dict of (name, signal_source) pairs Typically the signal sources will be output ports, but they can actually be any `SystemCallback` object in the system. Parameters: | Name | Type | Description | Default | | ------------------ | ----------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `SystemBase` | The hybrid dynamical system to simulate. | *required* | | `context` | `ContextBase` | The initial state and parameters of the system. | *required* | | `tspan` | `tuple[float, float]` | The start and end times of the simulation. | `None` | | `options` | `SimulatorOptions` | Options for the simulation process and ODE solver. | `None` | | `results_options` | `ResultsOptions` | Options related to how the outputs are stored, interpolated, and returned. | `None` | | `recorded_signals` | `dict[str, OutputPort]` | Dictionary of ports for which the time series should be recorded. Each recorded series is read back as results.outputs[name], sampled at results.time — there is no results.time_series attribute. If the recording buffer (SimulatorOptions. buffer_length) fills mid-run, the series is kept at reduced resolution (uniform decimation, still spanning the whole run) and a UserWarning recommends a larger buffer_length. | `None` | Returns: | Name | Type | Description | | ------------------- | ------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `SimulationResults` | `SimulationResults` | A named tuple containing the recorded signals (results.time and results.outputs, a dict keyed by the names given in recorded_signals) and the final context (if options.return_context is True). | Notes If `recorded_signals` is provided as a kwarg, it will override any entry in `options.recorded_signals`. This will be deprecated in the future in favor of only passing via `options`. A *repeated* call reuses the compiled kernel instead of re-tracing and re-compiling: the trace and XLA compile are a fixed cost that scales with block count, not with the length of the simulated span, and repeating a call used to pay it again in full. Reuse requires the same system, options, `t_span` and `context`, all of which are baked into the traced program. A call that varies the span or the context therefore still compiles. If you need reuse *across* spans or initial states — a snapshot walk, an interactive stepper, an MPC inner loop — construct a :class:`Simulator` once and call its `advance_to`, which takes the end time and context as traced arguments by design and so reuses one compiled kernel across all of them. Set `SimulatorOptions(reuse_compiled_kernel=False)` to opt out per call, or call :func:`clear_simulate_cache` to drop the memo. ### `simulate_batch(diagram, t_span, param_batches, options=None, recorded_signals=None, results_options=None, use_vmap=False, _force_loop=False, lazy=False)` Run `N` simulations differing only by parameters given in `param_batches`. **Execution paths**: - **Kernel path** (default for pure-JAX diagrams): builds the simulator once, compiles a single JIT kernel, and injects each batch element's parameters directly into a context pytree (no `ParameterCache` mutations). This eliminates N−1 recompilations and is substantially faster for moderate to large N. - **vmap path** (opt-in, `use_vmap=True`, pure-JAX only): further vectorises over the batch dimension with `jax.vmap` so all N simulations run as a single XLA call. Requires that all parameter values have compatible shapes and that the simulation fits in device memory N-fold. **CPU note (updated by T-019-followup, 2026-07-10).** The post-vmap finalize is now fully vectorised (batched trim + batched binary-search linear resampling instead of a per-row host loop), which removed the old CPU penalty: on the CPU damped-oscillator sweep at `N=1000` the vmap path improved from ~1.28 s to ~0.41 s against ~0.33 s for the kernel path (naive loop ~130 s, FastRestart ~0.30 s). CPU kernel-path wins are now marginal; on GPU / TPU vmap wins decisively. The old CPU+small-batch `UserWarning` was removed along with the penalty it warned about. - **Loop path** (forced when `CustomPythonBlock` or FMU blocks are present, or when `_force_loop=True`): the safe fallback — N independent calls to `simulate` + `with_parameters`. Parameters: | Name | Type | Description | Default | | ------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `diagram` | `Diagram` | Template diagram (unchanged). | *required* | | `t_span` | `tuple[float, float]` | (t_start, t_stop). | *required* | | `param_batches` | `dict[str, Any]` | Dot-path keys mapping to 1-D arrays of length N (same N for every entry), e.g. {"gain.gain": jnp.linspace(0.5, 2.0, 16)}. | *required* | | `options` | \`SimulatorOptions | None\` | :class:SimulatorOptions with math_backend="jax" and max_major_steps set (required). | | `recorded_signals` | \`dict[str, OutputPort] | None\` | Same convention as :func:simulate (ports refer to the template diagram; they are remapped per updated diagram for the loop path but used directly for the kernel / vmap path). | | `results_options` | \`ResultsOptions | None\` | Optional :class:ResultsOptions passed through. | | `use_vmap` | `bool` | If True, attempt vectorisation via jax.vmap (pure-JAX diagrams only). Raises ValueError if the diagram is not pure-JAX. | `False` | | `_force_loop` | `bool` | If True, always use the loop path regardless of diagram type (useful for testing / debugging). | `False` | Returns: | Type | Description | | ------------------------ | -------------------------------------------------------------- | | `BatchSimulationResults` | class:BatchSimulationResults with outputs[name].shape[0] == N. | Raises: | Type | Description | | ------------ | -------------------------------------------------------------- | | `ValueError` | Inconsistent batch sizes, missing options, or invalid backend. | | `TypeError` | diagram is not a :class:~jaxonomy.framework.diagram.Diagram. | ### `simulate_cloud(*args, **kwargs)` Run a batch of simulations on a remote execution backend. Not available in this build: no cloud execution backend is bundled. Use the local :func:`jaxonomy.simulate` / batch / distributed runners instead. This entry point is reserved and will be implemented, and documented, once the backend ships. Raises: | Type | Description | | --------------------- | ---------------------- | | `NotImplementedError` | always, in this build. | ### `simulate_jacfwd(system, context_fn, t_span, params, output_fn=None, *, options=None, record_provenance=False)` Forward-mode Jacobian of a simulation w.r.t. parameters (T-100). Wraps `jax.jacfwd` over a parametrised simulation. Use this when the parameter count is small compared to the output count (`n_params < n_outputs / 5` is a useful heuristic) — forward-mode AD scales linearly with input dim; reverse-mode (`jax.grad` / `jax.jacrev`) scales with output dim. The implementation uses `enable_autodiff=False` to bypass the custom-VJP `simulate` defines for reverse-mode (custom_vjp blocks forward-mode trace with a clear `TypeError`); the underlying simulator's natural JAX trace carries the tangent. Forward-mode plumbing is already exercised internally by `linearize`, the BDF Jacobian solve, and the Kalman/EKF blocks — this function exposes that plumbing as a stable public surface. Parameters: | Name | Type | Description | Default | | ------------------- | ---------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `system` | `SystemBase` | a Diagram or LeafSystem to simulate. | *required* | | `context_fn` | `Callable[..., ContextBase]` | a callable context_fn(params) -> Context that constructs an initial context with the given parameter pytree applied. | *required* | | `t_span` | `tuple[float, float]` | (t0, tf) simulation interval. | *required* | | `params` | `Any` | parameter pytree (the differentiation argument). | *required* | | `output_fn` | `Callable[[Any], Any]` | callable applied to the final Context to produce a scalar or array output. Defaults to extracting the final continuous state. | `None` | | `options` | `SimulatorOptions` | SimulatorOptions. enable_autodiff is forced to False for the JVP path; pass rtol / atol / ode_solver_method to control accuracy. | `None` | | `record_provenance` | `bool` | when True, return (jacobian, manifest) with a populated :class:~jaxonomy.simulation.provenance.ProvenanceManifest describing the run. Default False keeps the historical single-return contract byte-equivalent. The manifest is computed in plain Python around the :func:jax.jacfwd call — never inside the trace — so the default-off path adds zero work. See T-110-followup-attach-on-jacfwd. | `False` | Returns: | Type | Description | | ---- | ------------------------------------------------------ | | | J = ∂output/∂params with shape determined by | | | jax.jacfwd's output convention (output × params). When | | | record_provenance=True, returns (J, manifest) instead. | Example > > > def make_ctx(a): ... ctx = sys.create_context() ... ctx.parameters['a'] = a ... return ctx J = simulate_jacfwd(sys, make_ctx, (0., 2.), jnp.array(1.5)) J, m = simulate_jacfwd(sys, make_ctx, (0., 2.), jnp.array(1.5), ... record_provenance=True) ### `simulate_static_sweep(diagram_factory, t_span, static_param_grid, options, recorded_signals_factory, results_options=None, mode='zip')` Sweep over **static** parameters by rebuilding the diagram per element. Unlike :func:`simulate_batch`, which patches a single diagram's context with different dynamic parameter values, this helper accepts a factory that produces a *fresh* :class:`Diagram` for each combination of static param values. Each element is simulated independently in a Python loop; outputs are stacked into a :class:`BatchSimulationResults`-shaped struct. Because each diagram is fresh, port references are also per-diagram; `recorded_signals_factory` is invoked with the freshly-built diagram and must return the same kind of `{name: OutputPort}` dict that :func:`simulate` accepts. No `vmap` or shared JIT cache: static parameters change the diagram's structure (e.g. state-space dimensions of a :class:`TransferFunction`) and cannot compose with `jax.vmap` by definition. Each element pays a JIT compilation cost. Parameters: | Name | Type | Description | Default | | -------------------------- | --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------- | | `diagram_factory` | `Callable[..., Diagram]` | Callable taking the static-param keyword arguments specified by static_param_grid and returning a :class:Diagram. | *required* | | `t_span` | `tuple[float, float]` | (t_start, t_stop) — same for every element. | *required* | | `static_param_grid` | `dict[str, Sequence[Any]]` | Mapping parameter name -> sequence of values. Every list must have the same length under mode="zip"; or any lengths under mode="product" (cartesian product). | *required* | | `options` | `SimulatorOptions` | :class:SimulatorOptions. max_major_steps must be set if using math_backend="jax". | *required* | | `recorded_signals_factory` | `Callable[[Diagram], dict]` | Callable (diagram) -> {name: OutputPort}. Invoked once per grid element with the freshly-built diagram. | *required* | | `results_options` | \`ResultsOptions | None\` | Optional :class:ResultsOptions passed to :func:simulate. | | `mode` | `str` | "zip" (default — pair lists element-wise) or "product" (cartesian product over all keys). | `'zip'` | Returns: | Type | Description | | ------------------------ | ---------------------------------------------------------------- | | `BatchSimulationResults` | class:BatchSimulationResults with outputs[name].shape == (N, T) | | `BatchSimulationResults` | and time.shape == (T,) where N is the number of grid elements | | `BatchSimulationResults` | and T is the time-vector length of the first run (other runs are | | `BatchSimulationResults` | linearly interpolated onto this grid). The contexts attribute is | | `BatchSimulationResults` | attached to the returned object as a list of per-element final | | `BatchSimulationResults` | contexts. | Raises: | Type | Description | | ------------ | ------------------------------------------------------------------------------ | | `ValueError` | empty grid, mismatched zip lengths, unknown mode, or missing required options. | | `TypeError` | diagram_factory did not return a :class:Diagram. | ### `simulate_variant_sweep(diagram, t_span, *, param_batches=None, options=None, recorded_signals=None, results_options=None, use_vmap=False)` Sweep every variant configuration of `diagram`; for each, optionally sweep a parameter batch. Parameters: | Name | Type | Description | Default | | ------------------ | ---------------------------------------------- | ------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `diagram` | `Diagram` | A built :class:Diagram containing one or more :class:~jaxonomy.framework.variants.Variant nodes. | *required* | | `t_span` | `tuple[float, float]` | (t_start, t_stop) forwarded to the per-variant call. | *required* | | `param_batches` | \`dict[str, Any] | None\` | Optional dot-path → (N,)-shaped array dict forwarded to :func:simulate_batch once per variant. When None, a single :func:simulate is run per variant (equivalent to N=1 but without the batch axis). | | `options` | \`SimulatorOptions | None\` | :class:SimulatorOptions forwarded per variant. | | `recorded_signals` | \`Callable\[[Diagram], dict[str, OutputPort]\] | dict[str, OutputPort] | None\` | | `results_options` | \`ResultsOptions | None\` | Forwarded to :func:simulate_batch. | | `use_vmap` | `bool` | Forwarded to :func:simulate_batch when param_batches is supplied; ignored otherwise. | `False` | Returns: | Name | Type | Description | | ---- | ------------------------------------------------------------- | --------------------- | | | \`dict\[tuple\[tuple[str, Any], ...\], BatchSimulationResults | SimulationResults\]\` | | | \`dict\[tuple\[tuple[str, Any], ...\], BatchSimulationResults | SimulationResults\]\` | | | \`dict\[tuple\[tuple[str, Any], ...\], BatchSimulationResults | SimulationResults\]\` | | `or` | \`dict\[tuple\[tuple[str, Any], ...\], BatchSimulationResults | SimulationResults\]\` | Example .. code-block:: python ``` results = simulate_variant_sweep( diagram, t_span=(0.0, 1.0), param_batches={"plant.gain": jnp.linspace(0.5, 2.0, 8)}, recorded_signals=lambda diag: { "y": diag["plant"].output_ports[0], }, options=opts, ) for cfg, batch_results in results.items(): print(dict(cfg), batch_results.outputs["y"].shape) ``` Notes Each variant configuration triggers an independent JIT compile of the simulator. For a sweep over `V` variants and `N` parameter batches the cost is `V` compiles + `V * N` simulations (with the parameter axis vectorised inside each variant). Variant-axis vmap is genuinely not possible because the pytree shape is not stable across configurations — see the module docstring. ### `simulate_with_event_time_grad(diagram, ctx, t_span, params, event_index, guard_fn, ode_rhs_fn, state_at_event_fn, options=None, *, sim_runner=None, eps=1e-30)` Differentiable wrapper around `simulate` for event-time gradients. Returns the scalar firing time `t_event` of the FIRST recorded firing of `event_index` and registers a `jax.custom_vjp` rule that uses the implicit-function theorem (T-125 phase 1) for the reverse-mode gradient. As a consequence:: ``` jax.grad(simulate_with_event_time_grad)(diagram, ctx, t_span, params, event_index, guard_fn, ode_rhs_fn, state_at_event_fn) ``` yields `∂t_event/∂params` without the caller having to invoke :func:`event_time_gradient` manually. Parameters: | Name | Type | Description | Default | | ------------------- | -------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `diagram` | | SystemBase passed straight to :func:simulate. | *required* | | `ctx` | | ContextBase passed straight to :func:simulate. The caller is responsible for injecting params into ctx (e.g. via ctx.with_parameter(...)) so the forward sim uses the requested parameter values. | *required* | | `t_span` | | (t0, t1) tuple passed to :func:simulate. | *required* | | `params` | | Parameter PyTree to differentiate with respect to. Same semantics as :func:event_time_gradient — the wrapper does not modify ctx from this value; it is used only for the backward rule. | *required* | | `event_index` | `int` | Integer event slot whose firing time is returned. | *required* | | `guard_fn` | `Callable[[float, Any, Any], ndarray]` | (t, state, params) -> scalar — zero-crossing guard used by the implicit-function backward rule. | *required* | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt — continuous RHS evaluated at the event boundary. | *required* | | `state_at_event_fn` | \`Callable\[[Any], Any\] | Any\` | Either * (t_e, params) -> state — preferred signature, matches :func:event_times_gradient. The wrapper passes the recorded t_e as a concrete Python float so the implicit-function-theorem chain rule sees a non-trivial ∂x_e/∂p (in particular, y(t_e_fixed, h0) = h0 - g t_e²/2 has ∂/∂h0 = 1 even though y(t_e(h0), h0) ≡ 0). * params -> state — single-arg form, identical to the one accepted by :func:event_time_gradient. Useful when the caller has already bound t_e into a closure. The wrapper auto-detects which form was passed by argument count. See :func:event_time_gradient for the full contract on the constant-state-PyTree form. | | `options` | | Optional :class:SimulatorOptions. The wrapper forwards a copy with record_event_times=True to :func:simulate; options is None (default) constructs a fresh SimulatorOptions(record_event_times=True). | `None` | | `sim_runner` | \`Callable[..., float] | None\` | (diagram, ctx, t_span, params, event_index, options) -> float — optional override for the forward simulate call. Defaults to the standard :func:simulate path. Tests use this hook to substitute analytic forward trajectories where wiring a full simulate call would be disproportionate. | | `eps` | `float` | Floor for the implicit-function denominator (forwarded to :func:event_time_gradient). | `1e-30` | Returns: | Type | Description | | --------- | ----------------------------------- | | `ndarray` | Scalar jnp.ndarray holding t_event. | Notes Composes with `jax.jit` and `jax.vmap`: the forward pass runs as a `jax.pure_callback` (black-box w.r.t. JAX), and the backward pass uses :func:`event_time_gradient` which is itself JAX-traceable. Default-off byte-equivalence is preserved — the existing :func:`event_time_gradient` and :func:`simulate` are not touched by this wrapper. ### `verify_manifest(actual, expected, *, ignore_fields=None)` Assert that `actual` matches `expected` field-by-field. Convenience wrapper around :func:`compare_manifests` that raises :class:`ManifestMismatch` (an :class:`AssertionError` subclass) if any field drifted. The exception message lists every differing field on its own line; the `.differences` attribute carries the same data structurally for programmatic introspection. Composes naturally with `pytest` (`ManifestMismatch` is an `AssertionError`, so test runners will treat it like any other assertion failure). Parameters: | Name | Type | Description | Default | | --------------- | -------------------- | --------------------------------------------------- | ---------- | | `actual` | `ProvenanceManifest` | the manifest produced by the run being checked. | *required* | | `expected` | `ProvenanceManifest` | the reference manifest. | *required* | | `ignore_fields` | `Optional[set[str]]` | see :func:compare_manifests; default {"timestamp"}. | `None` | Raises: | Type | Description | | ------------------ | ------------------------------ | | `ManifestMismatch` | if any compared field differs. | ### `vmap_event_time_gradient(guard_fn, ode_rhs_fn, t_event_array, state_at_event_fn, params_batch, *, eps=1e-30, use_python_loop=False)` Vectorised event-time gradient over a batch of parameter samples. For `N` samples, computes `∂t_event/∂params` for each in turn and stacks the results along the leading axis — the same shape contract Monte-Carlo / Sobol workflows expect from :func:`simulate_batch`. Parameters: | Name | Type | Description | Default | | ------------------- | -------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `guard_fn` | `Callable[[float, Any, Any], ndarray]` | (t, state, params) -> scalar — zero-crossing guard. Same contract as :func:event_time_gradient; shared across the batch. | *required* | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt — RHS at the event boundary; shared across the batch. | *required* | | `t_event_array` | `ndarray` | (N,) array of per-sample firing instants. Treated as a constant w.r.t. the differentiation parameter inside the wrapper (jax.lax.stop_gradient) so the implicit-function chain rule is taken at the recorded instant — same convention as :func:simulate_with_event_time_grad. | *required* | | `state_at_event_fn` | `Callable[[Any, Any], Any]` | (t_e, params) -> state — reconstructs the trajectory state at firing time t_e parametrised by a single-sample params slice. Identical signature to the one accepted by :func:event_times_gradient; the wrapper composes it with each t_event_array[i] and the i-th slice of params_batch under jax.vmap. | *required* | | `params_batch` | `Any` | Batched parameter PyTree. All leaves must share a leading axis of length N matching t_event_array. May be a scalar batch (shape (N,) ndarray), a vector batch (shape (N, n_p)), or a PyTree thereof. | *required* | | `eps` | `float` | Forwarded to :func:event_time_gradient — denominator floor for grazing crossings. | `1e-30` | | `use_python_loop` | `bool` | When True, iterate explicitly over the sample axis instead of using jax.vmap. Slower but byte-identical; useful as a fallback if vmap composition ever breaks (e.g. under future JAX versions where a closure inside :func:event_time_gradient becomes non-vmap-friendly). | `False` | Returns: | Type | Description | | ----- | -------------------------------------------------------- | | `Any` | Per-sample gradients with the same leading axis as | | `Any` | params_batch. PyTree structure of each sample's gradient | | `Any` | matches the single-sample :func:event_time_gradient. | Notes Default-off: the wrapper is purely additive and does not modify the simulator path. Composes cleanly with `jax.jit` and downstream `jax.grad` of a scalar cost over the batch axis. ### `vmap_event_times_gradient(results, params_batch, guards, ode_rhs_fn, state_at_event_fn, *, event_indices=None, eps=1e-30, use_python_loop=False)` Cross-product of multi-event + batched-parameter event-time gradient. For each event index recorded in `results.event_times`, computes the implicit-function-theorem gradient `∂t_event/∂params` at every (sample, firing) pair and returns the result keyed by event index. Output contract:: ``` {event_index: gradient_batch} ``` where `gradient_batch` has leading axes `(N, n_firings, ...)` for array-valued `params_batch` leaves and is itself a PyTree mirroring the structure of `params_batch` for nested batches. `N` is the sample-axis length (shared across all batch leaves); `n_firings` is the per-event firing count read from `results.event_times[idx]`. Parameters: | Name | Type | Description | Default | | ------------------- | ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `results` | `Any` | A :class:SimulationResults whose event_times is populated (i.e., the simulation was run with SimulatorOptions(record_event_times=True)). Same requirement as :func:event_times_gradient. A results whose event_times is None raises ValueError with the remediation hint. | *required* | | `params_batch` | `Any` | Batched parameter PyTree. All leaves must share a leading axis of length N. Same contract as :func:vmap_event_time_gradient. | *required* | | `guards` | `Any` | Either a single guard callable (t, state, params) -> scalar applied to every event, or a mapping {event_index: guard_fn}. Same semantics as :func:event_times_gradient. | *required* | | `ode_rhs_fn` | `Callable[[float, Any, Any], Any]` | (t, state, params) -> dstate/dt — shared across firings and samples. | *required* | | `state_at_event_fn` | `Callable[[float, Any], Any]` | (t_e, params) -> state — reconstructs the trajectory state at firing time t_e parametrised by a single-sample params slice. Same signature as :func:event_times_gradient and :func:vmap_event_time_gradient. | *required* | | `event_indices` | `Any` | Optional iterable of event indices to compute gradients for. When None, every recorded event index is processed. Indices not present in results.event_times raise KeyError. | `None` | | `eps` | `float` | Forwarded to :func:event_time_gradient — denominator floor for grazing crossings. | `1e-30` | | `use_python_loop` | `bool` | When True, iterate explicitly over both the firing and sample axes in Python. Slower but byte-identical; honest fallback when vmap composition is invasive. | `False` | Returns: | Type | Description | | ------ | ------------------------------------------------------------ | | `dict` | {event_index: gradient_batch} — one entry per processed | | `dict` | event index. For each entry, leaves carry leading axes | | `dict` | (N, n_firings, ...). Empty firing lists yield a structurally | | `dict` | correct (N, 0, ...) leading-axis pair. | Notes Default-off: purely additive. Composes with `jax.jit`. The firing times read from `results.event_times` are treated as constants w.r.t. `params_batch` (the implicit-function theorem is applied at the recorded instants — same convention as :func:`vmap_event_time_gradient` and :func:`simulate_with_event_time_grad`). # Optimization ## `jaxonomy.optimization` ### `AutoTuner` PID autotuning (without a measurement filter) with constraints in the frequency domain. Supports only SISO systems. Supports only continuous-time plants (TODO: extend to discrete-time systems) Parameters: | Name | Type | Description | Default | | ------------------------ | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------ | | `plant` | | LeafSystem or a Diagram. If plant is not an LTISystem, operating points x_op and u_op must be provided for linearization. | *required* | | `n` | | int, optional Filter coefficient for the continuous-time PID controller | `100` | | `sim_time` | | float, optional Simulation time for computation of the error metric | `2.0` | | `metric` | | str, optional Error metric to be minimized. Options are "IAE" and "IE" "IAE": Integral of the absolute error "IE": Integral of the error | `'IAE'` | | `x_op` | | np.ndarray, optional Operating point of state vector for linearization | `None` | | `u_op` | | np.ndarray, optional Operating point of control vector for linearization | `None` | | `pid_gains_0` | | list or Array, optional Initial guess for PID gains [kp, ki, kd] | `[1.0, 10.0, 0.1]` | | `pid_gains_upper_bounds` | | list or Array, optional Upper bounds for PID gains [kp, ki, kd]. Lower bounds are set to 0 | `None` | | `Ms` | | float, optional Maximum sensitivity | `100.0` | | `Mt` | | float, optional Maximum complementary sensitivity | `100.0` | | `add_filter` | | bool, optional Add measurement filter (currently not implemented) | `False` | | `method` | | str, optional The method for optimization. Available options are: - "scipy-slsqp" - "scipy-cobyla" - "scipy-trust-constr" - "ipopt" - "nlopt-slsqp" - "nlopt-cobyla" - "nlopt-ld_mma" - "nlopt-isres" - "nlopt-ags" - "nlopt-direct" | `'scipy-slsqp'` | Notes: The utilities `plot_freq_response`, `plot_time_response`, and `plot_freq_and_time_responses` can be used to visualize the frequency and time responses of the closed-loop system. Post initialization the `tune` method should be called to obtain the optimal PID gains. See `notebooks/opt_framework/pid_autotuning.ipynb` for an example. #### `circle_constraint_(kp, ki, kd, omega, c, r)` Deprecated: this is needed for `self.constraints_` which is deprecated and replaced by `self.constraints`. #### `constraints_(pid_params)` Deprecated: replaced by `self.constraints` ### `CompositeTransform` Bases: `Transform` A composite transformation that applies a list of transformations in sequence. ### `ConfidenceIntervalResult` Confidence intervals and covariance matrix from the Laplace approximation. All matrix/array attributes are plain `numpy.ndarray` for easy inspection and serialisation. ##### Attributes param_names : list[str] Flat parameter names (array params expanded to `"theta[0]"`, etc.). opt_params : dict Optimised parameter values in the **original** (un-transformed) space. covariance : ndarray, shape (n, n) Estimated parameter covariance matrix. correlation : ndarray, shape (n, n) Correlation matrix (covariance normalised by marginal standard deviations). standard_errors : ndarray, shape (n,) Marginal standard deviations `sqrt(diag(covariance))`. confidence_intervals : dict\[str, tuple[float, float]\] Per-parameter `(lower, upper)` bounds in the **original** space. Keys match `param_names`. confidence_level : float Nominal confidence level (e.g. `0.95` for 95 %). z_score : float Standard-normal quantile corresponding to `confidence_level`. hessian : ndarray, shape (n, n) Hessian of the objective evaluated at the optimum, in the (possibly transformed) optimisation space. hessian_eigenvalues : ndarray, shape (n,) Eigenvalues of the Hessian (ascending). hessian_condition_number : float Ratio max|λ| / min|λ|. Large values (> 1 000) signal near-collinear parameters or an ill-conditioned problem. is_positive_definite : bool `True` when the Hessian was positive definite at the supplied point (necessary condition for a true local minimum). residual_variance : float or None Residual variance `σ²` used to scale the covariance. `None` when `n_data` was not provided (pure MLE / default). n_data : int or None Number of observations used (for least-squares scaling). objective_value : float Loss at the optimum. hessian_method : str How the Hessian was computed: `"AD"` (automatic differentiation), `"FD"` (finite differences), `"provided"`, or `"failed"`. message : str Any warnings raised during computation (empty when all is well). #### `contains(param_name, value)` Return `True` when *value* lies within the CI for *param_name*. #### `interval(param_name)` Return `(lower, upper)` for a single parameter by name. Raises `KeyError` if the name is not found. For array parameters use the expanded name, e.g. `ci.interval("theta[0]")`. #### `summary()` Return a formatted human-readable summary table. ### `DistributionConfig` Structure of attributes for specifying distributions for stochastic variables ### `Evosax` Bases: `Optimizer` Population based global optimizers from Evosax. Parameters: | Name | Type | Description | Default | | ------------------- | --------------- | ---------------------------------------------------------------------------- | -------------------------------------- | | `optimizable` | `Optimizable` | The optimizable object. | *required* | | `opt_method` | `str` | The optimization method to use. See evosax.Strategies for available methods. | `'CMA_ES'` | | `opt_method_config` | `dict` | Configuration for the optimization method. | `None` | | `pop_size` | `int` | The population size. | `10` | | `num_generations` | `int` | The number of generations. | `100` | | `print_every` | `int` | Print progress every print_every generations. | `1` | | `metrics_writer` | \`MetricsWriter | None\` | Optional CSV file to write metrics to. | | `seed` | `int` | The random seed. | `None` | #### `optimize()` Run optimization ### `IPOPT` Bases: `Optimizer` Interior Point Optimizer (IPOPT) for optimization of the objective function with optional constraints and bounds. Parameters: | Name | Type | Description | Default | | ------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------- | | `optimizable` | `Optimizable` | The optimizable object. | *required* | | `options` | `dict` | Options forwarded to cyipopt.minimize_ipopt. See https://coin-or.github.io/Ipopt/OPTIONS.html for the full list. Commonly used keys: maxiter (int, default 3000) Maximum number of IPOPT iterations. disp (int, default 5) Verbosity level (0 = silent). tol (float) Convergence tolerance on the NLP optimality conditions. acceptable_tol (float) Looser acceptable-solution tolerance. | `{'disp': 5}` | #### `optimize()` Run optimisation and return an :class:`~jaxonomy.optimization.OptimizationResult`. Gradients of the objective and constraint Jacobians are computed with JAX automatic differentiation (`jax.grad` / `jax.jacrev`). **Hessian strategy** — JAX's `jax.hessian` requires forward-mode automatic differentiation (`jacfwd`) through the gradient, but the jaxonomy ODE solver uses `custom_vjp` which only supports reverse-mode. Attempting to compute `jax.hessian` of a simulation objective therefore raises a runtime error. IPOPT is instead configured with `hessian_approximation = "limited-memory"` (L-BFGS approximation) which only requires first-order gradient information and converges super-linearly. For problems where you *know* the objective is twice-differentiable and do not use the jaxonomy ODE integrator you can override this by passing `options={"hessian_approximation": "exact", ...}` and providing `hess` via the `_hess_fn` constructor argument. ### `IdentityTransform` Bases: `Transform` A transformation that does nothing: `y = x`. ### `LogTransform` Bases: `Transform` A transformation that applies the natural logarithm to the values of the parameters. `y = log(x)`. ### `LogitTransform` Bases: `Transform` The logit transformation, defined as: `y = log(x / (1 - x))` ### `MultiStart` Multi-start wrapper for any jaxonomy optimizer. Runs `n_starts` optimizations from different initial points and returns all results as well as the best one (lowest `final_loss`). ##### Parameters optimizable : Optimizable The problem to optimize. Must be a jaxonomy `Optimizable` instance. optimizer_factory : Callable\[[Optimizable], Optimizer\] A *factory* function that takes an `Optimizable` (potentially with different initial parameters) and returns a ready-to-run optimizer. Example:: ``` factory = lambda opt: Scipy(opt, "L-BFGS-B", opt_method_config={"maxiter": 40}, use_autodiff_grad=True) ms = MultiStart(optimizable, factory, n_starts=8, seed=0) result = ms.run() ``` int Number of random restarts (default 10). init_sampler : Callable or None Custom sampling function with signature `(n_starts: int, params_0_flat: np.ndarray) -> np.ndarray` returning an array of shape `(n_starts, n_params)`. Row 0 is always replaced with the original `params_0_flat` when `include_initial=True`. If `None` (default), uniform sampling around `params_0` is used. sample_scale : float Scale factor for the default uniform sampler. The search window for each parameter is `[p0 ± sample_scale * max(|p0|, 1)]` (default 1.0). seed : int or None Random seed for reproducibility. include_initial : bool When `True` (default), the first start always uses the original `params_0`, regardless of the sampler output. #### `results` Results from the last :meth:`run` call (empty before first run). #### `run()` Execute all starts sequentially and return a :class:`MultiStartResult`. Each start clones the optimizable with a new `params_0_flat`, calls `optimizer_factory(clone)` to get a fresh optimizer, and runs `optimizer.optimize()`. Failed starts (exceptions) are recorded as unsuccessful `OptimizationResult` entries with `success=False`. ###### Returns MultiStartResult ### `MultiStartResult` Results from a multi-start optimization run. Attributes: | Name | Type | Description | | ------------------ | -------------------------- | ------------------------------------------------------------ | | `results` | `list[OptimizationResult]` | All OptimizationResult objects — one per start. | | `best_result` | `OptimizationResult` | The result with the lowest final_loss among successful runs. | | `best_start_index` | `int` | Index into results of the best run. | | `n_starts` | `int` | Total number of starts attempted. | | `n_successful` | `int` | Number of starts that reported success=True. | ### `NLopt` Bases: `Optimizer` Optimizers using the NLopt library. Parameters: | Name | Type | Description | Default | | ------------- | ------------- | ---------------------------------------------- | ---------- | | `optimizable` | `Optimizable` | The optimizable object. | *required* | | `opt_method` | `str` | The optimization method to use. | *required* | | `ftol_rel` | `float` | Relative tolerance on function value. | `1e-06` | | `ftol_abs` | `float` | Absolute tolerance on function value. | `1e-06` | | `xtol_rel` | `float` | Relative tolerance on optimization parameters. | `1e-06` | | `xtol_abs` | `float` | Absolute tolerance on optimization parameters. | `1e-06` | | `cons_tol` | `float` | Tolerance on constraints. | `1e-06` | | `maxeval` | `int` | Maximum number of function evaluations. | `500` | | `maxtime` | `float` | Maximum time in seconds. | `0` | #### `optimize()` Run optimization ### `NegativeNegativeLogTransform` Bases: `Transform` A transformation that applies the negative of the natural logarithm of the negative of the values of the parameters. `y = -log(-x)` ### `NormalizeTransform` Bases: `Transform` A transformation that normalizes the values of the parameters to the range [0, 1]. `y = (x - min) / (max - min)` Paramteters: - params_min: dict with the minimum values for each parameter. - params_max: dict with the maximum values for each parameter. ### `Optax` Bases: `Optimizer` Optax optimizer without support for stochastic variables. Paramters optimizable (Optimizable): The optimizable object. opt_method (str): The optimization method to use. learning_rate (float): The learning rate. opt_method_config (dict): Configuration for the optimization method. num_epochs (int): The number of epochs. clip_range (tuple): The range to clip the gradients. print_every (int): Print progress every `print_every` epochs. metrics_writer (MetricsWriter|None): Optional CSV file to write metrics to. #### `optimize()` Run optimization #### `step(params, opt_state)` Take a single optimization step ### `OptaxWithStochasticVars` Bases: `Optimizer` Optax optimizer with support for stochastic variables. Parameters: | Name | Type | Description | Default | | ------------------- | ------------------------------- | ------------------------------------------ | -------------------------------------- | | `optimizable` | `OptimizableWithStochasticVars` | The optimizable object. | *required* | | `opt_method` | `str` | The optimization method to use. | *required* | | `learning_rate` | `float` | The learning rate. | *required* | | `opt_method_config` | `dict` | Configuration for the optimization method. | *required* | | `num_epochs` | `int` | The number of epochs. | `100` | | `batch_size` | `int` | The batch size. | `1` | | `num_batches` | `int` | The number of batches. | `1` | | `clip_range` | `tuple` | The range to clip the gradients. | `None` | | `print_every` | `int` | Print progress every print_every epochs. | `None` | | `metrics_writer` | \`MetricsWriter | None\` | Optional CSV file to write metrics to. | #### `batched_objective_flat(params, stochastic_vars_batch_flat)` Mean of the objective function over a batch #### `optimize()` Run optimization #### `step(params, opt_state, stochastic_vars_batch)` Take a single optimization step over one batch ### `Optimizable` Bases: `OptimizableBase` Base class for all optimizables with no stochastic variables. For parameters, see `OptimizableBase`. The abstract method `prepare_context` should update the context to incorporate the optimization parameters. This classs creates methods for evaluation of the objective and constraints from the concrete implementation of the abstract methods. This class also creates methods for batched evaluation of the objective and constraints, which are useful for optimizers that can work with batches (eg. Optax), and population-based optimizers. #### `constraints(params)` Constraints function for optimization with dict parameters input #### `constraints_flat(params)` Constraints function for optimization with flattened parameters input #### `objective(params)` Objective function for optimization with dict parameters input #### `objective_flat(params)` Objective function for optimization with flattened parameters input #### `prepare_context(context, params)` Model-specific updates to incorporate the sample data and parameters. Return the updated context. #### `run_simulation(params)` Run simulation and return final results context. ### `OptimizableWithStochasticVars` Bases: `OptimizableBase` Base class for all optimizables with stochastic variables. This is designed only for Optax optimizers and without constraints. Other optimizers are unlikely to work well with stochastic variables. This class is similar to `Optimizable` with the key difference that both `params` and `vars` (stochastic variables) need to be updated as opposed to `params` alone Parameters: | Name | Type | Description | Default | | -------------------------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------- | ------- | | `vars_0` | | dict Initial stochastic variable values. If not provided, the stochastic_vars method will be used to extract these from the base context. | `None` | | `distribution_config_vars` | | DistributionConfig Configuration for stochastic variables. If not provided, standard normal distribution is used. | `None` | #### `generate_batches(data, num_batches, batch_size)` Given all samples `data`, generate `num_batches` random batches of size `batch_size` each #### `objective(params, vars)` Objective function for optimization with dict parameters and vars input #### `objective_flat(params, vars)` Objective function for optimization with flattened parameters and vars input #### `prepare_context(context, params, vars)` Model-specific updates to incorporate the parameters and stochastic vars. Return the updated context. #### `run_simulation(params, vars)` Run simulation and return final results context. #### `sample_random_vars(num_samples)` Generate random samples of the stochastic variables #### `stochastic_vars(context)` Extract stochastic `vars` from the context. These should be in the form of a dict of Pytrees. ### `OptimizationResult` Unified result returned by all jaxonomy optimizers. Supports dict-like access (`result["param"]`) for backward compatibility with code that treated the old return value as a plain parameter dict. Attributes: | Name | Type | Description | | -------------- | ---------------- | --------------------------------------------------------------------------------------------------------- | | `params` | `dict[str, Any]` | dict mapping parameter name → optimized value (same as the dict that optimizers used to return directly). | | `success` | `bool` | True if the optimizer reported convergence. | | `nit` | `int` | Number of iterations (or epochs / generations). | | `nfev` | `int` | Number of objective-function evaluations. | | `message` | `str` | Human-readable status message from the optimizer. | | `final_loss` | \`float | None\` | | `loss_history` | `list[float]` | Sequence of objective values recorded during optimization (one per epoch / generation). | ### `RLEnv` Bases: `Env` Base class for reinforcement learning environments in Jaxonomy. #### `get_done(pipeline_state, obs)` Return a boolean indicating whether the episode is done. #### `get_reward(pipeline_state, obs, act)` Return the reward for the current state and observation. #### `randomize(pipeline_state, rng)` Randomize the initial states, parameters, etc. #### `render(trajectory, height=240, width=320, camera=None)` Render the trajectory ### `Scipy` Bases: `Optimizer` Scipy/JAX-scipy optimizers. Parameters: | Name | Type | Description | Default | | ------------------- | ------------- | ----------------------------------------------------------------------- | ---------- | | `optimizable` | `Optimizable` | The optimizable object. | *required* | | `opt_method` | `str` | The optimization method to use. | *required* | | `tol` | `float` | Tolerance for termination. For detailed control, use opt_method_config. | `None` | | `opt_method_config` | `dict` | Configuration for the optimization method. | `None` | | `use_autodiff_grad` | `bool` | Whether to use autodiff for gradient computation. | `True` | | `use_jax_scipy` | `bool` | Whether to use JAX's version of optimize.minimize. | `False` | #### `optimize()` Run optimization ### `SensitivityResult` Result of a parameter sensitivity / identifiability analysis. All arrays are plain `numpy.ndarray` for easy inspection. ##### Attributes param_names : list[str] Parameter names, in the same order as the flat parameter vector. params_0 : dict[str, Any] The parameter values at which the analysis was performed. objective_value : float Objective value at `params_0`. gradients : ndarray, shape (n_params,) Gradient of the objective w.r.t. each parameter. normalized_sensitivity : ndarray, shape (n_params,) `|p_i * ∂L/∂p_i|` — relative sensitivity. Dimensionless and comparable across parameters with different scales. hessian : ndarray, shape (n_params, n_params) Hessian of the objective (FIM approximation). `NaN`-filled when `compute_hessian=False`. hessian_diagonal : ndarray, shape (n_params,) Diagonal of the Hessian. eigenvalues : ndarray, shape (n_params,) Eigenvalues of the Hessian (ascending). condition_number : float Ratio of largest to smallest non-negligible eigenvalue. Large values (> 1e6) indicate near-collinear parameters. unidentifiable_params : list[str] Parameter names whose normalised sensitivity is below `sensitivity_threshold * max_sensitivity`. sensitivity_threshold : float Relative threshold used to flag unidentifiable parameters. #### `summary()` Return a human-readable summary table. ### `Trainer` Base class for optimizing model parameters via simulation. Should probably get a more descriptive name once we're doing other kinds of training... #### `evaluate_cost(context)` Model-specific cost function, evaluated on final context #### `make_forward(start_time, stop_time)` Create a generic forward pass through the simulation, returning loss #### `make_loss_fn(forward, params)` Create a loss function based on a forward pass of the simulation `params` here can be any PyTree - it will get flattened to a single array #### `optimizable_parameters(context)` Extract optimizable model-specific parameters from the context. These should be in the form of a PyTree (e.g. tuple, dict, array, etc) and should be the first arguments to `prepare_context`. #### `prepare_context(context, *data, key=None)` Model-specific updates to incorporate the sample data and parameters. `data` should be the combination of the output of `optimizable_parameters` along with all the per-simulation "training data". Parameters will update once per epoch, and training data will update once per sample. #### `train(training_data, sim_start_time, sim_stop_time, epochs=100, key=None, params=None, opt_state=None)` Run the optimization loop over the training data ### `Transform` Bases: `ABC` Base class for transformations. #### `inverse_transform(params)` Take transformed parameters dict {key:value} and output a dict with identical keys but inverse-transformed `values`. #### `transform(params)` Take original parameters dict {key:value} and output a dict with identical keys but transformed `values`. ### `TuningResult` Result of a `tune_parameters` call. Attributes: | Name | Type | Description | | ----------- | ------------------------------ | --------------------------------------------------------------------------------------------------------- | | `params` | `Dict[str, Array]` | Optimal parameter values as a dict {name: jax.Array}. | | `objective` | `float` | Final objective value (scalar). | | `history` | `list` | List of (iteration, objective) tuples if tracking enabled. | | `success` | `bool` | True if the optimizer reported successful convergence. | | `message` | `str` | Human-readable status from the underlying optimizer. | | `raw` | `Optional[OptimizationResult]` | Underlying OptimizationResult from the optimizer framework (for inspection of optimizer-specific fields). | ### `compute_confidence_intervals(optimizable, opt_params, confidence_level=0.95, n_data=None, hessian=None, eps_fd=0.0001, regularize=True)` Compute Wald-type confidence intervals for optimised parameters. Uses the **Laplace approximation**: the parameter posterior is approximated as a Gaussian centred at the optimum θ\* with covariance `H⁻¹`, where `H = ∇²L(θ*)` is the Hessian of the loss. ##### Parameters optimizable : Optimizable The jaxonomy optimizable whose `objective_flat` is used. opt_params : OptimizationResult | dict | array-like Optimised parameters at which to evaluate the Hessian. Can be: ``` * An :class:`~jaxonomy.optimization.OptimizationResult` returned by any jaxonomy optimizer — the ``params`` dict is extracted and flattened automatically. * A plain ``dict`` mapping parameter names to values. * A flat 1-D array matching ``optimizable.params_0_flat``. ``` confidence_level : float Nominal confidence level (default `0.95` for 95 % CIs). n_data : int or None Number of observations. When provided, the covariance is scaled by the residual variance estimate ``` ``σ² = 2 · L(θ*) / max(n_data − n_params, 1)`` This is appropriate for **sum-of-squares objectives** ``L = ½ Σ rᵢ²``. For maximum-likelihood objectives leave ``None``. ``` hessian : ndarray or None Pre-computed Hessian matrix (e.g. from :func:`compute_sensitivity`). When `None` (default) the Hessian is computed automatically using JAX AD (with a finite-difference fallback for ODE-based objectives). eps_fd : float Step size used for the finite-difference Hessian fallback (default `1e-4`). Ignored when `hessian` is provided or when AD succeeds. regularize : bool When `True` (default), negative eigenvalues of the Hessian are clipped to a small positive value before inversion. This makes the covariance well-defined even when the supplied point is not a true local minimum. A warning is recorded in `result.message`. ##### Returns ConfidenceIntervalResult Dataclass containing the covariance matrix, standard errors, and per-parameter confidence intervals in the **original** (physical) parameter space. ##### Notes **Parameter transformations**: if the `Optimizable` uses a `transformation` (e.g. :class:`LogTransform`), the Hessian is computed in the *transformed* space and the resulting CI bounds are mapped back to the original space via `transformation.inverse_transform`. **Validity**: the Laplace approximation requires the objective to be smooth and the optimum to be a true interior local minimum (positive- definite Hessian). If `is_positive_definite` is `False` in the result, the CIs are computed but should be treated with caution. **Profile likelihood**: the Laplace approximation is a first-order Gaussian approximation. For strongly nonlinear models or highly non-Gaussian posteriors, profile likelihood confidence intervals are more accurate but require repeated re-optimisation. ##### Examples > > > from jaxonomy.optimization import Scipy, compute_confidence_intervals opt = Scipy(my_opt, method="L-BFGS-B", use_autodiff_grad=True) result = opt.optimize() ci = compute_confidence_intervals(my_opt, result, confidence_level=0.95) print(ci.summary()) lo, hi = ci.interval("c") ### `compute_sensitivity(optimizable, params_0_flat=None, sensitivity_threshold=0.001, compute_hessian=True)` Compute gradient-based parameter sensitivity at a given operating point. Uses JAX automatic differentiation — no finite differences, no extra simulations beyond two JIT-compiled evaluations (gradient + optional Hessian). ##### Parameters optimizable : Optimizable The jaxonomy optimizable whose `objective_flat` is differentiated. params_0_flat : array-like or None Flat parameter vector to evaluate at. Defaults to `optimizable.params_0_flat`. sensitivity_threshold : float Relative threshold (0–1) for flagging parameters as low-sensitivity. A parameter is flagged when its normalised sensitivity is less than `sensitivity_threshold × max(all normalised sensitivities)`. Default `1e-3`. compute_hessian : bool Whether to compute the full Hessian / FIM. Can be expensive for many parameters (O(n²) simulations). Default `True`. ##### Returns SensitivityResult ### `implicit_solver(solver, residual, linear_solve=None)` Make an iterative solver reverse-mode differentiable via the IFT. Parameters: | Name | Type | Description | Default | | -------------- | -------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `solver` | `Callable` | solver(theta) -> x_star — any (jit-compatible) function returning a solution as a flat array (shape (n,) or scalar). Internal control flow is unrestricted; it is never differentiated. theta may be any pytree. | *required* | | `residual` | `Callable` | residual(x, theta) -> r with r the same shape as x and residual(solver(theta), theta) ≈ 0. Must be JAX-differentiable — this is the equation the solution satisfies, used to construct both sides of the IFT. | *required* | | `linear_solve` | `Optional[Callable]` | optional linear_solve(A, b) -> w used for the adjoint system (∂g/∂x)ᵀ w = b. Defaults to the dense :func:jnp.linalg.solve — right for the small systems that appear inside dynamics callbacks (constraint dimensions of tens). Supply a matrix-free solver (e.g. CG) for large n. | `None` | Returns: | Type | Description | | ---- | ----------------------------------------------------------- | | | A function wrapped(theta) -> x_star that is byte-equivalent | | | to solver in the forward pass and reverse-differentiable. | Example — an implicit velocity law solved by Newton iteration:: ``` def solve_v(theta): # while_loop inside def newton(v): g = v + jnp.tanh(theta * v) - 1.0 dg = 1.0 + theta / jnp.cosh(theta * v) ** 2 return v - g / dg def cond(carry): v, i = carry return (jnp.abs(v + jnp.tanh(theta * v) - 1.0) > 1e-12) & (i < 50) def body(carry): v, i = carry return newton(v), i + 1 v, _ = jax.lax.while_loop(cond, body, (jnp.asarray(0.5), 0)) return v def residual(v, theta): return v + jnp.tanh(theta * v) - 1.0 solve_v_diff = implicit_solver(solve_v, residual) jax.grad(solve_v_diff)(0.3) # works; matches FD ``` ### `ise_objective(builder, signal_port, reference_port=None, weight=1.0, initial_cost=0.0, name='ise')` Add blocks to compute the **Integral of Squared Error** (ISE). .. math:: ``` J = \int_0^T w \, \| \text{signal}(t) - \text{reference}(t) \|^2 \, dt ``` When `reference_port` is `None` the reference is implicitly zero, so the objective is :math:`\int_0^T w \, \|\text{signal}(t)\|^2 \, dt`. The function adds the following blocks to *builder*: - (optional) :class:`~jaxonomy.library.Adder` computing `signal − reference` - :class:`~jaxonomy.library.Power` `(2.0)` - :class:`~jaxonomy.library.SumOfElements` (handles both scalar and vector signals transparently) - (optional) :class:`~jaxonomy.library.Gain` if `weight ≠ 1` - :class:`~jaxonomy.library.Integrator` accumulating the cost ##### Parameters builder: The :class:`~jaxonomy.DiagramBuilder` to add blocks to. signal_port: Output port of the signal to penalise. reference_port: Output port of the reference signal. `None` → reference is 0. weight: Scalar multiplier applied to the squared norm before integration. For per-component or matrix weighting use :func:`lqr_objective`. initial_cost: Initial value of the accumulating integrator (default `0.0`). name: Name prefix for the added blocks. ##### Returns OutputPort Scalar port whose value at the end of simulation equals *J*. ##### Examples Minimise oscillation energy of a spring-mass system:: ``` obj = ise_objective(b, x.output_ports[0]) # ∫ x² dt # later: return obj.eval(ctx) ``` Multi-signal ISE with a shared reference of zero:: ``` cost_x = ise_objective(b, x.output_ports[0], name="ise_x") cost_v = ise_objective(b, v.output_ports[0], name="ise_v") total = weighted_sum(b, [cost_x, cost_v], weights=[1.0, 0.5]) ``` ### `lqr_objective(builder, state_port, Q, control_port=None, R=None, initial_cost=0.0, name='lqr')` Add blocks to compute an **LQR-style quadratic cost**. .. math:: ``` J = \int_0^T \bigl( x(t)^\top Q\, x(t) \;+\; u(t)^\top R\, u(t) \bigr)\, dt ``` When `control_port` or `R` is `None` only the state cost :math:`\int x^\top Q x\, dt` is computed. The function adds a single-input :class:`~jaxonomy.library.ReduceBlock` for :math:`x^\top Q x` (and optionally one for :math:`u^\top R u`), an optional :class:`~jaxonomy.library.Adder`, and an :class:`~jaxonomy.library.Integrator`. ##### Parameters builder: The :class:`~jaxonomy.DiagramBuilder` to add blocks to. state_port: Output port of the state vector :math:`x`. Q: Positive semi-definite state weight matrix, shape `(nx, nx)`. control_port: Output port of the control vector :math:`u`. `None` → no control penalty. R: Positive definite control weight matrix, shape `(nu, nu)`. Required when *control_port* is provided. initial_cost: Initial value of the accumulating integrator (default `0.0`). name: Name prefix for the added blocks. ##### Returns OutputPort Scalar port whose value at the end of simulation equals *J*. ##### Examples Pendulum regulation:: ``` # ∫ θ²·Q[0,0] + ω²·Q[1,1] dt (diagonal Q) Q = jnp.diag(jnp.array([10.0, 1.0])) R = jnp.array([[0.1]]) cost = lqr_objective(b, x.output_ports[0], Q, u.output_ports[0], R) ``` ### `tracking_mse(builder, signal_port, t_data, y_data, weight=1.0, interpolation='linear', initial_cost=0.0, name='tracking_mse')` Add blocks to compute the **dataset tracking MSE**. Computes .. math:: ``` J = \int_0^T w \, \| \text{signal}(t) - y_{\text{ref}}(t) \|^2 \, dt ``` where :math:`y_{\text{ref}}(t)` is the reference signal *interpolated* from the dataset `(t_data, y_data)` at every simulation time step. The function wires: 1. :class:`~jaxonomy.library.Clock` → current simulation time 1. :class:`~jaxonomy.library.LookupTable1d` → interpolated reference 1. :func:`ise_objective` → squared error integrator ##### Parameters builder: The :class:`~jaxonomy.DiagramBuilder` to add blocks to. signal_port: Output port of the simulated signal to compare against the data. t_data: 1-D array of reference time points (must be strictly increasing). y_data: Array of reference values. Shape `(N,)` for scalar signals or `(N, ny)` for vector signals. Extrapolation clamps to the nearest endpoint value. weight: Scalar multiplier applied before integration. interpolation: Interpolation method passed to :class:`~jaxonomy.library.LookupTable1d`: `"linear"` (default), `"nearest"`, or `"flat"`. initial_cost: Initial value of the integrator. name: Name prefix for the added blocks. ##### Returns OutputPort Scalar port equal to *J* at the end of simulation. ##### Examples Fit a model to measured step-response data:: ``` import numpy as np t_meas = np.linspace(0, 5, 50) y_meas = 1 - np.exp(-t_meas) # first-order step response cost = tracking_mse(b, plant.output_ports[0], t_meas, y_meas) ``` ### `tune_parameters(diagram, base_context, sim_t_span, params_0, set_params, objective_fn, bounds=None, optimizer='scipy-lbfgs', n_iter=100, learning_rate=0.05, sim_options=None, verbose=True)` Tune scalar parameters of a jaxonomy diagram to minimize an objective. The simulator is differentiated through using JAX autodiff; the gradient of `objective_fn` with respect to each entry of `params_0` is computed automatically, and an optimizer minimizes the objective. Parameters: | Name | Type | Description | Default | | --------------- | ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------- | | `diagram` | | A built jaxonomy diagram. | *required* | | `base_context` | | A Context created from the diagram. The optimizer calls set_params(base_context, params) each iteration to inject the current parameter values, then advances the simulator over sim_t_span, then evaluates objective_fn on the final context. | *required* | | `sim_t_span` | `Tuple[float, float]` | (t0, tf) simulation interval. | *required* | | `params_0` | `Dict[str, Any]` | Dict of initial parameter values. Each value must be a scalar or a JAX-compatible array. | *required* | | `set_params` | `Callable[[Any, Dict[str, Array]], Any]` | Callback (context, params_dict) -> updated_context. Typical implementations write parameter values into LeafSystem parameters via context.with_parameter(...), or modify initial states. | *required* | | `objective_fn` | `Callable[[Any], Array]` | Callback (results_context) -> scalar. The optimizer minimizes this. Must be differentiable through JAX. | *required* | | `bounds` | `Optional[Dict[str, Tuple[float, float]]]` | Optional dict {param_name: (lb, ub)} for box constraints. Only honoured by box-constrained optimizers (l-bfgs-b, slsqp, trust-constr); ignored otherwise. | `None` | | `optimizer` | `str` | Optimizer alias or scipy method name. Defaults to "scipy-lbfgs". See \_OPTIMIZER_ALIASES for shortcuts. | `'scipy-lbfgs'` | | `n_iter` | `int` | Maximum number of optimizer iterations. | `100` | | `learning_rate` | `float` | Learning rate for optax optimizers (ignored by scipy). | `0.05` | | `sim_options` | `Optional[SimulatorOptions]` | Optional SimulatorOptions. If None, a default with autodiff enabled is used. | `None` | | `verbose` | `bool` | If True, log progress to the jaxonomy logger. | `True` | Returns: | Type | Description | | -------------- | ------------------------------------------------------------------ | | `TuningResult` | A TuningResult with optimal parameter values, the final objective, | | `TuningResult` | and a reference to the raw optimizer result. | Notes - **Discrete parameters** (e.g., a horizon length `N`, a state-machine guard threshold) are not differentiable through the simulator and should be left as fixed hyperparameters. If you need to sweep them, wrap `tune_parameters` in an outer loop or grid search. - **Saturation regions** (`jnp.clip`, `jax.lax.cond` with hard switches) have zero gradient. Tuning parameters whose value determines a saturation region may be impossible from a starting point already saturated; consider warm-starting away from the saturation boundary. - **Bounds enforcement**: with `scipy-lbfgs` / `slsqp` / `trust-constr`, bounds are honoured by the solver. With `optax` optimizers, bounds are not enforced; clip parameters yourself in `set_params` if you need them. See also `jaxonomy.optimization.Optimizable` — the lower-level interface this function wraps. Use it directly if you need stochastic variables, constraints, or batched evaluations. ### `weighted_sum(builder, objectives, weights=None, name='total_cost')` Combine multiple objective ports into a **weighted sum**. .. math:: ``` J_{\text{total}} = \sum_{i} w_i \, J_i ``` ##### Parameters builder: The :class:`~jaxonomy.DiagramBuilder` to add blocks to. objectives: Sequence of scalar output ports, one per term. weights: Scalar weights :math:`w_i`. `None` → uniform weight `1.0`. Must have the same length as *objectives* when provided. name: Name of the final :class:`~jaxonomy.library.Adder` block (and prefix for :class:`~jaxonomy.library.Gain` blocks when weights differ from 1). ##### Returns OutputPort Scalar port equal to :math:`J_{\text{total}}`. ##### Raises ValueError If *objectives* is empty or *weights* has a different length. ##### Examples Combine two ISE objectives with different priorities:: ``` cost_pos = ise_objective(b, x.output_ports[0], name="ise_x") cost_vel = ise_objective(b, v.output_ports[0], name="ise_v") total = weighted_sum(b, [cost_pos, cost_vel], weights=[10.0, 1.0]) ``` # Analysis ## `jaxonomy.analysis` Whole-model analysis on top of the framework's dependency structure. The influence graph merges the model's leaf-level dependency DAG (which says *whether* information flows) with autodiff Jacobians (which say *how much*), giving quantitative model slicing, chain-rule path attribution, bottleneck detection, and dead-edge diagnostics on one queryable object. ### `InfluenceGraph` A model's dependency structure with autodiff-computed edge weights. Build with :func:`influence_graph`; see this module's docstring for the weighting conventions the numbers obey. Attributes: | Name | Type | Description | | ------------- | --------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `system` | `Any` | The analyzed Diagram or LeafSystem. | | `graph` | `DiGraph` | The underlying networkx.DiGraph. Node attributes describe the signal (kind, block, port, size, value, units, sample_time, hybrid); edge attributes carry kind, jacobian, relative, weight, magnitude, local_gradient, note and — in trajectory mode — profile. | | `tau` | `float` | Time scale applied to continuous-state-rate edges, in seconds. | | `normalize` | `str` | "relative" or "none". | | `scale_floor` | `float` | Lower bound on a signal's magnitude when normalizing. | | `at` | `str` | "operating_point" or "trajectory". | | `times` | `Optional[ndarray]` | Snapshot times in trajectory mode, else None. | | `reduce` | `str` | How a trajectory profile became the scalar weight. | | `block_notes` | `Dict[str, Dict[str, str]]` | Per-block explanations for anything not differentiated. | #### `attribute(target, source, *, threshold=1e-06, max_depth=32, max_paths=512)` Decompose `source`'s influence on `target` path by path. Each path's contribution is the chain-rule product of its edge weights; the signed sum over paths is the end-to-end sensitivity, which is where cancellation between two routes shows up as a total far below the largest single path. Parameters: | Name | Type | Description | Default | | ----------- | ------- | ---------------------------------------------------------------------------------------------------------- | ---------- | | `target` | | Destination node (id, port, or fragment). | *required* | | `source` | | Origin node. | *required* | | `threshold` | `float` | Prune a path once | product | | `max_depth` | `int` | Maximum path length. | `32` | | `max_paths` | `int` | Stop after this many paths and mark the result truncated, rather than enumerating a combinatorial blow-up. | `512` | #### `bottlenecks(target, *, threshold=0.01, max_depth=32)` Nodes every influential path to `target` must pass through. Computed on the slice at `threshold`: a node is a bottleneck when deleting it disconnects at least one slice origin from `target`. These are the signals worth instrumenting, and the single points of failure in a redundancy argument. Returns a bare list, so it has nowhere to report that the underlying slice was truncated — a truncated slice is missing paths, and a missing path is exactly what turns a non-bottleneck into an apparent one. That case warns instead; take the slice yourself and check :attr:`InfluenceSlice.truncated` if you need to handle it. #### `dead_edges(threshold=0.0)` Structural edges that transmit no influence at this operating point. A wire the model declares and the mathematics ignores: a gain of zero, a saturated nonlinearity, a term that cancels. This is the quantitative form of a dead-store warning — the connection is real, the influence is not. Edges with no local gradient are excluded (unknown is not dead), and so are the state self-loops, whose zero A block is the *definition* of a plain integrator rather than a defect. #### `dominant_paths(target, k=5, *, source=None, threshold=1e-06, max_depth=32, max_paths=512)` The `k` strongest paths into `target` (optionally from `source`). With no `source`, every node with no in-edges inside the search — the model's genuine independent inputs and states — is used as an origin. #### `nodes_at_scale_floor()` Signals whose normalizer came from `scale_floor`, not from a value. A relative weight divides by the signal's magnitude, so a signal that is (near) zero at the operating point — an error signal at equilibrium, an integrator state at `t=0` — produces an elasticity governed by `scale_floor` rather than by the model. Those weights are not wrong so much as meaningless, and they are large, so they dominate any ranking. #### `relative_threshold(target, fraction=0.01, *, direction='backward', max_depth=32, floor=1e-12)` A threshold set at `fraction` of the strongest influence on `target`. An absolute threshold only reads as a percentage when `tau` is comparable to the time constants on the paths involved (see the module docstring). Scaling to the strongest score makes "keep what carries at least 1% of what the dominant contributor carries" mean the same thing at any `tau`. Parameters: | Name | Type | Description | Default | | ----------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------ | | `target` | | Node id, port object, or name fragment. | *required* | | `fraction` | `float` | Fraction of the strongest score to keep. | `0.01` | | `direction` | `str` | As in :meth:slice. | `'backward'` | | `max_depth` | `int` | As in :meth:slice. | `32` | | `floor` | `float` | Threshold the reference sweep runs at, and the value returned when nothing upstream carries influence. It is passed to the search rather than left at zero so the sweep stays pruned; a model whose strongest contributor falls below it would yield floor itself. | `1e-12` | Returns: | Type | Description | | ------- | ------------------------------------------------------- | | `float` | A threshold to pass to :meth:slice / :meth:bottlenecks. | #### `resolve(spec)` Turn a port object, locator, or name fragment into a node id. Accepts an exact node id, an `InputPort` / `OutputPort`, a `(system, port_index)` locator, or any unambiguous suffix of a node id (`"integ:out:out_0"`, `"integ"`, `"out:y"`). #### `slice(target, threshold=0.01, *, direction='backward', max_depth=32)` Quantitative model slice: what influences `target` by ≥ `threshold`. The boolean answer — everything structurally upstream — is :meth:`structural_slice`; this one keeps only what lies on a path carrying at least `threshold` of the influence, in the relative-sensitivity sense described in the module docstring. `0.01` reads as "1%" only when `tau` is comparable to the time constants on the paths involved — the threshold is absolute, and a path across *k* integrators carries a factor of `tau**k`, so the same cutoff means different things at different `tau`. When the strongest contributor scores 95, `threshold=0.01` retains everything down to ~0.01% of it, not 1%. Use :meth:`relative_threshold` to get the cutoff that means a fraction *of the dominant contributor*. Two kinds of node are kept, and the distinction is load-bearing. A node is **influential** when its own best path to `target` clears the threshold. It is a **connector** when it merely lies on some influential node's best route: a relative weight is an elasticity, so a signal can pass through a junction that nearly cancels it and be amplified back afterwards, leaving a mid-route node with a small score of its own. Keeping only the influential ones would punch holes in the result — naming a block as influential while the route from it to the target ran through blocks that had been dropped, leaving :attr:`InfluenceSlice.subgraph` disconnected and :meth:`bottlenecks` meaningless. Connectors are read off the routes the search actually found, so nothing is added that no real path uses. `scores` reports every retained node's own best product to the target, which is the number to rank by. Parameters: | Name | Type | Description | Default | | ----------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------ | | `target` | | Node id, port object, or name fragment (see :meth:resolve). | *required* | | `threshold` | `float` | Minimum | path product | | `direction` | `str` | "backward" (default, what influences the target) or "forward" (what the target influences). | `'backward'` | | `max_depth` | `int` | Hard bound on path length, and the only hard bound — a partial product is not a bound on the whole path's (see :meth:\_reach), so nothing may be pruned on the running value. | `32` | Returns: | Name | Type | Description | | ---- | ---------------- | --------------------- | | `An` | `InfluenceSlice` | class:InfluenceSlice. | #### `structural_slice(target, *, direction='backward')` Boolean slice: every block structurally connected to `target`. The over-approximation :meth:`slice` improves on, computed from the model's declared connectivity rather than from the weighted graph — so it stays a genuine bound even where a Jacobian could not be taken. Provided so the two can be compared directly on a real model. #### `summary()` Human-readable overview: size, conventions, and honesty labels. ### `InfluenceSlice` A quantitative model slice: what actually reaches a target. Attributes: | Name | Type | Description | | --------------- | ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `target` | `str` | Node id the slice was taken to (or from). | | `threshold` | `float` | Influence cutoff a path had to clear to be included. | | `direction` | `str` | "backward" (what influences the target) or "forward" (what the target influences). | | `scores` | `Dict[str, float]` | {node_id: best | | `edges` | `List[Tuple[str, str]]` | (src, dst) pairs retained. | | `blocks` | `List[str]` | Block name paths touched — the block-level slice. | | `unknown_nodes` | `List[str]` | Nodes that some retained path reaches across an edge with no local gradient. Their score accounts only for the measurable routes, so it is not the whole story — treat it as a partial reading rather than a measurement. | | `truncated` | `bool` | True if the path search hit its expansion budget, in which case the scores are lower bounds and the slice may be missing contributors. | | `graph` | `'InfluenceGraph'` | The originating :class:InfluenceGraph. | #### `block_scores` `{block name path: score}`, ranked, for the block-level answer. :attr:`scores` is keyed by *signal* (one node per input port, output port, and state group), which is the right granularity for tracing a route but the wrong one for "which block matters most". This reduces a block's nodes to one number by taking the **maximum**, so a block's score is that of its most influential signal. Max is the reducer because a block's input and output nodes lie on the *same* path — summing them would count one route twice, and a block's influence is not the sum of its ports' influences. The trade-off is that a block reached by several genuinely independent routes reads as its strongest one, not their total; use :meth:`attribute` when the split between routes is the question. The dict is ordered by descending score. Blocks holding a node in :attr:`unknown_nodes` are present with a score covering only their measurable routes — check that list before reading a rank as complete. #### `subgraph` The retained portion of the influence graph. Built from the retained nodes *and* edges rather than as an edge-induced view, so a node with no retained edge — the target of a slice that keeps nothing else — is still present. #### `unknown_paths` True if any retained path crosses an edge with no local gradient. #### `report(by='node')` Human-readable ranking. Parameters: | Name | Type | Description | Default | | ---- | ----- | --------------------------------------------------------------------------------------- | -------- | | `by` | `str` | "node" (default) ranks individual signals; "block" ranks blocks via :attr:block_scores. | `'node'` | ### `LeafJacobians` Local Jacobian blocks for one leaf at one operating point. Attributes: | Name | Type | Description | | ------- | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `leaf` | `Any` | The LeafSystem these Jacobians describe. | | `u0` | `List[Any]` | Operating-point value of each input port, in port order. | | `y0` | `List[Any]` | Operating-point value of each output port, in port order. | | `x0` | `Dict[str, Any]` | Operating-point value of each state kind present, keyed by "xc" / "xd". | | `d` | `Dict[Tuple[int, int], ndarray]` | {(out_i, in_j): ndarray(m_i, n_j)} — direct feedthrough. | | `c` | `Dict[Tuple[str, int], ndarray]` | {(kind, out_i): ndarray(m_i, n_x)} — state → output. | | `b` | `Dict[Tuple[str, int], ndarray]` | {(kind, in_j): ndarray(n_x, n_j)} — input → state rate/update. | | `a` | `Dict[Tuple[str, str], ndarray]` | {(src_kind, dst_kind): ndarray(n_dst, n_src)} — state → state rate/update, including the cross terms (an ODE reading discrete state, a periodic update reading continuous state). | | `notes` | `Dict[str, str]` | {subject: reason} for every quantity that could not be differentiated, e.g. {"out:mode": "non-inexact dtype int32"}. Callers turn these into local_gradient=None edge labels rather than silently reporting a zero. | ### `PathAttribution` Chain-rule decomposition of one source's influence on one target. Attributes: | Name | Type | Description | | ----------------- | ---------------------- | --------------------------------------------------------------------------------------------------------------- | | `target` | `str` | Destination node id. | | `source` | `str` | Origin node id. | | `paths` | `List[Dict[str, Any]]` | One entry per path, ranked by | | `total` | `Optional[float]` | Signed sum of path products when every path is signed, else None — a sum of magnitudes would hide cancellation. | | `total_magnitude` | `float` | Sum of | | `truncated` | `bool` | True if enumeration hit max_paths or max_depth. | ### `format_influence_subgraph(graph, focus, edges, types=None, rates=None, dropped_for_budget=0)` Render a node/edge selection as compact, citable text. The footer distinguishes the two reasons a block can be absent, because to a reader they mean opposite things. A block left out because its influence fell below the threshold is *known to be negligible* — that is an answer. A block left out because the budget ran out is simply *unknown*, and treating it as negligible would be a fabrication. Without the footer both look identical: missing. ### `influence_graph(system, context=None, *, at='operating_point', results=None, times=None, n_snapshots=5, tau=1.0, normalize='relative', scale_floor=1e-06, probe=None, reduce='max', simulator_options=None)` Build the sensitivity-weighted influence graph of a model. Parameters: | Name | Type | Description | Default | | ------------------- | --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------- | | `system` | | A Diagram or a single LeafSystem. | *required* | | `context` | | Root context fixing the operating point. Defaults to system.create_context(). | `None` | | `at` | `str` | "operating_point" weights every edge once, at context. "trajectory" weights at several snapshots and stores per-edge profiles — the honest answer when a nonlinearity means one number per edge cannot be right everywhere (a block saturated at the operating point has a zero local gradient there and a large one elsewhere). | `'operating_point'` | | `results` | | A SimulationResults supplying the snapshot times for at="trajectory". The states are re-derived by advancing context, because recorded signals do not pin down every stateful leaf — which costs one simulate call per snapshot. Budget for that on a large model: simulate's fixed setup cost scales with block count and dominates the integration itself (a 1 µs span costs the same as a 4 s one), so n_snapshots=6 on a 2500-block model is minutes rather than seconds. Building at a single operating point is linear in block count and stays in seconds at that size. | `None` | | `times` | `Optional[Sequence[float]]` | Explicit snapshot times, used instead of results. | `None` | | `n_snapshots` | `int` | How many times to take from results.time. | `5` | | `tau` | `float` | Seconds of integration represented by a continuous-state-rate edge; only affects edges into ẋc. Set it from the fastest state on the paths you care about — every integrator on a path contributes a factor of tau, so a value taken from the slow dynamics of a stiff model inflates multi-integrator path products (see the module docstring). | `1.0` | | `normalize` | `str` | "relative" (default, dimensionless elasticities) or "none" (raw partial derivatives in model units). | `'relative'` | | `scale_floor` | `float` | Floor on a signal's operating-point magnitude when normalizing, so a signal that happens to sit at zero does not produce an infinite elasticity. Nodes at the floor are visible via their value attribute. | `1e-06` | | `probe` | `Optional[float]` | When set to a relative step size (0.05 = 5% of each signal's magnitude), every edge whose exact derivative is zero is re-checked with a central-difference secant, and the secant is used instead when it is non-zero. This is the cross-check for the one thing an exact local derivative gets wrong: a quantizer between steps, a saturation at its rail, or a dead zone inside the zone is locally flat while still transmitting information, and would otherwise be reported dead. Costs two extra block evaluations per signal component; None (default) skips it. | `None` | | `reduce` | `str` | How a trajectory profile collapses to the scalar weight used by queries: "max" (default, conservative — never hides an influence that appears at some point), "mean", or "final". | `'max'` | | `simulator_options` | | SimulatorOptions for the trajectory-mode re-integration. | `None` | Returns: | Name | Type | Description | | ---- | ---------------- | --------------------- | | `An` | `InfluenceGraph` | class:InfluenceGraph. | Example > > > import jaxonomy from jaxonomy.library import Constant, Gain, Integrator from jaxonomy.analysis import influence_graph builder = jaxonomy.DiagramBuilder() source = builder.add(Constant(1.0, name="src")) gain = builder.add(Gain(3.0, name="gain")) plant = builder.add(Integrator(1.0, name="plant")) builder.connect(source.output_ports[0], gain.input_ports[0]) builder.connect(gain.output_ports[0], plant.input_ports[0]) diagram = builder.build(name="root") graph = influence_graph(diagram) graph.slice("plant:xc", threshold=0.01).blocks ['gain', 'plant', 'src'] ### `influence_subgraph(graph, focus, *, budget_tokens=1500, hops=4, threshold=0.0, direction='both')` A bounded, budgeted, citable neighbourhood of `focus`. Parameters: | Name | Type | Description | Default | | --------------- | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `graph` | `InfluenceGraph` | An :class:~jaxonomy.analysis.influence.InfluenceGraph. | *required* | | `focus` | | One or more focus points — node ids, port objects, name fragments, or a block name (which expands to all of that block's signals). | *required* | | `budget_tokens` | `int` | Approximate ceiling on the rendered text, at :data:CHARS_PER_TOKEN characters per token. Edges are dropped weakest-first to fit; the result reports what was dropped. | `1500` | | `hops` | `int` | How many graph edges out from the focus to expand. Nodes are signals, so crossing one block costs two hops (wire in, block Jacobian out) — the default of 4 reaches roughly two blocks. | `4` | | `threshold` | `float` | Minimum edge | weight | | `direction` | `str` | "both" (default), "backward" (what influences the focus) or "forward" (what it influences). | `'both'` | Returns: | Type | Description | | ---------------- | ------------------------------------------------------ | | `Dict[str, Any]` | A dict with text (the rendered context), nodes, edges, | | `Dict[str, Any]` | blocks, focus, estimated_tokens, dropped_edges and | | `Dict[str, Any]` | conventions. The dict is JSON-serializable. | ### `leaf_jacobians(leaf, root_context)` Compute every local Jacobian block of `leaf` at `root_context`. Parameters: | Name | Type | Description | Default | | -------------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------- | | `leaf` | | A LeafSystem belonging to the system root_context was created from. | *required* | | `root_context` | | Root context supplying the operating point — time, parameters, this leaf's state, and (via upstream evaluation) the values arriving on its input ports. | *required* | Returns: | Name | Type | Description | | ---- | --------------- | ---------------------------------------------------------------- | | `A` | `LeafJacobians` | class:LeafJacobians. Blocks that could not be differentiated are | | | `LeafJacobians` | absent from d / c / b / a and explained in notes; | | | `LeafJacobians` | this function does not raise on a non-differentiable block. | # Guides and scope # Tutorials These notebooks introduce **Jaxonomy** step by step: block diagrams, simulation, customization, and optimization. | Notebook | What you will learn | | ---------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ | | [Getting started](https://py.jaxonomy.com/tutorials/01-getting-started/index.md) | Build a simple diagram with `DiagramBuilder`, run `simulate`, inspect results. | | [Creating custom blocks](https://py.jaxonomy.com/tutorials/02-creating-custom-blocks/index.md) | Implement your own `LeafSystem` blocks and wire them into diagrams. | | [Creating custom acausal components](https://py.jaxonomy.com/tutorials/03-creating-custom-acausal-components/index.md) | Extend the acausal modeling layer with new components. | | [Wrappers and decorators](https://py.jaxonomy.com/tutorials/04-wrappers/index.md) | Patterns for wrapping and decorating systems. | | [Automatic differentiation and optimization](https://py.jaxonomy.com/tutorials/05-automatic-differentiation-optimization/index.md) | Use AD with models for optimization workflows. | For more topic-specific notebooks (MPC, estimation, hardware-inspired demos, etc.), see **[Examples](https://py.jaxonomy.com/examples/index.md)**. # Example notebooks ## Introductory examples If you haven't already, check out the [tutorials](https://py.jaxonomy.com/tutorials/index.md), which explain how to build and simulate models in Jaxonomy. ### [Primitive blocks and composability](https://py.jaxonomy.com/examples/primitives/index.md) Shows how to build systems with primitive blocks and how to compose them into larger diagrams. ### [Block diagram visualization](https://py.jaxonomy.com/examples/diagram_visualization/index.md) Render a `DiagramBuilder` composition as a block diagram to inspect wiring, ports, and hierarchy before simulating. ### [Custom block authoring: from `LeafSystem` to gradient-correctness CI](https://py.jaxonomy.com/examples/custom_block_authoring/index.md) Walks through authoring a custom `LeafSystem` block end-to-end — declaring state, parameters, and ports; wiring the callback signatures — and standing up the gradient-correctness test that guards it. ### [Bouncing ball](https://py.jaxonomy.com/examples/bouncing_ball/index.md) Shows hybrid dynamics modeling of a bouncing ball. ### [Hybrid thermostat (marimo tutorial)](https://py.jaxonomy.com/examples/hybrid_thermostat_tutorial/index.md) A [marimo](https://marimo.io/) notebook implementing a thermostat-controlled room: continuous temperature dynamics with discrete HEAT/OFF modes and hysteresis. Run with `marimo edit hybrid_thermostat_tutorial.py` or `marimo run hybrid_thermostat_tutorial.py` (requires `marimo` and `jaxonomy`). ### [Triple inverted pendulum: Jaxonomy LQR + MuJoCo (marimo)](https://py.jaxonomy.com/examples/triple_inverted_pendulum_mujoco_marimo/index.md) A [marimo](https://marimo.io/) app that stabilizes a triple pendulum on a cart. Jaxonomy's `LinearQuadraticRegulator` synthesizes the stabilizing gain (a continuous-time algebraic-Riccati solve on the MuJoCo-linearized plant), and MuJoCo `mj_step` is the high-fidelity physics it is validated against: a short open-loop kick-up impulse, then closed-loop LQR balance, with a matplotlib + MuJoCo animation. Run with `marimo run triple_inverted_pendulum_mujoco_marimo.py` (requires `marimo`, `jaxonomy`, `mujoco`). ### [Bouncing balls on stairs: Jaxonomy dynamics + MuJoCo visualization (marimo)](https://py.jaxonomy.com/examples/bouncing_ball_stairs_marimo/index.md) A [marimo](https://marimo.io/) app where three balls (bocce, golf, ping-pong) cascade down a five-step staircase. **Jaxonomy** computes the motion as a hybrid dynamical system — each ball a point mass in free fall with a zero-crossing event at every tread and a Newtonian restitution reset ((\\dot z \\mapsto -e,\\dot z)) — and **MuJoCo renders it kinematically** (ball positions set from the Jaxonomy trajectory each frame; `mj_forward`, no contact solver). The restitution coefficient (e) alone (0.85 / 0.70 / 0.55) produces the light/medium/heavy bounce characters, and the energy panel shows each bounce keeping a fraction (e^2) of the vertical KE. A multi-ball, staircase counterpart to the Jaxonomy-native [bouncing ball](https://py.jaxonomy.com/examples/bouncing_ball/index.md) example above. Run with `marimo run bouncing_ball_stairs_marimo.py` (requires `marimo`, `jaxonomy`, `mujoco`). ### [Linear Quadratic Regulator (LQR)](https://py.jaxonomy.com/examples/lqr/index.md) Demonstrates the LQR for a pendulum and a planar quadrotor model. ### [Energy shaping and LQR stabilization](https://py.jaxonomy.com/examples/energy_shaping_and_lqr/index.md) Demonstrates energy shaping control to swing a pendulum to the vertically 'up' orientation and then stabilize it in the 'up' orientation via LQR. ### [Linear Model Predictive Control (MPC)](https://py.jaxonomy.com/examples/linear_mpc/index.md) Demonstrates MPC on a linearized model of the Cessna Citation aircraft and a pendulum model. ### [Differentiable Predictive Control (DPC) of a two-tank system](https://py.jaxonomy.com/examples/dpc_two_tank_reference_tracking/index.md) Trains a neural control policy to make the bottom level of a cascaded two-tank rig track a setpoint, by differentiating straight through the simulated ODE. Shows the policy-as-a-block API (`PolicyBlock`, `PlantBlock`, `build_closed_loop`, `simulate_closed_loop`): the same policy object trains as a differentiable function and deploys as a composable `LeafSystem` under `jaxonomy.simulate`. Deployment note for policies imported from *discrete-time* training loops (torch/NEUROMANCER-style, via `ONNXJax` or the predictor blocks): sample-and-hold controllers must be followed by `ZeroOrderHold(dt=ts)` with `SimulatorOptions(max_major_step_length=ts, max_minor_step_size=ts)` for step-grid parity — otherwise the policy is silently re-evaluated at every RK4 stage and step-for-step parity with the exporting framework is lost (with the ZOH: parity ~4e-8 over 400 steps on this benchmark). ### [Multi-layer perceptron (MLP)](https://py.jaxonomy.com/examples/MLP_training/index.md) Demonstrates training of a multi-layer perceptron (MLP), a class of feedforward artificial neural networks, for a regression task. ## Advanced examples ### [Trajectory optimization and stabilization](https://py.jaxonomy.com/examples/trajectory_optimization_and_stabilization/index.md) Shows trajectory optimization for the problem of swinging an Acrobot to the vertically 'up' orientation and then stabilizing the trajectory via finite-horizon LQR. ### [Robotic arm control](https://py.jaxonomy.com/examples/mujoco/pick_and_place/index.md) Implement a controller for a "pick-and-place" task with a robotic arm using MuJoCo as a multibody physics engine. Download the necessary files from [here](https://py.jaxonomy.com/examples/mujoco/pick_and_place/assets/franka_emika_panda.zip). ### [Automatic tuning of a PID controller](https://py.jaxonomy.com/examples/pid_tuning/index.md) Demonstrates automatic differentiation and optimization capabilities of Jaxonomy to automatically tune the gains of a discrete-time PID controller. ### [Interactive and automatic tuning of a PID controller with sensitivity constraints](https://py.jaxonomy.com/examples/pid_autotuning_interactive/index.md) Showcases fast compiled simulations in Jaxonomy for interactive applications and automatic tuning of a continuous-time PID controller with maximum sensitivity and complementary sensitivity constraints. ### [Finding limit cycles](https://py.jaxonomy.com/examples/limit_cycles/index.md) Demonstrates how to find limit cycles and assess their stability by leveraging the automatic differentiation capabilities of Jaxonomy. ### [Kalman Filters: linear and nonlinear extensions](https://py.jaxonomy.com/examples/state_estimation_with_Kalman_filters/index.md) Demonstrates the use of Kalman filters (finite and infinite-horizon) and nonlinear extensions (Extended Kalman Filter and Unscented Kalman Filter) for state estimation in a pendulum model. Where necessary, the nonlinear Pendulum plant is automatically linearized and discretized by Jaxonomy for the construction of the filters. ### [Engine map fitting to MPC: differentiable lookup tables end-to-end](https://py.jaxonomy.com/examples/engine_map_fitting_to_mpc/index.md) Fit a noisy 2-D engine torque map with `fit_lookup_table_2d`, optimise breakpoint placement with `fit_table_1d_with_grid`, drop the fitted `LookupTable2d` into a 1-DOF vehicle plant, run a short-horizon shooting MPC tracking a velocity reference, and take a single `jax.grad` of closed-loop tracking RMSE with respect to the table values, the grid placement, **and** the MPC weights — all in one call. Showcases the differentiable lookup-table family. ### [Linearization workflow: from findop to Bode / Nyquist / pzmap and empirical FRE](https://py.jaxonomy.com/examples/linearization_workflow/index.md) The full linearization toolchain on one canonical plant: trim with `findop`, take Jacobians with `linearize`, evaluate frequency response with `frequency_response` / `bode_data`, plot Nyquist contours with `nyquist_data`, scatter poles and zeros with `pole_zero_map`, and overlay closed-form `step_response` / `impulse_response` against the nonlinear simulator. Finishes with `estimate_frequency_response` driven by a chirp (matching the analytic Bode to within ~1 dB) and a `jax.grad` of the Bode-peak magnitude w.r.t. damping — the differentiability a JAX-native stack adds on top of a classical linearization workflow. ### [Quantitative model slicing: the influence graph](https://py.jaxonomy.com/examples/quantitative_model_slicing/index.md) The tutorial treatment of `jaxonomy.analysis.influence_graph`, built on a 19-block PI-controlled DC motor whose boolean slice is all 19 blocks. Derives the two conventions that make the numbers readable — relative (elasticity) weights, which *telescope* so that a path product is the relative end-to-end sensitivity, and `tau`, which is a frequency choice rather than a fudge factor — and then verifies both: path attribution matches a central difference taken through the whole diagram to ~(9\\times10^{-12}) relative, and a (\\tau)-scaled path product through two integrators matches (|G(j\\omega)|) at (\\omega = 1/\\tau) to machine precision across four decades. A (\\tau) sweep on the motor shows each block's influence tracking (\\tau^k) for its integrator depth, with the setpoint's dominant route visibly switching from the proportional branch to the integral one. Slicing at 1% drops exactly the negligible stiction term; `dominant_paths` separates the integral route from the proportional one; `bottlenecks` names the ten blocks with no redundancy. With the driver saturated, the slice returns 13 of 19 blocks at *any* threshold — the control path's derivative is exactly zero, not small — and `probe=0.9` recovers it while `at="trajectory"` shows the edge switching on as the loop leaves the rail. Closes with a worked demonstration that a vector state's weight is an upper bound (2.0 reported for a true 0.0), a short and deliberately-bounded note on serializing a neighbourhood for a language model, six failure modes, and five exercises. ~30 s runtime. ### [Quantitative model slicing: the influence graph (runnable script)](https://py.jaxonomy.com/examples/influence_graph_model_slicing/index.md) A boolean dependency graph answers "what affects the shaft speed?" with "everything in the loop" — correct and useless. `jaxonomy.analysis.influence_graph` weights the same graph with exact local Jacobians, and the question becomes quantitative. On a 19-block PI-controlled DC motor: path attribution reproduces the analytic (-1/J) load-torque sensitivity and agrees with central differences through the whole diagram to ~(9\\times10^{-12}) relative; the 1% slice drops exactly the negligible stiction term the boolean slice keeps; `dominant_paths` separates the route through the integral term from the route through the proportional one — a distinction a boolean graph cannot draw. With the driver saturated, the 1% slice returns 13 of 19 blocks and `dead_edges` names the inert connection; `probe=0.9` recovers the six upstream control blocks whose *local* derivative is genuinely zero, and `at="trajectory"` shows that edge switching on as the loop leaves the rail. Closes with a token-budgeted, citable serialization of the graph neighbourhood and a plain statement of the four caveats (`tau` is a frequency choice; a signal sitting at zero is normalized by `scale_floor`, which trajectory mode avoids; vector states report an upper bound; a hybrid block's weights describe one mode). Run with `python docs/examples/influence_graph_model_slicing.py`. ### [Aleatoric vs epistemic uncertainty: Sobol decomposition on a noisy plant](https://py.jaxonomy.com/examples/aleatoric_vs_epistemic_uq/index.md) End-to-end tour of `jaxonomy.uq`: sanity-check Sobol indices on the Ishigami benchmark, compare IID / Latin-hypercube / quasi-Monte-Carlo convergence on a damped pendulum, rank parameters by first- and total-order Sobol index, screen with Morris elementary effects, and finish with the headline `decompose_variance_sobol` call that splits the QoI variance into aleatoric (irreducible) and epistemic (reducible-with-better-data) shares plus their cross-interaction — a distinction that is often hard to surface cleanly. Closes with a value-at-risk / CVaR computation that connects the variance budget to a concrete design margin. ### [Bit-exact reproducibility: capturing and replaying a simulation](https://py.jaxonomy.com/examples/reproducibility_manifest/index.md) Run a damped oscillator with `SimulatorOptions(record_provenance=True)`, persist the resulting `ProvenanceManifest` and parameter pytree as sibling JSON files, replay the simulation on the same machine and byte-compare the outputs via `array.tobytes()` equality, then deliberately tamper with the model — first by a single-ULP parameter perturbation, then by a sign flip in the dynamics — and watch the manifest's stable fingerprint and the parameter SHA each catch the right failure mode. Honest about the contract: shows precisely what the manifest fingerprints (versions, options, system type, sorted parameter names), what it deliberately doesn't (parameter values, source code, cross-device floats), and how to layer a parameter hash next to the manifest for full value-drift detection. Few simulation tools emit this artifact by default — it's the trust-building layer for reproducible results. ### [Hybrid trajectory optimization through events: gradients across the bounce](https://py.jaxonomy.com/examples/hybrid_trajopt_through_events/index.md) A bouncing-ball plant with `record_event_times=True` exposes touchdown instants to autodiff: `jax.grad` flows through the bounce via the saltation rule, so a scalar objective like "ball energy at (t=T)" becomes differentiable w.r.t. the restitution coefficient. Demonstrates closed-form vs simulator bounce-time agreement to (\\sim 10^{-8}) s, a single-point `dt_event/dh0 = 0.2258 s/m` matching analytic (1/\\sqrt{2 g h_0}), a multi-event tilted-floor case with one gradient per bounce, and `vmap_event_time_gradient` for batched gradients across a 16-point initial-height sweep. Honest about two cracks surfaced while authoring: `with_parameter("e", float(e))` triggers a fresh JIT trace per call (filed as a follow-up finding), and the multi-event saltation gradient disagrees with finite differences by 0.16 s/(grade) on the tilted-floor first bounce — a correctness regression also filed. The headline capability: gradients across discrete contact events, which most block-diagram tools don't offer. ### [Networked control: identifying actuator delay from bench data](https://py.jaxonomy.com/examples/actuator_delay_identification/index.md) Synthesize noisy open-loop bench data on a 1st-order servo with a known transport delay `τ_true`, define a mean-squared prediction error against the bench data as a function of `τ`, and watch `jax.grad` flow cleanly through the `VariableTransportDelay` block (the linear-interpolation lookup is what makes the gradient finite and nonzero). A handful of scipy iterations recover `τ` to within ~5% of ground truth; a higher-noise run shows the estimator gets noisier but stays unbiased; and a closing closed-loop demo contrasts the unmodeled-delay baseline (oscillating, saturated) against a Smith-predictor controller built around the identified delay. Showcases the continuous + variable transport delay blocks and the differentiability story behind the saltation rule. ### [Multirate controller: 1 kHz inner / 100 Hz outer / 10 Hz supervisor](https://py.jaxonomy.com/examples/multirate_controller/index.md) Build the canonical embedded-control cascade on a PM DC motor — 1 kHz current loop, 100 Hz velocity loop, 10 Hz position controller plus state-machine supervisor — and route a four-field measurement bus between them with `BusCreator` / `BusSelector`. Place `RateTransition` blocks explicitly between the layers and also exercise `DiagramBuilder(auto_insert_rate_transitions=True)`; inspect what the scheduler actually built with `Diagram.print_schedule()` and `rate_summary_dot()`. Closes with `analyze_saturation` and `analyze_phase_activity` on the closed-loop trajectory, calling out the benign FAULT-never-fired warning and pointing the reader at the FAULT-trip exercise. Showcases multirate scheduling + auto-inserted rate transitions, `RateTransition` / `Decimator`, and `BusCreator` / `BusSelector`. ### [Real-time fixed-step controller: the embedded deployment story](https://py.jaxonomy.com/examples/realtime_fixed_step_controller/index.md) The same notebook that runs the offline simulator is the reference for the deployed binary — fixed step, deterministic byte-for-byte, microseconds per tick. Builds a 1 kHz discrete PID plus a Luenberger velocity observer wrapped around a lightly-damped 2nd-order spring-mass plant with noisy position measurement, then compares `SimulatorOptions(ode_solver_method="rk4")` against the adaptive `dopri5` solver on three embedded-engineering metrics: per-tick wall-clock jitter (the WCET budget the deadline scheduler must accommodate), bit-identical determinism on re-runs, and post-JIT throughput. **Headline live numbers**: RK4 mean per-tick wall-clock = **24.1 μs**, 99/1 jitter ratio = **1.60×**; Dopri5 jitter ratio under the same `max_major_step_length=DT` constraint = **1.52×** (the host-callback synchronisation floor dominates per-tick timing on a developer laptop — the real algorithmic wedge is in the RHS-evaluation count). Under natural adaptive operation, **Dopri5 takes 1.20× more RHS evaluations than RK4** for the same controller horizon (6.01 vs 5.00 calls/tick) — and this is the embedded WCET metric that translates faithfully to an MCU because RHS count is fixed floating-point work. Post-JIT throughput in steady state: **45.7 μs / tick** (= ~21 kHz on a developer laptop, 4.6% CPU budget at the 1 kHz control rate). Bit-identical re-runs verified on both solvers via `tobytes()`; RK4 vs Dopri5 trajectories agree to **1.5 μm at final time** (= 0.015% of the 10 mm setpoint). Closes with `analyze_saturation` on the actuator (0.01% saturated, well under the 50% threshold) and five exercises (10 kHz rate, float32 opt-out, fault-detection state machine, XLA HLO inspection, per-tick budget breakdown on a real embedded project). Honest about the JIT-warmup cost, the XLA-CPU vs bare-metal-MCU gap, and the difference between bit-exact reproducibility *of the host development sim* and *of the deployed binary*. Showcases fixed-step RK4 via `SimulatorOptions(ode_solver_method="rk4")`, byte-exact reproducibility, and the precision policy (exercise-only). ### [Fast restart and batched sweeps: when JIT amortization pays off](https://py.jaxonomy.com/examples/fast_restart_and_batched_sweeps/index.md) Time a 1000-trial damping sweep on a damped harmonic oscillator three ways — a naive `simulate()` loop, `FastRestartSimulator`'s warm-cached kernel, and `simulate_batch(use_vmap=True)`'s vectorised XLA launch — and watch the JIT-amortisation wedge open up across two orders of magnitude. On CPU, `FastRestartSimulator` recovers the closed-form ISE-optimal damping ((\\zeta = 0.5)) from a brute-force grid in ~0.3 s versus ~130 s for the naive loop — a ~430× speedup with the four code paths agreeing bit-for-bit on the per-trial ISE. The tutorial is candid about where vmap surprises a CPU user (the per-row finalize is linear in (N), so vmap *loses* to the kernel path on CPU at these scales but wins on GPU/TPU where the parallel launch dominates), exercises three concrete failure modes — `FastRestartSimulator`'s structural-change warning on a dtype swap, `simulate_batch(use_vmap=True)`'s host-callback `ValueError`, and the `simulate(...).buffer_length` ring-buffer pitfall (now warned-on as of `7a24e31`) — and closes with a decision table that names the regime where each API wins. Showcases `FastRestartSimulator` + `simulate_batch` (kernel and vmap paths). ### [Product-family modeling with Variants: one diagram, three drivetrains](https://py.jaxonomy.com/examples/product_family_variants/index.md) Ship three trim levels — gasoline ICE, parallel hybrid, battery-electric — of the same longitudinal vehicle model from one `DiagramBuilder` and one `Variant` at the powertrain attachment point. `select_variant` resolves the topology at build time (unselected branches never enter the JIT trace); `Diagram.with_config(powertrain=...)` swaps the realised choice post-build; `dump_variant_config_to_json` / `load_variant_config_from_json` round-trip the active binding (bytes-equal on replay); the `python -m jaxonomy.cli.run_variants` CLI surfaces the same data from release-pipeline scripts. Closes with a 3 × 5 Cartesian-product sweep of 0-60 time over `(drivetrain, driver-aggressiveness)` via `expand_all_variant_configs` + `simulate_batch`, all 15 sims under one JIT-compiled kernel per variant. Showcases the Variants system (phases 1–4) plus the `with_config` / introspection / runtime-switch follow-ups, and cross-links the substrate `submodel_function` the variant DSL extends. ### [Truth tables and state machines: vehicle gear-selection logic](https://py.jaxonomy.com/examples/truth_table_gear_logic/index.md) A four-bit eligibility truth table (`brake`, `stopped`, `req_b1`, `req_b0`) drives a four-state Park / Reverse / Neutral / Drive FSM with ten prioritised transitions and — deliberately — *no* Drive ↔ Reverse direct edge. Builds the table with the fluent `TruthTable.builder(...).row(...).build()` API (wildcards on the driver-request bits collapse 16 input combinations to 3 rows), runs `validate(strict_completeness=True)` and watches the strict check catch a deliberately-removed row, byte-equal round-trips the table through `to_csv` / `from_csv`, and runs a 14-second scripted driving trace where the FSM silently refuses an illegal Drive → Reverse request (caught post-hoc by `analyze_phase_activity` on a truncated trace). Closes with `jax.grad` flowing through a standalone gain-schedule table whose row outputs are callables of the raw inputs (the numeric-output extension). Honest about the four DX papercuts the tutorial surfaced — TruthTable being branchless not vectorised (a follow-up finding), `TruthTableBuilder(input_names=...)` labels not propagating to port names, no `time_mode=` knob on `StateMachineBuilder.build()`, and `overlapping_pairs=[]` being counter-intuitive on wildcard tables. Showcases the `TruthTable` API + the existing `StateMachineBuilder` API. ### [Two-degree-of-freedom PID with classical auto-tuning rules](https://py.jaxonomy.com/examples/pid_2dof_classical_tuning/index.md) Hand the practicing engineer a starting point from two minutes of bench time. Identify the ultimate-cycle parameters ((K_u, T_u)) via a relay-feedback experiment, fit the FOPDT triple ((K, \\tau, \\theta)) off the open-loop step's inflection tangent, then apply `PIDController2DOF.ziegler_nichols(Ku, Tu, dt, mode="PID")`, `.cohen_coon(K, tau, theta, dt, mode="PID")`, and `.tyreus_luyben(Ku, Tu, dt)` to the same FOPDT plant ((K=2), (\\tau=10\\ \\mathrm{s}), (\\theta=2\\ \\mathrm{s})). Compare on a setpoint step + load disturbance with quantitative metrics (rise time, overshoot, 2% settling time, windowed ISE) — Z-N moderate, Cohen-Coon most aggressive (best disturbance ISE, worst overshoot), Tyreus-Luyben most conservative (no overshoot, slowest recovery). The 2-DOF wedge then shows the same Z-N gains under three ((b, c)) settings: setpoint overshoot drops from 48% → 16% → 12% while the disturbance maximum stays pinned at 1.194 to four significant figures — the empirical demonstration that ((b, c)) change only the setpoint-tracking numerator, not the closed-loop poles. Closes with `jax.grad` refinement of the Cohen-Coon starting point (25 steps, 22% ISE reduction in ~2.5 s), `analyze_saturation` / `analyze_control_oscillation` diagnostics on every controller (all silent), and a concrete failure-mode demo (Cohen-Coon on a misidentified second-order plant overshoots 68% with visible ringing). Showcases `PIDController2DOF` + the three classmethod tuning rules, and bridges to the gradient-based and sensitivity-constrained stories in [`pid_tuning.ipynb`](https://py.jaxonomy.com/examples/pid_tuning/index.md) and [`pid_autotuning_interactive.ipynb`](https://py.jaxonomy.com/examples/pid_autotuning_interactive/index.md). ### [Unit-safe wiring: dimensional consistency at build time](https://py.jaxonomy.com/examples/unit_safe_wiring/index.md) Annotate ports with `Unit` (`meter`, `newton`, `joule`, …) and watch `DiagramBuilder.connect()` refuse a force-into-displacement wire before any kernel launches; auto-convert `mm`→`m` at the wire (with `unit_conversion="auto"` / `"warn"` / `"error"` modes); run `propagate_diagram_units(diagram)` and see the math-block algebra catch `Constant(units=newton) → Integrator → output declared as joule` because `N·s ≠ J` — the headline beat that closes Simulink's "units dropped after computational blocks" gap. Tags `BusCreator(field_units=...)` for compound bus signals, demonstrates offset-aware °C↔K via `convert_offset_aware`, sets a USD→EUR FX rate to exercise the currency axis, and shows the acausal connector library's canonical `flow_units` / `pot_units` per domain plus the `physical_quantity` tag that distinguishes `N·m@torque` from `N·m@energy`. Closes with JSON round-trip of unit annotations. Showcases the units system (phases 1–3), the `BusUnit` cross-link, and the currency / temperature / pint / source-block follow-ups. ### [Conservation laws as CI: patterns for ensuring physical validity](https://py.jaxonomy.com/examples/conservation_laws_as_ci/index.md) Every real-world simulator drifts. Energy bleeds from an oscillator, angular momentum tilts on a tumbling body, probability mass leaks out of a diffusion — the drift is usually invisible in the trajectory but real and shipping. Pedagogical guide to jaxonomy's property-test framework for asserting that physical invariants hold within tolerance over long simulation horizons. Walks four canonical patterns (energy on an undamped SHO, angular momentum on a torque-free rigid body, probability mass on a 5000-path Brownian ensemble, electrical energy on an acausal LC oscillator under BDF) and shows the test catching deliberately-broken plants in each one. **Headline numbers**: SHO energy conserved to relative drift **(1.45 \\times 10^{-9})** over 50 oscillation periods under Dopri5 at `rtol=1e-10`; a stray (-cv) damping bug ((c = 10^{-3})) pushes the drift to **(1.87 \\times 10^{-2})** at the endpoint — seven decades above the floor — and the test catches it. Rigid-body (\\lVert \\mathbf{H} \\rVert^2) drift under the same solver = **(5.75 \\times 10^{-10})** over 20 s of intermediate-axis tumble; a (5%) asymmetric scaling on one cross-product component pushes drift to **(3.0 \\times 10^{-2})** (caught). Brownian variance: empirical (\\langle x^2 \\rangle(T) = 4.975) vs analytic (2DT = 5.000) for (N = 5000) paths (z-score (-0.25), within (\\pm 3\\sigma)); a (10%) noise-amplitude scaling bug produces (z = +10.2) (caught at the (3\\sigma) envelope). LC oscillator under BDF: energy drift **(1.94 \\times 10^{-8})** over 10 oscillation periods (the BDF (10^{-4}) envelope is loose because BDF dissipates). The **pedagogical climax**: in section 4 a small visible-bug, invisible-trajectory comparison forces the reader to internalise why visual inspection is insufficient and the property test is strictly stronger. Closes with the "where to put these tests in your project" pattern (`test//test_conservation.py`), five failure modes (tolerance-too-tight, statistical false positives, damped-system mis-test, chaotic-system non-invariants, endpoint-only vs max-over-trajectory), and five exercises (anti-conservation test for a damped SHO, charge conservation on the LC, `jax.grad` of drift w.r.t. solver tolerance, Lorenz divergence as a flow invariant, planar-pendulum constraint residual under Baumgarte/SSP/none). The marketing wedge: the property-test framework is to physical validity what gradient-correctness CI is to autodiff — defense-in-depth against silent failure modes the simulator cannot detect on its own. Runtime: ~6 s end-to-end. **Cells**: 21 code + 29 markdown = 50 total, 4 matplotlib figures. Showcases the conservation-test framework (`test/conservation/_framework.py::assert_conserved`), the acausal LC oscillator under BDF, and the stochastic-source / probability-mass conservation pattern. ### [Differentiable acausal DAEs: learning unmodeled drag inside a constrained pendulum](https://py.jaxonomy.com/examples/neural_dae_pendulum_drag/index.md) The differentiable-acausal capability in one notebook. A `PlanarPendulum` on a rigid link is a holonomic constraint (x^2+y^2=L^2) — an index-3 DAE compiled through Pantelides index reduction — and the real rig has an **unmodeled quadratic aerodynamic drag (-c,v,|v|)** the modeler never wrote. We bolt a small MLP (f\_{NN}(v;\\boldsymbol\\theta)) onto the compiled DAE's differential rows as a learned correction via the phase-2 `NeuralDAEBlock` (`targets=[(pend, "v")]`, gather-in/scatter-out, injected *after* index reduction so the non-symbolic neural term never enters the symbolic Pantelides path), then fit (\\boldsymbol\\theta) by gradient descent **through the BDF-DAE adjoint** — `jax.grad` flows from a terminal-state loss, through three implicit constrained-DAE integrations, into the network weights. This is the wedge no single tool spans: Modelica expresses the constraint but can't autodiff through it; causal UDE tools (Neuromancer, DiffEqFlux) autodiff but can't express the algebraic constraint coupling. **Headline live numbers** (no checkpoint — the fit runs live to demonstrate the adjoint): the multi-horizon terminal-state loss drops from **(4.97\\times10^{-2}) to (1.85\\times10^{-4}), a ~269× reduction in 25 Adam steps** (~7.3 s/grad through 3 BDF-DAE solves, ~3 min fit on laptop CPU); the fitted model's angle RMS vs the damped truth is **0.0118 rad vs 0.1088 rad** for the no-drag baseline (9.2× better) over a 6 s window it was never trained on. The **conservation validation** confirms the compile is correct: total mechanical energy is conserved without drag and decays monotonically with it, and the holonomic constraint residual (|x^2+y^2-L^2|) holds at **(\\sim10^{-14}) m²** (the DAE-projection invariant). Closes with an **honest identifiability beat** — the learned (f\_{NN}(v)) tracks the true drag inside the velocity band the trajectory explored (RMS 0.07 m/s²) but diverges outside it (RMS 1.35 m/s²), because you can only learn the dynamics where the data took you. Surfaces the genuine sharp edge it had to work around: index reduction's choice of which symmetric velocity ((v_x) vs (v_y)) survives as the *differential* state is `PYTHONHASHSEED`-dependent, so the demo resolves it at runtime by following `alias_map` into `sed.x` (filed as a follow-up finding); plus the `recorded_signals` ⊥ autodiff constraint that forces the terminal-state loss, the `buffer_length` overflow on the visualization path, and the Adam-overshoot early-stop. Five exercises (linear-drag swap, harder excitation to widen the identifiable band, hand-derive (\\lambda) from the acceleration-level constraint, compose a second block on the position row, symbolic-regress (f\_{NN}) back to closed form). **Runtime**: ~4–5 min end-to-end (live fit). Showcases the phase-2 `NeuralDAEBlock` + `AcausalDiagram.add_neural_correction_block`, the acausal Pantelides → BDF pipeline, and the autodiff-through-the-DAE-adjoint story; sibling to the causal-ODE [UDE notebook](https://py.jaxonomy.com/examples/ude_and_sr_lotka_volterra/index.md) and the [conservation-laws-as-CI notebook](https://py.jaxonomy.com/examples/conservation_laws_as_ci/index.md). ### [Stiff chemistry: the Robertson problem under BDF](https://py.jaxonomy.com/examples/stiff_robertson_bdf/index.md) The canonical stiff-ODE test case in numerical analysis (Robertson 1966; Hairer & Wanner 1996 Vol II §IV.1; Shampine 1994 Ch. 6): three coupled species with rate constants (k_1=0.04), (k_2=3 \\times 10^7), (k_3=10^4) spanning 9 orders of magnitude. Walks through what "stiff" actually means (the *stability* step is much smaller than the *accuracy* step — Curtiss & Hirschfelder 1952), why explicit Dopri5 catastrophically over-resolves the fast B-mode, and why implicit BDF cruises to steady state in a few hundred adaptive steps. **Headline numbers**: BDF integrates the full (t \\in [0, 10^{11}]) Robertson horizon in **~0.4 s** wall-time (854 recorded samples, mass conservation (|y_1+y_2+y_3-1| < 10^{-15})); Dopri5 at (T=1) s already needs 713 minor steps (~0.13 s wall), giving a linear extrapolation of ~600 *years* of wall-time to clear the full horizon. The stiffness crossover where Dopri5 step count first exceeds BDF sits at (k_2 \\approx 10^7) (ratio (k_2/k_1 \\sim 2.5 \\times 10^8)). Closes with a Jacobian-spectrum visualization (`jax.jacrev` on the RHS, six probe points, spectral ratio from 1 to (10^{19})), the canonical Robertson log-log trajectory plot, BDF-vs-Dopri5 horizon scaling (BDF flat across 11 decades, Dopri5 linear in (T)), (k_2)-stiffness sweep, a one-paragraph "why BDF works" explanation with the Newton-iteration math, three practical rules of thumb, a central-difference gradient cross-check ((\\partial y_3(10^4) / \\partial k_1 = 3.438)), and five exercises (less-stiff Dopri5 sweep, source-flux extension, multi-parameter gradient, HIRES/OREGO port, Radau IIA implementation). **3 follow-up findings**: (i) `int_time_scale` not auto-promoted for (t > 9 \\times 10^6) s; (ii) `max_major_step_length` is JIT-static so horizon sweeps re-compile per iteration; (iii) BDF's recorded `n_samples` measures major-step recording cadence, not solver step count — solver-agnostic step counting requires instrumentation. **Cells**: 12 code + 22 markdown = 34 total, 4 matplotlib figures. **Runtime**: ~10 s on warm JIT cache, ~60-90 s cold. Showcases BDF as a first-class jaxonomy solver + the canonical stiff-ODE diagnosis workflow + the gradient-through-implicit-solver story. ### [Container blocks tour: `ForEach`, `EnabledSubsystem`, `TriggeredSubsystem`](https://py.jaxonomy.com/examples/container_blocks_tour/index.md) Three canonical container-block patterns end-to-end on three real-engineering motivations: (a) `ForEach`-vectorised 100-cell battery sweep with per-cell capacity/resistance tolerance, (b) `EnabledSubsystem`-wrapped 4-wheel ADCS where one reaction wheel fails mid-mission and the remaining three redistribute disturbance-rejection torque, (c) `TriggeredSubsystem` gating a Kalman measurement update on the rising edge of a 1 Hz star-tracker pulse. Each pattern closes with a `jax.grad` through the container boundary — the autodiff wedge over Simulink. **Headline numbers**: (A) per-cell `∂(mean SOC)/∂C_i = 1.528e-3` matches the closed-form Coulomb counter to 0.01% in one backward pass; (B) post-failure per-wheel torque grows from 10.5 mN·m to 13.3 mN·m (ratio 1.27, theoretical 4/3 = 1.333), failed wheel held at exactly 0 mN·m by `EnabledSubsystem(mode="reset")`; (C) `d(ISE)/dK = −7.39` autodiff matches central-difference within 0.10%. Honest about the phase-1 `TriggeredSubsystem` running its child on every sample step (gates the latch, not the kernel) and the `requires_inputs=False` declaration needed to break the cycle checker on no-feedthrough output ports — both written up as follow-up findings. Runtime: ~7 s end-to-end. Showcases `EnabledSubsystem` + `TriggeredSubsystem` + `ForEach`. ### [Multi-domain HVAC: a heat pump heating a Stockholm apartment](https://py.jaxonomy.com/examples/multi_domain_hvac/index.md) A 100 m² apartment in Stockholm in mid-January under an air-source heat pump whose COP swings 1.5x across the diurnal ambient cycle. Builds the two-state RC envelope as an acausal thermal network (`HeatCapacitor` + `Insulator` + `TemperatureSource` + `HeatflowSource(enable_heat_port=True)`), wraps the heat pump as a causal `LeafSystem` whose temperature-dependent COP comes from an `npa.interp` lookup table, and closes the loop with a discrete-time PI controller updating at 1/60 Hz. **Headline unit-safety beat**: a deliberate `m^3/s` (volumetric refrigerant flow) -> `kg/s` (mass flow) wiring mistake at the causal/acausal seam raises `UnitMismatchError` at `DiagramBuilder.connect()` time, before any kernel launches; the fix is an explicit density-conversion `Gain(rho)` with the algebra verified on both ends. Honest about the units follow-up (canonical acausal units exposed but the Pantelides pass does not yet consume them — filed as a follow-up finding); the causal-seam check is the production-grade boundary that catches the bug today. 24 hours integrate in ~1 s under BDF (~16k samples — see ring-buffer note), tracking RMSE drops to mK during occupied hours, total electricity bill comes in at **12.97 kWh / day**, and a four-parameter central-difference sensitivity analysis ranks **wall R-value (\\gg) COP (>) set-back depth (\\gg) thermal mass** for retrofit ROI (the autodiff workaround caveat from the battery tutorial applies). Closes with `analyze_saturation` + `analyze_control_oscillation` (both silent), failure-modes section (deep-cold backup heater unmodelled, manufacturer-specific COP curve, two-state envelope vs ASHRAE 90.1 multi-zone, simplified refrigerant cycle, mean-ambient wall coupling), and five exercises (COP-curve swap, dawn pre-heat, multi-zone, optimal schedule, hybrid heater). Showcases unit-safety at the cross-domain seam + acausal thermal library + LeafSystem composition + autodiff sensitivities. ### [FMI export round-trip: build a controller in jaxonomy, co-simulate with an external tool](https://py.jaxonomy.com/examples/fmi_export_roundtrip/index.md) Build a discrete-time PI controller as a `DiagramBuilder` composition, export it as a binary FMI 2.0 Co-Simulation `.fmu` via `jaxonomy.library.fmu_export.build_fmu`, then re-import the binary three ways: as a `jaxonomy.library.ModelicaFMU` block inside a fresh diagram (Architecture B), and via raw `fmpy.fmi2.FMU2Slave` calls orchestrated by a hand-rolled Python loop (Architecture C). Headlines: the manual `fmpy` orchestration agrees with a pure-Python reference loop on the same difference equations to **max abs error 7.8e-4 on the plant state and 4.0e-3 on the control output over 400 steps** — the residual is FMI 2.0 §4.2.4's first-step initialization protocol, not a numerical artifact (the per-step error drops to ~1e-4 on x and ~1.6e-5 on u after k=1). Demonstrates `JaxonomyDiagramSlave` wiring, the auto-exposed `Constant`-block input convention, and the structural-XML round-trip over the `modelDescription.xml` jaxonomy generated, which is what an FMI host parses at import time (running it is a separate question — see the tool-coupling note in `KNOWN_GAPS.md`). The publication/fast-mode pattern caches the heavy `ModelicaFMU` JIT compile (~22 s) under `media/fmi_export_publication.npz` so the reader's notebook runs in ~3 seconds. Honest about three sharp edges: the `ModelicaFMU` block's periodic-update `offset=dt` introduces a 1-sample phase lag relative to in-process discrete blocks (Architecture A vs B mismatch in the transient), pythonfmu's embedded-Python single-instance dylib limit forces a `subprocess` for the manual loop, and the FMU boundary is not differentiable (FMI export is for *delivery*, not for *training*). One more honest number: each co-simulation `doStep` re-enters `jaxonomy.simulate` per segment, costing ~100+ ms per step versus ~tens of µs per step in-process on a small system — the FMI boundary is for validation and interop, not for training loops. Scope: exports are FMI 2.0 co-simulation and tool-coupled — the importing side needs Python with jaxonomy on its path, and a non-Python master needs the wrapper from `scripts/build_pythonfmu_wrapper.sh`, after which `fmusim` runs them. OpenModelica imports model exchange only, so it cannot host a jaxonomy export at all. Showcases binary FMU export, the Darwin pythonfmu wrapper, auto-exposed Constant block inputs, and FMI 2.0/3.0 import via `ModelicaFMU`. ### [Physics-informed learning across stacks, part 1: policy export via ONNX](https://py.jaxonomy.com/examples/pinn_across_stacks_part_1_policy_export/index.md) First of a three-part series pairing jaxonomy with [NEUROMANCER](https://github.com/pnnl/neuromancer) (PNNL's PyTorch differentiable-programming/DPC library) around one physical system — a two-tank pump-and-valve network from NEUROMANCER's own `psl` benchmark set. This part trains a bounded-MLP DPC policy in NEUROMANCER, exports it to ONNX, and embeds it in a jaxonomy diagram via the differentiable `ONNXJax` block plus a `ZeroOrderHold` (the load-bearing detail: without it the policy is re-evaluated at every RK4 stage instead of once per sample). **Headline numbers** (executed outputs): torch↔onnxruntime export parity **2.68e-7**; closed-loop parity against the exporting framework **3.84e-8 over 400 steps**; and the payoff neither stack does alone — a `jit`+`vmap` Monte-Carlo robustness sweep of 512 plant-parameter samples × 200-step rollouts in **~7 ms steady-state**, with the one fragile sample explained exactly by an analytic pump-authority boundary overlaid on the scatter. Ships the trained policy artifacts in `media/` so the notebook runs in ~5 s without torch installed; with NEUROMANCER present the training reproduces bit-for-bit. Showcases `ONNXJax`, `ZeroOrderHold`, adapter `LeafSystem`s, and `jaxonomy.diagnostics` on the closed loop. ### [Physics-informed learning across stacks, part 2: a neural DAE and gradients across the framework boundary](https://py.jaxonomy.com/examples/pinn_across_stacks_part_2_neural_dae/index.md) The series centerpiece. Builds the tank network as it really is — reservoir, pump, three-way manifold, two gravity tanks — by writing two custom acausal hydraulic components (~30 lines each) and letting `AcausalCompiler` produce a genuine semi-explicit DAE: 2 differential states plus **18 algebraic unknowns** (the manifold pressure among them), integrated by BDF with a singular mass matrix. Then gradients cross the stack boundary in both directions. Direction 1: the tank1→tank2 orifice is secretly clogged (55% of nominal); a physics-structured neural correction — one scalar flow, scattered with the fixed −1/ρA₁ : +1/ρA₂ ratio — is trained by `jax.value_and_grad` **through the implicit BDF solve** and recovers the true flow-deficit law to **1.43% relative error, correlation 0.996** (AD matches finite differences to ~5 significant digits; the unstructured 2-in/2-out variant drops the loss ~90× yet fails to identify the residual — the rank-1 physics prior is the lesson). Direction 2: the DAE plant becomes a NEUROMANCER `Node` via a `torch.autograd.Function` whose backward is the JAX VJP over zero-copy dlpack — steady batched DAE steps of **~3 ms** — and NEUROMANCER's unmodified Trainer optimizes a policy through it, with honest hold-state/NaN-sanitization guard counters (24 events, all one knife-edge point, root-caused to projection non-convergence). Cites Koch, Shapiro, Sharma, Vrabie & Drgoňa (CDC 2025) for the operator-splitting approach to the same problem class. Showcases custom acausal components, `add_neural_correction`, `project_constraints`, `dae_initial_projection`, and torch↔JAX interop. ### [Physics-informed learning across stacks, part 3: the plant as an FMI co-simulation FMU](https://py.jaxonomy.com/examples/pinn_across_stacks_part_3_fmi_cosim/index.md) Closes the series at the tool-neutral boundary: the two-tank plant is exported as a binary FMI 2.0 Co-Simulation FMU via `build_fmu` (passes `fmpy.validate_fmu` with **zero findings**, tri-platform binaries) and driven in lockstep by the part-1 policy — rebuilt master-side in plain torch from the weight file, parity-checked against the shipped ONNX at 4.2e-7. **Headline numbers**: 200-step lockstep loop matches the fully in-process closed loop to **max |Δh| = 1.783e-8 m**, with the residual explained by solver independence (adaptive integration inside the FMU vs a single RK4 step outside) — which is precisely what co-simulation buys. Honest cost accounting: ~140 ms per `doStep` through the FMI boundary vs ~69 µs in-process (~2,000×) — the boundary is for validation, certification workflows, and HIL prep, not training loops (part 2 exists because gradients don't cross FMI). Documents the slave sharp edges (output priming, `Constant`-block input convention, one-instance-per-process) and runs `jaxonomy.diagnostics` on both actuators. Showcases `JaxonomyDiagramSlave`, `build_fmu`, `fmpy` mastering, and the gained/lost trade-offs of the FMI boundary. ### [Hybrid ML + physics: dropping a pre-trained predictor into a control loop](https://py.jaxonomy.com/examples/hybrid_ml_physics_predictor/index.md) The drop-in story for teams who own a PyTorch / TensorFlow training pipeline. Trains a small residual MLP on a damped pendulum with hidden Coulomb friction, loads it as a `jaxonomy.library.PyTorch` predictor block when torch is installed (falls back to `TensorFlow`, then to a JAX/Equinox MLP — the public CI configuration), and closes a PD loop around three predictors: pure-physics (no Coulomb), pure-ML (no inductive bias), and hybrid (physics + learned residual). **Headline three-way RMSE comparison**: in-distribution pure-physics 0.0146 rad, pure-ML 0.0008 rad, **hybrid 0.0004 rad** (35.5x better than pure-physics); out-of-distribution (release from (\\theta_0 = 2.5) rad, 5x training envelope) pure-physics 0.0292 rad, pure-ML 0.0330 rad (catastrophic extrapolation), **hybrid 0.0019 rad** (~15x better than either alone — the structural argument for the hybrid). The autodiff bonus beat takes `jax.grad` of closed-loop tracking ISE w.r.t. PD gain `Kp` against the (differentiable) hybrid plant, cross-checks against central-difference to **0.0000% relative error**. Honest about the limit: `jax.grad` does NOT flow through the `PyTorch` / `TensorFlow` predictor blocks (same `jax.pure_callback` boundary that breaks gradient flow through `ModelicaFMU`); the workaround is to fine-tune against an in-process JAX surrogate of the predictor and validate on the real predictor in a second pass. The publication/fast-mode pattern caches the 4000-epoch training + 4 closed-loop simulations under `media/hybrid_ml_physics_publication.npz` so the notebook runs in ~5 s; in fast mode the inline retrain runs at 1500 epochs in ~30 s. Marketing wedge: **"trained your model in PyTorch on your team's data? Drop it in."** Closes Wave 4. Showcases the `PyTorch` / `TensorFlow` predictor blocks + the `jaxonomy.library.MLP` Equinox block as the JAX-native equivalent + the hybrid-ML-physics pattern from Rackauckas et al. 2020 and Karniadakis et al. 2021. ### [Co-simulating a Modelica plant FMU under a jaxonomy controller](https://py.jaxonomy.com/examples/openmodelica_plant_fmu_cosim/index.md) The **inverse direction** of `fmi_export_roundtrip.ipynb`: instead of exporting a jaxonomy controller as a binary FMU, *import* a Modelica-style plant FMU as the closed-loop plant under a jaxonomy PI controller. The combined story (controller-out + plant-in) covers both directions of the boundary; the import direction is the broadly applicable one, since jaxonomy consumes co-simulation and model-exchange FMUs from any exporter while its own exports carry the tool-coupling limit documented in `KNOWN_GAPS.md`. Builds the plant — a damped 2nd-order mass-spring ((m\\ddot{x} + c\\dot{x} + kx = F), (\\omega_n = 1) rad/s, (\\zeta = 0.25)) — as a `JaxonomyDiagramSlave`-wrapped `LeafSystem` and `build_fmu`'s it to a binary FMI 2.0 Co-Simulation `.fmu` (~21 ms / 899 KiB on darwin); then closes the loop with the same jaxonomy PI controller around it via `ModelicaFMU`. **Headline numbers** (all live or NPZ-cached): closed-loop position trajectory of the FMU-plant Architecture B matches the in-process Architecture A to **max abs error 2.46e-2 m over the full 12 s horizon**, dominated by the documented `ModelicaFMU` `offset=dt` first-step phase lag. The **autodiff bonus beat** — the marketing wedge over Simulink's FMI Import block — uses the cost-as-`Integrator` pattern to take `jax.grad` of the closed-loop tracking ISE w.r.t. the controller gains: `dJ/dKp = -0.2437` and `dJ/dKi = -1.4586`, cross-checked against central-difference to **0.179% and 0.018% relative error respectively**. Honest about the limit: `jax.grad` does NOT flow through the `ModelicaFMU` block (the FMU is a host C call); the autodiff path therefore runs against an in-process surrogate of the plant, and the tuned gains are then validated on the FMU plant in a second simulation step. The publication/fast-mode pattern caches the heavy Architecture B JIT-compile (~65 s offline) under `media/openmodelica_plant_fmu_publication.npz` so the reader's notebook runs in ~5 seconds. Three production routes documented: Reference-FMU corpus (Dahlquist / VanDerPol / BouncingBall — exercise 1), OpenModelica `.mo` → `buildModelFMU` export (exercise 2; `translateModelFMU` generates sources without packaging an `.fmu` on OpenModelica 1.27), and the always-runnable pythonfmu-built synthetic plant (the path the tutorial takes). Showcases `ModelicaFMU` import + `JaxonomyDiagramSlave` / `build_fmu` + the autodiff-across-the-FMU-boundary wedge. ### [Universal Differential Equations (UDEs) and symbolic regression (SR)](https://py.jaxonomy.com/examples/ude_and_sr_lotka_volterra/index.md) Demonstrates training a Universal Differential Equation (UDE) to fit the observations produced by the Lotka-Volterra predator-prey system. Subsequently, the UDE is symbolically regressed to learn a closed-form model. ### [Nonlinear MPC](#nmpc) See thematic series on modeling and control of 3D quadcopter [below](#nmpc), which showcases trajectory tracking by nonlinear MPC. ### [Wind turbine control: MPPT below rated, blade-pitch regulation above rated](https://py.jaxonomy.com/examples/dfig_wind_turbine/index.md) A variable-speed ~1.75 MW turbine built from primitive `LeafSystem` blocks: a standard (C_p(\\lambda,\\beta)) aerodynamic map, a two-mass torsional drivetrain, and a controller that does the two jobs wind-turbine control is made of. **Below rated wind** the maximum-power-point-tracking torque law (T_e=K\_{opt}\\omega_g^2) sits the rotor at the aerodynamic optimum — the validation gate confirms the steady tip-speed ratio reaches (\\lambda\_{opt}=8.10) and the captured (C_p) reaches (C\_{p,max}=0.480) exactly. **Above rated wind** the torque saturates and an anti-windup PI pitch loop feathers the blades to hold the generator within 2% of rated speed (1.79 MW). The MPPT gain is *derived* from the aerodynamics and gearbox, not guessed; a failure-mode cell shows a 2× mistuned gain costs 25% of capture. Notes the slow-rotor `buffer_length` trap (the adaptive recorder silently truncates to the tail if undersized). ### [HL-20 lifting body: an angle-of-attack-hold glide autopilot](https://py.jaxonomy.com/examples/hl20_glide_autopilot/index.md) The NASA HL-20's longitudinal flight dynamics (a 3-DOF model: range, altitude, pitch) with a representative analytical stability-derivative aero set, flown by an AoA-hold autopilot ((\\delta_e=\\delta\_{e,trim}+K\_\\alpha(\\alpha\_{cmd}-\\alpha)+K_q q)). The lesson is the lifting body's low lift-to-drag ratio ((L/D\_{max}\\approx5.8)): the shallowest steady glide it can hold is (\\gamma=-\\arctan(1/(L/D\_{max}))\\approx-10^\\circ). Commanding best-L/D angle of attack, the body settles at the analytic trim angle ((\\gamma\\approx-9.1^\\circ), glide ratio 6.25) — validated against (-\\arctan(C_D/C_L)) averaged over the lightly-damped phugoid. The failure-mode cell makes the point concrete: a naive controller commanding a shallow (-3^\\circ) glide slope rails the elevator at (-25^\\circ) and runs the angle of attack to (66^\\circ), because (-3^\\circ) is aerodynamically infeasible. ### [Differentiable vehicle handling: yaw-plane dynamics and design sensitivity](https://py.jaxonomy.com/examples/vehicle_handling_autodiff/index.md) A yaw-plane (single-track) handling model with a Pacejka magic-formula tyre, validated against classical understeer theory — at low lateral acceleration the nonlinear step-steer yaw rate matches (r/\\delta=U/(L+K\_{us}U^2)) to 1.3%. The headline is what a JAX-native simulator gives a chassis engineer: the steady yaw rate's sensitivity to a design knob (rear tyre stiffness) is differentiated **through the whole simulation** — `simulate_jacfwd` (forward) and `jax.grad` (reverse) both return (\\partial r/\\partial k_r=-0.0824), matching central differences to five significant figures, and the gradient is drawn as a tangent on the design sweep. The same gradient is a stability early-warning: softening the rear drives (K\_{us}) negative into oversteer, where the step-steer yaw rate diverges. ### [Satellite ADCS: attitude control with an EKF and reaction wheels](https://py.jaxonomy.com/examples/aerospace_adcs_ekf_wheels/index.md) A spacecraft attitude-determination-and-control loop: reaction-wheel actuation, an Extended Kalman Filter for attitude/rate estimation from noisy sensors, and momentum management under disturbance torques. ### [An artificial pancreas: closed-loop glucose control under uncertainty](https://py.jaxonomy.com/examples/artificial_pancreas_mpc/index.md) Model-predictive insulin dosing around a glucose-insulin plant, closing the loop under meal disturbances and parameter uncertainty — a safety-critical biomedical control example. ### [A grid-forming microgrid: droop control and fault ride-through](https://py.jaxonomy.com/examples/grid_forming_microgrid/index.md) Grid-forming inverter droop control on an islanded microgrid, including fault ride-through behaviour — an example from the power-electronics / energy-systems domain. ### [Reinforcement learning environment from a jaxonomy diagram](https://py.jaxonomy.com/examples/rl_environment_from_diagram/index.md) Wrap a `DiagramBuilder` model as a vectorised RL environment, using JAX `vmap`/`jit` to batch rollouts for policy training. ### [Differentiable digital audio: parametric EQ + dynamic-range compressor](https://py.jaxonomy.com/examples/differentiable_audio_dsp/index.md) A parametric equaliser and a dynamic-range compressor built as jaxonomy blocks, with `jax.grad` flowing through the DSP chain for gradient-based audio parameter fitting. ### [DAE constraint projection on a 1-hour pendulum simulation](https://py.jaxonomy.com/examples/dae_projection_pendulum/index.md) Long-horizon constrained-DAE integration with constraint projection, showing how the holonomic constraint residual is held near machine precision over an hour of simulated time. ## Thematic examples ### Battery modeling 1. [Equivalent circuit model (ECM) for a battery](https://py.jaxonomy.com/examples/battery_part_1_ecm_model/index.md) 1. [ECM parameter estimation: synthetic data](https://py.jaxonomy.com/examples/battery_part_2_parameter_estimation_synthetic_data/index.md) 1. [ECM parameter estimation: experimental data](https://py.jaxonomy.com/examples/battery_part_3_parameter_estimation_real_data/index.md) 1. [Data-driven modeling: Dynamic Mode Decomposition (DMD)](https://py.jaxonomy.com/examples/battery_part_4_data_driven_models_DMDc/index.md) 1. [Data-driven modeling: Extended DMD](https://py.jaxonomy.com/examples/battery_part_5_data_driven_models_eDMDc/index.md) 1. [Data-driven modeling: SINDy with control](https://py.jaxonomy.com/examples/battery_part_6_data_driven_models_SINDyc/index.md) 1. [Data-driven modeling: Neural Networks](https://py.jaxonomy.com/examples/battery_part_7_data_driven_models_Neural_Networks/index.md) 1. [Pack-level modeling: cell, module, and thermally-coupled pack](https://py.jaxonomy.com/examples/battery_pack_thermal/index.md) — compose `BatteryCellECM` cells into series modules and a parallel pack via the acausal electrical layer, wrap each cell in a `HeatCapacitor` with `Insulator`-based lateral conduction and `TemperatureSource` ambient boundaries to couple electrical and thermal domains, and run projected gradient descent on a fixed cooling budget (gradient via parameter-rebinding finite differences, since the BDF-DAE adjoint returns the wrong sign on this DAE — filed as a follow-up finding) to redistribute cooling capacity onto the bottlenecked module. Showcases `BatteryCellECM` + `BatteryCellTabular` and the cross-domain acausal libraries; the marketing wedge is production-grade pack modelling in open source, no toolbox license required. 1. [Scaling a battery pack from 8 cells to 100,000](https://py.jaxonomy.com/examples/battery_pack_10k_scaling/index.md) — takes the same acausal-ECM pack construction to a 100k-cell pack, documenting what scales cleanly under JAX (`vmap`ed cell dynamics, JIT-compiled kernels) and what breaks (compile time, memory footprint, DAE solve cost) as the cell count grows five orders of magnitude. ### Electric drives (PMSM) A six-part series taking an interior permanent-magnet synchronous motor from a differentiable electromagnetic model to an Arm Cortex-M binary — modelling, control, thermal derating, calibration-from-data, robustness, and embedded deployment. 1. [Modelling an interior PMSM you can differentiate](https://py.jaxonomy.com/examples/motor_part_1_pmsm_modeling/index.md) 1. [Field-oriented control](https://py.jaxonomy.com/examples/motor_part_2_field_oriented_control/index.md) 1. [Thermal coupling and torque derating](https://py.jaxonomy.com/examples/motor_part_3_thermal_and_derating/index.md) 1. [Calibrating the machine from data](https://py.jaxonomy.com/examples/motor_part_4_system_identification/index.md) 1. [Design margins under uncertainty](https://py.jaxonomy.com/examples/motor_part_5_design_margins/index.md) 1. [From JAX to an Arm Cortex-M binary](https://py.jaxonomy.com/examples/motor_part_6_embedded_deployment/index.md) ### 3D quadcopter modeling and control 1. [3D quadcopter modelling](https://py.jaxonomy.com/examples/01_quadcopter_modelling/index.md) 1. [Trajectory generation through differentially flat outputs](https://py.jaxonomy.com/examples/02_quadcopter_trajectory_generation/index.md) 1. [Control with nonlinear MPC](https://py.jaxonomy.com/examples/03_quadcopter_nonlinear_mpc/index.md) ### Quanser Qube Servo hardware control 1. [Qube Servo modeling](https://py.jaxonomy.com/examples/quanser/01-plant-model/index.md) 1. [Linear control](https://py.jaxonomy.com/examples/quanser/02-lqg/index.md) 1. [Nonlinear swing-up control](https://py.jaxonomy.com/examples/quanser/03-energy-shaping/index.md) 1. [Trajectory optimization](https://py.jaxonomy.com/examples/quanser/04-trajopt/index.md) 1. [Neural network control](https://py.jaxonomy.com/examples/quanser/05-nn-control/index.md) ### Returning rocket booster (Falcon-9-class) A six-part series modelling the propulsive landing of a returning rocket booster from atmospheric apogee to a soft touchdown on a drone-ship pad. Each part progressively adds fidelity — atmospheric aerodynamics and multi-phase guidance, a multi-engine cluster with bandwidth-limited actuators and variable inertia, noisy sensors with Kalman filtering — and culminates in cinematic MuJoCo rendering. Part 6 fills the GNC validation tooling that production aerospace work expects: linearised Bode / Nyquist / eigenvalue analysis at hover, a 1000-trial Monte Carlo dispersion sweep, the autodiff-vs-finite-difference timing comparison that quantifies what `jaxonomy` uniquely buys you over MATLAB / Modelica / hand-rolled scipy, plus prose on real-time scheduling, lossless convex powered-descent guidance, and the ITAR realities of production flight software. 1. [6-DOF dynamics and open-loop trajectory optimisation](https://py.jaxonomy.com/examples/booster_part_1_modeling/index.md) 1. [Closed-loop MPC with MuJoCo rendering](https://py.jaxonomy.com/examples/booster_part_2_mpc_and_render/index.md) 1. [Atmosphere, multi-phase guidance, and autodiff parameter tuning](https://py.jaxonomy.com/examples/booster_part_3_atmosphere_and_phases/index.md) 1. [High-fidelity propulsion: engine cluster, variable inertia, actuator dynamics, engine-out](https://py.jaxonomy.com/examples/booster_part_4_high_fidelity_propulsion/index.md) 1. [Imperfect sensing and EKF state estimation](https://py.jaxonomy.com/examples/booster_part_5_sensing_and_estimation/index.md) 1. [GNC validation, analysis, and the autodiff advantage](https://py.jaxonomy.com/examples/booster_part_6_gnc_validation_and_analysis/index.md) Bonus: the [cinematic Falcon-9-class landing demo](https://py.jaxonomy.com/examples/media/booster_landing_cinematic.mp4) — a 14-second 1280×720 video produced by [`render_booster.py`](https://py.jaxonomy.com/examples/media/render_booster/index.md), which solves Part 1's trajectory optimisation and renders it through a richly-instrumented MuJoCo scene (drone-ship pad, ocean, sun + fill lighting, four landing legs, four grid fins, four-layer plume with visible Mach diamonds, and a hinged engine assembly that visibly tilts on the gimbal command — the visceral cue that there is *active control* of the descent). ### F1 race car (6 parts) A six-part series taking a 2022+-era ground-effect Formula 1 car from a credible whole-lap simulator (Part 1) to a lap-time-aware aero shape-optimisation *architecture* that co-simulates an external CFD adjoint solver (Parts 5–6, demonstrated live with an in-process placeholder solver — the full SU2 run needs an external toolchain; see the note under Series C). The property that runs end-to-end: *every layer of the stack — chassis, tire, powertrain, driver, aero map, CFD — is differentiable, so `jax.grad(lap_time)` w.r.t. any setup or aero parameter is one backward pass*. That end-to-end differentiability is what sets this apart from conventional lap-time simulators, which fall back on finite differences that scale poorly as the parameter count grows. The series is structured as three sub-arcs: - **Series A — Lap-time-aware setup optimisation** (pure jaxonomy). Part 1 builds the bicycle + Pacejka + powertrain + QSS hot-lap driver, validates against analytic cornering equilibrium via `findop`, and renders the lap in MuJoCo. Part 2 wraps it as `lap_time(setup) → float` and takes `jax.grad` through it: 20 Adam steps drop the lap by ~0.4 s from a deliberately-bad baseline. - **Series B — CFD budget allocation under the FIA ATR** (synthetic aero map). Part 3 fits a noisy 5-D aero map from sparse CFD samples. Part 4 runs Sobol-decomposition over the map to allocate the FIA-ATR-limited CFD budget (the leader gets the least testing-hours; the optimisation question is *where* to spend them). - **Series C — Co-sim with SU2 for lap-time-aware shape optimisation** (3D adjoint CFD external). Part 5 closes the loop on a NACA airfoil for the tractable proof. Part 6 swaps it for a parametrised rear-wing assembly on the Perrinn 424 / DrivAer baseline, with PyVista + Blender Cycles for the hero shape-optimisation MP4. **Note:** Parts 5–6 demonstrate the co-sim *architecture* live with an in-process solver; the real SU2 RANS + discrete-adjoint results and the Blender hero render require external executables (`pysu2`, `SU2_AD`, OpenVSP, ParaView, Blender) not bundled with the repo, so those specific CFD figures are not yet included (each notebook flags this to the reader). - [Lap-time simulator: vehicle dynamics, Pacejka tire, powertrain](https://py.jaxonomy.com/examples/f1_part_1_lap_time_simulator/index.md) — a 4-state longitudinal–lateral–yaw bicycle with a Pacejka 5.2 magic-formula tire and a friction-ellipse closure; a `LookupTable1d` engine map + 7-speed gearbox with finite-time shifts; a synthetic 4-corner GP-style track expressed as (\\kappa(s)); a quasi-steady-state hot-lap driver that respects the friction-ellipse forward and backward; integration-to-steady-state validation against the analytic cornering equilibrium (V^2 = \\mu\_{\\text{eff}}(V),g,R) (agreement within 1.8% on a (R = 100) m corner; `findop` filed as failing on this stiff Jacobian + passive-integrators case as a follow-up finding); and a MuJoCo lap render with overlaid HUD. The whole stack is `jit`- and `grad`-able, sized for Part 2's setup-optimisation gradient. - [Setup optimisation via `jax.grad` through lap time](https://py.jaxonomy.com/examples/f1_part_2_setup_optimization/index.md) — wraps Part 1's LTS as `lap_time(setup) → float` via the cost-as-`Integrator` pattern under `SimulatorOptions(enable_autodiff=True)` and `with_parameters` on the car block. Headline live beat: `jax.grad(forward)(SETUP_BASELINE)` prints all 8 setup sensitivities in one backward pass at a tractable `T_END = 15` s horizon; FD validation on 2 components (`k_f`, `h_f`) cross-checks within 5%. The compute-heavy beats (29-iter L-BFGS-B from a deliberately-bad starting setup, 3^6 = 729-point FD-vs-AD grid head-to-head, 16-start LHS multi-start) follow the *publication / fast-mode pattern*: default `MODE = "publication"` loads from `media/f1_part_2_publication.npz` (currently placeholder, awaiting `media/f1_part_2_publication_offline.py`'s offline run at full `T_END = 60` fidelity — a lengthy offline run that did not complete on the reference hardware); set `MODE = "fast"` for a 3-step projected gradient descent verifier in ~30 s. The full-fidelity optimisation figures await a completed offline run (`media/f1_part_2_publication_offline.py` at `T_END = 60`, which did not finish on the reference hardware), so the notebook ships the fast-mode verifier above rather than unverified headline numbers. MuJoCo before/after render: two cars (baseline red + optimised yellow) through corner 1 with HUD overlay. - [Fitting a 5-D aero map from sparse CFD samples](https://py.jaxonomy.com/examples/f1_part_3_aero_map_fitting/index.md) — kicks off Series B. Synthesises a 5-D ground-truth aero map ((h_F, h_R, \\phi, \\beta, \\delta) \\to (C_L A, C_D A, x\_{\\text{CoP}})) from a closed-form drag-bucket + downforce-vs-rake + asymmetric yaw response, adds heteroscedastic CFD noise (~2% on (C_L), ~3% on (C_D)), draws 64 Latin-hypercube probes (the realistic CFD budget for one F1 design iteration under the FIA ATR), and fits a multilinear `LookupTableND` surrogate. Then drops the surrogate into a slim copy of the Part-1 LTS and takes `jax.grad(-arc_length)` w.r.t. the 5-D aero state at the nominal trim point. Headline live numbers: (C_L A) RMS-vs-truth ~ 0.23 m² (bias-dominated by the multilinear-on-quadratic-truth fit, *not* the 0.06 m² noise floor — honest framing in the prose); dominant-component gradient through the fit matches truth to within ~28% (the multilinear-fit bias dominates; the *sign* and *direction* of every component are preserved, which is what matters for setup-search). Pareto sweep over (N \\in {8, 16, 32, 64, 128, 256}) follows the *publication / fast-mode pattern*: default loads `media/f1_part_3_publication.npz`, whose fit-error curve is real (measured against a 5000-point validation set) but whose gradient-error curve is a disclosed structural estimate (flagged in-notebook as an estimate) pending a full-fidelity LTS gradient sweep — offline script at `media/f1_part_3_publication_offline.py`. The fit-error curve flattens to a bias floor past (N \\sim 64), while the gradient-error curve continues falling — the structural Pareto trade-off Part 4 will exploit. Surfaces 3 follow-up findings: (a) no `fit_lookup_table_nd` shipped (every author hand-rolls the multilinear LS design matrix), (b) `LapTimeAccumulator`'s integral-of-smoothed-indicator readout gives zero gradient when the lap doesn't finish within `T_END` (Part 2's live cell shows all-zeros for this reason; Part 3 swaps to a `-arc_length` proxy), (c) `analyze_phase_activity` lacks the `name=` kwarg that `analyze_saturation` / `analyze_control_oscillation` accept. - [Sobol-driven CFD budget allocation under the FIA ATR](https://py.jaxonomy.com/examples/f1_part_4_sobol_cfd_budget/index.md) — closes Series B. Inherits Part 3's fitted 5-D `LookupTableND` aero surrogate and Part 1's LTS substrate, then promotes the surrogate from a *gradient oracle* (Part 2) to a *spend-allocation oracle* via `sobol_indices`, `decompose_variance_sobol`, `vmap_qoi`, and `morris_screening`. A live Sobol pass ranks the 5 aero axes ((h_F) dominant at (S_T \\approx 0.77), then (h_R \\approx 0.32), (\\delta \\approx 0.18), (\\phi \\approx 0.11), with (\\beta \\approx 0.06) smallest), the `decompose_variance_sobol` call splits `Var(lap_time)` into aleatoric (~1%, CFD noise floor — irreducible) and epistemic (~99%, fit residual — reducible-with-more-CFD), and a per-cell variance-reduction-per-CFD-hour heatmap on the dominant ((h_F, \\beta)) slice points the next batch at the informative corners. **Strategy comparison** (the counter-intuitive headline): greedy Sobol-weighted allocation (Strategy B) *loses* to uniform-LHS (Strategy A) at the 50-probe budget — roughly ~12% vs ~34% variance reduction (~3× worse) — because chasing the single high-(S_T) axis starves the fit everywhere else. The Sobol indices stay the right diagnostic of *what matters* ((h_F) dominates); they are the wrong rule for *where to sample*. Tied back to the FIA ATR sliding scale (positions 1 through 10). Uses the *publication / fast-mode pattern* (default loads the real `media/f1_part_4_publication.npz`, generated by `media/f1_part_4_publication_offline.py`). Closes with Morris screening at (n\_\\text{traj} = 16) (96 evaluations — 10× cheaper than Sobol) that recovers the same headline ranking (the Series-C tractability lever for the SU2-coupled pipeline), a proxy-vs-Sobol sanity check on the §11 axis-sweep variance, five exercises (bootstrap-CI on Sobol, sixth-axis rake, multi-objective (\\alpha = w_V \\Delta V + w_G |\\nabla T|) acquisition), and the explicit *Series B is complete* hand-off to Series C. Surfaces 1 follow-up finding on `sobol_indices` ranking inconsistency under tiny-(N) MC noise (small enough at (N = 4096) to be benign, but worth a documented `n_samples` floor). - [NACA airfoil SU2 co-sim: `jax.custom_vjp` wrapping an external CFD solver](https://py.jaxonomy.com/examples/f1_part_5_naca_su2_cosim/index.md) — demonstrates the co-simulation *architecture* end-to-end: a `jax.custom_vjp + jax.pure_callback` wrapper lets `jax.grad(lap_time)` flow through an external CFD solver on a NACA airfoil. Runs live with an in-process panel-method stand-in solver; the SU2 v8.5.0 RANS + discrete-adjoint numbers need an external toolchain (`pysu2`, `SU2_AD`) not bundled here, so those specific figures are not yet included and the notebook flags this to the reader. - [Full lap-time-aware aero shape optimisation: DrivAerML + OpenVSP + SU2 + hero MP4](https://py.jaxonomy.com/examples/f1_part_6_drivaerml_hero/index.md) — the same architecture on a parametrised rear-wing assembly. The full pipeline (OpenVSP geometry → SU2_DEF/CFD/adjoint → ParaView/Blender hero render) needs external executables not available here; the headline shape-optimisation numbers are therefore not yet included, pending an offline run of `media/f1_part_6_publication_offline.sh` (the notebook flags this to the reader). # Scope: PINNs and PDE surrogates **TL;DR — classical physics-informed neural networks (PINNs) for PDEs are out of scope for Jaxonomy.** Jaxonomy is a simulation engine for systems governed by ODEs and DAEs evolving in *time*; it has no spatial discretization, no collocation-point sampling, and no PDE residual machinery, and we do not plan to add them. ## What "PINN" means here The term is used for two quite different things. Only one of them belongs in Jaxonomy: | | Classical PDE PINN | Physics-informed *dynamics* learning | | ------------------- | ---------------------------------------------------------------------------------- | ----------------------------------------------------------------------- | | Governing equations | PDEs over space(-time): Burgers, Navier–Stokes, heat equation | ODEs / DAEs over time: mechanics, circuits, thermal networks, chemistry | | Unknown | A neural field `u(x, t)` trained to satisfy the PDE residual at collocation points | Parameters and/or a neural correction term inside a *simulated* model | | Core machinery | Spatial sampling, residual losses, boundary/initial-condition penalties | A differentiable time-stepping simulator | | In Jaxonomy? | **No** | **Yes — this is a core capability** | ## Out of scope (use these instead) If you want to train `u(x, t)` against a PDE residual — surrogate models for fluid fields, heat maps over a plate, wave propagation — use a library built for it: - [DeepXDE](https://github.com/lululxvi/deepxde) — the reference PINN library (PDEs, IDEs, fractional PDEs; TensorFlow/PyTorch/JAX backends). - [NVIDIA PhysicsNeMo (formerly Modulus)](https://developer.nvidia.com/physicsnemo) — industrial-scale physics-ML, including PINNs and neural operators. - [Neuromancer](https://github.com/pnnl/neuromancer) — differentiable programming for constrained optimization and physics-informed system identification in PyTorch. A spatially discretized PDE (method of lines) *can* be simulated in Jaxonomy — a finite-volume battery-electrode model or a discretized heat rod is just a large ODE/DAE system — but Jaxonomy does not own the discretization, and we will not add collocation/residual training utilities for neural fields. ## In scope (what Jaxonomy does instead) Physics-informed learning where the physics enters through a **differentiable simulation in time**: - **Universal differential equations (UDE)** — a neural term inside an ODE right-hand side, trained end-to-end through `simulate` (see the [UDE + symbolic regression example](https://py.jaxonomy.com/examples/ude_and_sr_lotka_volterra/index.md)). - **Neural DAE** — a neural correction inside an *acausal, constrained* DAE (`NeuralDAEBlock`; see the [constrained pendulum drag-recovery example](https://py.jaxonomy.com/examples/neural_dae_pendulum_drag/index.md)). The index-reduction pipeline (Pantelides) runs unchanged with the neural term in place. - **Neural ODE blocks** — `MLP` (Equinox) and imported `PyTorch` / `TensorFlow` / `ONNX` networks as blocks inside a diagram. - **SINDy** — sparse symbolic regression of dynamics from data (`Sindy` block). - **Differentiable parameter estimation** — `fit_parameters`, lookup-table fitting, and the whole autodiff/optimization workflow. The dividing line: **if the "physics" constraint is enforced by simulating a dynamical system forward in time, it belongs here; if it is enforced by a residual loss over a spatial domain, it does not.** # Scope: reduced-order modeling and surrogates **TL;DR — reduced-order modeling (ROM) and data-driven surrogates *of dynamical systems* are in scope for Jaxonomy and live in `jaxonomy.library.rom`.** Given a full-order model (an ODE/DAE diagram or snapshot data from one), Jaxonomy can build a cheaper reduced model that is still a first-class, differentiable, simulatable block. What is **out of scope** is the same thing as for PINNs — surrogates of *spatial PDE fields* (neural operators, `u(x, t)` collocation); see [`pinn.md`](https://py.jaxonomy.com/scope/pinn/index.md). ## What "ROM" covers here A reduced-order model approximates a high-dimensional or expensive dynamical system with a low-order one that is fast to simulate. Jaxonomy supports the three families that matter for control and simulation engineers, plus statistical surrogates for the input→output map: | Family | Methods | When to reach for it | | ---------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | **Linear MOR** | balanced truncation (`balred`), minimal realization (`minreal`), modal truncation, singular-perturbation residualization, Krylov/IRKA *(planned)* | You have (or can linearize to) an LTI model and want a smaller LTI with a certified error bound | | **Projection ROM** | POD–Galerkin, Petrov–Galerkin/LSPG, DEIM hyper-reduction | You have the *equations* of a large nonlinear ODE/DAE (e.g. a method-of-lines PDE) and snapshots, and want an intrusive, physics-preserving reduction | | **Data-driven operator ROM** | DMD, DMDc, ERA, eDMD / Koopman | You have *data* (snapshots), maybe no equations, and want a linear predictor — including a lifted-linear (Koopman) model you can drop straight into linear MPC/LQR | | **Statistical surrogates** | Gaussian process / kriging, polynomial chaos (PCE), RBF response surface | You want a cheap, optionally uncertainty-aware surrogate of an input→output map (design maps, UQ, calibration) | ## Choosing a method - **Do you have the model equations, or only data?** Equations → linear MOR (if linear) or POD–Galerkin/DEIM (if nonlinear). Data only → DMD/DMDc, ERA, or eDMD/Koopman. - **Linear or nonlinear?** Linear and you want a guaranteed error → balanced truncation (a priori H∞ bound). Nonlinear with equations → POD–Galerkin, and add DEIM so the per-step cost stops scaling with the full state dimension. Nonlinear with only data → eDMD/Koopman with a lifting dictionary. - **Is the reduced model for a controller?** Koopman/DMDc produce a *linear* reduced model in (possibly lifted) coordinates — a good basis for linear MPC / LQR-style control (design in lifted coordinates, de-lift with `C`). A lifted model wants a *terminal-cost* MPC rather than a hard terminal-equality one; the `rom_dmdc_koopman_mpc` example shows the pattern. - **Is it an input→output map, not a trajectory?** Use a statistical surrogate (GP/PCE/RBF). PCE additionally yields analytic Sobol indices and moments, so it doubles as an accelerated-UQ path into `jaxonomy.uq`. ## In scope — what Jaxonomy does Every reducer returns a first-class Jaxonomy object: linear MOR returns a reduced `LinearizedSystem`/`LTISystem`; projection and operator ROMs return a differentiable, `jit`/`vmap`-able `LeafSystem` you can compose in a diagram and drive through `jaxonomy.simulate`; statistical surrogates are feedthrough blocks. ROM quality metrics (relative trajectory error, retained energy, projection error, held-out cross-validation) live alongside the reducers. ## Out of scope (use these instead) - **Spatial PDE field surrogates / neural operators** (`u(x, t)` over a domain, Fourier/DeepONet operators) — see [`pinn.md`](https://py.jaxonomy.com/scope/pinn/index.md). A spatially *discretized* PDE (method of lines) is a large ODE and *can* be reduced with POD–Galerkin/DEIM here; Jaxonomy just does not own the discretization or train neural fields. - **Mesh generation, CFD/FEA solvers.** Bring your own high-fidelity solver; Jaxonomy reduces the resulting dynamical system or its snapshots. The dividing line is the same as for PINNs: **if the reduced object is a dynamical system evolving in time (or a map you sample), it belongs here; if it is a neural field trained against a spatial PDE residual, it does not.** # Public benchmarks What Jaxonomy looks like compared to MuJoCo Playground and JaxSim on standard control / simulation problems, with an honest accounting of where the comparison is fair and where it isn't. All Jaxonomy CPU numbers come from `benchmarks/public.py` — median of 3 runs on a GitHub-hosted Linux x86_64 runner / JAX 0.9.2 / jaxonomy 3.0.0 / CPU. Reproduce with: ``` python benchmarks/public.py # write fresh baseline JSON python benchmarks/public.py --check # compare vs locked baseline ``` ## What we benchmark Five problems, each exercising a different part of the stack: 1. **Cartpole throughput** — single-env continuous-time ODE, t_end = 10 s, fixed input. Stresses cold-compile + sim-loop on a small dense-Jacobian system. 1. **Quadruped throughput at N=1000 parallel envs** — `simulate_batch` over a simplified 4-leg single-joint pendulum (8 states), t_end = 5 s. Stresses the vmapped kernel path. 1. **Articulated quadruped throughput** — single-env MJX continuous-time rollout of a hand-authored 12-DoF quadruped (free- joint trunk + 4 legs × 3 hinges, nq=19, nv=18) onto a floor with contact, t_end = 5 s, no controller. Measures multi-body physics throughput on an articulated body — *not* locomotion / RL env-step rate. Two quadruped benchmarks coexist deliberately: `quadruped_throughput` (4 independent damped pendulums, a `simulate_batch` throughput stand-in) and `articulated_quadruped_throughput` (real 12-DoF MJX body, multi- body dynamics with contact). 1. **System-ID convergence** — fit `(b, k)` of a spring-damper from noisy synthetic data using L-BFGS-B + finite-diff gradients on the simulator. Characterises the optimisation stack. 1. **Linearization vs analytic** — `jaxonomy.linearize()` on a nonlinear oscillator at the equilibrium x=0 vs the closed-form Jacobians. Measures wall time *and* numerical accuracy. ## Headline numbers CPU numbers are jaxonomy 3.0.0 on the Linux x86_64 baseline runner. The T4 column was collected separately on an NVIDIA T4 host running jaxonomy **2.2.0** (see "Hardware notes") — treat it as indicative, not a version-matched comparison, and note that these small `float64` problems are often *slower* on a T4 than on CPU. | problem | metric | Jaxonomy CPU | T4 (2.2.0) | A100 | H100 | | ------------------------------------------------------------- | ------------------------ | --------------------- | --------------------- | --------- | --------- | | cartpole_throughput (t_end=10s) | wall-s / sim-s | 0.049 | 0.103 | *pending* | *pending* | | quadruped_throughput (N=1000, t_end=5s) | env-s / wall-s | 3,967 | 101 | *pending* | *pending* | | articulated_quadruped_throughput (N=1, nq=19/nv=18, t_end=5s) | wall-s / sim-s | 0.796 | 0.968 | *pending* | *pending* | | sysid_convergence | iters / wall / param-MSE | 12 / 20.7 s / 5.2e-06 | 12 / 45.9 s / 5.2e-06 | *pending* | *pending* | | linearization_vs_analytic | warm s / err vs analytic | 0.068 / 0.0 | 0.166 / 0.0 | *pending* | *pending* | A100 / H100 columns are populated by re-running `benchmarks/public.py` on the corresponding NVIDIA runner (runner contract below). Empty cells mean we have not yet collected the number on that device — they are *not* zeros and not interpolated from CPU. Cartpole simulates 10 simulated-seconds in ~490 ms on a single CPU core (~20 simulated-s per wall-s). The simplified quadruped batch at N=1000 runs 5,000 env-seconds in ~1.26 s wall (~3,967 env-s/wall-s, ≈0.4M env-steps/s at a 100 Hz step). The articulated quadruped (12-DoF body + floor contact, MJX continuous-time, no controller) simulates 5 s of multi-body physics in ~4.0 s wall (env-s/wall-s ≈ 1.26). Linearization matches the analytic Jacobians to machine precision (the cubic vanishes at x=0). ## How to think about comparisons Jaxonomy, MuJoCo Playground, and JaxSim each optimise for different things. Direct head-to-head numbers are only meaningful on a subset of workloads, and misleading on the rest. | library | primary target | fair comparison axis | | --------------------- | ------------------------------------------------------------ | ------------------------------------------------------------------------------ | | **Jaxonomy** | accuracy-first ODE/DAE simulation, controls, autodiff, sysid | small/medium ODE throughput, linearisation correctness, batch parameter sweeps | | **MuJoCo Playground** | RL training throughput, contact-rich rigid bodies | massively-parallel RL env-steps/s, GPU/TPU only | | **JaxSim** | robotics multi-body dynamics with contact | articulated rigid bodies, jit-friendly featherstone/RBDL | **Fair**: single-env continuous ODE (Jaxonomy vs JaxSim on a smooth contact-free system); batched parameter sweeps with `vmap` on the same body; linearisation accuracy at an equilibrium; single-env articulated multi-body throughput on a fixed reference body (`articulated_quadruped_throughput` is fair vs JaxSim's articulated- body benchmarks on the same body, modulo solver choice). **Not fair**: contact-rich RL throughput on GPU (MuJoCo Playground's ~10⁵ env-steps/s for cartpole at 1000 envs is a T4-GPU MJX number — our CPU number isn't comparable); MuJoCo Playground RL-policy-step throughput, which includes a learned controller and contact-rich gait that our zero-control passive drop deliberately excludes; massively-parallel articulated bodies on GPU (our articulated quadruped is single-env; a batched MJX ensemble is a filed follow-up); MPC inner-loop solve at 1 kHz (out of scope for all three). ## Competitor numbers and sourcing External-reference numbers used here: - **MuJoCo Playground cartpole-1000-env on T4 GPU**: ~10⁵ env-steps/s (source: MuJoCo Playground README). - **JaxSim multi-body throughput**: comparable order of magnitude on GPU for articulated bodies; we have not run a like-for-like benchmark of our simplified quadruped against a JaxSim pendulum-array equivalent. If a number isn't here, it's because we don't have a defensible measurement, not because we cherry-picked. **No competitor numbers were estimated, interpolated, or derived from secondary sources.** ## Hardware notes The CPU baseline is produced on a GitHub-hosted Linux x86_64 runner (Azure-backed `ubuntu-latest`) / 16.8 GB RAM / JAX 0.9.2 / jaxonomy 3.0.0 / CPU only. CI re-runs weekly on `ubuntu-latest` and uploads the JSON; cross-runner CPU variance is ~10-20 %. The T4 column was collected on a separate NVIDIA T4 host (Linux x86_64, 31.5 GB RAM, JAX 0.9.2) running **jaxonomy 2.2.0** — it predates the current 3.0.0 CPU baseline, so it is indicative rather than version-matched. A100 / H100 numbers are populated opportunistically when a matching NVIDIA runner is available — see "GPU runner contract" below for the recipe. The exact driver / CUDA / cuDNN fingerprint for each device run is recorded under `hardware.` in `public_baseline.json`. ## GPU runner contract Recipe for filling in the T4 / A100 / H100 columns when you have access to an NVIDIA box. ### JAX platform names JAX accepts both `gpu` and `cuda` for NVIDIA backends. `cuda` is the canonical name as of JAX 0.4.x+; `gpu` is the documented alias and what `jax.default_backend()` prints on a CUDA host. Either string is valid input to `--device`. Use `jax.devices()` to confirm: ``` import jax print(jax.default_backend()) # "gpu" on a CUDA host print(jax.devices()) # [CudaDevice(id=0)] or [GpuDevice(id=0)] ``` The benchmark script forces `JAX_PLATFORMS=cuda,cpu` early at module import when `--device gpu` (or env var `JAXONOMY_BENCH_DEVICE=gpu`) is set, so JAX's lazy backend init picks the GPU before any `jax.numpy` arrays are allocated. ### Populating a device column ``` # A T4 host (Colab / GCP n1-standard-4 + nvidia-tesla-t4 / GHA gpu runner) JAXONOMY_BENCH_DEVICE=gpu python benchmarks/public.py \ --device gpu --update-baseline gpu_t4 # A100 host (e.g. AWS p4d, Lambda Labs, Modal A100 sandbox) JAXONOMY_BENCH_DEVICE=gpu python benchmarks/public.py \ --device gpu --update-baseline gpu_a100 # H100 host (AWS p5, Modal H100, GCP A3) JAXONOMY_BENCH_DEVICE=gpu python benchmarks/public.py \ --device gpu --update-baseline gpu_h100 ``` `--update-baseline gpu_t4` writes only the `gpu_t4` column under each `cases.` entry; the existing `cpu` column and other GPU columns are preserved. The hardware fingerprint for the run lands at `hardware.gpu_t4` so reviewers can verify driver / CUDA / cuDNN versions after the fact. ### Verifying without overwriting ``` JAXONOMY_BENCH_DEVICE=gpu python benchmarks/public.py \ --device gpu --check ``` If `gpu` (or the resolved key — `gpu_t4` / `gpu_a100` / `gpu_h100`) isn't in the baseline yet, every case prints `NEW` and the run exits 0. This is the backwards-compatibility guarantee: missing device columns do not fail `--check`. ### Caveats for GPU runs - **Warmup.** The first `simulate` call pays the XLA AOT compile + PTX-to-SASS lowering cost. The script already takes the *warm* number (second call) for `cartpole_throughput`. For ensemble cases (`quadruped_throughput`) the kernel JIT and the `simulate_batch` scan body each cost ~hundreds of ms cold; the recorded `compile_s` isolates that. - **Persistent JIT cache.** Setting `JAXONOMY_PERSISTENT_JIT_CACHE=1` (see `jit_cache.md`) makes re-runs on the same machine essentially free of compile cost — useful for back-to-back `--check` runs but irrelevant to the published numbers, which always use the `compile_s_median` / `wall_per_simsec_median` fields. - **Default dtype.** Jaxonomy enforces `float64` by default. Some GPUs (especially consumer ones) are massively faster in `float32`; do *not* override the precision policy when populating the public baseline — that would compare apples to oranges. It is also why the T4 numbers above are slower than CPU on these small problems: a T4 has little `float64` throughput. If you want a separate float32 column, add a new device key (e.g. `gpu_a100_f32`). - **Batch-size scaling.** `quadruped_throughput` is fixed at N=1000 and `articulated_quadruped_throughput` at N=1 deliberately, so numbers are comparable across devices. Don't bump N to "show off" a bigger GPU — that breaks the cross-device comparison. Larger-N is a filed follow-up. - **Persistent NVIDIA driver state.** If the runner is shared (HF Spaces, Colab), free GPU memory before invoking via `nvidia-smi --gpu-reset` or by restarting the kernel. Compile time jitter from leftover allocations is the most common false-regression signal. ### Runner specs Use these as the canonical hosts for each column. Numbers must come from a host that matches the spec — driver / CUDA / cuDNN versions should be recorded in `hardware.` automatically. | device_key | runner | accelerator | host CPU | RAM | notes | | ---------- | ------------------------------------------------------------- | ----------------- | -------- | ------- | -------------------------------------------------------------------------- | | `gpu_t4` | NVIDIA T4 host (Colab / GCP n1 + nvidia-tesla-t4) | NVIDIA T4 (16 GB) | 4 vCPU | 31.5 GB | measured on jaxonomy 2.2.0; parity re-run on 3.0.0 pending | | `gpu_a100` | AWS p4d.24xlarge, Modal a100-40gb sandbox (pending) | NVIDIA A100 40 GB | 8 vCPU | 64 GB+ | mid-tier reference for a fair head-to-head with JaxSim's published numbers | | `gpu_h100` | AWS p5.48xlarge, Modal h100 sandbox, GCP a3-highgpu (pending) | NVIDIA H100 80 GB | 8 vCPU | 128 GB+ | upper bound; JAX 0.9+ is required for full H100 SM_90 codegen | When a number lands, record the exact runner identifier (e.g. `Modal a100-40gb-`, \`GHA self-hosted [linux, gpu, t4] runner # 3\`) so future regressions can be triaged against the same hardware. ### Status **T4 numbers are measured** (on jaxonomy 2.2.0 — a version-matched 3.0.0 re-run is still pending), and are populated in the headline table above. The `gpu_a100` / `gpu_h100` columns in `public_baseline.json` are still explicit `null`s: those runs are pending A100 / H100 access on GitHub Actions GPU self-hosted runners or a Modal / Replicate notebook (recipe above). ## Caveats - **Simplified quadruped is a stand-in.** Four independent damped pendulums, not multi-body locomotion — the benchmark characterises `simulate_batch` throughput, not feature parity. The companion `articulated_quadruped_throughput` case covers the multi-body story. - **Articulated quadruped is single-env (N=1)**, no controller, just a passive-drop onto a floor under gravity. Honest physics- throughput measurement on an articulated body, *not* a locomotion / RL benchmark. A batched MJX ensemble is a filed follow-up. - **Sysid fits 2 of 3 parameters** (mass anchored); three-param identifiability under noise is a separate study. - **Cartpole input held at zero**, no controller. Apples-to-apples vs other libraries on the same setup; absolute number would shift with an LQR in the loop. - **A100 / H100 columns not yet measured.** The schema, CLI flag, and runner contract for GPU runs are in place (see the "GPU runner contract" section above). `gpu_a100` / `gpu_h100` are currently explicit `null`s in `public_baseline.json`; numbers land when a matching NVIDIA runner is provisioned. ## Reproduction checklist 1. `pip install -e ".[test]"` from a clean checkout. 1. `python benchmarks/public.py` writes `public_baseline.json` (includes hardware fingerprint). 1. CI runs `--check` weekly (Mon 06:00 UTC) plus on `workflow_dispatch`. *Not* on PR/push. ## Filed follow-ups - **GPU benchmark numbers (A100 / H100).** Schema + CLI flag + runner contract shipped (see "GPU runner contract" above); the `gpu_a100` / `gpu_h100` rows in `public_baseline.json` are `null` pending an A100 / H100 run. The `gpu_t4` row is populated (jaxonomy 2.2.0). - **Real articulated-quadruped vs JaxSim** on a fixed URDF. # Distributed ensemble execution `simulate_batch` runs an N-element parameter sweep on a single JAX device. `simulate_distributed` shards that same sweep across the host's local devices using `jax.pmap`. Cross-host fan-out (multi-machine clusters) is intentionally out of scope — see "External orchestration" below. ## When to reach for which entry point | your setup | recommended entry point | | ------------------------------------------- | ---------------------------------------------------------------- | | one CPU, one GPU, or one TPU | `simulate_batch` | | one host, multiple GPUs / TPUs | `simulate_distributed` | | many hosts (cluster scale) | external orchestration (Modal etc.) | | diagram contains `CustomPythonBlock` or FMU | `simulate_batch` (loop path); orchestrate externally for fan-out | `simulate_distributed` is a thin wrapper around the same single-JIT kernel that `simulate_batch` already builds. Numerics are bit-equivalent up to XLA reduction-order rounding (`rtol ~1e-10` in the cases we've measured). ## API ``` import jax import jaxonomy res = jaxonomy.simulate_distributed( diagram, t_span=(0.0, 5.0), param_batches={"leaf.k": ks}, # leading axis N must divide len(devices) options=opts, recorded_signals={"x": diagram["leaf"].output_ports[0]}, devices=None, # default: jax.devices() ) ``` Constraints: - The diagram must be pure-JAX (no `CustomPythonBlock` or FMU blocks). `simulate_distributed` raises with a clear error otherwise — wrap `simulate_batch` in an external orchestrator (below) for those. - `N` (leading axis of every entry in `param_batches`) must be divisible by `len(devices)`. The 1-device degenerate case defers to `simulate_batch`'s kernel path so numerics match exactly. ## Single-host multi-device recipe For multi-GPU / multi-TPU, just call `simulate_distributed`. To exercise the path on a CPU-only dev machine, fake a multi-device mesh with an `XLA_FLAGS` environment variable *before* importing JAX: ``` import os os.environ.setdefault( "XLA_FLAGS", "--xla_force_host_platform_device_count=4" ) import jax import jaxonomy # jax.devices() now returns 4 CpuDevice's ``` A runnable end-to-end example lives at `docs/examples/distributed_ensemble.py`. It runs `simulate_distributed` against a 4-device fake mesh and checks the output against a serial `simulate_batch` call. ## External orchestration (cluster scale) For workloads that exceed a single host (long-running ensembles, very large `N`, or a mix of pure-JAX and `CustomPythonBlock` diagrams), wrap `simulate_batch` in an external orchestrator. Jaxonomy deliberately does **not** ship its own job queue / worker fleet. The recommended options: ### Modal [Modal](https://modal.com) gives a Python-native `@app.function` decorator that runs the wrapped function in a managed container with optional GPU. Distribute via `.map()` (or `.starmap()` for multi-arg shards): ``` import modal import jax.numpy as jnp import jaxonomy app = modal.App("jaxonomy-sweep") image = modal.Image.debian_slim().pip_install("jaxonomy") @app.function(image=image, gpu="T4") def run_one_shard(k_values): diagram = build_my_diagram() return jaxonomy.simulate_batch( diagram, t_span=(0.0, 5.0), param_batches={"plant.k": jnp.asarray(k_values)}, options=opts, recorded_signals=rec, ) @app.local_entrypoint() def main(): shards = [list(jnp.linspace(0.1 + 0.5 * i, 0.5 + 0.5 * i, 100)) for i in range(20)] results = list(run_one_shard.map(shards)) # combine N=2000 result across 20 containers ``` ### Replicate, SkyPilot, Ray Serve The shape is identical: define one Python function that takes a parameter shard and returns a `BatchSimulationResults`, then use the platform's native fan-out (`replicate.run`, `sky launch`, Ray's `ray.remote` / `Serve.deployment`) to dispatch shards across worker replicas. ### Why not in-library cluster orchestration? Jaxonomy deliberately omits in-library worker-fleet / job-queue infrastructure. It would be strictly worse than purpose-built tools (Modal, Replicate, SkyPilot, Ray) on every axis (autoscaling, retries, cost reporting, multi-cloud). The in-library wrapper would only add a thin dispatching layer that those tools already provide. ## When to consider extending distributed simulation `simulate_distributed` covers the multi-device-on-one-host case. Extend the API only if a real workload demonstrates need; candidate follow-ups: - **Asynchronous fan-out across hosts** without an external orchestrator (would essentially rebuild Modal / Ray — not recommended). - **`shard_map` over the time axis** to parallelise long-horizon single simulations across devices. Useful only when one simulation exceeds a single device's memory; not a current bottleneck. - **Hybrid pmap + vmap** with a non-trivial mesh (e.g. shard over GPUs and vmap within each GPU). The current recipe already does this — pmap outer, vmap inner; if a more complex mesh is needed, switch to `shard_map` with an explicit `Mesh`. File a follow-up if a concrete workload surfaces. # Persistent JIT compilation cache Jaxonomy compiles each diagram + solver combination into XLA on first use. Small single-shot ODE simulations compile in well under a second (roughly 60–250 ms, depending on whether the model has zero-crossings or an acausal DAE); large diagrams and `simulate_batch` ensembles take longer. JAX provides a persistent on-disk cache for the XLA-compile share of that cost — Jaxonomy ships a one-call helper. ``` import jaxonomy jaxonomy.enable_persistent_jit_cache() # ~/.cache/jaxonomy/jit/ # or jaxonomy.enable_persistent_jit_cache("/scratch/jit") # custom dir ``` The first run after enabling pays the normal compile cost and writes the artefact to disk. Subsequent processes (with matching JAX version, JAXPR, and target device) read the compiled executable from disk instead of re-running XLA compilation. ## What it does and does not buy you Two structural limits bound the win, both measured (2026-07, jax 0.9.2, arm64 CPU): - **Small models are below the write threshold.** The helper sets `jax_persistent_cache_min_compile_time_secs = 1.0`, so kernels that compile faster than 1 s are never written. A 4-block PID loop and the bouncing-ball model with `record_event_times=True` (first `simulate` ≈ 0.21–0.25 s) produced **zero cache entries** — warm-cache startup was identical to cold. For these models the cache is a harmless no-op; the compile is cheap anyway. - **Python tracing is not cacheable.** The cache stores compiled XLA executables, not traces. On a 160-block diagram (first `simulate` ≈ 2.4 s cold), a warm cache cut the first call to ≈ 1.2 s — a genuine ~2× — but the remaining ~1.1 s is JAX tracing/lowering, which every process pays regardless. A side benefit on larger models: repeated *bare* `simulate` calls in the same process re-trace each time (see below), and with the cache enabled each re-trace's recompile hits the disk cache too (measured ≈ 2.2 s → ≈ 1.1 s per repeat call on the 160-block diagram). **The cache does not fix Python-loop parameter sweeps.** In `for v in grid: simulate(..., ctx.with_parameter("p", float(v)))` each float is baked into the HLO as a constant, so every value is a compulsory cache miss (measured: a 3-value sweep against a warm cache added 3 fresh entries and every iteration paid full re-trace + compile). Use `simulate_batch` or traced parameters instead — see the sweep entry in `KNOWN_GAPS.md`. ## What is cached The cache keys on the JAXPR, JAX version, and target device. Bumping JAX, switching CPU↔GPU, or changing a static parameter that flips a control-flow branch produces a fresh entry. Stale entries persist; periodically delete the cache directory on big version bumps to reclaim disk space. ## Tunables `enable_persistent_jit_cache` configures three JAX options: | option | value | rationale | | -------------------------------------------- | ------------------- | -------------------------------------------------------- | | `jax_compilation_cache_dir` | the cache directory | required | | `jax_persistent_cache_min_compile_time_secs` | `1.0` | trivial computations recompile faster than disk read | | `jax_persistent_cache_min_entry_size_bytes` | `-1` | size threshold disabled — gating on time alone is enough | For different thresholds, call `jax.config.update(...)` directly afterwards. ## When to enable Enable it when your model's compile time is noticeable — large diagrams (≳100 blocks), big acausal packs, `simulate_batch` ensembles, or jitted gradient loops (next section), where it roughly halves cold-start and recompile cost. For small models it is a harmless no-op (compiles under 1 s are never written). See `benchmarks/compile_time.py` for per-case compile timings. ## Repeated gradients: hoist the `jit` Every bare call to `jaxonomy.simulate` builds a *fresh* traced closure, so JAX's in-process jit cache misses on function identity and you pay a full re-trace + XLA compile **per call** — the numeric solve itself is milliseconds. This dominates design loops that differentiate through the simulator repeatedly: ``` # SLOW — each value_and_grad call re-traces + recompiles the whole # forward + adjoint (seconds per call; ~30 s on a 24-state acausal pack): def objective(theta): ctx = base_context.with_parameter("g0", theta) res = jaxonomy.simulate(model, ctx, (0.0, tf), options=opts) return res.context.continuous_state[1] for step in range(5): J, dJ = jax.value_and_grad(objective)(theta) # re-traces every time ... ``` The fix is to define the objective **once** as a pure function of `(theta, context)` and wrap the *outer* `value_and_grad` in `jax.jit`, so tracing happens exactly once: ``` # FAST — one compile, then ~milliseconds per call: @jax.jit def value_and_grad_fn(theta, context): def objective(theta): ctx = context.with_parameter("g0", theta) res = jaxonomy.simulate(model, ctx, (0.0, tf), options=opts) return res.context.continuous_state[1] return jax.value_and_grad(objective)(theta) for step in range(5): J, dJ = value_and_grad_fn(theta, base_context) # cached after call 1 ... ``` This works for the implicit **BDF/DAE** path too (with `SimulatorOptions(enable_autodiff=True)`): the simulation context and BDF solver state are ordinary pytrees and trace cleanly. Measured on the index-2 pendulum DAE (9 states, BDF, 2 s horizon, CPU): unjitted ~1.8 s **per call**; jitted 1.8 s once, then **~10 ms per call** (~180×). Cost envelope: per-call cost of the naive pattern is almost entirely trace+compile and scales with model size (states, blocks, solver machinery), not with the horizon; the compiled kernel's runtime scales with horizon and stiffness. Combine with the persistent cache above to also amortise the one-time compile across processes. Requirements for the pattern: pass the context (and any other non-differentiated inputs) as *arguments* of the jitted function rather than closing over mutable state, keep `t_span` and options static, and don't rebuild the diagram inside the traced function (acausal compilation is not jit-safe — build once, outside). # Memory footprint and large parameter sweeps How much resident memory (RSS) a Jaxonomy simulation will use, so you can size a parallel sweep without OOMing your machine. All numbers come from `benchmarks/memory.py`, median of 3 subprocess runs on a GitHub-hosted Linux x86_64 runner / JAX 0.9.2 / jaxonomy 3.0.0 / 16.8 GB RAM. ## How much memory does an N-element parameter sweep cost? For the 5-second exponential-decay benchmark (single integrator + gain, the canonical `simulate_batch` test case): | N elements | peak RSS | | ---------- | -------- | | 1 | ~356 MB | | 10 | ~358 MB | | 100 | ~358 MB | Linear fit gives **~0.01 MB per added batch element** on top of the ~356 MB process baseline (Python interpreter + JAX/XLA runtime + the compiled simulator kernel) — essentially flat for this small case, because the fixed-size recorded-signal buffer dominates the per-element state. A 1000-element sweep is roughly **356 + 1000 × 0.01 ≈ 366 MB**. For larger systems the slope grows with the per-element recorded-signal buffer size. Multiply the per-element overhead by `(buffer_length / 200) × (n_signals × dtype_bytes)` for a rough scaling. ## Long-horizon growth A single harmonic-oscillator simulation at increasing `t_end`: | t_end | peak RSS | | ------ | -------- | | 10 s | ~356 MB | | 100 s | ~358 MB | | 1000 s | ~356 MB | **Essentially flat** (measured slope is ≈0 MB per simulated second). Dopri5 records into a fixed-size buffer (`SimulatorOptions.max_major_steps`, default 200), so memory does not grow with simulated time, only with the buffer length you ask for. Plan for ~5 MB per million steps per scalar recorded signal. ## Compile-time peak memory Cold-compiling each benchmark case (subprocess isolated): | case | peak RSS | | --------------------------- | --------------------------------- | | scalar_exponential_decay | ~298 MB | | state_machine_three_state | ~299 MB | | harmonic_oscillator | ~300 MB | | bouncing_ball_zc | ~311 MB | | simulate_batch_decay (N=10) | ~358 MB | | rc_acausal_dae | ~422 MB | | pid_first_order_plant | *skipped (control not installed)* | None cross 500 MB. **Allow ~500 MB headroom for compile spikes on commodity 8 GB machines.** The acausal DAE case (`rc_acausal_dae`) is the current peak — same root cause as its top-tier compile-time slowdown in `jit_cache.md`. The PID case was skipped in the baseline run because the optional `control` dependency wasn't installed; install `.[safe]` or `python-control` to measure it. ## `simulate_batch` vs serial loop For the N=100 sweep above, the per-element overhead (~0.01 MB/element, ≈2 MB at N=100) is negligible next to the JAX/XLA process baseline (~356 MB). vmapped batch and a serial Python loop pay the same compile baseline; the trade-off is **wall-clock, not memory**. vmapped batch is 3–10× faster on CPU at N=100 because it amortises Python dispatch and lets XLA fuse across elements. Use `simulate_batch` whenever your sweep fits in memory; the only reason to fall back to a serial loop is N × per-element-MB exceeding RAM. ## Reproducing ``` python benchmarks/memory.py # write a fresh baseline JSON python benchmarks/memory.py --check # compare a future run vs baseline ``` `--check` tolerates +50 % variance — peak-RSS measurements on shared CI runners are noisy. See `benchmarks/README.md` for the threshold rationale. ## Honest limitations `ru_maxrss` reports peak RSS since process start, not a delta. To get clean per-workload numbers, the benchmark forks a subprocess per measurement — each datapoint pays the ~300 MB JAX-startup tax. For *additional* memory (the slope of N → MB), subtract the smallest value in the column. We don't attribute memory to specific JAX device buffers vs the C++ heap vs Python objects; `tracemalloc` only sees the Python heap (\<5 % of the real total) and the XLA debug allocators are not portably accessible. For finer-grained attribution, run under `mprof` or `psrecord`. # About **Jaxonomy** is a Python simulation library focused on **composable block-diagram models**, **hybrid dynamics** (continuous + discrete + events), and a **JAX-first** execution path for performance and differentiation. This documentation site is generated with **[MkDocs](https://www.mkdocs.org/)**, **[Material for MkDocs](https://squidfunk.github.io/mkdocs-material/)**, **[mkdocstrings](https://mkdocstrings.github.io/)** (API reference), and **[mkdocs-jupyter](https://github.com/danielfrg/mkdocs-jupyter)** (notebook pages). To build it locally from the repository: ``` pip install -r requirements.docs.txt mkdocs build --clean # or: mkdocs serve ``` General project information is at [www.jaxonomy.com](https://www.jaxonomy.com), and the source code at [github.com/machinavitalis/jaxonomy](https://github.com/machinavitalis/jaxonomy).