Simulation

jaxonomy.simulation

BatchSimulationResults dataclass

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	Optional :class:ProvenanceManifest capturing library versions, options and the system fingerprint at batch-start (T-110-followup-attach-on-batch). None unless SimulatorOptions.record_provenance=True. One manifest is shared across all batch replicas — they all share the same system + options, only the parameter values differ (and those are already in param_batches).

Source code in jaxonomy/simulation/batch.py

@dataclasses.dataclass
class BatchSimulationResults:
    """Results from :func:`simulate_batch`.

    Attributes:
        time: 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: Mapping ``signal_name -> array`` with shape ``(N, T, ...)`` where
            ``N`` is batch size.
        used_vmap: ``True`` if the vectorised vmap path was used.
        provenance: Optional :class:`ProvenanceManifest` capturing library
            versions, options and the system fingerprint at batch-start
            (T-110-followup-attach-on-batch).  ``None`` unless
            ``SimulatorOptions.record_provenance=True``.  One manifest is
            shared across all batch replicas — they all share the same
            system + options, only the parameter values differ (and those
            are already in ``param_batches``).
    """

    time: Any
    outputs: dict[str, Any]
    used_vmap: bool = False
    provenance: ProvenanceManifest | None = None

    def mean(self, signal: str) -> Any:
        """Mean trajectory across the batch (axis 0)."""
        return jnp.mean(self.outputs[signal], axis=0)

    def std(self, signal: str) -> Any:
        """Standard deviation across the batch (axis 0)."""
        return jnp.std(self.outputs[signal], axis=0)

    def percentile(self, signal: str, p: float) -> Any:
        """``p``-th percentile across the batch at each time index; ``p`` in ``[0, 100]``."""
        return jnp.percentile(self.outputs[signal], p, axis=0)

    def to_simulation_results(self, idx: int) -> SimulationResults:
        """Slice one batch index into a :class:`SimulationResults` (no final context)."""
        return SimulationResults(
            None,
            time=self.time,
            outputs={k: v[idx] for k, v in self.outputs.items()},
            parameters=None,
        )

mean(signal)

Mean trajectory across the batch (axis 0).

Source code in jaxonomy/simulation/batch.py

def mean(self, signal: str) -> Any:
    """Mean trajectory across the batch (axis 0)."""
    return jnp.mean(self.outputs[signal], axis=0)

percentile(signal, p)

p-th percentile across the batch at each time index; p in [0, 100].

Source code in jaxonomy/simulation/batch.py

def percentile(self, signal: str, p: float) -> Any:
    """``p``-th percentile across the batch at each time index; ``p`` in ``[0, 100]``."""
    return jnp.percentile(self.outputs[signal], p, axis=0)

std(signal)

Standard deviation across the batch (axis 0).

Source code in jaxonomy/simulation/batch.py

def std(self, signal: str) -> Any:
    """Standard deviation across the batch (axis 0)."""
    return jnp.std(self.outputs[signal], axis=0)

to_simulation_results(idx)

Slice one batch index into a :class:SimulationResults (no final context).

Source code in jaxonomy/simulation/batch.py

def to_simulation_results(self, idx: int) -> SimulationResults:
    """Slice one batch index into a :class:`SimulationResults` (no final context)."""
    return SimulationResults(
        None,
        time=self.time,
        outputs={k: v[idx] for k, v in self.outputs.items()},
        parameters=None,
    )

Decay

Bases: LeafSystem

xdot = -k * x; output = x.

Source code in jaxonomy/simulation/testing_systems.py

class Decay(LeafSystem):
    """xdot = -k * x; output = x."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("k", 1.0)
        self.declare_continuous_state(
            default_value=jnp.array(1.0), ode=self._ode,
        )
        self.declare_continuous_state_output(name="x")

    def _ode(self, t, state, **p):
        return -p["k"] * state.continuous_state

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.	None
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.	None

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).

Source code in jaxonomy/simulation/fast_restart.py

class 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.

    Args:
        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.
        t_span: ``(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.
        options: :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: 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).
    """

    def __init__(
        self,
        system,
        t_span: tuple[float, float],
        options: SimulatorOptions | None = None,
        recorded_signals: dict[str, OutputPort] | None = None,
    ):
        if recorded_signals is None or not recorded_signals:
            raise ValueError(
                "FastRestartSimulator requires a non-empty recorded_signals "
                "dict (same convention as simulate()).  The set of recorded "
                "signals must be fixed at construction so the kernel buffer "
                "shape can be locked in before JIT compilation."
            )
        if options is None:
            options = SimulatorOptions()
        if options.math_backend != "jax":
            raise ValueError(
                "FastRestartSimulator requires math_backend='jax' to benefit "
                "from JIT cache reuse; got "
                f"math_backend={options.math_backend!r}."
            )
        if not options.enable_tracing:
            raise ValueError(
                "FastRestartSimulator requires enable_tracing=True (the "
                "default) to JIT-compile the kernel; got "
                "enable_tracing=False."
            )

        self.system = system
        self.t_span = (float(t_span[0]), float(t_span[1]))
        self._user_options = options
        self.recorded_signals = dict(recorded_signals)

        # Lazy init — set on first ``run`` so we pay the build cost
        # exactly once, after the user has finalised their setup.
        self._sim = None
        self._base_ctx = None
        self._opts_resolved: SimulatorOptions | None = None
        self._kernel = None

        # Cached abstract-pytree signature of the context the kernel
        # was last driven with — used to emit a one-time warning when
        # a :meth:`run` invocation will force a JIT recompile.  The
        # JIT cache key is the context's treedef + leaf shapes/dtypes;
        # mismatched signature ⇒ recompile.  ``None`` until the
        # kernel has been driven at least once.
        self._cached_ctx_sig: tuple | None = None

        # Per-diagram-identity kernel cache for
        # :meth:`run_with_diagram` (T-112-followup-multi-system).
        # Keyed on ``id(diagram)``; each entry holds the resolved
        # simulator bundle for that diagram.  We also hold strong
        # references to the cached Diagram objects in
        # ``_diagram_pool`` so their ``id``s remain valid for the
        # lifetime of the FastRestartSimulator (Python may otherwise
        # reuse the id of a garbage-collected object).
        self._diagram_kernel_cache: dict[int, dict[str, Any]] = {}
        self._diagram_pool: dict[int, Diagram] = {}

        # Track call counts for diagnostics / tests.
        self.n_runs: int = 0

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def __enter__(self) -> "FastRestartSimulator":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def close(self) -> None:
        """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.
        """
        self._kernel = None
        self._base_ctx = None
        self._sim = None
        self._opts_resolved = None
        self._cached_ctx_sig = None
        self._diagram_kernel_cache.clear()
        self._diagram_pool.clear()

    def reset(self, diagram: Diagram | None = None) -> 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.

        Args:
            diagram: 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".
        """
        if diagram is not None:
            self.system = diagram
        self.close()

    # ------------------------------------------------------------------
    # First-run build
    # ------------------------------------------------------------------

    def _build(self) -> None:
        """Lazily build the Simulator, ODE solver, and JIT kernel."""
        # Local imports to avoid circular module dependencies at import time.
        import jax

        from ..backend import ODESolver, set_backend
        from .simulator import Simulator

        set_backend("jax")

        opts = _check_options(
            self.system,
            self._user_options,
            self.t_span,
            self.recorded_signals,
        )
        self._opts_resolved = opts

        ode_solver = ODESolver(self.system, options=opts.ode_options)
        sim = Simulator(self.system, ode_solver=ode_solver, options=opts)
        self._sim = sim

        self._base_ctx = self.system.create_context()

        t0, tf = self.t_span

        @jax.jit
        def _kernel(context):
            return sim.advance_to(tf, context.with_time(t0))

        self._kernel = _kernel

    # ------------------------------------------------------------------
    # Per-run execution
    # ------------------------------------------------------------------

    def run(
        self,
        parameters: dict[str, Any] | None = None,
        initial_state: Any | None = None,
    ) -> SimulationResults:
        """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.

        Args:
            parameters: 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: 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:
            A :class:`SimulationResults` populated with ``time``,
            ``outputs``, and (when
            ``options.return_context=True``) ``context``.
        """
        if self._kernel is None:
            self._build()

        # Patch the base context with the per-run parameter overrides.
        ctx = self._base_ctx
        if parameters:
            updates = {k: jnp.asarray(v) for k, v in parameters.items()}
            ctx = _pure_patch_context(ctx, updates)

        # Apply the initial-state override (if any) on top of the
        # parameter patch.  We do this *after* the parameter patch so
        # that user-supplied initial state always wins regardless of
        # whether the diagram exposed the IC as a parameter as well.
        if initial_state is not None:
            ctx = self._apply_initial_state(ctx, initial_state)

        # Emit a one-time structural-change warning if this run will
        # force a JIT recompile.  Compute the signature *after* all
        # context patches are applied so we compare against what the
        # kernel actually sees.
        self._maybe_warn_structural_change(ctx)

        # Drive the JIT-compiled kernel.
        sim_state = self._kernel(ctx)
        results_data = sim_state.results_data
        time, outputs = results_data.finalize()

        final_context = (
            sim_state.context if self._opts_resolved.return_context else None
        )

        self.n_runs += 1

        return SimulationResults(
            final_context,
            time=time,
            outputs=outputs,
            parameters=dict(parameters) if parameters else None,
        )

    # ------------------------------------------------------------------
    # Batched parameter sweep over the warm-cached kernel
    # ------------------------------------------------------------------

    def run_batch(
        self,
        parameters_batch: dict[str, Any],
        initial_states_batch: Any | None = None,
    ) -> "BatchSimulationResults":
        """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.

        Args:
            parameters_batch: ``{path: (N, ...) array}`` — same convention
                as :func:`simulate_batch`'s ``param_batches``.  Every value
                must have the same leading batch size ``N``.
            initial_states_batch: 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:
            A :class:`BatchSimulationResults` with
            ``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.
        """
        from .batch import BatchSimulationResults, _infer_batch_size

        if not parameters_batch:
            raise ValueError(
                "FastRestartSimulator.run_batch: parameters_batch must be "
                "non-empty.  For a single warm-restart simulation pass "
                "parameters=... to .run() instead."
            )

        # Ensure the kernel + base context are built (lazy on first call).
        if self._kernel is None:
            self._build()

        n = _infer_batch_size(parameters_batch)
        stacked = {path: jnp.asarray(arr) for path, arr in parameters_batch.items()}

        # Build a batched context: broadcast every leaf of the base ctx to
        # shape (N, ...), then patch in the per-batch parameter values
        # (which already have the leading N axis).
        batched_ctx = jax.tree_util.tree_map(
            lambda x: jnp.broadcast_to(
                jnp.asarray(x)[None], (n,) + jnp.asarray(x).shape
            ),
            self._base_ctx,
        )
        batched_ctx = _pure_patch_context(batched_ctx, stacked)

        # Optional batched initial-state override.  Apply *after* the
        # parameter patch so the user's IC always wins.
        if initial_states_batch is not None:
            batched_ctx = self._apply_batched_initial_state(
                batched_ctx, initial_states_batch, n,
            )

        # vmap the cached kernel.  ``axis_name="batch"`` mirrors
        # :func:`simulate_batch`'s vmap path so any
        # ``fold_in_batch_index=True`` blocks (T-122-followup) work
        # identically here.
        batched_kernel = jax.vmap(self._kernel, axis_name="batch")
        batch_sim_states = batched_kernel(batched_ctx)
        results_data_batch = batch_sim_states.results_data

        # Finalize each batch row (results_data carries variable-length
        # buffers; finalize must run per-element).  Mirror the
        # numpy-host interpolation pattern from
        # :func:`simulate_batch._vmap_path` so output shapes line up
        # across rows with different step counts.
        import warnings as _warnings

        import numpy as _np

        from .batch import _interp_on_time_np

        time_ref = None
        out_lists: dict[str, list] = {k: [] for k in self.recorded_signals}

        for i in range(n):
            rd_i = jax.tree_util.tree_map(lambda x: x[i], results_data_batch)
            time_i, outputs_i = rd_i.finalize()
            if time_ref is None:
                time_ref = time_i
                for sig_name in self.recorded_signals:
                    out_lists[sig_name].append(_np.asarray(outputs_i[sig_name]))
                continue

            t_ref_end = float(time_ref[-1])
            t_run_end = float(time_i[-1])
            if abs(t_run_end - t_ref_end) / max(abs(t_ref_end), 1e-10) > 0.01:
                _warnings.warn(
                    f"FastRestartSimulator.run_batch: row {i} ended at "
                    f"t={t_run_end:.4g} but reference row ended at "
                    f"t={t_ref_end:.4g}.  Outputs will be interpolated "
                    "(clamped) to fill the time grid.",
                    UserWarning,
                    stacklevel=3,
                )
            same_grid = (
                time_i.shape == time_ref.shape
                and _np.array_equal(time_i, time_ref)
            )
            for sig_name in self.recorded_signals:
                if same_grid:
                    out_lists[sig_name].append(_np.asarray(outputs_i[sig_name]))
                else:
                    out_lists[sig_name].append(
                        _interp_on_time_np(outputs_i[sig_name], time_i, time_ref)
                    )

        stacked_out = {k: _np.stack(vs, axis=0) for k, vs in out_lists.items()}
        self.n_runs += n
        return BatchSimulationResults(
            time=time_ref, outputs=stacked_out, used_vmap=True,
        )

    def _apply_batched_initial_state(self, batched_ctx, initial_states_batch, n: int):
        """Apply a batched IC override onto a context already broadcast to ``(N, ...)``.

        Mirrors :meth:`_apply_initial_state` but expects per-element
        arrays to carry the leading batch axis ``N``.  The
        ``with_continuous_state`` API accepts arrays whose shape may be
        ``(N, ...)`` as long as the tree structure matches the context's
        continuous-state slot, so we delegate to it.
        """
        from ..framework.context import DiagramContext, LeafContext

        if isinstance(batched_ctx, LeafContext):
            xc_new = jnp.asarray(initial_states_batch)
            if xc_new.shape[0] != n:
                raise ValueError(
                    "FastRestartSimulator.run_batch(initial_states_batch=...): "
                    f"expected leading batch size {n}, got shape {xc_new.shape}."
                )
            return batched_ctx.with_continuous_state(xc_new)

        if isinstance(batched_ctx, DiagramContext):
            n_xc = len(batched_ctx.continuous_subcontexts)
            if n_xc == 0:
                raise ValueError(
                    "FastRestartSimulator.run_batch(initial_states_batch=...): "
                    "diagram has no continuous-state blocks; nothing to override."
                )
            if isinstance(initial_states_batch, (list, tuple)):
                xs = [jnp.asarray(x) for x in initial_states_batch]
            else:
                if n_xc != 1:
                    raise ValueError(
                        "FastRestartSimulator.run_batch(initial_states_batch=...): "
                        f"diagram has {n_xc} continuous-state blocks but a "
                        "single array was passed; pass a list/tuple of "
                        "arrays matching ctx.continuous_state ordering."
                    )
                xs = [jnp.asarray(initial_states_batch)]
            if len(xs) != n_xc:
                raise ValueError(
                    "FastRestartSimulator.run_batch(initial_states_batch=...): "
                    f"expected {n_xc} arrays (one per continuous-state "
                    f"block); got {len(xs)}."
                )
            for k, x in enumerate(xs):
                if x.shape[0] != n:
                    raise ValueError(
                        "FastRestartSimulator.run_batch(initial_states_batch=...): "
                        f"array {k} has leading dim {x.shape[0]}, expected "
                        f"batch size {n}."
                    )
            return batched_ctx.with_continuous_state(xs)

        return batched_ctx.with_continuous_state(initial_states_batch)

    # ------------------------------------------------------------------
    # Multi-system (per-diagram-identity) kernel cache
    # ------------------------------------------------------------------

    def _build_for_diagram(
        self,
        diagram: Diagram,
        recorded_signals: dict[str, OutputPort],
    ) -> dict[str, Any]:
        """Build a Simulator + JIT kernel bundle for ``diagram``.

        Returns a dict with keys ``sim``, ``base_ctx``, ``opts_resolved``,
        ``kernel``, ``recorded_signals``, ``cached_ctx_sig`` (the latter
        is initially ``None`` and is populated on first kernel drive,
        mirroring the structural-change-warning behaviour of the
        ``self`` kernel).
        """
        import jax

        from ..backend import ODESolver, set_backend
        from .simulator import Simulator

        set_backend("jax")

        opts = _check_options(
            diagram,
            self._user_options,
            self.t_span,
            recorded_signals,
        )

        ode_solver = ODESolver(diagram, options=opts.ode_options)
        sim = Simulator(diagram, ode_solver=ode_solver, options=opts)
        base_ctx = diagram.create_context()

        t0, tf = self.t_span

        @jax.jit
        def _kernel(context):
            return sim.advance_to(tf, context.with_time(t0))

        return {
            "sim": sim,
            "base_ctx": base_ctx,
            "opts_resolved": opts,
            "kernel": _kernel,
            "recorded_signals": dict(recorded_signals),
            "cached_ctx_sig": None,
        }

    def run_with_diagram(
        self,
        diagram: Diagram,
        parameters: dict[str, Any] | None = None,
        initial_state: Any | None = None,
        recorded_signals: dict[str, OutputPort] | None = None,
    ) -> SimulationResults:
        """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.

        Args:
            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.
            parameters: Optional dot-path mapping ``{"block.param":
                value, ...}`` — same convention as :meth:`run`.
            initial_state: Optional override for the simulator's
                continuous state at ``t = t_span[0]`` — same convention
                as :meth:`run`.
            recorded_signals: 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:
            A :class:`SimulationResults` populated with ``time``,
            ``outputs``, and (when
            ``options.return_context=True``) ``context``.
        """
        cache_key = id(diagram)
        bundle = self._diagram_kernel_cache.get(cache_key)
        if bundle is None:
            # First call for this diagram — build the kernel.
            sigs = recorded_signals if recorded_signals is not None else self.recorded_signals
            if not sigs:
                raise ValueError(
                    "FastRestartSimulator.run_with_diagram(...): no "
                    "recorded_signals available for this diagram.  Pass "
                    "recorded_signals={'name': port, ...} on the first "
                    "call for each diagram in your pool — the ports are "
                    "diagram-specific and must match the diagram passed "
                    "in this call."
                )
            bundle = self._build_for_diagram(diagram, sigs)
            self._diagram_kernel_cache[cache_key] = bundle
            # Hold a strong reference so ``id(diagram)`` stays valid
            # (Python may otherwise reuse the id of a GC'd object).
            self._diagram_pool[cache_key] = diagram

        # Patch the cached base context with per-run overrides.
        ctx = bundle["base_ctx"]
        if parameters:
            updates = {k: jnp.asarray(v) for k, v in parameters.items()}
            ctx = _pure_patch_context(ctx, updates)
        if initial_state is not None:
            ctx = self._apply_initial_state(ctx, initial_state)

        # Per-diagram structural-change warning.
        self._maybe_warn_structural_change_for_bundle(bundle, ctx)

        # Drive the cached JIT kernel.
        sim_state = bundle["kernel"](ctx)
        results_data = sim_state.results_data
        time, outputs = results_data.finalize()

        final_context = (
            sim_state.context if bundle["opts_resolved"].return_context else None
        )

        self.n_runs += 1

        return SimulationResults(
            final_context,
            time=time,
            outputs=outputs,
            parameters=dict(parameters) if parameters else None,
        )

    def _maybe_warn_structural_change_for_bundle(
        self, bundle: dict[str, Any], ctx
    ) -> None:
        """Per-diagram analogue of :meth:`_maybe_warn_structural_change`."""
        ctx_sig = self._signature(ctx)
        cached = bundle["cached_ctx_sig"]
        if cached is None:
            bundle["cached_ctx_sig"] = ctx_sig
            return
        if ctx_sig != cached:
            warnings.warn(
                "FastRestartSimulator: detected a structural change in "
                "the context (parameter pytree or continuous-state "
                "shape/dtype) between runs of the same cached diagram — "
                "the JIT cache will miss and the kernel will be "
                "recompiled, defeating the fast-restart benefit.  "
                "Ensure parameter / initial_state values keep the same "
                "shapes/dtypes across runs.",
                UserWarning,
                stacklevel=3,
            )
            bundle["cached_ctx_sig"] = ctx_sig

    # ------------------------------------------------------------------
    # Initial-state override + structural-change detection
    # ------------------------------------------------------------------

    def _apply_initial_state(self, ctx, initial_state: Any):
        """Return ``ctx`` with its continuous state replaced by ``initial_state``.

        Accepts either a single array (auto-wrapped for the common
        single-block case on a :class:`DiagramContext`) or a sequence of
        arrays (one per continuous-state subcontext).  For a
        :class:`LeafContext`, a single array is passed through directly.
        """
        from ..framework.context import DiagramContext, LeafContext

        if isinstance(ctx, LeafContext):
            xc_new = jnp.asarray(initial_state)
            return ctx.with_continuous_state(xc_new)

        if isinstance(ctx, DiagramContext):
            # ``with_continuous_state`` expects one array per
            # continuous-state subcontext.
            n_xc = len(ctx.continuous_subcontexts)
            if n_xc == 0:
                raise ValueError(
                    "FastRestartSimulator.run(initial_state=...): the "
                    "diagram has no continuous-state blocks; nothing to "
                    "override."
                )
            if isinstance(initial_state, (list, tuple)):
                xs = [jnp.asarray(x) for x in initial_state]
            else:
                # Convenience: a single array for a single-block
                # diagram.
                if n_xc != 1:
                    raise ValueError(
                        "FastRestartSimulator.run(initial_state=...): "
                        "diagram has "
                        f"{n_xc} continuous-state blocks but a single "
                        "array was passed; pass a list/tuple of arrays "
                        "matching ctx.continuous_state ordering."
                    )
                xs = [jnp.asarray(initial_state)]
            if len(xs) != n_xc:
                raise ValueError(
                    "FastRestartSimulator.run(initial_state=...): "
                    f"expected {n_xc} arrays (one per continuous-state "
                    f"block); got {len(xs)}."
                )
            return ctx.with_continuous_state(xs)

        # Fall through: unknown context type — defer to the duck-typed
        # ``with_continuous_state`` and hope for the best.
        return ctx.with_continuous_state(initial_state)

    @staticmethod
    def _signature(value) -> tuple:
        """Return a hashable ``(treedef, leaf-shape/dtype tuple)`` signature.

        Two values share a signature iff a ``jax.jit`` keyed on their
        pytree structure would reuse the same compiled program.  Used to
        detect structural changes that will force a recompile.
        """
        leaves, treedef = jax.tree_util.tree_flatten(value)
        leaf_sig = tuple(
            (
                tuple(getattr(x, "shape", ())),
                str(getattr(x, "dtype", type(x).__name__)),
            )
            for x in leaves
        )
        return (str(treedef), leaf_sig)

    def _maybe_warn_structural_change(self, ctx) -> None:
        """Emit a ``UserWarning`` if this run will force a JIT recompile.

        The kernel is JIT-compiled with ``ctx`` as its single argument,
        so the cache key is the abstract signature (treedef + leaf
        shapes/dtypes) of the patched context.  If that signature
        differs from what the kernel was last driven with, JAX will
        recompile — defeating the fast-restart benefit.

        On the first invocation we just record the baseline signature
        (no warning — the first call is expected to compile).
        """
        ctx_sig = self._signature(ctx)

        if self._cached_ctx_sig is None:
            self._cached_ctx_sig = ctx_sig
            return

        if ctx_sig != self._cached_ctx_sig:
            warnings.warn(
                "FastRestartSimulator: detected a structural change in "
                "the context (parameter pytree or continuous-state "
                "shape/dtype) between runs — the JIT cache will miss "
                "and the kernel will be recompiled, defeating the "
                "fast-restart benefit.  Call sim.reset(diagram=...) "
                "and rebuild with the new structure if this is "
                "intentional, or ensure parameter / initial_state "
                "values keep the same shapes/dtypes across runs.",
                UserWarning,
                stacklevel=3,
            )
            # Update the cached signature so we only warn once per
            # structural change (not on every subsequent run with the
            # new shape).
            self._cached_ctx_sig = ctx_sig

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.

Source code in jaxonomy/simulation/fast_restart.py

def close(self) -> None:
    """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.
    """
    self._kernel = None
    self._base_ctx = None
    self._sim = None
    self._opts_resolved = None
    self._cached_ctx_sig = None
    self._diagram_kernel_cache.clear()
    self._diagram_pool.clear()

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".	None

Source code in jaxonomy/simulation/fast_restart.py

def reset(self, diagram: Diagram | None = None) -> 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.

    Args:
        diagram: 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".
    """
    if diagram is not None:
        self.system = diagram
    self.close()

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.	None
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).	None

Returns:

Name	Type	Description
A	SimulationResults	class:SimulationResults populated with time,
	SimulationResults	outputs, and (when
	SimulationResults	options.return_context=True) context.

Source code in jaxonomy/simulation/fast_restart.py

def run(
    self,
    parameters: dict[str, Any] | None = None,
    initial_state: Any | None = None,
) -> SimulationResults:
    """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.

    Args:
        parameters: 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: 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:
        A :class:`SimulationResults` populated with ``time``,
        ``outputs``, and (when
        ``options.return_context=True``) ``context``.
    """
    if self._kernel is None:
        self._build()

    # Patch the base context with the per-run parameter overrides.
    ctx = self._base_ctx
    if parameters:
        updates = {k: jnp.asarray(v) for k, v in parameters.items()}
        ctx = _pure_patch_context(ctx, updates)

    # Apply the initial-state override (if any) on top of the
    # parameter patch.  We do this *after* the parameter patch so
    # that user-supplied initial state always wins regardless of
    # whether the diagram exposed the IC as a parameter as well.
    if initial_state is not None:
        ctx = self._apply_initial_state(ctx, initial_state)

    # Emit a one-time structural-change warning if this run will
    # force a JIT recompile.  Compute the signature *after* all
    # context patches are applied so we compare against what the
    # kernel actually sees.
    self._maybe_warn_structural_change(ctx)

    # Drive the JIT-compiled kernel.
    sim_state = self._kernel(ctx)
    results_data = sim_state.results_data
    time, outputs = results_data.finalize()

    final_context = (
        sim_state.context if self._opts_resolved.return_context else None
    )

    self.n_runs += 1

    return SimulationResults(
        final_context,
        time=time,
        outputs=outputs,
        parameters=dict(parameters) if parameters else None,
    )

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.	None

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.

Source code in jaxonomy/simulation/fast_restart.py

def run_batch(
    self,
    parameters_batch: dict[str, Any],
    initial_states_batch: Any | None = None,
) -> "BatchSimulationResults":
    """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.

    Args:
        parameters_batch: ``{path: (N, ...) array}`` — same convention
            as :func:`simulate_batch`'s ``param_batches``.  Every value
            must have the same leading batch size ``N``.
        initial_states_batch: 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:
        A :class:`BatchSimulationResults` with
        ``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.
    """
    from .batch import BatchSimulationResults, _infer_batch_size

    if not parameters_batch:
        raise ValueError(
            "FastRestartSimulator.run_batch: parameters_batch must be "
            "non-empty.  For a single warm-restart simulation pass "
            "parameters=... to .run() instead."
        )

    # Ensure the kernel + base context are built (lazy on first call).
    if self._kernel is None:
        self._build()

    n = _infer_batch_size(parameters_batch)
    stacked = {path: jnp.asarray(arr) for path, arr in parameters_batch.items()}

    # Build a batched context: broadcast every leaf of the base ctx to
    # shape (N, ...), then patch in the per-batch parameter values
    # (which already have the leading N axis).
    batched_ctx = jax.tree_util.tree_map(
        lambda x: jnp.broadcast_to(
            jnp.asarray(x)[None], (n,) + jnp.asarray(x).shape
        ),
        self._base_ctx,
    )
    batched_ctx = _pure_patch_context(batched_ctx, stacked)

    # Optional batched initial-state override.  Apply *after* the
    # parameter patch so the user's IC always wins.
    if initial_states_batch is not None:
        batched_ctx = self._apply_batched_initial_state(
            batched_ctx, initial_states_batch, n,
        )

    # vmap the cached kernel.  ``axis_name="batch"`` mirrors
    # :func:`simulate_batch`'s vmap path so any
    # ``fold_in_batch_index=True`` blocks (T-122-followup) work
    # identically here.
    batched_kernel = jax.vmap(self._kernel, axis_name="batch")
    batch_sim_states = batched_kernel(batched_ctx)
    results_data_batch = batch_sim_states.results_data

    # Finalize each batch row (results_data carries variable-length
    # buffers; finalize must run per-element).  Mirror the
    # numpy-host interpolation pattern from
    # :func:`simulate_batch._vmap_path` so output shapes line up
    # across rows with different step counts.
    import warnings as _warnings

    import numpy as _np

    from .batch import _interp_on_time_np

    time_ref = None
    out_lists: dict[str, list] = {k: [] for k in self.recorded_signals}

    for i in range(n):
        rd_i = jax.tree_util.tree_map(lambda x: x[i], results_data_batch)
        time_i, outputs_i = rd_i.finalize()
        if time_ref is None:
            time_ref = time_i
            for sig_name in self.recorded_signals:
                out_lists[sig_name].append(_np.asarray(outputs_i[sig_name]))
            continue

        t_ref_end = float(time_ref[-1])
        t_run_end = float(time_i[-1])
        if abs(t_run_end - t_ref_end) / max(abs(t_ref_end), 1e-10) > 0.01:
            _warnings.warn(
                f"FastRestartSimulator.run_batch: row {i} ended at "
                f"t={t_run_end:.4g} but reference row ended at "
                f"t={t_ref_end:.4g}.  Outputs will be interpolated "
                "(clamped) to fill the time grid.",
                UserWarning,
                stacklevel=3,
            )
        same_grid = (
            time_i.shape == time_ref.shape
            and _np.array_equal(time_i, time_ref)
        )
        for sig_name in self.recorded_signals:
            if same_grid:
                out_lists[sig_name].append(_np.asarray(outputs_i[sig_name]))
            else:
                out_lists[sig_name].append(
                    _interp_on_time_np(outputs_i[sig_name], time_i, time_ref)
                )

    stacked_out = {k: _np.stack(vs, axis=0) for k, vs in out_lists.items()}
    self.n_runs += n
    return BatchSimulationResults(
        time=time_ref, outputs=stacked_out, used_vmap=True,
    )

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.	None
initial_state	Any | None	Optional override for the simulator's continuous state at t = t_span[0] — same convention as :meth:run.	None
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.	None

Returns:

Name	Type	Description
A	SimulationResults	class:SimulationResults populated with time,
	SimulationResults	outputs, and (when
	SimulationResults	options.return_context=True) context.

Source code in jaxonomy/simulation/fast_restart.py

def run_with_diagram(
    self,
    diagram: Diagram,
    parameters: dict[str, Any] | None = None,
    initial_state: Any | None = None,
    recorded_signals: dict[str, OutputPort] | None = None,
) -> SimulationResults:
    """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.

    Args:
        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.
        parameters: Optional dot-path mapping ``{"block.param":
            value, ...}`` — same convention as :meth:`run`.
        initial_state: Optional override for the simulator's
            continuous state at ``t = t_span[0]`` — same convention
            as :meth:`run`.
        recorded_signals: 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:
        A :class:`SimulationResults` populated with ``time``,
        ``outputs``, and (when
        ``options.return_context=True``) ``context``.
    """
    cache_key = id(diagram)
    bundle = self._diagram_kernel_cache.get(cache_key)
    if bundle is None:
        # First call for this diagram — build the kernel.
        sigs = recorded_signals if recorded_signals is not None else self.recorded_signals
        if not sigs:
            raise ValueError(
                "FastRestartSimulator.run_with_diagram(...): no "
                "recorded_signals available for this diagram.  Pass "
                "recorded_signals={'name': port, ...} on the first "
                "call for each diagram in your pool — the ports are "
                "diagram-specific and must match the diagram passed "
                "in this call."
            )
        bundle = self._build_for_diagram(diagram, sigs)
        self._diagram_kernel_cache[cache_key] = bundle
        # Hold a strong reference so ``id(diagram)`` stays valid
        # (Python may otherwise reuse the id of a GC'd object).
        self._diagram_pool[cache_key] = diagram

    # Patch the cached base context with per-run overrides.
    ctx = bundle["base_ctx"]
    if parameters:
        updates = {k: jnp.asarray(v) for k, v in parameters.items()}
        ctx = _pure_patch_context(ctx, updates)
    if initial_state is not None:
        ctx = self._apply_initial_state(ctx, initial_state)

    # Per-diagram structural-change warning.
    self._maybe_warn_structural_change_for_bundle(bundle, ctx)

    # Drive the cached JIT kernel.
    sim_state = bundle["kernel"](ctx)
    results_data = sim_state.results_data
    time, outputs = results_data.finalize()

    final_context = (
        sim_state.context if bundle["opts_resolved"].return_context else None
    )

    self.n_runs += 1

    return SimulationResults(
        final_context,
        time=time,
        outputs=outputs,
        parameters=dict(parameters) if parameters else None,
    )

LazyResults dataclass

A deferred-evaluation wrapper around a :class:SimulationResults.

Construct via :meth:SimulationResults.lazy rather than directly.

Source code in jaxonomy/simulation/lazy_results.py

@dataclass
class LazyResults:
    """A deferred-evaluation wrapper around a :class:`SimulationResults`.

    Construct via :meth:`SimulationResults.lazy` rather than directly.
    """

    _outputs: dict[str, np.ndarray]
    _time: np.ndarray
    _ops: list[_Op] = field(default_factory=list)
    _use_polars: bool = False
    _use_duckdb: bool = False
    # Stored as ``Any`` to avoid a hard import of duckdb at module load.
    _duckdb_conn: Optional[Any] = None
    # T-108 phase 1: per-signal native sample-time vectors carried
    # over from :class:`SimulationResults.per_signal_times` (T-013 / T-013a).
    # ``None`` means "every signal shares ``self._time``" — the legacy
    # default that keeps the fluent API byte-equivalent.
    _per_signal_times: Optional[dict[str, np.ndarray]] = None

    # ── factory ──────────────────────────────────────────────────────────

    @classmethod
    def from_results(cls, results: "SimulationResults") -> "LazyResults":
        if results.outputs is None:
            raise ValueError(
                "LazyResults: SimulationResults has no outputs.  Pass "
                "recorded_signals= to simulate() first."
            )
        outputs = {k: np.asarray(v) for k, v in results.outputs.items()}
        time = np.asarray(results.time) if results.time is not None else np.zeros(0)
        per_signal_times: Optional[dict[str, np.ndarray]] = None
        if getattr(results, "per_signal_times", None) is not None:
            per_signal_times = {
                k: np.asarray(v) for k, v in results.per_signal_times.items()
            }
        return cls(_outputs=outputs, _time=time, _per_signal_times=per_signal_times)

    @classmethod
    def from_parquet(cls, path, backend: str = "polars") -> "LazyResults":
        """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.
        """
        if backend not in {"polars", "duckdb"}:
            raise ValueError(
                f"LazyResults.from_parquet: unknown backend {backend!r}; "
                f"expected 'polars' or 'duckdb'."
            )
        try:
            import polars as pl
        except ImportError as e:
            raise ImportError(
                "LazyResults.from_parquet requires polars.  "
                "Install with `pip install polars`."
            ) from e

        df = pl.read_parquet(str(path))
        cols = df.columns
        if "time" not in cols:
            raise ValueError(
                f"LazyResults.from_parquet: file {path!r} has no 'time' column "
                f"(columns={cols})."
            )
        time = np.asarray(df["time"].to_numpy())
        # Re-collapse name__i columns into vector-valued arrays.
        outputs: dict[str, np.ndarray] = {}
        groups: dict[str, dict[int, str]] = {}
        scalars: list[str] = []
        for c in cols:
            if c == "time":
                continue
            if "__" in c:
                base, _, idx_str = c.rpartition("__")
                try:
                    idx = int(idx_str)
                except ValueError:
                    scalars.append(c)
                    continue
                groups.setdefault(base, {})[idx] = c
            else:
                scalars.append(c)
        for c in scalars:
            outputs[c] = np.asarray(df[c].to_numpy())
        for base, idx_map in groups.items():
            ordered = [df[idx_map[i]].to_numpy() for i in sorted(idx_map)]
            outputs[base] = np.stack([np.asarray(a) for a in ordered], axis=-1)
        if backend == "duckdb":
            try:
                import duckdb  # noqa: F401
            except ImportError as e:
                raise ImportError(
                    "LazyResults.from_parquet(backend='duckdb') requires "
                    "duckdb.  Install with `pip install duckdb`."
                ) from e
            conn = duckdb.connect()
            return cls(
                _outputs=outputs,
                _time=time,
                _use_duckdb=True,
                _duckdb_conn=conn,
            )
        return cls(_outputs=outputs, _time=time, _use_polars=True)

    # ── backend opt-in (T-015a) ──────────────────────────────────────────

    def with_polars_backend(self) -> "LazyResults":
        """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.
        """
        return LazyResults(
            _outputs=self._outputs,
            _time=self._time,
            _ops=list(self._ops),
            _use_polars=True,
            _per_signal_times=self._per_signal_times,
        )

    def with_duckdb_backend(self, connection=None) -> "LazyResults":
        """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.
        """
        if connection is None:
            try:
                import duckdb
            except ImportError as e:
                raise ImportError(
                    "LazyResults.with_duckdb_backend: duckdb is not "
                    "installed.  Install with `pip install duckdb`."
                ) from e
            connection = duckdb.connect()
        return LazyResults(
            _outputs=self._outputs,
            _time=self._time,
            _ops=list(self._ops),
            _use_duckdb=True,
            _duckdb_conn=connection,
            _per_signal_times=self._per_signal_times,
        )

    # ── lazy operations ──────────────────────────────────────────────────

    def _chain(self, op: _Op) -> "LazyResults":
        return LazyResults(
            _outputs=self._outputs,
            _time=self._time,
            _ops=self._ops + [op],
            _use_polars=self._use_polars,
            _use_duckdb=self._use_duckdb,
            _duckdb_conn=self._duckdb_conn,
            _per_signal_times=self._per_signal_times,
        )

    def select(self, *signals: str) -> "LazyResults":
        """Project to a subset of signals (defers)."""
        names = list(signals)

        def _apply(out: dict, t: np.ndarray):
            missing = [s for s in names if s not in out]
            if missing:
                raise KeyError(
                    f"LazyResults.select: unknown signal(s) {missing!r}.  "
                    f"Available: {list(out)}"
                )
            return {s: out[s] for s in names}, t

        def _polars_apply(lf, expanded_cols):
            # Keep "time" plus every expanded variant of the requested
            # signals (e.g. select("v") keeps both v__0 and v__1 for a
            # vector-valued v).
            keep = ["time"]
            existing = lf.collect_schema().names() if hasattr(lf, "collect_schema") else lf.columns
            for s in names:
                hit = [c for c in existing if c == s or c.startswith(f"{s}__")]
                if not hit:
                    raise KeyError(
                        f"LazyResults.select: unknown signal {s!r}.  "
                        f"Available: {[c for c in existing if c != 'time']}"
                    )
                keep.extend(hit)
            return lf.select(keep)

        def _duckdb_apply(plan, expanded_cols):
            keep = ["time"]
            existing = list(expanded_cols)
            for s in names:
                hit = [c for c in existing if c == s or c.startswith(f"{s}__")]
                if not hit:
                    raise KeyError(
                        f"LazyResults.select: unknown signal {s!r}.  "
                        f"Available: {[c for c in existing if c != 'time']}"
                    )
                keep.extend(hit)
            return plan.with_select(keep)

        return self._chain(
            _Op(
                name=f"select{tuple(names)!r}",
                apply=_apply,
                polars_apply=_polars_apply,
                duckdb_apply=_duckdb_apply,
            )
        )

    def where(self, mask) -> "LazyResults":
        """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"``).
        """

        def _apply(out: dict, t: np.ndarray):
            if callable(mask):
                m = mask(t, out)
            elif isinstance(mask, str):
                # Restrict to a known-safe globals dict.  Each signal is
                # available by name; ``t`` is the time vector.
                env = {"t": t, **out, "np": np}
                m = eval(mask, {"__builtins__": {}}, env)  # noqa: S307
            else:
                m = mask
            m = np.asarray(m, dtype=bool)
            if m.shape != t.shape:
                raise ValueError(
                    f"LazyResults.where: mask shape {m.shape} does not match "
                    f"time shape {t.shape}."
                )
            new_out = {k: v[m] if v.ndim == 1 else v[m, ...] for k, v in out.items()}
            return new_out, t[m]

        # Polars translator — only available for string expressions and
        # boolean-array masks.  Callable predicates fall back to eager-
        # numpy with a RuntimeWarning at materialise time.
        polars_apply: Optional[Callable[[Any, list[str]], Any]]
        duckdb_apply: Optional[Callable[[Any, list[str]], Any]]
        if callable(mask):
            polars_apply = None
            duckdb_apply = None
        elif isinstance(mask, str):
            expr_str = mask

            def polars_apply(lf, expanded_cols, _expr_str=expr_str):  # type: ignore[misc]
                import re

                import polars as pl

                # ``t`` is the time column in polars-land. Use a word-boundary
                # substitution so signal names that merely end in 't' (e.g.
                # ``out``, ``count``) are not mangled — matches the duckdb
                # path's _python_predicate_to_sql.
                sql_expr = re.sub(r"\bt\b", "time", _expr_str)
                # Common case: simple "x > 0.5" — let polars.sql_expr
                # handle it.  Raise a clear error on failure.
                try:
                    return lf.filter(pl.sql_expr(sql_expr))
                except Exception:
                    raise RuntimeError(
                        f"LazyResults.where: polars cannot translate "
                        f"expression {_expr_str!r}; use a numpy mask or "
                        f"omit .with_polars_backend()."
                    )

            def duckdb_apply(plan, expanded_cols, _expr_str=expr_str):  # type: ignore[misc]
                # Translate Python-style operators in the expression to
                # SQL: ``&`` / ``|`` -> ``AND`` / ``OR``; lone ``t`` -> ``time``.
                # We deliberately do not try to be exhaustive — DuckDB's
                # SQL parser already accepts ``>``, ``<``, ``>=``, ``<=``,
                # ``==`` (folded to ``=``), ``!=``, ``+``, ``-``, ``*``,
                # ``/`` directly.
                sql_expr = _python_predicate_to_sql(_expr_str)
                return plan.with_where(sql_expr)
        else:
            mask_arr = np.asarray(mask, dtype=bool)

            def polars_apply(lf, expanded_cols, _mask=mask_arr):  # type: ignore[misc]
                import polars as pl

                return lf.filter(pl.Series("__mask__", _mask))

            def duckdb_apply(plan, expanded_cols, _mask=mask_arr):  # type: ignore[misc]
                # Boolean-array masks: register a row-aligned mask
                # column on the connection and AND it into the WHERE.
                return plan.with_mask_array(_mask)

        return self._chain(
            _Op(
                name=f"where({mask!r})",
                apply=_apply,
                polars_apply=polars_apply,
                duckdb_apply=duckdb_apply,
            )
        )

    def resample(
        self,
        t_new,
        *,
        method: str = "linear",
    ) -> "LazyResults":
        """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).
        """
        from ..library.lookup_table import interp_1d

        t_new_arr = np.asarray(t_new)

        def _interp_channel(t: np.ndarray, v: np.ndarray) -> np.ndarray:
            """Per-channel interpolator. ``method=='linear'`` stays on
            ``np.interp`` for byte-equivalence with phase 1; everything
            else routes through the T-106 backend."""
            if method == "linear":
                return np.interp(t_new_arr, t, v)
            return np.asarray(interp_1d(t_new_arr, t, v, method=method))

        def _apply(out: dict, t: np.ndarray):
            if t.size == 0:
                raise ValueError(
                    "LazyResults.resample: cannot resample an empty result "
                    "(an upstream .where() may have removed all rows)."
                )
            t_min, t_max = float(t[0]), float(t[-1])
            if np.any(t_new_arr < t_min - 1e-12) or np.any(t_new_arr > t_max + 1e-12):
                raise ValueError(
                    f"LazyResults.resample: requested times outside "
                    f"[{t_min}, {t_max}]."
                )
            new_out = {}
            for k, v in out.items():
                if v.ndim == 1:
                    new_out[k] = _interp_channel(t, v)
                else:
                    new_out[k] = np.stack(
                        [_interp_channel(t, v[:, i]) for i in range(v.shape[1])],
                        axis=-1,
                    )
            return new_out, t_new_arr

        def _polars_apply(lf, expanded_cols, _t_new=t_new_arr, _method=method):
            if _method == "linear":
                return _polars_resample(lf, _t_new)
            # Non-linear methods aren't expressible as a single polars
            # expression — collect, route through the eager path, and
            # re-promote. The pushdown on prior ops still ran lazily.
            import polars as pl

            df = lf.collect()
            t = df["time"].to_numpy()
            out = _collapse_vectorized(df)
            new_out, new_t = _apply(out, t)
            return _eager_to_lazyframe(new_t, new_out)

        return self._chain(
            _Op(
                name=f"resample(len={len(t_new_arr)}, method={method!r})",
                apply=_apply,
                polars_apply=_polars_apply,
            )
        )

    def with_signal(self, name: str, fn: Callable) -> "LazyResults":
        """Derive a new signal ``name`` from existing ones (defers).

        ``fn`` receives ``(t, outputs)`` and returns an array shaped
        like ``time``.
        """

        def _apply(out: dict, t: np.ndarray):
            new_out = dict(out)
            new_out[name] = np.asarray(fn(t, out))
            return new_out, t

        def _polars_apply(lf, expanded_cols, _name=name, _fn=fn):
            # Polars cannot express an arbitrary Python user function
            # natively; collect to compute, then re-lazy.  We still
            # benefit from polars's pushdown on prior ops in the chain
            # (they ran lazily before this point).
            import polars as pl

            df = lf.collect()
            t = df["time"].to_numpy()
            # Reconstruct the outputs dict — vector-valued signals are
            # exploded as ``base__i`` so we re-collapse them.
            out = _collapse_vectorized(df)
            new_col = np.asarray(_fn(t, out))
            return df.with_columns(pl.Series(_name, new_col)).lazy()

        return self._chain(
            _Op(
                name=f"with_signal({name!r})",
                apply=_apply,
                polars_apply=_polars_apply,
            )
        )

    # ── per-signal native cadence (T-108 phase 1) ─────────────────────

    def signal(self, name: str) -> tuple[np.ndarray, np.ndarray]:
        """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:
            KeyError: if ``name`` is not a recorded signal.
        """
        if name not in self._outputs:
            raise KeyError(
                f"LazyResults.signal: unknown signal {name!r}.  "
                f"Available: {list(self._outputs)}"
            )
        value = np.asarray(self._outputs[name])
        time = self._native_time_for(name)
        if value.shape[0] != time.shape[0]:
            # Mode B "default" signal: outputs is full length, time is
            # deduplicated.  Project value onto the deduplicated times.
            global_t = np.asarray(self._time)
            idx = np.searchsorted(global_t, time)
            if global_t.shape[0] > 0:
                idx = np.clip(idx, 0, global_t.shape[0] - 1)
            value = value[idx]
        return time, value

    def cadence_of(self, name: str) -> str:
        """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:
            KeyError: if ``name`` is not a recorded signal.
        """
        if name not in self._outputs:
            raise KeyError(
                f"LazyResults.cadence_of: unknown signal {name!r}.  "
                f"Available: {list(self._outputs)}"
            )
        if (
            self._per_signal_times is None
            or name not in self._per_signal_times
        ):
            return "default"
        global_t = np.asarray(self._time)
        sig_t = np.asarray(self._per_signal_times[name])
        sig_v = np.asarray(self._outputs[name])
        if global_t.shape[0] == 0:
            return "default"
        if sig_t.shape[0] == global_t.shape[0]:
            return "continuous"
        if sig_v.shape[0] == sig_t.shape[0]:
            # Mode A: both arrays were trimmed to the schedule.
            return "periodic"
        # Mode B fallback: times deduplicated, values still full length.
        return "event-driven"

    def align_to(self, name: str) -> "LazyResults":
        """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:
            KeyError: if ``name`` is not a recorded signal.
        """
        if name not in self._outputs:
            raise KeyError(
                f"LazyResults.align_to: unknown signal {name!r}.  "
                f"Available: {list(self._outputs)}"
            )
        target_time = self._native_time_for(name)
        return self.resample(target_time)

    def _native_time_for(self, name: str) -> np.ndarray:
        """Return the per-signal native time vector for ``name``.

        Falls back to the global :attr:`_time` vector when the lazy
        results object was constructed from a :class:`SimulationResults`
        without ``per_signal_times`` (the legacy default-off path) or
        when ``name`` is not in the per-signal map.  Mirrors
        :meth:`SimulationResults.time_for`.
        """
        if (
            self._per_signal_times is not None
            and name in self._per_signal_times
        ):
            return np.asarray(self._per_signal_times[name])
        return np.asarray(self._time)

    # ── terminals (eager-numpy path) ─────────────────────────────────────

    def _collect_eager(self) -> dict:
        outputs, time = self._outputs, self._time
        for op in self._ops:
            outputs, time = op.apply(outputs, time)
        return {"time": time, **outputs}

    def collect(self) -> dict:
        """Materialise the chain.  Returns ``{"time": t, **signals}``."""
        if self._use_duckdb:
            return self._collect_duckdb()
        if self._use_polars:
            df = self._to_polars_df()
            time = np.asarray(df["time"].to_numpy())
            outputs = _collapse_vectorized(df)
            return {"time": time, **outputs}
        return self._collect_eager()

    def to_numpy(self) -> dict:
        """Alias for :meth:`collect`."""
        return self.collect()

    def to_pandas(self):
        """Materialise to a ``pandas.DataFrame`` (requires pandas).

        Vector-valued signals are exploded into ``name__0``, ``name__1`` columns.
        """
        if self._use_duckdb:
            return self._duckdb_to_pandas()
        if self._use_polars:
            return self._to_polars_df().to_pandas()
        try:
            import pandas as pd
        except ImportError as e:
            raise ImportError(
                "LazyResults.to_pandas: pandas is not installed.  "
                "Install with `pip install pandas` or use .to_numpy()."
            ) from e

        materialised = self._collect_eager()
        time = materialised.pop("time")
        cols = {"time": time}
        for k, v in materialised.items():
            if v.ndim == 1:
                cols[k] = v
            else:
                for i in range(v.shape[-1]):
                    cols[f"{k}__{i}"] = v[..., i]
        return pd.DataFrame(cols)

    def to_polars(self):
        """Materialise to a ``polars.DataFrame`` (requires polars)."""
        try:
            import polars as pl  # noqa: F401
        except ImportError as e:
            raise ImportError(
                "LazyResults.to_polars: polars is not installed.  "
                "Install with `pip install polars` or use .to_pandas()."
            ) from e

        if self._use_duckdb:
            return self._duckdb_to_polars()
        if self._use_polars:
            return self._to_polars_df()

        materialised = self._collect_eager()
        return _eager_dict_to_polars(materialised)

    def to_hdf5(
        self,
        path,
        key: str = "results",
        chunk_size: int = 10_000,
    ) -> None:
        """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/<name>__<i>`` 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.
        """
        try:
            import h5py
        except ImportError as e:
            raise ImportError(
                "LazyResults.to_hdf5: h5py is not installed.  "
                "Install with `pip install h5py`."
            ) from e

        del key  # reserved; see docstring
        if chunk_size <= 0:
            raise ValueError(
                f"LazyResults.to_hdf5: chunk_size must be positive (got {chunk_size})."
            )

        with h5py.File(str(path), "w") as f:
            out_grp = f.create_group("outputs")
            time_ds: Optional[Any] = None
            sig_dsets: dict[str, Any] = {}
            written = 0
            for chunk in self._iter_chunks(chunk_size):
                time_chunk = chunk["time"]
                n = time_chunk.shape[0]
                if n == 0:
                    continue
                if time_ds is None:
                    time_ds = f.create_dataset(
                        "time",
                        shape=(0,),
                        maxshape=(None,),
                        dtype=time_chunk.dtype,
                        chunks=(min(chunk_size, max(n, 1)),),
                    )
                time_ds.resize((written + n,))
                time_ds[written : written + n] = time_chunk

                for name, arr in chunk.items():
                    if name == "time":
                        continue
                    if name not in sig_dsets:
                        sig_dsets[name] = out_grp.create_dataset(
                            name,
                            shape=(0,),
                            maxshape=(None,),
                            dtype=arr.dtype,
                            chunks=(min(chunk_size, max(n, 1)),),
                        )
                    ds = sig_dsets[name]
                    ds.resize((written + n,))
                    ds[written : written + n] = arr
                written += n

    def to_zarr(self, path, chunk_size: int = 10_000) -> None:
        """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/<name>__<i>``).  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.
        """
        try:
            import zarr
        except ImportError as e:
            raise ImportError(
                "LazyResults.to_zarr: zarr is not installed.  "
                "Install with `pip install zarr`."
            ) from e

        if chunk_size <= 0:
            raise ValueError(
                f"LazyResults.to_zarr: chunk_size must be positive (got {chunk_size})."
            )

        root = zarr.open_group(str(path), mode="w")
        out_grp = root.create_group("outputs")
        time_arr: Optional[Any] = None
        sig_arrs: dict[str, Any] = {}
        written = 0
        for chunk in self._iter_chunks(chunk_size):
            time_chunk = chunk["time"]
            n = time_chunk.shape[0]
            if n == 0:
                continue
            if time_arr is None:
                time_arr = root.create_array(
                    "time",
                    shape=(0,),
                    chunks=(min(chunk_size, max(n, 1)),),
                    dtype=time_chunk.dtype,
                )
            time_arr.resize((written + n,))
            time_arr[written : written + n] = time_chunk

            for name, arr in chunk.items():
                if name == "time":
                    continue
                if name not in sig_arrs:
                    sig_arrs[name] = out_grp.create_array(
                        name,
                        shape=(0,),
                        chunks=(min(chunk_size, max(n, 1)),),
                        dtype=arr.dtype,
                    )
                za = sig_arrs[name]
                za.resize((written + n,))
                za[written : written + n] = arr
            written += n

    def _iter_chunks(self, chunk_size: int):
        """Yield ``{column_name: np.ndarray}`` dicts of at most ``chunk_size``
        rows each.

        For the default eager-numpy and the DuckDB / polars backends we
        materialise once and slice the result.  This is "honest
        streaming" only in the sense that the writer never holds more
        than one chunk *as its own copy* — the underlying source frame
        may already be in memory.  True out-of-core streaming would
        require a polars ``sink_batches`` plan, which the writer can
        layer on top of this iterator in a follow-up.

        Vector-valued signals are exploded into ``name__i`` keys so the
        consumer can treat every column as a 1-D array.
        """
        materialised = self.collect()
        time = np.asarray(materialised.pop("time"))
        n = time.shape[0]
        # Pre-explode vector-valued signals so we don't allocate
        # ``name__i`` arrays on every slice.
        exploded: dict[str, np.ndarray] = {}
        for k, v in materialised.items():
            arr = np.asarray(v)
            if arr.ndim == 1:
                exploded[k] = arr
            else:
                for i in range(arr.shape[-1]):
                    exploded[f"{k}__{i}"] = arr[..., i]

        if n == 0:
            yield {"time": time, **exploded}
            return

        for start in range(0, n, chunk_size):
            stop = min(start + chunk_size, n)
            chunk: dict[str, np.ndarray] = {"time": time[start:stop]}
            for k, arr in exploded.items():
                chunk[k] = arr[start:stop]
            yield chunk

    def to_parquet(self, path, batch_size: Optional[int] = 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.
        """
        if self._use_duckdb:
            self._sink_parquet_duckdb(path, batch_size=batch_size)
            return
        if self._use_polars:
            self._sink_parquet_polars(path, batch_size=batch_size)
            return

        try:
            df = self.to_pandas()
            df.to_parquet(path)
            return
        except ImportError:
            pass
        df = self.to_polars()
        df.write_parquet(str(path))

    # ── polars-backend internals (T-015a) ────────────────────────────────

    def _build_lazyframe(self):
        """Build the initial polars LazyFrame from numpy outputs."""
        try:
            import polars as pl
        except ImportError as e:
            raise ImportError(
                "LazyResults: polars backend requested but polars is not "
                "installed.  Install with `pip install polars`."
            ) from e

        data: dict[str, np.ndarray] = {"time": np.asarray(self._time)}
        for k, v in self._outputs.items():
            arr = np.asarray(v)
            if arr.ndim == 1:
                data[k] = arr
            else:
                for i in range(arr.shape[-1]):
                    data[f"{k}__{i}"] = arr[..., i]
        return pl.DataFrame(data).lazy()

    def _to_polars_df(self):
        """Run the op chain through polars; return a materialised DataFrame.

        Per-op fallback: when an op declares no ``polars_apply``
        translator (or that translator declines), the chain is collected
        to numpy, the eager apply runs, and the chain re-enters the
        polars LazyFrame.  Emits :class:`RuntimeWarning` on fallback.
        """
        lf = self._build_lazyframe()
        for op in self._ops:
            if op.polars_apply is None:
                warnings.warn(
                    f"LazyResults.with_polars_backend: op {op.name!r} has no "
                    f"polars equivalent; falling back to eager-numpy for this "
                    f"step.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                df = lf.collect()
                outputs, time = _df_to_eager(df)
                outputs, time = op.apply(outputs, time)
                lf = _eager_to_lazyframe(time, outputs)
            else:
                expanded = (
                    lf.collect_schema().names()
                    if hasattr(lf, "collect_schema")
                    else lf.columns
                )
                lf = op.polars_apply(lf, expanded)
        return lf.collect()

    def _sink_parquet_polars(self, path, batch_size: Optional[int]) -> None:
        """Stream the LazyFrame to parquet via ``sink_parquet``.

        With ``batch_size=None`` we let polars stream the whole plan to
        a single file.  With ``batch_size=N`` we materialise once,
        partition by row count, and write multiple files.
        """
        # Build the full LazyFrame (running any per-op fallbacks).
        # _to_polars_df runs the chain end-to-end and returns a
        # DataFrame; for the "true streaming" case we want to keep the
        # plan lazy.  When every op has a polars translator, we can
        # build lazily and sink_parquet.  Otherwise we fall back to
        # collect+write.
        from pathlib import Path

        path = str(path)
        if batch_size is None and all(op.polars_apply is not None for op in self._ops):
            lf = self._build_lazyframe()
            for op in self._ops:
                expanded = (
                    lf.collect_schema().names()
                    if hasattr(lf, "collect_schema")
                    else lf.columns
                )
                lf = op.polars_apply(lf, expanded)
            try:
                lf.sink_parquet(path)
                return
            except Exception:
                # Some plans (e.g. those built via map_batches) cannot
                # stream — fall through to the materialise path.
                pass

        df = self._to_polars_df()
        if batch_size is None:
            df.write_parquet(path)
            return

        n = df.height
        base = Path(path)
        stem = base.with_suffix("")
        suffix = base.suffix or ".parquet"
        for i, start in enumerate(range(0, n, batch_size)):
            chunk = df.slice(start, batch_size)
            chunk.write_parquet(f"{stem}.{i}{suffix}")

    # ── debug helpers ────────────────────────────────────────────────────

    def explain(self) -> str:
        """Render the deferred operation chain as a human-readable string."""
        if not self._ops:
            tag = "<identity>"
        else:
            tag = " | ".join(op.name for op in self._ops)
        if self._use_duckdb:
            return f"[duckdb] {tag}"
        if self._use_polars:
            return f"[polars] {tag}"
        return tag

    # ── duckdb-backend internals (T-015a-followup-resample-pushdown-duckdb) ─

    def _build_duckdb_plan(self) -> "_DuckDBPlan":
        """Register the (time, outputs) arrays as a DuckDB table and
        return an initial :class:`_DuckDBPlan` selecting all columns."""
        if self._duckdb_conn is None:  # pragma: no cover — defensive
            raise RuntimeError(
                "LazyResults: DuckDB backend requested but no connection "
                "is attached."
            )

        data: dict[str, np.ndarray] = {"time": np.asarray(self._time)}
        for k, v in self._outputs.items():
            arr = np.asarray(v)
            if arr.ndim == 1:
                data[k] = arr
            else:
                for i in range(arr.shape[-1]):
                    data[f"{k}__{i}"] = arr[..., i]

        # Use a unique table name per build so re-registration after a
        # mid-chain fallback does not collide with prior state on the
        # same connection.
        table = f"jaxonomy_lazy_{id(self)}_{_DuckDBPlan._next_id()}"
        self._duckdb_conn.register(table, data)
        cols = list(data.keys())
        return _DuckDBPlan(
            conn=self._duckdb_conn,
            table=table,
            columns=cols,
            select_cols=list(cols),
            where_clauses=[],
            mask_arrays=[],
        )

    def _materialize_duckdb(self) -> "_DuckDBPlan":
        """Run the op chain through DuckDB; return the final plan.

        Per-op fallback: when an op declares no ``duckdb_apply``
        translator, the chain materialises to numpy via fetchnumpy,
        the eager apply runs, and the chain re-registers a fresh
        DuckDB table.  Emits :class:`RuntimeWarning` on fallback.
        """
        plan = self._build_duckdb_plan()
        for op in self._ops:
            if op.duckdb_apply is None:
                warnings.warn(
                    f"LazyResults.with_duckdb_backend: op {op.name!r} has no "
                    f"DuckDB SQL equivalent; falling back to eager-numpy for "
                    f"this step.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                outputs, time = plan.fetch_eager()
                outputs, time = op.apply(outputs, time)
                plan = _eager_to_duckdb_plan(self._duckdb_conn, time, outputs)
            else:
                plan = op.duckdb_apply(plan, plan.select_cols)
        return plan

    def _collect_duckdb(self) -> dict:
        plan = self._materialize_duckdb()
        outputs, time = plan.fetch_eager()
        return {"time": time, **outputs}

    def _duckdb_to_pandas(self):
        plan = self._materialize_duckdb()
        return plan.fetch_pandas()

    def _duckdb_to_polars(self):
        plan = self._materialize_duckdb()
        return plan.fetch_polars()

    def _sink_parquet_duckdb(self, path, batch_size: Optional[int]) -> None:
        from pathlib import Path

        path_str = str(path)
        plan = self._materialize_duckdb()
        if batch_size is None:
            plan.copy_to_parquet(path_str)
            return

        # Batched: fall back to materialise + slice.  DuckDB supports
        # partitioned writes via PARTITION_BY, but only on a column,
        # not a row-count chunk size; row-chunked output is rare enough
        # that we don't try to optimise it.
        outputs, time = plan.fetch_eager()
        n = len(time)
        base = Path(path_str)
        stem = base.with_suffix("")
        suffix = base.suffix or ".parquet"
        for i, start in enumerate(range(0, n, batch_size)):
            stop = min(start + batch_size, n)
            sub_t = time[start:stop]
            sub_out = {k: v[start:stop] for k, v in outputs.items()}
            sub_plan = _eager_to_duckdb_plan(self._duckdb_conn, sub_t, sub_out)
            sub_plan.copy_to_parquet(f"{stem}.{i}{suffix}")

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.

Source code in jaxonomy/simulation/lazy_results.py

def align_to(self, name: str) -> "LazyResults":
    """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:
        KeyError: if ``name`` is not a recorded signal.
    """
    if name not in self._outputs:
        raise KeyError(
            f"LazyResults.align_to: unknown signal {name!r}.  "
            f"Available: {list(self._outputs)}"
        )
    target_time = self._native_time_for(name)
    return self.resample(target_time)

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.

Source code in jaxonomy/simulation/lazy_results.py

def cadence_of(self, name: str) -> str:
    """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:
        KeyError: if ``name`` is not a recorded signal.
    """
    if name not in self._outputs:
        raise KeyError(
            f"LazyResults.cadence_of: unknown signal {name!r}.  "
            f"Available: {list(self._outputs)}"
        )
    if (
        self._per_signal_times is None
        or name not in self._per_signal_times
    ):
        return "default"
    global_t = np.asarray(self._time)
    sig_t = np.asarray(self._per_signal_times[name])
    sig_v = np.asarray(self._outputs[name])
    if global_t.shape[0] == 0:
        return "default"
    if sig_t.shape[0] == global_t.shape[0]:
        return "continuous"
    if sig_v.shape[0] == sig_t.shape[0]:
        # Mode A: both arrays were trimmed to the schedule.
        return "periodic"
    # Mode B fallback: times deduplicated, values still full length.
    return "event-driven"

collect()

Materialise the chain. Returns {"time": t, **signals}.

Source code in jaxonomy/simulation/lazy_results.py

def collect(self) -> dict:
    """Materialise the chain.  Returns ``{"time": t, **signals}``."""
    if self._use_duckdb:
        return self._collect_duckdb()
    if self._use_polars:
        df = self._to_polars_df()
        time = np.asarray(df["time"].to_numpy())
        outputs = _collapse_vectorized(df)
        return {"time": time, **outputs}
    return self._collect_eager()

explain()

Render the deferred operation chain as a human-readable string.

Source code in jaxonomy/simulation/lazy_results.py

def explain(self) -> str:
    """Render the deferred operation chain as a human-readable string."""
    if not self._ops:
        tag = "<identity>"
    else:
        tag = " | ".join(op.name for op in self._ops)
    if self._use_duckdb:
        return f"[duckdb] {tag}"
    if self._use_polars:
        return f"[polars] {tag}"
    return tag

from_parquet(path, backend='polars') classmethod

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.

Source code in jaxonomy/simulation/lazy_results.py

@classmethod
def from_parquet(cls, path, backend: str = "polars") -> "LazyResults":
    """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.
    """
    if backend not in {"polars", "duckdb"}:
        raise ValueError(
            f"LazyResults.from_parquet: unknown backend {backend!r}; "
            f"expected 'polars' or 'duckdb'."
        )
    try:
        import polars as pl
    except ImportError as e:
        raise ImportError(
            "LazyResults.from_parquet requires polars.  "
            "Install with `pip install polars`."
        ) from e

    df = pl.read_parquet(str(path))
    cols = df.columns
    if "time" not in cols:
        raise ValueError(
            f"LazyResults.from_parquet: file {path!r} has no 'time' column "
            f"(columns={cols})."
        )
    time = np.asarray(df["time"].to_numpy())
    # Re-collapse name__i columns into vector-valued arrays.
    outputs: dict[str, np.ndarray] = {}
    groups: dict[str, dict[int, str]] = {}
    scalars: list[str] = []
    for c in cols:
        if c == "time":
            continue
        if "__" in c:
            base, _, idx_str = c.rpartition("__")
            try:
                idx = int(idx_str)
            except ValueError:
                scalars.append(c)
                continue
            groups.setdefault(base, {})[idx] = c
        else:
            scalars.append(c)
    for c in scalars:
        outputs[c] = np.asarray(df[c].to_numpy())
    for base, idx_map in groups.items():
        ordered = [df[idx_map[i]].to_numpy() for i in sorted(idx_map)]
        outputs[base] = np.stack([np.asarray(a) for a in ordered], axis=-1)
    if backend == "duckdb":
        try:
            import duckdb  # noqa: F401
        except ImportError as e:
            raise ImportError(
                "LazyResults.from_parquet(backend='duckdb') requires "
                "duckdb.  Install with `pip install duckdb`."
            ) from e
        conn = duckdb.connect()
        return cls(
            _outputs=outputs,
            _time=time,
            _use_duckdb=True,
            _duckdb_conn=conn,
        )
    return cls(_outputs=outputs, _time=time, _use_polars=True)

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).

Source code in jaxonomy/simulation/lazy_results.py

def resample(
    self,
    t_new,
    *,
    method: str = "linear",
) -> "LazyResults":
    """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).
    """
    from ..library.lookup_table import interp_1d

    t_new_arr = np.asarray(t_new)

    def _interp_channel(t: np.ndarray, v: np.ndarray) -> np.ndarray:
        """Per-channel interpolator. ``method=='linear'`` stays on
        ``np.interp`` for byte-equivalence with phase 1; everything
        else routes through the T-106 backend."""
        if method == "linear":
            return np.interp(t_new_arr, t, v)
        return np.asarray(interp_1d(t_new_arr, t, v, method=method))

    def _apply(out: dict, t: np.ndarray):
        if t.size == 0:
            raise ValueError(
                "LazyResults.resample: cannot resample an empty result "
                "(an upstream .where() may have removed all rows)."
            )
        t_min, t_max = float(t[0]), float(t[-1])
        if np.any(t_new_arr < t_min - 1e-12) or np.any(t_new_arr > t_max + 1e-12):
            raise ValueError(
                f"LazyResults.resample: requested times outside "
                f"[{t_min}, {t_max}]."
            )
        new_out = {}
        for k, v in out.items():
            if v.ndim == 1:
                new_out[k] = _interp_channel(t, v)
            else:
                new_out[k] = np.stack(
                    [_interp_channel(t, v[:, i]) for i in range(v.shape[1])],
                    axis=-1,
                )
        return new_out, t_new_arr

    def _polars_apply(lf, expanded_cols, _t_new=t_new_arr, _method=method):
        if _method == "linear":
            return _polars_resample(lf, _t_new)
        # Non-linear methods aren't expressible as a single polars
        # expression — collect, route through the eager path, and
        # re-promote. The pushdown on prior ops still ran lazily.
        import polars as pl

        df = lf.collect()
        t = df["time"].to_numpy()
        out = _collapse_vectorized(df)
        new_out, new_t = _apply(out, t)
        return _eager_to_lazyframe(new_t, new_out)

    return self._chain(
        _Op(
            name=f"resample(len={len(t_new_arr)}, method={method!r})",
            apply=_apply,
            polars_apply=_polars_apply,
        )
    )

select(*signals)

Project to a subset of signals (defers).

Source code in jaxonomy/simulation/lazy_results.py

def select(self, *signals: str) -> "LazyResults":
    """Project to a subset of signals (defers)."""
    names = list(signals)

    def _apply(out: dict, t: np.ndarray):
        missing = [s for s in names if s not in out]
        if missing:
            raise KeyError(
                f"LazyResults.select: unknown signal(s) {missing!r}.  "
                f"Available: {list(out)}"
            )
        return {s: out[s] for s in names}, t

    def _polars_apply(lf, expanded_cols):
        # Keep "time" plus every expanded variant of the requested
        # signals (e.g. select("v") keeps both v__0 and v__1 for a
        # vector-valued v).
        keep = ["time"]
        existing = lf.collect_schema().names() if hasattr(lf, "collect_schema") else lf.columns
        for s in names:
            hit = [c for c in existing if c == s or c.startswith(f"{s}__")]
            if not hit:
                raise KeyError(
                    f"LazyResults.select: unknown signal {s!r}.  "
                    f"Available: {[c for c in existing if c != 'time']}"
                )
            keep.extend(hit)
        return lf.select(keep)

    def _duckdb_apply(plan, expanded_cols):
        keep = ["time"]
        existing = list(expanded_cols)
        for s in names:
            hit = [c for c in existing if c == s or c.startswith(f"{s}__")]
            if not hit:
                raise KeyError(
                    f"LazyResults.select: unknown signal {s!r}.  "
                    f"Available: {[c for c in existing if c != 'time']}"
                )
            keep.extend(hit)
        return plan.with_select(keep)

    return self._chain(
        _Op(
            name=f"select{tuple(names)!r}",
            apply=_apply,
            polars_apply=_polars_apply,
            duckdb_apply=_duckdb_apply,
        )
    )

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.

Source code in jaxonomy/simulation/lazy_results.py

def signal(self, name: str) -> tuple[np.ndarray, np.ndarray]:
    """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:
        KeyError: if ``name`` is not a recorded signal.
    """
    if name not in self._outputs:
        raise KeyError(
            f"LazyResults.signal: unknown signal {name!r}.  "
            f"Available: {list(self._outputs)}"
        )
    value = np.asarray(self._outputs[name])
    time = self._native_time_for(name)
    if value.shape[0] != time.shape[0]:
        # Mode B "default" signal: outputs is full length, time is
        # deduplicated.  Project value onto the deduplicated times.
        global_t = np.asarray(self._time)
        idx = np.searchsorted(global_t, time)
        if global_t.shape[0] > 0:
            idx = np.clip(idx, 0, global_t.shape[0] - 1)
        value = value[idx]
    return time, value

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/<name>__<i> 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.

Source code in jaxonomy/simulation/lazy_results.py

def to_hdf5(
    self,
    path,
    key: str = "results",
    chunk_size: int = 10_000,
) -> None:
    """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/<name>__<i>`` 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.
    """
    try:
        import h5py
    except ImportError as e:
        raise ImportError(
            "LazyResults.to_hdf5: h5py is not installed.  "
            "Install with `pip install h5py`."
        ) from e

    del key  # reserved; see docstring
    if chunk_size <= 0:
        raise ValueError(
            f"LazyResults.to_hdf5: chunk_size must be positive (got {chunk_size})."
        )

    with h5py.File(str(path), "w") as f:
        out_grp = f.create_group("outputs")
        time_ds: Optional[Any] = None
        sig_dsets: dict[str, Any] = {}
        written = 0
        for chunk in self._iter_chunks(chunk_size):
            time_chunk = chunk["time"]
            n = time_chunk.shape[0]
            if n == 0:
                continue
            if time_ds is None:
                time_ds = f.create_dataset(
                    "time",
                    shape=(0,),
                    maxshape=(None,),
                    dtype=time_chunk.dtype,
                    chunks=(min(chunk_size, max(n, 1)),),
                )
            time_ds.resize((written + n,))
            time_ds[written : written + n] = time_chunk

            for name, arr in chunk.items():
                if name == "time":
                    continue
                if name not in sig_dsets:
                    sig_dsets[name] = out_grp.create_dataset(
                        name,
                        shape=(0,),
                        maxshape=(None,),
                        dtype=arr.dtype,
                        chunks=(min(chunk_size, max(n, 1)),),
                    )
                ds = sig_dsets[name]
                ds.resize((written + n,))
                ds[written : written + n] = arr
            written += n

to_numpy()

Alias for :meth:collect.

Source code in jaxonomy/simulation/lazy_results.py

def to_numpy(self) -> dict:
    """Alias for :meth:`collect`."""
    return self.collect()

to_pandas()

Materialise to a pandas.DataFrame (requires pandas).

Vector-valued signals are exploded into name__0, name__1 columns.

Source code in jaxonomy/simulation/lazy_results.py

def to_pandas(self):
    """Materialise to a ``pandas.DataFrame`` (requires pandas).

    Vector-valued signals are exploded into ``name__0``, ``name__1`` columns.
    """
    if self._use_duckdb:
        return self._duckdb_to_pandas()
    if self._use_polars:
        return self._to_polars_df().to_pandas()
    try:
        import pandas as pd
    except ImportError as e:
        raise ImportError(
            "LazyResults.to_pandas: pandas is not installed.  "
            "Install with `pip install pandas` or use .to_numpy()."
        ) from e

    materialised = self._collect_eager()
    time = materialised.pop("time")
    cols = {"time": time}
    for k, v in materialised.items():
        if v.ndim == 1:
            cols[k] = v
        else:
            for i in range(v.shape[-1]):
                cols[f"{k}__{i}"] = v[..., i]
    return pd.DataFrame(cols)

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.

Source code in jaxonomy/simulation/lazy_results.py

def to_parquet(self, path, batch_size: Optional[int] = 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.
    """
    if self._use_duckdb:
        self._sink_parquet_duckdb(path, batch_size=batch_size)
        return
    if self._use_polars:
        self._sink_parquet_polars(path, batch_size=batch_size)
        return

    try:
        df = self.to_pandas()
        df.to_parquet(path)
        return
    except ImportError:
        pass
    df = self.to_polars()
    df.write_parquet(str(path))

to_polars()

Materialise to a polars.DataFrame (requires polars).

Source code in jaxonomy/simulation/lazy_results.py

def to_polars(self):
    """Materialise to a ``polars.DataFrame`` (requires polars)."""
    try:
        import polars as pl  # noqa: F401
    except ImportError as e:
        raise ImportError(
            "LazyResults.to_polars: polars is not installed.  "
            "Install with `pip install polars` or use .to_pandas()."
        ) from e

    if self._use_duckdb:
        return self._duckdb_to_polars()
    if self._use_polars:
        return self._to_polars_df()

    materialised = self._collect_eager()
    return _eager_dict_to_polars(materialised)

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/<name>__<i>). 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.

Source code in jaxonomy/simulation/lazy_results.py

def to_zarr(self, path, chunk_size: int = 10_000) -> None:
    """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/<name>__<i>``).  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.
    """
    try:
        import zarr
    except ImportError as e:
        raise ImportError(
            "LazyResults.to_zarr: zarr is not installed.  "
            "Install with `pip install zarr`."
        ) from e

    if chunk_size <= 0:
        raise ValueError(
            f"LazyResults.to_zarr: chunk_size must be positive (got {chunk_size})."
        )

    root = zarr.open_group(str(path), mode="w")
    out_grp = root.create_group("outputs")
    time_arr: Optional[Any] = None
    sig_arrs: dict[str, Any] = {}
    written = 0
    for chunk in self._iter_chunks(chunk_size):
        time_chunk = chunk["time"]
        n = time_chunk.shape[0]
        if n == 0:
            continue
        if time_arr is None:
            time_arr = root.create_array(
                "time",
                shape=(0,),
                chunks=(min(chunk_size, max(n, 1)),),
                dtype=time_chunk.dtype,
            )
        time_arr.resize((written + n,))
        time_arr[written : written + n] = time_chunk

        for name, arr in chunk.items():
            if name == "time":
                continue
            if name not in sig_arrs:
                sig_arrs[name] = out_grp.create_array(
                    name,
                    shape=(0,),
                    chunks=(min(chunk_size, max(n, 1)),),
                    dtype=arr.dtype,
                )
            za = sig_arrs[name]
            za.resize((written + n,))
            za[written : written + n] = arr
        written += n

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").

Source code in jaxonomy/simulation/lazy_results.py

def where(self, mask) -> "LazyResults":
    """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"``).
    """

    def _apply(out: dict, t: np.ndarray):
        if callable(mask):
            m = mask(t, out)
        elif isinstance(mask, str):
            # Restrict to a known-safe globals dict.  Each signal is
            # available by name; ``t`` is the time vector.
            env = {"t": t, **out, "np": np}
            m = eval(mask, {"__builtins__": {}}, env)  # noqa: S307
        else:
            m = mask
        m = np.asarray(m, dtype=bool)
        if m.shape != t.shape:
            raise ValueError(
                f"LazyResults.where: mask shape {m.shape} does not match "
                f"time shape {t.shape}."
            )
        new_out = {k: v[m] if v.ndim == 1 else v[m, ...] for k, v in out.items()}
        return new_out, t[m]

    # Polars translator — only available for string expressions and
    # boolean-array masks.  Callable predicates fall back to eager-
    # numpy with a RuntimeWarning at materialise time.
    polars_apply: Optional[Callable[[Any, list[str]], Any]]
    duckdb_apply: Optional[Callable[[Any, list[str]], Any]]
    if callable(mask):
        polars_apply = None
        duckdb_apply = None
    elif isinstance(mask, str):
        expr_str = mask

        def polars_apply(lf, expanded_cols, _expr_str=expr_str):  # type: ignore[misc]
            import re

            import polars as pl

            # ``t`` is the time column in polars-land. Use a word-boundary
            # substitution so signal names that merely end in 't' (e.g.
            # ``out``, ``count``) are not mangled — matches the duckdb
            # path's _python_predicate_to_sql.
            sql_expr = re.sub(r"\bt\b", "time", _expr_str)
            # Common case: simple "x > 0.5" — let polars.sql_expr
            # handle it.  Raise a clear error on failure.
            try:
                return lf.filter(pl.sql_expr(sql_expr))
            except Exception:
                raise RuntimeError(
                    f"LazyResults.where: polars cannot translate "
                    f"expression {_expr_str!r}; use a numpy mask or "
                    f"omit .with_polars_backend()."
                )

        def duckdb_apply(plan, expanded_cols, _expr_str=expr_str):  # type: ignore[misc]
            # Translate Python-style operators in the expression to
            # SQL: ``&`` / ``|`` -> ``AND`` / ``OR``; lone ``t`` -> ``time``.
            # We deliberately do not try to be exhaustive — DuckDB's
            # SQL parser already accepts ``>``, ``<``, ``>=``, ``<=``,
            # ``==`` (folded to ``=``), ``!=``, ``+``, ``-``, ``*``,
            # ``/`` directly.
            sql_expr = _python_predicate_to_sql(_expr_str)
            return plan.with_where(sql_expr)
    else:
        mask_arr = np.asarray(mask, dtype=bool)

        def polars_apply(lf, expanded_cols, _mask=mask_arr):  # type: ignore[misc]
            import polars as pl

            return lf.filter(pl.Series("__mask__", _mask))

        def duckdb_apply(plan, expanded_cols, _mask=mask_arr):  # type: ignore[misc]
            # Boolean-array masks: register a row-aligned mask
            # column on the connection and AND it into the WHERE.
            return plan.with_mask_array(_mask)

    return self._chain(
        _Op(
            name=f"where({mask!r})",
            apply=_apply,
            polars_apply=polars_apply,
            duckdb_apply=duckdb_apply,
        )
    )

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.

Source code in jaxonomy/simulation/lazy_results.py

def with_duckdb_backend(self, connection=None) -> "LazyResults":
    """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.
    """
    if connection is None:
        try:
            import duckdb
        except ImportError as e:
            raise ImportError(
                "LazyResults.with_duckdb_backend: duckdb is not "
                "installed.  Install with `pip install duckdb`."
            ) from e
        connection = duckdb.connect()
    return LazyResults(
        _outputs=self._outputs,
        _time=self._time,
        _ops=list(self._ops),
        _use_duckdb=True,
        _duckdb_conn=connection,
        _per_signal_times=self._per_signal_times,
    )

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.

Source code in jaxonomy/simulation/lazy_results.py

def with_polars_backend(self) -> "LazyResults":
    """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.
    """
    return LazyResults(
        _outputs=self._outputs,
        _time=self._time,
        _ops=list(self._ops),
        _use_polars=True,
        _per_signal_times=self._per_signal_times,
    )

with_signal(name, fn)

Derive a new signal name from existing ones (defers).

fn receives (t, outputs) and returns an array shaped like time.

Source code in jaxonomy/simulation/lazy_results.py

def with_signal(self, name: str, fn: Callable) -> "LazyResults":
    """Derive a new signal ``name`` from existing ones (defers).

    ``fn`` receives ``(t, outputs)`` and returns an array shaped
    like ``time``.
    """

    def _apply(out: dict, t: np.ndarray):
        new_out = dict(out)
        new_out[name] = np.asarray(fn(t, out))
        return new_out, t

    def _polars_apply(lf, expanded_cols, _name=name, _fn=fn):
        # Polars cannot express an arbitrary Python user function
        # natively; collect to compute, then re-lazy.  We still
        # benefit from polars's pushdown on prior ops in the chain
        # (they ran lazily before this point).
        import polars as pl

        df = lf.collect()
        t = df["time"].to_numpy()
        # Reconstruct the outputs dict — vector-valued signals are
        # exploded as ``base__i`` so we re-collapse them.
        out = _collapse_vectorized(df)
        new_col = np.asarray(_fn(t, out))
        return df.with_columns(pl.Series(_name, new_col)).lazy()

    return self._chain(
        _Op(
            name=f"with_signal({name!r})",
            apply=_apply,
            polars_apply=_polars_apply,
        )
    )

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.

Source code in jaxonomy/simulation/provenance.py

class ManifestMismatch(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.
    """

    def __init__(self, differences: list[tuple[str, Any, Any]]):
        self.differences: list[tuple[str, Any, Any]] = list(differences)
        lines = [f"{len(self.differences)} manifest field(s) drifted:"]
        for path, actual, expected in self.differences:
            lines.append(f"  {path}: actual={actual!r} expected={expected!r}")
        super().__init__("\n".join(lines))

ODESolverOptions dataclass

Options for the ODE solver.

See documentation for simulate for details on these options.

Source code in jaxonomy/backend/ode_solver.py

@dataclasses.dataclass
class ODESolverOptions:
    """Options for the ODE solver.

    See documentation for `simulate` for details on these options.
    """

    rtol: float = 1e-3
    atol: float = 1e-6
    min_step_size: float = None
    max_step_size: float = None
    method: str = "auto"  # Dopri5 (jax/scipy) or BDF (jax)
    enable_autodiff: bool = False
    max_checkpoints: int = None  # Only used for checkpointing in autodiff

ProvenanceManifest dataclass

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.

Source code in jaxonomy/simulation/provenance.py

@dataclasses.dataclass(frozen=True)
class 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.
    """

    jaxonomy_version: str
    jax_version: str
    numpy_version: str
    precision_info: dict[str, Any]
    options: dict[str, Any]
    system: dict[str, Any]
    timestamp: str
    git_head: Optional[str] = None
    # T-110-followup-git-revision: richer git metadata.
    git_head_sha: Optional[str] = None
    git_branch: Optional[str] = None
    git_dirty: Optional[bool] = None
    git_head_commit_time: Optional[str] = None
    # T-110-followup-config-hash: deterministic run-identity hash.
    config_hash: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly dict representation of the manifest."""
        return {
            "jaxonomy_version": self.jaxonomy_version,
            "jax_version": self.jax_version,
            "numpy_version": self.numpy_version,
            "precision_info": dict(self.precision_info),
            "options": dict(self.options),
            "system": dict(self.system),
            "timestamp": self.timestamp,
            "git_head": self.git_head,
            "git_head_sha": self.git_head_sha,
            "git_branch": self.git_branch,
            "git_dirty": self.git_dirty,
            "git_head_commit_time": self.git_head_commit_time,
            "config_hash": self.config_hash,
        }

    def to_json(self, *, indent: Optional[int] = None) -> str:
        """Serialise :meth:`to_dict` via ``json.dumps``."""
        return json.dumps(self.to_dict(), indent=indent, sort_keys=True, default=repr)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ProvenanceManifest":
        """Construct a :class:`ProvenanceManifest` from a dict produced by
        :meth:`to_dict` (round-trip helper for serialisation tests)."""
        git_dirty = data.get("git_dirty")
        if git_dirty is not None:
            git_dirty = bool(git_dirty)
        return cls(
            jaxonomy_version=str(data.get("jaxonomy_version", "")),
            jax_version=str(data.get("jax_version", "")),
            numpy_version=str(data.get("numpy_version", "")),
            precision_info=dict(data.get("precision_info", {}) or {}),
            options=dict(data.get("options", {}) or {}),
            system=dict(data.get("system", {}) or {}),
            timestamp=str(data.get("timestamp", "")),
            git_head=data.get("git_head"),
            git_head_sha=data.get("git_head_sha"),
            git_branch=data.get("git_branch"),
            git_dirty=git_dirty,
            git_head_commit_time=data.get("git_head_commit_time"),
            config_hash=str(data.get("config_hash", "") or ""),
        )

from_dict(data) classmethod

Construct a :class:ProvenanceManifest from a dict produced by :meth:to_dict (round-trip helper for serialisation tests).

Source code in jaxonomy/simulation/provenance.py

@classmethod
def from_dict(cls, data: dict[str, Any]) -> "ProvenanceManifest":
    """Construct a :class:`ProvenanceManifest` from a dict produced by
    :meth:`to_dict` (round-trip helper for serialisation tests)."""
    git_dirty = data.get("git_dirty")
    if git_dirty is not None:
        git_dirty = bool(git_dirty)
    return cls(
        jaxonomy_version=str(data.get("jaxonomy_version", "")),
        jax_version=str(data.get("jax_version", "")),
        numpy_version=str(data.get("numpy_version", "")),
        precision_info=dict(data.get("precision_info", {}) or {}),
        options=dict(data.get("options", {}) or {}),
        system=dict(data.get("system", {}) or {}),
        timestamp=str(data.get("timestamp", "")),
        git_head=data.get("git_head"),
        git_head_sha=data.get("git_head_sha"),
        git_branch=data.get("git_branch"),
        git_dirty=git_dirty,
        git_head_commit_time=data.get("git_head_commit_time"),
        config_hash=str(data.get("config_hash", "") or ""),
    )

to_dict()

Return a JSON-friendly dict representation of the manifest.

Source code in jaxonomy/simulation/provenance.py

def to_dict(self) -> dict[str, Any]:
    """Return a JSON-friendly dict representation of the manifest."""
    return {
        "jaxonomy_version": self.jaxonomy_version,
        "jax_version": self.jax_version,
        "numpy_version": self.numpy_version,
        "precision_info": dict(self.precision_info),
        "options": dict(self.options),
        "system": dict(self.system),
        "timestamp": self.timestamp,
        "git_head": self.git_head,
        "git_head_sha": self.git_head_sha,
        "git_branch": self.git_branch,
        "git_dirty": self.git_dirty,
        "git_head_commit_time": self.git_head_commit_time,
        "config_hash": self.config_hash,
    }

to_json(*, indent=None)

Serialise :meth:to_dict via json.dumps.

Source code in jaxonomy/simulation/provenance.py

def to_json(self, *, indent: Optional[int] = None) -> str:
    """Serialise :meth:`to_dict` via ``json.dumps``."""
    return json.dumps(self.to_dict(), indent=indent, sort_keys=True, default=repr)

ResultsWithProvenance dataclass

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.

Source code in jaxonomy/simulation/provenance.py

@dataclasses.dataclass(frozen=True)
class 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.
    """

    results: Any
    provenance: Any

    # ``__getattr__`` is only invoked when normal attribute lookup
    # fails, so ``self.results`` / ``self.provenance`` always resolve
    # against the dataclass fields (no infinite recursion).
    def __getattr__(self, name: str) -> Any:
        # ``object.__getattribute__`` reaches the dataclass slot
        # directly and raises ``AttributeError`` if (somehow) the field
        # isn't yet set — which is the right signal for ``hasattr``.
        results = object.__getattribute__(self, "results")
        try:
            return getattr(results, name)
        except AttributeError as exc:
            # Re-raise with the wrapper's own type in the message so
            # the user sees that the lookup went through us.
            raise AttributeError(
                f"{type(self).__name__!s} has no attribute {name!r} "
                f"(neither does the wrapped {type(results).__name__})"
            ) from exc

    def __repr__(self) -> str:
        # Show both sides explicitly so the wrapper is self-describing
        # even when the underlying results' ``repr`` is verbose.
        return (
            f"ResultsWithProvenance(results={self.results!r}, "
            f"provenance={self.provenance!r})"
        )

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.

Source code in jaxonomy/simulation/errors.py

class SimulationError(JaxonomyError):
    """Raised when a simulator entry point fails at trace or run time.

    Attributes:
        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``.
    """

    def __init__(
        self,
        message: str,
        *,
        cause: BaseException | None = None,
        block: str | None = None,
        port: str | None = None,
    ):
        super().__init__(message)
        self.cause = cause
        self.block = block
        self.port = port

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.

Source code in jaxonomy/simulation/types.py

class SimulationResults(NamedTuple):
    """Data structure for the results of a simulation.

    Attributes:
        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.
    """

    context: ContextBase
    time: Array = None
    outputs: dict[str, Array] = None
    parameters: dict[str, Any] = None
    # T-013: optional per-signal timestamp vectors.  When non-None,
    # ``per_signal_times[name]`` is the native time vector for the
    # corresponding recorded signal ``outputs[name]``.  When None, every
    # signal shares ``self.time`` (backwards-compatible default).
    per_signal_times: Optional[dict[str, Array]] = None

    # T-012a: optional marker that ``query`` should use a higher-order
    # interpolant rather than linear.  Today this field carries the literal
    # string ``"pchip"`` when ``SimulatorOptions.record_solver_states=True``
    # (the PCHIP fallback shipped in T-012a partial); a future T-012a-followup
    # will populate it with the per-major-step solver-state pytree needed for
    # the solver's native dense interpolant.  When ``None`` (default and
    # legacy), ``query`` uses ``jnp.interp`` linear interpolation — fully
    # backwards-compatible.
    solver_states: Optional[Any] = None

    # T-110 Phase 1: optional reproducibility manifest populated by
    # ``simulate(...)`` when ``SimulatorOptions.record_provenance=True``.
    # When ``None`` (default), the path is byte-equivalent to the pre-
    # T-110 behaviour.  See :mod:`jaxonomy.simulation.provenance`.
    provenance: Optional[Any] = None

    # T-113 Phase 1: optional per-major-step DAE constraint drift trace
    # populated by ``simulate(...)`` when
    # ``SimulatorOptions.record_dae_drift=True``.  When non-None, a dict
    # of ``{"time": np.ndarray, "residual": np.ndarray}`` recording the
    # post-projection ``||f_a||_∞`` value at each major step (in step
    # order; chronological).  When ``None`` (default), the trace was
    # not recorded — the path is byte-equivalent to the pre-T-113
    # behaviour.  Pure-ODE systems (no mass matrix) yield ``None``
    # even when the option is True.
    dae_drift_trace: Optional[dict] = None

    # T-125-followup-record-event-times: optional dict of zero-crossing
    # event firing times populated by ``simulate(...)`` when
    # ``SimulatorOptions.record_event_times=True``.  When non-None,
    # ``event_times[i]`` is the 1-D ``np.ndarray`` of firing times for the
    # ``i``-th zero-crossing event (matching the order of
    # ``system.zero_crossing_events.events``).  Events that never fired
    # have an empty array; events that fired multiple times have a
    # monotonically-increasing array.  When ``None`` (default), the
    # capture was not requested — the path is byte-equivalent to the
    # pre-followup behaviour.  Diagrams with no zero-crossing events
    # yield ``None`` even when the option is True.  Designed to feed
    # straight into :func:`jaxonomy.event_time_gradient` so callers do
    # not have to track event times manually.
    event_times: Optional[dict] = None

    def time_for(self, signal: str):
        """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.
        """
        if self.per_signal_times is not None and signal in self.per_signal_times:
            return self.per_signal_times[signal]
        return self.time

    def align(self, 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.

        Args:
            time_vector: 1-D array of times to sample at.
            signals: Optional iterable of signal names to include.
                Defaults to all recorded signals.

        Returns:
            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.
        """
        import jax.numpy as _jnp
        import numpy as _np

        if self.outputs is None:
            raise ValueError(
                "SimulationResults.align: no recorded signals to align."
            )
        signals = list(signals) if signals is not None else list(self.outputs)
        time_vector = _jnp.asarray(time_vector)
        t_q = _np.asarray(time_vector)

        new_outputs = {}
        for name in signals:
            if name not in self.outputs:
                raise ValueError(
                    f"SimulationResults.align: unknown signal {name!r}.  "
                    f"Recorded: {list(self.outputs)}"
                )
            t_src = _np.asarray(self.time_for(name))
            y_src = _np.asarray(self.outputs[name])
            # T-013a: when ``per_signal_times`` is populated by Mode B,
            # the outputs array is full-resolution while ``t_src`` is
            # deduplicated.  Reconstruct the matching value vector by
            # picking the leading-axis indices from ``self.time`` that
            # equal the per-signal timestamps, so interp's xp/fp pair
            # is the same length.
            if (
                self.per_signal_times is not None
                and name in self.per_signal_times
                and t_src.shape[0] != y_src.shape[0]
            ):
                global_t = _np.asarray(self.time)
                # Match each t_src entry to its index in the global
                # vector via searchsorted (both are monotonic).
                idx = _np.searchsorted(global_t, t_src)
                # Clamp in case of floating-point drift.
                idx = _np.clip(idx, 0, global_t.shape[0] - 1)
                y_src = y_src[idx]
            t_min, t_max = float(t_src[0]), float(t_src[-1])
            if _np.any(t_q < t_min - 1e-12) or _np.any(t_q > t_max + 1e-12):
                raise ValueError(
                    f"SimulationResults.align: query times out of range "
                    f"for signal {name!r} (covered [{t_min}, {t_max}])."
                )
            if y_src.ndim == 1:
                new_outputs[name] = _jnp.interp(time_vector, t_src, y_src)
            else:
                new_outputs[name] = _jnp.stack(
                    [_jnp.interp(time_vector, t_src, y_src[:, i])
                     for i in range(y_src.shape[1])],
                    axis=-1,
                )

        return SimulationResults(
            context=self.context,
            time=time_vector,
            outputs=new_outputs,
            parameters=self.parameters,
            per_signal_times=None,
            solver_states=self.solver_states,
            provenance=self.provenance,
            dae_drift_trace=self.dae_drift_trace,
            event_times=self.event_times,
        )

    def query(self, t, signal: Optional[str] = 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.

        Args:
            t: Scalar time, or 1-D array of times.
            signal: Optional signal name.  If provided, return only
                that signal's interpolated value.  If None, return a
                dict of all recorded signals.

        Returns:
            - If ``signal`` is provided: the interpolated array (scalar
              when ``t`` is scalar, 1-D otherwise).
            - Otherwise: ``dict[str, Array]`` matching ``self.outputs``.

        Raises:
            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``.
        """
        import jax.numpy as _jnp
        import numpy as _np

        if self.outputs is None or self.time is None:
            raise ValueError(
                "SimulationResults.query: no recorded signals.  Pass "
                "recorded_signals= to simulate() first."
            )

        t_vec = _np.asarray(self.time)
        t_arr = _np.asarray(t)

        # Bound check — a single violated endpoint fails the whole call.
        t_min, t_max = float(t_vec[0]), float(t_vec[-1])
        if _np.any(t_arr < t_min - 1e-12) or _np.any(t_arr > t_max + 1e-12):
            raise ValueError(
                f"SimulationResults.query: t out of range.  "
                f"Simulation covered [{t_min}, {t_max}]; got {t_arr!r}."
            )

        # T-012a / T-012a-followup: select interpolant.
        #   ``solver_states is None`` → linear (legacy + load-from-disk).
        #   ``"pchip"`` sentinel → PCHIP cubic-Hermite fallback.
        #   ``NativeInterpolant`` → native solver polynomial (sub-ULP).
        native_interp = (
            self.solver_states
            if isinstance(self.solver_states, NativeInterpolant)
            else None
        )
        use_pchip = self.solver_states == "pchip" and t_vec.shape[0] >= 2
        if native_interp is not None and t_vec.shape[0] >= 2:
            # PCHIP is the per-signal fallback when the native polynomial
            # doesn't match the recorded signal (e.g. a discrete output,
            # not a state passthrough).
            use_pchip = True

        def _is_zoh(col: "_np.ndarray") -> bool:
            """Detect zero-order-hold-style signals: long constant runs.

            PCHIP is shape-preserving but a discrete signal that holds
            a value across many samples and then steps is best served
            by step interpolation — PCHIP would still smooth the corner
            slightly.  Heuristic: if more than half the consecutive
            differences are exactly zero, treat as ZOH.
            """
            if col.shape[0] < 3:
                return False
            d = _np.diff(col)
            return _np.count_nonzero(d == 0) > col.shape[0] / 2

        def _native_eval(col: "_np.ndarray"):
            """T-012a-followup: evaluate the solver's polynomial at t_arr.

            Returns ``(values,)`` matching ``t_arr`` shape if the column
            is a continuous-state passthrough — values match the
            polynomial at every recorded segment endpoint to within
            float64 round-off.  Returns ``None`` otherwise so the caller
            falls back to PCHIP/linear.
            """
            ni = native_interp
            t_prev = _np.asarray(ni.t_prev)
            t_step = _np.asarray(ni.t_step)
            coeffs = _np.asarray(ni.interp_coeff)
            n_seg = t_prev.shape[0]
            if n_seg == 0:
                return None
            n_y = coeffs.shape[2]
            # End-point values per segment via polyval at theta=1.
            end_vals = _np.empty((n_seg, n_y), dtype=coeffs.dtype)
            for i in range(n_seg):
                end_vals[i] = _np.polyval(coeffs[i], 1.0)
            # Match each segment's t_step to the col index.
            seg_end_idx = _np.searchsorted(t_vec, t_step)
            seg_end_idx = _np.clip(seg_end_idx, 0, t_vec.shape[0] - 1)
            # Pick the state-component whose polynomial endpoint best
            # matches the recorded col over all segments.  If no
            # component agrees within 1e-6, the col isn't a state
            # passthrough — abort.
            recorded = col[seg_end_idx]
            best_comp = -1
            best_err = _np.inf
            for c in range(n_y):
                err = _np.max(_np.abs(end_vals[:, c] - recorded))
                if err < best_err:
                    best_err = err
                    best_comp = c
            if best_err > 1e-6 or best_comp < 0:
                return None
            # Locate each query time in the segments.  ``side="left"``
            # plus clip lands t == t_step[i] in segment i (good — the
            # endpoint is the polynomial's right edge).
            t_arr_1d = _np.atleast_1d(t_arr).astype(_np.float64)
            seg_idx = _np.searchsorted(t_step, t_arr_1d, side="left")
            seg_idx = _np.clip(seg_idx, 0, n_seg - 1)
            tp = t_prev[seg_idx]
            ts = t_step[seg_idx]
            dt = ts - tp
            dt = _np.where(dt == 0.0, 1.0, dt)
            theta = (t_arr_1d - tp) / dt
            # Vectorised Horner over the picked component.
            picked = coeffs[seg_idx, :, best_comp]  # (n_q, n_coeff)
            n_coeff = picked.shape[-1]
            result = _np.zeros_like(theta)
            for k in range(n_coeff):
                result = result * theta + picked[..., k]
            # Snap exact-recorded-time queries to recorded values to
            # remove residual round-off (the polynomial is a near-exact
            # interpolant but not bit-exact at the endpoints).
            t_match_idx = _np.searchsorted(t_vec, t_arr_1d)
            t_match_idx = _np.clip(t_match_idx, 0, t_vec.shape[0] - 1)
            on_recorded = _np.isclose(
                t_vec[t_match_idx], t_arr_1d, atol=1e-15, rtol=0.0,
            )
            result = _np.where(on_recorded, col[t_match_idx], result)
            return result.reshape(t_arr.shape) if t_arr.ndim > 0 else result[0]

        def _interp_column(col: "_np.ndarray") -> "_np.ndarray":
            # T-012a-followup: try the native polynomial first.
            if native_interp is not None and not _is_zoh(col):
                native_result = _native_eval(col)
                if native_result is not None:
                    return _np.asarray(native_result)
            if use_pchip and not _is_zoh(col):
                # PCHIP requires strictly-increasing x.  Recorded times
                # are monotonic-non-decreasing (zero-crossing handler
                # may inject a sample at the same instant); collapse
                # any duplicates by keeping the first.
                _, uniq_idx = _np.unique(t_vec, return_index=True)
                uniq_idx = _np.sort(uniq_idx)
                if uniq_idx.shape[0] >= 2:
                    from scipy.interpolate import PchipInterpolator
                    interp = PchipInterpolator(
                        t_vec[uniq_idx], col[uniq_idx], extrapolate=False,
                    )
                    return _np.asarray(interp(t_arr))
            # Linear fallback (legacy and ZOH path).
            return _np.asarray(_jnp.interp(t_arr, t_vec, col))

        def _interp_one(arr):
            arr = _np.asarray(arr)
            if arr.ndim == 1:
                return _jnp.asarray(_interp_column(arr))
            # vector-valued signal: interp each component
            return _jnp.stack(
                [_jnp.asarray(_interp_column(arr[:, i]))
                 for i in range(arr.shape[1])],
                axis=-1,
            )

        if signal is not None:
            if signal not in self.outputs:
                raise ValueError(
                    f"SimulationResults.query: unknown signal {signal!r}.  "
                    f"Recorded: {list(self.outputs)}"
                )
            return _interp_one(self.outputs[signal])
        return {name: _interp_one(arr) for name, arr in self.outputs.items()}

    def lazy(self):
        """Return a :class:`LazyResults` wrapper for fluent / deferred queries.

        See :mod:`jaxonomy.simulation.lazy_results` for the full API.
        """
        from .lazy_results import LazyResults
        return LazyResults.from_results(self)

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.

Source code in jaxonomy/simulation/types.py

def align(self, 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.

    Args:
        time_vector: 1-D array of times to sample at.
        signals: Optional iterable of signal names to include.
            Defaults to all recorded signals.

    Returns:
        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.
    """
    import jax.numpy as _jnp
    import numpy as _np

    if self.outputs is None:
        raise ValueError(
            "SimulationResults.align: no recorded signals to align."
        )
    signals = list(signals) if signals is not None else list(self.outputs)
    time_vector = _jnp.asarray(time_vector)
    t_q = _np.asarray(time_vector)

    new_outputs = {}
    for name in signals:
        if name not in self.outputs:
            raise ValueError(
                f"SimulationResults.align: unknown signal {name!r}.  "
                f"Recorded: {list(self.outputs)}"
            )
        t_src = _np.asarray(self.time_for(name))
        y_src = _np.asarray(self.outputs[name])
        # T-013a: when ``per_signal_times`` is populated by Mode B,
        # the outputs array is full-resolution while ``t_src`` is
        # deduplicated.  Reconstruct the matching value vector by
        # picking the leading-axis indices from ``self.time`` that
        # equal the per-signal timestamps, so interp's xp/fp pair
        # is the same length.
        if (
            self.per_signal_times is not None
            and name in self.per_signal_times
            and t_src.shape[0] != y_src.shape[0]
        ):
            global_t = _np.asarray(self.time)
            # Match each t_src entry to its index in the global
            # vector via searchsorted (both are monotonic).
            idx = _np.searchsorted(global_t, t_src)
            # Clamp in case of floating-point drift.
            idx = _np.clip(idx, 0, global_t.shape[0] - 1)
            y_src = y_src[idx]
        t_min, t_max = float(t_src[0]), float(t_src[-1])
        if _np.any(t_q < t_min - 1e-12) or _np.any(t_q > t_max + 1e-12):
            raise ValueError(
                f"SimulationResults.align: query times out of range "
                f"for signal {name!r} (covered [{t_min}, {t_max}])."
            )
        if y_src.ndim == 1:
            new_outputs[name] = _jnp.interp(time_vector, t_src, y_src)
        else:
            new_outputs[name] = _jnp.stack(
                [_jnp.interp(time_vector, t_src, y_src[:, i])
                 for i in range(y_src.shape[1])],
                axis=-1,
            )

    return SimulationResults(
        context=self.context,
        time=time_vector,
        outputs=new_outputs,
        parameters=self.parameters,
        per_signal_times=None,
        solver_states=self.solver_states,
        provenance=self.provenance,
        dae_drift_trace=self.dae_drift_trace,
        event_times=self.event_times,
    )

lazy()

Return a :class:LazyResults wrapper for fluent / deferred queries.

See :mod:jaxonomy.simulation.lazy_results for the full API.

Source code in jaxonomy/simulation/types.py

def lazy(self):
    """Return a :class:`LazyResults` wrapper for fluent / deferred queries.

    See :mod:`jaxonomy.simulation.lazy_results` for the full API.
    """
    from .lazy_results import LazyResults
    return LazyResults.from_results(self)

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.

Source code in jaxonomy/simulation/types.py

def query(self, t, signal: Optional[str] = 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.

    Args:
        t: Scalar time, or 1-D array of times.
        signal: Optional signal name.  If provided, return only
            that signal's interpolated value.  If None, return a
            dict of all recorded signals.

    Returns:
        - If ``signal`` is provided: the interpolated array (scalar
          when ``t`` is scalar, 1-D otherwise).
        - Otherwise: ``dict[str, Array]`` matching ``self.outputs``.

    Raises:
        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``.
    """
    import jax.numpy as _jnp
    import numpy as _np

    if self.outputs is None or self.time is None:
        raise ValueError(
            "SimulationResults.query: no recorded signals.  Pass "
            "recorded_signals= to simulate() first."
        )

    t_vec = _np.asarray(self.time)
    t_arr = _np.asarray(t)

    # Bound check — a single violated endpoint fails the whole call.
    t_min, t_max = float(t_vec[0]), float(t_vec[-1])
    if _np.any(t_arr < t_min - 1e-12) or _np.any(t_arr > t_max + 1e-12):
        raise ValueError(
            f"SimulationResults.query: t out of range.  "
            f"Simulation covered [{t_min}, {t_max}]; got {t_arr!r}."
        )

    # T-012a / T-012a-followup: select interpolant.
    #   ``solver_states is None`` → linear (legacy + load-from-disk).
    #   ``"pchip"`` sentinel → PCHIP cubic-Hermite fallback.
    #   ``NativeInterpolant`` → native solver polynomial (sub-ULP).
    native_interp = (
        self.solver_states
        if isinstance(self.solver_states, NativeInterpolant)
        else None
    )
    use_pchip = self.solver_states == "pchip" and t_vec.shape[0] >= 2
    if native_interp is not None and t_vec.shape[0] >= 2:
        # PCHIP is the per-signal fallback when the native polynomial
        # doesn't match the recorded signal (e.g. a discrete output,
        # not a state passthrough).
        use_pchip = True

    def _is_zoh(col: "_np.ndarray") -> bool:
        """Detect zero-order-hold-style signals: long constant runs.

        PCHIP is shape-preserving but a discrete signal that holds
        a value across many samples and then steps is best served
        by step interpolation — PCHIP would still smooth the corner
        slightly.  Heuristic: if more than half the consecutive
        differences are exactly zero, treat as ZOH.
        """
        if col.shape[0] < 3:
            return False
        d = _np.diff(col)
        return _np.count_nonzero(d == 0) > col.shape[0] / 2

    def _native_eval(col: "_np.ndarray"):
        """T-012a-followup: evaluate the solver's polynomial at t_arr.

        Returns ``(values,)`` matching ``t_arr`` shape if the column
        is a continuous-state passthrough — values match the
        polynomial at every recorded segment endpoint to within
        float64 round-off.  Returns ``None`` otherwise so the caller
        falls back to PCHIP/linear.
        """
        ni = native_interp
        t_prev = _np.asarray(ni.t_prev)
        t_step = _np.asarray(ni.t_step)
        coeffs = _np.asarray(ni.interp_coeff)
        n_seg = t_prev.shape[0]
        if n_seg == 0:
            return None
        n_y = coeffs.shape[2]
        # End-point values per segment via polyval at theta=1.
        end_vals = _np.empty((n_seg, n_y), dtype=coeffs.dtype)
        for i in range(n_seg):
            end_vals[i] = _np.polyval(coeffs[i], 1.0)
        # Match each segment's t_step to the col index.
        seg_end_idx = _np.searchsorted(t_vec, t_step)
        seg_end_idx = _np.clip(seg_end_idx, 0, t_vec.shape[0] - 1)
        # Pick the state-component whose polynomial endpoint best
        # matches the recorded col over all segments.  If no
        # component agrees within 1e-6, the col isn't a state
        # passthrough — abort.
        recorded = col[seg_end_idx]
        best_comp = -1
        best_err = _np.inf
        for c in range(n_y):
            err = _np.max(_np.abs(end_vals[:, c] - recorded))
            if err < best_err:
                best_err = err
                best_comp = c
        if best_err > 1e-6 or best_comp < 0:
            return None
        # Locate each query time in the segments.  ``side="left"``
        # plus clip lands t == t_step[i] in segment i (good — the
        # endpoint is the polynomial's right edge).
        t_arr_1d = _np.atleast_1d(t_arr).astype(_np.float64)
        seg_idx = _np.searchsorted(t_step, t_arr_1d, side="left")
        seg_idx = _np.clip(seg_idx, 0, n_seg - 1)
        tp = t_prev[seg_idx]
        ts = t_step[seg_idx]
        dt = ts - tp
        dt = _np.where(dt == 0.0, 1.0, dt)
        theta = (t_arr_1d - tp) / dt
        # Vectorised Horner over the picked component.
        picked = coeffs[seg_idx, :, best_comp]  # (n_q, n_coeff)
        n_coeff = picked.shape[-1]
        result = _np.zeros_like(theta)
        for k in range(n_coeff):
            result = result * theta + picked[..., k]
        # Snap exact-recorded-time queries to recorded values to
        # remove residual round-off (the polynomial is a near-exact
        # interpolant but not bit-exact at the endpoints).
        t_match_idx = _np.searchsorted(t_vec, t_arr_1d)
        t_match_idx = _np.clip(t_match_idx, 0, t_vec.shape[0] - 1)
        on_recorded = _np.isclose(
            t_vec[t_match_idx], t_arr_1d, atol=1e-15, rtol=0.0,
        )
        result = _np.where(on_recorded, col[t_match_idx], result)
        return result.reshape(t_arr.shape) if t_arr.ndim > 0 else result[0]

    def _interp_column(col: "_np.ndarray") -> "_np.ndarray":
        # T-012a-followup: try the native polynomial first.
        if native_interp is not None and not _is_zoh(col):
            native_result = _native_eval(col)
            if native_result is not None:
                return _np.asarray(native_result)
        if use_pchip and not _is_zoh(col):
            # PCHIP requires strictly-increasing x.  Recorded times
            # are monotonic-non-decreasing (zero-crossing handler
            # may inject a sample at the same instant); collapse
            # any duplicates by keeping the first.
            _, uniq_idx = _np.unique(t_vec, return_index=True)
            uniq_idx = _np.sort(uniq_idx)
            if uniq_idx.shape[0] >= 2:
                from scipy.interpolate import PchipInterpolator
                interp = PchipInterpolator(
                    t_vec[uniq_idx], col[uniq_idx], extrapolate=False,
                )
                return _np.asarray(interp(t_arr))
        # Linear fallback (legacy and ZOH path).
        return _np.asarray(_jnp.interp(t_arr, t_vec, col))

    def _interp_one(arr):
        arr = _np.asarray(arr)
        if arr.ndim == 1:
            return _jnp.asarray(_interp_column(arr))
        # vector-valued signal: interp each component
        return _jnp.stack(
            [_jnp.asarray(_interp_column(arr[:, i]))
             for i in range(arr.shape[1])],
            axis=-1,
        )

    if signal is not None:
        if signal not in self.outputs:
            raise ValueError(
                f"SimulationResults.query: unknown signal {signal!r}.  "
                f"Recorded: {list(self.outputs)}"
            )
        return _interp_one(self.outputs[signal])
    return {name: _interp_one(arr) for name, arr in self.outputs.items()}

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.

Source code in jaxonomy/simulation/types.py

def time_for(self, signal: str):
    """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.
    """
    if self.per_signal_times is not None and signal in self.per_signal_times:
        return self.per_signal_times[signal]
    return self.time

Simulator

Class for orchestrating simulations of hybrid dynamical systems.

See the simulate function for more details.

Source code in jaxonomy/simulation/simulator.py

class Simulator:
    """Class for orchestrating simulations of hybrid dynamical systems.

    See the `simulate` function for more details.
    """

    def __init__(
        self,
        system: SystemBase,
        ode_solver: ODESolverBase = None,
        options: SimulatorOptions = None,
    ):
        """Initialize the simulator.

        Args:
            system (SystemBase): The hybrid dynamical system to simulate.
            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.
            options (SimulatorOptions):
                Options for the simulation process.  See `simulate` for details.
        """
        self.system = system

        if options is None:
            options = SimulatorOptions()

        # Determine whether JAX tracing can be used (jit, grad, vmap, etc)
        math_backend, self.enable_tracing = _check_backend(options)

        # Set the math backend
        set_backend(math_backend)

        # Should the simulation be run with autodiff enabled?  This will override
        # the `advance_to` method with a custom autodiff rule.
        self.enable_autodiff = options.enable_autodiff

        if ode_solver is None:
            ode_solver = ODESolver(system, options=options.ode_options)

        # Store configuration options
        self.max_major_steps = options.max_major_steps
        # Honor max_major_steps as a hard cap whenever it was explicitly provided
        # (either via _check_simulate_options or directly by the caller).
        self._explicit_max_major_steps = options._explicit_max_major_steps or (
            options.max_major_steps is not None and options.max_major_steps > 0
        )
        self.max_major_step_length = options.max_major_step_length
        self.zc_bisection_loop_count = options.zc_bisection_loop_count
        self.major_step_callback = options.major_step_callback

        # T-003a: opt-in DAE constraint projection at the end of each major step.
        self.dae_projection_enabled = getattr(
            options, "dae_projection_enabled", False,
        )
        self.dae_projection_tol = getattr(options, "dae_projection_tol", 1e-8)
        self.dae_projection_max_iter = getattr(
            options, "dae_projection_max_iter", 3,
        )

        # T-113-followup-event-reprojection: opt-in projection immediately
        # after discrete-event resets *within* a major step (top-of-step
        # ``_handle_discrete_update`` and triggered ZC resets inside
        # ``_advance_continuous_time``).  Default ``False`` so the hot
        # path is byte-equivalent to the pre-followup code.  Reuses the
        # T-003a tolerance / max-iter knobs to avoid surface-area churn.
        self.dae_reproject_after_events = getattr(
            options, "dae_reproject_after_events", False,
        )

        # T-003b: opt-in DAE drift threshold (post-projection check).
        # ``None`` (default) disables the check entirely — no overhead.
        self.dae_drift_threshold = getattr(
            options, "dae_drift_threshold", None,
        )

        # T-132: declared per-block state projections
        # (``declare_continuous_state(project=...)``, e.g. unit-quaternion
        # renormalization).  Collected once, statically — when empty (the
        # default), ``_major_step`` skips the block at trace time and the
        # hot path is byte-equivalent.
        _leaves = getattr(system, "leaf_systems", None)
        if _leaves is None:
            _leaves = [system]
        self._state_projection_leaves = [
            s
            for s in _leaves
            if getattr(s, "_continuous_projection", None) is not None
        ]

        # T-113 Phase 1: opt-in per-major-step DAE drift trace.
        # ``False`` (default) disables the trace entirely — no monitor
        # constructed and the simulator's ``_major_step`` skips the
        # trace block at trace time.  When True AND the system has a
        # mass matrix, attach a host-side ``_DAEDriftMonitor`` so the
        # trace block forwards each major step's ``(time, residual)``
        # via ``jax.debug.callback``.  Non-DAE systems get no monitor
        # (the diagnostic is mass-matrix-specific by definition).
        self.record_dae_drift = getattr(options, "record_dae_drift", False)
        self._dae_drift_monitor: _DAEDriftMonitor | None = None
        if self.record_dae_drift and getattr(
            self.system, "has_mass_matrix", False,
        ):
            self._dae_drift_monitor = _DAEDriftMonitor()

        # T-125-followup-record-event-times: opt-in capture of zero-
        # crossing event firing times.  Construct a host-side recorder
        # only when the option is True AND the diagram has at least one
        # zero-crossing event — diagrams with no events get no recorder
        # so ``_advance_continuous_time`` short-circuits the callback at
        # trace time, preserving the byte-equivalent default-off path.
        # ``n_zero_crossing_events`` is set further down in __init__ but
        # the snapshot below makes the count available without re-
        # importing the system; we settle for re-counting cheaply here
        # to avoid reordering the existing init blocks.
        self.record_event_times = getattr(options, "record_event_times", False)
        self._event_time_recorder: _EventTimeRecorder | None = None
        if self.record_event_times:
            n_zc = len(system.zero_crossing_events.events)
            if n_zc > 0:
                self._event_time_recorder = _EventTimeRecorder(n_zc)

        # T-038a-followup-bdf-condition-check: opt-in BDF Newton
        # condition-number diagnostic.  When the threshold is set AND
        # the active solver is a BDF solver, attach a
        # ``_BDFConditionMonitor`` to it as a side-channel — the BDF
        # solver checks ``getattr(self, "_cond_monitor", None)`` inside
        # ``newton_iteration`` and forwards the cond estimate via
        # ``jax.debug.callback`` only when set.  Default-off path is
        # byte-equivalent (no monitor → no-op in BDF).  Non-BDF
        # solvers silently ignore the option (the diagnostic is BDF-
        # specific by definition).
        self.bdf_condition_warning_threshold = getattr(
            options, "bdf_condition_warning_threshold", None,
        )
        self._bdf_cond_monitor: _BDFConditionMonitor | None = None
        if self.bdf_condition_warning_threshold is not None:
            # Only attach to BDF — non-BDF solvers don't have a Newton
            # iteration to monitor, and we don't want to silently
            # mislead users who set the option on a non-BDF run.
            try:
                from ..backend._jax.bdf import BDFSolver as _BDFSolver
            except Exception:  # pragma: no cover — import-time guard only
                _BDFSolver = None
            if _BDFSolver is not None and isinstance(ode_solver, _BDFSolver):
                self._bdf_cond_monitor = _BDFConditionMonitor(
                    self.bdf_condition_warning_threshold,
                )
                # Attach as an instance attribute so the BDF solver
                # picks it up via ``getattr(self, "_cond_monitor", None)``
                # without needing a constructor change.
                ode_solver._cond_monitor = self._bdf_cond_monitor

        # Detailed non-finite abort diagnostics: stamp the opt-in flag on
        # the BDF solver (checked at trace time; default path compiles no
        # callback ops).
        if getattr(options, "bdf_nonfinite_diagnostics", False):
            try:
                from ..backend._jax.bdf import BDFSolver as _BDFSolver2
            except Exception:  # pragma: no cover — import-time guard only
                _BDFSolver2 = None
            if _BDFSolver2 is not None and isinstance(ode_solver, _BDFSolver2):
                ode_solver._nonfinite_diagnostics = True

        # T-027a-followup: simulator-level Zeno protection options.  All
        # default-off — the recovery probe and the latch are skipped at
        # Python level when ``zeno_protection_enabled=False``, keeping the
        # default hot path byte-equivalent.
        self.zeno_protection_enabled = getattr(
            options, "zeno_protection_enabled", False,
        )
        self.zeno_tolerance = getattr(options, "zeno_tolerance", 1e-6)
        self.zeno_recovery_period = getattr(
            options, "zeno_recovery_period", 10,
        )
        # T-027a-followup-vector-tprev: count zero-crossing events at
        # construction time so ``initialize`` can allocate per-event
        # ``zeno_tprev`` / ``zeno_active`` vectors of the correct shape.
        # ``event_system_ids`` records each event's owning leaf — kept
        # for the per-leaf freeze gate (T-027a-followup-per-leaf-freeze).
        # Both are static (system topology doesn't change at runtime),
        # so this is a one-shot pass at __init__.
        zc_events_static = system.zero_crossing_events.events
        self.n_zero_crossing_events = len(zc_events_static)
        self.event_system_ids = tuple(
            getattr(ev, "system_id", None) for ev in zc_events_static
        )
        # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        # snapshot each zero-crossing event's static ``direction`` string
        # so the per-event recovery-probe nudge in ``_apply_recovery_w0_nudge``
        # can pick the right side of the threshold to push ``w0`` to.
        # Mirrors ``event_system_ids`` — same order as
        # ``system.zero_crossing_events.events`` and the per-event Zeno
        # carry vectors built in ``initialize``.
        self.event_directions = tuple(
            getattr(ev, "direction", "crosses_zero") for ev in zc_events_static
        )

        # T-027a-followup-per-leaf-freeze: build a static map from each
        # leaf's ``system_id`` to its index in
        # ``DiagramContext.continuous_state`` (which is a list ordered by
        # the iteration of subcontexts that have continuous state).  This
        # lets ``_major_step`` know which slot of the list to roll back
        # when an event owned by that leaf has its Zeno latch engaged.
        # For single-LeafSystem simulations the list collapses to one
        # entry (``LeafContext.continuous_state`` is a single Array, not
        # a list — handled separately at the freeze site).
        if isinstance(system, Diagram):
            _leaves = list(system.leaf_systems)
        else:
            _leaves = [system]
        self._sysid_to_cs_idx = {}
        _cs_idx = 0
        for _leaf in _leaves:
            if getattr(_leaf, "has_continuous_state", False):
                self._sysid_to_cs_idx[_leaf.system_id] = _cs_idx
                _cs_idx += 1
        self._n_continuous_leaves = _cs_idx
        # Reverse map: cs_index -> tuple of event-vector positions whose
        # ``zeno_active[i]`` should freeze that leaf.  Built once so the
        # per-step freeze logic is a flat scatter, no per-step Python
        # iteration over event_system_ids.
        self._cs_idx_to_event_positions: dict[int, tuple[int, ...]] = {}
        for i, sid in enumerate(self.event_system_ids):
            cs_i = self._sysid_to_cs_idx.get(sid)
            if cs_i is None:
                continue
            self._cs_idx_to_event_positions.setdefault(cs_i, []).append(i)
        self._cs_idx_to_event_positions = {
            k: tuple(v) for k, v in self._cs_idx_to_event_positions.items()
        }

        # T-027a-followup-per-leaf-solver-state: compute each continuous
        # leaf's flat slice into the raveled ODE state vector ``y`` so the
        # per-leaf freeze gate can decompose ``Dopri5State.{y, f,
        # interp_coeff}`` (and ``BDFState.{y, f, D}``) along the last
        # axis.  The flat layout matches ``ravel_pytree(context.
        # continuous_state)`` exactly: leaves with continuous state are
        # iterated in ``DiagramContext.continuous_subcontexts`` order
        # (subcontexts.values() filtered by has_continuous_state), which
        # is the same order as ``Diagram.leaf_systems`` filtered by
        # ``has_continuous_state``.  Each leaf's flat size is the sum of
        # its ``_default_continuous_state`` pytree-leaf sizes — usually
        # a single Array, but pytree-valued continuous states are also
        # handled.  ``_leaf_flat_slices`` is a tuple of ``(start, end)``
        # int pairs ordered by ``cs_idx``; total length is the flat
        # ODE state dimension ``_n_y_total``.  Default-off path
        # (``zeno_protection_enabled=False``) never reads these; they
        # are purely metadata.
        self._leaf_flat_slices: tuple[tuple[int, int], ...] = ()
        self._n_y_total: int = 0
        if self._n_continuous_leaves > 0:
            _slices: list[tuple[int, int]] = []
            _offset = 0
            # Re-scan ``_leaves`` in the same order used for ``_sysid_to_cs_idx``.
            for _leaf in _leaves:
                if not getattr(_leaf, "has_continuous_state", False):
                    continue
                _xc0 = getattr(_leaf, "_default_continuous_state", None)
                if _xc0 is None:
                    _size = 0
                else:
                    _size = int(sum(
                        int(np.prod(np.shape(_l))) if np.shape(_l) else 1
                        for _l in jax.tree_util.tree_leaves(_xc0)
                    ))
                _slices.append((_offset, _offset + _size))
                _offset += _size
            self._leaf_flat_slices = tuple(_slices)
            self._n_y_total = _offset

        # T-013a-followup-mode-a-buffers: when the user opts into the
        # "buffers" mode, classify each recorded signal's cadence
        # statically here and pass the result through to the recorder.
        # The classification is reused at the per-step decision in
        # ``JaxResultsData.update`` to skip writes for unfired periodic
        # signals.  Default ``"auto"`` does NOT enable buffers — it
        # remains the post-finalize schedule trim path.
        psts_mode = getattr(options, "per_signal_timestamps_mode", "auto")
        psts_enabled = getattr(options, "per_signal_timestamps", False)
        per_signal_buffers_classifications = None
        if (
            psts_enabled
            and psts_mode == "buffers"
            and options.recorded_signals is not None
        ):
            per_signal_buffers_classifications = (
                ResultsRecorder.classify_signal_cadence(options.recorded_signals)
            )

        # T-012a-followup: thread record_solver_states through to the
        # recorder so the JaxResultsData allocates a per-step interpolant
        # ring and ``save`` snapshots ``Dopri5State.interp_coeff`` per
        # call.  Default-off path is byte-equivalent.
        self.record_solver_states = getattr(
            options, "record_solver_states", False,
        )
        # T-002b-followup-buffer-overflow-auto-size — when the user
        # constructs ``Simulator`` directly (bypassing ``simulate``), the
        # ``_check_options`` auto-sizing path is skipped, so ``options.
        # buffer_length`` may still be ``None``. Fall back to
        # ``max_major_steps`` (the natural cap) or a legacy 1000-sample
        # default when neither is available.
        if options.buffer_length is not None:
            recorder_buffer_length = options.buffer_length
        elif self.max_major_steps is not None and self.max_major_steps > 0:
            recorder_buffer_length = max(
                int(self.max_major_steps), _MIN_AUTO_BUFFER_LENGTH
            )
        else:
            recorder_buffer_length = _MIN_AUTO_BUFFER_LENGTH
        self.results_recorder = ResultsRecorder(
            save_time_series=options.save_time_series,
            recorded_outputs=options.recorded_signals,
            buffer_length=recorder_buffer_length,
            per_signal_buffers_classifications=per_signal_buffers_classifications,
            record_solver_states=self.record_solver_states,
        )

        # Zero-crossing handler encapsulates guard evaluation and bisection logic
        self.zc_handler = ZeroCrossingHandler(
            system,
            self.zc_bisection_loop_count,
            lower_triangular_discrete_update=getattr(
                options, "lower_triangular_discrete_update", False,
            ),
        )

        if self.max_major_step_length is None:
            self.max_major_step_length = np.inf

        logger.debug("Simulator created with enable_tracing=%s", self.enable_tracing)

        self.ode_solver = ode_solver

        # T-113-followup-baumgarte-and-ssp: opt-in Baumgarte stabilization.
        # When ``baumgarte_alpha`` and/or ``baumgarte_beta`` are set, wrap
        # the solver's ``ode_rhs`` to add ``-2α·ġ - β²·g`` to the
        # algebraic rows of the rhs.  ``baumgarte_augment_ode_rhs`` is a
        # no-op (returns the input rhs unchanged) when both gains are
        # ``None`` or when the system has no algebraic constraints — the
        # disabled hot path's JIT trace graph is byte-equivalent to the
        # pre-followup behaviour.  Wraps before any ``ode_solver.initialize``
        # call so ``flat_ode_rhs = ravel_first_arg(self.ode_rhs, ...)`` in
        # the JAX impl picks up the augmented version.
        b_alpha = getattr(options, "baumgarte_alpha", None)
        b_beta = getattr(options, "baumgarte_beta", None)
        if (b_alpha is not None or b_beta is not None) and getattr(
            self.system, "has_mass_matrix", False,
        ):
            from .dae_projection import baumgarte_augment_ode_rhs
            ode_solver.ode_rhs = baumgarte_augment_ode_rhs(
                ode_solver.ode_rhs, self.system, b_alpha, b_beta,
            )

        from .autodiff_rules import make_advance_to_vjp, make_guarded_integrate_vjp
        # Modify the default autodiff rule slightly to correctly capture variations
        # in end time of the simulation interval.
        self.has_terminal_events = system.zero_crossing_events.has_terminal_events
        # T-006: wrap advance_to so direct callers (not going through
        # simulate()) also get JAX-error remapping with block/port context.
        # T-A2-followup-advance-to-jit-cache: jit the inner advance_to so a
        # *persistent* Simulator (construct once, call ``advance_to`` many
        # times — interactive stepping, MPC inner loops) reuses the compiled
        # kernel instead of re-tracing op-by-op on every call. The jit is a
        # stable instance attribute, so JAX's cache hits across calls with the
        # same context aval. Only the non-autodiff path is wrapped: the
        # autodiff path returns a ``custom_vjp`` callable that ``simulate``
        # already jits at the outer ``_wrapped_simulate`` boundary, and we
        # keep ``remap_simulation_errors`` on the *outside* so runtime errors
        # are still remapped at the call boundary (not just at trace time).
        _advance_to_impl = make_advance_to_vjp(self)
        if self.enable_tracing and not self.enable_autodiff:
            _advance_to_impl = jax.jit(_advance_to_impl)
        self.advance_to = remap_simulation_errors(_advance_to_impl)

        # Also override the guarded ODE integration with a custom autodiff rule
        # to capture variations due to zero-crossing time.
        self.guarded_integrate = make_guarded_integrate_vjp(self)

    def compile(self, tf: float, context: ContextBase):
        """Warm up / pre-compile the simulation advance_to method on the device."""
        if self.enable_tracing and not self.enable_autodiff:
            self.advance_to(tf, context)

    def while_loop(self, 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.
        """
        use_bounded = self.enable_autodiff or self._explicit_max_major_steps
        if use_bounded:
            return _bounded_while_loop(cond_fun, body_fun, val, self.max_major_steps)
        else:
            return backend.while_loop(cond_fun, body_fun, val)

    def initialize(self, context: ContextBase) -> SimulatorState:
        """Perform initial setup for the simulation."""
        logger.debug("Initializing simulator")
        # context.state.pprint(logger.debug)

        # Initial simulation time as integer (picoseconds)
        initial_int_time = IntegerTime.from_decimal(context.time)

        # Ensure that _next_update_time() can return the current time by perturbing
        # current time as slightly toward negative infinity as possible
        time_of_next_timed_event, timed_events = _next_update_time(
            self.system.periodic_events, initial_int_time - 1
        )

        # timed_events is now marked with the active events at the next update time
        logger.debug("Time of next timed event (int): %s", time_of_next_timed_event)
        logger.debug(
            "Time of next event (sec): %s",
            IntegerTime.as_decimal(time_of_next_timed_event),
        )
        timed_events.pprint(logger.debug)

        end_reason = npa.where(
            time_of_next_timed_event == initial_int_time,
            StepEndReason.TimeTriggered,
            StepEndReason.NothingTriggered,
        )

        # Initialize the results data that will hold recorded time series data.
        results_data = self.results_recorder.initialize(context)

        # T-027a-followup-vector-tprev: when simulator-level Zeno
        # protection is enabled, allocate per-event ``zeno_tprev`` /
        # ``zeno_active`` vectors so each event tracks its own last-
        # firing time independently.  ``zeno_tprev`` is initialised to
        # ``-inf`` so the first firing is never inside tolerance.  When
        # disabled, leave the carry as the scalar defaults from the
        # ``SimulatorState`` declaration so the default-off path's
        # pytree is byte-equivalent.
        if self.zeno_protection_enabled:
            n = max(self.n_zero_crossing_events, 1)
            zeno_tprev = jnp.full((n,), -jnp.inf)
            zeno_active = jnp.zeros((n,), dtype=jnp.bool_)
            # T-027a-followup-per-event-recovery: ``zeno_frozen_steps``
            # vectorises to ``(N_events,)`` so each event independently
            # counts its own consecutive-frozen-step streak.  When event
            # ``i`` hits ``zeno_recovery_period``, only its own latch
            # clears; other events keep cascading.  Default-off path
            # keeps the scalar default in ``SimulatorState`` so the
            # disabled pytree shape is byte-equivalent.
            zeno_frozen_steps = jnp.zeros((n,), dtype=jnp.int32)
            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # per-event mask of "the previous major step's recovery probe
            # just fired for this event".  Initialised to all-False so the
            # first ODE step has no nudge applied.  Shape matches the
            # other per-event carry vectors so the elementwise compare/
            # scatter inside ``_apply_recovery_w0_nudge`` aligns.
            zeno_recovery_just_cleared = jnp.zeros((n,), dtype=jnp.bool_)
            return SimulatorState(
                context=context,
                timed_events=timed_events,
                step_end_reason=end_reason,
                int_time=initial_int_time,
                results_data=results_data,
                ode_solver_state=self.ode_solver.initialize(context),
                zeno_tprev=zeno_tprev,
                zeno_active=zeno_active,
                zeno_frozen_steps=zeno_frozen_steps,
                zeno_recovery_just_cleared=zeno_recovery_just_cleared,
            )

        return SimulatorState(
            context=context,
            timed_events=timed_events,
            step_end_reason=end_reason,
            int_time=initial_int_time,
            results_data=results_data,
            ode_solver_state=self.ode_solver.initialize(context),
        )





    def _guarded_integrate(
        self,
        solver_state: ODESolverState,
        results_data: ResultsData,
        tf: float,
        context: ContextBase,
        zc_events: EventCollection,
        recovery_just_cleared=None,
        prior_zeno_active=None,
    ) -> tuple[bool, ODESolverState, ContextBase, ResultsData, EventCollection]:
        """Guarded ODE integration.

        Advance continuous time using an ODE solver, localizing any zero-crossing events
        that occur during the requested interval.  If any zero-crossing events trigger,
        the dense interpolant is used to localize the events and the associated reset maps
        are handled.  The method then returns, guaranteeing that the major step terminates
        either at the end of the requested interval or at the time of a zero-crossing
        event.

        Args:
            solver_state (ODESolverState): The current state of the ODE solver.
            results_data (ResultsData): The results data that will hold recorded time
                series data.
            tf (float): The end time of the integration interval.
            context (ContextBase): The current state of the system.
            zc_events (EventCollection): The current zero-crossing events.
            recovery_just_cleared: Optional per-event boolean mask flagging
                events whose Zeno latch was cleared by the recovery probe
                on the previous major step.  When non-None, a direction-
                aware nudge is applied to ``w0`` after ``record_interval_start``
                to recover the trigger semantics for events whose host
                leaf is at the post-reset rest condition.  See
                ``_apply_recovery_w0_nudge`` for the rationale.  Default
                ``None`` is the byte-equivalent legacy path (no nudge).

        Returns:
            tuple[bool, ODESolverState, ContextBase, ResultsData, EventCollection]:
                A tuple containing the following:
                - A boolean indicating whether the major step was terminated early due to
                  a zero-crossing event.
                - The updated state of the ODE solver.
                - The updated state of the system.
                - The updated results data.
                - The updated zero-crossing events.
        """
        solver = self.ode_solver
        func = solver.flat_ode_rhs  # Raveled ODE RHS function

        # Close over the additional arguments so that the RHS function has the
        # signature `func(y, t)`.
        def _func(y, t):
            return func(y, t, context)

        def _localize_zc_minor(
            solver_state, context_t0, context_tf, zc_events, results_data
        ):
            # Use the ZeroCrossingHandler to localize via bisection
            int_t1 = IntegerTime.from_decimal(context_tf.time)
            int_t0 = IntegerTime.from_decimal(context_t0.time)
            context_tf, zc_events = self.zc_handler.localize(
                solver_state, context_tf, zc_events, int_t0, int_t1
            )

            # record results sample for the ZC having 'occurred'
            minor_step_end_time = IntegerTime.as_decimal(int_t1)
            minor_step_start_time = IntegerTime.as_decimal(int_t0)
            zc_occur_time = context_tf.time - (
                minor_step_end_time - minor_step_start_time
            ) / (2 ** (self.zc_bisection_loop_count + 1))
            context_zc_time = context_tf.with_time(zc_occur_time)
            context_zc_time = context_zc_time.refresh_port_cache()
            # T-012a-followup: pre-localization solver_state's interp_coeff
            # spans the bracket that contained the ZC time; pass it through
            # so query() can later evaluate the polynomial at any t inside.
            results_data = self.results_recorder.save(
                results_data, context_zc_time, ode_solver_state=solver_state,
            )

            # Handle any triggered zero-crossing events
            context_tf = self.zc_handler.handle_events(zc_events, context_tf)

            # Re-initialize the solver since state may have been reset
            solver_state = solver.initialize(context_tf)
            return solver_state, context_tf, zc_events, results_data

        def _no_events_fun(
            solver_state, context_t0, context_tf, zc_events, results_data
        ):
            return solver_state, context_tf, zc_events, results_data

        # T-017b: when the system has no zero-crossing events, the
        # ``backend.cond(triggered, _localize_zc_minor, _no_events_fun, ...)``
        # below would still trace ``_localize_zc_minor`` (which retraces
        # ``solver.initialize`` — a measurable XLA cost for BDF/DAE
        # systems with mass matrices).  Skip that branch entirely at
        # Python level when the system declares no zero-crossings, and
        # also skip ``zc_handler.check_triggered``.  Numerically
        # bit-equivalent to the original path because ``triggered``
        # would always be False in that case.
        has_zero_crossings = self.system.has_zero_crossing_events

        def _ode_step(carry):
            _, solver_state, context_t0, results_data, zc_events = carry

            # Save results at the top of the loop. This will save data at t=t0,
            # but not at t=tf.  This is okay, since we will save the results at
            # the top of the next major step, as well as at the end of the main
            # simulation loop.
            context_t0 = context_t0.refresh_port_cache()
            # T-012a-followup: pass solver_state so the recorder snapshots
            # the per-step interpolant coefficients alongside (time, outputs).
            # Default-off path: ``record_solver_states=False`` means the
            # recorder ignores the kwarg — byte-equivalent to legacy.
            results_data = self.results_recorder.save(
                results_data, context_t0, ode_solver_state=solver_state,
            )

            zc_events = self.zc_handler.record_interval_start(zc_events, context_t0)

            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # direction-aware ``w0`` nudge for events whose Zeno latch was
            # just cleared by the recovery probe (or whose recorded ``w0``
            # is at the threshold).  Default-off path
            # (``recovery_just_cleared`` is None) skips at Python level
            # so the legacy hot path is byte-equivalent.  When applied,
            # the nudge is gated per-event by ``mask[i]`` so events not
            # in recovery flow through unchanged.  See
            # ``_apply_recovery_w0_nudge`` for the rationale.
            if recovery_just_cleared is not None:
                zc_events = self._apply_recovery_w0_nudge(
                    zc_events, recovery_just_cleared,
                    prior_zeno_active=prior_zeno_active,
                )

            # Advance ODE solver
            solver_state = solver.step(_func, tf, solver_state)
            xc = solver_state.unraveled_state
            context = context_t0.with_time(solver_state.t).with_continuous_state(xc)

            context = context.refresh_port_cache()

            if not has_zero_crossings:
                # No-ZC fast path: skip both ``check_triggered`` and the
                # ``cond(_localize_zc_minor, _no_events_fun)`` branch.
                return (False, solver_state, context, results_data, zc_events)

            # Check for zero-crossing events
            zc_events = self.zc_handler.check_triggered(zc_events, context)

            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # mask out triggers for events whose simulator-level Zeno
            # latch is engaged — the per-leaf freeze rollback at
            # ``_major_step`` already handles those leaves' continuous
            # state, and letting their ``triggered`` flag propagate here
            # would terminate the ODE step at a tiny dt and stall the
            # still-bouncing leaf's natural progression.
            if prior_zeno_active is not None:
                zc_events = self._mask_triggered_for_active_latch(
                    zc_events, prior_zeno_active,
                )

            triggered = zc_events.has_triggered

            args = (solver_state, context_t0, context, zc_events, results_data)
            solver_state, context, zc_events, results_data = backend.cond(
                triggered, _localize_zc_minor, _no_events_fun, *args
            )

            return (triggered, solver_state, context, results_data, zc_events)

        def _cond_fun(carry):
            triggered, solver_state, _, _, _ = carry
            return (solver_state.t < tf) & (~triggered)

        carry = (False, solver_state, context, results_data, zc_events)
        triggered, solver_state, context, results_data, zc_events = backend.while_loop(
            _cond_fun,
            _ode_step,
            carry,
        )

        return triggered, solver_state, context, results_data, zc_events



    def _advance_continuous_time(
        self,
        cdata: ContinuousIntervalData,
    ) -> ContinuousIntervalData:
        """Advance the simulation to the next discrete update or zero-crossing event.

        This stores the values of all active guard functions and advances the
        continuous-time component of the system to the next discrete update or
        zero-crossing event, whichever comes first.  Zero-crossing events are
        localized using a bisection search defined by `_trigger_search`, which will
        also record the final guard function values at the end of the search interval
        and determine which (if any) zero-crossing events were triggered.
        """

        # Unpack inputs
        int_tf = cdata.tf
        context = cdata.context
        results_data = cdata.results_data

        context = context.refresh_port_cache()
        zc_events = self.zc_handler.evaluate_guards(context)

        if self.system.has_continuous_state:
            solver_state = cdata.ode_solver_state
            tf = IntegerTime.as_decimal(int_tf)

            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # plumb the per-event ``recovery_just_cleared`` mask down into
            # ``_guarded_integrate`` so the inner ``_ode_step`` can apply
            # the direction-aware ``w0`` nudge.  When the autodiff custom
            # VJP wrapper is in play, ``self.guarded_integrate`` is a
            # ``custom_vjp`` callable with a fixed 5-arg signature — call
            # the unwrapped ``_guarded_integrate`` directly in that case
            # since autodiff users don't go through the Zeno cascade
            # pathway in practice and the nudge would have no effect on
            # the gradient (it is an idempotent ``where`` on a value
            # that, in normal operation, already satisfies the trigger
            # direction inequality).  Default ``recovery_just_cleared
            # is None`` keeps the legacy 5-arg call so the byte-
            # equivalent default-off path is preserved.
            rjc = cdata.recovery_just_cleared
            pza = cdata.prior_zeno_active
            if rjc is None and pza is None:
                (
                    triggered,
                    solver_state,
                    context,
                    results_data,
                    zc_events,
                ) = self.guarded_integrate(
                    solver_state,
                    results_data,
                    tf,
                    context,
                    zc_events,
                )
            else:
                (
                    triggered,
                    solver_state,
                    context,
                    results_data,
                    zc_events,
                ) = self._guarded_integrate(
                    solver_state,
                    results_data,
                    tf,
                    context,
                    zc_events,
                    recovery_just_cleared=rjc,
                    prior_zeno_active=pza,
                )

            context = context.with_time(solver_state.t)
            context = context.with_continuous_state(solver_state.unraveled_state)

            # Converting from decimal -> integer time incurs a loss of precision.  This is
            # okay for unscheduled zero-crossing events, but problematic for timed events.
            # So only do this conversion if a zero-crossing was triggered.  Otherwise we
            # know we have reached the end of the interval and can keep the requested end
            # time.
            int_tf = npa.where(
                triggered,
                IntegerTime.from_decimal(context.time),
                int_tf,
            )

        else:
            # Skip the ODE solver for systems without continuous state.  We still
            # have to check for triggered events here in case there are any
            # transitions triggered by time that need to be handled before the
            # periodic discrete update at the top of the next major step
            triggered = False
            solver_state = cdata.ode_solver_state

            zc_events = self.zc_handler.record_interval_start(zc_events, context)
            results_data = self.results_recorder.save(results_data, context)

            # Advance time to the end of the interval
            context = context.with_time(IntegerTime.as_decimal(int_tf))
            context = context.refresh_port_cache()

            # Record guard values after the discrete update and check if anything
            # triggered as a result of advancing time
            zc_events = self.zc_handler.record_interval_end(zc_events, context)
            zc_events = self.zc_handler.check_triggered(zc_events, context)

            # Handle any triggered zero-crossing events
            context = self.zc_handler.handle_events(zc_events, context)

        # Even though the zero-crossing events have already been "handled", the
        # information about whether a terminal event has been triggered is still in
        # the events collection (since "triggered" has not been cleared by a call
        # to determine_triggered_guards).
        terminate_early = zc_events.has_active_terminal

        # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        # extract a per-event triggered mask from the post-step events
        # collection so the simulator-level Zeno tracker can update each
        # event's ``tprev[i]`` independently.  Without this, the scalar
        # ``triggered`` (any event fired) broadcasts to all events,
        # which spuriously latches unrelated events when a single
        # leaf's cascade drives sub-tolerance major-step ends — e.g.
        # ball A's bouncing cascade causing ball B's latch to engage.
        # Default-off path: ``zeno_protection_enabled=False`` ignores
        # this field (cdata.per_event_triggered carries through as
        # ``None``), so the byte-equivalent legacy path is preserved.
        # The terminal-early branch in ``_major_step`` keeps whatever
        # placeholder cdata had at construction time, so when Zeno is
        # enabled we always populate this with a fixed-shape array
        # whether or not events fired (the all-False zeros from the
        # placeholder remain semantically valid in that case).
        if cdata.per_event_triggered is not None and self.n_zero_crossing_events > 0:
            per_event_triggered = self._extract_per_event_triggered(zc_events)
        else:
            per_event_triggered = cdata.per_event_triggered

        # T-125-followup-record-event-times: tee ``(time, per_event_mask)``
        # to the host-side recorder when the option is on AND the diagram
        # has zero-crossing events.  The recorder ignores all-False masks
        # host-side so non-triggering major steps add nothing.  Default-
        # off path: ``self._event_time_recorder is None`` skips the
        # entire block at trace time, preserving byte-equivalence.
        # ``context.time`` is already the localized event firing time
        # when ``triggered`` is True (``_advance_continuous_time`` sets
        # ``context.time = solver_state.t`` after ``guarded_integrate``
        # which clamps to the bisection root); when ``triggered`` is
        # False the mask is all-False and the host-side guard short-
        # circuits without appending anything.
        if (
            self._event_time_recorder is not None
            and self.n_zero_crossing_events > 0
        ):
            event_mask = self._extract_per_event_triggered(zc_events)
            jax.debug.callback(
                self._event_time_recorder.update,
                context.time,
                event_mask,
            )

        return cdata._replace(
            triggered=triggered,
            terminate_early=terminate_early,
            context=context,
            tf=int_tf,
            results_data=results_data,
            ode_solver_state=solver_state,
            per_event_triggered=per_event_triggered,
        )

    def _extract_per_event_triggered(self, zc_events):
        """T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        flatten ``zc_events`` into a per-event boolean array of shape
        ``(N_events,)`` aligned with ``self.event_directions`` /
        ``self.event_system_ids`` / per-event Zeno carry vectors.

        Walks the events tree in flatten order — the same order as
        ``system.zero_crossing_events.events``.  Each
        ``event.event_data.triggered`` is a JAX scalar that may be a
        tracer or a concrete bool; we collect them into a stack and
        return a ``(N_events,)`` array.

        Used by ``_major_step`` to plumb per-event trigger info up to
        ``_update_zeno_tracking`` so each event's ``tprev[i]`` and
        engagement check is per-event rather than scalar-broadcast.
        """
        triggered_list: list = []

        def _collect(event):
            if isinstance(event, ZeroCrossingEvent):
                triggered_list.append(
                    jnp.asarray(event.event_data.triggered, dtype=jnp.bool_),
                )
                # Returning ``event`` keeps the tree walk happy; the
                # actual collection happens via the side-effect closure.
                return event
            return event

        jax.tree_util.tree_map(
            _collect, zc_events,
            is_leaf=lambda x: isinstance(x, ZeroCrossingEvent),
        )
        if not triggered_list:
            return jnp.zeros((0,), dtype=jnp.bool_)
        return jnp.stack(triggered_list)

    def _handle_discrete_update(
        self, context: ContextBase, timed_events: EventCollection
    ) -> tuple[ContextBase, bool]:
        """Handle discrete updates triggered by time.

        This method is called at the beginning of each major step to handle any
        discrete updates that are triggered by time.  This includes both discrete
        updates that are triggered by time and any zero-crossing events that are
        triggered by the discrete update.

        This will also work when there are no zero-crossing events: the zero-crossing
        collection will be empty and only the periodic discrete update will happen.

        Args:
            context (ContextBase): The current state of the system.
            timed_events (EventCollection):
                The collection of timed events, with the active events marked.

        Returns:
            ContextBase: The updated state of the system.
            bool: Whether the simulation should terminate early as a result of a
                triggered terminal condition.
        """
        return self.zc_handler.check_after_discrete_update(context, timed_events)

    def _update_zeno_tracking(
        self,
        zeno_tprev,
        zeno_active,
        zeno_frozen_steps,
        triggered,
        current_time,
    ):
        """T-027a-followup: simulator-level Zeno tracker with recovery probe.

        Updates the carry triple ``(tprev, active, frozen_steps)`` after a
        major step.  Logic, in order:

        1. If the major step ended on a guard trigger AND the time since
           the last recorded trigger is below ``zeno_tolerance``, latch
           the protection on (``active=True``).  The latch is sticky —
           once on, only the recovery probe (step 3) clears it.
        2. Update ``tprev`` to the current step's end time whenever a
           guard fired, so the tolerance check in step 1 reflects the
           most recent triggering history.
        3. Recovery probe: after ``zeno_recovery_period`` consecutive
           frozen major steps, clear ``active`` for one step.  The next
           guard-trigger check then either re-engages Zeno (cascade
           ongoing — tolerance check fires again) or leaves it cleared
           (cascade resolved).  The frozen-step counter resets on every
           ``active`` flip.

        T-027a-followup-vector-tprev: ``zeno_tprev`` and ``zeno_active``
        are per-event vectors of shape ``(N_events,)``; ``triggered`` may
        be a scalar (broadcast to the full vector) or a same-shape per-
        event mask.  Each event's last-firing time is tracked
        independently, so unrelated events do not contaminate each
        other's tolerance check.

        T-027a-followup-per-event-recovery: ``zeno_frozen_steps`` is a
        per-event vector of shape ``(N_events,)``.  Each event's
        counter increments while THAT event's ``active[i]`` is True and
        resets when it clears.  When a single event ``i`` hits
        ``zeno_recovery_period``, only ``active[i]`` (and its own
        counter) clear — other events' latches and counters are
        untouched, so a still-cascading event A no longer gets a free
        probe just because event B's cascade ended.  When a scalar
        ``zeno_frozen_steps`` is passed in (the default-disabled
        SimulatorState carry), it is broadcast to the per-event shape;
        the byte-equivalent default-off path never enters this branch.

        The actual freeze (zero out ode_rhs while ``active``) is not
        wired in here.  ``zeno_active`` is observational for now;
        callers can inspect ``sim_state.zeno_active`` to diagnose Zeno
        cascades without changing simulation outputs.  The per-leaf
        freeze gate that consumes ``zeno_active[i]`` and the static
        event→leaf map ``self.event_system_ids`` is filed under
        T-027a-followup-per-leaf-freeze.
        """
        # Treat ``triggered`` / ``active`` as JAX-friendly booleans so this
        # works under jit/vmap as well as the Python/numpy backend.  Match
        # the dtypes of the in/out carry fields to the SimulatorState
        # defaults so ``lax.cond`` true/false-branch dtype checks pass.
        tprev = jnp.asarray(zeno_tprev)
        active_b = jnp.asarray(zeno_active, dtype=jnp.bool_)
        # T-027a-followup-per-event-recovery: ``frozen`` is a per-event
        # ``(N_events,)`` vector when Zeno protection is enabled.  A
        # scalar input (legacy default-off carry, or the V-005 helper
        # ``_step_zeno`` which passes ``int(frozen)``) is broadcast up
        # to ``tprev.shape`` so the elementwise compare/where below
        # works for both cases.
        frozen = jnp.broadcast_to(
            jnp.asarray(zeno_frozen_steps, dtype=jnp.int32), tprev.shape,
        )
        # Broadcast a scalar ``triggered`` up to the per-event tprev /
        # active shape so the masks align elementwise.  Callers passing
        # a properly-shaped per-event mask flow through unchanged.
        triggered_b = jnp.broadcast_to(
            jnp.asarray(triggered, dtype=jnp.bool_), tprev.shape,
        )
        t = jnp.asarray(current_time, dtype=tprev.dtype)
        tol = jnp.asarray(self.zeno_tolerance, dtype=tprev.dtype)
        period = jnp.asarray(self.zeno_recovery_period, dtype=frozen.dtype)
        one_step = jnp.asarray(1, dtype=frozen.dtype)
        zero_step = jnp.asarray(0, dtype=frozen.dtype)

        # Engagement condition: a guard fired AND inter-event time was
        # below tolerance.  Per-event: each ``tprev[i]`` is checked
        # against ``t`` independently.
        dt_since_prev = t - tprev
        engage = triggered_b & (dt_since_prev < tol) & (dt_since_prev >= 0)

        # Per-event active value before the recovery probe.
        active_after_engage = active_b | engage

        # T-027a-followup-per-event-recovery: per-event frozen-step
        # counter.  Each event's ``frozen[i]`` increments while its own
        # ``active_after_engage[i]`` is True and resets to zero when
        # cleared.  The recovery probe fires per-event: when
        # ``frozen[i] >= K``, only ``active[i]`` (and ``frozen[i]``)
        # clear — other events keep their latches and counters.  This
        # prevents a still-cascading event from getting a free probe
        # just because an unrelated event's cascade ended.
        new_frozen = jnp.where(active_after_engage, frozen + one_step, zero_step)
        should_probe = active_after_engage & (new_frozen >= period)
        new_active = jnp.where(should_probe, False, active_after_engage)
        new_frozen = jnp.where(should_probe, zero_step, new_frozen)

        # Update ``tprev[i]`` whenever event ``i`` fired so each event's
        # next tolerance check sees its own most-recent firing time.
        new_tprev = jnp.where(triggered_b, t, tprev)

        return new_tprev, new_active, new_frozen

    def _mask_triggered_for_active_latch(self, zc_events, prior_zeno_active):
        """T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        suppress per-event ``triggered`` when the simulator-level Zeno
        latch is engaged for that event.

        When a leaf is in Zeno hold, the per-leaf freeze rollback at
        the end of ``_major_step`` re-pins its continuous state to the
        snapshot.  But the inner ``_ode_step`` loop checks
        ``zc_events.has_triggered`` and exits early on a trigger, so a
        latched-but-still-firing event would terminate the ODE step at
        a tiny dt — preventing any unfrozen leaf (e.g. ball B in a
        staggered cascade) from advancing.  Masking the latched
        events' triggered flag to False lets the ODE step run to its
        natural termination (next unfrozen-leaf trigger or interval
        end), giving the staggered cascade case a path forward.

        Default-off path (``zeno_protection_enabled=False``) never
        calls this helper, so the legacy semantics are unchanged.
        """
        if self.n_zero_crossing_events == 0 or prior_zeno_active is None:
            return zc_events

        active_mask = jnp.asarray(prior_zeno_active, dtype=jnp.bool_)
        idx_box = [0]

        def _update(event):
            if not isinstance(event, ZeroCrossingEvent):
                return event
            i = idx_box[0]
            idx_box[0] += 1
            old_triggered = jnp.asarray(
                event.event_data.triggered, dtype=jnp.bool_,
            )
            new_triggered = old_triggered & ~active_mask[i]
            return dataclasses.replace(
                event,
                event_data=dataclasses.replace(
                    event.event_data, triggered=new_triggered,
                ),
            )

        return jax.tree_util.tree_map(
            _update, zc_events,
            is_leaf=lambda x: isinstance(x, ZeroCrossingEvent),
        )

    def _apply_recovery_w0_nudge(
        self,
        zc_events,
        recovery_just_cleared,
        prior_zeno_active=None,
    ):
        """T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        direction-aware ``w0`` nudge so events whose host leaf is at the
        post-reset rest condition can still trigger.

        Background: with reset maps that clamp the continuous state to
        the threshold (e.g. ``h := max(h, 0)`` for a bouncing ball), the
        post-reset state sits exactly on the guard's threshold.  The
        next major step's ``_ode_step`` records ``w0 = guard(t_start) = 0``
        in ``record_interval_start``.  When the post-reset velocity is
        small enough that the ODE step jumps the state to the other
        side of the threshold inside one solver step, the trigger
        condition for ``positive_then_non_positive``,
        ``(w0 > 0) & (w1 <= 0)``, evaluates to ``False & True = False``
        because ``w0 = 0`` is not strictly positive.  No trigger fires,
        the bounce reset is missed, and the leaf is left in free-fall
        below the threshold.  Subsequent steps see ``w0 < 0`` so the
        guard never refires.

        Originally surfaced via the per-event recovery probe (the
        cleared-step has the same ``(h=0, v ≈ 0+)`` initial condition);
        in practice it also fires on the staggered multi-leaf cascade
        BEFORE the latch ever engages, when ball A's natural cascade
        convergence drives the post-reset velocity below the solver's
        step-size resolution.  The nudge therefore fires whenever the
        recorded ``w0`` is on the "wrong side" of the trigger direction's
        active region by at most ``zeno_tolerance``, irrespective of
        whether the recovery probe just fired:

        - ``positive_then_non_positive``: active region is ``w > 0``;
          fire when ``w0 < eps`` (i.e. at or below threshold).  Set
          ``w0 := +eps``.
        - ``negative_then_non_negative``: active region is ``w < 0``;
          fire when ``w0 > -eps``.  Set ``w0 := -eps``.
        - ``crosses_zero``: symmetric trigger; fire when
          ``|w0| < eps`` and pick ``+eps`` (the side the typical
          contact-penetration constraint retreats into).
        - ``none`` / ``edge_detection``: no nudge.

        The ``recovery_just_cleared`` mask is preserved as a "force
        nudge" override — when set, the nudge fires regardless of the
        ``w0``-side check.  This mirrors the original (c) intent and
        gives a known-safe path even if the at-threshold check were
        too conservative for some pathological case.

        Idempotent: when ``w0`` is already strictly on the active side
        with magnitude above ``eps``, the nudge is a no-op.  Default-
        off path: ``zeno_protection_enabled=False`` skips this entirely
        at Python level (``recovery_just_cleared`` is ``None`` and the
        call site short-circuits) — byte-equivalent to legacy.

        Implementation note: a Python list closure counter walks the
        ``tree_map`` in flatten order — the same order as
        ``system.zero_crossing_events.events`` and the per-event
        ``recovery_just_cleared`` mask.  ``LeafEventCollection`` and
        ``DiagramEventCollection`` both register flatten orders that
        match the ``events`` property iteration order.  Under JIT/AD,
        the trace runs once, so the counter sees each event exactly
        once in the right order.
        """
        # Empty-events fast path.  Without this guard, the closure counter
        # below would still execute but produce no work; this avoids the
        # tree_map call entirely.
        if self.n_zero_crossing_events == 0:
            return zc_events

        # Static directions tuple captured at trace time — Python strings
        # so the per-event branch is a Python `if` (no traced control
        # flow), not a `lax.switch`.
        directions = self.event_directions
        mask = jnp.asarray(recovery_just_cleared, dtype=jnp.bool_)
        # When ``prior_zeno_active[i]`` is True, event ``i``'s host leaf
        # is in Zeno hold and the per-leaf freeze rollback is the
        # mechanism keeping its state pinned.  Skipping the nudge for
        # frozen events prevents a tight loop where the nudge fires
        # the trigger, the reset re-pins state at threshold, and the
        # next ODE step would re-fire — burning major steps without
        # advancing time.  Default ``None`` is treated as "no event
        # frozen", which is the correct behaviour pre-engagement and
        # for systems without Zeno protection.
        if prior_zeno_active is None:
            frozen_mask = jnp.zeros_like(mask)
        else:
            frozen_mask = jnp.asarray(prior_zeno_active, dtype=jnp.bool_)
        # Mirror the zeno tolerance for the nudge magnitude.  This is
        # intentionally larger than machine epsilon: it has to be big
        # enough that the next ODE step's ``w1`` lands on the OTHER
        # side of zero so the trigger condition ``w0 > 0 & w1 <= 0``
        # fires, but small enough that it does not perturb the bounce-
        # reset's localised trigger time meaningfully.  ``zeno_tolerance``
        # is the natural scale here — it is the same threshold the
        # latch engagement uses to decide "this firing is sub-tolerance
        # of the previous one".
        eps = jnp.asarray(self.zeno_tolerance)

        idx_box = [0]

        def _update(event):
            if not isinstance(event, ZeroCrossingEvent):
                return event
            i = idx_box[0]
            idx_box[0] += 1
            d = directions[i]
            if d == "positive_then_non_positive":
                # Trigger needs ``w0 > 0`` strictly.  Fire the nudge
                # when the recorded ``w0`` is at or below threshold
                # (``w0 < eps``) OR when the recovery probe just fired
                # for this event.  Skip when the leaf is in Zeno hold
                # (the per-leaf freeze rollback handles it).
                old_w0 = jnp.asarray(event.event_data.w0)
                eps_typed = jnp.asarray(eps, dtype=old_w0.dtype)
                should_nudge = (
                    (mask[i] | (old_w0 < eps_typed)) & ~frozen_mask[i]
                )
                new_w0 = jnp.where(should_nudge, eps_typed, old_w0)
            elif d == "negative_then_non_negative":
                # Trigger needs ``w0 < 0`` strictly.  Fire the nudge
                # when ``w0 > -eps`` OR mask is set.  Skip when frozen.
                old_w0 = jnp.asarray(event.event_data.w0)
                eps_typed = jnp.asarray(eps, dtype=old_w0.dtype)
                neg_eps = -eps_typed
                should_nudge = (
                    (mask[i] | (old_w0 > neg_eps)) & ~frozen_mask[i]
                )
                new_w0 = jnp.where(should_nudge, neg_eps, old_w0)
            elif d == "crosses_zero":
                # Symmetric trigger.  Fire the nudge when ``|w0| < eps``
                # OR mask is set; pick ``+eps`` as the active side.
                # Skip when frozen.
                old_w0 = jnp.asarray(event.event_data.w0)
                eps_typed = jnp.asarray(eps, dtype=old_w0.dtype)
                should_nudge = (
                    (mask[i] | (jnp.abs(old_w0) < eps_typed)) & ~frozen_mask[i]
                )
                new_w0 = jnp.where(should_nudge, eps_typed, old_w0)
            else:
                # ``"none"`` / ``"edge_detection"`` — no nudge applies.
                return event
            return dataclasses.replace(
                event,
                event_data=dataclasses.replace(event.event_data, w0=new_w0),
            )

        return jax.tree_util.tree_map(
            _update, zc_events, is_leaf=lambda x: isinstance(x, ZeroCrossingEvent),
        )

    def _apply_per_leaf_zeno_freeze(
        self,
        pre_xc,
        post_context,
        pre_solver_state,
        post_solver_state,
        zeno_active,
    ):
        """T-027a-followup-per-leaf-freeze: per-leaf rollback of continuous state.

        When one or more events have their Zeno latch engaged, the host
        leaf of each engaged event has its slice of the post-ODE
        continuous state replaced with the pre-ODE snapshot.  Other
        leaves keep their advanced state.  The choice is per-element:
        ``jnp.where(leaf_frozen, pre_leaf, post_leaf)``.

        For a single-LeafSystem simulation, ``pre_xc`` /
        ``post_context.continuous_state`` are scalar Arrays (one leaf,
        the "list" collapses), so the gate becomes a single
        ``jnp.where``.

        T-027a-followup-per-leaf-solver-state: the solver-state
        rollback now decomposes per leaf for the fields whose last
        axis matches the flat ODE state dimension ``n_y``.  An element-
        wise mask of length ``n_y`` (True at indices owned by frozen
        leaves) is built from the per-leaf freeze decisions and the
        static ``_leaf_flat_slices`` map; ``jnp.where(mask, pre, post)``
        rolls the frozen-leaf slice of ``y``, ``f``, and the per-step
        interpolant table back to the pre-ODE snapshot, while non-
        frozen leaves' slices keep their post-ODE values.  Scalar
        integration-step fields (``t``, ``dt``, ``t_prev``, ``t_return``,
        ``n_acc``, ``n_rej``, ``accepted``, ``order``,
        ``n_equal_steps``, ``updated_jacobian``) are global to the
        integrator's adaptive step controller and stay at the post-
        step values.  BDF Jacobian/mass/LU matrices (``J``, ``M``,
        ``LU``, ``U``) are ``(n_y, n_y)``-shaped with cross-leaf
        coupling and are NOT split — they roll back all-or-nothing
        when any leaf is frozen, matching the prior all-or-nothing
        behaviour for those specific fields.  Per-leaf-decomposable:
        Dopri5 ``y``, ``f``, ``interp_coeff``; BDF ``y``, ``f``,
        ``D``.
        """
        active_b = jnp.asarray(zeno_active, dtype=jnp.bool_)
        any_active = jnp.any(active_b)

        # Per-leaf freeze decisions (one Python bool tracer per cs_idx).
        # Built from the static ``_cs_idx_to_event_positions`` map so
        # this loop runs at trace-time and produces a fixed-shape pytree.
        # Leaves with no events trivially stay un-frozen.
        per_leaf_frozen: list = []
        for cs_i in range(self._n_continuous_leaves):
            positions = self._cs_idx_to_event_positions.get(cs_i, ())
            if positions:
                leaf_mask = active_b[jnp.asarray(positions)]
                per_leaf_frozen.append(jnp.any(leaf_mask))
            else:
                per_leaf_frozen.append(jnp.asarray(False))

        # Per-leaf gate on continuous state.
        post_xc = post_context.continuous_state
        if isinstance(post_xc, list):
            # DiagramContext path: per-leaf gate.
            new_xc_list = list(post_xc)
            for cs_i in range(len(post_xc)):
                positions = self._cs_idx_to_event_positions.get(cs_i, ())
                if not positions:
                    continue
                leaf_frozen = per_leaf_frozen[cs_i]
                pre_leaf = pre_xc[cs_i]
                post_leaf = post_xc[cs_i]
                new_xc_list[cs_i] = jax.tree_util.tree_map(
                    lambda p, n, _frozen=leaf_frozen: jnp.where(_frozen, p, n),
                    pre_leaf, post_leaf,
                )
            new_context = post_context.with_continuous_state(new_xc_list)
        else:
            # LeafContext path: single Array (or pytree).  Any event
            # active = freeze.  When the host system IS the leaf, the
            # cs_i=0 freeze decision covers every event.
            if per_leaf_frozen:
                leaf_frozen = per_leaf_frozen[0]
            else:
                leaf_frozen = any_active
            new_xc = jax.tree_util.tree_map(
                lambda p, n: jnp.where(leaf_frozen, p, n),
                pre_xc, post_xc,
            )
            new_context = post_context.with_continuous_state(new_xc)

        # Per-leaf solver-state rollback.  Build a (n_y,) bool mask
        # True at indices belonging to frozen leaves; ``jnp.where``
        # along the last axis splits the per-leaf-decomposable fields.
        new_solver_state = self._per_leaf_solver_gate(
            pre_solver_state, post_solver_state, per_leaf_frozen, any_active,
        )
        return new_context, new_solver_state

    def _per_leaf_solver_gate(
        self,
        pre_solver_state,
        post_solver_state,
        per_leaf_frozen,
        any_active,
    ):
        """T-027a-followup-per-leaf-solver-state: decompose the solver-
        state rollback per leaf along the flat ``n_y`` axis.

        Fields gated per-leaf (last-axis ``n_y``):
          - Dopri5State: ``y``, ``f``, ``interp_coeff``
          - BDFState: ``y``, ``f``, ``D``

        Fields kept post-step (global to the adaptive step controller):
          - ``t``, ``dt``, ``t_prev``, ``t_return``,
            ``n_acc``, ``n_rej``, ``accepted``,
            ``order``, ``n_equal_steps``, ``updated_jacobian``

        Fields rolled back all-or-nothing under ``any_active`` (BDF
        Jacobian/mass/LU matrices are ``(n_y, n_y)``-shaped with
        cross-leaf coupling, not last-axis-decomposable):
          - ``J``, ``M``, ``LU``, ``U``

        When ``self._n_y_total == 0`` (no continuous state) or no leaf
        is frozen, returns ``post_solver_state`` unchanged.  When the
        flat-slice metadata is missing (a defensive fallback), reverts
        to the previous all-or-nothing rollback.
        """
        if self._n_y_total <= 0 or not per_leaf_frozen:
            # Nothing to gate per-leaf — fall back to the
            # all-or-nothing path so behaviour matches the prior
            # T-027a-followup-per-leaf-freeze contract.
            return backend.cond(
                any_active,
                lambda _ss: pre_solver_state,
                lambda _ss: _ss,
                post_solver_state,
            )

        # Build the (n_y,) elementwise mask: True where the leaf
        # owning this index is frozen.  Static ``_leaf_flat_slices``
        # bounds plus per-leaf freeze decisions.
        n_y = self._n_y_total
        leaf_indicators = []
        for cs_i, (start, end) in enumerate(self._leaf_flat_slices):
            seg_len = end - start
            if seg_len <= 0:
                continue
            frozen_i = per_leaf_frozen[cs_i] if cs_i < len(per_leaf_frozen) else jnp.asarray(False)
            seg = jnp.broadcast_to(
                jnp.asarray(frozen_i, dtype=jnp.bool_), (seg_len,),
            )
            leaf_indicators.append(seg)
        if not leaf_indicators:
            return post_solver_state
        mask = jnp.concatenate(leaf_indicators)  # shape (n_y,)
        # Defensive: if any leaf had a flat-size mismatch, keep the
        # all-or-nothing fallback to avoid accidentally mis-aligning
        # the slices.  Static check at trace time.
        if mask.shape[0] != n_y:
            return backend.cond(
                any_active,
                lambda _ss: pre_solver_state,
                lambda _ss: _ss,
                post_solver_state,
            )

        def _gate_last_axis(pre_arr, post_arr):
            # Broadcast the (n_y,) mask to ``post_arr`` shape — the
            # mask aligns with the trailing axis (n_y).
            return jnp.where(mask, pre_arr, post_arr)

        # Detect Dopri5State vs BDFState by attribute presence —
        # avoids importing the backend classes here (the simulator
        # is backend-agnostic).
        post = post_solver_state
        pre = pre_solver_state
        kwargs = {}
        # Always-present per-leaf-decomposable fields.
        kwargs["y"] = _gate_last_axis(pre.y, post.y)
        kwargs["f"] = _gate_last_axis(pre.f, post.f)
        # Dopri5: interp_coeff has shape (5, n_y).
        if hasattr(post, "interp_coeff") and post.interp_coeff is not None:
            kwargs["interp_coeff"] = _gate_last_axis(
                pre.interp_coeff, post.interp_coeff,
            )
        # BDF: D has shape (MAX_ORDER+3, n_y).
        if hasattr(post, "D") and post.D is not None:
            kwargs["D"] = _gate_last_axis(pre.D, post.D)
        # BDF Jacobian/mass/LU matrices are (n_y, n_y) with cross-leaf
        # coupling — roll back all-or-nothing under any_active.
        if hasattr(post, "J") and post.J is not None:
            kwargs["J"] = jnp.where(any_active, pre.J, post.J)
        if hasattr(post, "M") and post.M is not None:
            kwargs["M"] = jnp.where(any_active, pre.M, post.M)
        if hasattr(post, "LU") and post.LU is not None:
            kwargs["LU"] = jnp.where(any_active, pre.LU, post.LU)
        if hasattr(post, "U") and post.U is not None:
            kwargs["U"] = jnp.where(any_active, pre.U, post.U)

        # Replace only the gated fields.  Scalar / integration-step
        # fields stay at their post-step values automatically.
        return dataclasses.replace(post, **kwargs)

    def _major_step(
        self,
        sim_state: SimulatorState,
        int_boundary_time: int,
        int_max_step_length: int,
    ) -> SimulatorState:
        end_reason = sim_state.step_end_reason
        context = sim_state.context
        timed_events = sim_state.timed_events
        int_time = sim_state.int_time

        if not self.enable_tracing:
            logger.debug("Starting a simulation step at t=%s", context.time)
            logger.debug("   merged_events: %s", timed_events)

        # Handle any discrete updates that are triggered by time along with
        # any zero-crossing events that are triggered by the discrete update.
        context, terminate_early = self._handle_discrete_update(
            context, timed_events
        )
        logger.debug("Terminate early after discrete update: %s", terminate_early)

        # T-113-followup-event-reprojection: project the post-reset state
        # back onto the constraint manifold *before* continuous integration
        # resumes.  ``_handle_discrete_update`` may have applied a discrete
        # update or a ZC reset triggered by it — either can drop algebraic
        # states off the manifold.  Default-off path
        # (``dae_reproject_after_events=False``) skips the block at trace
        # time so the disabled hot path is byte-equivalent.  No-op for
        # systems without a mass matrix.  Composes with T-003a's end-of-
        # major-step projection (both can run — they target different
        # within-step instants).
        if self.dae_reproject_after_events and getattr(
            self.system, "has_mass_matrix", False,
        ):
            from .dae_projection import project_constraints
            context = project_constraints(
                self.system,
                context,
                tol=self.dae_projection_tol,
                max_iter=self.dae_projection_max_iter,
            )

        # How far can we go before we have to handle timed events?
        # The time returned here will be the integer time representation.
        time_of_next_timed_event, timed_events = _next_update_time(
            self.system.periodic_events, int_time
        )
        if not self.enable_tracing:
            logger.debug(
                "Next timed event at t=%s",
                IntegerTime.as_decimal(time_of_next_timed_event),
            )
            timed_events.pprint(logger.debug)

        # Determine whether the events include a timed update
        update_time = IntegerTime.max_int_time

        if timed_events.num_events > 0:
            update_time = time_of_next_timed_event

        # Limit the major step end time to the simulation end time, major step limit,
        # or next periodic update time.
        # This is the mechanism used to advance time for systems that have
        # no states and no periodic events.
        # Discrete systems] when there are discrete periodic events, we use those
        # to determine each major step end time.
        # Feedthrough system] when there are just feedthrough blocks (no states or
        # events), use max_major_step_length to determine each major step end time.
        int_tf_limit = int_time + int_max_step_length
        int_tf = npa.min(
            npa.array(
                [
                    int_boundary_time,
                    int_tf_limit,
                    update_time,
                ]
            )
        )
        if not self.enable_tracing:
            logger.debug(
                "Expecting to integrate to t=%s",
                IntegerTime.as_decimal(int_tf),
            )

        # T-027a-followup-per-leaf-freeze: snapshot the pre-ODE
        # continuous state and ODE solver state when simulator-level
        # Zeno protection is enabled.  These are restored *after*
        # ``_advance_continuous_time`` for leaves whose Zeno latch was
        # already engaged at the START of this major step (carried in
        # by ``sim_state.zeno_active`` from step N-1) — mirroring the
        # leaf-level pattern where a latch set at step N-1 prevents
        # continuous-state evolution at step N.  This avoids the
        # pathological "freeze the just-completed reset map" replay.
        # Default-off: no snapshot, no extra ops — byte-equivalent.
        if self.zeno_protection_enabled:
            pre_ode_xc = context.continuous_state
            pre_ode_solver_state = sim_state.ode_solver_state
            prior_zeno_active = sim_state.zeno_active
            prior_any_active = jnp.any(
                jnp.asarray(prior_zeno_active, dtype=jnp.bool_),
            )
            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # carry the prior step's recovery-probe-fired mask through to
            # the inner ODE step's guard-recording site.  When event
            # ``i``'s latch was cleared by the recovery probe at step
            # N-1, this step (N) is the "first post-recovery" step and
            # ``w0[i]`` may need a direction-aware nudge to compensate
            # for the resting-state guard value (e.g. ``h=0`` for a
            # bouncing ball with ``max(h, 0)`` reset clamping).
            prior_recovery_just_cleared = sim_state.zeno_recovery_just_cleared
            # Initial ``per_event_triggered`` placeholder for the cond
            # branches' pytree-shape consistency.  The terminal-early
            # branch passes cdata through unchanged so we need a fixed
            # array shape here; ``_advance_continuous_time`` overrides
            # this with the actual post-step per-event mask.
            n_events = max(self.n_zero_crossing_events, 1)
            per_event_triggered_initial = jnp.zeros((n_events,), dtype=jnp.bool_)
        else:
            pre_ode_xc = None
            pre_ode_solver_state = None
            prior_zeno_active = None
            prior_any_active = None
            prior_recovery_just_cleared = None
            per_event_triggered_initial = None

        # Normally we will advance continuous time to the end of the major step
        # here. However, if a terminal event was triggered as part of the discrete
        # update, we should respect that and skip the continuous update.
        #
        # Construct the container used to hold various data related to advancing
        # continuous time.  This is passed to ODE solvers, zero-crossing
        # localization, and related functions.
        if self.system.has_continuous_state:
            leaves = jax.tree.leaves(context.continuous_state)
            dtype = leaves[0].dtype if leaves else jnp.empty(0).dtype
            context = context.with_time(jnp.asarray(context.time, dtype=dtype))
            if sim_state.ode_solver_state is not None:
                def _cast_leaf(x):
                    if isinstance(x, (float, np.floating)):
                        return jnp.asarray(x, dtype=dtype)
                    if isinstance(x, jnp.ndarray) and jnp.issubdtype(x.dtype, jnp.floating):
                        return x.astype(dtype)
                    return x
                ode_solver_state = jax.tree.map(_cast_leaf, sim_state.ode_solver_state)
            else:
                ode_solver_state = sim_state.ode_solver_state
        else:
            ode_solver_state = sim_state.ode_solver_state

        cdata = ContinuousIntervalData(
            context=context,
            terminate_early=terminate_early,
            triggered=False,
            t0=int_time,
            tf=int_tf,
            results_data=sim_state.results_data,
            ode_solver_state=ode_solver_state,
            recovery_just_cleared=prior_recovery_just_cleared,
            per_event_triggered=per_event_triggered_initial,
            prior_zeno_active=prior_zeno_active,
        )
        cdata = backend.cond(
            (self.has_terminal_events & cdata.terminate_early),
            lambda cdata: cdata,  # Terminal event triggered - return immediately
            self._advance_continuous_time,  # Advance continuous time normally
            cdata,
        )

        # Unpack the results of the continuous time advance
        context = cdata.context
        terminate_early = cdata.terminate_early
        triggered = cdata.triggered
        int_tf = cdata.tf
        results_data = cdata.results_data
        ode_solver_state = cdata.ode_solver_state

        # Determine the reason why the major step ended.  Did a zero-crossing
        # trigger, did a timed event trigger, neither, or both?
        # terminate_early = terminate_early | zc_events.has_active_terminal
        logger.debug("Terminate early after major step: %s", terminate_early)
        end_reason = _determine_step_end_reason(
            triggered, terminate_early, int_tf, update_time
        )
        logger.debug("Major step end reason: %s", end_reason)

        # Conditionally activate timed events depending on whether the major step
        # ended as a result of a time trigger or zero-crossing event.
        timed_events = activate_timed_events(timed_events, end_reason)

        if self.major_step_callback:
            io_callback(self.major_step_callback, (), context.time)

        # T-113-followup-event-reprojection: project after the
        # continuous-integration phase.  When ``_advance_continuous_time``
        # ends on a localized ZC trigger, ``handle_events`` has just
        # applied the reset map and the algebraic states may have left
        # the manifold; projecting here re-establishes ``f_a = 0`` before
        # the next major step's discrete update sees the post-reset
        # state.  We project unconditionally — projection is a no-op
        # (zero Newton iterations, residual already below tol) on a
        # non-triggering step where state is already feasible, so the
        # extra cost on non-event steps is a single Newton residual
        # evaluation.  Default-off path skips the block at trace time.
        # No-op for systems without a mass matrix.  Runs before T-003a's
        # end-of-major-step projection so both hooks compose cleanly.
        if self.dae_reproject_after_events and getattr(
            self.system, "has_mass_matrix", False,
        ):
            from .dae_projection import project_constraints
            context = project_constraints(
                self.system,
                context,
                tol=self.dae_projection_tol,
                max_iter=self.dae_projection_max_iter,
            )

        # T-003a: opt-in DAE constraint projection.  Re-establishes
        # ``f_a(t, x, p) = 0`` after each major step by Newton-correcting
        # the algebraic component of the continuous state, holding the
        # differential component fixed.  No-op for systems without a mass
        # matrix or when the option is disabled (the default).
        if self.dae_projection_enabled and getattr(
            self.system, "has_mass_matrix", False,
        ):
            from .dae_projection import project_constraints
            context = project_constraints(
                self.system,
                context,
                tol=self.dae_projection_tol,
                max_iter=self.dae_projection_max_iter,
            )

        # T-132: declared per-block state projections (manifold
        # retractions, e.g. unit-quaternion renormalization) at the end
        # of every major step.  Static python loop over the (usually 0
        # or 1) declaring blocks; skipped entirely at trace time when no
        # block declares a projection.  Applied after the DAE projection
        # so both hooks compose; ordinary traced ops, so reverse-mode AD
        # flows through.
        if self._state_projection_leaves:

            def _project_component(proj, xc):
                # A leaf's context-level continuous state may arrive
                # wrapped as a single-element LeafStateComponent tuple /
                # list; the declared projection sees the state in its
                # ode-callback structure, so unwrap-and-rewrap.
                if isinstance(xc, (list, tuple)) and len(xc) == 1:
                    inner = _project_component(proj, xc[0])
                    return type(xc)((inner,))
                return proj(xc)

            for _sys in self._state_projection_leaves:
                if _sys is self.system:
                    context = context.with_continuous_state(
                        _project_component(
                            _sys._continuous_projection,
                            context.continuous_state,
                        )
                    )
                else:
                    _sub = context[_sys.system_id]
                    _sub = _sub.with_continuous_state(
                        _project_component(
                            _sys._continuous_projection,
                            _sub.continuous_state,
                        )
                    )
                    context = context.with_subcontext(_sys.system_id, _sub)

        # T-003b: opt-in DAE drift monitor.  Computes ``||f_a||_∞`` and
        # emits a ``UserWarning`` (via ``jax.debug.callback`` so it works
        # under jit/vmap) when above the threshold.  Default-off path is
        # byte-equivalent — the entire block is skipped at trace time
        # when ``dae_drift_threshold is None``.  Runs *after* projection
        # so the warning reflects the post-correction residual.  The
        # mask-and-max is done inline rather than via
        # ``constraint_residual_norm`` because boolean indexing of a
        # tracer is not jit-safe; ``jnp.where`` is.
        if self.dae_drift_threshold is not None and getattr(
            self.system, "has_mass_matrix", False,
        ):
            from .dae_drift import algebraic_row_mask
            mask_np = algebraic_row_mask(self.system)
            if mask_np is not None and mask_np.any():
                xcdot = self.system.eval_time_derivatives(context)
                xcdot_flat = jnp.concatenate(
                    [jnp.ravel(leaf) for leaf in jax.tree.leaves(xcdot)]
                )
                mask = jnp.asarray(mask_np)
                residual_max = jnp.max(jnp.where(mask, jnp.abs(xcdot_flat), 0.0))
                threshold = jnp.asarray(self.dae_drift_threshold)
                jax.debug.callback(
                    _emit_dae_drift_warning,
                    context.time,
                    residual_max,
                    threshold,
                )

        # T-113 Phase 1: opt-in per-major-step DAE drift trace.
        # Default-off path is byte-equivalent — entire block is skipped
        # at trace time when ``self._dae_drift_monitor is None``
        # (i.e. when ``record_dae_drift=False`` or the system has no
        # mass matrix).  Same residual computation as T-003b above; we
        # do not share the value because T-003b's block is itself
        # gated on ``dae_drift_threshold is not None``, and forcing
        # them to share would couple two independently-opt-in switches.
        # The cost of recomputing ``||f_a||_∞`` is one extra
        # ``eval_time_derivatives`` per major step, only paid when the
        # user opted in to the trace.  Runs *after* projection so the
        # trace reflects the post-correction residual (matching T-003b).
        if self._dae_drift_monitor is not None:
            from .dae_drift import algebraic_row_mask as _alg_mask
            mask_np = _alg_mask(self.system)
            if mask_np is not None and mask_np.any():
                xcdot = self.system.eval_time_derivatives(context)
                xcdot_flat = jnp.concatenate(
                    [jnp.ravel(leaf) for leaf in jax.tree.leaves(xcdot)]
                )
                mask = jnp.asarray(mask_np)
                residual_max = jnp.max(
                    jnp.where(mask, jnp.abs(xcdot_flat), 0.0)
                )
                jax.debug.callback(
                    self._dae_drift_monitor.update,
                    context.time,
                    residual_max,
                )

        # T-027a-followup: simulator-level Zeno tracker + recovery probe.
        # Default-off path: ``zeno_protection_enabled=False`` skips the
        # entire block at Python level — zero ops compiled in, the carry
        # is byte-equivalent to the pre-followup state.  When enabled,
        # ``_update_zeno_tracking`` returns the new ``(tprev, active,
        # frozen_steps)`` triple, including the recovery-probe behaviour:
        # after K=zeno_recovery_period consecutive frozen steps the latch
        # is cleared for one step so the next guard-trigger check
        # naturally re-engages Zeno if the cascade is still active, or
        # stays cleared otherwise.
        if self.zeno_protection_enabled:
            # T-027a-followup-per-leaf-freeze: apply the freeze gate
            # using the PRIOR-step latch (``sim_state.zeno_active``)
            # before updating it.  When a leaf's latch was engaged at
            # step N-1, step N's ODE/ZC advance is rolled back for that
            # leaf — the latch from step N-1 prevents continuous-state
            # evolution at step N, mirroring the leaf-level pattern
            # (``_wrap_ode_for_zeno`` zeros the rhs when the discrete
            # zeno flag is set).  Other leaves' post-ODE state are kept.
            context, ode_solver_state = self._apply_per_leaf_zeno_freeze(
                pre_ode_xc,
                context,
                pre_ode_solver_state,
                ode_solver_state,
                prior_zeno_active,
            )
            # Time-advance fix-up for the frozen-by-prior-step case:
            # ``_advance_continuous_time`` would have clamped ``int_tf``
            # to the just-fired ZC trigger time (essentially equal to
            # ``int_t0`` for a dense cascade), and ``context.time`` to
            # the same.  Spinning on that would never let the recovery
            # probe drain ``zeno_frozen_steps`` to ``K``.
            #
            # T-027a-followup-multi-leaf-cascade-bug (partial fix): the
            # prior gate used ``prior_any_active`` (any leaf frozen),
            # which in a multi-leaf diagram could force-advance global
            # ``int_tf`` past an unfrozen leaf's just-localised ZC
            # trigger time, dropping that leaf's pending event handling.
            # This gate is now ``prior_all_active`` over leaves with
            # continuous state — the only case where there is no
            # genuine ODE progress for any leaf and therefore no harm
            # in jumping global time forward.  When only SOME leaves
            # are frozen, the natural ``int_tf`` from
            # ``_advance_continuous_time`` is kept so the still-
            # bouncing leaf's slice retains the time it actually
            # integrated to; the per-leaf rollback above already
            # restored the frozen leaves' continuous state.  Single-
            # leaf and identical-leaves-cascading-simultaneously
            # cases collapse to ``all == any`` so behaviour is
            # unchanged for them.
            #
            # T-027a-followup-multi-leaf-cascade-architecture (candidate
            # (c), 2026-05-01): the original staggered-cascade failure
            # — leaf at post-reset rest (h=0, v≈0+) integrating under
            # gravity to h<0 with ``positive_then_non_positive`` guards
            # rejecting the trigger because ``w0=0`` is not strictly
            # positive — is now resolved by ``_apply_recovery_w0_nudge``
            # (direction-aware ``w0`` nudge), ``_extract_per_event_triggered``
            # (per-event tracker advance), and ``_mask_triggered_for_active_latch``
            # (suppress latched events' triggers inside the inner ODE
            # step so still-bouncing leaves can advance).
            int_tf_planned = npa.min(
                npa.array([int_boundary_time, int_tf_limit, update_time])
            )
            prior_active_b = jnp.asarray(prior_zeno_active, dtype=jnp.bool_)
            # Per-leaf frozen reduction restricted to leaves with
            # continuous state.  Built from the same static
            # ``_cs_idx_to_event_positions`` map used by the freeze
            # gate so the two stay in sync.  Leaves with NO events
            # (still continuous-state-bearing — e.g. an integrator
            # block in a diagram with a separate bouncing ball)
            # contribute ``False`` to the reduction so their
            # continued ODE progress also blocks the global time-
            # skip; this is the safe default because their state is
            # advancing and any global ``int_tf`` jump would lose it.
            if self._n_continuous_leaves > 0:
                _leaf_frozen_terms = []
                for cs_i in range(self._n_continuous_leaves):
                    positions = self._cs_idx_to_event_positions.get(cs_i, ())
                    if positions:
                        _leaf_frozen_terms.append(
                            jnp.any(prior_active_b[jnp.asarray(positions)]),
                        )
                    else:
                        _leaf_frozen_terms.append(jnp.asarray(False))
                prior_all_active = jnp.all(jnp.stack(_leaf_frozen_terms))
            else:
                # No continuous-state leaves: fall back to the prior
                # ``any_active`` semantics (no rollback to align with).
                prior_all_active = jnp.asarray(prior_any_active, dtype=jnp.bool_)
            should_skip = prior_all_active & (
                jnp.asarray(int_tf, dtype=int_tf_planned.dtype) < int_tf_planned
            )
            int_tf = jnp.where(should_skip, int_tf_planned, int_tf)
            context = context.with_time(IntegerTime.as_decimal(int_tf))
            # When the freeze fully skipped the ODE-advance side
            # effects, the post-ODE ``triggered`` is masked False so
            # the next call to ``_update_zeno_tracking`` sees no firing
            # event.  This lets the recovery probe drain to K and clear
            # the latch so ``int_tf`` resumes ZC-localised stepping.
            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # use the per-event triggered mask plumbed back from
            # ``_advance_continuous_time`` so each event's tracker
            # update fires only for events that actually triggered on
            # this major step.  Falls back to the legacy scalar broadcast
            # when the per-event mask is unavailable (e.g. no
            # continuous-state path or no events in the system).
            per_event_triggered_in = cdata.per_event_triggered
            if per_event_triggered_in is None:
                triggered_b_for_tracker = jnp.asarray(triggered, dtype=jnp.bool_)
            else:
                triggered_b_for_tracker = jnp.asarray(
                    per_event_triggered_in, dtype=jnp.bool_,
                )
            triggered_for_tracker = jnp.where(
                should_skip,
                jnp.zeros_like(triggered_b_for_tracker),
                triggered_b_for_tracker,
            )
            zeno_tprev, zeno_active, zeno_frozen_steps = (
                self._update_zeno_tracking(
                    sim_state.zeno_tprev,
                    sim_state.zeno_active,
                    sim_state.zeno_frozen_steps,
                    triggered_for_tracker,
                    context.time,
                )
            )
            # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
            # derive the per-event "recovery probe just fired" mask from
            # the latch transition.  ``should_probe`` is the only path
            # in ``_update_zeno_tracking`` that flips ``active[i]``
            # True->False on a single step; engagement only flips
            # False->True.  So ``prior_active & ~new_active`` is exactly
            # the recovery-probe-fires set.  This mask is consumed by
            # the NEXT major step's ``_advance_continuous_time`` to
            # apply the direction-aware ``w0`` nudge.  The mask self-
            # clears on the step after that because the latch is
            # already cleared (no new ``should_probe`` firing without
            # re-engagement).
            prior_active_b_for_rjc = jnp.asarray(
                sim_state.zeno_active, dtype=jnp.bool_,
            )
            new_active_b_for_rjc = jnp.asarray(zeno_active, dtype=jnp.bool_)
            zeno_recovery_just_cleared = (
                prior_active_b_for_rjc & ~new_active_b_for_rjc
            )
        else:
            zeno_tprev = sim_state.zeno_tprev
            zeno_active = sim_state.zeno_active
            zeno_frozen_steps = sim_state.zeno_frozen_steps
            zeno_recovery_just_cleared = sim_state.zeno_recovery_just_cleared

        return SimulatorState(
            step_end_reason=end_reason,
            context=context,
            timed_events=timed_events,
            int_time=int_tf,
            results_data=results_data,
            ode_solver_state=ode_solver_state,
            zeno_tprev=zeno_tprev,
            zeno_active=zeno_active,
            zeno_frozen_steps=zeno_frozen_steps,
            zeno_recovery_just_cleared=zeno_recovery_just_cleared,
        )

    # This method is marked private because it will be wrapped with a custom autodiff
    # rule to get the correct derivatives with respect to the end time of the
    # simulation interval using `_override_advance_to_vjp`.  This also copies the
    # docstring to the overridden function. Normally the wrapped attribute `advance_to`
    # is what should be called by users.
    def _advance_to(self, boundary_time: float, context: ContextBase) -> SimulatorState:
        """Core control flow logic for running a simulation.

        This is the main loop for advancing the simulation.  It is called by `simulate`
        or can be called directly if more fine-grained control is needed. This method
        essentially loops over "major steps" until the boundary time is reached. See
        the documentation for `simulate` for details on the order of operations in a
        major step.

        Args:
            boundary_time (float): The time to advance to.
            context (ContextBase): The current state of the system.

        Returns:
            SimulatorState:
                A named tuple containing the final state of the simulation, including
                the final context, a collection of pending timed events, and a flag
                indicating the reason that the most recent major step ended.

        Notes:
            API will change slightly as a result of WC-87, which will break out the
            initialization from the main loop so that `advance_to` can be called
            repeatedly.  See:
            https://jaxonomy.atlassian.net/browse/WC-87
        """

        system = self.system
        sim_state = self.initialize(context)
        end_reason = sim_state.step_end_reason
        context = sim_state.context
        timed_events = sim_state.timed_events
        int_boundary_time = IntegerTime.from_decimal(boundary_time)

        # We will be limiting each step by the max_major_step_length.  However, if this
        # is infinite we should just use the end time of the simulation to avoid
        # integer overflow.  This could be problematic if the end time of the
        # simulation is close to the maximum representable integer time, but we can come
        # back to that if it's an issue.
        int_max_step_length = IntegerTime.from_decimal(
            npa.minimum(self.max_major_step_length, boundary_time)
        )

        # Only activate timed events if the major step ended on a time trigger
        timed_events = activate_timed_events(timed_events, end_reason)

        # Called on the "True" branch of the conditional
        def _major_step(sim_state: SimulatorState) -> SimulatorState:
            return self._major_step(sim_state, int_boundary_time, int_max_step_length)

        def _cond_fun(sim_state: SimulatorState):
            return (sim_state.int_time < int_boundary_time) & (
                sim_state.step_end_reason != StepEndReason.TerminalEventTriggered
            )

        # Initialize the "carry" values for the main loop.
        if self.system.has_continuous_state:
            leaves = jax.tree.leaves(context.continuous_state)
            dtype = leaves[0].dtype if leaves else jnp.empty(0).dtype
            context = context.with_time(jnp.asarray(context.time, dtype=dtype))
            if sim_state.ode_solver_state is not None:
                def _cast_leaf(x):
                    if isinstance(x, (float, np.floating)):
                        return jnp.asarray(x, dtype=dtype)
                    if isinstance(x, jnp.ndarray) and jnp.issubdtype(x.dtype, jnp.floating):
                        return x.astype(dtype)
                    return x
                ode_solver_state = jax.tree.map(_cast_leaf, sim_state.ode_solver_state)
            else:
                ode_solver_state = sim_state.ode_solver_state
        else:
            ode_solver_state = sim_state.ode_solver_state

        sim_state = SimulatorState(
            context=context,
            timed_events=timed_events,
            step_end_reason=end_reason,
            int_time=sim_state.int_time,
            results_data=sim_state.results_data,
            ode_solver_state=ode_solver_state,
            zeno_tprev=sim_state.zeno_tprev,
            zeno_active=sim_state.zeno_active,
            zeno_frozen_steps=sim_state.zeno_frozen_steps,
            zeno_recovery_just_cleared=sim_state.zeno_recovery_just_cleared,
        )

        logger.debug(
            "Running simulation from t=%s to t=%s", context.time, boundary_time
        )

        try:
            # Main loop call
            sim_state = self.while_loop(_cond_fun, _major_step, sim_state)
            logger.debug("Simulation complete at t=%s", sim_state.context.time)
        except KeyboardInterrupt:
            # NOTE: flag simulation as interrupted somewhere in sim_state
            logger.info("Simulation interrupted at t=%s", sim_state.context.time)

        # At the end of the simulation we need to handle any pending discrete updates
        # and store the solution one last time.
        # NOTE (WC-87): The returned simulator state can't be used with advance_to again,
        # since the discrete updates have already been performed. Should be broken out
        # into a `finalize` method as part of WC-87.

        # update discrete state to x+ at the simulation end_time
        if self.results_recorder.save_time_series:
            logger.debug("Finalizing solution...")
            # 1] do discrete update (will skip if the simulation was terminated early)
            context, _terminate_early = self._handle_discrete_update(
                sim_state.context, sim_state.timed_events
            )
            # 2] do update solution
            context = context.refresh_port_cache()
            # T-012a-followup: pass solver_state so the final sample also
            # snapshots the interpolant covering [t_prev, t_end].
            results_data = self.results_recorder.save(
                sim_state.results_data,
                context,
                ode_solver_state=sim_state.ode_solver_state,
            )
            sim_state = sim_state._replace(
                context=context,
                results_data=results_data,
            )
            logger.debug("Done finalizing solution")

        return sim_state

__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

Source code in jaxonomy/simulation/simulator.py

def __init__(
    self,
    system: SystemBase,
    ode_solver: ODESolverBase = None,
    options: SimulatorOptions = None,
):
    """Initialize the simulator.

    Args:
        system (SystemBase): The hybrid dynamical system to simulate.
        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.
        options (SimulatorOptions):
            Options for the simulation process.  See `simulate` for details.
    """
    self.system = system

    if options is None:
        options = SimulatorOptions()

    # Determine whether JAX tracing can be used (jit, grad, vmap, etc)
    math_backend, self.enable_tracing = _check_backend(options)

    # Set the math backend
    set_backend(math_backend)

    # Should the simulation be run with autodiff enabled?  This will override
    # the `advance_to` method with a custom autodiff rule.
    self.enable_autodiff = options.enable_autodiff

    if ode_solver is None:
        ode_solver = ODESolver(system, options=options.ode_options)

    # Store configuration options
    self.max_major_steps = options.max_major_steps
    # Honor max_major_steps as a hard cap whenever it was explicitly provided
    # (either via _check_simulate_options or directly by the caller).
    self._explicit_max_major_steps = options._explicit_max_major_steps or (
        options.max_major_steps is not None and options.max_major_steps > 0
    )
    self.max_major_step_length = options.max_major_step_length
    self.zc_bisection_loop_count = options.zc_bisection_loop_count
    self.major_step_callback = options.major_step_callback

    # T-003a: opt-in DAE constraint projection at the end of each major step.
    self.dae_projection_enabled = getattr(
        options, "dae_projection_enabled", False,
    )
    self.dae_projection_tol = getattr(options, "dae_projection_tol", 1e-8)
    self.dae_projection_max_iter = getattr(
        options, "dae_projection_max_iter", 3,
    )

    # T-113-followup-event-reprojection: opt-in projection immediately
    # after discrete-event resets *within* a major step (top-of-step
    # ``_handle_discrete_update`` and triggered ZC resets inside
    # ``_advance_continuous_time``).  Default ``False`` so the hot
    # path is byte-equivalent to the pre-followup code.  Reuses the
    # T-003a tolerance / max-iter knobs to avoid surface-area churn.
    self.dae_reproject_after_events = getattr(
        options, "dae_reproject_after_events", False,
    )

    # T-003b: opt-in DAE drift threshold (post-projection check).
    # ``None`` (default) disables the check entirely — no overhead.
    self.dae_drift_threshold = getattr(
        options, "dae_drift_threshold", None,
    )

    # T-132: declared per-block state projections
    # (``declare_continuous_state(project=...)``, e.g. unit-quaternion
    # renormalization).  Collected once, statically — when empty (the
    # default), ``_major_step`` skips the block at trace time and the
    # hot path is byte-equivalent.
    _leaves = getattr(system, "leaf_systems", None)
    if _leaves is None:
        _leaves = [system]
    self._state_projection_leaves = [
        s
        for s in _leaves
        if getattr(s, "_continuous_projection", None) is not None
    ]

    # T-113 Phase 1: opt-in per-major-step DAE drift trace.
    # ``False`` (default) disables the trace entirely — no monitor
    # constructed and the simulator's ``_major_step`` skips the
    # trace block at trace time.  When True AND the system has a
    # mass matrix, attach a host-side ``_DAEDriftMonitor`` so the
    # trace block forwards each major step's ``(time, residual)``
    # via ``jax.debug.callback``.  Non-DAE systems get no monitor
    # (the diagnostic is mass-matrix-specific by definition).
    self.record_dae_drift = getattr(options, "record_dae_drift", False)
    self._dae_drift_monitor: _DAEDriftMonitor | None = None
    if self.record_dae_drift and getattr(
        self.system, "has_mass_matrix", False,
    ):
        self._dae_drift_monitor = _DAEDriftMonitor()

    # T-125-followup-record-event-times: opt-in capture of zero-
    # crossing event firing times.  Construct a host-side recorder
    # only when the option is True AND the diagram has at least one
    # zero-crossing event — diagrams with no events get no recorder
    # so ``_advance_continuous_time`` short-circuits the callback at
    # trace time, preserving the byte-equivalent default-off path.
    # ``n_zero_crossing_events`` is set further down in __init__ but
    # the snapshot below makes the count available without re-
    # importing the system; we settle for re-counting cheaply here
    # to avoid reordering the existing init blocks.
    self.record_event_times = getattr(options, "record_event_times", False)
    self._event_time_recorder: _EventTimeRecorder | None = None
    if self.record_event_times:
        n_zc = len(system.zero_crossing_events.events)
        if n_zc > 0:
            self._event_time_recorder = _EventTimeRecorder(n_zc)

    # T-038a-followup-bdf-condition-check: opt-in BDF Newton
    # condition-number diagnostic.  When the threshold is set AND
    # the active solver is a BDF solver, attach a
    # ``_BDFConditionMonitor`` to it as a side-channel — the BDF
    # solver checks ``getattr(self, "_cond_monitor", None)`` inside
    # ``newton_iteration`` and forwards the cond estimate via
    # ``jax.debug.callback`` only when set.  Default-off path is
    # byte-equivalent (no monitor → no-op in BDF).  Non-BDF
    # solvers silently ignore the option (the diagnostic is BDF-
    # specific by definition).
    self.bdf_condition_warning_threshold = getattr(
        options, "bdf_condition_warning_threshold", None,
    )
    self._bdf_cond_monitor: _BDFConditionMonitor | None = None
    if self.bdf_condition_warning_threshold is not None:
        # Only attach to BDF — non-BDF solvers don't have a Newton
        # iteration to monitor, and we don't want to silently
        # mislead users who set the option on a non-BDF run.
        try:
            from ..backend._jax.bdf import BDFSolver as _BDFSolver
        except Exception:  # pragma: no cover — import-time guard only
            _BDFSolver = None
        if _BDFSolver is not None and isinstance(ode_solver, _BDFSolver):
            self._bdf_cond_monitor = _BDFConditionMonitor(
                self.bdf_condition_warning_threshold,
            )
            # Attach as an instance attribute so the BDF solver
            # picks it up via ``getattr(self, "_cond_monitor", None)``
            # without needing a constructor change.
            ode_solver._cond_monitor = self._bdf_cond_monitor

    # Detailed non-finite abort diagnostics: stamp the opt-in flag on
    # the BDF solver (checked at trace time; default path compiles no
    # callback ops).
    if getattr(options, "bdf_nonfinite_diagnostics", False):
        try:
            from ..backend._jax.bdf import BDFSolver as _BDFSolver2
        except Exception:  # pragma: no cover — import-time guard only
            _BDFSolver2 = None
        if _BDFSolver2 is not None and isinstance(ode_solver, _BDFSolver2):
            ode_solver._nonfinite_diagnostics = True

    # T-027a-followup: simulator-level Zeno protection options.  All
    # default-off — the recovery probe and the latch are skipped at
    # Python level when ``zeno_protection_enabled=False``, keeping the
    # default hot path byte-equivalent.
    self.zeno_protection_enabled = getattr(
        options, "zeno_protection_enabled", False,
    )
    self.zeno_tolerance = getattr(options, "zeno_tolerance", 1e-6)
    self.zeno_recovery_period = getattr(
        options, "zeno_recovery_period", 10,
    )
    # T-027a-followup-vector-tprev: count zero-crossing events at
    # construction time so ``initialize`` can allocate per-event
    # ``zeno_tprev`` / ``zeno_active`` vectors of the correct shape.
    # ``event_system_ids`` records each event's owning leaf — kept
    # for the per-leaf freeze gate (T-027a-followup-per-leaf-freeze).
    # Both are static (system topology doesn't change at runtime),
    # so this is a one-shot pass at __init__.
    zc_events_static = system.zero_crossing_events.events
    self.n_zero_crossing_events = len(zc_events_static)
    self.event_system_ids = tuple(
        getattr(ev, "system_id", None) for ev in zc_events_static
    )
    # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
    # snapshot each zero-crossing event's static ``direction`` string
    # so the per-event recovery-probe nudge in ``_apply_recovery_w0_nudge``
    # can pick the right side of the threshold to push ``w0`` to.
    # Mirrors ``event_system_ids`` — same order as
    # ``system.zero_crossing_events.events`` and the per-event Zeno
    # carry vectors built in ``initialize``.
    self.event_directions = tuple(
        getattr(ev, "direction", "crosses_zero") for ev in zc_events_static
    )

    # T-027a-followup-per-leaf-freeze: build a static map from each
    # leaf's ``system_id`` to its index in
    # ``DiagramContext.continuous_state`` (which is a list ordered by
    # the iteration of subcontexts that have continuous state).  This
    # lets ``_major_step`` know which slot of the list to roll back
    # when an event owned by that leaf has its Zeno latch engaged.
    # For single-LeafSystem simulations the list collapses to one
    # entry (``LeafContext.continuous_state`` is a single Array, not
    # a list — handled separately at the freeze site).
    if isinstance(system, Diagram):
        _leaves = list(system.leaf_systems)
    else:
        _leaves = [system]
    self._sysid_to_cs_idx = {}
    _cs_idx = 0
    for _leaf in _leaves:
        if getattr(_leaf, "has_continuous_state", False):
            self._sysid_to_cs_idx[_leaf.system_id] = _cs_idx
            _cs_idx += 1
    self._n_continuous_leaves = _cs_idx
    # Reverse map: cs_index -> tuple of event-vector positions whose
    # ``zeno_active[i]`` should freeze that leaf.  Built once so the
    # per-step freeze logic is a flat scatter, no per-step Python
    # iteration over event_system_ids.
    self._cs_idx_to_event_positions: dict[int, tuple[int, ...]] = {}
    for i, sid in enumerate(self.event_system_ids):
        cs_i = self._sysid_to_cs_idx.get(sid)
        if cs_i is None:
            continue
        self._cs_idx_to_event_positions.setdefault(cs_i, []).append(i)
    self._cs_idx_to_event_positions = {
        k: tuple(v) for k, v in self._cs_idx_to_event_positions.items()
    }

    # T-027a-followup-per-leaf-solver-state: compute each continuous
    # leaf's flat slice into the raveled ODE state vector ``y`` so the
    # per-leaf freeze gate can decompose ``Dopri5State.{y, f,
    # interp_coeff}`` (and ``BDFState.{y, f, D}``) along the last
    # axis.  The flat layout matches ``ravel_pytree(context.
    # continuous_state)`` exactly: leaves with continuous state are
    # iterated in ``DiagramContext.continuous_subcontexts`` order
    # (subcontexts.values() filtered by has_continuous_state), which
    # is the same order as ``Diagram.leaf_systems`` filtered by
    # ``has_continuous_state``.  Each leaf's flat size is the sum of
    # its ``_default_continuous_state`` pytree-leaf sizes — usually
    # a single Array, but pytree-valued continuous states are also
    # handled.  ``_leaf_flat_slices`` is a tuple of ``(start, end)``
    # int pairs ordered by ``cs_idx``; total length is the flat
    # ODE state dimension ``_n_y_total``.  Default-off path
    # (``zeno_protection_enabled=False``) never reads these; they
    # are purely metadata.
    self._leaf_flat_slices: tuple[tuple[int, int], ...] = ()
    self._n_y_total: int = 0
    if self._n_continuous_leaves > 0:
        _slices: list[tuple[int, int]] = []
        _offset = 0
        # Re-scan ``_leaves`` in the same order used for ``_sysid_to_cs_idx``.
        for _leaf in _leaves:
            if not getattr(_leaf, "has_continuous_state", False):
                continue
            _xc0 = getattr(_leaf, "_default_continuous_state", None)
            if _xc0 is None:
                _size = 0
            else:
                _size = int(sum(
                    int(np.prod(np.shape(_l))) if np.shape(_l) else 1
                    for _l in jax.tree_util.tree_leaves(_xc0)
                ))
            _slices.append((_offset, _offset + _size))
            _offset += _size
        self._leaf_flat_slices = tuple(_slices)
        self._n_y_total = _offset

    # T-013a-followup-mode-a-buffers: when the user opts into the
    # "buffers" mode, classify each recorded signal's cadence
    # statically here and pass the result through to the recorder.
    # The classification is reused at the per-step decision in
    # ``JaxResultsData.update`` to skip writes for unfired periodic
    # signals.  Default ``"auto"`` does NOT enable buffers — it
    # remains the post-finalize schedule trim path.
    psts_mode = getattr(options, "per_signal_timestamps_mode", "auto")
    psts_enabled = getattr(options, "per_signal_timestamps", False)
    per_signal_buffers_classifications = None
    if (
        psts_enabled
        and psts_mode == "buffers"
        and options.recorded_signals is not None
    ):
        per_signal_buffers_classifications = (
            ResultsRecorder.classify_signal_cadence(options.recorded_signals)
        )

    # T-012a-followup: thread record_solver_states through to the
    # recorder so the JaxResultsData allocates a per-step interpolant
    # ring and ``save`` snapshots ``Dopri5State.interp_coeff`` per
    # call.  Default-off path is byte-equivalent.
    self.record_solver_states = getattr(
        options, "record_solver_states", False,
    )
    # T-002b-followup-buffer-overflow-auto-size — when the user
    # constructs ``Simulator`` directly (bypassing ``simulate``), the
    # ``_check_options`` auto-sizing path is skipped, so ``options.
    # buffer_length`` may still be ``None``. Fall back to
    # ``max_major_steps`` (the natural cap) or a legacy 1000-sample
    # default when neither is available.
    if options.buffer_length is not None:
        recorder_buffer_length = options.buffer_length
    elif self.max_major_steps is not None and self.max_major_steps > 0:
        recorder_buffer_length = max(
            int(self.max_major_steps), _MIN_AUTO_BUFFER_LENGTH
        )
    else:
        recorder_buffer_length = _MIN_AUTO_BUFFER_LENGTH
    self.results_recorder = ResultsRecorder(
        save_time_series=options.save_time_series,
        recorded_outputs=options.recorded_signals,
        buffer_length=recorder_buffer_length,
        per_signal_buffers_classifications=per_signal_buffers_classifications,
        record_solver_states=self.record_solver_states,
    )

    # Zero-crossing handler encapsulates guard evaluation and bisection logic
    self.zc_handler = ZeroCrossingHandler(
        system,
        self.zc_bisection_loop_count,
        lower_triangular_discrete_update=getattr(
            options, "lower_triangular_discrete_update", False,
        ),
    )

    if self.max_major_step_length is None:
        self.max_major_step_length = np.inf

    logger.debug("Simulator created with enable_tracing=%s", self.enable_tracing)

    self.ode_solver = ode_solver

    # T-113-followup-baumgarte-and-ssp: opt-in Baumgarte stabilization.
    # When ``baumgarte_alpha`` and/or ``baumgarte_beta`` are set, wrap
    # the solver's ``ode_rhs`` to add ``-2α·ġ - β²·g`` to the
    # algebraic rows of the rhs.  ``baumgarte_augment_ode_rhs`` is a
    # no-op (returns the input rhs unchanged) when both gains are
    # ``None`` or when the system has no algebraic constraints — the
    # disabled hot path's JIT trace graph is byte-equivalent to the
    # pre-followup behaviour.  Wraps before any ``ode_solver.initialize``
    # call so ``flat_ode_rhs = ravel_first_arg(self.ode_rhs, ...)`` in
    # the JAX impl picks up the augmented version.
    b_alpha = getattr(options, "baumgarte_alpha", None)
    b_beta = getattr(options, "baumgarte_beta", None)
    if (b_alpha is not None or b_beta is not None) and getattr(
        self.system, "has_mass_matrix", False,
    ):
        from .dae_projection import baumgarte_augment_ode_rhs
        ode_solver.ode_rhs = baumgarte_augment_ode_rhs(
            ode_solver.ode_rhs, self.system, b_alpha, b_beta,
        )

    from .autodiff_rules import make_advance_to_vjp, make_guarded_integrate_vjp
    # Modify the default autodiff rule slightly to correctly capture variations
    # in end time of the simulation interval.
    self.has_terminal_events = system.zero_crossing_events.has_terminal_events
    # T-006: wrap advance_to so direct callers (not going through
    # simulate()) also get JAX-error remapping with block/port context.
    # T-A2-followup-advance-to-jit-cache: jit the inner advance_to so a
    # *persistent* Simulator (construct once, call ``advance_to`` many
    # times — interactive stepping, MPC inner loops) reuses the compiled
    # kernel instead of re-tracing op-by-op on every call. The jit is a
    # stable instance attribute, so JAX's cache hits across calls with the
    # same context aval. Only the non-autodiff path is wrapped: the
    # autodiff path returns a ``custom_vjp`` callable that ``simulate``
    # already jits at the outer ``_wrapped_simulate`` boundary, and we
    # keep ``remap_simulation_errors`` on the *outside* so runtime errors
    # are still remapped at the call boundary (not just at trace time).
    _advance_to_impl = make_advance_to_vjp(self)
    if self.enable_tracing and not self.enable_autodiff:
        _advance_to_impl = jax.jit(_advance_to_impl)
    self.advance_to = remap_simulation_errors(_advance_to_impl)

    # Also override the guarded ODE integration with a custom autodiff rule
    # to capture variations due to zero-crossing time.
    self.guarded_integrate = make_guarded_integrate_vjp(self)

compile(tf, context)

Warm up / pre-compile the simulation advance_to method on the device.

Source code in jaxonomy/simulation/simulator.py

def compile(self, tf: float, context: ContextBase):
    """Warm up / pre-compile the simulation advance_to method on the device."""
    if self.enable_tracing and not self.enable_autodiff:
        self.advance_to(tf, context)

initialize(context)

Perform initial setup for the simulation.

Source code in jaxonomy/simulation/simulator.py

def initialize(self, context: ContextBase) -> SimulatorState:
    """Perform initial setup for the simulation."""
    logger.debug("Initializing simulator")
    # context.state.pprint(logger.debug)

    # Initial simulation time as integer (picoseconds)
    initial_int_time = IntegerTime.from_decimal(context.time)

    # Ensure that _next_update_time() can return the current time by perturbing
    # current time as slightly toward negative infinity as possible
    time_of_next_timed_event, timed_events = _next_update_time(
        self.system.periodic_events, initial_int_time - 1
    )

    # timed_events is now marked with the active events at the next update time
    logger.debug("Time of next timed event (int): %s", time_of_next_timed_event)
    logger.debug(
        "Time of next event (sec): %s",
        IntegerTime.as_decimal(time_of_next_timed_event),
    )
    timed_events.pprint(logger.debug)

    end_reason = npa.where(
        time_of_next_timed_event == initial_int_time,
        StepEndReason.TimeTriggered,
        StepEndReason.NothingTriggered,
    )

    # Initialize the results data that will hold recorded time series data.
    results_data = self.results_recorder.initialize(context)

    # T-027a-followup-vector-tprev: when simulator-level Zeno
    # protection is enabled, allocate per-event ``zeno_tprev`` /
    # ``zeno_active`` vectors so each event tracks its own last-
    # firing time independently.  ``zeno_tprev`` is initialised to
    # ``-inf`` so the first firing is never inside tolerance.  When
    # disabled, leave the carry as the scalar defaults from the
    # ``SimulatorState`` declaration so the default-off path's
    # pytree is byte-equivalent.
    if self.zeno_protection_enabled:
        n = max(self.n_zero_crossing_events, 1)
        zeno_tprev = jnp.full((n,), -jnp.inf)
        zeno_active = jnp.zeros((n,), dtype=jnp.bool_)
        # T-027a-followup-per-event-recovery: ``zeno_frozen_steps``
        # vectorises to ``(N_events,)`` so each event independently
        # counts its own consecutive-frozen-step streak.  When event
        # ``i`` hits ``zeno_recovery_period``, only its own latch
        # clears; other events keep cascading.  Default-off path
        # keeps the scalar default in ``SimulatorState`` so the
        # disabled pytree shape is byte-equivalent.
        zeno_frozen_steps = jnp.zeros((n,), dtype=jnp.int32)
        # T-027a-followup-multi-leaf-cascade-architecture (candidate (c)):
        # per-event mask of "the previous major step's recovery probe
        # just fired for this event".  Initialised to all-False so the
        # first ODE step has no nudge applied.  Shape matches the
        # other per-event carry vectors so the elementwise compare/
        # scatter inside ``_apply_recovery_w0_nudge`` aligns.
        zeno_recovery_just_cleared = jnp.zeros((n,), dtype=jnp.bool_)
        return SimulatorState(
            context=context,
            timed_events=timed_events,
            step_end_reason=end_reason,
            int_time=initial_int_time,
            results_data=results_data,
            ode_solver_state=self.ode_solver.initialize(context),
            zeno_tprev=zeno_tprev,
            zeno_active=zeno_active,
            zeno_frozen_steps=zeno_frozen_steps,
            zeno_recovery_just_cleared=zeno_recovery_just_cleared,
        )

    return SimulatorState(
        context=context,
        timed_events=timed_events,
        step_end_reason=end_reason,
        int_time=initial_int_time,
        results_data=results_data,
        ode_solver_state=self.ode_solver.initialize(context),
    )

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.

Source code in jaxonomy/simulation/simulator.py

def while_loop(self, 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.
    """
    use_bounded = self.enable_autodiff or self._explicit_max_major_steps
    if use_bounded:
        return _bounded_while_loop(cond_fun, body_fun, val, self.max_major_steps)
    else:
        return backend.while_loop(cond_fun, body_fun, val)

SimulatorOptions dataclass

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.

Source code in jaxonomy/simulation/types.py

@dataclasses.dataclass
class 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.
    """

    math_backend: str = dataclasses.field(
        default_factory=lambda: numpy_api.active_backend
    )
    enable_tracing: bool = True
    enable_autodiff: bool = False
    precision: str = "auto"  # "auto", "float32", or "float64"

    # diff-mode follow-up: explicit, self-documenting differentiation-mode
    # selector. ``enable_autodiff`` conflates two things — "make the sim
    # differentiable at all" and "install the reverse-mode adjoint" — which
    # makes forward-mode autodiff counterintuitive: it requires
    # ``enable_autodiff=False`` so JAX traces the real solver ops (the
    # reverse-mode ``custom_vjp`` intercepts ``jax.jacfwd`` / ``jvp`` and does
    # NOT forward-differentiate the solver). Set ``diff_mode`` instead of
    # reasoning about the boolean; it resolves into ``enable_autodiff`` (the
    # canonical flag every downstream site reads):
    #   - "reverse"   → reverse-mode adjoint (``jax.grad`` / ``jacrev``);
    #                   sets ``enable_autodiff=True``.
    #   - "forward"   → forward-mode (``jax.jacfwd`` / ``jvp``); sets
    #                   ``enable_autodiff=False`` so the real ops are traced.
    #   - "none"      → no autodiff; sets ``enable_autodiff=False``.
    #   - "auto"/None → leave ``enable_autodiff`` as given (back-compat
    #                   default).
    # Resolved once in ``__post_init__`` and then cleared back to ``None`` so a
    # later ``dataclasses.replace(enable_autodiff=...)`` is never re-clobbered.
    diff_mode: str | None = None

    # If autodiff is enabled, max_major_steps must be set in order to bound the number
    # of iterations in the while loop.  When running a simulation using the `simulate`
    # function, this can typically be determined automatically based on the number of
    # periodic events in the system.  However, it should be specified manually in the
    # following cases:
    #   - When running a simulation by creating a `Simulator` object and calling the
    #     `advance_to` method directly. In this case the `Simulator` object does not
    #     attempt to automatically determine a bound on the number of major steps.
    #   - When autodiff is used to compute the sensitivity with respect to simulation
    #     end time, for example when computing periodic limit cycles. In this case the
    #     time variables passed to `estimate_max_major_steps` are JAX tracers and cannot
    #     be used to determine a fixed (static) bound on the number of major steps.
    #   - When the system has frequent zero-crossing events.  In this case the "safety
    #     factor" in the heuristic for estimating the number of major steps may be too
    #     small, underestimating the bound on the number of major steps.
    # In any case, `estimate_max_major_steps` can still be called statically ahead
    # of time to determine a reasonable value for `max_major_steps`, using for instance
    # a conservative bound on end time and safety factor.
    max_major_steps: int = None
    # NOTE (T-A3): ``max_major_step_length`` is *JIT-static* — it participates
    # in deriving ``max_major_steps`` (the bounded-loop trip count), which is a
    # Python int baked into the compiled kernel. Changing it between calls to a
    # jitted ``simulate`` / ``advance_to`` therefore forces a recompile. For
    # parameter sweeps, hold it fixed and vary traced quantities (initial state,
    # dynamic parameters) instead.
    max_major_step_length: float = None
    # T-A3 follow-up: ``max_major_step_size`` is an accepted alias for
    # ``max_major_step_length`` — it reads more naturally next to
    # ``max_minor_step_size`` and is the spelling careful users reach for first.
    # Reconciled in ``__post_init__``: if only one is set it populates the
    # other; setting both to conflicting values raises.
    max_major_step_size: float = None

    # Length of the recording ring buffer for the time series.  When ``None``
    # (the default), ``_check_options`` auto-sizes this to ``max_major_steps``.
    #
    # IMPORTANT (T-B3/B8): the recorder saves one sample per *accepted minor
    # (solver) step*, not per major step. Adaptive solvers — Dopri5 and the
    # "auto" default — take many minor steps per major step, and *more* of them
    # as ``rtol`` / ``atol`` tighten. A tight-tolerance Dopri5 run can therefore
    # record far more samples than the major-step-derived auto size, overrunning
    # the ring buffer and silently dropping the *head* of the trajectory
    # (``results.time`` then starts mid-run). The simulator detects this after
    # the fact and emits a loud, solver/tolerance-aware ``UserWarning``
    # recommending a concrete larger ``buffer_length``. To avoid it up front:
    # set ``buffer_length`` explicitly for long fine-grained recordings, loosen
    # the tolerances, or use the fixed-step ``ode_solver_method="rk4"`` (whose
    # sample count is predictable from ``max_minor_step_size``). Set a small
    # fixed N for memory-constrained streaming.
    buffer_length: int | None = None

    # ODE solver options
    ode_solver_method: str = "auto"  # Dopri5 (jax/scipy) or BDF (jax)
    rtol: float = 1e-6  # Relative tolerance for adaptive solvers
    atol: float = 1e-8  # Absolute tolerance for adaptive solvers
    min_minor_step_size: float = None
    max_minor_step_size: float = None

    # This is used to bound the number of "checkpoints" in the adjoint solver and
    # is used only when autodiff is enabled.  Increasing this may improve the
    # accuracy of the adjoint solver (especially over long integration times), but
    # will also increase memory usage.  Whether or not the resulting adjoint solve
    # is faster depends on the details of the problem, for instance on the number of
    # major steps and the ODE solver tolerance.  This can also be set to None to
    # disable checkpointing altogether.
    max_checkpoints: int = 16

    # 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.
    # NOTE: remove this and use `recorded_signals` instead. There are usecases
    # where simulate() is not used and we use the Simulator's advance_to function
    # directly. In those cases, recorded_signals can be set while save_time_series
    # is False which is confusing.
    save_time_series: bool = False

    # 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.
    recorded_signals: dict[str, SystemCallback] = None

    # 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.
    return_context: bool = True

    # Validate the diagram before simulating to check for common errors and unsupported
    # feature interactions like autodiff through python-only blocks.
    validate: bool = True

    # Zero crossings are localized in time using the ODE solver interpolant,
    # which provides state values for any time value in the previous integration
    # time interval.
    # Bisection is used to search the time interval. Rather than run bisection
    # in a while loop until the time interval is _small_, bisection is run for
    # fixed number of iterations, as this results in localizing zero crossings in
    # time within some small fraction of the integrated time interval.
    # e.g. if the major step length is 1.0 second, and bisection is run for 40
    # loops, the zero crossing time tolerance is approx. 1e-12, a.k.a. picosecond.
    zc_bisection_loop_count: int = 40

    # Scale of integer time used for event synchronization.
    #   - "auto" (default): pick the finest power-of-ten scale that still
    #     represents ``t_span[1]`` with headroom. Short simulations keep
    #     picosecond resolution (1e-12, max ~0.3 years); longer horizons
    #     transparently coarsen (1e-9 ns, 1e-6 µs, ...) so a multi-year
    #     simulation just runs instead of raising a representability error.
    #   - a float (e.g. 1e-9): pin the scale explicitly.
    #   - None: leave the global IntegerTime scale untouched (legacy escape
    #     hatch; not recommended — relies on process-global state).
    int_time_scale: float | str | None = "auto"

    # Called at the end of each major step with the current time as an argument.
    major_step_callback: Callable[[Scalar]] = None

    # T-022a: opt into the lower-triangular discrete-update scheduler.  When
    # True, Phase 2 of `handle_discrete_update` evaluates state updates in
    # the topological order of the discrete dependency graph: a block reading
    # an upstream block's discrete state sees the post-update x⁺ rather than
    # the snapshotted x⁻.  Cycles in the dependency graph raise
    # `DependencyCycleError`.  Default `False` (diagonal Drake-style update,
    # preserving the cross-block-swap atomicity documented on
    # `SystemBase.handle_discrete_update`).
    lower_triangular_discrete_update: bool = False

    # T-105 Phase 1: opt-in multirate consistency check.  When set to
    # ``"warn"`` or ``"error"``, ``simulate`` runs
    # ``rate_groups.detect_rate_mismatches`` over the diagram immediately
    # after ``validate_diagram``.  A "warn" run logs each mismatched
    # connection through ``warnings.warn(RateMismatchWarning, ...)`` but
    # otherwise lets the simulation proceed (back-compat for existing
    # multirate models that work today thanks to per-block periodic
    # events).  An "error" run raises ``RateMismatchError`` on the first
    # offender.  Default ``None`` keeps the path completely off so
    # single-rate diagrams stay byte-equivalent.  Phase 2 (T-123) will
    # add auto-insertion of ``RateTransition`` blocks; until then this
    # is a diagnostic-only switch.
    check_rate_transitions: str | None = None

    # T-003a: opt-in DAE constraint projection (Newton's method on the
    # algebraic states) at the end of each major step.  The differential
    # states are held fixed; only the algebraic component is corrected.
    # Default `False` (no projection — backwards-compatible).  Has no
    # effect on systems without a mass matrix; the simulator skips the
    # projection cleanly in that case.  `dae_projection_tol` is the
    # max-abs threshold above which the corrector iterates;
    # `dae_projection_max_iter` caps the Newton loop.  The loop is a
    # ``lax.while_loop`` with an early exit, so unused iterations cost
    # nothing at runtime; the default of 20 covers cold starts far from
    # the manifold (near-manifold post-step corrections converge in 1-2
    # iterations regardless).  A non-converged projection emits a
    # ``UserWarning``.  See ``jaxonomy.simulation.dae_projection``.
    #
    # ``dae_initial_projection`` projects the *caller-supplied* context
    # once, before stepping begins.  Use it whenever the initial
    # continuous state was constructed rather than produced by a prior
    # solve — e.g. ``with_continuous_state`` on a DAE system leaves the
    # algebraic rows inconsistent and the first implicit step fails
    # (NaN) without this.
    dae_projection_enabled: bool = False
    dae_projection_tol: float = 1e-8
    dae_projection_max_iter: int = 20
    dae_initial_projection: bool = False

    # T-113-followup-baumgarte-and-ssp: opt-in Baumgarte stabilization of
    # the algebraic constraint residual.  When ``baumgarte_alpha`` and/or
    # ``baumgarte_beta`` are non-None, the simulator wraps
    # ``ode_solver.ode_rhs`` to add ``-2α·ġ - β²·g`` to each algebraic
    # row of the rhs, where ``g = f_a(x)`` is the algebraic-row residual
    # at the current state.  This drives drift to zero exponentially —
    # critically damped at α = β = 1/τ (τ = relaxation time).
    #
    # Default ``None`` for both → no augmentation, the disabled hot path
    # is byte-equivalent (the wrapper short-circuits and returns the
    # original ``ode_rhs`` unchanged when both gains are None).  Has no
    # effect on systems without a mass matrix.  Composes cleanly with
    # ``dae_projection_enabled`` (projection at major-step boundaries
    # kills accumulated drift; Baumgarte damps drift continuously
    # between projections).
    #
    # See :func:`jaxonomy.simulation.dae_projection.baumgarte_augment_ode_rhs`
    # for the augmentation details and the index-reduction caveat.
    baumgarte_alpha: float | None = None
    baumgarte_beta: float | None = None

    # T-113-followup-event-reprojection: opt-in DAE constraint projection
    # immediately after each discrete event reset *within* a major step.
    # T-003a's ``dae_projection_enabled`` projects only at the end of a
    # major step — after the ODE integration plus any localized ZC reset
    # have already happened.  Discrete updates handled at the *top* of
    # ``_major_step`` (``_handle_discrete_update``) modify state before
    # continuous integration runs; if the reset map drops state off the
    # constraint manifold, the subsequent ODE step integrates on infeasible
    # state until the next major-step boundary projection (T-003a) catches
    # up.  Setting this to ``True`` runs ``project_constraints`` right
    # after the discrete-update reset and again after a triggered ZC reset
    # within ``_advance_continuous_time``, so continuous integration always
    # resumes on feasible state.  Default ``False`` (byte-equivalent to
    # the pre-followup hot path).  Has no effect on systems without a
    # mass matrix; the simulator skips the hook cleanly in that case.
    # Composes with ``dae_projection_enabled`` (both can run — major-step
    # boundary projection still fires) and with ``baumgarte_*`` (continuous
    # damping between projections).  Reuses the same ``dae_projection_tol``
    # / ``dae_projection_max_iter`` knobs.
    dae_reproject_after_events: bool = False

    # T-003b: opt-in DAE constraint-residual drift monitor.  When set,
    # the simulator computes ``||f_a||_∞`` at each major step; values
    # above the threshold emit a ``UserWarning`` naming the step time
    # and the measured residual.  Default ``None`` disables the check
    # (no overhead — the default-off path is byte-equivalent to the
    # pre-T-003b code).  Disable projection
    # (``dae_projection_enabled=False``) and enable just this threshold
    # to monitor drift without correcting it.  Has no effect on systems
    # without a mass matrix; the simulator skips the check cleanly in
    # that case.
    dae_drift_threshold: float | None = None

    # T-113 Phase 1: opt-in per-major-step DAE constraint drift trace.
    # Companion to ``dae_drift_threshold`` — that option emits a
    # ``UserWarning`` per violating step but does not retain the raw
    # samples.  When ``record_dae_drift=True``, the simulator tees the
    # post-projection residual ``||f_a||_∞`` plus the step time to a
    # Python-side accumulator via ``jax.debug.callback`` (mirroring the
    # ``_BDFConditionMonitor`` pattern from
    # T-038a-followup-bdf-condition-check) and surfaces the captured
    # ``(time, residual)`` arrays on
    # ``SimulationResults.dae_drift_trace`` — a small dict
    # ``{"time": np.ndarray, "residual": np.ndarray}`` post-finalize.
    # Default ``False`` disables the trace entirely; the default-off
    # path is byte-equivalent (no extra ops compiled in, no monitor
    # constructed) and ``SimulationResults.dae_drift_trace is None``.
    # Has no effect on systems without a mass matrix; the simulator
    # skips the trace cleanly in that case.
    record_dae_drift: bool = False

    # T-125-followup-record-event-times: opt-in capture of zero-crossing
    # event firing times during ``simulate``.  When ``True``, the simulator
    # tees ``(event_index, t_event)`` from each major step that ends on a
    # guard trigger to a Python-side ``_EventTimeRecorder`` via
    # ``jax.debug.callback`` (mirrors the ``_BDFConditionMonitor`` /
    # ``_DAEDriftMonitor`` pattern from T-038a-followup-bdf-condition-check
    # / T-113 phase 1) and surfaces the captured firing-time arrays on
    # ``SimulationResults.event_times`` — a dict ``{event_index:
    # np.ndarray}`` post-finalize.  Default ``False`` disables the
    # capture entirely (no monitor constructed, no ops compiled in) and
    # ``SimulationResults.event_times is None`` — preserves the byte-
    # equivalent default-off path.  Diagrams without zero-crossing events
    # yield ``None`` even when the option is True.  Pairs with
    # :func:`jaxonomy.event_time_gradient` for the implicit-function
    # gradient: feed an entry from ``results.event_times`` straight in as
    # the recorded ``t_event`` rather than tracking it manually.
    record_event_times: bool = False

    # T-038a-followup-bdf-condition-check: opt-in BDF Newton-iteration
    # condition-number diagnostic.  When non-None, the BDF solver's
    # ``newton_iteration`` computes ``jnp.linalg.cond(M - c*J)`` once
    # per major step (cheap on small Newton matrices — one extra SVD
    # on an ``n_states × n_states`` matrix) and forwards the estimate
    # plus the current time to a Python-side aggregator via
    # ``jax.debug.callback``.  The simulator tracks the *maximum*
    # condition number observed across the whole trajectory along
    # with the time at which it occurred, and, on ``simulate`` exit,
    # emits ONE ``UserWarning`` naming the threshold, the max value,
    # and the time of occurrence.  The aggregated warning surface is
    # deliberately *not* per-step — a per-step warning would be too
    # noisy on stiff systems where every step is poorly conditioned
    # — but the underlying max-tracker is per-step so transient
    # ill-conditioning is still surfaced.  Default ``None`` disables
    # the diagnostic entirely (no extra ops compiled in, the BDF hot
    # path is byte-equivalent to the pre-followup code).  Has no
    # effect on non-BDF solvers; the simulator skips attaching the
    # monitor when ``ode_solver`` is not a BDF solver.
    bdf_condition_warning_threshold: float | None = None

    # Compile the detailed BDF non-finite abort diagnostic (failure time,
    # collapsed dt, offending state rows) into the solver.  Default off:
    # the in-graph host callback costs ~0.3 s of XLA compile time per BDF
    # model.  The default path still warns generically after the run when
    # the final state is non-finite, pointing at this flag.
    bdf_nonfinite_diagnostics: bool = False

    # T-027a-followup: simulator-level Zeno protection toggles.  When
    # ``zeno_protection_enabled=False`` (the default), the simulator's
    # ``_major_step`` skips the Zeno tracker entirely — the carry is
    # byte-equivalent to the pre-followup hot path, no extra ops compiled
    # in, no test-suite churn.  When True, every major step that ends on
    # a guard trigger consults ``(time - sim_state.zeno_tprev) <
    # zeno_tolerance`` to decide whether to engage a global Zeno-hold; on
    # engagement, the simulator latches ``zeno_active=True`` and pauses
    # continuous-time integration.  ``zeno_recovery_period`` is the
    # number of consecutive frozen major steps after which the simulator
    # briefly probes for recovery: it clears ``zeno_active`` for one step,
    # and the next guard-trigger check will naturally re-engage Zeno if
    # the cascade is still active (because ``(time - tprev) <
    # zeno_tolerance`` will fire again), or stay cleared otherwise — so
    # transient cascades release the latch automatically while persistent
    # ones stay frozen.  The simulator-level path complements (does not
    # replace) the per-leaf ``declare_zero_crossing(zeno_tolerance=...)``
    # protection from T-027/T-027a.
    #
    # T-027a-followup-vector-tprev: ``zeno_tprev`` and ``zeno_active`` are
    # per-event vectors of shape ``(N_events,)``, one slot per
    # ``ZeroCrossingEvent``.  Each event's last-firing time is tracked
    # independently, so an unrelated event's cascade does not poison the
    # tolerance check for another event.
    # T-027a-followup-per-event-recovery: ``zeno_frozen_steps`` is also a
    # ``(N_events,)`` vector, so the recovery probe fires per-event —
    # event ``i`` clears its own latch when ``frozen[i] >= K`` without
    # affecting any other event's counter or latch.
    zeno_protection_enabled: bool = False
    zeno_tolerance: float = 1e-6
    zeno_recovery_period: int = 10

    # T-013a: opt-in per-signal timestamp capture.  The recording pipeline
    # stores every recorded signal at every major step (legacy global-vector
    # behaviour).  When ``per_signal_timestamps=True``, an out-of-JIT post-
    # processor in ``simulate`` populates
    # ``SimulationResults.per_signal_times`` with each signal's native cadence
    # so ``time_for(name)`` switches to the per-signal vector.
    #
    # ``per_signal_timestamps_mode`` selects the strategy:
    #   - ``"auto"`` (default when the option is on): Mode A — classify each
    #     signal from its source ``OutputPort`` (continuous / periodic /
    #     default) and trim BOTH ``per_signal_times[name]`` AND
    #     ``outputs[name]`` to the schedule for periodic signals.  Genuine
    #     storage savings for downstream consumers (a 1 Hz signal in a 10 s
    #     simulation produces ~11 stored samples instead of ~1001).  Per-
    #     signal fallback to Mode B for signals the classifier cannot place.
    #   - ``"schedule"``: alias for ``"auto"`` — same Mode A path.
    #   - ``"diff"``: Mode B — value-diff dedup of times only, outputs stay
    #     at full length.  Bit-equivalent to the previous Mode-B behaviour.
    #   - ``"buffers"``: T-013a-followup-mode-a-buffers — true in-JIT Mode A.
    #     The simulator allocates per-signal ``(times, values, count)`` rings
    #     at init and each major step's recording write only consumes a slot
    #     in a signal's ring when that signal's cadence classification fires
    #     at the current time.  Cuts both peak buffer memory (vs.
    #     ``"auto"``'s post-finalize trim) and finalize-time post-processing.
    #     Falls back transparently to ``"auto"`` semantics for the global
    #     ``outputs`` shape; the storage saving lives in
    #     ``SimulationResults.outputs[name]`` and ``per_signal_times[name]``.
    #
    # ``per_signal_timestamps_atol`` is the absolute tolerance used when
    # detecting "the signal changed since the last sample" (Mode B) or "the
    # current time matches a tick of the period schedule" (Mode A).  The
    # default 1e-12 catches genuine zero-order-hold plateaus / synchronises
    # to integer-time picosecond precision while staying below any realistic
    # float64 round-off.  Mode A scales the time-tolerance by the period so
    # long simulations don't drift off-schedule.
    #
    # Note: Mode A operates on the trimmed numpy arrays out-of-JIT — the
    # in-JIT recording buffer is unchanged from the legacy path, so peak
    # buffer memory is the same.  The savings are in the post-finalize
    # arrays (``outputs[name]`` and ``per_signal_times[name]``) handed to
    # the user.
    per_signal_timestamps: bool = False
    per_signal_timestamps_atol: float = 1e-12
    per_signal_timestamps_mode: str = "auto"

    # T-110 Phase 1: opt-in provenance/reproducibility manifest.  When
    # ``True``, ``simulate`` calls
    # :func:`jaxonomy.simulation.provenance.compute_provenance` (entirely
    # outside the JIT-traced kernel) and attaches the resulting
    # :class:`ProvenanceManifest` to ``SimulationResults.provenance``.
    # Default ``False`` → ``SimulationResults.provenance is None`` and the
    # simulate path stays byte-equivalent to the pre-followup behaviour.
    record_provenance: bool = False

    # T-012a: opt-in higher-order interpolant for ``SimulationResults.query``.
    # When ``False`` (default) ``query`` uses ``jnp.interp`` linear
    # interpolation — preserves legacy behaviour exactly.  When ``True`` the
    # results pipeline marks ``SimulationResults.solver_states = "pchip"`` (a
    # sentinel placeholder for the eventual native solver-state plumbing) and
    # ``query`` falls back to a PCHIP cubic-Hermite interpolant built from the
    # recorded ``(time, outputs)`` samples.  PCHIP is shape-preserving (no
    # spurious overshoot at zero-order-hold plateaus) and ~3 orders of
    # magnitude more accurate than linear on smooth signals.  The native
    # solver-state path (storing ``Dopri5State.interp_coeff`` per major step
    # for sub-ULP accuracy) remains a deferred follow-up.
    record_solver_states: bool = False

    # Internal flag: True when max_major_steps was explicitly set by the caller rather
    # than auto-estimated by _check_simulate_options.  When True, the bounded fori_loop
    # is used even without enable_autodiff=True so that max_major_steps is honored as a
    # hard simulation budget (useful for Zeno-protection and for non-autodiff workflows
    # that still want a step-count cap).
    _explicit_max_major_steps: bool = dataclasses.field(
        default=False, repr=False, compare=False
    )

    def __post_init__(self):
        # T-A3 follow-up: reconcile the ``max_major_step_size`` alias with
        # the canonical ``max_major_step_length``. When only the alias is set it
        # populates the canonical field; the canonical field otherwise wins. The
        # alias is always re-synced so reads of either are consistent — this
        # also keeps ``dataclasses.replace`` (which re-runs ``__post_init__``)
        # well-behaved regardless of which spelling the caller overrode.
        if self.max_major_step_length is None and self.max_major_step_size is not None:
            self.max_major_step_length = self.max_major_step_size
        self.max_major_step_size = self.max_major_step_length

        # diff-mode follow-up: resolve the explicit differentiation-mode
        # selector into the canonical ``enable_autodiff`` flag, then clear it so
        # re-runs of __post_init__ (via ``dataclasses.replace``) do not re-apply
        # it and clobber an explicit ``enable_autodiff`` override.
        if self.diff_mode is not None:
            valid_modes = ("auto", "forward", "reverse", "none")
            if self.diff_mode not in valid_modes:
                raise ValueError(
                    f"diff_mode={self.diff_mode!r} is not valid; expected one "
                    f"of {valid_modes} (or None)."
                )
            if self.diff_mode == "reverse":
                self.enable_autodiff = True
            elif self.diff_mode in ("forward", "none"):
                # enable_autodiff defaults to False, so a True here is an
                # explicit, contradictory override — fail loudly rather than
                # silently picking one (the footgun this option exists to kill).
                if self.enable_autodiff:
                    raise ValueError(
                        f"Conflicting differentiation settings: diff_mode="
                        f"{self.diff_mode!r} requests no reverse-mode adjoint, "
                        "but enable_autodiff=True installs one. Forward-mode "
                        "autodiff (jax.jacfwd / jvp) must not use the reverse "
                        "adjoint. Use diff_mode='forward' on its own (leave "
                        "enable_autodiff unset)."
                    )
                self.enable_autodiff = False
            # "auto" leaves enable_autodiff untouched.
            self.diff_mode = None

    @property
    def ode_options(self) -> ODESolverOptions:
        return ODESolverOptions(
            rtol=self.rtol,
            atol=self.atol,
            min_step_size=self.min_minor_step_size,
            max_step_size=self.max_minor_step_size,
            method=self.ode_solver_method,
            enable_autodiff=self.enable_autodiff,
            max_checkpoints=self.max_checkpoints,
        )

    def __repr__(self) -> str:
        return (
            f"SimulatorOptions("
            f"math_backend={self.math_backend}, "
            f"enable_tracing={self.enable_tracing}, "
            f"max_major_step_length={self.max_major_step_length}, "
            f"max_major_steps={self.max_major_steps}, "
            f"ode_solver_method={self.ode_solver_method}, "
            f"rtol={self.rtol}, "
            f"atol={self.atol}, "
            f"min_minor_step_size={self.min_minor_step_size}, "
            f"max_minor_step_size={self.max_minor_step_size}, "
            f"zc_bisection_loop_count={self.zc_bisection_loop_count}, "
            f"save_time_series={self.save_time_series}, "
            f"recorded_signals={len(self.recorded_signals or [])}, "  # changed
            f"return_context={self.return_context}, "
            f"validate={self.validate}"
            f")"
        )

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.

Source code in jaxonomy/simulation/dae_drift.py

def algebraic_row_mask(system: "SystemBase") -> np.ndarray | None:
    """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``.
    """
    if not system.has_mass_matrix:
        return None

    # Diagram.mass_matrix is a list of per-leaf matrices; flatten to a block
    # diagonal to get the combined (n, n) mass matrix.
    from scipy.linalg import block_diag

    mm_tree = system.mass_matrix
    leaves = jax.tree.leaves(mm_tree)
    if not leaves:
        return None
    mm = block_diag(*[np.asarray(leaf) for leaf in leaves])
    # Row is algebraic if all entries are zero (well below eps scale).
    return ~np.any(np.abs(mm) > 1e-12, axis=1)

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).	required
options	SimulatorOptions | None	The :class:SimulatorOptions used for the batch run.	required

Returns:

Type	Description
BatchSimulationResults	The same results instance, with results.provenance
BatchSimulationResults	populated.

Source code in jaxonomy/simulation/batch.py

def attach_provenance_to_batch(
    results: BatchSimulationResults,
    system: Diagram | None,
    options: SimulatorOptions | None,
) -> BatchSimulationResults:
    """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.

    Args:
        results: A :class:`BatchSimulationResults` to mutate.
        system: The diagram that was simulated (passed through to
            :func:`compute_provenance`).
        options: The :class:`SimulatorOptions` used for the batch run.

    Returns:
        The same ``results`` instance, with ``results.provenance``
        populated.
    """
    results.provenance = compute_provenance(system, options)
    return results

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.

Source code in jaxonomy/simulation/provenance.py

def bundle_results(results: Any) -> Any:
    """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.
    """
    provenance = getattr(results, "provenance", None)
    if provenance is None:
        return results
    return ResultsWithProvenance(results=results, provenance=provenance)

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.

Source code in jaxonomy/simulation/provenance.py

def compare_manifests(
    actual: ProvenanceManifest,
    expected: ProvenanceManifest,
    *,
    ignore_fields: Optional[set[str]] = None,
) -> list[tuple[str, Any, Any]]:
    """Diff two manifests field-by-field.

    Args:
        actual: the manifest produced by the run being checked.
        expected: the reference manifest (e.g. loaded from a published
            release-tag artifact via :func:`load_manifest`).
        ignore_fields: 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.

    Returns:
        A flat list of ``(dotted_path, actual_value, expected_value)``
        triples — one per differing leaf.  An empty list means the two
        manifests agree on every compared field.
    """
    if ignore_fields is None:
        ignore_fields = set(_DEFAULT_IGNORE_FIELDS)

    actual_dict = actual.to_dict()
    expected_dict = expected.to_dict()
    for field in ignore_fields:
        actual_dict.pop(field, None)
        expected_dict.pop(field, None)
    return _diff_value("", actual_dict, expected_dict)

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).

Source code in jaxonomy/simulation/dae_drift.py

def compute_constraint_residual(
    system: "SystemBase",
    context: "ContextBase",
) -> jnp.ndarray | None:
    """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).
    """
    mask = algebraic_row_mask(system)
    if mask is None or not mask.any():
        return None

    xcdot = system.eval_time_derivatives(context)
    xcdot_flat = jnp.concatenate(
        [jnp.ravel(leaf) for leaf in jax.tree.leaves(xcdot)]
    )
    return xcdot_flat[jnp.asarray(mask)]

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.

Source code in jaxonomy/simulation/provenance.py

def compute_provenance(
    system: Optional["SystemBase"],
    options: Optional["SimulatorOptions"] = None,
    *,
    include_git: bool = True,
    timestamp: Optional[str] = None,
) -> ProvenanceManifest:
    """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.

    Args:
        system: the system being simulated (may be ``None`` for tests
            or pre-built recordings).
        options: the active :class:`SimulatorOptions`; ``None`` records
            an empty options dict.
        include_git: when ``False``, skip the git-HEAD lookup (useful
            when the caller knows it isn't in a git checkout or wants a
            faster path).
        timestamp: optional override (ISO-8601 string).  Defaults to the
            current UTC time.  Override is useful for deterministic
            tests.

    Returns:
        A populated :class:`ProvenanceManifest`.
    """
    versions = _capture_versions()
    precision = _capture_precision_info()
    options_dict = _capture_options(options)
    sys_fp = _system_fingerprint(system)
    if timestamp is None:
        timestamp = datetime.datetime.now(datetime.timezone.utc).isoformat()
    # T-110-followup-git-revision: gather richer git metadata in a single
    # mockable call.  When ``include_git=False`` the entire git capture
    # is skipped (manifest stays None across all git_* fields).
    if include_git:
        git_info = gather_git_info()
    else:
        git_info = {"sha": None, "branch": None, "dirty": None, "commit_time": None}
    # T-110-followup-config-hash: deterministic run-identity hash —
    # excludes timestamp and git metadata so the same config produces
    # the same hash across commits and re-runs.
    config_hash = _compute_config_hash(
        options=options_dict,
        system=sys_fp,
        jaxonomy_version=versions.get("jaxonomy", ""),
        jax_version=versions.get("jax", ""),
    )
    return ProvenanceManifest(
        jaxonomy_version=versions.get("jaxonomy", ""),
        jax_version=versions.get("jax", ""),
        numpy_version=versions.get("numpy", ""),
        precision_info=precision,
        options=options_dict,
        system=sys_fp,
        timestamp=timestamp,
        git_head=git_info.get("sha"),
        git_head_sha=git_info.get("sha"),
        git_branch=git_info.get("branch"),
        git_dirty=git_info.get("dirty"),
        git_head_commit_time=git_info.get("commit_time"),
        config_hash=config_hash,
    )

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.

Source code in jaxonomy/simulation/dae_drift.py

def constraint_residual_norm(
    system: "SystemBase",
    context: "ContextBase",
) -> float | None:
    """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.
    """
    residual = compute_constraint_residual(system, context)
    if residual is None:
        return None
    return float(jnp.max(jnp.abs(residual)))

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

Source code in jaxonomy/simulation/simulator.py

def estimate_max_major_steps(
    system: SystemBase,
    tspan: tuple[float, float],
    max_major_step_length: float = None,
    safety_factor: int = 2,
) -> int:
    """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`.

    Args:
        system (SystemBase): The system to simulate.
        tspan (tuple[float, float]): The time interval to simulate over.
        max_major_step_length (float, optional): 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.
        safety_factor (int, optional): The safety factor to apply to the number of
            major steps.  Defaults to 2.
    """
    # For autodiff of jaxonomy.simulate, this path is not possible, JAX
    # throws an error. To work around this, create:
    #   options = SimulatorOptions(max_major_steps=<my value>)
    # outside jaxonomy.simulate, and pass in like this:
    #   jaxonomy.simulate(my_model, options=options)

    # Find the smallest period amongst the periodic events of the system
    if system.periodic_events.has_events or max_major_step_length is not None:
        # Initialize to infinity - will be overwritten by at least one conditional
        min_discrete_step = np.inf

        # Bound the number of major steps based on the smallest discrete period in
        # the system.
        if system.periodic_events.has_events:
            event_periods = jax.tree_util.tree_map(
                lambda event_data: event_data.period,
                system.periodic_events,
                is_leaf=is_event_data,
            )
            min_discrete_step = jax.tree_util.tree_reduce(min, event_periods)

        # Also bound the number of major steps based on the max major step length
        # in case that is shorter than any of the update periods.
        if max_major_step_length is not None:
            min_discrete_step = min(min_discrete_step, max_major_step_length)

        # in this case, we assume that, on average, major steps triggered by
        # zero crossing event, will be as frequent or less frequent than major steps
        # triggered by the smallest discrete period.
        # anything less than 100 is considered inadequate. user can override if they want this.
        max_major_steps = max(100, safety_factor * int(tspan[1] // min_discrete_step))
        logger.info(
            "max_major_steps=%s based on smallest discrete period=%s",
            max_major_steps,
            min_discrete_step,
        )
    else:
        # in this case we really have no valuable information on which to make an
        # educated guess. who knows how many events might occurr!!!
        # users will have to iterate.
        max_major_steps = 200
        logger.info(
            "max_major_steps=%s by default since no discrete period in system",
            max_major_steps,
        )
    return max_major_steps

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: <constant> and is useful when the user only wants the explicit ∂g/∂p contribution.	required
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.

Source code in jaxonomy/simulation/event_gradient.py

def event_time_gradient(
    guard_fn: Callable[[float, Any, Any], jnp.ndarray],
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    t_event: jnp.ndarray,
    state_at_event_fn: Callable[[Any], Any] | Any,
    params: Any,
    *,
    eps: float = 1e-30,
) -> Any:
    """Compute ``∂t_event/∂params`` via the implicit-function theorem.

    Args:
        guard_fn: ``(t, state, params) -> scalar`` — the zero-crossing
            guard.  Must be JAX-traceable in all three arguments.
        ode_rhs_fn: ``(t, state, params) -> dstate/dt`` — the continuous
            RHS evaluated at the event boundary.  Same PyTree structure
            as the state.
        t_event: Scalar time at which the guard fires.
        state_at_event_fn: 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: <constant>`` and is useful
            when the user only wants the explicit ``∂g/∂p`` contribution.
        params: Parameter PyTree to differentiate with respect to.  May
            be a scalar, ndarray, or any nested container.
        eps: Floor used to clip the denominator
            ``(∂g/∂x · ẋ + ∂g/∂t)`` away from zero before division — keeps
            ``jax.grad`` finite at grazing crossings.  Sign-preserving.

    Returns:
        The PyTree of ``∂t_event/∂params`` with the same structure as
        ``params``.
    """
    t_e = jnp.asarray(t_event)

    # Normalise state_at_event_fn into a callable.
    if callable(state_at_event_fn):
        _state_fn = state_at_event_fn
    else:
        _const = state_at_event_fn
        def _state_fn(_p):  # noqa: ANN001
            return _const

    # Evaluate state at event for use in the rhs.
    x_e = _state_fn(params)

    # Denominator: dg/dx · ẋ + dg/dt, evaluated at (t_e, x_e, params).
    f_at_event = ode_rhs_fn(t_e, x_e, params)

    def _g_of_t(t):
        return guard_fn(t, x_e, params)

    dg_dt = jax.grad(_g_of_t)(t_e)

    # Compute ∂g/∂x · f via a JVP — avoids materialising the full Jacobian.
    _, inner = jax.jvp(
        lambda x: guard_fn(t_e, x, params),
        (x_e,),
        (f_at_event,),
    )
    denom = inner + dg_dt

    # Clip denominator away from exact zero — keeps grad finite at
    # grazing crossings.  Sign-preserving floor.
    safe_denom = jnp.where(
        jnp.abs(denom) < eps,
        jnp.where(denom >= 0, eps, -eps),
        denom,
    )

    # Numerator: total derivative of the residual ``R(p) = guard(t_e,
    # state_fn(p), p)`` w.r.t. ``p`` — JAX handles the chain rule across
    # the ``state_fn`` dependency and the explicit ``params`` dependency
    # uniformly.
    def _residual(p):
        return guard_fn(t_e, _state_fn(p), p)

    dR_dp = jax.grad(_residual)(params)

    # dt_e/dp = -dR_dp / safe_denom.
    return jax.tree_util.tree_map(lambda r: -r / safe_denom, dR_dp)

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,).

Source code in jaxonomy/simulation/event_gradient.py

def event_time_jacobian(
    guard_fn: Callable[[float, Any, Any], jnp.ndarray],
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    t_event: jnp.ndarray,
    state_at_event_fn: Callable[[Any], Any] | Any,
    params: jnp.ndarray,
    *,
    eps: float = 1e-30,
) -> jnp.ndarray:
    """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,)``.
    """
    grad = event_time_gradient(
        guard_fn,
        ode_rhs_fn,
        t_event,
        state_at_event_fn,
        params,
        eps=eps,
    )
    return npa.asarray(grad)

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.

Source code in jaxonomy/simulation/event_gradient.py

def event_times_gradient(
    results: Any,
    params: Any,
    guards: Any,
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    state_at_event_fn: Callable[[float, Any], Any],
    *,
    event_indices: Any = None,
    eps: float = 1e-30,
) -> dict:
    """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.

    Args:
        results: 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.
        params: Parameter PyTree to differentiate with respect to.
            Same semantics as :func:`event_time_gradient`.
        guards: 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.
        ode_rhs_fn: ``(t, state, params) -> dstate/dt`` — same as in
            :func:`event_time_gradient`.  Reused across all firings.
        state_at_event_fn: ``(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.
        event_indices: 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``.
        eps: Forwarded to :func:`event_time_gradient` — denominator
            floor for grazing crossings.

    Returns:
        ``{event_index: stacked_gradient}`` — for each event index,
        the per-firing gradients stacked along a leading axis (so a
        gradient that is itself a PyTree leaf of shape ``S`` becomes
        an array of shape ``(n_firings,) + S``; PyTree containers are
        preserved by mapping the stack over leaves).  Events that
        fired zero times yield an empty leading axis.
    """
    event_times_dict = getattr(results, "event_times", None)
    if event_times_dict is None:
        raise ValueError(
            "event_times_gradient: results.event_times is None. "
            "Re-run simulate(...) with "
            "SimulatorOptions(record_event_times=True) so the firing "
            "instants are captured."
        )

    if callable(guards):
        # Single guard for every event index.
        def _guard_for(_idx):  # noqa: ANN001
            return guards
    else:
        # Per-event mapping.
        guards_map = dict(guards)
        def _guard_for(idx):
            if idx not in guards_map:
                raise KeyError(
                    f"event_times_gradient: no guard supplied for event "
                    f"index {idx}.  Provide guards[{idx}] = <fn> or pass "
                    f"a single callable to apply uniformly."
                )
            return guards_map[idx]

    if event_indices is None:
        selected = list(event_times_dict.keys())
    else:
        selected = list(event_indices)
        for idx in selected:
            if idx not in event_times_dict:
                raise KeyError(
                    f"event_times_gradient: event index {idx} not present "
                    f"in results.event_times (have: "
                    f"{sorted(event_times_dict.keys())})."
                )

    out: dict = {}
    for idx in selected:
        firings = jnp.asarray(event_times_dict[idx])
        guard_fn = _guard_for(idx)
        n_firings = int(firings.shape[0]) if firings.ndim >= 1 else 0

        if n_firings == 0:
            # Empty firing set — surface a structurally-correct empty
            # leading axis by computing a single dummy gradient at
            # ``t_e=0`` and slicing it off.  Avoids special-casing the
            # PyTree shape downstream.
            template = event_time_gradient(
                guard_fn,
                ode_rhs_fn,
                jnp.asarray(0.0),
                lambda p: state_at_event_fn(jnp.asarray(0.0), p),
                params,
                eps=eps,
            )
            out[idx] = jax.tree_util.tree_map(lambda leaf: leaf[None][:0], template)
            continue

        per_firing_grads = []
        for k in range(n_firings):
            t_e = firings[k]
            # Bind the firing time into the state callable so the
            # T-125 helper sees the standard ``state_fn(p) -> state``
            # signature.
            def _state_fn_for_firing(p, _t_e=t_e):
                return state_at_event_fn(_t_e, p)
            g_k = event_time_gradient(
                guard_fn,
                ode_rhs_fn,
                t_e,
                _state_fn_for_firing,
                params,
                eps=eps,
            )
            per_firing_grads.append(g_k)

        # Stack per-firing gradients leafwise so the PyTree structure
        # is preserved with a new leading axis of length n_firings.
        out[idx] = jax.tree_util.tree_map(
            lambda *leaves: jnp.stack(leaves, axis=0),
            *per_firing_grads,
        )

    return out

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.

Source code in jaxonomy/simulation/provenance.py

def load_manifest(path) -> ProvenanceManifest:
    """Load a :class:`ProvenanceManifest` from a JSON file written by
    :meth:`ProvenanceManifest.save`.

    Args:
        path: filesystem path (str or pathlib.Path) of the saved manifest.

    Returns:
        The reconstructed :class:`ProvenanceManifest`.

    Raises:
        FileNotFoundError: if ``path`` does not exist.
        json.JSONDecodeError: if the file is not valid JSON.
    """
    import pathlib

    data = json.loads(pathlib.Path(path).read_text())
    return ProvenanceManifest.from_dict(data)

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.	required
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.	required
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).	required
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.

Source code in jaxonomy/simulation/event_gradient.py

def multi_event_time_gradient(
    guard_fn: Callable[[float, Any, Any], jnp.ndarray] | Any,
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    reset_map_fn: Callable[[float, Any, Any], Any] | Any,
    initial_state: Callable[[Any], Any] | Any,
    event_times: Any,
    params: Any,
    *,
    t0: float = 0.0,
    eps: float = 1e-30,
    rtol: float = 1e-10,
    atol: float = 1e-12,
    return_state_sensitivity: bool = False,
) -> Any:
    """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.

    Args:
        guard_fn: ``(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: ``(t, state, params) -> dstate/dt`` continuous RHS,
            shared across all arcs.  Must be JAX-traceable.
        reset_map_fn: ``(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: 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: 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``.
        params: parameter PyTree to differentiate with respect to.
        t0: trajectory start time (default ``0.0``).
        eps: sign-preserving floor on the implicit-function denominator
            ``(∂g/∂x · ẋ + ∂g/∂t)`` — guards grazing crossings.
        rtol: relative tolerance for the augmented (state + sensitivity)
            arc integration.
        atol: absolute tolerance for the augmented arc integration.
        return_state_sensitivity: when ``True``, also return the list of
            pre-event forward sensitivities ``S⁻(t_e)`` (flat ``(n_x, n_p)``
            arrays) for inspection / debugging.

    Returns:
        The per-firing ``dt_e/dp`` stacked along a leading axis of length
        ``n`` and shaped like ``params`` (a scalar parameter yields shape
        ``(n,)``; a PyTree parameter yields the same PyTree with each leaf
        carrying a leading firing axis).  If ``return_state_sensitivity`` is
        ``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.
    """
    from jax.flatten_util import ravel_pytree
    from jax.experimental.ode import odeint

    t_list = [jnp.asarray(t) for t in jnp.asarray(event_times)]
    n_events = len(t_list)

    # Resolve per-firing guard / reset callables.
    def _as_per_event(obj, name):
        if callable(obj):
            return [obj] * n_events
        seq = list(obj)
        if len(seq) != n_events:
            raise ValueError(
                f"multi_event_time_gradient: {name} has {len(seq)} entries "
                f"but there are {n_events} event_times; supply one callable "
                f"to apply uniformly or a sequence of matching length."
            )
        return seq

    guards = _as_per_event(guard_fn, "guard_fn")
    resets = _as_per_event(reset_map_fn, "reset_map_fn")

    # Normalise the initial-state spec into a callable.
    if callable(initial_state):
        _x0_fn = initial_state
    else:
        _const_x0 = initial_state
        def _x0_fn(_p):  # noqa: ANN001
            return _const_x0

    p_flat, unravel_p = ravel_pytree(params)
    n_p = p_flat.shape[0]

    # Seed state + sensitivity.  Flatten the state PyTree and capture the
    # unravel so guard/rhs/reset can be called in their native structure.
    x0 = _x0_fn(unravel_p(p_flat))
    x_flat, unravel_x = ravel_pytree(x0)
    n_x = x_flat.shape[0]

    def _x0_flat(pf):
        return ravel_pytree(_x0_fn(unravel_p(pf)))[0]

    # S(t0) = ∂x0/∂p  (zeros when initial_state is constant).
    S = jax.jacfwd(_x0_flat)(p_flat)  # (n_x, n_p)

    # Flat-coordinate adapters around the user callables.
    def _f_flat(t, xf, pf):
        return ravel_pytree(ode_rhs_fn(t, unravel_x(xf), unravel_p(pf)))[0]

    def _g_flat(guard, t, xf, pf):
        return jnp.asarray(guard(t, unravel_x(xf), unravel_p(pf)))

    def _r_flat(reset, t, xf, pf):
        return ravel_pytree(reset(t, unravel_x(xf), unravel_p(pf)))[0]

    def _integrate_arc(xf, Smat, ta, tb):
        z0 = jnp.concatenate([xf, Smat.reshape(-1)])

        def _aug(z, t):
            xx = z[:n_x]
            SS = z[n_x:].reshape(n_x, n_p)
            f = _f_flat(t, xx, p_flat)
            # Variational equation: Ṡ = (∂f/∂x) S + ∂f/∂p.
            Jx = jax.jacfwd(lambda a: _f_flat(t, a, p_flat))(xx)
            Jp = jax.jacfwd(lambda b: _f_flat(t, xx, b))(p_flat)
            dS = Jx @ SS + Jp
            return jnp.concatenate([f, dS.reshape(-1)])

        zf = odeint(_aug, z0, jnp.stack([ta, tb]), rtol=rtol, atol=atol)[-1]
        return zf[:n_x], zf[n_x:].reshape(n_x, n_p)

    grads_flat: list = []
    S_minus_list: list = []
    t_prev = jnp.asarray(t0)
    for k in range(n_events):
        t_e = t_list[k]
        guard = guards[k]
        reset = resets[k]

        # Advance state + sensitivity to the firing instant.
        x_flat, S = _integrate_arc(x_flat, S, t_prev, t_e)
        S_minus_list.append(S)

        # Implicit-function-theorem denominator + numerator at the event,
        # using the CORRECTLY propagated S⁻(t_e).
        f_minus = _f_flat(t_e, x_flat, p_flat)
        g_x = jax.grad(lambda a: _g_flat(guard, t_e, a, p_flat))(x_flat)
        g_p = jax.grad(lambda b: _g_flat(guard, t_e, x_flat, b))(p_flat)
        g_t = jax.grad(lambda tt: _g_flat(guard, tt, x_flat, p_flat))(t_e)

        denom = g_x @ f_minus + g_t
        safe_denom = jnp.where(
            jnp.abs(denom) < eps,
            jnp.where(denom >= 0, eps, -eps),
            denom,
        )
        dtau_dp = -(g_x @ S + g_p) / safe_denom  # (n_p,)
        grads_flat.append(dtau_dp)

        # Apply the reset map and the saltation jump to S so the next arc
        # starts from the consistent post-event sensitivity.
        x_plus = _r_flat(reset, t_e, x_flat, p_flat)
        R_x = jax.jacfwd(lambda a: _r_flat(reset, t_e, a, p_flat))(x_flat)
        R_p = jax.jacfwd(lambda b: _r_flat(reset, t_e, x_flat, b))(p_flat)
        R_t = jax.jacfwd(lambda tt: _r_flat(reset, tt, x_flat, p_flat))(t_e)
        f_plus = _f_flat(t_e, x_plus, p_flat)
        S = R_x @ S + R_p + jnp.outer(R_t + R_x @ f_minus - f_plus, dtau_dp)
        x_flat = x_plus
        t_prev = t_e

    # Re-shape each per-firing flat gradient back into the params PyTree and
    # stack leafwise so the output mirrors event_time_gradient's structure
    # with a leading firing axis.
    if n_events == 0:
        template = jax.tree_util.tree_map(
            lambda leaf: leaf[None][:0], unravel_p(jnp.zeros(n_p))
        )
        grads = template
    else:
        per_firing = [unravel_p(g) for g in grads_flat]
        grads = jax.tree_util.tree_map(
            lambda *leaves: jnp.stack(leaves, axis=0), *per_firing
        )

    if return_state_sensitivity:
        return grads, S_minus_list
    return grads

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)

Source code in jaxonomy/simulation/simulator.py

def scalar_cost_simulate(
    system: SystemBase,
    context_fn: Callable[[Any], ContextBase],
    t_span: tuple[float, float],
    params: Any,
    cost_fn: Callable[[ContextBase], Any] = None,
    *,
    options: SimulatorOptions = None,
    return_grad: bool = 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.

    Args:
        system: the Diagram / LeafSystem to simulate.
        context_fn: ``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.
        t_span: ``(t0, tf)`` simulation interval.
        params: parameter pytree — the differentiation argument.
        cost_fn: ``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.
        options: ``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``.
        return_grad: 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).

    Returns:
        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)
    """
    if options is None:
        options = SimulatorOptions()
    # Recording is incompatible with autodiff; force the supported config.
    options = dataclasses.replace(
        options, enable_autodiff=True, recorded_signals=None,
    )

    if cost_fn is None:
        def cost_fn(ctx):
            import jax.numpy as _jnp
            return _jnp.sum(_jnp.asarray(ctx.continuous_state))

    def _cost(p):
        ctx = context_fn(p)
        res = simulate(system, ctx, t_span=t_span, options=options)
        return cost_fn(res.context)

    if return_grad:
        return jax.value_and_grad(_cost)(params)
    return _cost(params)

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.	None

Returns:

Name	Type	Description
SimulationResults	SimulationResults	A named tuple containing the 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.

This function is meant to best handle single independent simulations. Calling this function repeatedly will always trigger a recompilation of the model when using the JAX backend. To avoid this, call advance_to directly.

Source code in jaxonomy/simulation/simulator.py

@remap_simulation_errors
def simulate(
    system: SystemBase,
    context: ContextBase,
    t_span: tuple[float, float] = None,
    options: SimulatorOptions = None,
    results_options: ResultsOptions = None,
    recorded_signals: dict[str, OutputPort] = None,
    postprocess: bool = True,
    flatten: bool = False,
    *,
    tspan: tuple[float, float] = None,
) -> SimulationResults:
    """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.

    Args:
        system (SystemBase): The hybrid dynamical system to simulate.
        context (ContextBase): The initial state and parameters of the system.
        tspan (tuple[float, float]): The start and end times of the simulation.
        options (SimulatorOptions): Options for the simulation process and ODE solver.
        results_options (ResultsOptions): Options related to how the outputs are
            stored, interpolated, and returned.
        recorded_signals (dict[str, OutputPort]):
            Dictionary of ports for which the time series should be recorded.

    Returns:
        SimulationResults: A named tuple containing the 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`.

        This function is meant to best handle single independent simulations.
        Calling this function repeatedly will always trigger a recompilation of the
        model when using the JAX backend. To avoid this, call advance_to directly.
    """

    # Backward-compatibility shim: accept legacy `tspan` keyword argument.
    if tspan is not None:
        if t_span is not None:
            raise TypeError("Cannot specify both 't_span' and 'tspan'; 'tspan' is deprecated.")
        import warnings
        warnings.warn(
            "The 'tspan' argument is deprecated; use 't_span' instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        t_span = tspan

    if t_span is None:
        raise TypeError("simulate() missing required argument: 't_span'")

    options = _check_options(system, options, t_span, recorded_signals)

    import warnings
    from jaxonomy.framework.validation import validate_diagram

    if getattr(options, 'validate', True):
        result = validate_diagram(system, options)
        if result.warnings:
            for w in result.warnings:
                warnings.warn(w, UserWarning, stacklevel=2)
        result.raise_if_invalid()

    # T-105 Phase 1: opt-in multirate consistency check.  Default
    # ``None`` keeps the existing single-rate path byte-equivalent;
    # users opt in via ``SimulatorOptions.check_rate_transitions``.
    _rate_check = getattr(options, "check_rate_transitions", None)
    if _rate_check:
        from .rate_groups import detect_rate_mismatches  # local import: avoids cycle
        if isinstance(system, Diagram):
            detect_rate_mismatches(system, on_mismatch=_rate_check)

    # Optionally flatten nested Diagrams to a single depth for reduced overhead.
    if flatten and isinstance(system, Diagram):
        system = flatten_diagram(system)
        context = system.create_context(time=t_span[0])

    # Opt-in initial-consistency projection: a caller-supplied context
    # whose algebraic entries are stale (e.g. after
    # ``with_continuous_state`` on a DAE system) makes the first implicit
    # step fail; this Newton-projects them onto the constraint manifold
    # before stepping begins.  No-op for systems without a mass matrix.
    if getattr(options, "dae_initial_projection", False) and getattr(
        system, "has_mass_matrix", False,
    ):
        from .dae_projection import project_constraints
        context = project_constraints(
            system,
            context,
            tol=options.dae_projection_tol,
            max_iter=options.dae_projection_max_iter,
        )

    if results_options is None:
        results_options = ResultsOptions()

    if results_options.mode != ResultsMode.auto:
        raise NotImplementedError(
            f"Simulation output mode {results_options.mode.name} is not supported. "
            "Only 'auto' is presently supported."
        )

    if system.has_dirty_static_parameters:
        raise ValueError(
            "Some static parameters have been updated. Please create a new context."
        )

    # HACK: Jaxonomy presently does not use interpolant to produce
    # results sample between minor_step end times, so we clamp
    # the minor step size to the max_results_interval instead.
    if (
        results_options.max_results_interval is not None
        and results_options.max_results_interval > 0
        # max_minor_step_size is None by default (unbounded), which is
        # always larger than any finite results interval, so it needs
        # clamping too — guard the comparison against None either way.
        and (
            options.max_minor_step_size is None
            or results_options.max_results_interval < options.max_minor_step_size
        )
    ):
        options = dataclasses.replace(
            options,
            max_minor_step_size=results_options.max_results_interval,
        )
        logger.info(
            "max_minor_step_size reduced to %s to match max_results_interval",
            options.max_minor_step_size,
        )

    orig_x64 = jax.config.read("jax_enable_x64")
    enable_x64 = orig_x64
    if options and options.precision != "auto":
        enable_x64 = (options.precision == "float64")
    jax.config.update("jax_enable_x64", enable_x64)

    if options and options.precision != "auto":
        target_dtype = jnp.float64 if options.precision == "float64" else jnp.float32
        def cast_floats(x):
            if isinstance(x, (jax.Array, np.ndarray)):
                if jnp.issubdtype(x.dtype, jnp.floating):
                    return x.astype(target_dtype)
            return x
        context = jax.tree_util.tree_map(cast_floats, context)

    ode_solver = ODESolver(system, options=options.ode_options)

    sim = Simulator(system, ode_solver=ode_solver, options=options)
    logger.info("Simulator ready to start: %s, %s", options, ode_solver)

    # Define a function to be traced by JAX, if allowed, closing over the
    # arguments to `_simulate`.
    def _wrapped_simulate() -> tuple[ContextBase, ResultsData]:
        t0, tf = t_span
        initial_context = context.with_time(t0)
        sim_state = sim.advance_to(tf, initial_context)
        error_end_time_not_reached(
            tf, sim_state.context.time, sim_state.step_end_reason
        )
        final_context = sim_state.context if options.return_context else None
        return final_context, sim_state.results_data

    # JIT-compile the simulation, if allowed
    if options.enable_tracing:
        _wrapped_simulate = jax.jit(_wrapped_simulate)
        _wrapped_simulate = Profiler.jaxjit_profiledfunc(
            _wrapped_simulate, "_wrapped_simulate"
        )

    # Run the simulation
    try:
        system.cache_enabled = True
        final_context, results_data = _wrapped_simulate()

        if postprocess and results_data is not None:
            time, outputs = results_data.finalize()
            # The backend wrapper exposes ``finalize`` only; reach into
            # the inner ``_solution_data`` for the optional finalize
            # variants (Mode A buffers, native interpolant) when the JAX
            # backend supplied them.  Falls back to None for backends
            # that don't.
            _inner_results = getattr(results_data, "_solution_data", results_data)
            # T-012a-followup: pull the per-step interpolant ring (when
            # the buffer was allocated via ``record_solver_states=True``).
            _interpolant_finalized = None
            finalize_interp = getattr(
                _inner_results, "finalize_interpolant", None,
            )
            if finalize_interp is not None:
                _interpolant_finalized = finalize_interp()
            # T-013a-followup-mode-a-buffers: when per-signal buffers
            # were allocated (mode="buffers"), pull each signal's
            # trimmed (times, values) directly from its own ring rather
            # than reusing the legacy global trim.  This is where the
            # storage saving is realised — periodic signals' arrays are
            # already at the right cadence, no post-trim required.
            per_signal_finalized = None
            finalize_per_signal = getattr(
                _inner_results, "finalize_per_signal", None,
            )
            if finalize_per_signal is not None:
                per_signal_finalized = finalize_per_signal()
            if per_signal_finalized is not None:
                _global_t, per_outputs, per_times = per_signal_finalized
                # Keep ``time`` as the legacy global vector (some
                # downstream code expects it for ``align`` cross-
                # references); use the per-signal trimmed outputs.
                outputs = per_outputs
                # Stash the per-signal times so the post-processor at
                # the bottom of ``simulate`` can promote them to
                # ``SimulationResults.per_signal_times`` without
                # re-running the cadence classifier.
                _per_signal_times_from_buffers = per_times
            else:
                _per_signal_times_from_buffers = None
        else:
            time, outputs = None, None
            _per_signal_times_from_buffers = None
            _interpolant_finalized = None

    finally:
        system.post_simulation_finalize()
        system.cache_enabled = False
        jax.config.update("jax_enable_x64", orig_x64)

    # T-038a-followup-bdf-condition-check: emit ONE aggregated
    # ``UserWarning`` if the BDF Newton-iteration condition number ever
    # exceeded the threshold during the trajectory.  The monitor's
    # running max is updated host-side via ``jax.debug.callback`` from
    # inside the JIT'd BDF ``newton_iteration``.  ``maybe_warn`` is a
    # no-op when the option was unset or when the max stayed below
    # threshold — preserves the byte-equivalent default-off path.
    _bdf_cond_monitor = getattr(sim, "_bdf_cond_monitor", None)
    if _bdf_cond_monitor is not None:
        _bdf_cond_monitor.maybe_warn()

    # Free post-run non-finite check: when the returned final state contains
    # NaN/Inf, say so and point at the detailed opt-in diagnostic.  Runs
    # host-side once after the JIT'd kernel returns; skipped when the final
    # context isn't materialized (traced/vmap callers get the in-graph
    # diagnostic instead, if enabled).
    if final_context is not None:
        try:
            _xc = final_context.continuous_state
            _leaves = jax.tree.leaves(_xc)
            _nonfinite = any(
                not bool(jnp.all(jnp.isfinite(l))) for l in _leaves if l is not None
            )
        except Exception:
            _nonfinite = False
        if _nonfinite:
            import warnings as _warnings
            _warnings.warn(
                "Simulation ended with a non-finite continuous state. For the "
                "failure time, collapsed step size, and the offending state "
                "rows, rerun with "
                "SimulatorOptions(bdf_nonfinite_diagnostics=True) (BDF). "
                "Common causes: an inconsistent algebraic initial state (for "
                "DAEs, set dae_initial_projection=True), an ill-conditioned "
                "component equation, or a genuinely diverging solution.",
                UserWarning,
                stacklevel=2,
            )

    # T-138 — decimated-recording diagnostic.  When the recording buffer
    # filled, the JAX backend now degrades to uniform decimation (keep
    # every Nth sample spanning the whole trajectory) instead of the old
    # ring-wrap that kept only the tail.  Detection is exact: the
    # backend's ``record_stride`` ends > 1 iff at least one compaction
    # ran.  This Python-side check runs once after the JIT'd kernel
    # returns (so it doesn't disturb vmap) and tells the user the
    # results are complete but at reduced resolution.
    _final_stride = None
    _total_steps = None
    if results_data is not None:
        _inner_rd = getattr(results_data, "_solution_data", results_data)
        _stride_attr = getattr(_inner_rd, "record_stride", None)
        if _stride_attr is not None:
            try:
                _final_stride = int(np.max(np.asarray(_stride_attr)))
                _total_steps = int(np.max(np.asarray(_inner_rd.step_count)))
            except (TypeError, ValueError):
                _final_stride = None
    if (
        _final_stride is not None
        and _final_stride > 1
        and time is not None
        and not getattr(options, "enable_autodiff", False)
    ):
        n_kept = len(time)
        resolved_buffer = options.buffer_length
        buf_str = (
            f"buffer_length={resolved_buffer}"
            if resolved_buffer is not None
            else "buffer_length=None (auto)"
        )
        warnings.warn(
            f"jaxonomy.simulate: the recording buffer ({buf_str}) filled; "
            f"the trajectory was recorded at reduced resolution "
            f"({n_kept} of {_total_steps} samples, keeping every "
            f"{_final_stride}th). The recorded time-series still starts at "
            f"t0 and covers the whole trajectory (the head is no longer "
            f"dropped; the last kept sample may precede tf by up to "
            f"{_final_stride} steps). Set SimulatorOptions(buffer_length="
            f"{max(int(_total_steps or 0) + 1, 4000)}) or larger to capture "
            f"every sample, loosen rtol/atol, or reduce the number of "
            f"recorded signals.",
            UserWarning,
            stacklevel=2,
        )

    # T-002b-followup-buffer-overflow-warning (legacy signature, kept as a
    # backstop for backends without the T-138 decimation fields): a
    # ring-wrapped buffer surfaces as the *start* of the recorded
    # trajectory being well past the requested ``t_span[0]``; a
    # non-overflowed (or decimated) run always records its first sample
    # at ``t_span[0]``.
    if (
        _final_stride in (None, 1)
        and time is not None
        and len(time) >= 1
        and not getattr(options, "enable_autodiff", False)
        and t_span is not None
    ):
        t0 = float(t_span[0])
        tf = float(t_span[1])
        span = max(tf - t0, 1.0)  # avoid divide-by-zero on degenerate spans
        first_recorded = float(time[0])
        # 0.1% of the integration span is the threshold below which we
        # treat ``time[0]`` as "essentially t_span[0]"; above it the
        # buffer almost certainly wrapped.
        if first_recorded - t0 > span * 1e-3:
            import warnings as _warnings

            resolved_buffer = options.buffer_length
            n_kept = len(time)
            buf_str = (
                f"buffer_length={resolved_buffer}"
                if resolved_buffer is not None
                else "buffer_length=None (auto)"
            )
            # T-B3/B8-followup-buffer-dopri5-sizing: the recorder saves one
            # sample per *minor* (accepted) solver step. Adaptive solvers
            # (Dopri5 / the "auto" default) take far more minor steps than the
            # major-step count the auto-size is derived from — and they take
            # *more* minor steps as rtol/atol tighten. So the auto buffer
            # (sized to the major-step estimate) can be overrun by a tight-
            # tolerance Dopri5 run, silently dropping the trajectory head.
            # Make the recommendation concrete and solver/tolerance-aware.
            solver_method = getattr(options, "ode_solver_method", "auto")
            rtol = getattr(options, "rtol", None)
            atol = getattr(options, "atol", None)
            # The buffer held at least the kept-tail samples; the full run
            # needed more. Recommend a healthy multiple of what we know was
            # already exceeded.
            if resolved_buffer is not None:
                suggested = max(int(resolved_buffer) * 4, n_kept * 4, 4000)
            else:
                suggested = max(n_kept * 4, 4000)
            adaptive_note = ""
            if solver_method in ("auto", "dopri5", "Dopri5", "RK45", "DOP853"):
                adaptive_note = (
                    f" This is an adaptive solver "
                    f"(ode_solver_method={solver_method!r}, rtol={rtol}, "
                    f"atol={atol}); it records one sample per accepted minor "
                    f"step, and tightening rtol/atol increases the step count. "
                    f"A fixed-step solver (ode_solver_method='rk4') records a "
                    f"predictable sample count if you need a tight buffer."
                )
            _warnings.warn(
                f"jaxonomy.simulate: recording buffer overflow detected — "
                f"the returned ``results.time`` starts at "
                f"{first_recorded!r} (requested t_span[0]={t0!r}); samples "
                f"from earlier in the trajectory were overwritten because "
                f"the simulator's recording ring buffer ({buf_str}) "
                f"filled (kept the last {n_kept} samples).{adaptive_note} "
                f"Set SimulatorOptions(buffer_length={suggested}) (or larger) "
                f"to capture the full trajectory, loosen rtol/atol, or reduce "
                f"the number of recorded signals.",
                UserWarning,
                stacklevel=2,
            )

    # Reset the integer time scale to the default value in case we decreased precision
    # to reach the end time of a long simulation.  Typically this won't do anything.
    if options.int_time_scale is not None:
        IntegerTime.set_default_scale()

    # T-013a: opt-in per-signal timestamp capture.  Runs as a post-
    # processor on the trimmed numpy arrays — no schema change to the
    # recording buffer.  Two modes:
    #
    #   - Mode A (``per_signal_timestamps_mode="auto"`` or ``"schedule"``,
    #     the default when the option is enabled): classify each signal's
    #     natural cadence from its source ``OutputPort`` (continuous /
    #     periodic / default) and trim BOTH the times and the outputs of
    #     periodic signals down to the schedule.  Genuine storage savings
    #     for the per-signal arrays.  Falls back to Mode B per signal
    #     when classification is "default".
    #   - Mode B (``per_signal_timestamps_mode="diff"``): legacy value-
    #     diff dedup on times only; outputs stay at full length.
    per_signal_times = None
    if (
        getattr(options, "per_signal_timestamps", False)
        and time is not None
        and outputs is not None
    ):
        atol = getattr(options, "per_signal_timestamps_atol", 1e-12)
        mode = getattr(options, "per_signal_timestamps_mode", "auto")
        if mode == "buffers" and _per_signal_times_from_buffers is not None:
            # T-013a-followup-mode-a-buffers: in-JIT per-signal buffers
            # already produced trimmed arrays during ``finalize_per_signal``
            # — promote them straight to the result.  No post-finalize
            # classification or trim runs.
            per_signal_times = _per_signal_times_from_buffers
        elif mode == "diff":
            per_signal_times = ResultsRecorder.compute_per_signal_times(
                time, outputs, atol=atol,
            )
        else:  # "auto" or "schedule"
            classifications = ResultsRecorder.classify_signal_cadence(
                options.recorded_signals,
            )
            per_signal_times, outputs = (
                ResultsRecorder.compute_per_signal_schedule(
                    time, outputs, classifications, atol=atol,
                )
            )

    # T-012a-followup: when ``record_solver_states=True`` AND the JAX
    # backend produced per-step interpolant data, build a
    # ``NativeInterpolant`` so ``query`` evaluates the solver's own
    # polynomial at sub-ULP accuracy.  Fall back to the PCHIP sentinel
    # otherwise — preserves the T-012a partial behaviour for solvers /
    # backends that don't expose a fixed-shape ``interp_coeff``.
    solver_states_marker = None
    if (
        getattr(options, "record_solver_states", False)
        and time is not None
        and outputs is not None
    ):
        if _interpolant_finalized is not None:
            from .types import NativeInterpolant
            t_prev_arr, t_step_arr, coeff_arr = _interpolant_finalized
            if t_prev_arr.shape[0] >= 1:
                # Build a _StableUnravel from the final context's continuous
                # state so query() can restore the polyval result to the
                # original pytree shape.  When the final_context is None
                # (return_context=False), we still pull the original
                # context the caller passed in via ``context`` — same shape.
                try:
                    from ..backend._jax.ode_solver_impl import _StableUnravel
                    cs_template = (
                        final_context.continuous_state
                        if final_context is not None
                        else context.continuous_state
                    )
                    unravel = _StableUnravel(cs_template)
                except Exception:
                    unravel = None
                solver_states_marker = NativeInterpolant(
                    t_prev=np.asarray(t_prev_arr),
                    t_step=np.asarray(t_step_arr),
                    interp_coeff=np.asarray(coeff_arr),
                    unravel=unravel,
                    solver="dopri5",
                )
            else:
                solver_states_marker = "pchip"
        else:
            solver_states_marker = "pchip"

    # T-110 Phase 1: opt-in provenance/reproducibility manifest.
    # Computed entirely in Python after the JIT-traced kernel returns —
    # never inside ``_wrapped_simulate`` — so the default-off path is
    # byte-equivalent.  See :mod:`jaxonomy.simulation.provenance`.
    provenance_manifest = None
    if getattr(options, "record_provenance", False):
        from .provenance import compute_provenance
        provenance_manifest = compute_provenance(system, options)

    # T-113 Phase 1: opt-in per-major-step DAE drift trace.  Pull the
    # captured (time, residual) lists from the host-side monitor and
    # promote to numpy arrays.  ``finalize`` returns ``None`` when no
    # samples were collected (e.g. zero major steps, or pure-ODE
    # system).  Default-off (no monitor attached) leaves the field
    # ``None`` — byte-equivalent.
    dae_drift_trace = None
    _dae_drift_monitor = getattr(sim, "_dae_drift_monitor", None)
    if _dae_drift_monitor is not None:
        dae_drift_trace = _dae_drift_monitor.finalize()

    # T-125-followup-record-event-times: pull captured ``(event_index,
    # firing_time)`` samples from the host-side recorder.  ``finalize``
    # returns a dict keyed by event index (numpy arrays of times) when
    # the recorder was constructed, or ``None`` when no recorder is
    # attached (default-off path or zero-event diagrams).  Byte-
    # equivalent default-off behaviour: option unset → no recorder →
    # ``event_times is None``.
    event_times = None
    _event_time_recorder = getattr(sim, "_event_time_recorder", None)
    if _event_time_recorder is not None:
        event_times = _event_time_recorder.finalize()

    return SimulationResults(
        final_context,
        time=time,
        outputs=outputs,
        per_signal_times=per_signal_times,
        solver_states=solver_states_marker,
        provenance=provenance_manifest,
        dae_drift_trace=dae_drift_trace,
        event_times=event_times,
    )

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).	None
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).	None
results_options	ResultsOptions | None	Optional :class:ResultsOptions passed through.	None
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.

Source code in jaxonomy/simulation/batch.py

@remap_simulation_errors
def simulate_batch(
    diagram: Diagram,
    t_span: tuple[float, float],
    param_batches: dict[str, Any],
    options: SimulatorOptions | None = None,
    recorded_signals: dict[str, OutputPort] | None = None,
    results_options: ResultsOptions | None = None,
    use_vmap: bool = False,
    _force_loop: bool = False,
    lazy: bool = False,
) -> BatchSimulationResults:
    """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``.

    Args:
        diagram: Template diagram (unchanged).
        t_span: ``(t_start, t_stop)``.
        param_batches: 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)}``.
        options: :class:`SimulatorOptions` with ``math_backend="jax"`` and
            ``max_major_steps`` set (required).
        recorded_signals: 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: Optional :class:`ResultsOptions` passed through.
        use_vmap: If ``True``, attempt vectorisation via ``jax.vmap`` (pure-JAX
            diagrams only). Raises ``ValueError`` if the diagram is not pure-JAX.
        _force_loop: If ``True``, always use the loop path regardless of diagram
            type (useful for testing / debugging).

    Returns:
        :class:`BatchSimulationResults` with ``outputs[name].shape[0] == N``.

    Raises:
        ValueError: Inconsistent batch sizes, missing options, or invalid backend.
        TypeError: ``diagram`` is not a :class:`~jaxonomy.framework.diagram.Diagram`.
    """
    # Error-message remapping is handled by the @remap_simulation_errors
    # decorator on simulate_batch itself.
    if not isinstance(diagram, Diagram):
        raise TypeError(f"simulate_batch expects a Diagram, got {type(diagram)}")
    if recorded_signals is None:
        raise ValueError("simulate_batch requires recorded_signals (same as simulate).")
    if options is None:
        raise ValueError(
            "simulate_batch requires SimulatorOptions with math_backend='jax' and "
            "max_major_steps set."
        )
    if options.max_major_steps is None or options.max_major_steps <= 0:
        raise ValueError(
            "simulate_batch requires options.max_major_steps to be set to a positive int."
        )
    if options.math_backend != "jax":
        raise ValueError(
            f"simulate_batch only supports math_backend='jax', got {options.math_backend!r}."
        )

    n = _infer_batch_size(param_batches)
    opts = dataclasses.replace(options, enable_autodiff=False)

    vmap_safe = _is_vmap_safe(diagram)

    if use_vmap and not vmap_safe:
        raise ValueError(
            "use_vmap=True requires a pure-JAX diagram (no CustomPythonBlock / FMU). "
            "The diagram contains non-traceable blocks. Use use_vmap=False or the "
            "loop path."
        )

    # T-019-followup (2026-07-10): the CPU+small-batch UserWarning that
    # used to fire here was removed together with the per-row host-loop
    # finalize it warned about — the finalize is now vectorised and the
    # vmap path is within ~25% of the kernel path on CPU at N=1000
    # (0.41 s vs 0.33 s on the reference damped-oscillator sweep).

    use_kernel = vmap_safe and not _force_loop

    # T-110-followup-attach-on-batch: capture provenance ONCE at the start
    # of the batch (shared across all replicas — the only thing that
    # differs per-replica is the parameter dict, which is already in
    # ``param_batches``).  Captured before dispatch so the manifest's
    # timestamp records when the batch began; default-off path stays
    # byte-equivalent because ``compute_provenance`` is never called when
    # ``record_provenance=False``.
    provenance_manifest: ProvenanceManifest | None = None
    if getattr(options, "record_provenance", False):
        provenance_manifest = compute_provenance(diagram, options)

    if use_kernel:
        result = _simulate_batch_kernel(
            diagram, t_span, param_batches, opts, recorded_signals,
            results_options, use_vmap, n, lazy=lazy,
        )
    else:
        result = _simulate_batch_loop(
            diagram, t_span, param_batches, opts, recorded_signals,
            results_options, n,
        )

    if provenance_manifest is not None:
        result.provenance = provenance_manifest
    return result

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.

Source code in jaxonomy/simulation/cloud_runner.py

def simulate_cloud(*args: Any, **kwargs: Any):
    """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:
        NotImplementedError: always, in this build.
    """
    raise NotImplementedError(
        "simulate_cloud() is not available in this build: no cloud execution "
        "backend is bundled. Use the local simulate()/batch runners instead."
    )

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)

Source code in jaxonomy/simulation/simulator.py

def simulate_jacfwd(
    system: SystemBase,
    context_fn: Callable[..., ContextBase],
    t_span: tuple[float, float],
    params: Any,
    output_fn: Callable[[Any], Any] = None,
    *,
    options: SimulatorOptions = None,
    record_provenance: bool = 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.

    Args:
        system: a Diagram or LeafSystem to simulate.
        context_fn: a callable ``context_fn(params) -> Context`` that
            constructs an initial context with the given parameter
            pytree applied.
        t_span: ``(t0, tf)`` simulation interval.
        params: parameter pytree (the differentiation argument).
        output_fn: callable applied to the final ``Context`` to produce
            a scalar or array output. Defaults to extracting the final
            continuous state.
        options: ``SimulatorOptions``. ``enable_autodiff`` is forced to
            False for the JVP path; pass ``rtol`` / ``atol`` /
            ``ode_solver_method`` to control accuracy.
        record_provenance: 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.

    Returns:
        ``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)
    """
    if options is None:
        options = SimulatorOptions()
    options = dataclasses.replace(options, enable_autodiff=False)

    if output_fn is None:
        def output_fn(ctx):
            return ctx.continuous_state

    def _fwd(p):
        ctx = context_fn(p)
        res = simulate(system, ctx, t_span=t_span, options=options)
        return output_fn(res.context)

    jacobian = jax.jacfwd(_fwd)(params)

    # T-110-followup-attach-on-jacfwd: when opt-in, compute the
    # provenance manifest in plain Python around the jacfwd call and
    # return it as the second element of a tuple.  ``simulate_jacfwd``'s
    # natural return type is a JAX array (or pytree of arrays) — it
    # has no ``.provenance`` field to attach to, so the tuple shape is
    # the cleanest stable surface.  The Python-side ``record_provenance``
    # kwarg never reaches the traced ``_fwd`` body so the default-off
    # path is byte-equivalent to the pre-followup code.
    if record_provenance:
        from .provenance import compute_provenance
        # Force a Python-side option snapshot that reflects what was
        # actually requested for provenance purposes — the inner
        # ``options`` we passed to ``simulate`` has ``enable_autodiff``
        # already flipped to False, which is the value we want to
        # record (it's what produced the result).  We also flip
        # ``record_provenance`` to True on the snapshot so an auditor
        # can tell the manifest was opt-in (mirrors simulate_batch).
        options_for_manifest = dataclasses.replace(
            options, record_provenance=True,
        )
        manifest = compute_provenance(system, options_for_manifest)
        return jacobian, manifest

    return jacobian

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.	None
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.

Source code in jaxonomy/simulation/static_sweep.py

@remap_simulation_errors
def simulate_static_sweep(
    diagram_factory: Callable[..., Diagram],
    t_span: tuple[float, float],
    static_param_grid: dict[str, Sequence[Any]],
    options: SimulatorOptions,
    recorded_signals_factory: Callable[[Diagram], dict],
    results_options: ResultsOptions | None = None,
    mode: str = "zip",
) -> BatchSimulationResults:
    """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.

    Args:
        diagram_factory: Callable taking the static-param keyword arguments
            specified by ``static_param_grid`` and returning a :class:`Diagram`.
        t_span: ``(t_start, t_stop)`` — same for every element.
        static_param_grid: Mapping parameter name -> sequence of values. Every
            list must have the same length under ``mode="zip"``; or any
            lengths under ``mode="product"`` (cartesian product).
        options: :class:`SimulatorOptions`. ``max_major_steps`` must be set if
            using ``math_backend="jax"``.
        recorded_signals_factory: Callable ``(diagram) -> {name: OutputPort}``.
            Invoked once per grid element with the freshly-built diagram.
        results_options: Optional :class:`ResultsOptions` passed to
            :func:`simulate`.
        mode: ``"zip"`` (default — pair lists element-wise) or ``"product"``
            (cartesian product over all keys).

    Returns:
        :class:`BatchSimulationResults` with ``outputs[name].shape == (N, T)``
        and ``time.shape == (T,)`` where ``N`` is the number of grid elements
        and ``T`` is the time-vector length of the first run (other runs are
        linearly interpolated onto this grid). The ``contexts`` attribute is
        attached to the returned object as a list of per-element final
        contexts.

    Raises:
        ValueError: empty grid, mismatched zip lengths, unknown mode, or
            missing required options.
        TypeError: ``diagram_factory`` did not return a :class:`Diagram`.
    """
    if options is None:
        raise ValueError("simulate_static_sweep requires a SimulatorOptions.")
    if not callable(diagram_factory):
        raise TypeError(
            f"simulate_static_sweep: diagram_factory must be callable, "
            f"got {type(diagram_factory)}"
        )
    if not callable(recorded_signals_factory):
        raise TypeError(
            "simulate_static_sweep: recorded_signals_factory must be callable "
            f"((diagram) -> dict), got {type(recorded_signals_factory)}"
        )

    grid = _expand_grid(static_param_grid, mode)
    n = len(grid)

    time_ref = None
    out_lists: dict[str, list] = {}
    contexts: list = []
    signal_names: list[str] | None = None

    for i, combo in enumerate(grid):
        d = diagram_factory(**combo)
        if not isinstance(d, Diagram):
            raise TypeError(
                f"simulate_static_sweep: diagram_factory(**{combo!r}) returned "
                f"{type(d)}, expected a Diagram."
            )
        sig = recorded_signals_factory(d)
        if not isinstance(sig, dict) or not sig:
            raise ValueError(
                "simulate_static_sweep: recorded_signals_factory must return a "
                f"non-empty dict of {{name: OutputPort}}; got {type(sig)} for "
                f"combo {combo!r}."
            )
        if signal_names is None:
            signal_names = list(sig.keys())
            out_lists = {k: [] for k in signal_names}
        elif list(sig.keys()) != signal_names:
            raise ValueError(
                f"simulate_static_sweep: recorded_signals_factory returned "
                f"different signal names at element {i} ({list(sig.keys())!r}) "
                f"vs element 0 ({signal_names!r})."
            )

        ctx = d.create_context()
        res = simulate(
            d,
            ctx,
            t_span,
            options=options,
            results_options=results_options,
            recorded_signals=sig,
        )
        if res.outputs is None:
            raise RuntimeError(
                f"simulate_static_sweep: simulate returned no outputs at "
                f"element {i} (combo={combo!r})."
            )
        contexts.append(res.context)

        if time_ref is None:
            time_ref = res.time
            for name in signal_names:
                out_lists[name].append(res.outputs[name])
        else:
            t_ref_end = float(time_ref[-1])
            t_run_end = float(res.time[-1])
            if abs(t_run_end - t_ref_end) / max(abs(t_ref_end), 1e-10) > 0.01:
                warnings.warn(
                    f"simulate_static_sweep: run {i} ended at t={t_run_end:.4g} "
                    f"but reference run ended at t={t_ref_end:.4g}. Outputs "
                    "will be interpolated (clamped) to fill the time grid.",
                    UserWarning,
                    stacklevel=2,
                )
            # Trim to the shorter time grid if lengths differ.
            t_ref_arr = jnp.asarray(time_ref)
            t_run_arr = jnp.asarray(res.time)
            if t_run_arr.shape[0] < t_ref_arr.shape[0]:
                # Shrink the reference grid and re-trim previously-collected
                # signals so all rows share the same length.
                new_len = int(t_run_arr.shape[0])
                time_ref = t_run_arr
                for name in signal_names:
                    out_lists[name] = [
                        jnp.asarray(y)[:new_len] for y in out_lists[name]
                    ]
                for name in signal_names:
                    out_lists[name].append(jnp.asarray(res.outputs[name])[:new_len])
            else:
                for name in signal_names:
                    out_lists[name].append(
                        _interp_on_time(res.outputs[name], res.time, time_ref)
                    )

    stacked = {k: jnp.stack(vs, axis=0) for k, vs in out_lists.items()}
    result = BatchSimulationResults(time=time_ref, outputs=stacked, used_vmap=False)
    # Attach per-element final contexts as a public attribute (dataclass with
    # default_factory would require redefining BatchSimulationResults; a simple
    # post-hoc attribute is sufficient and documented in the docstring).
    object.__setattr__(result, "contexts", contexts)
    return result

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).	None
options	SimulatorOptions | None	:class:SimulatorOptions forwarded per variant.	None
recorded_signals	Callable[[Diagram], dict[str, OutputPort]] | dict[str, OutputPort] | None	Either a static {name: OutputPort} dict — in which case the port references must be valid on every generated per-variant diagram — or a callable recorded_signals(variant_diagram) -> {name: OutputPort} that re-resolves the ports against the concrete diagram. The callable form is the safer choice when variants substitute entire subdiagrams.	None
results_options	ResultsOptions | None	Forwarded to :func:simulate_batch.	None
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 keyed by the variant configuration (a sorted tuple of
	dict[tuple[tuple[str, Any], ...], BatchSimulationResults | SimulationResults]	(path, choice) pairs) whose values are
	dict[tuple[tuple[str, Any], ...], BatchSimulationResults | SimulationResults]	class:BatchSimulationResults (when param_batches is set)
or	dict[tuple[tuple[str, Any], ...], BatchSimulationResults | SimulationResults]	class:SimulationResults (when it isn't).

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.

Source code in jaxonomy/simulation/simulate_variants.py

def simulate_variant_sweep(
    diagram: Diagram,
    t_span: tuple[float, float],
    *,
    param_batches: dict[str, Any] | None = None,
    options: SimulatorOptions | None = None,
    recorded_signals: Callable[[Diagram], dict[str, OutputPort]] | dict[str, OutputPort] | None = None,
    results_options: ResultsOptions | None = None,
    use_vmap: bool = False,
) -> dict[tuple[tuple[str, Any], ...], BatchSimulationResults | SimulationResults]:
    """Sweep every variant configuration of ``diagram``; for each, optionally
    sweep a parameter batch.

    Args:
        diagram: A built :class:`Diagram` containing one or more
            :class:`~jaxonomy.framework.variants.Variant` nodes.
        t_span: ``(t_start, t_stop)`` forwarded to the per-variant call.
        param_batches: 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: :class:`SimulatorOptions` forwarded per variant.
        recorded_signals: Either a static ``{name: OutputPort}`` dict —
            in which case the port references must be valid on every
            generated per-variant diagram — or a callable
            ``recorded_signals(variant_diagram) -> {name: OutputPort}``
            that re-resolves the ports against the concrete diagram. The
            callable form is the safer choice when variants substitute
            entire subdiagrams.
        results_options: Forwarded to :func:`simulate_batch`.
        use_vmap: Forwarded to :func:`simulate_batch` when
            ``param_batches`` is supplied; ignored otherwise.

    Returns:
        Dict keyed by the variant configuration (a sorted tuple of
        ``(path, choice)`` pairs) whose values are
        :class:`BatchSimulationResults` (when ``param_batches`` is set)
        or :class:`SimulationResults` (when it isn't).

    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.
    """
    if not isinstance(diagram, Diagram):
        raise TypeError(
            f"simulate_variant_sweep expects a Diagram, got {type(diagram)}"
        )

    out: dict[tuple[tuple[str, Any], ...], Any] = {}
    for config, variant_diagram in iter_variant_configurations(diagram):
        # Resolve recorded_signals against the concrete per-variant diagram.
        if callable(recorded_signals):
            rs_for_variant = recorded_signals(variant_diagram)
        else:
            rs_for_variant = recorded_signals

        if param_batches is None:
            ctx = variant_diagram.create_context()
            res = simulate(
                variant_diagram,
                ctx,
                t_span,
                options=options,
                recorded_signals=rs_for_variant,
                results_options=results_options,
            )
        else:
            res = simulate_batch(
                variant_diagram,
                t_span,
                param_batches=param_batches,
                options=options,
                recorded_signals=rs_for_variant,
                results_options=results_options,
                use_vmap=use_vmap,
            )
        out[_config_key(config)] = res
    return out

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.	required
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.	None
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.

Source code in jaxonomy/simulation/event_gradient.py

def simulate_with_event_time_grad(
    diagram,
    ctx,
    t_span,
    params,
    event_index: int,
    guard_fn: Callable[[float, Any, Any], jnp.ndarray],
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    state_at_event_fn: Callable[[Any], Any] | Any,
    options=None,
    *,
    sim_runner: Callable[..., float] | None = None,
    eps: float = 1e-30,
) -> jnp.ndarray:
    """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.

    Args:
        diagram: SystemBase passed straight to :func:`simulate`.
        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.
        t_span: ``(t0, t1)`` tuple passed to :func:`simulate`.
        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.
        event_index: Integer event slot whose firing time is returned.
        guard_fn: ``(t, state, params) -> scalar`` — zero-crossing
            guard used by the implicit-function backward rule.
        ode_rhs_fn: ``(t, state, params) -> dstate/dt`` — continuous
            RHS evaluated at the event boundary.
        state_at_event_fn: 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)``.
        sim_runner: ``(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: Floor for the implicit-function denominator (forwarded to
            :func:`event_time_gradient`).

    Returns:
        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.
    """
    runner = sim_runner if sim_runner is not None else _default_sim_runner

    # Bind all non-pytree / non-traced args via closure; ``custom_vjp``
    # only differentiates with respect to ``params``.
    @jax.custom_vjp
    def _wrapped(params_):
        return _fwd_event_time(params_, diagram, ctx, t_span,
                               event_index, options, runner)

    def _fwd(params_):
        t_e = _fwd_event_time(params_, diagram, ctx, t_span,
                              event_index, options, runner)
        # Save everything the backward rule needs.  ``params_`` flows in
        # as a JAX value; ``t_e`` is the scalar firing time.
        return t_e, (params_, t_e)

    def _bwd(residuals, cotangent):
        params_, t_e = residuals
        # Normalize the ``state_at_event_fn`` argument count so that
        # :func:`event_time_gradient` always sees its standard
        # ``state_fn(params) -> state`` form.  Critically, we bind ``t_e``
        # in as a *concrete* Python float when the user supplied the
        # ``(t_e, params)`` form — this guarantees ``∂x_e/∂p`` is taken
        # at the recorded firing instant rather than along the
        # parameter-dependent ``t_e(p)`` curve (which would zero out
        # the implicit dependence; see the note in the API docstring).
        if callable(state_at_event_fn):
            try:
                import inspect
                n_args = len(inspect.signature(state_at_event_fn).parameters)
            except (TypeError, ValueError):
                n_args = 1
            if n_args >= 2:
                # Capture ``t_e`` by closure; ``jax.lax.stop_gradient``
                # severs any cotangent path *through* ``t_e`` so that
                # the implicit-function gradient is taken at the
                # recorded firing instant rather than along the
                # ``t_e(p)`` curve.  Works whether ``t_e`` is a concrete
                # scalar (eager) or a JIT tracer (under ``jax.jit`` /
                # ``jax.vmap``) — the standalone
                # :func:`event_time_gradient` already treats ``t_event``
                # as a constant w.r.t. the differentiation parameter
                # for the same reason.
                t_e_bound = jax.lax.stop_gradient(t_e)
                def _state_fn_bound(_p, _t=t_e_bound):
                    return state_at_event_fn(_t, _p)
            else:
                _state_fn_bound = state_at_event_fn
        else:
            _state_fn_bound = state_at_event_fn  # constant state PyTree

        # Implicit-function-theorem gradient (T-125 phase 1).
        dt_dp = event_time_gradient(
            guard_fn,
            ode_rhs_fn,
            t_e,
            _state_fn_bound,
            params_,
            eps=eps,
        )
        # Chain rule: upstream cotangent (scalar) times dt_e/dp.
        scaled = jax.tree_util.tree_map(lambda g: cotangent * g, dt_dp)
        return (scaled,)

    _wrapped.defvjp(_fwd, _bwd)
    return _wrapped(params)

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.

Source code in jaxonomy/simulation/provenance.py

def verify_manifest(
    actual: ProvenanceManifest,
    expected: ProvenanceManifest,
    *,
    ignore_fields: Optional[set[str]] = None,
) -> 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).

    Args:
        actual: the manifest produced by the run being checked.
        expected: the reference manifest.
        ignore_fields: see :func:`compare_manifests`; default
            ``{"timestamp"}``.

    Raises:
        ManifestMismatch: if any compared field differs.
    """
    differences = compare_manifests(actual, expected, ignore_fields=ignore_fields)
    if differences:
        raise ManifestMismatch(differences)

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.

Source code in jaxonomy/simulation/event_gradient.py

def vmap_event_time_gradient(
    guard_fn: Callable[[float, Any, Any], jnp.ndarray],
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    t_event_array: jnp.ndarray,
    state_at_event_fn: Callable[[Any, Any], Any],
    params_batch: Any,
    *,
    eps: float = 1e-30,
    use_python_loop: bool = False,
) -> Any:
    """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`.

    Args:
        guard_fn: ``(t, state, params) -> scalar`` — zero-crossing guard.
            Same contract as :func:`event_time_gradient`; shared across
            the batch.
        ode_rhs_fn: ``(t, state, params) -> dstate/dt`` — RHS at the
            event boundary; shared across the batch.
        t_event_array: ``(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`.
        state_at_event_fn: ``(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``.
        params_batch: 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.
        eps: Forwarded to :func:`event_time_gradient` — denominator floor
            for grazing crossings.
        use_python_loop: 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).

    Returns:
        Per-sample gradients with the same leading axis as
        ``params_batch``.  PyTree structure of each sample's gradient
        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.
    """
    t_event_array = jnp.asarray(t_event_array)

    def _single_sample(t_e, params_i):
        # Bind ``t_e`` into the state callable so the inner helper sees
        # the canonical ``state_fn(params) -> state`` form.  Use
        # ``stop_gradient`` so the implicit-function gradient is taken
        # at the recorded firing instant rather than along the
        # ``t_e(params)`` curve — same convention as the custom-VJP
        # wrapper.
        t_e_const = jax.lax.stop_gradient(t_e)

        def _state_fn(p, _t=t_e_const):
            return state_at_event_fn(_t, p)

        return event_time_gradient(
            guard_fn,
            ode_rhs_fn,
            t_e,
            _state_fn,
            params_i,
            eps=eps,
        )

    if use_python_loop:
        # Honest fallback path — iterate the sample axis in Python and
        # stack the per-sample gradients leafwise.  Slower than vmap but
        # robust against any future vmap-composition regression inside
        # :func:`event_time_gradient`.
        n = int(t_event_array.shape[0])
        per_sample = []
        for i in range(n):
            params_i = jax.tree_util.tree_map(lambda leaf, _i=i: leaf[_i], params_batch)
            per_sample.append(_single_sample(t_event_array[i], params_i))
        if n == 0:
            # Empty batch — produce a structurally-correct empty leading
            # axis by computing a single dummy gradient and slicing it
            # off.  Mirrors :func:`event_times_gradient`'s empty-firing
            # convention.
            template_params = jax.tree_util.tree_map(
                lambda leaf: leaf[:1].reshape((1,) + leaf.shape[1:]) if leaf.ndim >= 1 else leaf,
                params_batch,
            )
            template_params_0 = jax.tree_util.tree_map(lambda leaf: leaf[0], template_params)
            template = _single_sample(jnp.asarray(0.0), template_params_0)
            return jax.tree_util.tree_map(lambda leaf: leaf[None][:0], template)
        return jax.tree_util.tree_map(
            lambda *leaves: jnp.stack(leaves, axis=0),
            *per_sample,
        )

    # Default path: jax.vmap.  Sample axis is the leading axis of every
    # leaf of ``params_batch`` and of ``t_event_array``.
    return jax.vmap(_single_sample, in_axes=(0, 0))(t_event_array, params_batch)

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).

Source code in jaxonomy/simulation/event_gradient.py

def vmap_event_times_gradient(
    results: Any,
    params_batch: Any,
    guards: Any,
    ode_rhs_fn: Callable[[float, Any, Any], Any],
    state_at_event_fn: Callable[[float, Any], Any],
    *,
    event_indices: Any = None,
    eps: float = 1e-30,
    use_python_loop: bool = False,
) -> dict:
    """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]``.

    Args:
        results: 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.
        params_batch: Batched parameter PyTree.  All leaves must share a
            leading axis of length ``N``.  Same contract as
            :func:`vmap_event_time_gradient`.
        guards: 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`.
        ode_rhs_fn: ``(t, state, params) -> dstate/dt`` — shared across
            firings and samples.
        state_at_event_fn: ``(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`.
        event_indices: 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``.
        eps: Forwarded to :func:`event_time_gradient` — denominator floor
            for grazing crossings.
        use_python_loop: When ``True``, iterate explicitly over both the
            firing and sample axes in Python.  Slower but byte-identical;
            honest fallback when vmap composition is invasive.

    Returns:
        ``{event_index: gradient_batch}`` — one entry per processed
        event index.  For each entry, leaves carry leading axes
        ``(N, n_firings, ...)``.  Empty firing lists yield a structurally
        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`).
    """
    event_times_dict = getattr(results, "event_times", None)
    if event_times_dict is None:
        raise ValueError(
            "vmap_event_times_gradient: results.event_times is None. "
            "Re-run simulate(...) with "
            "SimulatorOptions(record_event_times=True) so the firing "
            "instants are captured."
        )

    if callable(guards):
        def _guard_for(_idx):  # noqa: ANN001
            return guards
    else:
        guards_map = dict(guards)
        def _guard_for(idx):
            if idx not in guards_map:
                raise KeyError(
                    f"vmap_event_times_gradient: no guard supplied for "
                    f"event index {idx}.  Provide guards[{idx}] = <fn> "
                    f"or pass a single callable to apply uniformly."
                )
            return guards_map[idx]

    if event_indices is None:
        selected = list(event_times_dict.keys())
    else:
        selected = list(event_indices)
        for idx in selected:
            if idx not in event_times_dict:
                raise KeyError(
                    f"vmap_event_times_gradient: event index {idx} not "
                    f"present in results.event_times (have: "
                    f"{sorted(event_times_dict.keys())})."
                )

    # Validate that ``params_batch`` carries a sample axis on every leaf
    # and infer ``N`` from the first leaf.  Mirrors the in_axes=0 contract
    # of :func:`vmap_event_time_gradient`.
    leaves = jax.tree_util.tree_leaves(params_batch)
    if not leaves:
        raise ValueError(
            "vmap_event_times_gradient: params_batch has no array leaves; "
            "cannot infer batch size."
        )
    n_samples = int(jnp.asarray(leaves[0]).shape[0])

    out: dict = {}
    for idx in selected:
        firings = jnp.asarray(event_times_dict[idx])
        guard_fn = _guard_for(idx)
        n_firings = int(firings.shape[0]) if firings.ndim >= 1 else 0

        if n_firings == 0:
            # Empty firing set — produce a structurally-correct
            # ``(N, 0, ...)`` leading pair by computing a single dummy
            # gradient via :func:`vmap_event_time_gradient` at ``t_e=0``
            # for every sample, then slicing out the firing axis.
            dummy_t = jnp.zeros((n_samples,), dtype=jnp.float64)
            template_batch = vmap_event_time_gradient(
                guard_fn,
                ode_rhs_fn,
                dummy_t,
                state_at_event_fn,
                params_batch,
                eps=eps,
                use_python_loop=use_python_loop,
            )
            # template_batch leaves have shape (N, ...).  Insert an empty
            # firing axis at position 1 -> (N, 0, ...).
            out[idx] = jax.tree_util.tree_map(
                lambda leaf: leaf[:, None, ...][:, :0, ...],
                template_batch,
            )
            continue

        # Per-firing pass: for each firing instant, vmap the single-shot
        # gradient over the sample axis.  Stack the per-firing results
        # leafwise along axis=1 so the final leaf shape is
        # ``(N, n_firings, ...)``.
        per_firing_batches = []
        for k in range(n_firings):
            t_e_k = firings[k]
            # Broadcast the scalar firing time across the sample axis so
            # the inner vmap sees a ``(N,)`` ``t_event_array`` (matching
            # the leaf shape contract of :func:`vmap_event_time_gradient`).
            t_e_broadcast = jnp.broadcast_to(t_e_k, (n_samples,))

            grad_batch_k = vmap_event_time_gradient(
                guard_fn,
                ode_rhs_fn,
                t_e_broadcast,
                state_at_event_fn,
                params_batch,
                eps=eps,
                use_python_loop=use_python_loop,
            )
            per_firing_batches.append(grad_batch_k)

        # Stack along the new firing axis (position 1, since position 0 is
        # the sample axis).  ``tree_map(*pytrees)`` operates leafwise.
        out[idx] = jax.tree_util.tree_map(
            lambda *leaves: jnp.stack(leaves, axis=1),
            *per_firing_batches,
        )

    return out