Framework

`jaxonomy.framework`

`BlockInitializationError`

Bases: JaxonomyError

A generic error to be thrown when a block fails at init time, but the full exceptions are known to cause issues, eg. with ray serialization.

Source code in jaxonomy/framework/error.py

class BlockInitializationError(JaxonomyError):
    """A generic error to be thrown when a block fails at init time, but
    the full exceptions are known to cause issues, eg. with ray serialization.
    """

    pass

`BlockParameterError`

Bases: StaticError

Block parameters are missing or have invalid values.

Source code in jaxonomy/framework/error.py

class BlockParameterError(StaticError):
    """Block parameters are missing or have invalid values."""

    pass

`BlockRuntimeError`

Bases: JaxonomyError

A generic error to be thrown when a block fails at runtime, but the full exceptions are known to cause issues, eg. with ray serialization.

Source code in jaxonomy/framework/error.py

class BlockRuntimeError(JaxonomyError):
    """A generic error to be thrown when a block fails at runtime, but
    the full exceptions are known to cause issues, eg. with ray serialization.
    """

    pass

`BusUnit` `dataclass`

Compound unit carrying one :class:Unit per named bus field.

Attached to the output port of a :class:BusCreator (and the matching input port of a :class:BusSelector) so that the connect-time consistency check can verify each field's unit individually.

Attributes:

Name	Type	Description
`fields`	`Mapping[str, Unit]`	Mapping from bus field name to its :class:`Unit`. Stored as a plain `dict` (insertion order preserved).

Source code in jaxonomy/framework/units.py

@dataclass(frozen=True, eq=False)
class BusUnit:
    """Compound unit carrying one :class:`Unit` per named bus field.

    Attached to the output port of a :class:`BusCreator` (and the
    matching input port of a :class:`BusSelector`) so that the
    connect-time consistency check can verify each field's unit
    individually.

    Attributes:
        fields: Mapping from bus field name to its :class:`Unit`.
            Stored as a plain ``dict`` (insertion order preserved).
    """

    fields: Mapping[str, Unit] = field(default_factory=dict)

    def __post_init__(self):
        # Normalise to a plain dict so the class is hashable and the
        # ordering is deterministic. Validate each entry is a Unit.
        normalised: dict[str, Unit] = {}
        for k, v in dict(self.fields).items():
            if not isinstance(k, str):
                raise TypeError(
                    f"BusUnit field name must be a str, got {type(k).__name__}: {k!r}"
                )
            if not isinstance(v, Unit):
                raise TypeError(
                    f"BusUnit field {k!r} must map to a Unit instance, "
                    f"got {type(v).__name__}: {v!r}"
                )
            normalised[k] = v
        object.__setattr__(self, "fields", normalised)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, BusUnit):
            return NotImplemented
        if set(self.fields.keys()) != set(other.fields.keys()):
            return False
        return all(self.fields[k] == other.fields[k] for k in self.fields)

    def __hash__(self) -> int:
        return hash(tuple(sorted(self.fields.items(), key=lambda kv: kv[0])))

    def __repr__(self) -> str:
        body = ", ".join(f"{k}={v!r}" for k, v in self.fields.items())
        return f"BusUnit({{{body}}})"

    def field_unit(self, name: str) -> Unit | None:
        """Return the :class:`Unit` for ``name``, or ``None`` if absent.

        Used by :class:`BusSelector` to look up its output-port unit
        when wired downstream of a unit-tagged bus.
        """
        return self.fields.get(name)

`field_unit(name)`

Return the :class:Unit for name, or None if absent.

Used by :class:BusSelector to look up its output-port unit when wired downstream of a unit-tagged bus.

Source code in jaxonomy/framework/units.py

def field_unit(self, name: str) -> Unit | None:
    """Return the :class:`Unit` for ``name``, or ``None`` if absent.

    Used by :class:`BusSelector` to look up its output-port unit
    when wired downstream of a unit-tagged bus.
    """
    return self.fields.get(name)

`ContextBase` `dataclass`

Context object containing state, parameters, etc for a system.

NOTE: Type hints in ContextBase indicate the union between what would be returned by a LeafContext and a DiagramContext. See type hints of the subclasses for the specific argument and return types.

Attributes:

Name	Type	Description
`owning_system`	`SystemBase`	The owning system of the context.
`time`	`Scalar`	The time associated with the context. Will be None unless the context is the root context.
`is_initialized`	`bool`	Flag indicating if the context is initialized. This should only be set by the ContextFactory during creation.

Source code in jaxonomy/framework/context.py

@dataclasses.dataclass(frozen=True)
class ContextBase(metaclass=abc.ABCMeta):
    """Context object containing state, parameters, etc for a system.

    NOTE: Type hints in ContextBase indicate the union between what would be returned
    by a LeafContext and a DiagramContext. See type hints of the subclasses for
    the specific argument and return types.

    Attributes:
        owning_system (SystemBase):
            The owning system of the context.
        time (Scalar):
            The time associated with the context. Will be None unless the context
            is the root context.
        is_initialized (bool):
            Flag indicating if the context is initialized. This should only be set
            by the ContextFactory during creation.
    """

    owning_system: SystemBase
    time: Optional[Scalar] = None
    is_initialized: bool = False
    parameters: Optional[Mapping[str, Array]] = None
    port_cache: Optional[PortCache] = None

    @abc.abstractmethod
    def __getitem__(self, key: Hashable) -> LeafContext:
        """Get the subcontext associated with the given system ID.

        For leaf contexts, this will return `self`, but the method is provided
        so that there is a consistent interface for working with either an
        individual LeafSystem or tree-structured Diagram.

        For nested diagrams, intermediate diagrams do not have associated contexts,
        so indexing will fail.
        """
        pass

    @abc.abstractmethod
    def with_subcontext(self, key: Hashable, ctx: LeafContext) -> ContextBase:
        """Create a copy of this context, replacing the specified subcontext."""
        pass

    def with_time(self, value: Scalar) -> ContextBase:
        """Create a copy of this context, replacing time with the given value.

        This should only be called on the root context, since it is expected that all
        subcontexts will have a time value of None to avoid any conflicts.
        """
        return dataclasses.replace(self, time=value)

    @property
    @abc.abstractmethod
    def state(self) -> State:
        pass

    @abc.abstractmethod
    def with_state(self, state: State) -> ContextBase:
        """Create a copy of this context, replacing the entire state."""
        pass

    @abc.abstractmethod
    def with_new_state(self) -> ContextBase:
        """Create a copy of this context, replacing the state with a new state."""
        pass

    @property
    @abc.abstractmethod
    def continuous_state(self) -> StateComponent:
        pass

    @abc.abstractmethod
    def with_continuous_state(self, value: StateComponent) -> ContextBase:
        """Create a copy of this context, replacing the continuous state."""
        pass

    @property
    @abc.abstractmethod
    def num_continuous_states(self) -> int:
        pass

    @property
    @abc.abstractmethod
    def has_continuous_state(self) -> bool:
        pass

    @property
    @abc.abstractmethod
    def discrete_state(self) -> StateComponent:
        pass

    @abc.abstractmethod
    def with_discrete_state(self, value: StateComponent) -> ContextBase:
        """Create a copy of this context, replacing the discrete state."""
        pass

    @property
    @abc.abstractmethod
    def num_discrete_states(self) -> int:
        pass

    @property
    @abc.abstractmethod
    def has_discrete_state(self) -> bool:
        pass

    @property
    @abc.abstractmethod
    def mode(self) -> Mode:
        pass

    @property
    @abc.abstractmethod
    def has_mode(self) -> bool:
        pass

    @abc.abstractmethod
    def with_mode(self, value: Mode) -> ContextBase:
        """Create a copy of this context, replacing the mode."""
        pass

    def mark_initialized(self) -> ContextBase:
        return dataclasses.replace(self, is_initialized=True)

    @abc.abstractmethod
    def with_updated_parameters(self) -> ContextBase:
        """Create a copy of this context, updating all parameters to their current values."""
        pass

    def with_parameter(self, name: str, value: ArrayLike) -> ContextBase:
        """Create a copy of this context, replacing the specified parameter."""
        return self.with_parameters({name: value})

    @abc.abstractmethod
    def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
        """Create a copy of this context, replacing only the specified parameters."""
        pass

    def with_port_cache(self, cache: PortCache) -> ContextBase:
        return dataclasses.replace(self, port_cache=cache)

    def with_port_cache_entry(self, key: Hashable, val: Array) -> ContextBase:
        return dataclasses.replace(self, port_cache={**self.port_cache, key: val})

    def refresh_port_cache(self) -> ContextBase:
        if not self.owning_system.cache_enabled:
            return self
        return self.owning_system.recompute_port_cache(self)

`getitem(key)` `abstractmethod`

Get the subcontext associated with the given system ID.

For leaf contexts, this will return self, but the method is provided so that there is a consistent interface for working with either an individual LeafSystem or tree-structured Diagram.

For nested diagrams, intermediate diagrams do not have associated contexts, so indexing will fail.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def __getitem__(self, key: Hashable) -> LeafContext:
    """Get the subcontext associated with the given system ID.

    For leaf contexts, this will return `self`, but the method is provided
    so that there is a consistent interface for working with either an
    individual LeafSystem or tree-structured Diagram.

    For nested diagrams, intermediate diagrams do not have associated contexts,
    so indexing will fail.
    """
    pass

`with_continuous_state(value)` `abstractmethod`

Create a copy of this context, replacing the continuous state.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_continuous_state(self, value: StateComponent) -> ContextBase:
    """Create a copy of this context, replacing the continuous state."""
    pass

`with_discrete_state(value)` `abstractmethod`

Create a copy of this context, replacing the discrete state.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_discrete_state(self, value: StateComponent) -> ContextBase:
    """Create a copy of this context, replacing the discrete state."""
    pass

`with_mode(value)` `abstractmethod`

Create a copy of this context, replacing the mode.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_mode(self, value: Mode) -> ContextBase:
    """Create a copy of this context, replacing the mode."""
    pass

`with_new_state()` `abstractmethod`

Create a copy of this context, replacing the state with a new state.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_new_state(self) -> ContextBase:
    """Create a copy of this context, replacing the state with a new state."""
    pass

`with_parameter(name, value)`

Create a copy of this context, replacing the specified parameter.

Source code in jaxonomy/framework/context.py

def with_parameter(self, name: str, value: ArrayLike) -> ContextBase:
    """Create a copy of this context, replacing the specified parameter."""
    return self.with_parameters({name: value})

`with_parameters(new_parameters)` `abstractmethod`

Create a copy of this context, replacing only the specified parameters.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
    """Create a copy of this context, replacing only the specified parameters."""
    pass

`with_state(state)` `abstractmethod`

Create a copy of this context, replacing the entire state.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_state(self, state: State) -> ContextBase:
    """Create a copy of this context, replacing the entire state."""
    pass

`with_subcontext(key, ctx)` `abstractmethod`

Create a copy of this context, replacing the specified subcontext.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_subcontext(self, key: Hashable, ctx: LeafContext) -> ContextBase:
    """Create a copy of this context, replacing the specified subcontext."""
    pass

`with_time(value)`

Create a copy of this context, replacing time with the given value.

This should only be called on the root context, since it is expected that all subcontexts will have a time value of None to avoid any conflicts.

Source code in jaxonomy/framework/context.py

def with_time(self, value: Scalar) -> ContextBase:
    """Create a copy of this context, replacing time with the given value.

    This should only be called on the root context, since it is expected that all
    subcontexts will have a time value of None to avoid any conflicts.
    """
    return dataclasses.replace(self, time=value)

`with_updated_parameters()` `abstractmethod`

Create a copy of this context, updating all parameters to their current values.

Source code in jaxonomy/framework/context.py

@abc.abstractmethod
def with_updated_parameters(self) -> ContextBase:
    """Create a copy of this context, updating all parameters to their current values."""
    pass

`DependencyTicket`

Singleton class for managing unique dependency tickets.

Source code in jaxonomy/framework/dependency_graph.py

class DependencyTicket:
    """Singleton class for managing unique dependency tickets."""

    nothing = 0  # Indicates "not dependent on anything".
    time = 1  # Time.
    xc = 2  # All continuous state variables.
    xd = 3  # All discrete state variables
    mode = 4  # All modes.
    x = 5  # All state variables x = {xc, xd, mode}.
    p = 6  # All parameters
    all_sources_except_input_ports = 7  # Everything except input ports.
    u = 8  # All input ports u.
    all_sources = 9  # All of the above.
    xcdot = 10  # Continuous state time derivative

    _next_available = 11  # This will get incremented by next_available_ticket().

    @classmethod
    def next_available_ticket(cls):
        cls._next_available += 1
        return cls._next_available

`Diagram` `dataclass`

Bases: SystemBase

Composite block-diagram representation of a dynamical system.

A Diagram is a collection of Systems connected together to form a larger hybrid dynamical system. Diagrams can be nested to any depth, creating a tree-structured block diagram.

NOTE: The Diagram class is not intended to be constructed directly. Instead, use the DiagramBuilder to construct a Diagram, which will pass the appropriate information to this constructor.

Source code in jaxonomy/framework/diagram.py

@dataclasses.dataclass
class Diagram(SystemBase):
    """Composite block-diagram representation of a dynamical system.

    A Diagram is a collection of Systems connected together to form a larger hybrid
    dynamical system. Diagrams can be nested to any depth, creating a tree-structured
    block diagram.

    NOTE: The Diagram class is not intended to be constructed directly.  Instead,
    use the `DiagramBuilder` to construct a Diagram, which will pass the appropriate
    information to this constructor.
    """

    # Redefine here to make pylint happy
    system_id: Hashable = dataclasses.field(default_factory=next_system_id, init=False)
    name: str = None  # Human-readable name for this system (optional)
    ui_id: str = None  # UUID of the block when loaded from JSON (optional)

    # None of these attributes are intended to be modified or accessed directly after
    # construction.  Instead, use the interface defined by `SystemBase`.

    # Direct children of this Diagram
    nodes: List[SystemBase] = dataclasses.field(default_factory=list)

    # Mapping from input ports to output ports of child subsystems
    connection_map: Mapping[InputPortLocator, OutputPortLocator] = dataclasses.field(
        default_factory=dict,
    )

    # Optional identifier for "reference diagrams"
    ref_id: str = None

    # for serialization
    instance_parameters: set[str] = dataclasses.field(default_factory=set)

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self.name}, {len(self.nodes)} nodes)"

    def _pprint(self, prefix="", fancy=True) -> str:
        if fancy:
            return pprint_fancy(prefix, self)
        return f"{prefix}|-- {self.name}\n"

    def _pprint_helper(self, prefix="", fancy=True) -> str:
        s = self._pprint(prefix=prefix, fancy=fancy)
        for _, substate in enumerate(self.nodes):
            s += substate._pprint_helper(prefix=f"{prefix}    ", fancy=fancy)
        return s

    def __hash__(self) -> Hashable:
        return hash(self.system_id)

    def __getitem__(self, name: str) -> SystemBase:
        # Access by name - for convenient user interface only.  Programmatic
        #  access should use the nodes directly, e.g. `self.nodes[idx]`
        lookup = {node.name: node for node in self.nodes}
        return lookup[name]

    def __iter__(self) -> Iterator[SystemBase]:
        return iter(self.nodes)

    def print_schedule(
        self,
        *,
        format: str = "text",
        file=None,
        ensure_initialized: bool = True,
    ) -> None:
        """Print the inferred sample-time schedule of this diagram.

        Renders one entry per rate group — period (for discrete blocks),
        the leaves that fire at that rate, any detected rate mismatches,
        and the deterministic execution order. Pure inspection helper;
        does not mutate the diagram or its contexts.

        T-105-followup-print-schedule-pre-context: by default,
        ``print_schedule()`` lazily calls :meth:`create_context()` first
        so that discrete blocks whose periodic events are registered in
        their :meth:`initialize` hook (e.g. :class:`PIDDiscrete`,
        :class:`Decimator`, :class:`UnitDelay`, :class:`ZeroOrderHold`)
        are bucketed into the correct rate group. Pre-fix those blocks
        showed up as ``constant`` when ``print_schedule()`` ran before
        any context had been created, because the periodic event hadn't
        yet been declared. The lazy ``create_context()`` call is
        idempotent and cheap on already-initialised diagrams; if it
        fails (e.g. the diagram has missing connections), the schedule
        is rendered anyway with a one-line warning explaining the
        rate-group output may be incomplete. Pass
        ``ensure_initialized=False`` to opt out and use the pre-fix
        behaviour (useful for debugging the initialisation path
        itself).

        Args:
            format: ``"text"`` (default), ``"markdown"``, or ``"json"``.
            file: Destination file-like object. Defaults to ``sys.stdout``
                when None.
            ensure_initialized: If True (default), call
                :meth:`create_context` once before rendering so all
                discrete blocks have registered their periodic events.
                Set to False to render against the current
                (possibly pre-init) state.

        Examples:
            >>> diagram.print_schedule()
            rate groups:
              discrete(period=0.001, offset=0.0): ctrl_inner
              discrete(period=0.01,  offset=0.0): ctrl_outer
              continuous: plant
            execution order: ...

            >>> with open("schedule.md", "w") as f:
            ...     diagram.print_schedule(format="markdown", file=f)

        See also:
            :func:`jaxonomy.simulation.rate_groups.rate_summary` — the
                underlying string formatter, useful when you want the
                result as a string for embedding in a manifest or PR body.
            :func:`jaxonomy.simulation.rate_groups.rate_summary_dot` — the
                DOT-format companion for graphviz visualization.
        """
        import sys
        import warnings
        from ..simulation.rate_groups import rate_summary

        if ensure_initialized:
            try:
                # Drop the returned context — we don't need it; we just
                # want the side effect of running every leaf's
                # ``initialize()`` so periodic events register.
                self.create_context()
            except Exception as exc:  # noqa: BLE001
                warnings.warn(
                    "print_schedule(): create_context() failed during "
                    f"the lazy initialisation step ({type(exc).__name__}: "
                    f"{exc}). Rendering the rate-group summary anyway, "
                    "but blocks whose periodic events are registered in "
                    "initialize() (PIDDiscrete, Decimator, UnitDelay, "
                    "ZeroOrderHold) may be misbucketed as 'constant'. "
                    "Fix the underlying error and re-run, or pass "
                    "ensure_initialized=False to silence this warning.",
                    UserWarning,
                    stacklevel=2,
                )

        text = rate_summary(self, format=format)
        print(text, file=file if file is not None else sys.stdout)

    def __post_init__(self):
        super().__post_init__()

        # Set parent for all the immediate child systems
        for node in self.nodes:
            node.parent = self

        # The map of subsystem inputs/outputs to inputs/outputs of this Diagram.
        self._input_port_map: Mapping[InputPortLocator, int] = {}
        self._output_port_map: Mapping[OutputPortLocator, int] = {}

        # Also need the inverse output map, for determining feedthrough paths.
        self._inv_output_port_map: Mapping[int, OutputPortLocator] = {}

        # Leaves of the system tree (not necessarily the same as the direct
        # children of this Diagram, which may themselves be Diagrams)
        self.leaf_systems: List[LeafSystem] = []
        for sys in self.nodes:
            if isinstance(sys, Diagram):
                # FIXME: In case of 'param estimation' optimization run, we end
                # up here and sys.leaf_systems is now None. Using or [] fixes
                # the crash but something is a bit fishy.
                self.leaf_systems.extend(sys.leaf_systems or [])
                # No longer need the child leaf systems, since methods using this
                # should only be called from the top level.
                sys.leaf_systems = None
            else:
                self.leaf_systems.append(sys)

    def post_simulation_finalize(self) -> None:
        """Perform any post-simulation cleanup for this system."""
        for system in self.nodes:
            system.post_simulation_finalize()

    # Inherits docstrings from SystemBase
    @property
    def has_feedthrough_side_effects(self) -> bool:
        # See explanation in `SystemBase.has_feedthrough_side_effects`.
        return any(sys.has_feedthrough_side_effects for sys in self.nodes)

    # Inherits docstrings from SystemBase
    @property
    def has_ode_side_effects(self) -> bool:
        # Return true if either of the following are true:
        # 1. At least one subsystem has ODE side effects
        # 2. At least one subsystem has feedthrough side effects and the output
        #    ports of the diagram are used as ODE inputs.

        if self.dependency_graph is None:
            raise ValueError("Must create dependency graph first.")

        # If no subsystems have feedthrough side effects, we're done.
        if not self.has_feedthrough_side_effects:
            return False

        # If any subsystem is already known to have this property, we're done.
        if any(sys.has_ode_side_effects for sys in self.nodes):
            return True

        # If we get here, we need to actually test the dependency graph.
        for sys in self.nodes:
            if sys.has_feedthrough_side_effects:
                for port in sys.output_ports:
                    tracker = port.tracker
                    if tracker.is_prerequisite_of([DependencyTicket.xcdot]):
                        return True
        return False

    @property
    def has_continuous_state(self) -> bool:
        return any(sys.has_continuous_state for sys in self.nodes)

    @property
    def has_discrete_state(self) -> bool:
        return any(sys.has_discrete_state for sys in self.nodes)

    @property
    def has_zero_crossing_events(self) -> bool:
        return any(sys.has_zero_crossing_events for sys in self.nodes)

    @property
    def num_systems(self) -> int:
        # Number of subsystems _at this level_
        return len(self.nodes)

    def check_types(
        self,
        context: DiagramContext,
        error_collector: ErrorCollector = None,
    ) -> None:
        """Perform any system-specific static analysis."""
        for system in self.nodes:
            system.check_types(
                context,
                error_collector=error_collector,
            )

    #
    # Simulation interface
    #

    # Inherits docstrings from SystemBase
    def eval_time_derivatives(self, root_context: DiagramContext) -> List[Array]:
        leaf_systems = [
            subctx.owning_system for subctx in root_context.continuous_subcontexts
        ]
        return [sys.eval_time_derivatives(root_context) for sys in leaf_systems]

    @property
    def mass_matrix(self) -> List[Array]:
        return [sys.mass_matrix for sys in self.leaf_systems]

    @property
    def has_mass_matrix(self) -> bool:
        return any(sys.has_mass_matrix for sys in self.leaf_systems)

    @property
    def continuous_substep_vector(self) -> List[Array]:
        """T-133: per-leaf multirate substep vectors, in ``leaf_systems``
        order (the same ordering ``mass_matrix`` relies on for alignment
        with the flattened continuous state)."""
        return [sys.continuous_substep_vector for sys in self.leaf_systems]

    @property
    def has_multirate_substeps(self) -> bool:
        return any(sys.has_multirate_substeps for sys in self.leaf_systems)

    #
    # Event handling
    #
    @property
    def state_update_events(self) -> FlatEventCollection:
        assert self.parent is None, (
            "Can only get periodic events from top-level Diagram, not "
            f"{self.system_id} with parent {self.parent.system_id}"
        )
        events = sum(
            [sys.state_update_events for sys in self.leaf_systems],
            start=FlatEventCollection(),
        )
        return events

    @property
    def zero_crossing_events(self) -> DiagramEventCollection:
        assert self.parent is None, (
            "Can only get zero-crossing events from top-level Diagram, not "
            f"{self.system_id} with parent {self.parent.system_id}"
        )
        return DiagramEventCollection(
            OrderedDict(
                {sys.system_id: sys.zero_crossing_events for sys in self.leaf_systems}
            )
        )

    # Inherits docstrings from SystemBase
    def determine_active_guards(
        self, root_context: DiagramContext
    ) -> DiagramEventCollection:
        assert self.parent is None, (
            "Can only get zero-crossing events from top-level Diagram, not "
            f"{self.system_id} with parent {self.parent.system_id}"
        )
        return DiagramEventCollection(
            OrderedDict(
                {
                    sys.system_id: sys.determine_active_guards(root_context)
                    for sys in self.leaf_systems
                }
            )
        )

    # Inherits docstrings from SystemBase
    def eval_zero_crossing_updates(
        self,
        root_context: DiagramContext,
        events: DiagramEventCollection,
    ) -> dict[Hashable, LeafState]:
        substates = OrderedDict()
        for system_id, subctx in root_context.subcontexts.items():
            sys = subctx.owning_system
            substates[system_id] = sys.eval_zero_crossing_updates(root_context, events)

        return substates

    #
    # I/O ports
    #
    @property
    def _flat_callbacks(self) -> List[SystemCallback]:
        """Return a flat list of all SystemCallbacks in the Diagram."""
        return [cb for sys in self.nodes for cb in sys._flat_callbacks]

    @property
    def exported_input_ports(self):
        return self._input_port_map

    @property
    def exported_output_ports(self):
        return self._output_port_map

    def eval_subsystem_input_port(
        self, context: DiagramContext, port_locator: InputPortLocator
    ) -> Array:
        """Evaluate the input port for a child of this system given the root context.

        Args:
            context (ContextBase): root context for this system
            port_locator (InputPortLocator): tuple of (system, port_index) identifying
                the input port to evaluate

        Returns:
            Array: Value returned from evaluating the subsystem port.

        Raises:
            InputNotConnectedError: if the input port is not connected
        """

        is_exported = port_locator in self._input_port_map
        if is_exported:
            # The upstream source is an input to this whole Diagram; evaluate that
            # input port and use the result as the value for this one.
            port_index = self._input_port_map[port_locator]  # Diagram-level index
            return self.input_ports[port_index].eval(context)  # Return upstream value

        is_connected = port_locator in self.connection_map
        if is_connected:
            # The upstream source is an output port of one of this Diagram's child
            # subsystems; evaluate the upstream output.
            upstream_locator = self.connection_map[port_locator]

            # This will return the value of the upstream port
            return self.eval_subsystem_output_port(context, upstream_locator)

        block, port_index = port_locator
        raise InputNotConnectedError(
            system=block,
            port_index=port_index,
            port_direction="in",
            message=f"Input port {block.name}[{port_index}] is not connected",
        )

    def eval_subsystem_output_port(
        self, context: DiagramContext, port_locator: OutputPortLocator
    ) -> Array:
        """ "Evaluate the output port for a child of this system given the root context.

        Args:
            context (ContextBase): root context for this system
            port_locator (OutputPortLocator): tuple of (system, port_index) identifying
                the output port to evaluate

        Returns:
            Array: Value returned from evaluating the subsystem port.
        """
        system, port_index = port_locator
        port = system.output_ports[port_index]

        # During simulation all we should need to do is evaluate the port.
        if context.is_initialized:
            return port.eval(context)

        # If the context is not initialized, we have to determine the signal data type.
        # In the easy case, the port has a default value, so we can just use that.
        if port.default_value is not None:
            logger.debug(
                "Using default output value of %s for %s",
                port.default_value,
                port_locator[0].name,
            )
            return port.default_value

        logger.debug(
            "Evaluating output port %s for system %s. Context initialized: %s",
            port_locator,
            port_locator[0].name,
            context.is_initialized,
        )

        # If there is no default value, try to evaluate the port to pull a "template"
        # value with an appropriate data type from upstream.  This will return None if
        # the port is not yet connected (e.g. if its upstream is an exported input of)
        # a Diagram, so we can defer evaluation.

        # Try again to evaluate the port
        val = port.eval(context)
        logger.debug(
            "  ---> %s returns %s", (port_locator[0].name, port_locator[1]), val
        )

        # If there is still no value, the port is not connected to anything.
        # Post-initialization this would be an error, but pre-initialization
        # it may be the case that the upstream is an exported input port of
        # the Diagram, so we can defer evaluation. Expect the block that is
        # doing this to handle the UpstreamEvalError appropriately.
        if val is None:
            system_name = system.name_path_str
            logger.debug(
                "Upstream evaluation of %s.out[%s] returned None. Deferring evaluation.",
                system_name,
                port_index,
            )
            raise UpstreamEvalError(port_locator=(system, "out", port_index))
        return val

    #
    # System-level declarations (should be done via DiagramBuilder)
    #
    def export_input(self, locator: InputPortLocator, port_name: str) -> int:
        """Export a subsystem input port as a diagram-level input.

        This should typically only be called during construction by DiagramBuilder.
        The standard workflow will be to call export_input on the _builder_ object,
        which will automatically call this method on the Diagram once created.

        Args:
            locator (InputPortLocator): tuple of (system, port_index) identifying
                the input port to export
            port_name (str): name of the new exported input port

        Returns:
            int: index of the exported input port in the diagram input_ports list
        """
        diagram_port_index = self.declare_input_port(name=port_name)
        self._input_port_map[locator] = diagram_port_index

        return diagram_port_index

    def export_output(self, locator: OutputPortLocator, port_name: str) -> int:
        """Export a subsystem output port as a diagram-level output.

        This should typically only be called during construction by DiagramBuilder.
        The standard workflow will be to call export_input on the _builder_ object,
        which will automatically call this method on the Diagram once created.

        Args:
            locator (OutputPortLocator): tuple of (system, port_index) identifying
                the output port to export
            port_name (str): name of the new exported output port

        Returns:
            int: index of the exported output port in the diagram output_ports list
        """
        subsystem, subsystem_port_index = locator
        source_port = subsystem.output_ports[subsystem_port_index]
        diagram_port_index = self.declare_output_port(
            source_port.eval,
            name=port_name,
            prerequisites_of_calc=[source_port.ticket],
        )
        self._output_port_map[locator] = diagram_port_index
        self._inv_output_port_map[diagram_port_index] = locator

        return diagram_port_index

    #
    # Initialization
    #
    @property
    def context_factory(self) -> DiagramContextFactory:
        return DiagramContextFactory(self)

    @property
    def dependency_graph_factory(self) -> DiagramDependencyGraphFactory:
        return DiagramDependencyGraphFactory(self)

    def initialize_static_data(self, context: DiagramContext) -> DiagramContext:
        """Perform any system-specific static analysis."""
        for system in self.nodes:
            context = system.initialize_static_data(context)
        return context

    def _has_feedthrough(self, input_port_index: int, output_port_index: int) -> bool:
        """Check if there is a direct-feedthrough path from the input port to the output port.

        Internal function used by `get_feedthrough`.  Should not typically need to
        be called directly.
        """
        # TODO: Would this be simpler if the input port map was inverted?
        input_ids = []
        for locator, index in self._input_port_map.items():
            if index == input_port_index:
                input_ids.append(locator)

        input_ids = set(input_ids)

        # Search graph for a direct-feedthrough connection from the output_port
        # to the input_port.  Maintain a set of the output port identifiers that
        # are known to have a direct-feedthrough path to the output_port
        active_set: Set[OutputPortLocator] = set()
        active_set.add(self._inv_output_port_map[output_port_index])

        while len(active_set) > 0:
            sys, sys_output = active_set.pop()
            for u, v in sys.get_feedthrough():
                if v == sys_output:
                    curr_input_id = (sys, u)
                    if curr_input_id in input_ids:
                        # Found a direct-feedthrough path to the input_port
                        return True
                    elif curr_input_id in self.connection_map:
                        # Intermediate input port has a direct-feedthrough path to
                        # output_port. Add the upstream output port (if there
                        # is one) to the active set.
                        active_set.add(self.connection_map[curr_input_id])

        # If there are no intermediate output ports with a direct-feedthrough path
        # to the output port, there is no direct feedthrough from the input port
        return False

    # Inherits docstring from SystemBase.get_feedthrough
    def get_feedthrough(self) -> List[Tuple[int, int]]:
        if self.feedthrough_pairs is not None:
            return self.feedthrough_pairs

        pairs = []
        for u in range(self.num_input_ports):
            for v in range(self.num_output_ports):
                if self._has_feedthrough(u, v):
                    pairs.append((u, v))

        self.feedthrough_pairs = pairs
        return self.feedthrough_pairs

    def find_system_with_path(self, path: str | list[str]) -> SystemBase:
        if isinstance(path, str):
            path = path.split(".")

        def _find_in_children():
            for child in self.nodes:
                if child.name == path[0]:
                    if len(path) == 1:
                        return child
                    if isinstance(child, Diagram):
                        return child.find_system_with_path(path[1:])
                    return None
            return None

        if self.parent is None:
            return _find_in_children()

        if self.name == path[0] and len(path) == 1:
            return self

        return _find_in_children()

    def _get_block_by_name(self, name: str):
        """Return the immediate child subsystem with the given name, or ``None``."""
        for system in self.nodes:
            if system.name == name:
                return system
        return None

    def get_parameter(self, path: str):
        """Get a parameter by dot-separated path (child blocks and nested diagrams).

        For a path ``"block.param"``, ``block`` must be a direct child name of
        this diagram; the remainder is resolved on that child (recursively for
        nested diagrams). A single segment refers to this diagram's own
        parameters (same as :meth:`SystemBase.get_parameter`).

        Examples:
            ``diagram.get_parameter("gain.gain")`` for a child named ``gain``
            with parameter ``gain``.

        Raises:
            KeyError: If a segment does not match a child or parameter.
        """
        parts = path.split(".", 1)
        if len(parts) == 1:
            return super().get_parameter(path)

        block_name, remainder = parts
        block = self._get_block_by_name(block_name)
        if block is None:
            available = [s.name for s in self.nodes]
            raise KeyError(
                f"Block {block_name!r} not found in diagram {self.name!r}. "
                f"Available blocks: {available}"
            )
        return block.get_parameter(remainder)

    def with_parameters(self, updates: dict[str, Any]) -> Diagram:
        """Return a new diagram with parameters replaced (dot-notation paths).

        Grouping is by top-level block name; nested paths are forwarded recursively.
        The original diagram is unchanged.

        Args:
            updates: Map from dot paths to new values, e.g.
                ``{"motor.R": jnp.array(2.3), "controller.Kp": jnp.array(1.5)}``.

        Returns:
            New :class:`Diagram` instance.

        Raises:
            KeyError: Unknown block or parameter.
            TypeError: Attempt to replace a static parameter.
        """
        new = copy.deepcopy(self)
        new.system_id = next_system_id()
        new.parent = None
        new._dependency_graph = None
        new.feedthrough_pairs = None
        new._cache_update_events = None

        if not updates:
            _diagram_rebuild_leaf_systems(new)
            return new

        own: dict[str, Any] = {}
        by_child: dict[str, dict[str, Any]] = {}
        for path, val in updates.items():
            parts = path.split(".", 1)
            if len(parts) == 1:
                own[parts[0]] = val
            else:
                block_name, remainder = parts
                by_child.setdefault(block_name, {})[remainder] = val

        for pname, val in own.items():
            if pname in new._static_parameters:
                raise TypeError(
                    f"Parameter {pname!r} is static on {new.name!r}; static "
                    "parameters cannot be replaced at runtime without recompilation."
                )
            if pname not in new._dynamic_parameters:
                available = sorted(
                    {*new._static_parameters.keys(), *new._dynamic_parameters.keys()}
                )
                raise KeyError(
                    f"Parameter {pname!r} not found on diagram {new.name!r}. "
                    f"Available: {available}"
                )
            old_p = new._dynamic_parameters[pname]
            try:
                val = _check_values_compatible(Parameter.unwrap(old_p), val)
            except ValueError as e:
                raise ValueError(
                    f"{e} (parameter {pname!r} on {new.name!r})"
                ) from None
            if isinstance(old_p, Parameter):
                # Mutate the (copied) Parameter in place rather than swapping
                # in a fresh object: blocks that reference this parameter as a
                # shared alias (directly, or via a deserialized string
                # expression) are registered as ParameterCache dependents of
                # *this* object, and set() is what invalidates them. Replacing
                # the dict entry would update the diagram-level name only and
                # silently leave every referencing block at its stale value
                # (T-141).
                old_p.set(val)
            else:
                new._dynamic_parameters[pname] = Parameter(value=val, name=pname)

        for block_name, subupdates in by_child.items():
            idx = None
            for i, node in enumerate(self.nodes):
                if node.name == block_name:
                    idx = i
                    break
            if idx is None:
                available = [s.name for s in self.nodes]
                raise KeyError(
                    f"Block {block_name!r} not found in diagram {self.name!r}. "
                    f"Available blocks: {available}"
                )

            old_child = new.nodes[idx]
            orig_child = self.nodes[idx]
            if isinstance(orig_child, Diagram):
                repl = orig_child.with_parameters(subupdates)
            elif isinstance(orig_child, LeafSystem):
                repl = orig_child
                for subpath, v in subupdates.items():
                    if "." in subpath:
                        raise KeyError(
                            f"Invalid parameter path {block_name!r}.{subpath!r}: "
                            f"block {block_name!r} is a leaf system."
                        )
                    repl = repl.with_parameter(subpath, v)
            else:
                raise TypeError(
                    f"Unsupported system type for with_parameters: {type(orig_child)}"
                )

            new.nodes[idx] = repl
            repl.parent = new
            _diagram_rewrite_child_refs(new, old_child, repl)
            _diagram_refresh_exported_outputs_for_child(new, repl)

        _diagram_rebuild_leaf_systems(new)
        return new

    def list_parameters(self, prefix: str = "") -> dict:
        """Flatten parameters under this diagram with dot-notation keys.

        Includes this diagram's own parameters (if any), then each child's
        parameters prefixed by ``child_name.``. Nested diagrams recurse.

        Args:
            prefix: Internal use: prepend to every key (non-empty when called
                recursively from a parent diagram).
        """
        result = {}
        for name, value in super().list_parameters().items():
            key = f"{prefix}{name}" if prefix else name
            result[key] = value

        for system in self.nodes:
            sub_prefix = f"{prefix}{system.name}." if prefix else f"{system.name}."
            if isinstance(system, Diagram):
                result.update(system.list_parameters(prefix=sub_prefix))
            else:
                for pname, pval in system.list_parameters().items():
                    result[f"{sub_prefix}{pname}"] = pval

        return result

    def declare_dynamic_parameter(
        self, name: str, parameter: Array | Parameter
    ) -> None:
        """Declare a parameter for this system.

        Parameters:
            name (str): The name of the parameter.
            parameter (Parameter): The parameter object.
        """
        # Force the parameter to have the correct name, all diagram parameters
        # should be named.
        parameter.name = name
        # do not wrap in a new Parameter object like we do for LeafSystem because
        # the parameter could be used in multiple places.
        self._dynamic_parameters[name] = parameter

    # TODO: move this to context? it can't be called without the context first
    # being created (which creates the dependency graph)
    def check_no_algebraic_loops(self):
        """Check for algebraic loops in the diagram.

        This is a more or less direct port of the Drake method
        DiagramBuilder::ThrowIfAlgebraicLoopExists. Some comments are verbatim
        explanations of the algorithm implemented there.
        """

        # The nodes in the graph are the input/output ports defined as part of
        # the diagram's internal connections.  Ports that are not internally
        # connected cannot participate in a cycle at this level, so we don't include them
        # in the nodes set.
        nodes: Set[PortBase] = set()

        # For each `value` in `edges[key]`, the `key` directly influences `value`.
        edges: Mapping[PortBase, Set[PortBase]] = {}

        # Add the diagram's internal connections to the digraph nodes and edges
        for input_port_locator, output_port_locator in self.connection_map.items():
            # Directly using the port locator does not result in a unique identifier
            # since (sys, 0) represents both input port 0 and output port 0.  Instead,
            # use the port directly as a key, since it is a unique hashable object.
            input_system, input_index = input_port_locator
            input_port = input_system.input_ports[input_index]
            logger.debug(f"Adding locator {input_port} to nodes")
            nodes.add(input_port)

            output_system, output_index = output_port_locator
            output_port = output_system.output_ports[output_index]
            logger.debug(f"Adding locator {output_port} to nodes")
            nodes.add(output_port)

            if output_port not in edges:
                edges[output_port] = set()

            logger.debug(f"Adding edge[{output_port}] = {input_port}")
            edges[output_port].add(input_port)

        # Add more edges based on each System's direct feedthrough.
        # input -> output port iff there is direct feedthrough from input -> output
        # If a feedthrough edge refers to a port not in `nodes`, omit it because ports
        # that are not connected inside the diagram cannot participate in a cycle at
        # the level of this diagram (higher-level diagrams will test for cycles at
        # their level).
        for system in self.nodes:
            logger.debug(f"Checking feedthrough for system {system.name}")
            for input_index, output_index in system.get_feedthrough():
                input_port = system.input_ports[input_index]
                output_port = system.output_ports[output_index]
                logger.debug(f"Feedthrough from {input_port} to {output_port}")
                if input_port in nodes and output_port in nodes:
                    if input_port not in edges:
                        edges[input_port] = set()
                    edges[input_port].add(output_port)

        def _graph_has_cycle(
            node: PortBase,
            visited: Set[DirectedPortLocator],
            stack: List[DirectedPortLocator],
        ) -> bool:
            # Helper to do the algebraic loop test by depth-first search on the graph
            # to find cycles. Modifies `visited` and `stack` in place.

            logger.debug(f"Checking node {node}")

            assert node.directed_locator not in visited
            visited.add(node.directed_locator)

            if node in edges:
                assert node not in stack
                stack.append(node.directed_locator)
                edge_iter = edges[node]
                for target in edge_iter:
                    if target.directed_locator not in visited and _graph_has_cycle(
                        target, visited, stack
                    ):
                        logger.debug(f"Found cycle at {target}")
                        return True
                    elif target.directed_locator in stack:
                        logger.debug(f"Found target {target} in stack {stack}")
                        return True
                stack.pop()

            # If we get this far there is no cycle
            return False

        # Evaluate the graph for cycles
        visited: Set[DirectedPortLocator] = set()
        stack: List[DirectedPortLocator] = []
        for node in nodes:
            if node.directed_locator in visited:
                continue
            if _graph_has_cycle(node, visited, stack):
                raise AlgebraicLoopError(self.name, stack)

    @property
    def has_dirty_static_parameters(self) -> bool:
        """Check if any static parameters have been modified."""
        return any(n.has_dirty_static_parameters for n in self.nodes)

`continuous_substep_vector` `property`

T-133: per-leaf multirate substep vectors, in leaf_systems order (the same ordering mass_matrix relies on for alignment with the flattened continuous state).

`has_dirty_static_parameters` `property`

Check if any static parameters have been modified.

`check_no_algebraic_loops()`

Check for algebraic loops in the diagram.

This is a more or less direct port of the Drake method DiagramBuilder::ThrowIfAlgebraicLoopExists. Some comments are verbatim explanations of the algorithm implemented there.

Source code in jaxonomy/framework/diagram.py

def check_no_algebraic_loops(self):
    """Check for algebraic loops in the diagram.

    This is a more or less direct port of the Drake method
    DiagramBuilder::ThrowIfAlgebraicLoopExists. Some comments are verbatim
    explanations of the algorithm implemented there.
    """

    # The nodes in the graph are the input/output ports defined as part of
    # the diagram's internal connections.  Ports that are not internally
    # connected cannot participate in a cycle at this level, so we don't include them
    # in the nodes set.
    nodes: Set[PortBase] = set()

    # For each `value` in `edges[key]`, the `key` directly influences `value`.
    edges: Mapping[PortBase, Set[PortBase]] = {}

    # Add the diagram's internal connections to the digraph nodes and edges
    for input_port_locator, output_port_locator in self.connection_map.items():
        # Directly using the port locator does not result in a unique identifier
        # since (sys, 0) represents both input port 0 and output port 0.  Instead,
        # use the port directly as a key, since it is a unique hashable object.
        input_system, input_index = input_port_locator
        input_port = input_system.input_ports[input_index]
        logger.debug(f"Adding locator {input_port} to nodes")
        nodes.add(input_port)

        output_system, output_index = output_port_locator
        output_port = output_system.output_ports[output_index]
        logger.debug(f"Adding locator {output_port} to nodes")
        nodes.add(output_port)

        if output_port not in edges:
            edges[output_port] = set()

        logger.debug(f"Adding edge[{output_port}] = {input_port}")
        edges[output_port].add(input_port)

    # Add more edges based on each System's direct feedthrough.
    # input -> output port iff there is direct feedthrough from input -> output
    # If a feedthrough edge refers to a port not in `nodes`, omit it because ports
    # that are not connected inside the diagram cannot participate in a cycle at
    # the level of this diagram (higher-level diagrams will test for cycles at
    # their level).
    for system in self.nodes:
        logger.debug(f"Checking feedthrough for system {system.name}")
        for input_index, output_index in system.get_feedthrough():
            input_port = system.input_ports[input_index]
            output_port = system.output_ports[output_index]
            logger.debug(f"Feedthrough from {input_port} to {output_port}")
            if input_port in nodes and output_port in nodes:
                if input_port not in edges:
                    edges[input_port] = set()
                edges[input_port].add(output_port)

    def _graph_has_cycle(
        node: PortBase,
        visited: Set[DirectedPortLocator],
        stack: List[DirectedPortLocator],
    ) -> bool:
        # Helper to do the algebraic loop test by depth-first search on the graph
        # to find cycles. Modifies `visited` and `stack` in place.

        logger.debug(f"Checking node {node}")

        assert node.directed_locator not in visited
        visited.add(node.directed_locator)

        if node in edges:
            assert node not in stack
            stack.append(node.directed_locator)
            edge_iter = edges[node]
            for target in edge_iter:
                if target.directed_locator not in visited and _graph_has_cycle(
                    target, visited, stack
                ):
                    logger.debug(f"Found cycle at {target}")
                    return True
                elif target.directed_locator in stack:
                    logger.debug(f"Found target {target} in stack {stack}")
                    return True
            stack.pop()

        # If we get this far there is no cycle
        return False

    # Evaluate the graph for cycles
    visited: Set[DirectedPortLocator] = set()
    stack: List[DirectedPortLocator] = []
    for node in nodes:
        if node.directed_locator in visited:
            continue
        if _graph_has_cycle(node, visited, stack):
            raise AlgebraicLoopError(self.name, stack)

`check_types(context, error_collector=None)`

Perform any system-specific static analysis.

Source code in jaxonomy/framework/diagram.py

def check_types(
    self,
    context: DiagramContext,
    error_collector: ErrorCollector = None,
) -> None:
    """Perform any system-specific static analysis."""
    for system in self.nodes:
        system.check_types(
            context,
            error_collector=error_collector,
        )

`declare_dynamic_parameter(name, parameter)`

Declare a parameter for this system.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the parameter.	required
`parameter`	`Parameter`	The parameter object.	required

Source code in jaxonomy/framework/diagram.py

def declare_dynamic_parameter(
    self, name: str, parameter: Array | Parameter
) -> None:
    """Declare a parameter for this system.

    Parameters:
        name (str): The name of the parameter.
        parameter (Parameter): The parameter object.
    """
    # Force the parameter to have the correct name, all diagram parameters
    # should be named.
    parameter.name = name
    # do not wrap in a new Parameter object like we do for LeafSystem because
    # the parameter could be used in multiple places.
    self._dynamic_parameters[name] = parameter

`eval_subsystem_input_port(context, port_locator)`

Evaluate the input port for a child of this system given the root context.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	root context for this system	required
`port_locator`	`InputPortLocator`	tuple of (system, port_index) identifying the input port to evaluate	required

Returns:

Name	Type	Description
`Array`	`Array`	Value returned from evaluating the subsystem port.

Raises:

Type	Description
`InputNotConnectedError`	if the input port is not connected

Source code in jaxonomy/framework/diagram.py

def eval_subsystem_input_port(
    self, context: DiagramContext, port_locator: InputPortLocator
) -> Array:
    """Evaluate the input port for a child of this system given the root context.

    Args:
        context (ContextBase): root context for this system
        port_locator (InputPortLocator): tuple of (system, port_index) identifying
            the input port to evaluate

    Returns:
        Array: Value returned from evaluating the subsystem port.

    Raises:
        InputNotConnectedError: if the input port is not connected
    """

    is_exported = port_locator in self._input_port_map
    if is_exported:
        # The upstream source is an input to this whole Diagram; evaluate that
        # input port and use the result as the value for this one.
        port_index = self._input_port_map[port_locator]  # Diagram-level index
        return self.input_ports[port_index].eval(context)  # Return upstream value

    is_connected = port_locator in self.connection_map
    if is_connected:
        # The upstream source is an output port of one of this Diagram's child
        # subsystems; evaluate the upstream output.
        upstream_locator = self.connection_map[port_locator]

        # This will return the value of the upstream port
        return self.eval_subsystem_output_port(context, upstream_locator)

    block, port_index = port_locator
    raise InputNotConnectedError(
        system=block,
        port_index=port_index,
        port_direction="in",
        message=f"Input port {block.name}[{port_index}] is not connected",
    )

`eval_subsystem_output_port(context, port_locator)`

"Evaluate the output port for a child of this system given the root context.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	root context for this system	required
`port_locator`	`OutputPortLocator`	tuple of (system, port_index) identifying the output port to evaluate	required

Returns:

Name	Type	Description
`Array`	`Array`	Value returned from evaluating the subsystem port.

Source code in jaxonomy/framework/diagram.py

def eval_subsystem_output_port(
    self, context: DiagramContext, port_locator: OutputPortLocator
) -> Array:
    """ "Evaluate the output port for a child of this system given the root context.

    Args:
        context (ContextBase): root context for this system
        port_locator (OutputPortLocator): tuple of (system, port_index) identifying
            the output port to evaluate

    Returns:
        Array: Value returned from evaluating the subsystem port.
    """
    system, port_index = port_locator
    port = system.output_ports[port_index]

    # During simulation all we should need to do is evaluate the port.
    if context.is_initialized:
        return port.eval(context)

    # If the context is not initialized, we have to determine the signal data type.
    # In the easy case, the port has a default value, so we can just use that.
    if port.default_value is not None:
        logger.debug(
            "Using default output value of %s for %s",
            port.default_value,
            port_locator[0].name,
        )
        return port.default_value

    logger.debug(
        "Evaluating output port %s for system %s. Context initialized: %s",
        port_locator,
        port_locator[0].name,
        context.is_initialized,
    )

    # If there is no default value, try to evaluate the port to pull a "template"
    # value with an appropriate data type from upstream.  This will return None if
    # the port is not yet connected (e.g. if its upstream is an exported input of)
    # a Diagram, so we can defer evaluation.

    # Try again to evaluate the port
    val = port.eval(context)
    logger.debug(
        "  ---> %s returns %s", (port_locator[0].name, port_locator[1]), val
    )

    # If there is still no value, the port is not connected to anything.
    # Post-initialization this would be an error, but pre-initialization
    # it may be the case that the upstream is an exported input port of
    # the Diagram, so we can defer evaluation. Expect the block that is
    # doing this to handle the UpstreamEvalError appropriately.
    if val is None:
        system_name = system.name_path_str
        logger.debug(
            "Upstream evaluation of %s.out[%s] returned None. Deferring evaluation.",
            system_name,
            port_index,
        )
        raise UpstreamEvalError(port_locator=(system, "out", port_index))
    return val

`export_input(locator, port_name)`

Export a subsystem input port as a diagram-level input.

This should typically only be called during construction by DiagramBuilder. The standard workflow will be to call export_input on the builder object, which will automatically call this method on the Diagram once created.

Parameters:

Name	Type	Description	Default
`locator`	`InputPortLocator`	tuple of (system, port_index) identifying the input port to export	required
`port_name`	`str`	name of the new exported input port	required

Returns:

Name	Type	Description
`int`	`int`	index of the exported input port in the diagram input_ports list

Source code in jaxonomy/framework/diagram.py

def export_input(self, locator: InputPortLocator, port_name: str) -> int:
    """Export a subsystem input port as a diagram-level input.

    This should typically only be called during construction by DiagramBuilder.
    The standard workflow will be to call export_input on the _builder_ object,
    which will automatically call this method on the Diagram once created.

    Args:
        locator (InputPortLocator): tuple of (system, port_index) identifying
            the input port to export
        port_name (str): name of the new exported input port

    Returns:
        int: index of the exported input port in the diagram input_ports list
    """
    diagram_port_index = self.declare_input_port(name=port_name)
    self._input_port_map[locator] = diagram_port_index

    return diagram_port_index

`export_output(locator, port_name)`

Export a subsystem output port as a diagram-level output.

This should typically only be called during construction by DiagramBuilder. The standard workflow will be to call export_input on the builder object, which will automatically call this method on the Diagram once created.

Parameters:

Name	Type	Description	Default
`locator`	`OutputPortLocator`	tuple of (system, port_index) identifying the output port to export	required
`port_name`	`str`	name of the new exported output port	required

Returns:

Name	Type	Description
`int`	`int`	index of the exported output port in the diagram output_ports list

Source code in jaxonomy/framework/diagram.py

def export_output(self, locator: OutputPortLocator, port_name: str) -> int:
    """Export a subsystem output port as a diagram-level output.

    This should typically only be called during construction by DiagramBuilder.
    The standard workflow will be to call export_input on the _builder_ object,
    which will automatically call this method on the Diagram once created.

    Args:
        locator (OutputPortLocator): tuple of (system, port_index) identifying
            the output port to export
        port_name (str): name of the new exported output port

    Returns:
        int: index of the exported output port in the diagram output_ports list
    """
    subsystem, subsystem_port_index = locator
    source_port = subsystem.output_ports[subsystem_port_index]
    diagram_port_index = self.declare_output_port(
        source_port.eval,
        name=port_name,
        prerequisites_of_calc=[source_port.ticket],
    )
    self._output_port_map[locator] = diagram_port_index
    self._inv_output_port_map[diagram_port_index] = locator

    return diagram_port_index

`get_parameter(path)`

Get a parameter by dot-separated path (child blocks and nested diagrams).

For a path "block.param", block must be a direct child name of this diagram; the remainder is resolved on that child (recursively for nested diagrams). A single segment refers to this diagram's own parameters (same as :meth:SystemBase.get_parameter).

Examples:

diagram.get_parameter("gain.gain") for a child named gain with parameter gain.

Raises:

Type	Description
`KeyError`	If a segment does not match a child or parameter.

Source code in jaxonomy/framework/diagram.py

def get_parameter(self, path: str):
    """Get a parameter by dot-separated path (child blocks and nested diagrams).

    For a path ``"block.param"``, ``block`` must be a direct child name of
    this diagram; the remainder is resolved on that child (recursively for
    nested diagrams). A single segment refers to this diagram's own
    parameters (same as :meth:`SystemBase.get_parameter`).

    Examples:
        ``diagram.get_parameter("gain.gain")`` for a child named ``gain``
        with parameter ``gain``.

    Raises:
        KeyError: If a segment does not match a child or parameter.
    """
    parts = path.split(".", 1)
    if len(parts) == 1:
        return super().get_parameter(path)

    block_name, remainder = parts
    block = self._get_block_by_name(block_name)
    if block is None:
        available = [s.name for s in self.nodes]
        raise KeyError(
            f"Block {block_name!r} not found in diagram {self.name!r}. "
            f"Available blocks: {available}"
        )
    return block.get_parameter(remainder)

`initialize_static_data(context)`

Perform any system-specific static analysis.

Source code in jaxonomy/framework/diagram.py

def initialize_static_data(self, context: DiagramContext) -> DiagramContext:
    """Perform any system-specific static analysis."""
    for system in self.nodes:
        context = system.initialize_static_data(context)
    return context

`list_parameters(prefix='')`

Flatten parameters under this diagram with dot-notation keys.

Includes this diagram's own parameters (if any), then each child's parameters prefixed by child_name.. Nested diagrams recurse.

Parameters:

Name	Type	Description	Default
`prefix`	`str`	Internal use: prepend to every key (non-empty when called recursively from a parent diagram).	`''`

Source code in jaxonomy/framework/diagram.py

def list_parameters(self, prefix: str = "") -> dict:
    """Flatten parameters under this diagram with dot-notation keys.

    Includes this diagram's own parameters (if any), then each child's
    parameters prefixed by ``child_name.``. Nested diagrams recurse.

    Args:
        prefix: Internal use: prepend to every key (non-empty when called
            recursively from a parent diagram).
    """
    result = {}
    for name, value in super().list_parameters().items():
        key = f"{prefix}{name}" if prefix else name
        result[key] = value

    for system in self.nodes:
        sub_prefix = f"{prefix}{system.name}." if prefix else f"{system.name}."
        if isinstance(system, Diagram):
            result.update(system.list_parameters(prefix=sub_prefix))
        else:
            for pname, pval in system.list_parameters().items():
                result[f"{sub_prefix}{pname}"] = pval

    return result

`post_simulation_finalize()`

Perform any post-simulation cleanup for this system.

Source code in jaxonomy/framework/diagram.py

def post_simulation_finalize(self) -> None:
    """Perform any post-simulation cleanup for this system."""
    for system in self.nodes:
        system.post_simulation_finalize()

`print_schedule(*, format='text', file=None, ensure_initialized=True)`

Print the inferred sample-time schedule of this diagram.

Renders one entry per rate group — period (for discrete blocks), the leaves that fire at that rate, any detected rate mismatches, and the deterministic execution order. Pure inspection helper; does not mutate the diagram or its contexts.

T-105-followup-print-schedule-pre-context: by default, print_schedule() lazily calls :meth:create_context() first so that discrete blocks whose periodic events are registered in their :meth:initialize hook (e.g. :class:PIDDiscrete, :class:Decimator, :class:UnitDelay, :class:ZeroOrderHold) are bucketed into the correct rate group. Pre-fix those blocks showed up as constant when print_schedule() ran before any context had been created, because the periodic event hadn't yet been declared. The lazy create_context() call is idempotent and cheap on already-initialised diagrams; if it fails (e.g. the diagram has missing connections), the schedule is rendered anyway with a one-line warning explaining the rate-group output may be incomplete. Pass ensure_initialized=False to opt out and use the pre-fix behaviour (useful for debugging the initialisation path itself).

Parameters:

Name	Type	Description	Default
`format`	`str`	`"text"` (default), `"markdown"`, or `"json"`.	`'text'`
`file`		Destination file-like object. Defaults to `sys.stdout` when None.	`None`
`ensure_initialized`	`bool`	If True (default), call :meth:`create_context` once before rendering so all discrete blocks have registered their periodic events. Set to False to render against the current (possibly pre-init) state.	`True`

Examples:

>>> diagram.print_schedule()
rate groups:
  discrete(period=0.001, offset=0.0): ctrl_inner
  discrete(period=0.01,  offset=0.0): ctrl_outer
  continuous: plant
execution order: ...

>>> with open("schedule.md", "w") as f:
...     diagram.print_schedule(format="markdown", file=f)

`with_parameters(updates)`

Return a new diagram with parameters replaced (dot-notation paths).

Grouping is by top-level block name; nested paths are forwarded recursively. The original diagram is unchanged.

Parameters:

Name	Type	Description	Default
`updates`	`dict[str, Any]`	Map from dot paths to new values, e.g. `{"motor.R": jnp.array(2.3), "controller.Kp": jnp.array(1.5)}`.	required

Returns:

Name	Type	Description
`New`	`Diagram`	class:`Diagram` instance.

Raises:

Type	Description
`KeyError`	Unknown block or parameter.
`TypeError`	Attempt to replace a static parameter.

Source code in jaxonomy/framework/diagram.py

def with_parameters(self, updates: dict[str, Any]) -> Diagram:
    """Return a new diagram with parameters replaced (dot-notation paths).

    Grouping is by top-level block name; nested paths are forwarded recursively.
    The original diagram is unchanged.

    Args:
        updates: Map from dot paths to new values, e.g.
            ``{"motor.R": jnp.array(2.3), "controller.Kp": jnp.array(1.5)}``.

    Returns:
        New :class:`Diagram` instance.

    Raises:
        KeyError: Unknown block or parameter.
        TypeError: Attempt to replace a static parameter.
    """
    new = copy.deepcopy(self)
    new.system_id = next_system_id()
    new.parent = None
    new._dependency_graph = None
    new.feedthrough_pairs = None
    new._cache_update_events = None

    if not updates:
        _diagram_rebuild_leaf_systems(new)
        return new

    own: dict[str, Any] = {}
    by_child: dict[str, dict[str, Any]] = {}
    for path, val in updates.items():
        parts = path.split(".", 1)
        if len(parts) == 1:
            own[parts[0]] = val
        else:
            block_name, remainder = parts
            by_child.setdefault(block_name, {})[remainder] = val

    for pname, val in own.items():
        if pname in new._static_parameters:
            raise TypeError(
                f"Parameter {pname!r} is static on {new.name!r}; static "
                "parameters cannot be replaced at runtime without recompilation."
            )
        if pname not in new._dynamic_parameters:
            available = sorted(
                {*new._static_parameters.keys(), *new._dynamic_parameters.keys()}
            )
            raise KeyError(
                f"Parameter {pname!r} not found on diagram {new.name!r}. "
                f"Available: {available}"
            )
        old_p = new._dynamic_parameters[pname]
        try:
            val = _check_values_compatible(Parameter.unwrap(old_p), val)
        except ValueError as e:
            raise ValueError(
                f"{e} (parameter {pname!r} on {new.name!r})"
            ) from None
        if isinstance(old_p, Parameter):
            # Mutate the (copied) Parameter in place rather than swapping
            # in a fresh object: blocks that reference this parameter as a
            # shared alias (directly, or via a deserialized string
            # expression) are registered as ParameterCache dependents of
            # *this* object, and set() is what invalidates them. Replacing
            # the dict entry would update the diagram-level name only and
            # silently leave every referencing block at its stale value
            # (T-141).
            old_p.set(val)
        else:
            new._dynamic_parameters[pname] = Parameter(value=val, name=pname)

    for block_name, subupdates in by_child.items():
        idx = None
        for i, node in enumerate(self.nodes):
            if node.name == block_name:
                idx = i
                break
        if idx is None:
            available = [s.name for s in self.nodes]
            raise KeyError(
                f"Block {block_name!r} not found in diagram {self.name!r}. "
                f"Available blocks: {available}"
            )

        old_child = new.nodes[idx]
        orig_child = self.nodes[idx]
        if isinstance(orig_child, Diagram):
            repl = orig_child.with_parameters(subupdates)
        elif isinstance(orig_child, LeafSystem):
            repl = orig_child
            for subpath, v in subupdates.items():
                if "." in subpath:
                    raise KeyError(
                        f"Invalid parameter path {block_name!r}.{subpath!r}: "
                        f"block {block_name!r} is a leaf system."
                    )
                repl = repl.with_parameter(subpath, v)
        else:
            raise TypeError(
                f"Unsupported system type for with_parameters: {type(orig_child)}"
            )

        new.nodes[idx] = repl
        repl.parent = new
        _diagram_rewrite_child_refs(new, old_child, repl)
        _diagram_refresh_exported_outputs_for_child(new, repl)

    _diagram_rebuild_leaf_systems(new)
    return new

`DiagramBuilder`

Class for constructing block diagram systems.

The DiagramBuilder class is responsible for building a diagram by adding systems, connecting ports, and exporting inputs and outputs. It keeps track of the registered systems, input and output ports, and the connection map between input and output ports of the child systems.

Source code in jaxonomy/framework/diagram_builder.py

class DiagramBuilder:
    """Class for constructing block diagram systems.

    The `DiagramBuilder` class is responsible for building a diagram by adding systems, connecting ports,
    and exporting inputs and outputs. It keeps track of the registered systems, input and output ports,
    and the connection map between input and output ports of the child systems.
    """

    def __init__(
        self,
        *,
        validate_rates_at_connect: str | bool | None = None,
        unit_conversion: str = "auto",
        auto_insert_rate_transitions: bool = False,
    ):
        """Construct a DiagramBuilder.

        Args:
            validate_rates_at_connect: T-105 Phase 2 — opt-in connect-time
                multirate consistency check.  When set to ``"warn"`` (or
                ``True``, treated as ``"warn"``), each :meth:`connect`
                call routes through
                :func:`jaxonomy.simulation.rate_groups.check_connection_rate_compat`
                and emits a :class:`RateMismatchWarning` for adjacent
                blocks of incompatible discrete rates.  ``"error"``
                raises :class:`RateMismatchError` on the first offender.
                Default ``None`` keeps the legacy path completely off
                (byte-equivalent to the pre-T-105-Phase-2 behaviour).
            unit_conversion: T-104 followup — controls behaviour when two
                connected ports share base-dimensions but differ only by a
                scalar ``scale`` (e.g. ``meter`` vs ``kilometer``):
                  * ``"auto"`` (default) silently inserts the conversion
                    factor on the destination input port;
                  * ``"warn"`` inserts the factor and emits a
                    :class:`UserWarning`;
                  * ``"error"`` refuses the connection (preserves the
                    Phase-1 strict-equal behaviour).
                Genuine dimensional mismatches (e.g. ``meter`` vs
                ``second``) always raise regardless of mode.
            auto_insert_rate_transitions: T-105-followup-phase3 — when
                ``True``, :meth:`connect` automatically synthesises a
                :func:`jaxonomy.library.RateTransition` block (a
                ``ZeroOrderHold`` for slow→fast, a ``Decimator`` for
                fast→slow) between any two adjacent leaves whose
                inferred discrete sample times differ.  The rewritten
                wiring is ``src → rate_transition → dst`` and an
                informational log line documents the insertion.
                Composes with ``validate_rates_at_connect``: when both
                are enabled, the warning still fires and the transition
                still gets inserted.  Default ``False`` keeps the legacy
                code path byte-equivalent (the strict mode that surfaces
                rate mismatches rather than silently inserting transitions).
        """
        # Child input ports that are exported as diagram-level inputs
        self._input_port_ids: List[InputPortLocator] = []
        self._input_port_names: List[str] = []
        # Child output ports that are exported as diagram-level outputs
        self._output_port_ids: List[OutputPortLocator] = []
        self._output_port_names: List[str] = []

        # Connection map between input and output ports of the child systems
        self._connection_map: Mapping[InputPortLocator, OutputPortLocator] = {}

        # List of registered systems
        self._registered_systems: List[SystemBase] = []

        # Name lookup for input ports
        self._diagram_input_indices: Mapping[str, InputPortLocator] = {}

        # All input ports of child systems (for use in ensuring proper connectivity)
        self._all_input_ports: List[InputPortLocator] = []

        # Each DiagramBuilder can only be used to build a single diagram.  This is to
        # avoid creating multiple diagrams that reference the same LeafSystem. Doing so
        # may or may not actually lead to problems, since the LeafSystems themselves
        # should act like a collection of pure functions, but best practice is to have
        # each leaf system be fully unique.
        self._already_built = False
        self._built_as_name = None

        # T-105 Phase 2: normalise the connect-time-rate-validation flag
        # to the same string vocabulary that ``check_connection_rate_compat``
        # understands.  ``None`` means "off" (byte-equivalent default);
        # ``True`` is sugar for ``"warn"``.
        if validate_rates_at_connect is True:
            self._validate_rates_at_connect: str | None = "warn"
        elif validate_rates_at_connect in (None, False):
            self._validate_rates_at_connect = None
        elif validate_rates_at_connect in ("warn", "error"):
            self._validate_rates_at_connect = validate_rates_at_connect
        else:
            raise BuilderError(
                "validate_rates_at_connect must be None, True/False, "
                "'warn', or 'error'; got "
                f"{validate_rates_at_connect!r}"
            )

        # T-104 followup: unit-conversion mode.  Validated up front so
        # typos surface at construction time rather than from the first
        # ``connect`` call.
        if unit_conversion not in ("auto", "warn", "error"):
            raise BuilderError(
                "unit_conversion must be 'auto', 'warn', or 'error'; "
                f"got {unit_conversion!r}"
            )
        self._unit_conversion = unit_conversion

        # T-105-followup-phase3: opt-in auto-insertion of RateTransition
        # blocks at connect time.  Default ``False`` preserves the legacy
        # behaviour (byte-equivalent).  When enabled, ``connect`` will
        # synthesise a RateTransition block between any two adjacent
        # leaves whose inferred discrete sample times differ.  The
        # counter below gives auto-inserted blocks unique names without
        # leaking into the user-visible name space when the flag is off.
        self._auto_insert_rate_transitions = bool(auto_insert_rate_transitions)
        self._auto_rate_transition_counter = 0

    @overload
    def add(self, system: SystemBase) -> SystemBase: ...

    @overload
    def add(self, system: SystemBase, *systems: SystemBase) -> List[SystemBase]: ...

    def add(self, *systems: SystemBase) -> List[SystemBase] | SystemBase:
        """Add one or more systems to the diagram.

        Args:
            *systems SystemBase:
                System(s) to add to the diagram.

        Returns:
            List[SystemBase] | SystemBase:
                The added system(s). Will return a single system if there is only
                a single system in the argument list.

        Raises:
            BuilderError: If the diagram has already been built.
            BuilderError: If the system is already registered.
            BuilderError: If the system name is not unique.
        """
        for system in systems:
            self._check_not_already_built()
            self._check_system_not_registered(system)
            self._check_system_name_is_unique(system)
            self._registered_systems.append(system)

            # Add the system's input ports to the list of all input ports
            # So that we can make sure they're all connected before building.
            self._all_input_ports.extend([port.locator for port in system.input_ports])

            logger.debug("Added system %s to DiagramBuilder", system.name)
            logger.debug(
                "    Registered systems: %s",
                [s.name for s in self._registered_systems],
            )
        build_recorder.add_block(self, systems)

        return systems[0] if len(systems) == 1 else systems

    def connect(self, src: OutputPort, dest: InputPort):
        """Connect an output port to an input port.

        The input port and output port must both belong to systems that have
        already been added to the diagram.  The input port must not already be
        connected to another output port.

        Args:
            src (OutputPort): The output port to connect.
            dest (InputPort): The input port to connect.

        Raises:
            BuilderError: If the diagram has already been built.
            BuilderError: If the source system is not registered.
            BuilderError: If the destination system is not registered.
            BuilderError: If the input port is already connected.
            BuilderError: If src is an InputPort or dest is an OutputPort.
        """
        # Local imports to avoid a hard import cycle at module load.
        from .port import InputPort as _InputPort, OutputPort as _OutputPort

        # Direction validation -- catches the common (input, input) /
        # (output, output) miswiring at connect time rather than letting it
        # surface as an opaque "input not connected" error during simulation.
        src_is_input = isinstance(src, _InputPort)
        dest_is_output = isinstance(dest, _OutputPort)
        if src_is_input and isinstance(dest, _InputPort):
            raise BuilderError(
                f"Cannot connect input-to-input: "
                f"'{src.system.name}.in[{src.index}]' -> "
                f"'{dest.system.name}.in[{dest.index}]'. "
                f"The first argument must be an output port."
            )
        if isinstance(src, _OutputPort) and dest_is_output:
            raise BuilderError(
                f"Cannot connect output-to-output: "
                f"'{src.system.name}.out[{src.index}]' -> "
                f"'{dest.system.name}.out[{dest.index}]'. "
                f"The second argument must be an input port."
            )
        if src_is_input:
            raise BuilderError(
                f"connect() expected an OutputPort as the first argument, got "
                f"InputPort '{src.system.name}.in[{src.index}]'."
            )
        if dest_is_output:
            raise BuilderError(
                f"connect() expected an InputPort as the second argument, got "
                f"OutputPort '{dest.system.name}.out[{dest.index}]'."
            )

        self._check_not_already_built()
        self._check_system_is_registered(src.system)
        self._check_system_is_registered(dest.system)
        self._check_input_not_connected(dest.locator)

        # T-104 phase 1 / followup: connect-time unit consistency check.
        # Ports that never declared a unit are treated as dimensionless and
        # connect to anything (default-off byte-equivalence).  When both
        # sides declare units:
        #   * dimensional mismatch (e.g. m vs s) -> always UnitMismatchError;
        #   * scalar-scale mismatch (e.g. m vs km):
        #       - "error": raise (Phase-1 strict behaviour);
        #       - "warn":  apply factor + emit UserWarning;
        #       - "auto":  silently apply factor.
        src_units = getattr(src, "units", None)
        dst_units = getattr(dest, "units", None)
        src_label = f"'{src.system.name}.out[{src.index}]' ({src.name})"
        dst_label = f"'{dest.system.name}.in[{dest.index}]' ({dest.name})"

        if self._unit_conversion == "error":
            # Preserve the Phase-1 strict-equal behaviour.
            assert_unit_compatible(
                src_units,
                dst_units,
                src_label=src_label,
                dst_label=dst_label,
            )
        else:
            # Returns the multiplicative factor (1.0 for matched / wildcard
            # units, src.scale / dst.scale otherwise).  Raises on genuine
            # dimensional mismatch.
            factor = assert_units_compatible_with_scale(
                src_units,
                dst_units,
                src_label=src_label,
                dst_label=dst_label,
            )
            if factor != 1.0:
                src_u = resolve_unit(src_units)
                dst_u = resolve_unit(dst_units)
                msg = (
                    f"Unit conversion: applying factor {factor!r} to "
                    f"connection {src_label} ({src_u!r}) -> "
                    f"{dst_label} ({dst_u!r})."
                )
                if self._unit_conversion == "warn":
                    warnings.warn(msg, UserWarning, stacklevel=2)
                else:
                    # "auto": informational log line; no warning.
                    logger.info(msg)
                _install_unit_conversion(dest, factor)

        # T-105 Phase 2: opt-in connect-time multirate consistency check.
        # Default-off (``self._validate_rates_at_connect is None``) keeps
        # the legacy code path byte-equivalent.  When the builder was
        # constructed with ``validate_rates_at_connect=...``, route the
        # source/dest pair through ``check_connection_rate_compat`` which
        # honours the ``_jaxonomy_rate_transition`` marker (T-123) and
        # the universal-sample-time rule (constant / inherited bridge
        # any rates).
        if self._validate_rates_at_connect is not None:
            # Local import: avoids a hard import cycle between
            # ``framework`` and ``simulation`` at module load.
            from ..simulation.rate_groups import check_connection_rate_compat

            check_connection_rate_compat(
                src.system,
                src.index,
                dest.system,
                dest.index,
                on_mismatch=self._validate_rates_at_connect,
            )

        # T-105-followup-phase3: opt-in connect-time auto-insertion of
        # RateTransition blocks between adjacent leaves with differing
        # discrete sample times.  Default-off keeps the legacy code
        # path byte-equivalent.  Composes with the validate hook above:
        # when both are enabled, the warning still fires (above) and
        # the bridge still gets inserted (here).
        if self._auto_insert_rate_transitions:
            inserted = self._maybe_auto_insert_rate_transition(src, dest)
            if inserted is not None:
                # ``_maybe_auto_insert_rate_transition`` already wrote
                # both legs of ``src → bridge → dest`` into the
                # connection map (and recorded both connections with
                # the build recorder), so we are done here.
                return

        build_recorder.connect_ports(self, src, dest)

        self._connection_map[dest.locator] = src.locator

        logger.debug(
            f"Connected port {src.name} of system {src.system.name} to port {dest.name} of system {dest.system.name}"
        )

    def _maybe_auto_insert_rate_transition(self, src, dest):
        """T-105-followup-phase3 helper.

        Inspect the inferred sample times of ``src.system`` and
        ``dest.system``; if they are both discrete with different
        periods, synthesise a :func:`jaxonomy.library.RateTransition`
        block, register it with the builder, and rewrite the connection
        as ``src → rate_transition → dest``.

        Returns the freshly added bridge block, or ``None`` if no
        insertion was appropriate (matched rates, universal source/dest,
        or either side already a rate-transition bridge).
        """
        # Local imports: avoid hard import cycles between framework and
        # simulation/library at module load.
        from ..simulation.rate_groups import infer_block_sample_time

        # If either side is already a rate-transition bridge, do nothing
        # — the user (or a previous auto-insertion) has already handled
        # the transition.
        if getattr(src.system, "_jaxonomy_rate_transition", False):
            return None
        if getattr(dest.system, "_jaxonomy_rate_transition", False):
            return None

        src_st = infer_block_sample_time(src.system)
        dst_st = infer_block_sample_time(dest.system)

        # Only auto-insert when both sides are discrete with different
        # periods.  Universal sample times (constant/inherited) and
        # continuous-to-continuous matches do not need a bridge.
        if not (src_st.is_discrete() and dst_st.is_discrete()):
            return None
        if src_st.matches(dst_st):
            return None

        from ..library import RateTransition

        self._auto_rate_transition_counter += 1
        bridge_name = (
            f"_auto_rate_transition_{self._auto_rate_transition_counter}_"
            f"{src.system.name}_to_{dest.system.name}"
        )
        bridge = RateTransition(
            input_dt=src_st.period,
            output_dt=dst_st.period,
            name=bridge_name,
        )
        # Belt-and-suspenders: ``RateTransition`` already tags the
        # returned block with ``_jaxonomy_rate_transition = True`` on
        # the slow→fast (ZOH) and fast→slow (Decimator) branches.  Set
        # it again here so future auto-insertion calls definitely skip
        # this block even if the factory's tagging changes.
        bridge._jaxonomy_rate_transition = True

        self.add(bridge)

        logger.info(
            "T-105-followup-phase3 auto-inserted RateTransition '%s' "
            "between '%s.out[%d]' (dt=%s) and '%s.in[%d]' (dt=%s).",
            bridge_name,
            src.system.name,
            src.index,
            src_st.period,
            dest.system.name,
            dest.index,
            dst_st.period,
        )

        # Wire ``src → bridge.in[0]`` and ``bridge.out[0] → dest``.
        # We bypass the public ``connect`` method on these two legs to
        # avoid recursing into auto-insertion (the bridge is tagged so
        # it would short-circuit anyway, but going through the recorder
        # / connection-map directly is simpler and matches the
        # documented invariant that the auto-inserted block is exactly
        # one bridge between the two original ports).
        bridge_in = bridge.input_ports[0]
        bridge_out = bridge.output_ports[0]

        build_recorder.connect_ports(self, src, bridge_in)
        self._connection_map[bridge_in.locator] = src.locator
        logger.debug(
            f"Connected port {src.name} of system {src.system.name} to port {bridge_in.name} of system {bridge.name}"
        )

        build_recorder.connect_ports(self, bridge_out, dest)
        self._connection_map[dest.locator] = bridge_out.locator
        logger.debug(
            f"Connected port {bridge_out.name} of system {bridge.name} to port {dest.name} of system {dest.system.name}"
        )

        return bridge

    def export_input(self, port: InputPort, name: str = None) -> int:
        """Export an input port of a child system as a diagram-level input.

        The input port must belong to a system that has already been added to the
        diagram. The input port must not already be connected to another output port.

        Args:
            port (InputPort): The input port to export.
            name (str, optional):
                The name to assign to the exported input port. If not provided, a
                unique name will be generated.

        Returns:
            int: The index (in the to-be-built diagram) of the exported input port.

        Raises:
            BuilderError: If the diagram has already been built.
            BuilderError: If the system is not registered.
            BuilderError: If the input port is already connected.
            BuilderError: If the input port name is not unique.
        """
        self._check_not_already_built()
        self._check_system_is_registered(port.system)
        self._check_input_not_connected(port.locator)

        if name is None:
            # Since the system names are unique, auto-generated port names are also unique
            # at the level of _this_ diagram (subsystems can have ports with the same name)
            name = f"{port.system.name}_{port.name}"
        elif name in self._diagram_input_indices:
            raise BuilderError(
                f"Input port name {name} is not unique",
                system=port.system,
                port_index=port.index,
                port_direction="in",
            )

        # Index at the diagram (not subsystem) level
        port_index = len(self._input_port_ids)
        self._input_port_ids.append(port.locator)
        self._input_port_names.append(name)

        self._diagram_input_indices[name] = port_index

        build_recorder.export_port(self, port.system, "input", port.index, name)

        return port_index

    def export_output(self, port: OutputPort, name: str = None) -> int:
        """Export an output port of a child system as a diagram-level output.

        The output port must belong to a system that has already been added to the
        diagram.

        Args:
            port (OutputPort): The output port to export.
            name (str, optional):
                The name to assign to the exported output port. If not provided, a
                unique name will be generated.

        Returns:
            int: The index (in the to-be-built diagram) of the exported output port.

        Raises:
            BuilderError: If the diagram has already been built.
            BuilderError: If the system is not registered.
            BuilderError: If the output port name is not unique.
        """
        self._check_not_already_built()
        self._check_system_is_registered(port.system)

        if name is None:
            # Since the system names are unique, auto-generated port names are also unique
            # at the level of _this_ diagram (subsystems can have ports with the same name)
            name = f"{port.system.name}_{port.name}"
        elif name in self._output_port_names:
            raise BuilderError(
                f"Output port name {name} is not unique",
                system=port.system,
                port_index=port.index,
                port_direction="out",
            )

        # Index at the diagram (not subsystem) level
        port_index = len(self._output_port_ids)
        self._output_port_ids.append(port.locator)
        self._output_port_names.append(name)

        build_recorder.export_port(self, port.system, "output", port.index, name)

        return port_index

    def _check_not_already_built(self):
        if self._already_built:
            raise BuilderError(
                "DiagramBuilder: build has already been called to "
                "create a diagram; this DiagramBuilder may no longer be used: "
                f"{self._built_as_name}"
            )

    def _check_system_name_is_unique(self, system: SystemBase):
        if system.name in map(lambda s: s.name, self._registered_systems):
            raise SystemNameNotUniqueError(system)

    def _system_is_registered(self, system: SystemBase) -> bool:
        # return (system is not None) and (system in self._registered_systems)
        if system.system_id is None:  # system.__init__ is not done yet
            return False
        return system.system_id in map(lambda s: s.system_id, self._registered_systems)

    def _check_system_not_registered(self, system: SystemBase):
        if self._system_is_registered(system):
            raise BuilderError(
                f"System {system.name} is already registered",
                system=system,
            )

    def _check_system_is_registered(self, system: SystemBase):
        if not self._system_is_registered(system):
            raise BuilderError(
                f"System {system.name} is not registered",
                system=system,
            )

    def _check_input_not_connected(self, input_port_locator: InputPortLocator):
        if not (
            (input_port_locator not in self._input_port_ids)
            and (input_port_locator not in self._connection_map)
        ):
            system, port_index = input_port_locator
            raise BuilderError(
                f"Input port {port_index} for {system} is already connected",
                system=system,
                port_index=port_index,
                port_direction="in",
            )

    def _check_input_is_connected(self, input_port_locator: InputPortLocator):
        if not (
            (input_port_locator in self._input_port_ids)
            or (input_port_locator in self._connection_map)
        ):
            raise DisconnectedInputError(input_port_locator)

    def _check_contents_are_complete(self):
        # Make sure all the systems referenced in the builder attributes are registered

        # Check that systems and registered_systems have the same elements
        for system in self._registered_systems:
            self._check_system_is_registered(system)

        # Check that connection_map only refers to registered systems
        for (
            input_port_locator,
            output_port_locator,
        ) in self._connection_map.items():
            self._check_system_is_registered(input_port_locator[0])
            self._check_system_is_registered(output_port_locator[0])

        # Check that input_port_ids and output_port_ids only refer to registered systems
        for port_locator in [*self._input_port_ids, *self._output_port_ids]:
            self._check_system_is_registered(port_locator[0])

    def _check_ports_are_valid(self):
        for dst, src in self._connection_map.items():
            dst_sys, dst_idx = dst
            if (dst_idx < 0) or (dst_idx >= dst_sys.num_input_ports):
                raise BuilderError(
                    f"Input port index {dst_idx} is out of range "
                    f"(0-{dst_sys.num_input_ports-1})",
                    system=dst_sys,
                    port_index=dst_idx,
                    port_direction="in",
                )
            src_sys, src_idx = src
            if (src_idx < 0) or (src_idx >= src_sys.num_output_ports):
                raise BuilderError(
                    f"Output port index {src_idx} is out of range "
                    f"(0-{src_sys.num_output_ports-1})",
                    system=src_sys,
                    port_index=src_idx,
                    port_direction="out",
                )

    def build(
        self,
        name: str = "root",
        ui_id: str = None,
        parameters: dict[str, Parameter] = None,
    ) -> Diagram:
        """Builds a Diagram system with the specified name and system ID.

        Args:
            name (str, optional): The name of the diagram. Defaults to "root".
            ui_id (str, optional): The unique identifier for the diagram.
            parameters (dict[str, Parameter], optional):
                A dictionary of dynamic parameters to declare for the diagram.

        Returns:
            Diagram: The newly constructed diagram.

        Raises:
            EmptyDiagramError: If no systems are registered in the diagram.
            BuilderError: If the diagram has already been built.
            AlgebraicLoopError: If an algebraic loop is detected in the diagram.
            DisconnectedInputError: If an input port is not connected.
        """
        self._check_not_already_built()
        self._check_contents_are_complete()
        self._check_ports_are_valid()

        # Check that all internal input ports are connected
        for input_port_locator in self._input_port_ids:
            self._check_input_is_connected(input_port_locator)

        if len(self._registered_systems) == 0:
            raise EmptyDiagramError(name)

        diagram = Diagram(
            nodes=self._registered_systems,
            name=name,
            connection_map=self._connection_map,
            ui_id=ui_id,
        )

        build_recorder.build_diagram(self, diagram, parameters)

        if parameters:
            for name, parameter in parameters.items():
                diagram.declare_dynamic_parameter(name, parameter)
                diagram.instance_parameters.add(name)

        # Export diagram-level inputs
        for locator, port_name in zip(self._input_port_ids, self._input_port_names):
            diagram.export_input(locator, port_name)

        # Export diagram-level outputs
        assert len(self._output_port_ids) == len(self._output_port_names)
        for locator, port_name in zip(self._output_port_ids, self._output_port_names):
            diagram.export_output(locator, port_name)

        self._already_built = True  # Prevent further use of this builder
        self._built_as_name = name
        return diagram

`init(*, validate_rates_at_connect=None, unit_conversion='auto', auto_insert_rate_transitions=False)`

Construct a DiagramBuilder.

Parameters:

Name	Type	Description	Default
`validate_rates_at_connect`	`str \| bool \| None`	T-105 Phase 2 — opt-in connect-time multirate consistency check. When set to `"warn"` (or `True`, treated as `"warn"`), each :meth:`connect` call routes through :func:`jaxonomy.simulation.rate_groups.check_connection_rate_compat` and emits a :class:`RateMismatchWarning` for adjacent blocks of incompatible discrete rates. `"error"` raises :class:`RateMismatchError` on the first offender. Default `None` keeps the legacy path completely off (byte-equivalent to the pre-T-105-Phase-2 behaviour).	`None`
`unit_conversion`	`str`	T-104 followup — controls behaviour when two connected ports share base-dimensions but differ only by a scalar `scale` (e.g. `meter` vs `kilometer`): * `"auto"` (default) silently inserts the conversion factor on the destination input port; * `"warn"` inserts the factor and emits a :class:`UserWarning`; * `"error"` refuses the connection (preserves the Phase-1 strict-equal behaviour). Genuine dimensional mismatches (e.g. `meter` vs `second`) always raise regardless of mode.	`'auto'`
`auto_insert_rate_transitions`	`bool`	T-105-followup-phase3 — when `True`, :meth:`connect` automatically synthesises a :func:`jaxonomy.library.RateTransition` block (a `ZeroOrderHold` for slow→fast, a `Decimator` for fast→slow) between any two adjacent leaves whose inferred discrete sample times differ. The rewritten wiring is `src → rate_transition → dst` and an informational log line documents the insertion. Composes with `validate_rates_at_connect`: when both are enabled, the warning still fires and the transition still gets inserted. Default `False` keeps the legacy code path byte-equivalent (the strict mode that surfaces rate mismatches rather than silently inserting transitions).	`False`

Source code in jaxonomy/framework/diagram_builder.py

def __init__(
    self,
    *,
    validate_rates_at_connect: str | bool | None = None,
    unit_conversion: str = "auto",
    auto_insert_rate_transitions: bool = False,
):
    """Construct a DiagramBuilder.

    Args:
        validate_rates_at_connect: T-105 Phase 2 — opt-in connect-time
            multirate consistency check.  When set to ``"warn"`` (or
            ``True``, treated as ``"warn"``), each :meth:`connect`
            call routes through
            :func:`jaxonomy.simulation.rate_groups.check_connection_rate_compat`
            and emits a :class:`RateMismatchWarning` for adjacent
            blocks of incompatible discrete rates.  ``"error"``
            raises :class:`RateMismatchError` on the first offender.
            Default ``None`` keeps the legacy path completely off
            (byte-equivalent to the pre-T-105-Phase-2 behaviour).
        unit_conversion: T-104 followup — controls behaviour when two
            connected ports share base-dimensions but differ only by a
            scalar ``scale`` (e.g. ``meter`` vs ``kilometer``):
              * ``"auto"`` (default) silently inserts the conversion
                factor on the destination input port;
              * ``"warn"`` inserts the factor and emits a
                :class:`UserWarning`;
              * ``"error"`` refuses the connection (preserves the
                Phase-1 strict-equal behaviour).
            Genuine dimensional mismatches (e.g. ``meter`` vs
            ``second``) always raise regardless of mode.
        auto_insert_rate_transitions: T-105-followup-phase3 — when
            ``True``, :meth:`connect` automatically synthesises a
            :func:`jaxonomy.library.RateTransition` block (a
            ``ZeroOrderHold`` for slow→fast, a ``Decimator`` for
            fast→slow) between any two adjacent leaves whose
            inferred discrete sample times differ.  The rewritten
            wiring is ``src → rate_transition → dst`` and an
            informational log line documents the insertion.
            Composes with ``validate_rates_at_connect``: when both
            are enabled, the warning still fires and the transition
            still gets inserted.  Default ``False`` keeps the legacy
            code path byte-equivalent (the strict mode that surfaces
            rate mismatches rather than silently inserting transitions).
    """
    # Child input ports that are exported as diagram-level inputs
    self._input_port_ids: List[InputPortLocator] = []
    self._input_port_names: List[str] = []
    # Child output ports that are exported as diagram-level outputs
    self._output_port_ids: List[OutputPortLocator] = []
    self._output_port_names: List[str] = []

    # Connection map between input and output ports of the child systems
    self._connection_map: Mapping[InputPortLocator, OutputPortLocator] = {}

    # List of registered systems
    self._registered_systems: List[SystemBase] = []

    # Name lookup for input ports
    self._diagram_input_indices: Mapping[str, InputPortLocator] = {}

    # All input ports of child systems (for use in ensuring proper connectivity)
    self._all_input_ports: List[InputPortLocator] = []

    # Each DiagramBuilder can only be used to build a single diagram.  This is to
    # avoid creating multiple diagrams that reference the same LeafSystem. Doing so
    # may or may not actually lead to problems, since the LeafSystems themselves
    # should act like a collection of pure functions, but best practice is to have
    # each leaf system be fully unique.
    self._already_built = False
    self._built_as_name = None

    # T-105 Phase 2: normalise the connect-time-rate-validation flag
    # to the same string vocabulary that ``check_connection_rate_compat``
    # understands.  ``None`` means "off" (byte-equivalent default);
    # ``True`` is sugar for ``"warn"``.
    if validate_rates_at_connect is True:
        self._validate_rates_at_connect: str | None = "warn"
    elif validate_rates_at_connect in (None, False):
        self._validate_rates_at_connect = None
    elif validate_rates_at_connect in ("warn", "error"):
        self._validate_rates_at_connect = validate_rates_at_connect
    else:
        raise BuilderError(
            "validate_rates_at_connect must be None, True/False, "
            "'warn', or 'error'; got "
            f"{validate_rates_at_connect!r}"
        )

    # T-104 followup: unit-conversion mode.  Validated up front so
    # typos surface at construction time rather than from the first
    # ``connect`` call.
    if unit_conversion not in ("auto", "warn", "error"):
        raise BuilderError(
            "unit_conversion must be 'auto', 'warn', or 'error'; "
            f"got {unit_conversion!r}"
        )
    self._unit_conversion = unit_conversion

    # T-105-followup-phase3: opt-in auto-insertion of RateTransition
    # blocks at connect time.  Default ``False`` preserves the legacy
    # behaviour (byte-equivalent).  When enabled, ``connect`` will
    # synthesise a RateTransition block between any two adjacent
    # leaves whose inferred discrete sample times differ.  The
    # counter below gives auto-inserted blocks unique names without
    # leaking into the user-visible name space when the flag is off.
    self._auto_insert_rate_transitions = bool(auto_insert_rate_transitions)
    self._auto_rate_transition_counter = 0

`add(*systems)`

add(system: SystemBase) -> SystemBase

add(system: SystemBase, *systems: SystemBase) -> List[SystemBase]

Add one or more systems to the diagram.

Parameters:

Name	Type	Description	Default
`*systems SystemBase`		System(s) to add to the diagram.	required

Returns:

Type	Description
`List[SystemBase] \| SystemBase`	List[SystemBase] \| SystemBase: The added system(s). Will return a single system if there is only a single system in the argument list.

Raises:

Type	Description
`BuilderError`	If the diagram has already been built.
`BuilderError`	If the system is already registered.
`BuilderError`	If the system name is not unique.

Source code in jaxonomy/framework/diagram_builder.py

def add(self, *systems: SystemBase) -> List[SystemBase] | SystemBase:
    """Add one or more systems to the diagram.

    Args:
        *systems SystemBase:
            System(s) to add to the diagram.

    Returns:
        List[SystemBase] | SystemBase:
            The added system(s). Will return a single system if there is only
            a single system in the argument list.

    Raises:
        BuilderError: If the diagram has already been built.
        BuilderError: If the system is already registered.
        BuilderError: If the system name is not unique.
    """
    for system in systems:
        self._check_not_already_built()
        self._check_system_not_registered(system)
        self._check_system_name_is_unique(system)
        self._registered_systems.append(system)

        # Add the system's input ports to the list of all input ports
        # So that we can make sure they're all connected before building.
        self._all_input_ports.extend([port.locator for port in system.input_ports])

        logger.debug("Added system %s to DiagramBuilder", system.name)
        logger.debug(
            "    Registered systems: %s",
            [s.name for s in self._registered_systems],
        )
    build_recorder.add_block(self, systems)

    return systems[0] if len(systems) == 1 else systems

`build(name='root', ui_id=None, parameters=None)`

Builds a Diagram system with the specified name and system ID.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the diagram. Defaults to "root".	`'root'`
`ui_id`	`str`	The unique identifier for the diagram.	`None`
`parameters`	`dict[str, Parameter]`	A dictionary of dynamic parameters to declare for the diagram.	`None`

Returns:

Name	Type	Description
`Diagram`	`Diagram`	The newly constructed diagram.

Raises:

Type	Description
`EmptyDiagramError`	If no systems are registered in the diagram.
`BuilderError`	If the diagram has already been built.
`AlgebraicLoopError`	If an algebraic loop is detected in the diagram.
`DisconnectedInputError`	If an input port is not connected.

Source code in jaxonomy/framework/diagram_builder.py

def build(
    self,
    name: str = "root",
    ui_id: str = None,
    parameters: dict[str, Parameter] = None,
) -> Diagram:
    """Builds a Diagram system with the specified name and system ID.

    Args:
        name (str, optional): The name of the diagram. Defaults to "root".
        ui_id (str, optional): The unique identifier for the diagram.
        parameters (dict[str, Parameter], optional):
            A dictionary of dynamic parameters to declare for the diagram.

    Returns:
        Diagram: The newly constructed diagram.

    Raises:
        EmptyDiagramError: If no systems are registered in the diagram.
        BuilderError: If the diagram has already been built.
        AlgebraicLoopError: If an algebraic loop is detected in the diagram.
        DisconnectedInputError: If an input port is not connected.
    """
    self._check_not_already_built()
    self._check_contents_are_complete()
    self._check_ports_are_valid()

    # Check that all internal input ports are connected
    for input_port_locator in self._input_port_ids:
        self._check_input_is_connected(input_port_locator)

    if len(self._registered_systems) == 0:
        raise EmptyDiagramError(name)

    diagram = Diagram(
        nodes=self._registered_systems,
        name=name,
        connection_map=self._connection_map,
        ui_id=ui_id,
    )

    build_recorder.build_diagram(self, diagram, parameters)

    if parameters:
        for name, parameter in parameters.items():
            diagram.declare_dynamic_parameter(name, parameter)
            diagram.instance_parameters.add(name)

    # Export diagram-level inputs
    for locator, port_name in zip(self._input_port_ids, self._input_port_names):
        diagram.export_input(locator, port_name)

    # Export diagram-level outputs
    assert len(self._output_port_ids) == len(self._output_port_names)
    for locator, port_name in zip(self._output_port_ids, self._output_port_names):
        diagram.export_output(locator, port_name)

    self._already_built = True  # Prevent further use of this builder
    self._built_as_name = name
    return diagram

`connect(src, dest)`

Connect an output port to an input port.

The input port and output port must both belong to systems that have already been added to the diagram. The input port must not already be connected to another output port.

Parameters:

Name	Type	Description	Default
`src`	`OutputPort`	The output port to connect.	required
`dest`	`InputPort`	The input port to connect.	required

Raises:

Type	Description
`BuilderError`	If the diagram has already been built.
`BuilderError`	If the source system is not registered.
`BuilderError`	If the destination system is not registered.
`BuilderError`	If the input port is already connected.
`BuilderError`	If src is an InputPort or dest is an OutputPort.

Source code in jaxonomy/framework/diagram_builder.py

def connect(self, src: OutputPort, dest: InputPort):
    """Connect an output port to an input port.

    The input port and output port must both belong to systems that have
    already been added to the diagram.  The input port must not already be
    connected to another output port.

    Args:
        src (OutputPort): The output port to connect.
        dest (InputPort): The input port to connect.

    Raises:
        BuilderError: If the diagram has already been built.
        BuilderError: If the source system is not registered.
        BuilderError: If the destination system is not registered.
        BuilderError: If the input port is already connected.
        BuilderError: If src is an InputPort or dest is an OutputPort.
    """
    # Local imports to avoid a hard import cycle at module load.
    from .port import InputPort as _InputPort, OutputPort as _OutputPort

    # Direction validation -- catches the common (input, input) /
    # (output, output) miswiring at connect time rather than letting it
    # surface as an opaque "input not connected" error during simulation.
    src_is_input = isinstance(src, _InputPort)
    dest_is_output = isinstance(dest, _OutputPort)
    if src_is_input and isinstance(dest, _InputPort):
        raise BuilderError(
            f"Cannot connect input-to-input: "
            f"'{src.system.name}.in[{src.index}]' -> "
            f"'{dest.system.name}.in[{dest.index}]'. "
            f"The first argument must be an output port."
        )
    if isinstance(src, _OutputPort) and dest_is_output:
        raise BuilderError(
            f"Cannot connect output-to-output: "
            f"'{src.system.name}.out[{src.index}]' -> "
            f"'{dest.system.name}.out[{dest.index}]'. "
            f"The second argument must be an input port."
        )
    if src_is_input:
        raise BuilderError(
            f"connect() expected an OutputPort as the first argument, got "
            f"InputPort '{src.system.name}.in[{src.index}]'."
        )
    if dest_is_output:
        raise BuilderError(
            f"connect() expected an InputPort as the second argument, got "
            f"OutputPort '{dest.system.name}.out[{dest.index}]'."
        )

    self._check_not_already_built()
    self._check_system_is_registered(src.system)
    self._check_system_is_registered(dest.system)
    self._check_input_not_connected(dest.locator)

    # T-104 phase 1 / followup: connect-time unit consistency check.
    # Ports that never declared a unit are treated as dimensionless and
    # connect to anything (default-off byte-equivalence).  When both
    # sides declare units:
    #   * dimensional mismatch (e.g. m vs s) -> always UnitMismatchError;
    #   * scalar-scale mismatch (e.g. m vs km):
    #       - "error": raise (Phase-1 strict behaviour);
    #       - "warn":  apply factor + emit UserWarning;
    #       - "auto":  silently apply factor.
    src_units = getattr(src, "units", None)
    dst_units = getattr(dest, "units", None)
    src_label = f"'{src.system.name}.out[{src.index}]' ({src.name})"
    dst_label = f"'{dest.system.name}.in[{dest.index}]' ({dest.name})"

    if self._unit_conversion == "error":
        # Preserve the Phase-1 strict-equal behaviour.
        assert_unit_compatible(
            src_units,
            dst_units,
            src_label=src_label,
            dst_label=dst_label,
        )
    else:
        # Returns the multiplicative factor (1.0 for matched / wildcard
        # units, src.scale / dst.scale otherwise).  Raises on genuine
        # dimensional mismatch.
        factor = assert_units_compatible_with_scale(
            src_units,
            dst_units,
            src_label=src_label,
            dst_label=dst_label,
        )
        if factor != 1.0:
            src_u = resolve_unit(src_units)
            dst_u = resolve_unit(dst_units)
            msg = (
                f"Unit conversion: applying factor {factor!r} to "
                f"connection {src_label} ({src_u!r}) -> "
                f"{dst_label} ({dst_u!r})."
            )
            if self._unit_conversion == "warn":
                warnings.warn(msg, UserWarning, stacklevel=2)
            else:
                # "auto": informational log line; no warning.
                logger.info(msg)
            _install_unit_conversion(dest, factor)

    # T-105 Phase 2: opt-in connect-time multirate consistency check.
    # Default-off (``self._validate_rates_at_connect is None``) keeps
    # the legacy code path byte-equivalent.  When the builder was
    # constructed with ``validate_rates_at_connect=...``, route the
    # source/dest pair through ``check_connection_rate_compat`` which
    # honours the ``_jaxonomy_rate_transition`` marker (T-123) and
    # the universal-sample-time rule (constant / inherited bridge
    # any rates).
    if self._validate_rates_at_connect is not None:
        # Local import: avoids a hard import cycle between
        # ``framework`` and ``simulation`` at module load.
        from ..simulation.rate_groups import check_connection_rate_compat

        check_connection_rate_compat(
            src.system,
            src.index,
            dest.system,
            dest.index,
            on_mismatch=self._validate_rates_at_connect,
        )

    # T-105-followup-phase3: opt-in connect-time auto-insertion of
    # RateTransition blocks between adjacent leaves with differing
    # discrete sample times.  Default-off keeps the legacy code
    # path byte-equivalent.  Composes with the validate hook above:
    # when both are enabled, the warning still fires (above) and
    # the bridge still gets inserted (here).
    if self._auto_insert_rate_transitions:
        inserted = self._maybe_auto_insert_rate_transition(src, dest)
        if inserted is not None:
            # ``_maybe_auto_insert_rate_transition`` already wrote
            # both legs of ``src → bridge → dest`` into the
            # connection map (and recorded both connections with
            # the build recorder), so we are done here.
            return

    build_recorder.connect_ports(self, src, dest)

    self._connection_map[dest.locator] = src.locator

    logger.debug(
        f"Connected port {src.name} of system {src.system.name} to port {dest.name} of system {dest.system.name}"
    )

`export_input(port, name=None)`

Export an input port of a child system as a diagram-level input.

The input port must belong to a system that has already been added to the diagram. The input port must not already be connected to another output port.

Parameters:

Name	Type	Description	Default
`port`	`InputPort`	The input port to export.	required
`name`	`str`	The name to assign to the exported input port. If not provided, a unique name will be generated.	`None`

Returns:

Name	Type	Description
`int`	`int`	The index (in the to-be-built diagram) of the exported input port.

Raises:

Type	Description
`BuilderError`	If the diagram has already been built.
`BuilderError`	If the system is not registered.
`BuilderError`	If the input port is already connected.
`BuilderError`	If the input port name is not unique.

Source code in jaxonomy/framework/diagram_builder.py

def export_input(self, port: InputPort, name: str = None) -> int:
    """Export an input port of a child system as a diagram-level input.

    The input port must belong to a system that has already been added to the
    diagram. The input port must not already be connected to another output port.

    Args:
        port (InputPort): The input port to export.
        name (str, optional):
            The name to assign to the exported input port. If not provided, a
            unique name will be generated.

    Returns:
        int: The index (in the to-be-built diagram) of the exported input port.

    Raises:
        BuilderError: If the diagram has already been built.
        BuilderError: If the system is not registered.
        BuilderError: If the input port is already connected.
        BuilderError: If the input port name is not unique.
    """
    self._check_not_already_built()
    self._check_system_is_registered(port.system)
    self._check_input_not_connected(port.locator)

    if name is None:
        # Since the system names are unique, auto-generated port names are also unique
        # at the level of _this_ diagram (subsystems can have ports with the same name)
        name = f"{port.system.name}_{port.name}"
    elif name in self._diagram_input_indices:
        raise BuilderError(
            f"Input port name {name} is not unique",
            system=port.system,
            port_index=port.index,
            port_direction="in",
        )

    # Index at the diagram (not subsystem) level
    port_index = len(self._input_port_ids)
    self._input_port_ids.append(port.locator)
    self._input_port_names.append(name)

    self._diagram_input_indices[name] = port_index

    build_recorder.export_port(self, port.system, "input", port.index, name)

    return port_index

`export_output(port, name=None)`

Export an output port of a child system as a diagram-level output.

The output port must belong to a system that has already been added to the diagram.

Parameters:

Name	Type	Description	Default
`port`	`OutputPort`	The output port to export.	required
`name`	`str`	The name to assign to the exported output port. If not provided, a unique name will be generated.	`None`

Returns:

Name	Type	Description
`int`	`int`	The index (in the to-be-built diagram) of the exported output port.

Raises:

Type	Description
`BuilderError`	If the diagram has already been built.
`BuilderError`	If the system is not registered.
`BuilderError`	If the output port name is not unique.

Source code in jaxonomy/framework/diagram_builder.py

def export_output(self, port: OutputPort, name: str = None) -> int:
    """Export an output port of a child system as a diagram-level output.

    The output port must belong to a system that has already been added to the
    diagram.

    Args:
        port (OutputPort): The output port to export.
        name (str, optional):
            The name to assign to the exported output port. If not provided, a
            unique name will be generated.

    Returns:
        int: The index (in the to-be-built diagram) of the exported output port.

    Raises:
        BuilderError: If the diagram has already been built.
        BuilderError: If the system is not registered.
        BuilderError: If the output port name is not unique.
    """
    self._check_not_already_built()
    self._check_system_is_registered(port.system)

    if name is None:
        # Since the system names are unique, auto-generated port names are also unique
        # at the level of _this_ diagram (subsystems can have ports with the same name)
        name = f"{port.system.name}_{port.name}"
    elif name in self._output_port_names:
        raise BuilderError(
            f"Output port name {name} is not unique",
            system=port.system,
            port_index=port.index,
            port_direction="out",
        )

    # Index at the diagram (not subsystem) level
    port_index = len(self._output_port_ids)
    self._output_port_ids.append(port.locator)
    self._output_port_names.append(name)

    build_recorder.export_port(self, port.system, "output", port.index, name)

    return port_index

`DiagramContext` `dataclass`

Bases: ContextBase

Source code in jaxonomy/framework/context.py

@dataclasses.dataclass(frozen=True)
class DiagramContext(ContextBase):
    subcontexts: OrderedDict[Hashable, LeafContext] = dataclasses.field(
        default_factory=OrderedDict
    )

    def _check_key(self, key: Hashable) -> None:
        _reject_port_key(key)
        assert key == self.owning_system.system_id or key in self.subcontexts, (
            f"System ID {key} not found in DiagramContext {self}.\nIf this ID "
            "references an intermediate diagram, note that intermediate diagrams do "
            "not have associated contexts. Only the root diagram and leaf systems have "
            "contexts."
        )

    def __getitem__(self, key: Hashable) -> LeafContext:
        self._check_key(key)
        if key == self.owning_system.system_id:
            return self
        return self.subcontexts[key]

    def find_context_with_path(self, path: list[str]) -> ContextBase:
        system = self.owning_system.find_system_with_path(path)
        if system is None:
            raise ValueError(
                f"No system with path {path} found in {self.owning_system}"
            )
        return self[system.system_id]

    def with_subcontext(self, key: Hashable, ctx: LeafContext) -> DiagramContext:
        self._check_key(key)
        subcontexts = self.subcontexts.copy()
        subcontexts[key] = ctx
        return dataclasses.replace(self, subcontexts=subcontexts)

    #
    # Simulation interface
    #
    @property
    def state(self) -> Mapping[Hashable, LeafState]:
        return OrderedDict(
            {system_id: subctx.state for system_id, subctx in self.subcontexts.items()}
        )

    @property
    def continuous_subcontexts(self) -> List[LeafContext]:
        return [
            subctx
            for subctx in self.subcontexts.values()
            if subctx.has_continuous_state
        ]

    @property
    def continuous_state(self) -> List[Array]:
        return [subctx.continuous_state for subctx in self.continuous_subcontexts]

    def with_continuous_state(self, sub_xcs: List[Array]) -> DiagramContext:
        # Shallow copy the subcontexts - only modify the ones that have continuous states
        new_subcontexts = self.subcontexts.copy()
        for subctx, sub_xc in zip(self.continuous_subcontexts, sub_xcs):
            new_subcontexts[subctx.system_id] = subctx.with_continuous_state(sub_xc)
        return dataclasses.replace(self, subcontexts=new_subcontexts)

    @property
    def num_continuous_states(self) -> int:
        return sum(
            [subctx.num_continuous_states for subctx in self.subcontexts.values()]
        )

    @property
    def has_continuous_state(self) -> bool:
        return self.num_continuous_states > 0

    @property
    def discrete_subcontexts(self) -> List[LeafContext]:
        return [
            subctx for subctx in self.subcontexts.values() if subctx.has_discrete_state
        ]

    @property
    def discrete_state(self) -> List[List[Array]]:
        return [subctx.discrete_state for subctx in self.discrete_subcontexts]

    def with_discrete_state(self, sub_xds: List[List[Array]]) -> DiagramContext:
        # Shallow copy the subcontexts - only modify the ones that have discrete states
        new_subcontexts = self.subcontexts.copy()
        for subctx, sub_xd in zip(self.discrete_subcontexts, sub_xds):
            new_subcontexts[subctx.system_id] = subctx.with_discrete_state(sub_xd)
        return dataclasses.replace(self, subcontexts=new_subcontexts)

    @property
    def num_discrete_states(self) -> int:
        return sum([subctx.num_discrete_states for subctx in self.subcontexts.values()])

    @property
    def has_discrete_state(self) -> bool:
        return self.num_discrete_states > 0

    @property
    def mode_subcontexts(self) -> List[LeafContext]:
        return [subctx for subctx in self.subcontexts.values() if subctx.has_mode]

    @property
    def mode(self) -> List[int]:
        return [subctx.mode for subctx in self.mode_subcontexts]

    def with_mode(self, sub_modes: List[int]) -> DiagramContext:
        new_subcontexts = self.subcontexts.copy()
        for subctx, sub_mode in zip(self.mode_subcontexts, sub_modes):
            new_subcontexts[subctx.system_id] = subctx.with_mode(sub_mode)
        return dataclasses.replace(self, subcontexts=new_subcontexts)

    @property
    def has_mode(self) -> bool:
        return any([subctx.has_mode for subctx in self.subcontexts.values()])

    def with_state(self, sub_states: Mapping[Hashable, LeafState]) -> DiagramContext:
        new_subcontexts = OrderedDict()
        for system_id, sub_state in sub_states.items():
            new_subcontexts[system_id] = dataclasses.replace(
                self.subcontexts[system_id], state=sub_state
            )
        return dataclasses.replace(self, subcontexts=new_subcontexts)

    def with_new_state(self) -> ContextBase:
        new_subcontexts = OrderedDict()
        for system_id, subctx in self.subcontexts.items():
            new_subcontexts[system_id] = subctx.with_new_state()
        return dataclasses.replace(self, subcontexts=new_subcontexts)

    def with_updated_parameters(self) -> ContextBase:
        new_parameters = {
            name: param.get()
            for name, param in self.owning_system.dynamic_parameters.items()
        }
        new_subcontexts = {}
        for k, v in self.subcontexts.items():
            new_subcontexts[k] = v.with_updated_parameters()

        return dataclasses.replace(
            self, subcontexts=new_subcontexts, parameters=new_parameters
        )

    def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
        """Create a copy of this context, replacing only the specified parameters."""
        parameters = {**self.parameters}

        # First validate that all parameters exist and are dynamic
        for name, value in new_parameters.items():
            if name not in self.owning_system.dynamic_parameters:
                raise ValueError(
                    f"Parameter {name} not found in {self.owning_system.name}"
                )
            param = self.owning_system.dynamic_parameters[name]

            if param.static_dependents:
                static_dependents = ", ".join(
                    [f"{dep.system.name}" for dep in param.static_dependents]
                )
                raise StaticParameterError(
                    f"Parameter {name} is used in static parameters"
                    " and cannot be updated dynamically. Please create a new context."
                    f" Static dependents in blocks: {static_dependents}"
                )

        for name, value in new_parameters.items():
            value = _coerce_param_for_jit_cache(parameters.get(name), value)
            self.owning_system.dynamic_parameters[name].set(value)
            parameters[name] = value

        context = dataclasses.replace(self, parameters=parameters)
        return context.with_updated_parameters()

`with_parameters(new_parameters)`

Create a copy of this context, replacing only the specified parameters.

Source code in jaxonomy/framework/context.py

def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
    """Create a copy of this context, replacing only the specified parameters."""
    parameters = {**self.parameters}

    # First validate that all parameters exist and are dynamic
    for name, value in new_parameters.items():
        if name not in self.owning_system.dynamic_parameters:
            raise ValueError(
                f"Parameter {name} not found in {self.owning_system.name}"
            )
        param = self.owning_system.dynamic_parameters[name]

        if param.static_dependents:
            static_dependents = ", ".join(
                [f"{dep.system.name}" for dep in param.static_dependents]
            )
            raise StaticParameterError(
                f"Parameter {name} is used in static parameters"
                " and cannot be updated dynamically. Please create a new context."
                f" Static dependents in blocks: {static_dependents}"
            )

    for name, value in new_parameters.items():
        value = _coerce_param_for_jit_cache(parameters.get(name), value)
        self.owning_system.dynamic_parameters[name].set(value)
        parameters[name] = value

    context = dataclasses.replace(self, parameters=parameters)
    return context.with_updated_parameters()

`DiscreteUpdateEvent` `dataclass`

Bases: Event

Event representing a discrete update in a hybrid system.

Source code in jaxonomy/framework/event.py

@tree_util.register_pytree_node_class
@dataclasses.dataclass
class DiscreteUpdateEvent(Event):
    """Event representing a discrete update in a hybrid system."""

    # Supersede type hints in Event with the specific signature for discrete updates
    callback: Callable[[ContextBase], Array] = None
    passthrough: Callable[[ContextBase], Array] = None

    # Inherits docstring from Event. This is only needed to specialize type hints.
    def handle(self, context: ContextBase) -> Array:
        return super().handle(context)

`DtypeMismatchError`

Bases: StaticError

Block parameters or input/outputs have mismatched dtypes.

Source code in jaxonomy/framework/error.py

class DtypeMismatchError(StaticError):
    """Block parameters or input/outputs have mismatched dtypes."""

    def __init__(self, expected_dtype=None, actual_dtype=None, **kwargs):
        super().__init__(**kwargs)
        self.expected_dtype = expected_dtype
        self.actual_dtype = actual_dtype

    def __str__(self):
        if self.expected_dtype or self.actual_dtype:
            return (
                f"Data type mismatch: "
                f"expected {self.expected_dtype}, got {self.actual_dtype}"
                + self._context_info()
            )
        return f"Dtype mismatch{self._context_info()}"

`EnabledMode`

Allowed string values for EnabledSubsystem.mode.

Source code in jaxonomy/framework/containers.py

class EnabledMode:
    """Allowed string values for ``EnabledSubsystem.mode``."""

    RESET = "reset"
    HOLD = "hold"
    PASSTHROUGH = "passthrough"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.RESET, cls.HOLD, cls.PASSTHROUGH)

`EnabledStateMode`

Allowed string values for EnabledSubsystem.state_mode.

Controls how the continuous state (declared via state_dynamics) evolves while the enable signal is false:

HOLD (default): freeze the state at its current value (xdot = 0 while disabled). Resumes integration on re-enable.
RESET: snap the state back to initial_state on every disable→enable transition (so each enable window starts from the configured initial value). While disabled, the state is held.
FREE: the state evolves according to state_dynamics regardless of enable. Only the output is masked per mode=.

Source code in jaxonomy/framework/containers.py

class EnabledStateMode:
    """Allowed string values for ``EnabledSubsystem.state_mode``.

    Controls how the *continuous state* (declared via ``state_dynamics``)
    evolves while the enable signal is false:

    - ``HOLD`` (default): freeze the state at its current value
      (``xdot = 0`` while disabled). Resumes integration on re-enable.
    - ``RESET``: snap the state back to ``initial_state`` on every
      disable→enable transition (so each enable window starts from the
      configured initial value). While disabled, the state is held.
    - ``FREE``: the state evolves according to ``state_dynamics``
      regardless of enable. Only the *output* is masked per ``mode=``.
    """

    HOLD = "hold"
    RESET = "reset"
    FREE = "free"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.HOLD, cls.RESET, cls.FREE)

`EnabledSubsystem`

Bases: LeafSystem

Container block: run a submodel only while an enable signal is true.

This is the subsystem-framing wrapper around the existing :class:jaxonomy.library.Conditional primitive (T-009). It exists as a separate class so that:

The block-diagram-vocabulary name EnabledSubsystem is discoverable next to the rest of the container family.
We can later extend the mode="hold" path with subsystem-state semantics (per-block discrete-state binding) without disturbing the lighter Conditional primitive.

Parameters:

Name	Type	Description	Default
`submodel`	`Callable`	Callable `f(*inputs) -> output` (single output per phase 1). Must be JAX-traceable.	required
`n_inputs`	`int`	Number of submodel inputs (does NOT include the enable port). Input port 0 is always the enable signal; ports 1..n_inputs carry the submodel inputs.	`1`
`mode`	`Literal['reset', 'passthrough', 'hold']`	One of `"reset"` / `"passthrough"` / `"hold"`. `reset`: when disabled, output = `initial_value`. `passthrough`: when disabled, output = first user input (input port 1). Submodel and passthrough output must broadcast-compatibly. `hold`: when disabled, output holds the most recent snapshot taken at `hold_period`. Requires a positive `hold_period`.	`RESET`
`initial_value`		Output value when disabled in reset mode, and the seed for the held discrete state in hold mode. Used to infer output shape/dtype.	`0.0`
`hold_period`	`float \| None`	Sample period (seconds) for the held snapshot in hold mode. Required iff `mode == "hold"`.	`None`
`state_mode`	`Literal['hold', 'reset', 'free']`	One of `"hold"` / `"reset"` / `"free"`. Controls the continuous-state behaviour while disabled (independent of `mode=` which gates only the output). See :class:`EnabledStateMode`. Default `"hold"`. Only has an effect when `state_dynamics` is provided; for the stateless submodel default this kwarg is validated but otherwise a no-op (so the default-off path is byte-equivalent to phase 1).	`HOLD`
`state_dynamics`	`Callable \| None`	Optional callable `f(t, x, *user_inputs) -> xdot` defining a continuous state for the EnabledSubsystem itself. When provided, the block declares a continuous state seeded by `initial_state` (or `initial_value` if `initial_state` is None) and applies `state_mode` semantics around it. When omitted, the block has no continuous state and behaves exactly as in T-120 phase 1.	`None`
`initial_state`		Initial value of the continuous state. Required when `state_dynamics` is provided.	`None`
`name`		Optional block name.	required

Source code in jaxonomy/framework/containers.py

class EnabledSubsystem(LeafSystem):
    """Container block: run a submodel only while an enable signal is true.

    This is the subsystem-framing wrapper around the existing
    :class:`jaxonomy.library.Conditional` primitive (T-009). It exists
    as a separate class so that:

    - The block-diagram-vocabulary name ``EnabledSubsystem`` is discoverable
      next to the rest of the container family.
    - We can later extend the ``mode="hold"`` path with subsystem-state
      semantics (per-block discrete-state binding) without disturbing
      the lighter ``Conditional`` primitive.

    Args:
        submodel: Callable ``f(*inputs) -> output`` (single output per
            phase 1). Must be JAX-traceable.
        n_inputs: Number of submodel inputs (does NOT include the
            enable port). Input port 0 is always the enable signal;
            ports 1..n_inputs carry the submodel inputs.
        mode: One of ``"reset"`` / ``"passthrough"`` / ``"hold"``.

            - ``reset``: when disabled, output = ``initial_value``.
            - ``passthrough``: when disabled, output = first user input
              (input port 1). Submodel and passthrough output must
              broadcast-compatibly.
            - ``hold``: when disabled, output holds the most recent
              snapshot taken at ``hold_period``. Requires a positive
              ``hold_period``.
        initial_value: Output value when disabled in reset mode, and
            the seed for the held discrete state in hold mode. Used to
            infer output shape/dtype.
        hold_period: Sample period (seconds) for the held snapshot in
            hold mode. Required iff ``mode == "hold"``.
        state_mode: One of ``"hold"`` / ``"reset"`` / ``"free"``.
            Controls the *continuous-state* behaviour while disabled
            (independent of ``mode=`` which gates only the output).
            See :class:`EnabledStateMode`. Default ``"hold"``. Only has
            an effect when ``state_dynamics`` is provided; for the
            stateless submodel default this kwarg is validated but
            otherwise a no-op (so the default-off path is byte-equivalent
            to phase 1).
        state_dynamics: Optional callable
            ``f(t, x, *user_inputs) -> xdot`` defining a continuous
            state for the EnabledSubsystem itself. When provided, the
            block declares a continuous state seeded by ``initial_state``
            (or ``initial_value`` if ``initial_state`` is None) and
            applies ``state_mode`` semantics around it. When omitted,
            the block has no continuous state and behaves exactly as in
            T-120 phase 1.
        initial_state: Initial value of the continuous state. Required
            when ``state_dynamics`` is provided.
        name: Optional block name.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        mode: Literal["reset", "passthrough", "hold"] = EnabledMode.RESET,
        initial_value=0.0,
        hold_period: float | None = None,
        state_mode: Literal["hold", "reset", "free"] = EnabledStateMode.HOLD,
        state_dynamics: Callable | None = None,
        initial_state=None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if mode not in EnabledMode.valid():
            raise ValueError(
                f"EnabledSubsystem: mode must be one of "
                f"{EnabledMode.valid()!r}, got {mode!r}"
            )
        if state_mode not in EnabledStateMode.valid():
            raise ValueError(
                f"EnabledSubsystem: state_mode must be one of "
                f"{EnabledStateMode.valid()!r}, got {state_mode!r}"
            )
        if mode == EnabledMode.HOLD and not hold_period:
            raise ValueError(
                "EnabledSubsystem(mode='hold') requires a positive "
                "hold_period to determine the snapshot sample rate."
            )
        if n_inputs < 0:
            raise ValueError(
                f"EnabledSubsystem: n_inputs must be >= 0, got {n_inputs}"
            )
        if mode == EnabledMode.PASSTHROUGH and n_inputs < 1:
            raise ValueError(
                "EnabledSubsystem(mode='passthrough') requires at least "
                "one user input to use as the bypass signal."
            )
        if state_dynamics is not None and initial_state is None:
            raise ValueError(
                "EnabledSubsystem: state_dynamics= requires an "
                "initial_state= value (the seed for the continuous "
                "state)."
            )

        self._submodel = submodel
        self._mode = mode
        self._initial = jnp.asarray(initial_value)
        self._state_mode = state_mode
        self._state_dynamics = state_dynamics
        self._initial_state = (
            jnp.asarray(initial_state) if initial_state is not None else None
        )

        # Port 0 is always enable; remaining ports are submodel inputs.
        self.declare_input_port(name="enable")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        if mode == EnabledMode.HOLD:
            self.declare_discrete_state(default_value=self._initial)
            self.declare_periodic_update(
                self._hold_update,
                period=float(hold_period),
                offset=0.0,
            )

        # T-120-followup-enabled-cont-state: declare a continuous state on
        # this block when the user supplied a state_dynamics callable. The
        # state_mode kwarg controls how the state evolves vs. the enable
        # signal. The default-off path (state_dynamics=None) bypasses this
        # block entirely → byte-equivalent to phase 1.
        if self._state_dynamics is not None:
            self.declare_continuous_state(
                default_value=self._initial_state,
                ode=self._wrapped_ode,
            )

            if self._state_mode == EnabledStateMode.RESET:
                # Snap the continuous state back to its initial value on
                # every enable transition (rising or falling), so each
                # disable window leaves the state at the seed and each
                # re-enable starts from the same point. Pair this with
                # the "hold" ode-zeroing during the disabled window so
                # the state actually stays at the seed instead of drifting.
                #
                # Treat the enable signal as a centred continuous guard
                # (``enable - 0.5``): zero-crossings of this expression
                # match enable transitions in either direction. The
                # framework's continuous detector preserves the float
                # carry-type for the guard signal — using
                # ``direction="edge_detection"`` with a boolean guard
                # would conflict with the simulator's float-typed
                # internal carry, so we go through the continuous path.
                self.declare_zero_crossing(
                    guard=self._enable_guard,
                    reset_map=self._reset_continuous_state,
                    direction="crosses_zero",
                    name="enable_reset",
                )

        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _submodel_output(self, inputs):
        user_inputs = inputs[1:]  # skip enable
        return jnp.asarray(self._submodel(*user_inputs))

    def _hold_update(self, time, state, *inputs, **params):
        enable = jnp.asarray(inputs[0]).astype(bool)
        y_sub = self._submodel_output(inputs)
        return jnp.where(enable, y_sub, state.discrete_state)

    def _compute_output(self, time, state, *inputs, **params):
        enable = jnp.asarray(inputs[0]).astype(bool)
        y_sub = self._submodel_output(inputs)

        if self._mode == EnabledMode.RESET:
            return jnp.where(enable, y_sub, self._initial)

        if self._mode == EnabledMode.HOLD:
            return jnp.where(enable, y_sub, state.discrete_state)

        # passthrough
        bypass = jnp.asarray(inputs[1])
        return jnp.where(enable, y_sub, bypass)

    # ── continuous-state callbacks (T-120-followup-enabled-cont-state) ───

    def _wrapped_ode(self, time, state, *inputs, **params):
        """Apply ``state_mode`` semantics to the user's ``state_dynamics``.

        - ``hold``: multiply the user's xdot by the enable flag. While
          disabled the ode returns zero, so the integrator preserves the
          state across the disabled window.
        - ``reset``: same as ``hold`` for the per-step ode (state stays
          at the seed during the disabled window); the reset to the
          initial value is enforced by a zero-crossing reset_map on the
          enable-edge.
        - ``free``: pass the user's xdot through untouched. The state
          evolves regardless of enable (only the output port is masked).
        """
        enable = jnp.asarray(inputs[0]).astype(bool)
        user_inputs = inputs[1:]
        xc = state.continuous_state
        xdot = self._state_dynamics(time, xc, *user_inputs)
        xdot = jnp.asarray(xdot)

        if self._state_mode == EnabledStateMode.FREE:
            return xdot

        # HOLD and RESET both gate the per-step derivative. RESET layers
        # the additional snap-to-initial behaviour via the zero-crossing
        # event registered in __init__.
        scale = jnp.asarray(enable, dtype=xdot.dtype)
        return xdot * scale

    def _enable_guard(self, time, state, *inputs, **params):
        """Continuous guard for enable transitions.

        Returns ``enable - 0.5`` so any 0↔1 transition crosses zero. Used
        with ``direction="crosses_zero"`` to fire the continuous-state
        reset on either a rising or falling enable edge.
        """
        return jnp.asarray(inputs[0]) - 0.5

    def _reset_continuous_state(self, time, state, *inputs, **params):
        """Snap the continuous state back to its initial value."""
        return state.with_continuous_state(
            jnp.asarray(self._initial_state, dtype=state.continuous_state.dtype)
        )

`ErrorCollector`

Tool used to collect errors related to users model specification. Errors related to user model specification are identified during model static analysis, e.g. context creation, type checking, etc.

An instance of this tool can be created, and then passed down a tree of function calls to collect errors found any where in the tree. Locally in the tree it can be determined whether it is ok to continue or not. This tool enables collecting errors up until the point when continuation is no longer possible.

Note: this latter behavior, where sometimes there is early exit desired, and all other "pipeline" operations are "nullified", might better be implemented using pymonad:Either class.

Source code in jaxonomy/framework/error.py

class ErrorCollector:
    """
    Tool used to collect errors related to users model specification.
    Errors related to user model specification are identified during
    model static analysis, e.g. context creation, type checking, etc.

    An instance of this tool can be created, and then passed down a
    tree of function calls to collect errors found any where in
    the tree. Locally in the tree it can be determined whether it is
    ok to continue or not. This tool enables collecting errors up until
    the point when continuation is no longer possible.

    Note: this latter behavior, where sometimes there is early exit desired,
    and all other "pipeline" operations are "nullified", might better be
    implemented using pymonad:Either class.
    """

    def __init__(self):
        self._disable_collection = False
        self._parent: Optional["ErrorCollector"] = None
        self.errors: list[BaseException] = []

    def add_error(self, error: BaseException):
        """Add an error to the collection."""

        if self._parent is not None:
            self._parent.add_error(error)
            return

        if not self._disable_collection:
            self.errors.append(error)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # Return values: True to suppress the exception, False to propagate it

        if exc_type is not None:
            if self._parent is not None:
                self._parent.add_error(exc_value)
                return True

            self.add_error(exc_value)
            return False

        return True

    @classmethod
    def context(cls, parent: "ErrorCollector" = None) -> "ErrorCollector":
        """A context manager convenience to use when tracing errors.

        Use as:
        ```
        with ErrorCollector.trace(error_context) as ec:
            ...
        ```

        If the parent context is None, then exceptions will pass through without
        being collected. Else, exceptions will be collected in the parent context.
        """

        if parent is None:
            ctx = cls()
            ctx._disable_collection = True
            return ctx

        ctx = cls()
        ctx._parent = parent
        return ctx

`add_error(error)`

Add an error to the collection.

Source code in jaxonomy/framework/error.py

def add_error(self, error: BaseException):
    """Add an error to the collection."""

    if self._parent is not None:
        self._parent.add_error(error)
        return

    if not self._disable_collection:
        self.errors.append(error)

`context(parent=None)` `classmethod`

A context manager convenience to use when tracing errors.

Use as:

with ErrorCollector.trace(error_context) as ec:
    ...

If the parent context is None, then exceptions will pass through without being collected. Else, exceptions will be collected in the parent context.

Source code in jaxonomy/framework/error.py

@classmethod
def context(cls, parent: "ErrorCollector" = None) -> "ErrorCollector":
    """A context manager convenience to use when tracing errors.

    Use as:
    ```
    with ErrorCollector.trace(error_context) as ec:
        ...
    ```

    If the parent context is None, then exceptions will pass through without
    being collected. Else, exceptions will be collected in the parent context.
    """

    if parent is None:
        ctx = cls()
        ctx._disable_collection = True
        return ctx

    ctx = cls()
    ctx._parent = parent
    return ctx

`EventCollection`

A collection of events owned by a system.

Users should not need to interact with these objects directly. They are intended to be used internally by the simulation framework for handling events in hybrid system simulation.

These contain callback functions that update the context in various ways when the event is triggered. There will be different "collections" for each trigger type in simulation (e.g. periodic vs zero-crossing). Within the collections, events are broken out by function (e.g. discrete vs unrestricted updates).

There are separate implementations for leaf and diagram systems, where the DiagramCEventCollection preserves the tree structure of the underlying Diagram. However, the interface in both cases is the same and is identical to the interface defined by EventCollection.

Source code in jaxonomy/framework/event.py

class EventCollection(metaclass=abc.ABCMeta):
    """A collection of events owned by a system.

    Users should not need to interact with these objects directly. They are intended
    to be used internally by the simulation framework for handling events in hybrid
    system simulation.

    These contain callback functions that update the context in various ways
    when the event is triggered. There will be different "collections" for each
    trigger type in simulation (e.g. periodic vs zero-crossing). Within the
    collections, events are broken out by function (e.g. discrete vs unrestricted
    updates).

    There are separate implementations for leaf and diagram systems, where the
    DiagramCEventCollection preserves the tree structure of the underlying
    Diagram. However, the interface in both cases is the same and is identical to
    the interface defined by EventCollection.
    """

    @abc.abstractmethod
    def __getitem__(self, key: Hashable) -> EventCollection:
        pass

    @property
    @abc.abstractmethod
    def events(self) -> List[Event]:
        pass

    @property
    @abc.abstractmethod
    def num_events(self) -> int:
        pass

    @property
    def has_events(self) -> bool:
        return self.num_events > 0

    def __iter__(self):
        return iter(self.events)

    def __len__(self):
        return self.num_events

    @abc.abstractmethod
    def activate(self, activation_fn) -> EventCollection:
        pass

    def mark_all_active(self) -> EventCollection:
        return self.activate(lambda _: True)

    def mark_all_inactive(self) -> EventCollection:
        return self.activate(lambda _: False)

    @property
    def num_active(self) -> int:
        def _get_active(event_data: EventData) -> bool:
            return event_data.active

        active_tree = tree_util.tree_map(
            _get_active,
            self,
            is_leaf=is_event_data,
        )
        return sum(tree_util.tree_leaves(active_tree))

    @property
    def has_active(self) -> bool:
        return self.num_active > 0

    @property
    def has_triggered(self) -> bool:
        def _get_triggered(event_data: EventData) -> bool:
            return event_data.active & event_data.triggered

        triggered_tree = tree_util.tree_map(
            _get_triggered,
            self,
            is_leaf=is_event_data,
        )
        return sum(tree_util.tree_leaves(triggered_tree)) > 0

    @property
    @abc.abstractmethod
    def terminal_events(self) -> EventCollection:
        pass

    @property
    def has_terminal_events(self):
        return self.terminal_events.has_events

    @property
    def has_active_terminal(self) -> bool:
        return self.terminal_events.has_triggered

    def pprint(self, output=print):
        output(self._pprint_helper().strip())

    def _pprint_helper(self, prefix="") -> str:
        s = f"{prefix}|-- \n"
        if len(self.events) > 0:
            s += f"{prefix}    Events:\n"
            for event in self.events:
                s += f"{prefix}    |  {event}\n"
        return s

    def __repr__(self) -> str:
        s = f"{type(self).__name__}("
        if self.has_events:
            s += f"discrete_update: {self.events} "
        s += ")"

        return s

`ForLoop`

Bases: LeafSystem

Container block: run body_fn n_iter times per major step.

ForLoop wraps :func:jax.lax.fori_loop. The block declares a single input port carrying the initial carry value and a single output port returning the carry after n_iter iterations.

Parameters:

Name	Type	Description	Default
`body_fn`	`Callable`	Callable `(i: int, carry) -> carry`. Must be JAX-traceable. The carry pytree must have a fixed structure and shape across iterations (this is a `lax.fori_loop` requirement, not a Jaxonomy choice).	required
`n_iter`	`int`	Number of iterations. Must be a non-negative Python int (static); a runtime-traced `n_iter` would force `lax.while_loop` semantics and is not supported here — use :class:`WhileLoop` for that case.	required
`name`		Optional block name.	required

Differentiability

:func:jax.grad flows through body_fn's parameters and through the initial carry. The loop count n_iter is static and not differentiable.

Example

A body that accumulates i into the carry over 10 iterations yields carry_initial + (0+1+...+9) = carry + 45.

Notes

body_fn must close over any constants it needs; the i-th iteration receives only (i, carry).
Per T-005, default float dtype is float64 unless the active precision policy says otherwise; ForLoop does not cast.

Source code in jaxonomy/framework/containers.py

class ForLoop(LeafSystem):
    """Container block: run ``body_fn`` ``n_iter`` times per major step.

    ``ForLoop`` wraps :func:`jax.lax.fori_loop`. The block declares a
    single input port carrying the *initial carry value* and a single
    output port returning the carry after ``n_iter`` iterations.

    Args:
        body_fn: Callable ``(i: int, carry) -> carry``. Must be
            JAX-traceable. The carry pytree must have a fixed structure
            and shape across iterations (this is a ``lax.fori_loop``
            requirement, not a Jaxonomy choice).
        n_iter: Number of iterations. Must be a non-negative Python int
            (static); a runtime-traced ``n_iter`` would force
            ``lax.while_loop`` semantics and is not supported here —
            use :class:`WhileLoop` for that case.
        name: Optional block name.

    Differentiability:
        :func:`jax.grad` flows through ``body_fn``'s parameters and
        through the initial carry. The loop count ``n_iter`` is static
        and not differentiable.

    Example:
        A body that accumulates ``i`` into the carry over 10
        iterations yields ``carry_initial + (0+1+...+9) = carry + 45``.

    Notes:
        - ``body_fn`` must close over any constants it needs; the
          ``i``-th iteration receives only ``(i, carry)``.
        - Per T-005, default float dtype is float64 unless the active
          precision policy says otherwise; ``ForLoop`` does not cast.
    """

    def __init__(
        self,
        body_fn: Callable,
        n_iter: int,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if not isinstance(n_iter, int):
            raise TypeError(
                f"ForLoop: n_iter must be a Python int (static), got "
                f"{type(n_iter).__name__}"
            )
        if n_iter < 0:
            raise ValueError(
                f"ForLoop: n_iter must be >= 0, got {n_iter}"
            )

        self._body_fn = body_fn
        self._n_iter = int(n_iter)

        # Single input port: the initial carry value.
        self.declare_input_port(name="carry_init")
        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _compute_output(self, time, state, *inputs, **params):
        initial_carry = inputs[0]
        return jax.lax.fori_loop(
            0, self._n_iter, self._body_fn, initial_carry
        )

`IntegerTime`

Class for managing conversion between decimal and integer time.

Source code in jaxonomy/framework/event.py

class IntegerTime:
    """Class for managing conversion between decimal and integer time."""

    # TODO: Can we use this directly as an int?  Would need to implement __add__,
    # __sub__, etc.  Also, comparisons and floor divide.  Would make the code in
    # Simulator cleaner, but dealing with JAX tracers in `where` and the like
    # might be difficult.  See commit 043c8f757 for a previous attempt.

    #
    # Class variables
    #
    time_scale = DEFAULT_TIME_SCALE  # int -> float conversion factor
    inv_time_scale = 1 / time_scale  # float -> int conversion factor

    # Type of the integer time representation. Defaults to x64 unless explicitly disabled.
    dtype: DTypeLike = npa.intx

    # Largest time value representable by IntegerTime.dtype
    max_int_time = npa.iinfo(dtype).max

    # Floating point representation of max_int_time. Built with concrete
    # numpy (not the backend ``npa``) so it stays a host-side constant: it is
    # read via ``float(...)`` in the representability check and must never
    # become a JAX tracer when ``set_scale`` runs inside an autodiff-through-
    # ``simulate`` trace (T-B6-followup-int-time-scale-trace-safety).
    max_float_time = np.asarray(max_int_time * time_scale, dtype=dtype)

    #
    # Class methods
    #
    @classmethod
    def set_scale(cls, time_scale: float):
        cls.time_scale = time_scale
        cls.inv_time_scale = 1 / time_scale
        # Concrete numpy — keep host-side / tracer-free (see class attr note).
        cls.max_float_time = np.asarray(cls.max_int_time * time_scale, dtype=cls.dtype)

    @classmethod
    def set_default_scale(cls):
        cls.set_scale(DEFAULT_TIME_SCALE)

    @classmethod
    def from_decimal(cls, time: float) -> int:
        """Convert a floating-point time to an integer time."""
        # First limit to the max value to avoid overflow with inf or very large values.
        time = npa.minimum(time, cls.max_float_time)
        return npa.asarray(time * cls.inv_time_scale, dtype=cls.dtype)

    @classmethod
    def as_decimal(cls, time: int) -> float:
        """Convert an integer time to a floating-point time."""
        return time * cls.time_scale

`as_decimal(time)` `classmethod`

Convert an integer time to a floating-point time.

Source code in jaxonomy/framework/event.py

@classmethod
def as_decimal(cls, time: int) -> float:
    """Convert an integer time to a floating-point time."""
    return time * cls.time_scale

`from_decimal(time)` `classmethod`

Convert a floating-point time to an integer time.

Source code in jaxonomy/framework/event.py

@classmethod
def from_decimal(cls, time: float) -> int:
    """Convert a floating-point time to an integer time."""
    # First limit to the max value to avoid overflow with inf or very large values.
    time = npa.minimum(time, cls.max_float_time)
    return npa.asarray(time * cls.inv_time_scale, dtype=cls.dtype)

`JaxonomyError`

Bases: Exception

Base class for all custom jaxonomy errors.

Source code in jaxonomy/framework/error.py

class JaxonomyError(Exception):
    """Base class for all custom jaxonomy errors."""

    # Ideally we'd always have a system to pass but there are at least 2 cases
    # where we may not have one:
    # 1. parsing from json, the block hasn't been built yet
    # 2. other errors not specific to a block
    # In case 1, we should pass all name, path & ui_id info to the error

    def __init__(
        self,
        message=None,
        *,
        system: "SystemBase" = None,  # noqa
        system_id: Hashable = None,
        name_path: list[str] = None,
        ui_id_path: list[str] = None,
        port_index: int = None,
        port_name: str = None,
        port_direction: str = None,  # 'in' or 'out'
        parameter_name: str = None,
        loop: list["DirectedPortLocator"] = None,
    ):
        """Create a new JaxonomyError.

        Only `message` is a positional argument, all others are keyword arguments.

        Args:
            message: A custom error message, defaults to the error class name.
            system: The system that the error occurred in, if available.
            system_id: The id of the system that the error occurred in, use if system can't be passed.
            name_path: The name path of the block that the error occurred in, use if system can't be passed.
            ui_id_path: The ui_id (uuid) path of the block that the error occurred in, use if system can't be passed.
            port_index: The index of the port that the error occurred at.
            port_name: The name of the port that the error occurred at.
            port_direction: The direction of the port that the error occurred at.
            parameter_name: The name of the parameter that the error occurred at.
            loop: A list of I/O ports where the error occurred (eg. AlgebraicLoopError).
        """
        super().__init__(message)

        if system and system_id:
            warnings.warn(
                "Should not specify both system and system_id when raising exceptions"
            )

        if system:
            self.system_id = system.system_id
            self.name_path = name_path or system.name_path
            self.ui_id_path = ui_id_path or system.ui_id_path
        else:
            self.system_id = system_id
            self.name_path = name_path
            self.ui_id_path = ui_id_path

        self.message = message
        self.port_index = port_index
        self.port_name = port_name
        self.port_direction = port_direction
        self.parameter_name = parameter_name

        # Extract serializable info from loop
        # NOTE: we could compact it a bit if the JSON becomes too large...
        self.loop: list[LoopItem] = None
        if loop is not None:
            self.loop = [
                LoopItem(
                    name_path=loc[0].name_path,
                    ui_id_path=loc[0].ui_id_path,
                    port_direction=loc[1],
                    port_index=loc[2],
                )
                for loc in loop
            ]

    def __str__(self):
        message = self.message or self.default_message
        return f"{message}{self._context_info()}"

    def _context_info(self) -> str:
        strbuf = []

        if self.name_path:
            # FIXME: this is known to be too verbose when looking at errors from
            # the UI but makes it better when running pytest or from code.
            # For now, be verbose.
            name_path = ".".join(self.name_path)
            strbuf.append(f" in block {name_path}")
        elif self.system_id:  # Unnamed blocks, likely from code
            strbuf.append(f" in system {self.system_id}")

        if self.port_direction:
            strbuf.append(
                f" at {self.port_direction}put port {self.port_name or self.port_index}"
            )
        elif self.port_name:
            strbuf.append(f" at port {self.port_name}")
        elif self.port_index is not None:
            strbuf.append(f" at port {self.port_index}")
        if self.parameter_name:
            strbuf.append(f" with parameter {self.parameter_name}")
        if self.__cause__ is not None:
            strbuf.append(f": {self.__cause__}")

        return "".join(strbuf)

    @property
    def block_name(self):
        if self.name_path is None:
            return None
        if len(self.name_path) == 0:
            return "root"
        return self.name_path[-1]

    @property
    def default_message(self):
        return type(self).__name__

    def caused_by(self, exc_type: type):
        """Check if this error is or was caused by another error type.

        For instance, if a JaxonomyError is raised because of a TypeError,
        this method will return True when called with TypeError as exc_type.

        Args:
            exc_type: The type of exception to check for (eg. TypeError)

        Returns:
            bool: True if the error is or was caused by the given exception type.
        """

        def _is_or_caused_by(exc, cause_type) -> bool:
            if not exc or not cause_type:
                return False
            if isinstance(exc, cause_type):
                return True
            if not hasattr(self, "__cause__"):
                return False
            return _is_or_caused_by(exc.__cause__, cause_type)

        return _is_or_caused_by(self, exc_type)

`init(message=None, *, system=None, system_id=None, name_path=None, ui_id_path=None, port_index=None, port_name=None, port_direction=None, parameter_name=None, loop=None)`

Create a new JaxonomyError.

Only message is a positional argument, all others are keyword arguments.

Parameters:

Name	Type	Description	Default
`message`		A custom error message, defaults to the error class name.	`None`
`system`	`SystemBase`	The system that the error occurred in, if available.	`None`
`system_id`	`Hashable`	The id of the system that the error occurred in, use if system can't be passed.	`None`
`name_path`	`list[str]`	The name path of the block that the error occurred in, use if system can't be passed.	`None`
`ui_id_path`	`list[str]`	The ui_id (uuid) path of the block that the error occurred in, use if system can't be passed.	`None`
`port_index`	`int`	The index of the port that the error occurred at.	`None`
`port_name`	`str`	The name of the port that the error occurred at.	`None`
`port_direction`	`str`	The direction of the port that the error occurred at.	`None`
`parameter_name`	`str`	The name of the parameter that the error occurred at.	`None`
`loop`	`list[DirectedPortLocator]`	A list of I/O ports where the error occurred (eg. AlgebraicLoopError).	`None`

Source code in jaxonomy/framework/error.py

def __init__(
    self,
    message=None,
    *,
    system: "SystemBase" = None,  # noqa
    system_id: Hashable = None,
    name_path: list[str] = None,
    ui_id_path: list[str] = None,
    port_index: int = None,
    port_name: str = None,
    port_direction: str = None,  # 'in' or 'out'
    parameter_name: str = None,
    loop: list["DirectedPortLocator"] = None,
):
    """Create a new JaxonomyError.

    Only `message` is a positional argument, all others are keyword arguments.

    Args:
        message: A custom error message, defaults to the error class name.
        system: The system that the error occurred in, if available.
        system_id: The id of the system that the error occurred in, use if system can't be passed.
        name_path: The name path of the block that the error occurred in, use if system can't be passed.
        ui_id_path: The ui_id (uuid) path of the block that the error occurred in, use if system can't be passed.
        port_index: The index of the port that the error occurred at.
        port_name: The name of the port that the error occurred at.
        port_direction: The direction of the port that the error occurred at.
        parameter_name: The name of the parameter that the error occurred at.
        loop: A list of I/O ports where the error occurred (eg. AlgebraicLoopError).
    """
    super().__init__(message)

    if system and system_id:
        warnings.warn(
            "Should not specify both system and system_id when raising exceptions"
        )

    if system:
        self.system_id = system.system_id
        self.name_path = name_path or system.name_path
        self.ui_id_path = ui_id_path or system.ui_id_path
    else:
        self.system_id = system_id
        self.name_path = name_path
        self.ui_id_path = ui_id_path

    self.message = message
    self.port_index = port_index
    self.port_name = port_name
    self.port_direction = port_direction
    self.parameter_name = parameter_name

    # Extract serializable info from loop
    # NOTE: we could compact it a bit if the JSON becomes too large...
    self.loop: list[LoopItem] = None
    if loop is not None:
        self.loop = [
            LoopItem(
                name_path=loc[0].name_path,
                ui_id_path=loc[0].ui_id_path,
                port_direction=loc[1],
                port_index=loc[2],
            )
            for loc in loop
        ]

`caused_by(exc_type)`

Check if this error is or was caused by another error type.

For instance, if a JaxonomyError is raised because of a TypeError, this method will return True when called with TypeError as exc_type.

Parameters:

Name	Type	Description	Default
`exc_type`	`type`	The type of exception to check for (eg. TypeError)	required

Returns:

Name	Type	Description
`bool`		True if the error is or was caused by the given exception type.

Source code in jaxonomy/framework/error.py

def caused_by(self, exc_type: type):
    """Check if this error is or was caused by another error type.

    For instance, if a JaxonomyError is raised because of a TypeError,
    this method will return True when called with TypeError as exc_type.

    Args:
        exc_type: The type of exception to check for (eg. TypeError)

    Returns:
        bool: True if the error is or was caused by the given exception type.
    """

    def _is_or_caused_by(exc, cause_type) -> bool:
        if not exc or not cause_type:
            return False
        if isinstance(exc, cause_type):
            return True
        if not hasattr(self, "__cause__"):
            return False
        return _is_or_caused_by(exc.__cause__, cause_type)

    return _is_or_caused_by(self, exc_type)

`LeafContext` `dataclass`

Bases: ContextBase

Source code in jaxonomy/framework/context.py

@dataclasses.dataclass(frozen=True)
class LeafContext(ContextBase):
    state: Optional[LeafState] = None

    @property
    def system_id(self) -> Hashable:
        return self.owning_system.system_id

    def __getitem__(self, key: Hashable) -> LeafContext:
        """Dummy indexing for compatibility with DiagramContexts, returning self."""
        _reject_port_key(key)
        assert key == self.system_id, f"Attempting to get subcontext {key} from {self}"
        return self

    def with_subcontext(self, key: Hashable, ctx: LeafContext) -> LeafContext:
        """Dummy replacement for compatibility with DiagramContexts, returning ctx."""
        assert (
            key == self.system_id
        ), f"System ID {key} does not match leaf ID {self.system_id}"
        assert (
            key == ctx.system_id
        ), f"System ID {key} does not match leaf ID {ctx.system_id}"
        return ctx

    def __repr__(self) -> str:
        return f"{type(self).__name__}(sys={self.system_id})"

    def with_state(self, state: LeafState) -> LeafContext:
        return dataclasses.replace(self, state=state)

    @property
    def continuous_state(self) -> LeafStateComponent:
        return self.state.continuous_state

    def with_continuous_state(self, value: LeafStateComponent) -> LeafContext:
        return dataclasses.replace(self, state=self.state.with_continuous_state(value))

    @property
    def num_continuous_states(self) -> int:
        return self.state.num_continuous_states

    @property
    def has_continuous_state(self) -> bool:
        return self.state.has_continuous_state

    @property
    def discrete_state(self) -> LeafStateComponent:
        return self.state.discrete_state

    def with_discrete_state(self, value: LeafStateComponent) -> LeafContext:
        return dataclasses.replace(self, state=self.state.with_discrete_state(value))

    @property
    def num_discrete_states(self) -> int:
        return self.state.num_discrete_states

    @property
    def has_discrete_state(self) -> bool:
        return self.state.has_discrete_state

    @property
    def mode(self) -> int:
        return self.state.mode

    @property
    def has_mode(self) -> bool:
        return self.state.has_mode

    def with_mode(self, value: int) -> LeafContext:
        return dataclasses.replace(self, state=self.state.with_mode(value))

    @property
    def cache(self) -> tuple[Array]:
        return self.state.cache

    @property
    def num_cached_values(self) -> int:
        return self.state.num_cached_values

    @property
    def has_cache(self) -> bool:
        return self.state.has_cache

    def with_cached_value(self, index: int, value: Array) -> LeafContext:
        return dataclasses.replace(
            self, state=self.state.with_cached_value(index, value)
        )

    def with_updated_parameters(self) -> ContextBase:
        params = self.owning_system.dynamic_parameters
        new_parameters = {}
        for name, param in params.items():
            new_parameters[name] = param.get()
        return dataclasses.replace(self, parameters=new_parameters)

    def with_new_state(self) -> ContextBase:
        return dataclasses.replace(self, state=self.owning_system.create_state())

    def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
        """Create a copy of this context, replacing only the specified parameters."""
        parameters = {**self.parameters}
        for name, value in new_parameters.items():
            param = self.owning_system.dynamic_parameters[name]
            value = _coerce_param_for_jit_cache(parameters.get(name), value)
            param.set(value)
            parameters[name] = param.get()
        return dataclasses.replace(self, parameters=parameters)

`getitem(key)`

Dummy indexing for compatibility with DiagramContexts, returning self.

Source code in jaxonomy/framework/context.py

def __getitem__(self, key: Hashable) -> LeafContext:
    """Dummy indexing for compatibility with DiagramContexts, returning self."""
    _reject_port_key(key)
    assert key == self.system_id, f"Attempting to get subcontext {key} from {self}"
    return self

`with_parameters(new_parameters)`

Create a copy of this context, replacing only the specified parameters.

Source code in jaxonomy/framework/context.py

def with_parameters(self, new_parameters: Mapping[str, ArrayLike]) -> ContextBase:
    """Create a copy of this context, replacing only the specified parameters."""
    parameters = {**self.parameters}
    for name, value in new_parameters.items():
        param = self.owning_system.dynamic_parameters[name]
        value = _coerce_param_for_jit_cache(parameters.get(name), value)
        param.set(value)
        parameters[name] = param.get()
    return dataclasses.replace(self, parameters=parameters)

`with_subcontext(key, ctx)`

Dummy replacement for compatibility with DiagramContexts, returning ctx.

Source code in jaxonomy/framework/context.py

def with_subcontext(self, key: Hashable, ctx: LeafContext) -> LeafContext:
    """Dummy replacement for compatibility with DiagramContexts, returning ctx."""
    assert (
        key == self.system_id
    ), f"System ID {key} does not match leaf ID {self.system_id}"
    assert (
        key == ctx.system_id
    ), f"System ID {key} does not match leaf ID {ctx.system_id}"
    return ctx

`LeafState` `dataclass`

Container for state information for a leaf system.

Attributes:

Name	Type	Description
`name`	`str`	Name of the leaf system that owns this state.
`continuous_state`	`LeafStateComponent`	Continuous state of the system, i.e. the component of state that evolves in continuous time. If the system has no continuous state, this will be None.
`discrete_state`	`LeafStateComponent`	Discrete state of the system, i.e. one or more components of state that do not change continuously with ime (not necessarily discrete-valued). If the system has no discrete state, this will be None.
`mode`	`int`	An integer value indicating the current "mode", "stage", or discrete-valued state component of the system. Used for finite state machines or multi-stage hybrid systems. If the system has no mode, this will be None.
`cache`	`tuple[LeafStateComponent]`	The current values of sample-and-hold outputs from the system. In a pure discrete system these would not be state components (just results of feedthrough computations), but in a hybrid or multirate system they act as discrete state from the perspective of continuous or asynchronous discrete components of the system. Hence, they are stored in the state, but are maintained separately from the normal internal state of the system.

Notes

(1) This class is immutable. To modify a LeafState, use the with_* methods.

(2) The type annotations for state components are LeafStateComponent, which is a union of array, tuple, and named tuple. The most common case is arrays, but this allows for more flexibility in defining state components, e.g. a second-order system can define a named tuple of generalized coordinates and velocities rather than concatenating into a single array.

Source code in jaxonomy/framework/state.py

@dataclasses.dataclass(frozen=True)
class LeafState:
    """Container for state information for a leaf system.

    Attributes:
        name (str):
            Name of the leaf system that owns this state.
        continuous_state (LeafStateComponent):
            Continuous state of the system, i.e. the component of state that evolves in
            continuous time. If the system has no continuous state, this will be None.
        discrete_state (LeafStateComponent):
            Discrete state of the system, i.e. one or more components of state that do
            not change continuously with ime (not necessarily discrete-_valued_). If
            the system has no discrete state, this will be None.
        mode (int):
            An integer value indicating the current "mode", "stage", or discrete-valued
            state component of the system.  Used for finite state machines or
            multi-stage hybrid systems.  If the system has no mode, this will be None.
        cache (tuple[LeafStateComponent]):
            The current values of sample-and-hold outputs from the system.  In a pure
            discrete system these would not be state components (just results of
            feedthrough computations), but in a hybrid or multirate system they act as
            discrete state from the perspective of continuous or asynchronous discrete
            components of the system.  Hence, they are stored in the state, but are
            maintained separately from the normal internal state of the system.

    Notes:
        (1) This class is immutable.  To modify a LeafState, use the `with_*` methods.

        (2) The type annotations for state components are LeafStateComponent, which is
        a union of array, tuple, and named tuple. The most common case is arrays, but
        this allows for more flexibility in defining state components, e.g. a
        second-order system can define a named tuple of generalized coordinates and
        velocities rather than concatenating into a single array.
    """

    name: Optional[str] = None
    continuous_state: Optional[LeafStateComponent] = None
    discrete_state: Optional[LeafStateComponent] = None
    mode: Optional[int] = None
    cache: Optional[tuple[Array]] = None

    def __repr__(self) -> str:
        states = []
        if self.continuous_state is not None:
            states.append(f"xc={self.continuous_state}")
        if self.discrete_state is not None:
            states.append(f"xd={self.discrete_state}")
        if self.mode is not None:
            states.append(f"s={self.mode}")
        return f"{type(self).__name__}({', '.join(states)})"

    def with_continuous_state(self, value: LeafStateComponent) -> LeafState:
        """Create a copy of this LeafState with the continuous state replaced."""
        if value is not None and self.continuous_state is not None:
            value = tree_util.tree_map(self._reshape_like, value, self.continuous_state)

        return dataclasses.replace(self, continuous_state=value)

    def _component_size(self, component: LeafStateComponent) -> int:
        if component is None:
            return 0
        if isinstance(component, tuple):
            # return sum(x.size for x in component)
            return len(component)
        return component.size

    def _reshape_like(self, new_value: Array, current_value: Array) -> Array:
        """Helper function for tree-mapped type conversions.

        Ensures that the new components are array-like and have the same shape as
        the existing state to preserve PyTree structure.
        """
        return reshape(new_value, current_value.shape)

    @property
    def num_continuous_states(self) -> int:
        return self._component_size(self.continuous_state)

    @property
    def has_continuous_state(self) -> bool:
        return self.num_continuous_states > 0

    def with_discrete_state(self, value: LeafStateComponent) -> LeafState:
        """Create a copy of this LeafState with the discrete state replaced."""
        if value is not None and self.discrete_state is not None:
            value = tree_util.tree_map(self._reshape_like, value, self.discrete_state)

        return dataclasses.replace(self, discrete_state=value)

    @property
    def num_discrete_states(self) -> int:
        return self._component_size(self.discrete_state)

    @property
    def has_discrete_state(self) -> bool:
        return self.num_discrete_states > 0

    def with_mode(self, value: int) -> LeafState:
        """Create a copy of this LeafState with the mode replaced."""
        return dataclasses.replace(self, mode=value)

    @property
    def has_mode(self) -> bool:
        return self.mode is not None

    def with_cached_value(self, index: int, value: Array) -> LeafState:
        """Create a copy of this LeafState with the specified cache value replaced."""
        cache = list(self.cache)
        cache[index] = value
        return dataclasses.replace(self, cache=tuple(cache))

    def has_cache(self) -> bool:
        return self.cache is not None

    def num_cached_values(self) -> int:
        return len(self.cache)

`with_cached_value(index, value)`

Create a copy of this LeafState with the specified cache value replaced.

Source code in jaxonomy/framework/state.py

def with_cached_value(self, index: int, value: Array) -> LeafState:
    """Create a copy of this LeafState with the specified cache value replaced."""
    cache = list(self.cache)
    cache[index] = value
    return dataclasses.replace(self, cache=tuple(cache))

`with_continuous_state(value)`

Create a copy of this LeafState with the continuous state replaced.

Source code in jaxonomy/framework/state.py

def with_continuous_state(self, value: LeafStateComponent) -> LeafState:
    """Create a copy of this LeafState with the continuous state replaced."""
    if value is not None and self.continuous_state is not None:
        value = tree_util.tree_map(self._reshape_like, value, self.continuous_state)

    return dataclasses.replace(self, continuous_state=value)

`with_discrete_state(value)`

Create a copy of this LeafState with the discrete state replaced.

Source code in jaxonomy/framework/state.py

def with_discrete_state(self, value: LeafStateComponent) -> LeafState:
    """Create a copy of this LeafState with the discrete state replaced."""
    if value is not None and self.discrete_state is not None:
        value = tree_util.tree_map(self._reshape_like, value, self.discrete_state)

    return dataclasses.replace(self, discrete_state=value)

`with_mode(value)`

Create a copy of this LeafState with the mode replaced.

Source code in jaxonomy/framework/state.py

def with_mode(self, value: int) -> LeafState:
    """Create a copy of this LeafState with the mode replaced."""
    return dataclasses.replace(self, mode=value)

`LeafSystem` `dataclass`

Bases: SystemBase

Basic building block for dynamical systems.

A LeafSystem is a minimal component of a system model in jaxonomy, containing no subsystems. Inputs, outputs, state, parameters, updates, etc. can be added to the block using the various declare_* methods. The built-in blocks in jaxonomy.library are all subclasses of LeafSystem, as are any custom blocks defined by the user.

Source code in jaxonomy/framework/leaf_system.py

class LeafSystem(SystemBase, metaclass=InitializeParameterResolver):
    """Basic building block for dynamical systems.

    A LeafSystem is a minimal component of a system model in jaxonomy, containing no
    subsystems.  Inputs, outputs, state, parameters, updates, etc. can be added to the
    block using the various `declare_*` methods.  The built-in blocks in
    jaxonomy.library are all subclasses of LeafSystem, as are any custom blocks defined
    by the user."""

    # SystemBase is a dataclass, so we need to call __post_init__ explicitly
    def __post_init__(self):
        super().__post_init__()
        logger.debug(f"Initializing {self.name} [{self.system_id}]")

        # If not None, this defines the shape and data type of the continuous state
        # component.  This value will be used to initialize the context, so it will
        # also serve as the initial value unless explicitly overridden. It will
        # typically be an array, but it can be any PyTree-structured object (list,
        # dict, namedtuple, etc.), provided that the ODE function returns a PyTree
        # of the same structure.
        self._default_continuous_state: LeafStateComponent = None
        self._mass_matrix: Array = None
        self._continuous_state_output_port_idx: int = None

        # The SystemCallback associated with time derivatives of the continuous state.
        # This is initialized in the `declare_continuous_state` method.
        self.ode_callback: SystemCallback = None

        # If not empty, this defines the shape and data type of the discrete state.
        # This value will be used to initialize the context, so it will also serve
        # as the initial value unless explicitly overridden. This will often be an
        # array, but as for the continuous state it can be any PyTree-structured
        # object (list, dict, namedtuple, etc.), provided that the update functions
        # return a PyTree of the same structure.
        self._default_discrete_state: LeafStateComponent = None

        # If not None, the system has a "mode" or "stage" component of the state.
        # In a "state machine" paradigm, this represents the current state of the
        # system (although "state" is obviously used for other things in this case).
        # The mode is an integer value, and the system can declare transitions between
        # modes using the `declare_zero_crossing` method, which in addition to the
        # guard function and reset map also takes optional `start_mode` and `end_mode`
        # arguments.
        self._default_mode: int = None
        self._mode_output_port_idx: int = None

        # Set of "template" values for the sample-and-hold output ports, if known.
        # If not known, these will be `None`, in which case an appropriate value is
        # inferred from upstream during static analysis.
        self._default_cache: List[LeafStateComponent] = []

        # Transition map from (start_mode -> [*end_modes]) indicating which
        # transition events are active in each mode.  This is not used by
        # any logic in the system, but can be useful for debugging.
        self.transition_map: dict[int, List[Tuple[int, ZeroCrossingEvent]]] = {}

        # Set of events that updates at a fixed rate.  Each event has its own period
        # and offset, so "fires" independently of the other events. These can be
        # created using the `declare_periodic_update` method.
        self._state_update_events: List[DiscreteUpdateEvent] = []

        # Set of events that update when a zero-crossing occurs.  Each event has its
        # own guard function and, optionally, reset map, start mode, and end mode.
        # These can be created using the `declare_zero_crossing` method.
        self._zero_crossing_events: List[ZeroCrossingEvent] = []

        # T-115-followup-saturate-rate-classification: count only ZC events
        # that have *behavioral* side effects (a user-supplied ``reset_map``
        # or a mode transition via ``start_mode`` / ``end_mode``). Pure
        # solver-hint ZC events (e.g. :class:`Saturate` declaring clip
        # boundaries so the ODE integrator can localise the discontinuity)
        # should not flip the block's rate-group classification to
        # ``event_driven`` — they exist only to help the solver, not to
        # represent an asynchronous trigger. The rate-groups inference uses
        # this counter to distinguish the two cases.
        self._n_behavioral_zc_events: int = 0

        # T-027: per-event Zeno-hold tracking for `declare_zero_crossing(zeno_tolerance=...)`.
        # Each entry is a dict {slot, tol, name} where `slot` is the index into the
        # private Zeno discrete state. The discrete state is a NamedTuple with fields
        # `zeno` (bool array, one slot per protected event) and `tprev` (float array).
        # See `_install_zeno_protection` for details. Empty by default; the existing
        # `Integrator` Zeno path uses its own private discrete state and is unaffected.
        #
        # T-027a: when the user also calls `declare_discrete_state(...)`, the framework
        # packs the Zeno tracker alongside the user's value as `_DiscreteWithZeno(
        # user=..., zeno=_ZenoState(zeno=..., tprev=...))`. User callbacks (ode, guard,
        # reset) are wrapped to see only their own `user` slot via `state.discrete_state`,
        # and the framework re-packs on the way out. When the user has no discrete
        # state, the bare `_ZenoState` is stored as before (no wrapper, no overhead).
        self._zeno_protected_events: List[dict] = []
        self._zeno_state_type = None  # Set on first protected event.
        self._zeno_combined_type = None  # Set when user has discrete state too.
        self._zeno_user_default = None  # User's declared default (if any) at install time.
        self._zeno_ode_wrapped = False

    def initialize(self, **parameters):
        """Hook for initializing a system. Called during context creation.

        If the parameters are instances of Parameter, they will be resolved.
        If implemented, the function signature should contain all the declared
        parameters.

        This function should not be called directly. It will be called implicitly
        after __init__ with the resolved parameters.
        """
        pass

    @property
    def has_feedthrough_side_effects(self) -> bool:
        # See explanation in `SystemBase.has_feedthrough_side_effects`.  This will
        # almost always be False, but can be overridden in special cases where a
        # feedthrough output is computed via use of `io_callback`.
        return False

    @property
    def has_ode_side_effects(self) -> bool:
        # This will almost always be False for a LeafSystem - Diagram systems
        # have some special logic to do this determination.
        return False

    @property
    def has_continuous_state(self) -> bool:
        return self._default_continuous_state is not None

    @property
    def continuous_state_default(self) -> "LeafStateComponent":
        """The declared default continuous-state value (read-only).

        This is the ``default_value`` passed to ``declare_continuous_state``
        (or the array inferred from ``shape`` / ``dtype``), i.e. the value
        that seeds ``context.continuous_state`` before any user override.
        Returns ``None`` when the block has no continuous state. Exposed as a
        documented accessor so callers don't have to reach into the private
        ``_default_continuous_state`` attribute (T-C2-followup).
        """
        return self._default_continuous_state

    @property
    def has_discrete_state(self) -> bool:
        return self._default_discrete_state is not None

    @property
    def has_zero_crossing_events(self) -> bool:
        return len(self._zero_crossing_events) > 0

    #
    # Event handling
    #
    def wrap_callback(
        self, callback: Callable, collect_inputs: bool | list[int] = True
    ) -> Callable:
        """Wrap an update function to unpack local variables and block inputs.

        The callback should have the signature
        `callback(time, state, *inputs, **params) -> result`
        and will be wrapped to have the signature `callback(context) -> result`,
        as expected by the event handling logic.

        This is used internally for declaration methods like
        `declare_periodic_update` so that users can write more intuitive
        block-level update functions without worrying about the "context", and have
        them automatically wrapped to have the right interface.  It can also be
        called directly by users to wrap their own update functions, for example to
        create a callback function for `declare_output_port`.

        The context and state are strictly immutable, so the callback should not
        attempt to change any values in the context or state.  Even in cases where
        it is impossible to _enforce_ this (e.g. a state component is a list, which
        is always mutable in Python), the callback should be careful to avoid direct
        modification of the context or state, which may lead to unexpected behavior
        or JAX tracer errors.

        Args:
            callback (Callable):
                The (pure) function to be wrapped. See above for expected signature.
            collect_inputs (bool):
                If True, the callback will eval input ports to gather input values.
                Normally this should be True, but it can be set to False if the
                return value depends only on the state but not inputs, for
                instance. This helps reduce the number of expressions that need to
                be JIT compiled. Can also be specified as a list of integer port indices.
                Default is True (collect all inputs).

        Returns:
            Callable:
                The wrapped function, with signature `callback(context) -> result`.
        """
        return partial(
            _wrap_leaf_user_callback,
            owner=self,
            user_callback=callback,
            collect_inputs=collect_inputs,
        )

    def _passthrough(self, context: ContextBase) -> LeafState:
        """Dummy callback for inactive events."""
        return context[self.system_id].state

    @property
    def state_update_events(self) -> FlatEventCollection:
        return FlatEventCollection(tuple(self._state_update_events))

    @property
    def zero_crossing_events(self) -> LeafEventCollection:
        # The default is for all to be active. Use the `determine_active_guards`
        # method to determine which are active conditioned on the current "mode"
        # or "stage" of the system.
        return LeafEventCollection(tuple(self._zero_crossing_events)).mark_all_active()

    def with_parameter(self, name: str, value) -> LeafSystem:
        """Return a copy of this system with one dynamic parameter replaced.

        The returned system is a new instance. The original is unchanged.

        Args:
            name: Parameter name (must exist as a dynamic parameter).
            value: New value (typically a JAX array for ``jax.grad`` / ``jax.vmap``).

        Raises:
            KeyError: If ``name`` is not a dynamic parameter.
            TypeError: If ``name`` is a static parameter.
        """
        if name in self._static_parameters:
            raise TypeError(
                f"Parameter {name!r} is static on {self.name!r}; static parameters "
                "cannot be replaced at runtime without recompilation."
            )
        if name not in self._dynamic_parameters:
            available = sorted(
                {*self._static_parameters.keys(), *self._dynamic_parameters.keys()}
            )
            raise KeyError(
                f"Parameter {name!r} is not a dynamic parameter on {self.name!r}. "
                f"Available: {available}"
            )

        old_param = self._dynamic_parameters[name]
        old_val = Parameter.unwrap(old_param)
        try:
            value = _check_values_compatible(old_val, value)
        except ValueError as e:
            raise ValueError(f"{e} (parameter {name!r} on {self.name!r})") from None

        new = copy.deepcopy(self)
        new.parent = None
        new._dependency_graph = None
        new.feedthrough_pairs = None
        new._cache_update_events = None
        new._cached_input_ports.clear()
        new._cached_output_ports.clear()

        if isinstance(old_param, Parameter):
            new._dynamic_parameters[name] = dataclasses.replace(
                old_param,
                value=value,
                name=name,
                system=new,
            )
        else:
            new._dynamic_parameters[name] = Parameter(
                value=value,
                name=name,
                system=new,
            )

        return new

    # Inherits docstring from SystemBase
    def eval_zero_crossing_updates(
        self,
        context: ContextBase,
        events: LeafEventCollection,
    ) -> LeafState:
        local_events = events[self.system_id]
        state = context[self.system_id].state

        logger.debug(f"Eval update events for {self.name}")
        logger.debug(f"local events: {local_events}")

        for event in local_events:
            # This is evaluated conditionally on event_data.active
            state = event.handle(context)

            # Store the updated state in the context for this block
            leaf_context = context[self.system_id].with_state(state)

            # Update the context for this block in the overall context
            context = context.with_subcontext(self.system_id, leaf_context)

        # Now `context` contains the updated "plus" state for this block, but
        # this needs to be discarded so that other block updates can also be
        # processed using the "minus" state. This is done by simply returning the
        # "plus" state and discarding the rest of the updated context.
        return state

    # Inherits docstring from SystemBase
    def determine_active_guards(self, root_context: ContextBase) -> LeafEventCollection:
        mode = root_context[self.system_id].mode  # Current system mode

        def _conditionally_activate(
            event: ZeroCrossingEvent,
        ) -> ZeroCrossingEvent:
            # Check to see if the event corresponds to a mode transition
            # If not, just return the event unchanged (will be active)
            if event.active_mode is None:
                return event
            # If the event does correspond to a mode transition, check to see
            # if the event is active in the current mode
            return cond(
                mode == event.active_mode,
                lambda e: e.mark_active(),
                lambda e: e.mark_inactive(),
                event,
            )

        # Apply the conditional activation to all events
        zero_crossing_events = LeafEventCollection(
            tuple(_conditionally_activate(e) for e in self.zero_crossing_events)
        )

        logger.debug(f"Zero-crossing events for {self.name}: {zero_crossing_events}")
        return zero_crossing_events

    @property
    def _flat_callbacks(self) -> List[OutputPort]:
        """Return all of the sample-and-hold output ports in this system."""
        return self.callbacks

    def declare_cache(
        self,
        callback: Callable,
        period: float | Parameter = None,
        offset: float | Parameter = 0.0,
        name: str = None,
        prerequisites_of_calc: List[DependencyTicket] = None,
        default_value: Array = None,
        requires_inputs: bool = True,
    ) -> int:
        """Declare a stored computation for the system.

        This method accepts a callback function with the block-level signature
            `callback(time, state, *inputs, **parameters) -> value`
        and wraps it to have the signature
            `callback(context) -> value`

        This callback can optionally be used to define a periodic update event that
        refreshes the cached value.  Other calculations (e.g. sample-and-hold output
        ports) can then depend on the cached value.

        Args:
            callback (Callable):
                The callback function defining the cached computation.
            period (float, optional):
                If not None, the callback function will be used to define a periodic
                update event that refreshes the value. Defaults to None.
            offset (float, optional):
                The offset of the periodic update event. Defaults to 0.0.  Will be ignored
                unless `period` is not None.
            name (str, optional):
                The name of the cached value. Defaults to None.
            default_value (Array, optional):
                The default value of the result, if known. Defaults to None.
            requires_inputs (bool, optional):
                If True, the callback will eval input ports to gather input values.
                This will add a bit to compile time, so setting to False where possible
                is recommended. Defaults to True.
            prerequisites_of_calc (List[DependencyTicket], optional):
                The dependency tickets for the computation. Defaults to None, in which
                case the default is to assume dependency on either (inputs) if
                `requires_inputs` is True, or (nothing) otherwise.

        Returns:
            int: The index of the callback in `system.callbacks`.  The cache index can
                recovered from `system.callbacks[callback_index].cache_index`.
        """
        # The index in the list of system callbacks
        callback_index = len(self.callbacks)

        # This is the index that this cached value will have in state.cache
        cache_index = len(self._default_cache)
        self._default_cache.append(default_value)

        # To help avoid unnecessary flagging of algebraic loops, trim the inputs as a
        # default prereq if the update callback doesn't use them
        if prerequisites_of_calc is None:
            if requires_inputs:
                prerequisites_of_calc = [DependencyTicket.u]
            else:
                prerequisites_of_calc = [DependencyTicket.nothing]

        def _update_callback(
            time: Scalar, state: LeafState, *inputs, **parameters
        ) -> LeafState:
            output = callback(time, state, *inputs, **parameters)
            return state.with_cached_value(cache_index, output)

        _update_callback = self.wrap_callback(
            _update_callback, collect_inputs=requires_inputs
        )

        if period is None:
            event = None

        else:
            # The cache has a periodic event updating its value defined by the callback
            event = DiscreteUpdateEvent(
                system_id=self.system_id,
                event_data=PeriodicEventData(
                    period=period, offset=offset, active=False
                ),
                name=f"{self.name}:cache_update_{cache_index}_",
                callback=_update_callback,
                passthrough=self._passthrough,
            )

        if name is None:
            name = f"cache_{cache_index}"

        sys_callback = SystemCallback(
            callback=_update_callback,
            system=self,
            callback_index=callback_index,
            name=name,
            prerequisites_of_calc=prerequisites_of_calc,
            event=event,
            default_value=default_value,
            cache_index=cache_index,
        )
        self.callbacks.append(sys_callback)

        return callback_index

    # NOTE: we can only declare one continuous state per system because each
    # call will overwrite self._default_continuous_state
    def declare_continuous_state(
        self,
        shape: ShapeLike = None,
        default_value: Array = None,
        dtype: DTypeLike = None,
        ode: Callable = None,
        mass_matrix: Array = None,
        as_array: bool = True,
        requires_inputs: bool = True,
        prerequisites_of_calc: List[DependencyTicket] = None,
        substeps: int = 1,
        project: Callable = None,
    ):
        """Declare a continuous state component for the system.

        The continuous state value is read inside callbacks as
        ``state.continuous_state`` (the ``state`` argument of the ``ode`` /
        output callbacks). **Unpack contract** (T-C3-followup): the shape of
        ``state.continuous_state`` mirrors exactly what you passed as
        ``default_value`` (or the zeros array implied by ``shape`` /
        ``dtype``):

        - A scalar default (``jnp.array(0.0)``) gives a scalar
          ``state.continuous_state`` — read it directly, do **not** index.
        - A vector default (``jnp.zeros(3)``) gives a length-3 array — index
          / unpack as ``x, y, z = state.continuous_state`` or
          ``state.continuous_state[i]``.
        - A PyTree default (tuple / NamedTuple / dict) gives back the same
          PyTree structure; your ``ode`` must return ``xcdot`` with the
          identical structure.

        The ``ode`` callback's return value must match the
        ``default_value`` structure element-for-element, since it is added to
        the state during integration. A common error is declaring a scalar
        state but returning ``jnp.array([xdot])`` (shape ``(1,)``) from the
        ode — keep both scalar or both vector.

        Multirate substepping (T-133): ``substeps=N`` declares that this
        block's continuous dynamics have a fast time constant needing ``N``
        inner integration steps per outer solver step (e.g. a motor's
        electrical winding inside a 1 kHz control loop). Honored by the
        fixed-step ``rk4`` solver (``SimulatorOptions(ode_solver_method=
        "rk4")``): the block's states advance with ``N`` RK4 substeps of
        ``h/N`` while the rest of the diagram takes one step of ``h``,
        with first-order (zero-order-hold) coupling at the boundary —
        each side sees the other's start-of-step values, matching the
        semantics of a hand-rolled JIT-safe substep loop. Adaptive solvers
        (``dopri5``/``bdf``) ignore the declaration — they control
        stiffness through global step adaptation. ``N`` must be a static
        Python ``int >= 1``; the default 1 is byte-equivalent to the
        pre-T-133 behavior.

        Reverse-mode autodiff (``enable_autodiff=True``) is supported —
        the substep loop has a static trip count and the checkpointed
        adjoint substeps the costates alongside their primals. Gradient
        accuracy carries the scheme's first-order coupling error: the
        adjoint converges to the true sensitivity linearly in the outer
        step ``h`` (exact FD agreement is only recovered as ``h`` is
        refined), and for dynamics *unstable at the outer step* the
        adjoint's reverse-time primal re-integration further limits
        accuracy. Reduce the outer step when gradients through the
        coupling interface need to be tight.

        Declared state projection (T-132): ``project=fn`` declares that
        this block's continuous state lives on a manifold and supplies
        the retraction back onto it — e.g. unit-quaternion
        renormalization for an attitude state (``nq=4`` integrated
        componentwise drifts off the unit sphere under any one-step
        integrator). ``fn(x) -> x`` receives the state in its declared
        structure, must be shape-preserving and jit-safe, and is applied
        by the simulator **at the end of every major step** (composing
        with, and independent of, the T-003a DAE projection). Within-step
        drift is bounded by the step size; the recorded trajectory and
        all values other blocks see at major-step boundaries are on the
        manifold. Differentiable: the projection participates in
        reverse-mode AD as ordinary traced ops.
        """
        if not isinstance(substeps, (int, np.integer)) or isinstance(
            substeps, bool
        ) or substeps < 1:
            raise ValueError(
                f"declare_continuous_state: substeps must be a static Python "
                f"int >= 1, got {substeps!r}. (It sets a compile-time inner "
                "loop count and cannot be traced or fractional.)"
            )
        self._continuous_substeps = int(substeps)
        if project is not None and not callable(project):
            raise ValueError(
                f"declare_continuous_state: project must be a callable "
                f"x -> x (shape-preserving, jit-safe), got {project!r}."
            )
        self._continuous_projection = project

        self.ode_callback = SystemCallback(
            callback=None,
            system=self,
            callback_index=len(self.callbacks),
            name=f"{self.name}_ode",
            prerequisites_of_calc=prerequisites_of_calc,
        )
        self.callbacks.append(self.ode_callback)
        callback_idx = len(self.callbacks) - 1

        # FIXME: this is to preserve some backward compatibility while we decouple
        # declaration from configuration. Declaration should not have to call
        # configuration.
        if default_value is not None or shape is not None:
            self.configure_continuous_state(
                callback_idx,
                shape=shape,
                default_value=default_value,
                dtype=dtype,
                ode=ode,
                mass_matrix=mass_matrix,
                as_array=as_array,
                requires_inputs=requires_inputs,
                prerequisites_of_calc=prerequisites_of_calc,
            )

        return callback_idx

    def configure_continuous_state(
        self,
        callback_idx: int,
        shape: ShapeLike = None,
        default_value: Array = None,
        dtype: DTypeLike = None,
        ode: Callable = None,
        mass_matrix: Array = None,
        as_array: bool = True,
        requires_inputs: bool = True,
        prerequisites_of_calc: List[DependencyTicket] = None,
    ):
        """Configure a continuous state component for the system.

        The `ode` callback computes the time derivative of the continuous state based on the
        current time, state, and any additional inputs. If `ode` is not provided, a default
        zero vector of the same size as the continuous state is used. If provided, the `ode`
        callback should have the signature `ode(time, state, *inputs, **params) -> xcdot`.

        Args:
            callback_idx (int):
                The index of the callback in the system's callback list.
            shape (ShapeLike, optional):
                The shape of the continuous state vector. Defaults to None.
            default_value (Array, optional):
                The initial value of the continuous state vector. Defaults to None.
            dtype (DTypeLike, optional):
                The data type of the continuous state vector. Defaults to None.
            ode (Callable, optional):
                The callback for computing the time derivative of the continuous state.
                Should have the signature:
                    `ode(time, state, *inputs, **parameters) -> xcdot`.
                Defaults to None.
            mass_matrix (Array, optional):
                The mass matrix for the continuous state. Defaults to None. If
                provided, must be a square matrix with the same shape as the
                continuous state.  Using a mass matrix different from the identity
                in any LeafSystem will require the use of a compatible continuous-time
                solver (currently only BDF is supported).  Currently mass matrices are
                also only supported for scalar- or vector-valued continuous states (
                i.e. no matrices or other PyTree-structured states).
            as_array (bool, optional):
                If True, treat the default_value as an array-like (cast if necessary).
                Otherwise, it will be stored as the default state without modification.
            requires_inputs (bool, optional):
                If True, indicates that the ODE computation requires inputs.
            prerequisites_of_calc (List[DependencyTicket], optional):
                The dependency tickets for the ODE computation. Defaults to None, in
                which case the assumption is a dependency on either (time, continuous
                state) if `requires_inputs` is False, otherwise (time, continuous state,
                inputs.

        Raises:
            AssertionError:
                If neither shape nor default_value is provided, or if the mass matrix
                is inconsistent with the continuous state.

        Notes:
            (1) Only one of `shape` and `default_value` should be provided. If `default_value`
            is provided, it will be used as the initial value of the continuous state. If
            `shape` is provided, the initial value will be a zero vector of the given shape
            and specified dtype.
        """

        if prerequisites_of_calc is None:
            prerequisites_of_calc = [DependencyTicket.time, DependencyTicket.xc]
            if requires_inputs:
                prerequisites_of_calc.append(DependencyTicket.u)

        if as_array:
            default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

        logger.debug(f"In block {self.name} [{self.system_id}]: {default_value=}")

        # Tree-map the default value to ensure that it is an array-like with the
        # correct shape and dtype. This is necessary because the default value
        # may be a list, tuple, or other PyTree-structured object.
        default_value = tree_util.tree_map(npa.asarray, default_value)

        self._default_continuous_state = default_value
        if self._continuous_state_output_port_idx is not None:
            port = self.output_ports[self._continuous_state_output_port_idx]
            port.default_value = default_value
            self._default_cache[port.cache_index] = default_value

        if ode is None:
            # If no ODE is specified, return a zero vector of the same size as the
            # continuous state. This will break if the continuous state is
            # a named tuple, in which case a custom ODE must be provided.
            assert as_array, "Must provide custom ODE for non-array continuous state"

            def ode(time, state, *inputs, **parameters):
                return npa.zeros_like(default_value)

        # Wrap the ode function to accept a context and return the time derivatives.
        ode = self.wrap_callback(ode)

        # Declare the time derivative function as a system callback so that its
        # dependencies can be tracked in the system dependency graph
        self.ode_callback._callback = ode
        self.ode_callback.prerequisites_of_calc = prerequisites_of_calc

        # Override the default `eval_time_derivatives` to use the wrapped ODE function
        self.eval_time_derivatives = self.ode_callback.eval

        # T-027: if Zeno-protected events were registered before this, wrap the
        # ode so a Zeno-hold freezes the continuous state.
        if self._zeno_protected_events:
            self._zeno_ode_wrapped = False
            self._wrap_ode_for_zeno()

        if mass_matrix is not None:
            # Check that the state is a vector or scalar
            assert as_array, "Mass matrix only supported for array-valued states"
            assert (
                len(default_value.shape) <= 1
            ), "Mass matrix only supported for scalar or vector continuous states"
            n = default_value.size
            assert mass_matrix.shape in ((n, n), (n,)), (
                "Mass matrix must be either a square matrix or vector of the same "
                f"size as the continuous state, but got {mass_matrix.shape} for "
                f"continuous state of shape {default_value.shape}."
            )
            if len(mass_matrix.shape) == 1:
                mass_matrix = np.diag(mass_matrix)
            else:
                mass_matrix = np.asarray(mass_matrix)

            # If we end up with an identity matrix, we can just ignore the mass
            # matrix and use the default mass matrix (which is None).  This will
            # allow us to continue using explicit ODE solvers.
            nontrivial_mass_matrix = not np.allclose(mass_matrix, np.eye(n))
            if not nontrivial_mass_matrix:
                mass_matrix = None

        self._mass_matrix = mass_matrix

    @property
    def mass_matrix(self) -> Array:
        # When this is called, an array return value is expected, so we can safely
        # return the mass matrix as an array, even if the internal value is None.
        if self._default_continuous_state is None:
            return None

        if self._mass_matrix is not None:
            return self._mass_matrix

        # Currently only scalar- or vector-valued continuous states are supported,
        # so check that the continuous state (or all tree leaves if tree-structured)
        # is a scalar or vector, and return corresponding identity matrices.
        xc_leaves = tree_util.tree_leaves(self._default_continuous_state)
        if not all(len(xc.shape) <= 1 for xc in xc_leaves):
            raise ValueError(
                "Mass matrix DAEs are only supported when the continuous state is "
                f"scalar- or vector-valued.  System {self.name} has non-vector "
                "continuous state with default value "
                f"{self._default_continuous_state}."
            )

        # Now we are guaranteed that the continuous state is a scalar or vector, so
        # we can return the corresponding (tree-structured) identity matrix.
        return jax.tree.map(lambda x: np.eye(x.size), self._default_continuous_state)

    @property
    def has_mass_matrix(self) -> bool:
        # Does the system have a nontrivial mass matrix?  This will return
        # False if the mass matrix is None or the identity matrix, since
        # the internal _mass_matrix attribute is set to None during
        # continuous state creation in the case where the mass matrix is
        # the identity.
        return self._mass_matrix is not None

    @property
    def continuous_substep_vector(self):
        """T-133: per-entry multirate substep factors for this block.

        Returns pytree-structured int vectors aligned with the flattened
        continuous state (same leaves-concatenation ordering the ODE
        solvers use for ``mass_matrix``), or ``None`` when the block has
        no continuous state. Every entry carries the block-level factor
        declared via ``declare_continuous_state(substeps=N)`` (default 1).
        """
        if self._default_continuous_state is None:
            return None
        factor = int(getattr(self, "_continuous_substeps", 1))
        return jax.tree.map(
            lambda x: np.full((np.asarray(x).size,), factor, dtype=np.int32),
            self._default_continuous_state,
        )

    @property
    def has_multirate_substeps(self) -> bool:
        """True when this block declared ``substeps > 1`` (T-133)."""
        return (
            self._default_continuous_state is not None
            and int(getattr(self, "_continuous_substeps", 1)) > 1
        )

    def declare_discrete_state(
        self,
        shape: ShapeLike = None,
        default_value: Array | Parameter = None,
        dtype: DTypeLike = None,
        as_array: bool = True,
        name: str = None,
    ):
        """Declare a discrete state component for the system.

        The discrete state is a component of the system's state that can be updated
        at specific events, such as zero-crossings or periodic updates.

        .. note::
            Currently only **one** discrete state component is supported per
            ``LeafSystem``.  If ``declare_discrete_state`` is called more than once,
            the second call will silently overwrite the first.  To store several
            independent values, pack them into a single array and split inside your
            update callback.

        Args:
            shape (ShapeLike, optional):
                The shape of the discrete state. Defaults to None.
            default_value (Array, optional):
                The initial value of the discrete state. Defaults to None.
            dtype (DTypeLike, optional):
                The data type of the discrete state. Defaults to None.
            as_array (bool, optional):
                If True, treat the default_value as an array-like (cast if necessary).
                Otherwise, it will be stored as the default state without modification.
            name (str, optional):
                Readability label for the discrete state (parity with
                ``declare_continuous_state_output(name=...)``). Stored as
                ``self.discrete_state_name`` for diagnostics/debugging; it
                does not change runtime behaviour, and the state is still
                read as ``state.discrete_state``.

        Raises:
            AssertionError:
                If as_array is True and neither shape nor default_value is provided.

        Notes:
            (1) Only one of `shape` and `default_value` should be provided. If
            `default_value` is provided, it will be used as the initial value of the
            continuous state. If `shape` is provided, the initial value will be a
            zero vector of the given shape and specified dtype.

            (2) Use `declare_periodic_update` to declare an update event that
            modifies the discrete state at a recurring interval.
        """
        self.discrete_state_name = name
        if as_array:
            default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

        # Tree-map the default value to ensure that it is an array-like with the
        # correct shape and dtype. This is necessary because the default value
        # may be a list, tuple, or other PyTree-structured object.
        default_value = tree_util.tree_map(npa.asarray, default_value)

        # T-027a: if Zeno protection is already installed, pack the user's
        # value alongside the existing Zeno tracker rather than overwriting it.
        if self._zeno_protected_events:
            self._zeno_user_default = default_value
            current = self._default_discrete_state
            zeno_xd = (
                current.zeno
                if isinstance(current, self._zeno_combined_type)
                else current
            )
            self._default_discrete_state = self._zeno_combined_type(
                user=default_value, zeno=zeno_xd
            )
        else:
            self._default_discrete_state = default_value

    def configure_discrete_state_default_value(
        self, default_value: Array, as_array: bool = True
    ):
        if as_array:
            dtype = self._default_discrete_state.dtype
            shape = self._default_discrete_state.shape
            default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

        # Tree-map the default value to ensure that it is an array-like with the
        # correct shape and dtype. This is necessary because the default value
        # may be a list, tuple, or other PyTree-structured object.
        default_value = tree_util.tree_map(npa.asarray, default_value)

        _check_values_compatible(self._default_discrete_state, default_value)

        self._default_discrete_state = default_value

    #
    # I/O declaration
    #
    def _resolve_requires_inputs(
        self,
        requires_inputs: bool | list[int] | None,
        prerequisites_of_calc: List[DependencyTicket] | None,
    ) -> bool | list[int]:
        """Resolve the ``requires_inputs`` flag for an output port.

        Inference is deliberately conservative (T-A4-followup-requires-inputs-infer):
        we only infer ``requires_inputs=False`` when the caller explicitly
        declared ``prerequisites_of_calc=[DependencyTicket.nothing]`` — an
        unambiguous "this output depends on nothing" signal. Every other
        unset case keeps the legacy default of ``True`` (collect all inputs),
        because the established convention is that ``prerequisites_of_calc``
        may list *upstream / transitive* tickets (e.g. ``xcdot`` for a
        derivative output whose callback still reads ``u``, or ``xd`` for a
        sample-and-hold port whose *update* event reads ``u``) while
        ``requires_inputs`` independently controls input collection. Auto-
        flipping to ``False`` from a non-input prereq list would silently
        starve those callbacks of their inputs.
        """
        if requires_inputs is not None:
            return requires_inputs
        if prerequisites_of_calc is not None:
            # The only unambiguous "no inputs" declaration.
            if list(prerequisites_of_calc) == [DependencyTicket.nothing]:
                return False
        # Legacy default: collect all inputs.
        return True

    def declare_output_port(
        self,
        callback: Callable = None,
        period: float = None,
        offset: float = 0.0,
        name: str = None,
        prerequisites_of_calc: List[DependencyTicket] = None,
        default_value: Array = None,
        requires_inputs: bool | list[int] | None = None,
        units=None,
    ) -> int:
        """Declare an output port in the LeafSystem.

        This method accepts a callback function with the block-level signature
            `callback(time, state, *inputs, **parameters) -> value`
        and wraps it to the signature expected by SystemBase.declare_output_port:
            `callback(context) -> value`

        Args:
            callback (Callable):
                The callback function defining the output port.
            period (float, optional):
                If not None, the port will act as a "sample-and-hold", with the
                callback function used to define a periodic update event that refreshes
                the value that will be returned by the port. Typically this should
                match the update period of some associated update event in the system.
                Defaults to None.
            offset (float, optional):
                The offset of the periodic update event. Defaults to 0.0.  Will be ignored
                unless `period` is not None.
            name (str, optional):
                The name of the output port. Defaults to None.
            default_value (Array, optional):
                The default value of the output port, if known. Defaults to None.
            requires_inputs (bool | list[int] | None, optional):
                Whether the callback reads input port values.

                **Defaults to ``None``.** ``None`` resolves to ``True``
                (collect all inputs) in every case except the unambiguous
                ``prerequisites_of_calc=[DependencyTicket.nothing]``
                declaration, which resolves to ``False``
                (T-A4-followup-requires-inputs-infer). The inference is
                deliberately conservative: ``prerequisites_of_calc`` may list
                *upstream / transitive* tickets (e.g. ``xcdot`` for a
                derivative output whose callback still reads ``u``, or ``xd``
                for a sample-and-hold port whose *update* event reads ``u``),
                so a non-input prereq list does **not** imply the callback is
                input-free — only ``[nothing]`` does.

                **Set this to ``False`` explicitly whenever the output does NOT
                depend on any input port** (e.g. a ZOH output that returns a
                stored discrete state, or a CT output that only reads continuous
                state).  This serves two purposes:
                  1. **Eliminates false-positive algebraic-loop detection.**  The
                     diagram-level algebraic-loop checker conservatively assumes every
                     output with ``requires_inputs=True`` has direct feedthrough from
                     all connected inputs.  Declaring ``requires_inputs=False`` tells
                     the checker there is no feedthrough from inputs to this output,
                     which is required to break apparent cycles in discrete feedback
                     topologies (A→B→A) that are valid because updates use x⁻.
                  2. **Reduces compile time** by avoiding unnecessary input collection.

                Can also be specified as a list of integer port indices to declare
                selective feedthrough (only the listed inputs feed through to this
                output).  Defaults to ``True`` (collect all inputs, assume full
                feedthrough).
            prerequisites_of_calc (List[DependencyTicket], optional):
                The dependency tickets for the output port computation.  Defaults to
                None, in which case the assumption is a dependency on either (nothing)
                if `requires_inputs` is False otherwise (inputs).

        Returns:
            int: The index of the declared output port.
        """

        # T-A4-followup-requires-inputs-infer: when the caller leaves
        # ``requires_inputs`` unset (None) but supplies ``prerequisites_of_calc``,
        # infer whether inputs are needed from the prerequisites rather than
        # forcing the user to keep the two arguments consistent by hand.
        requires_inputs = self._resolve_requires_inputs(
            requires_inputs, prerequisites_of_calc
        )

        if default_value is not None:
            default_value = npa.array(default_value)

        cache_index = None
        if period is not None:
            # The output port will be of "sample-and-hold" type, so we have to declare a
            # periodic event to update the value.  The callback will be used to define the
            # update event, and the output callback will simply return the stored value.

            # This is the index that this port value will have in state.cache
            cache_index = len(self._default_cache)
            self._default_cache.append(default_value)

        output_port_idx = super().declare_output_port(
            callback, name=name, cache_index=cache_index, units=units
        )

        self.configure_output_port(
            output_port_idx,
            callback,
            period=period,
            offset=offset,
            prerequisites_of_calc=prerequisites_of_calc,
            default_value=default_value,
            requires_inputs=requires_inputs,
        )

        return output_port_idx

    def configure_output_port(
        self,
        port_index: int,
        callback: Callable,
        period: float = None,
        offset: float = 0.0,
        prerequisites_of_calc: List[DependencyTicket] = None,
        default_value: Array = None,
        requires_inputs: bool | list[int] | None = None,
    ):
        """Configure an output port in the LeafSystem.

        See `declare_output_port` for a description of the arguments.

        Args:
            port_index (int):
                The index of the output port to configure.

        Returns:
            None
        """
        if default_value is not None:
            default_value = npa.array(default_value)

        # Infer requires_inputs from prerequisites when left unset, so a
        # standalone configure_output_port call gets the same ergonomics as
        # declare_output_port. (T-A4-followup-requires-inputs-infer)
        requires_inputs = self._resolve_requires_inputs(
            requires_inputs, prerequisites_of_calc
        )

        # To help avoid unnecessary flagging of algebraic loops, trim the inputs as a
        # default prereq if the output callback doesn't use them
        if prerequisites_of_calc is None:
            if requires_inputs:
                prerequisites_of_calc = [DependencyTicket.u]
            else:
                prerequisites_of_calc = [DependencyTicket.nothing]

        if period is None:
            event = None
            _output_callback = self.wrap_callback(
                callback, collect_inputs=requires_inputs
            )
            cache_index = None

        else:
            # The output port will be of "sample-and-hold" type, so we have to declare a
            # periodic event to update the value.  The callback will be used to define the
            # update event, and the output callback will simply return the stored value.

            # This is the index that this port value will have in state.cache
            cache_index = self.output_ports[port_index].cache_index
            if cache_index is None:
                cache_index = len(self._default_cache)
                self._default_cache.append(default_value)

            def _output_callback(context: ContextBase) -> Array:
                state = context[self.system_id].state
                return state.cache[cache_index]

            def _update_callback(
                time: Scalar, state: LeafState, *inputs, **parameters
            ) -> LeafState:
                output = callback(time, state, *inputs, **parameters)
                return state.with_cached_value(cache_index, output)

            _update_callback = self.wrap_callback(
                _update_callback, collect_inputs=requires_inputs
            )

            # Create the associated update event
            event = DiscreteUpdateEvent(
                system_id=self.system_id,
                event_data=PeriodicEventData(
                    period=period, offset=offset, active=False
                ),
                name=f"{self.name}:output_{cache_index}",
                callback=_update_callback,
                passthrough=self._passthrough,
            )

            # Note that in this case the "prerequisites of calc" will correspond to the
            # prerequisites of the update event, not the literal output callback itself.
            # However, these can be used to determine dependencies for the update event
            # via the output port.

        super().configure_output_port(
            port_index,
            _output_callback,
            prerequisites_of_calc=prerequisites_of_calc,
            default_value=default_value,
            event=event,
            cache_index=cache_index,
        )

    def configure_continuous_state_default_value(
        self, callback_idx: int, default_value: Array, as_array: bool = True
    ):
        if as_array:
            dtype = self._default_continuous_state.dtype
            shape = self._default_continuous_state.shape
            default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

        # Tree-map the default value to ensure that it is an array-like with the
        # correct shape and dtype. This is necessary because the default value
        # may be a list, tuple, or other PyTree-structured object.
        default_value = tree_util.tree_map(npa.asarray, default_value)

        _check_values_compatible(self._default_continuous_state, default_value)

        self._default_continuous_state = default_value
        if self._continuous_state_output_port_idx is not None:
            port = self.output_ports[self._continuous_state_output_port_idx]
            port.default_value = default_value
            self._default_cache[port.cache_index] = default_value

    def configure_output_port_default_value(
        self,
        port_index: int,
        default_value: Array,
    ):
        port = self.output_ports[port_index]
        if port.event is None:
            # T-107-followup-transport-delay-warn-quiet — demoted from
            # ``logger.warning`` to ``logger.debug``. ``TransportDelay``
            # (and similar event-less ports) emit this on every
            # construction; the behaviour is correct (the default really
            # is unused because the port has no periodic event), but the
            # WARNING level reads like a user-facing bug. The information
            # is still available at DEBUG for anyone wiring up a new
            # event-less port that *intended* to use ``default_value``.
            logger.debug(
                "period is None so default_value is not used for port %d in block %s",
                port_index,
                self.name,
            )
            return
        default_value = npa.array(default_value)
        cache_index = self.output_ports[port_index].cache_index

        if cache_index is None:
            raise ValueError(
                "Output port does not have a cache index, so default value cannot be set"
            )

        _check_values_compatible(self._default_cache[cache_index], default_value)
        self._default_cache[cache_index] = default_value

    def declare_continuous_state_output(
        self,
        name: str = None,
    ) -> int:
        """Declare a continuous state output port in the system.

        This method creates a new block-level output port which returns the full
        continuous state of the system.

        Args:
            name (str, optional):
                The name of the output port. Defaults to None (autogenerate name).

        Returns:
            int: The index of the new output port.
        """
        if self._continuous_state_output_port_idx is not None:
            raise ValueError("Continuous state output port already declared")

        def _callback(time: Scalar, state: LeafState, *inputs, **parameters):
            return state.continuous_state

        self._continuous_state_output_port_idx = self.declare_output_port(
            _callback,
            name=name,
            prerequisites_of_calc=[DependencyTicket.xc],
            default_value=self._default_continuous_state,
            requires_inputs=False,
        )
        return self._continuous_state_output_port_idx

    def declare_mode_output(self, name: str = None) -> int:
        """Declare a mode output port in the system.

        This method creates a new block-level output port which returns the component
        of the system's state corresponding to the discrete "mode" or "stage".

        Args:
            name (str, optional):
                The name of the output port. Defaults to None.

        Returns:
            int:
                The index of the declared mode output port.
        """

        def _callback(time: Scalar, state: LeafState, *inputs, **parameters):
            return state.mode

        self._mode_output_port_idx = self.declare_output_port(
            _callback,
            name=name,
            prerequisites_of_calc=[DependencyTicket.mode],
            default_value=self._default_mode,
            requires_inputs=False,
        )

        return self._mode_output_port_idx

    #
    # Event declaration
    #
    def declare_periodic_update(
        self,
        callback: Callable = None,
        period: Scalar | Parameter = None,
        offset: Scalar | Parameter = None,
        enable_tracing: bool = None,
    ):
        self._state_update_events.append(None)
        event_idx = len(self._state_update_events) - 1

        # FIXME: this is to preserve some backward compatibility while we decouple
        # declaration from configuration. Declaration should not have to call
        # configuration.
        if callback is not None:
            # Default ``offset`` to 0.0 when only ``period`` is supplied. Leaving
            # offset=None would propagate into PeriodicEventData and trigger an
            # opaque ``npa.minimum(None, ...)`` TypeError at the first scheduler
            # tick rather than at construction.
            if period is not None and offset is None:
                offset = 0.0
            self.configure_periodic_update(
                event_idx,
                callback,
                period,
                offset,
                enable_tracing=enable_tracing,
            )
        return event_idx

    def configure_periodic_update(
        self,
        event_index: int,
        callback: Callable,
        period: Scalar | Parameter,
        offset: Scalar | Parameter,
        enable_tracing: bool = None,
    ):
        """Configure an existing periodic update event.

        The event will be triggered at regular intervals defined by the period and
        offset parameters. The callback should have the signature
        `callback(time, state, *inputs, **params) -> xd_plus`, where `xd_plus` is the
        updated value of the discrete state.

        This callback should be written to compute the "plus" value of the discrete
        state component given the "minus" values of all state components and inputs.

        Args:
            event_index (int):
                The index of the event to configure.
            callback (Callable):
                The callback function defining the update.
            period (Scalar):
                The period at which the update event occurs.
            offset (Scalar):
                The offset at which the first occurrence of the event is triggered.
            enable_tracing (bool, optional):
                If True, enable tracing for this event. Defaults to None.
        """
        _wrapped_callback = self.wrap_callback(callback)

        def _callback(context: ContextBase) -> LeafState:
            xd = _wrapped_callback(context)
            return context[self.system_id].state.with_discrete_state(xd)

        if enable_tracing is None:
            enable_tracing = True

        event = DiscreteUpdateEvent(
            system_id=self.system_id,
            name=f"{self.name}:periodic_update",
            event_data=PeriodicEventData(period=period, offset=offset, active=False),
            callback=_callback,
            passthrough=self._passthrough,
            enable_tracing=enable_tracing,
            is_state_update=True,
        )
        self._state_update_events[event_index] = event

    def declare_default_mode(self, mode: int):
        self._default_mode = mode

    def configure_default_mode(self, mode: int):
        self._default_mode = mode
        if self._mode_output_port_idx:
            self.configure_output_port_default_value(self._mode_output_port_idx, mode)

    def declare_zero_crossing(
        self,
        guard: Callable,
        reset_map: Callable = None,
        start_mode: int = None,
        end_mode: int = None,
        direction: str = "crosses_zero",
        terminal: bool = False,
        name: str = None,
        enable_tracing: bool = None,
        zeno_tolerance: float | None = None,
        grad_guard: Callable = None,
    ):
        """Declare an event triggered by a zero-crossing of a guard function.

        Optionally, the system can also transition between discrete modes
        If `start_mode` and `end_mode` are specified, the system will transition
        from `start_mode` to `end_mode` when the event is triggered according to `guard`.
        This event will be active conditionally on `state.mode == start_mode` and when
        triggered will result in applying the reset map. In addition, the mode will be
        updated to `end_mode`.

        If `start_mode` and `end_mode` are not specified, the event will always be active
        and will not result in a mode transition.

        The guard function should have the signature:
            `guard(time, state, *inputs, **parameters) -> float`

        and the reset map should have the signature of an unrestricted update:
            `reset_map(time, state, *inputs, **parameters) -> state`

        Args:
            guard (Callable):
                The guard function which triggers updates on zero crossing.
            reset_map (Callable, optional):
                The reset map which is applied when the event is triggered. If None
                (default), no reset is applied.
            start_mode (int, optional):
                The mode or stage of the system in which the guard will be
                actively monitored. If None (default), the event will always be
                active.
            end_mode (int, optional):
                The mode or stage of the system to which the system will transition
                when the event is triggered. If start_mode is None, this is ignored.
                Otherwise it _must_ be specified, though it can be the same as
                start_mode.
            direction (str, optional):
                The direction of the zero crossing. Options are "crosses_zero"
                (default), "positive_then_non_positive", "negative_then_non_negative",
                and "edge_detection".  All except edge detection operate on continuous
                signals; edge detection operates on boolean signals and looks for a
                jump from False to True or vice versa.
            terminal (bool, optional):
                If True, the event will halt simulation if and when the zero-crossing
                occurs. If this event is triggered the reset map will still be applied
                as usual prior to termination. Defaults to False.
            name (str, optional):
                The name of the event. Defaults to None.
            enable_tracing (bool, optional):
                If True, enable tracing for this event. Defaults to None.

        Notes:
            By default the system state does not have a "mode" component, so in
            order to declare "state transitions" with non-null start and end modes,
            the user must first call `declare_default_mode` to set the default mode
            to be some integer (initial condition for the system).
        """

        logger.debug(
            f"Declaring transition for {self.name} with guard {guard} and reset map {reset_map}"
        )

        if enable_tracing is None:
            enable_tracing = True

        if start_mode is not None or end_mode is not None:
            assert (
                self._default_mode is not None
            ), "System has no mode: call `declare_default_mode` before transitions."
            assert isinstance(start_mode, int) and isinstance(end_mode, int)

        # T-027: optional Zeno-hold protection. If `zeno_tolerance` is set,
        # wrap the user's reset_map to flag a Zeno entry, declare a companion
        # exit event, and freeze the continuous-state ODE while held.
        if zeno_tolerance is not None:
            assert (
                isinstance(zeno_tolerance, (float, int)) and float(zeno_tolerance) > 0.0
            ), "zeno_tolerance must be a positive float"
            reset_map, _zeno_companion = self._install_zeno_protection(
                reset_map=reset_map,
                guard=guard,
                direction=direction,
                tol=float(zeno_tolerance),
                name=name,
            )
        else:
            _zeno_companion = None

        # Wrap the reset map with a mode update if necessary
        def _reset_and_update_mode(
            time: Scalar, state: LeafState, *inputs, **parameters
        ) -> LeafState:
            if reset_map is not None:
                state = reset_map(time, state, *inputs, **parameters)
            logger.debug(f"Updating mode from {state.mode} to {end_mode}")

            # If the start and end modes are declared, update the mode
            if start_mode is not None:
                logger.debug(f"Updating mode from {state.mode} to {end_mode}")
                state = state.with_mode(end_mode)

            return state

        _wrapped_guard = self.wrap_callback(guard)
        # Optional smooth guard residual for the event-time (saltation) gradient
        # only — wrapped the same way as the trigger guard.  ``None`` keeps the
        # legacy behaviour (the saltation paths fall back to ``guard``).
        _wrapped_grad_guard = (
            self.wrap_callback(grad_guard) if grad_guard is not None else None
        )
        _wrapped_reset = _wrap_reset_map(
            self, _reset_and_update_mode, _wrapped_guard, terminal,
            grad_guard=_wrapped_grad_guard,
        )

        event = ZeroCrossingEvent(
            system_id=self.system_id,
            guard=_wrapped_guard,
            grad_guard=_wrapped_grad_guard,
            reset_map=_wrapped_reset,
            passthrough=self._passthrough,
            direction=direction,
            is_terminal=terminal,
            name=name,
            event_data=ZeroCrossingEventData(active=True, triggered=False),
            enable_tracing=enable_tracing,
            active_mode=start_mode,
        )

        event_index = len(self._zero_crossing_events)
        self._zero_crossing_events.append(event)

        # T-115-followup-saturate-rate-classification: bump the
        # behavioral-ZC counter only when this event actually does
        # something on trigger — has a user-supplied reset map or
        # participates in a mode transition. Pure guard-only events
        # (Saturate / DeadZone clip boundaries) are solver hints with
        # no behavioral effect, so they should not flip the block's
        # rate-group classification to ``event_driven``.
        if (
            reset_map is not None
            or start_mode is not None
            or end_mode is not None
        ):
            self._n_behavioral_zc_events += 1

        # Record the transition in the transition map (for debugging or analysis)
        if start_mode is not None:
            if start_mode not in self.transition_map:
                self.transition_map[start_mode] = []
            self.transition_map[start_mode].append((event_index, event))

        # T-027: register the companion `_exit_zeno` event AFTER the main event
        # so the slot index is finalized first.
        if _zeno_companion is not None:
            _zeno_companion(event_index)

    def _install_zeno_protection(
        self,
        reset_map: Callable | None,
        guard: Callable,
        direction: str,
        tol: float,
        name: str | None,
    ) -> Tuple[Callable, Callable]:
        """T-027: install per-event Zeno-hold tracking on a `declare_zero_crossing`.

        Returns ``(wrapped_reset_map, register_companion)``.

        - ``wrapped_reset_map`` runs the user's reset, then sets the per-event
          Zeno flag if `(time - tprev) < tol`.
        - ``register_companion(event_index)`` declares the partner ``_exit_zeno``
          event whose reset clears the flag. Called by ``declare_zero_crossing``
          once the main event index is known.

        Storage layout:
        - If the host LeafSystem has NO user discrete state, the discrete state
          slot holds the bare ``_ZenoState(zeno, tprev)`` NamedTuple of arrays
          sized to the number of protected events.
        - T-027a: if the host LeafSystem ALSO calls ``declare_discrete_state``,
          the discrete state is packed as ``_DiscreteWithZeno(user, zeno)``.
          User callbacks (ode, guard, reset) are wrapped to see only their own
          ``user`` slot via ``state.discrete_state``; the framework re-packs on
          the way out. Order of declarations does not matter.

        The host's continuous-state ode is wrapped exactly once: when ANY
        protected event has ``zeno=True``, the ode output is multiplied by
        0 to freeze the state.
        """
        from collections import namedtuple

        # Reverse the direction for the exit companion event.
        _reverse = {
            "positive_then_non_positive": "negative_then_non_negative",
            "negative_then_non_negative": "positive_then_non_positive",
            "crosses_zero": "crosses_zero",
        }
        if direction not in _reverse:
            raise ValueError(
                f"zeno_tolerance is not supported for direction={direction!r}"
            )
        exit_direction = _reverse[direction]

        # Allocate this event's slot. We grow the discrete state lazily so the
        # number of protected events does not need to be known up-front.
        slot = len(self._zeno_protected_events)
        self._zeno_protected_events.append(
            {"slot": slot, "tol": tol, "name": name}
        )

        # On first installation: set up types and capture any pre-existing
        # user discrete state so we can pack it alongside the Zeno tracker.
        if slot == 0:
            self._zeno_state_type = namedtuple("_ZenoState", ["zeno", "tprev"])
            self._zeno_combined_type = namedtuple(
                "_DiscreteWithZeno", ["user", "zeno"]
            )
            # If the user already declared discrete state, capture its default
            # so we can preserve it inside the combined wrapper. Subsequent
            # `declare_discrete_state` calls also flow through this path
            # (they update `_default_discrete_state` directly; we re-pack on
            # `create_state`).
            self._zeno_user_default = self._default_discrete_state
        # (Re)build the default value so it always matches the current count.
        n = slot + 1
        zeno_default = self._zeno_state_type(
            zeno=npa.zeros(n, dtype=bool),
            tprev=npa.zeros(n, dtype=float),
        )
        if self._zeno_user_default is not None:
            self._default_discrete_state = self._zeno_combined_type(
                user=self._zeno_user_default, zeno=zeno_default
            )
        else:
            self._default_discrete_state = zeno_default

        zeno_type = self._zeno_state_type
        combined_type = self._zeno_combined_type

        def _split_xd(xd):
            """Return (user_xd_or_None, zeno_xd) given the framework discrete state."""
            if isinstance(xd, combined_type):
                return xd.user, xd.zeno
            return None, xd

        def _pack_xd(user_xd, zeno_xd):
            """Re-pack the framework discrete state from updated parts."""
            if user_xd is None:
                return zeno_xd
            return combined_type(user=user_xd, zeno=zeno_xd)

        # Wrap the user's reset_map: present a "user view" of state.discrete_state
        # (just their own xd, not the wrapper), run the reset, and re-pack with
        # the updated Zeno tracker.
        user_reset = reset_map

        def _zeno_aware_reset(time, state, *inputs, **params):
            user_xd, zeno_xd = _split_xd(state.discrete_state)
            if user_reset is not None:
                if user_xd is not None:
                    # Present a user-facing view: replace the wrapper with just
                    # the user's slot. Bypass `with_discrete_state`'s
                    # tree-shape coercion (the wrapper has a different
                    # structure from the user's xd).
                    user_view = dataclasses.replace(state, discrete_state=user_xd)
                else:
                    user_view = state
                out_state = user_reset(time, user_view, *inputs, **params)
                new_user_xd = out_state.discrete_state if user_xd is not None else None
                # Carry over continuous_state / mode / cache from the user's return.
                state = out_state
            else:
                new_user_xd = user_xd
            dt = time - zeno_xd.tprev[slot]
            entered = (dt - tol) <= 0.0
            new_zeno = zeno_xd.zeno.at[slot].set(
                npa.logical_or(zeno_xd.zeno[slot], entered)
            )
            new_tprev = zeno_xd.tprev.at[slot].set(time)
            new_zeno_xd = zeno_type(zeno=new_zeno, tprev=new_tprev)
            return dataclasses.replace(
                state, discrete_state=_pack_xd(new_user_xd, new_zeno_xd)
            )

        # Companion `_exit_zeno` event: its guard is the user's guard re-used,
        # but with reversed direction so it fires when the trigger condition
        # disappears. The reset clears the slot's `zeno` flag (and leaves the
        # user's discrete state unchanged).
        def _register_companion(main_event_index: int):
            def _exit_reset(time, state, *_inputs, **_params):
                user_xd, zeno_xd = _split_xd(state.discrete_state)
                new_zeno = zeno_xd.zeno.at[slot].set(False)
                new_zeno_xd = zeno_type(zeno=new_zeno, tprev=zeno_xd.tprev)
                return dataclasses.replace(
                    state, discrete_state=_pack_xd(user_xd, new_zeno_xd)
                )

            self.declare_zero_crossing(
                guard=guard,
                reset_map=_exit_reset,
                direction=exit_direction,
                name=(f"{name}__exit_zeno" if name else "_exit_zeno"),
            )

        # Wrap the ode once: multiply by `(1 - any_zeno)` so any protected
        # event being held freezes the entire host system.
        if not self._zeno_ode_wrapped and self.ode_callback is not None:
            self._wrap_ode_for_zeno()
        # Mark a deferred wrap if `declare_continuous_state` happens later.
        self._zeno_ode_wrapped = self._zeno_ode_wrapped or (
            self.ode_callback is not None
        )

        return _zeno_aware_reset, _register_companion

    def _wrap_ode_for_zeno(self) -> None:
        """T-027: replace the ode_callback to freeze when any Zeno flag is set.

        T-027a: handle the combined ``_DiscreteWithZeno(user, zeno)`` layout.
        """
        if self.ode_callback is None or self.ode_callback._callback is None:
            return
        original = self.ode_callback._callback
        sys_id = self.system_id
        combined_type = self._zeno_combined_type

        def _frozen_ode(context):
            xdot = original(context)
            xd = context[sys_id].state.discrete_state
            zeno_xd = xd.zeno if isinstance(xd, combined_type) else xd
            any_zeno = npa.any(zeno_xd.zeno)
            scale = npa.where(any_zeno, 0.0, 1.0)
            return tree_util.tree_map(lambda v: v * scale, xdot)

        self.ode_callback._callback = _frozen_ode
        # Keep `eval_time_derivatives` pointing at the (now-wrapped) callback.
        self.eval_time_derivatives = self.ode_callback.eval
        self._zeno_ode_wrapped = True

    #
    # Initialization
    #
    @property
    def context_factory(self) -> LeafContextFactory:
        return LeafContextFactory(self)

    @property
    def dependency_graph_factory(self) -> LeafDependencyGraphFactory:
        return LeafDependencyGraphFactory(self)

    def create_state(self) -> LeafState:
        # Hook for context creation: get the default state for this system.
        # Users should not need to call this directly - the state will be created
        # as part of the context.  Generally, `system.create_context()` should
        # be all that's necessary for initialization.
        self.reset_default_values(**self.dynamic_parameters)
        return LeafState(
            name=self.name,
            continuous_state=self._default_continuous_state,
            discrete_state=self._default_discrete_state,
            mode=self._default_mode,
            cache=tuple(self._default_cache),
        )

    def initialize_static_data(self, context: ContextBase):
        # Try to infer any missing default values for "sample-and-hold" output ports
        # and any other cached computations.
        cached_callbacks: list[SystemCallback] = [
            cb for cb in self.callbacks if cb.cache_index is not None
        ]

        for callback in cached_callbacks:
            i = callback.cache_index
            if self._default_cache[i] is None:
                try:
                    if isinstance(callback, OutputPort):
                        # Try to eval the callback for the _event_ (not the output
                        # port return function), which would return a value of the
                        # right data type for the output port, provided it is connected
                        _eval = callback.event.callback
                    else:
                        # If it's not an output port, the callback function evaluation
                        # should return the correct data type.
                        _eval = callback.eval

                    state: LeafState = _eval(context)
                    y = state.cache[i]
                    self._default_cache[i] = y
                    local_context = context[self.system_id].with_cached_value(i, y)
                    context = context.with_subcontext(self.system_id, local_context)
                except UpstreamEvalError:
                    logger.debug(
                        "%s.initialize_static_data: UpstreamEvalError. "
                        "Continuing without default value initialization.",
                        self.name,
                    )

        return context

    def _create_dependency_cache(self) -> dict[int, CallbackTracer]:
        cache = {}
        for source in self.callbacks:
            cache[source.callback_index] = CallbackTracer(ticket=source.ticket)
        return cache

    # Inherits docstring from SystemBase.get_feedthrough
    def get_feedthrough(self) -> List[Tuple[int, int]]:
        # NOTE: This implementation is basically a direct port of the Drake algorithm

        if self.dependency_graph is None:
            raise ValueError("Must create dependency graph first.")

        # If we already did this or it was set manually, return the stored value
        if self.feedthrough_pairs is not None:
            return self.feedthrough_pairs

        feedthrough = []  # Confirmed feedthrough pairs (input, output)

        # First collect all possible feedthrough pairs
        unknown: Set[Tuple[int, int]] = set()
        for iport in self.input_ports:
            for oport in self.output_ports:
                unknown.add((iport.index, oport.index))

        if len(unknown) == 0:
            return feedthrough

        # Create a local context and "cache".  The cache here just contains CallbackTracer
        # objects that can be used to trace dependencies through the system, but
        # otherwise don't store any actual values.  This is different from any "cached"
        # computations that might be stored in the state for reuse by multiple ports or
        # downstream calculations within the system.
        #
        # This cache will only contain local sources - this is fine since we're just
        # testing local input -> output paths for this system.
        cache = self._create_dependency_cache()

        original_unknown = unknown.copy()
        for pair in original_unknown:
            u, v = pair
            output_port = self.output_ports[v]
            input_port = self.input_ports[u]

            # If output prerequisites are unspecified, this tells us nothing
            if DependencyTicket.all_sources in output_port.prerequisites_of_calc:
                continue

            # Determine feedthrough dependency via cache invalidation
            cache = _mark_up_to_date(cache, output_port.callback_index)

            # Notify subscribers of a value change in the input, invalidating all
            # downstream cache values
            input_tracker = self.dependency_graph[input_port.ticket]
            cache = input_tracker.notify_subscribers(
                cache, self.dependency_graph, local_only=True
            )

            # If the output cache is now out of date, this is a feedthrough path
            if cache[output_port.callback_index].is_out_of_date:
                feedthrough.append(pair)

            # Regardless of the result of the caching, the pair is no longer unknown
            unknown.remove(pair)

            # Reset the output cache to out-of-date in case other inputs also
            # feed through to this output.
            cache = _mark_out_of_date(cache, output_port.callback_index)

        logger.debug(f"{self.name} feedthrough pairs: {feedthrough}")

        # Conservatively assume everything still unknown is feedthrough
        for pair in unknown:
            feedthrough.append(pair)

        self.feedthrough_pairs = feedthrough
        return self.feedthrough_pairs

    def reset_default_values(self, **dynamic_parameters):
        """This function is used to reset default values for
        continuous/discrete states, ports and mode based on dynamic parameters.
        It is called in `create_state()` and used to reset states in ensemble sims
        and optimization with the context method `with_new_state()`.

        Note that dtypes and shapes can't be changed after initialization because
        the diagram may already have been jax-compiled. Only values may change.
        """
        pass

`continuous_state_default` `property`

The declared default continuous-state value (read-only).

This is the default_value passed to declare_continuous_state (or the array inferred from shape / dtype), i.e. the value that seeds context.continuous_state before any user override. Returns None when the block has no continuous state. Exposed as a documented accessor so callers don't have to reach into the private _default_continuous_state attribute (T-C2-followup).

`continuous_substep_vector` `property`

T-133: per-entry multirate substep factors for this block.

Returns pytree-structured int vectors aligned with the flattened continuous state (same leaves-concatenation ordering the ODE solvers use for mass_matrix), or None when the block has no continuous state. Every entry carries the block-level factor declared via declare_continuous_state(substeps=N) (default 1).

`has_multirate_substeps` `property`

True when this block declared substeps > 1 (T-133).

`configure_continuous_state(callback_idx, shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None)`

Configure a continuous state component for the system.

The ode callback computes the time derivative of the continuous state based on the current time, state, and any additional inputs. If ode is not provided, a default zero vector of the same size as the continuous state is used. If provided, the ode callback should have the signature ode(time, state, *inputs, **params) -> xcdot.

Parameters:

Name	Type	Description	Default
`callback_idx`	`int`	The index of the callback in the system's callback list.	required
`shape`	`ShapeLike`	The shape of the continuous state vector. Defaults to None.	`None`
`default_value`	`Array`	The initial value of the continuous state vector. Defaults to None.	`None`
`dtype`	`DTypeLike`	The data type of the continuous state vector. Defaults to None.	`None`
`ode`	`Callable`	The callback for computing the time derivative of the continuous state. Should have the signature: `ode(time, state, inputs, *parameters) -> xcdot`. Defaults to None.	`None`
`mass_matrix`	`Array`	The mass matrix for the continuous state. Defaults to None. If provided, must be a square matrix with the same shape as the continuous state. Using a mass matrix different from the identity in any LeafSystem will require the use of a compatible continuous-time solver (currently only BDF is supported). Currently mass matrices are also only supported for scalar- or vector-valued continuous states ( i.e. no matrices or other PyTree-structured states).	`None`
`as_array`	`bool`	If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification.	`True`
`requires_inputs`	`bool`	If True, indicates that the ODE computation requires inputs.	`True`
`prerequisites_of_calc`	`List[DependencyTicket]`	The dependency tickets for the ODE computation. Defaults to None, in which case the assumption is a dependency on either (time, continuous state) if `requires_inputs` is False, otherwise (time, continuous state, inputs.	`None`

Raises:

Type	Description
`AssertionError`	If neither shape nor default_value is provided, or if the mass matrix is inconsistent with the continuous state.

Notes

(1) Only one of shape and default_value should be provided. If default_value is provided, it will be used as the initial value of the continuous state. If shape is provided, the initial value will be a zero vector of the given shape and specified dtype.

Source code in jaxonomy/framework/leaf_system.py

def configure_continuous_state(
    self,
    callback_idx: int,
    shape: ShapeLike = None,
    default_value: Array = None,
    dtype: DTypeLike = None,
    ode: Callable = None,
    mass_matrix: Array = None,
    as_array: bool = True,
    requires_inputs: bool = True,
    prerequisites_of_calc: List[DependencyTicket] = None,
):
    """Configure a continuous state component for the system.

    The `ode` callback computes the time derivative of the continuous state based on the
    current time, state, and any additional inputs. If `ode` is not provided, a default
    zero vector of the same size as the continuous state is used. If provided, the `ode`
    callback should have the signature `ode(time, state, *inputs, **params) -> xcdot`.

    Args:
        callback_idx (int):
            The index of the callback in the system's callback list.
        shape (ShapeLike, optional):
            The shape of the continuous state vector. Defaults to None.
        default_value (Array, optional):
            The initial value of the continuous state vector. Defaults to None.
        dtype (DTypeLike, optional):
            The data type of the continuous state vector. Defaults to None.
        ode (Callable, optional):
            The callback for computing the time derivative of the continuous state.
            Should have the signature:
                `ode(time, state, *inputs, **parameters) -> xcdot`.
            Defaults to None.
        mass_matrix (Array, optional):
            The mass matrix for the continuous state. Defaults to None. If
            provided, must be a square matrix with the same shape as the
            continuous state.  Using a mass matrix different from the identity
            in any LeafSystem will require the use of a compatible continuous-time
            solver (currently only BDF is supported).  Currently mass matrices are
            also only supported for scalar- or vector-valued continuous states (
            i.e. no matrices or other PyTree-structured states).
        as_array (bool, optional):
            If True, treat the default_value as an array-like (cast if necessary).
            Otherwise, it will be stored as the default state without modification.
        requires_inputs (bool, optional):
            If True, indicates that the ODE computation requires inputs.
        prerequisites_of_calc (List[DependencyTicket], optional):
            The dependency tickets for the ODE computation. Defaults to None, in
            which case the assumption is a dependency on either (time, continuous
            state) if `requires_inputs` is False, otherwise (time, continuous state,
            inputs.

    Raises:
        AssertionError:
            If neither shape nor default_value is provided, or if the mass matrix
            is inconsistent with the continuous state.

    Notes:
        (1) Only one of `shape` and `default_value` should be provided. If `default_value`
        is provided, it will be used as the initial value of the continuous state. If
        `shape` is provided, the initial value will be a zero vector of the given shape
        and specified dtype.
    """

    if prerequisites_of_calc is None:
        prerequisites_of_calc = [DependencyTicket.time, DependencyTicket.xc]
        if requires_inputs:
            prerequisites_of_calc.append(DependencyTicket.u)

    if as_array:
        default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

    logger.debug(f"In block {self.name} [{self.system_id}]: {default_value=}")

    # Tree-map the default value to ensure that it is an array-like with the
    # correct shape and dtype. This is necessary because the default value
    # may be a list, tuple, or other PyTree-structured object.
    default_value = tree_util.tree_map(npa.asarray, default_value)

    self._default_continuous_state = default_value
    if self._continuous_state_output_port_idx is not None:
        port = self.output_ports[self._continuous_state_output_port_idx]
        port.default_value = default_value
        self._default_cache[port.cache_index] = default_value

    if ode is None:
        # If no ODE is specified, return a zero vector of the same size as the
        # continuous state. This will break if the continuous state is
        # a named tuple, in which case a custom ODE must be provided.
        assert as_array, "Must provide custom ODE for non-array continuous state"

        def ode(time, state, *inputs, **parameters):
            return npa.zeros_like(default_value)

    # Wrap the ode function to accept a context and return the time derivatives.
    ode = self.wrap_callback(ode)

    # Declare the time derivative function as a system callback so that its
    # dependencies can be tracked in the system dependency graph
    self.ode_callback._callback = ode
    self.ode_callback.prerequisites_of_calc = prerequisites_of_calc

    # Override the default `eval_time_derivatives` to use the wrapped ODE function
    self.eval_time_derivatives = self.ode_callback.eval

    # T-027: if Zeno-protected events were registered before this, wrap the
    # ode so a Zeno-hold freezes the continuous state.
    if self._zeno_protected_events:
        self._zeno_ode_wrapped = False
        self._wrap_ode_for_zeno()

    if mass_matrix is not None:
        # Check that the state is a vector or scalar
        assert as_array, "Mass matrix only supported for array-valued states"
        assert (
            len(default_value.shape) <= 1
        ), "Mass matrix only supported for scalar or vector continuous states"
        n = default_value.size
        assert mass_matrix.shape in ((n, n), (n,)), (
            "Mass matrix must be either a square matrix or vector of the same "
            f"size as the continuous state, but got {mass_matrix.shape} for "
            f"continuous state of shape {default_value.shape}."
        )
        if len(mass_matrix.shape) == 1:
            mass_matrix = np.diag(mass_matrix)
        else:
            mass_matrix = np.asarray(mass_matrix)

        # If we end up with an identity matrix, we can just ignore the mass
        # matrix and use the default mass matrix (which is None).  This will
        # allow us to continue using explicit ODE solvers.
        nontrivial_mass_matrix = not np.allclose(mass_matrix, np.eye(n))
        if not nontrivial_mass_matrix:
            mass_matrix = None

    self._mass_matrix = mass_matrix

`configure_output_port(port_index, callback, period=None, offset=0.0, prerequisites_of_calc=None, default_value=None, requires_inputs=None)`

Configure an output port in the LeafSystem.

See declare_output_port for a description of the arguments.

Parameters:

Name	Type	Description	Default
`port_index`	`int`	The index of the output port to configure.	required

Returns:

Type	Description
	None

Source code in jaxonomy/framework/leaf_system.py

def configure_output_port(
    self,
    port_index: int,
    callback: Callable,
    period: float = None,
    offset: float = 0.0,
    prerequisites_of_calc: List[DependencyTicket] = None,
    default_value: Array = None,
    requires_inputs: bool | list[int] | None = None,
):
    """Configure an output port in the LeafSystem.

    See `declare_output_port` for a description of the arguments.

    Args:
        port_index (int):
            The index of the output port to configure.

    Returns:
        None
    """
    if default_value is not None:
        default_value = npa.array(default_value)

    # Infer requires_inputs from prerequisites when left unset, so a
    # standalone configure_output_port call gets the same ergonomics as
    # declare_output_port. (T-A4-followup-requires-inputs-infer)
    requires_inputs = self._resolve_requires_inputs(
        requires_inputs, prerequisites_of_calc
    )

    # To help avoid unnecessary flagging of algebraic loops, trim the inputs as a
    # default prereq if the output callback doesn't use them
    if prerequisites_of_calc is None:
        if requires_inputs:
            prerequisites_of_calc = [DependencyTicket.u]
        else:
            prerequisites_of_calc = [DependencyTicket.nothing]

    if period is None:
        event = None
        _output_callback = self.wrap_callback(
            callback, collect_inputs=requires_inputs
        )
        cache_index = None

    else:
        # The output port will be of "sample-and-hold" type, so we have to declare a
        # periodic event to update the value.  The callback will be used to define the
        # update event, and the output callback will simply return the stored value.

        # This is the index that this port value will have in state.cache
        cache_index = self.output_ports[port_index].cache_index
        if cache_index is None:
            cache_index = len(self._default_cache)
            self._default_cache.append(default_value)

        def _output_callback(context: ContextBase) -> Array:
            state = context[self.system_id].state
            return state.cache[cache_index]

        def _update_callback(
            time: Scalar, state: LeafState, *inputs, **parameters
        ) -> LeafState:
            output = callback(time, state, *inputs, **parameters)
            return state.with_cached_value(cache_index, output)

        _update_callback = self.wrap_callback(
            _update_callback, collect_inputs=requires_inputs
        )

        # Create the associated update event
        event = DiscreteUpdateEvent(
            system_id=self.system_id,
            event_data=PeriodicEventData(
                period=period, offset=offset, active=False
            ),
            name=f"{self.name}:output_{cache_index}",
            callback=_update_callback,
            passthrough=self._passthrough,
        )

        # Note that in this case the "prerequisites of calc" will correspond to the
        # prerequisites of the update event, not the literal output callback itself.
        # However, these can be used to determine dependencies for the update event
        # via the output port.

    super().configure_output_port(
        port_index,
        _output_callback,
        prerequisites_of_calc=prerequisites_of_calc,
        default_value=default_value,
        event=event,
        cache_index=cache_index,
    )

`configure_periodic_update(event_index, callback, period, offset, enable_tracing=None)`

Configure an existing periodic update event.

The event will be triggered at regular intervals defined by the period and offset parameters. The callback should have the signature callback(time, state, *inputs, **params) -> xd_plus, where xd_plus is the updated value of the discrete state.

This callback should be written to compute the "plus" value of the discrete state component given the "minus" values of all state components and inputs.

Parameters:

Name	Type	Description	Default
`event_index`	`int`	The index of the event to configure.	required
`callback`	`Callable`	The callback function defining the update.	required
`period`	`Scalar`	The period at which the update event occurs.	required
`offset`	`Scalar`	The offset at which the first occurrence of the event is triggered.	required
`enable_tracing`	`bool`	If True, enable tracing for this event. Defaults to None.	`None`

Source code in jaxonomy/framework/leaf_system.py

def configure_periodic_update(
    self,
    event_index: int,
    callback: Callable,
    period: Scalar | Parameter,
    offset: Scalar | Parameter,
    enable_tracing: bool = None,
):
    """Configure an existing periodic update event.

    The event will be triggered at regular intervals defined by the period and
    offset parameters. The callback should have the signature
    `callback(time, state, *inputs, **params) -> xd_plus`, where `xd_plus` is the
    updated value of the discrete state.

    This callback should be written to compute the "plus" value of the discrete
    state component given the "minus" values of all state components and inputs.

    Args:
        event_index (int):
            The index of the event to configure.
        callback (Callable):
            The callback function defining the update.
        period (Scalar):
            The period at which the update event occurs.
        offset (Scalar):
            The offset at which the first occurrence of the event is triggered.
        enable_tracing (bool, optional):
            If True, enable tracing for this event. Defaults to None.
    """
    _wrapped_callback = self.wrap_callback(callback)

    def _callback(context: ContextBase) -> LeafState:
        xd = _wrapped_callback(context)
        return context[self.system_id].state.with_discrete_state(xd)

    if enable_tracing is None:
        enable_tracing = True

    event = DiscreteUpdateEvent(
        system_id=self.system_id,
        name=f"{self.name}:periodic_update",
        event_data=PeriodicEventData(period=period, offset=offset, active=False),
        callback=_callback,
        passthrough=self._passthrough,
        enable_tracing=enable_tracing,
        is_state_update=True,
    )
    self._state_update_events[event_index] = event

`declare_cache(callback, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=True)`

Declare a stored computation for the system.

This method accepts a callback function with the block-level signature callback(time, state, *inputs, **parameters) -> value and wraps it to have the signature callback(context) -> value

This callback can optionally be used to define a periodic update event that refreshes the cached value. Other calculations (e.g. sample-and-hold output ports) can then depend on the cached value.

Parameters:

Name	Type	Description	Default
`callback`	`Callable`	The callback function defining the cached computation.	required
`period`	`float`	If not None, the callback function will be used to define a periodic update event that refreshes the value. Defaults to None.	`None`
`offset`	`float`	The offset of the periodic update event. Defaults to 0.0. Will be ignored unless `period` is not None.	`0.0`
`name`	`str`	The name of the cached value. Defaults to None.	`None`
`default_value`	`Array`	The default value of the result, if known. Defaults to None.	`None`
`requires_inputs`	`bool`	If True, the callback will eval input ports to gather input values. This will add a bit to compile time, so setting to False where possible is recommended. Defaults to True.	`True`
`prerequisites_of_calc`	`List[DependencyTicket]`	The dependency tickets for the computation. Defaults to None, in which case the default is to assume dependency on either (inputs) if `requires_inputs` is True, or (nothing) otherwise.	`None`

Returns:

Name	Type	Description
`int`	`int`	The index of the callback in `system.callbacks`. The cache index can recovered from `system.callbacks[callback_index].cache_index`.

Source code in jaxonomy/framework/leaf_system.py

def declare_cache(
    self,
    callback: Callable,
    period: float | Parameter = None,
    offset: float | Parameter = 0.0,
    name: str = None,
    prerequisites_of_calc: List[DependencyTicket] = None,
    default_value: Array = None,
    requires_inputs: bool = True,
) -> int:
    """Declare a stored computation for the system.

    This method accepts a callback function with the block-level signature
        `callback(time, state, *inputs, **parameters) -> value`
    and wraps it to have the signature
        `callback(context) -> value`

    This callback can optionally be used to define a periodic update event that
    refreshes the cached value.  Other calculations (e.g. sample-and-hold output
    ports) can then depend on the cached value.

    Args:
        callback (Callable):
            The callback function defining the cached computation.
        period (float, optional):
            If not None, the callback function will be used to define a periodic
            update event that refreshes the value. Defaults to None.
        offset (float, optional):
            The offset of the periodic update event. Defaults to 0.0.  Will be ignored
            unless `period` is not None.
        name (str, optional):
            The name of the cached value. Defaults to None.
        default_value (Array, optional):
            The default value of the result, if known. Defaults to None.
        requires_inputs (bool, optional):
            If True, the callback will eval input ports to gather input values.
            This will add a bit to compile time, so setting to False where possible
            is recommended. Defaults to True.
        prerequisites_of_calc (List[DependencyTicket], optional):
            The dependency tickets for the computation. Defaults to None, in which
            case the default is to assume dependency on either (inputs) if
            `requires_inputs` is True, or (nothing) otherwise.

    Returns:
        int: The index of the callback in `system.callbacks`.  The cache index can
            recovered from `system.callbacks[callback_index].cache_index`.
    """
    # The index in the list of system callbacks
    callback_index = len(self.callbacks)

    # This is the index that this cached value will have in state.cache
    cache_index = len(self._default_cache)
    self._default_cache.append(default_value)

    # To help avoid unnecessary flagging of algebraic loops, trim the inputs as a
    # default prereq if the update callback doesn't use them
    if prerequisites_of_calc is None:
        if requires_inputs:
            prerequisites_of_calc = [DependencyTicket.u]
        else:
            prerequisites_of_calc = [DependencyTicket.nothing]

    def _update_callback(
        time: Scalar, state: LeafState, *inputs, **parameters
    ) -> LeafState:
        output = callback(time, state, *inputs, **parameters)
        return state.with_cached_value(cache_index, output)

    _update_callback = self.wrap_callback(
        _update_callback, collect_inputs=requires_inputs
    )

    if period is None:
        event = None

    else:
        # The cache has a periodic event updating its value defined by the callback
        event = DiscreteUpdateEvent(
            system_id=self.system_id,
            event_data=PeriodicEventData(
                period=period, offset=offset, active=False
            ),
            name=f"{self.name}:cache_update_{cache_index}_",
            callback=_update_callback,
            passthrough=self._passthrough,
        )

    if name is None:
        name = f"cache_{cache_index}"

    sys_callback = SystemCallback(
        callback=_update_callback,
        system=self,
        callback_index=callback_index,
        name=name,
        prerequisites_of_calc=prerequisites_of_calc,
        event=event,
        default_value=default_value,
        cache_index=cache_index,
    )
    self.callbacks.append(sys_callback)

    return callback_index

`declare_continuous_state(shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None, substeps=1, project=None)`

Declare a continuous state component for the system.

The continuous state value is read inside callbacks as state.continuous_state (the state argument of the ode / output callbacks). Unpack contract (T-C3-followup): the shape of state.continuous_state mirrors exactly what you passed as default_value (or the zeros array implied by shape / dtype):

A scalar default (jnp.array(0.0)) gives a scalar state.continuous_state — read it directly, do not index.
A vector default (jnp.zeros(3)) gives a length-3 array — index / unpack as x, y, z = state.continuous_state or state.continuous_state[i].
A PyTree default (tuple / NamedTuple / dict) gives back the same PyTree structure; your ode must return xcdot with the identical structure.

The ode callback's return value must match the default_value structure element-for-element, since it is added to the state during integration. A common error is declaring a scalar state but returning jnp.array([xdot]) (shape (1,)) from the ode — keep both scalar or both vector.

Multirate substepping (T-133): substeps=N declares that this block's continuous dynamics have a fast time constant needing N inner integration steps per outer solver step (e.g. a motor's electrical winding inside a 1 kHz control loop). Honored by the fixed-step rk4 solver (SimulatorOptions(ode_solver_method= "rk4")): the block's states advance with N RK4 substeps of h/N while the rest of the diagram takes one step of h, with first-order (zero-order-hold) coupling at the boundary — each side sees the other's start-of-step values, matching the semantics of a hand-rolled JIT-safe substep loop. Adaptive solvers (dopri5/bdf) ignore the declaration — they control stiffness through global step adaptation. N must be a static Python int >= 1; the default 1 is byte-equivalent to the pre-T-133 behavior.

Reverse-mode autodiff (enable_autodiff=True) is supported — the substep loop has a static trip count and the checkpointed adjoint substeps the costates alongside their primals. Gradient accuracy carries the scheme's first-order coupling error: the adjoint converges to the true sensitivity linearly in the outer step h (exact FD agreement is only recovered as h is refined), and for dynamics unstable at the outer step the adjoint's reverse-time primal re-integration further limits accuracy. Reduce the outer step when gradients through the coupling interface need to be tight.

Declared state projection (T-132): project=fn declares that this block's continuous state lives on a manifold and supplies the retraction back onto it — e.g. unit-quaternion renormalization for an attitude state (nq=4 integrated componentwise drifts off the unit sphere under any one-step integrator). fn(x) -> x receives the state in its declared structure, must be shape-preserving and jit-safe, and is applied by the simulator at the end of every major step (composing with, and independent of, the T-003a DAE projection). Within-step drift is bounded by the step size; the recorded trajectory and all values other blocks see at major-step boundaries are on the manifold. Differentiable: the projection participates in reverse-mode AD as ordinary traced ops.

Source code in jaxonomy/framework/leaf_system.py

def declare_continuous_state(
    self,
    shape: ShapeLike = None,
    default_value: Array = None,
    dtype: DTypeLike = None,
    ode: Callable = None,
    mass_matrix: Array = None,
    as_array: bool = True,
    requires_inputs: bool = True,
    prerequisites_of_calc: List[DependencyTicket] = None,
    substeps: int = 1,
    project: Callable = None,
):
    """Declare a continuous state component for the system.

    The continuous state value is read inside callbacks as
    ``state.continuous_state`` (the ``state`` argument of the ``ode`` /
    output callbacks). **Unpack contract** (T-C3-followup): the shape of
    ``state.continuous_state`` mirrors exactly what you passed as
    ``default_value`` (or the zeros array implied by ``shape`` /
    ``dtype``):

    - A scalar default (``jnp.array(0.0)``) gives a scalar
      ``state.continuous_state`` — read it directly, do **not** index.
    - A vector default (``jnp.zeros(3)``) gives a length-3 array — index
      / unpack as ``x, y, z = state.continuous_state`` or
      ``state.continuous_state[i]``.
    - A PyTree default (tuple / NamedTuple / dict) gives back the same
      PyTree structure; your ``ode`` must return ``xcdot`` with the
      identical structure.

    The ``ode`` callback's return value must match the
    ``default_value`` structure element-for-element, since it is added to
    the state during integration. A common error is declaring a scalar
    state but returning ``jnp.array([xdot])`` (shape ``(1,)``) from the
    ode — keep both scalar or both vector.

    Multirate substepping (T-133): ``substeps=N`` declares that this
    block's continuous dynamics have a fast time constant needing ``N``
    inner integration steps per outer solver step (e.g. a motor's
    electrical winding inside a 1 kHz control loop). Honored by the
    fixed-step ``rk4`` solver (``SimulatorOptions(ode_solver_method=
    "rk4")``): the block's states advance with ``N`` RK4 substeps of
    ``h/N`` while the rest of the diagram takes one step of ``h``,
    with first-order (zero-order-hold) coupling at the boundary —
    each side sees the other's start-of-step values, matching the
    semantics of a hand-rolled JIT-safe substep loop. Adaptive solvers
    (``dopri5``/``bdf``) ignore the declaration — they control
    stiffness through global step adaptation. ``N`` must be a static
    Python ``int >= 1``; the default 1 is byte-equivalent to the
    pre-T-133 behavior.

    Reverse-mode autodiff (``enable_autodiff=True``) is supported —
    the substep loop has a static trip count and the checkpointed
    adjoint substeps the costates alongside their primals. Gradient
    accuracy carries the scheme's first-order coupling error: the
    adjoint converges to the true sensitivity linearly in the outer
    step ``h`` (exact FD agreement is only recovered as ``h`` is
    refined), and for dynamics *unstable at the outer step* the
    adjoint's reverse-time primal re-integration further limits
    accuracy. Reduce the outer step when gradients through the
    coupling interface need to be tight.

    Declared state projection (T-132): ``project=fn`` declares that
    this block's continuous state lives on a manifold and supplies
    the retraction back onto it — e.g. unit-quaternion
    renormalization for an attitude state (``nq=4`` integrated
    componentwise drifts off the unit sphere under any one-step
    integrator). ``fn(x) -> x`` receives the state in its declared
    structure, must be shape-preserving and jit-safe, and is applied
    by the simulator **at the end of every major step** (composing
    with, and independent of, the T-003a DAE projection). Within-step
    drift is bounded by the step size; the recorded trajectory and
    all values other blocks see at major-step boundaries are on the
    manifold. Differentiable: the projection participates in
    reverse-mode AD as ordinary traced ops.
    """
    if not isinstance(substeps, (int, np.integer)) or isinstance(
        substeps, bool
    ) or substeps < 1:
        raise ValueError(
            f"declare_continuous_state: substeps must be a static Python "
            f"int >= 1, got {substeps!r}. (It sets a compile-time inner "
            "loop count and cannot be traced or fractional.)"
        )
    self._continuous_substeps = int(substeps)
    if project is not None and not callable(project):
        raise ValueError(
            f"declare_continuous_state: project must be a callable "
            f"x -> x (shape-preserving, jit-safe), got {project!r}."
        )
    self._continuous_projection = project

    self.ode_callback = SystemCallback(
        callback=None,
        system=self,
        callback_index=len(self.callbacks),
        name=f"{self.name}_ode",
        prerequisites_of_calc=prerequisites_of_calc,
    )
    self.callbacks.append(self.ode_callback)
    callback_idx = len(self.callbacks) - 1

    # FIXME: this is to preserve some backward compatibility while we decouple
    # declaration from configuration. Declaration should not have to call
    # configuration.
    if default_value is not None or shape is not None:
        self.configure_continuous_state(
            callback_idx,
            shape=shape,
            default_value=default_value,
            dtype=dtype,
            ode=ode,
            mass_matrix=mass_matrix,
            as_array=as_array,
            requires_inputs=requires_inputs,
            prerequisites_of_calc=prerequisites_of_calc,
        )

    return callback_idx

`declare_continuous_state_output(name=None)`

Declare a continuous state output port in the system.

This method creates a new block-level output port which returns the full continuous state of the system.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the output port. Defaults to None (autogenerate name).	`None`

Returns:

Name	Type	Description
`int`	`int`	The index of the new output port.

Source code in jaxonomy/framework/leaf_system.py

def declare_continuous_state_output(
    self,
    name: str = None,
) -> int:
    """Declare a continuous state output port in the system.

    This method creates a new block-level output port which returns the full
    continuous state of the system.

    Args:
        name (str, optional):
            The name of the output port. Defaults to None (autogenerate name).

    Returns:
        int: The index of the new output port.
    """
    if self._continuous_state_output_port_idx is not None:
        raise ValueError("Continuous state output port already declared")

    def _callback(time: Scalar, state: LeafState, *inputs, **parameters):
        return state.continuous_state

    self._continuous_state_output_port_idx = self.declare_output_port(
        _callback,
        name=name,
        prerequisites_of_calc=[DependencyTicket.xc],
        default_value=self._default_continuous_state,
        requires_inputs=False,
    )
    return self._continuous_state_output_port_idx

`declare_discrete_state(shape=None, default_value=None, dtype=None, as_array=True, name=None)`

Declare a discrete state component for the system.

The discrete state is a component of the system's state that can be updated at specific events, such as zero-crossings or periodic updates.

.. note:: Currently only one discrete state component is supported per LeafSystem. If declare_discrete_state is called more than once, the second call will silently overwrite the first. To store several independent values, pack them into a single array and split inside your update callback.

Parameters:

Name	Type	Description	Default
`shape`	`ShapeLike`	The shape of the discrete state. Defaults to None.	`None`
`default_value`	`Array`	The initial value of the discrete state. Defaults to None.	`None`
`dtype`	`DTypeLike`	The data type of the discrete state. Defaults to None.	`None`
`as_array`	`bool`	If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification.	`True`
`name`	`str`	Readability label for the discrete state (parity with `declare_continuous_state_output(name=...)`). Stored as `self.discrete_state_name` for diagnostics/debugging; it does not change runtime behaviour, and the state is still read as `state.discrete_state`.	`None`

Raises:

Type	Description
`AssertionError`	If as_array is True and neither shape nor default_value is provided.

Notes

(1) Only one of shape and default_value should be provided. If default_value is provided, it will be used as the initial value of the continuous state. If shape is provided, the initial value will be a zero vector of the given shape and specified dtype.

(2) Use declare_periodic_update to declare an update event that modifies the discrete state at a recurring interval.

Source code in jaxonomy/framework/leaf_system.py

def declare_discrete_state(
    self,
    shape: ShapeLike = None,
    default_value: Array | Parameter = None,
    dtype: DTypeLike = None,
    as_array: bool = True,
    name: str = None,
):
    """Declare a discrete state component for the system.

    The discrete state is a component of the system's state that can be updated
    at specific events, such as zero-crossings or periodic updates.

    .. note::
        Currently only **one** discrete state component is supported per
        ``LeafSystem``.  If ``declare_discrete_state`` is called more than once,
        the second call will silently overwrite the first.  To store several
        independent values, pack them into a single array and split inside your
        update callback.

    Args:
        shape (ShapeLike, optional):
            The shape of the discrete state. Defaults to None.
        default_value (Array, optional):
            The initial value of the discrete state. Defaults to None.
        dtype (DTypeLike, optional):
            The data type of the discrete state. Defaults to None.
        as_array (bool, optional):
            If True, treat the default_value as an array-like (cast if necessary).
            Otherwise, it will be stored as the default state without modification.
        name (str, optional):
            Readability label for the discrete state (parity with
            ``declare_continuous_state_output(name=...)``). Stored as
            ``self.discrete_state_name`` for diagnostics/debugging; it
            does not change runtime behaviour, and the state is still
            read as ``state.discrete_state``.

    Raises:
        AssertionError:
            If as_array is True and neither shape nor default_value is provided.

    Notes:
        (1) Only one of `shape` and `default_value` should be provided. If
        `default_value` is provided, it will be used as the initial value of the
        continuous state. If `shape` is provided, the initial value will be a
        zero vector of the given shape and specified dtype.

        (2) Use `declare_periodic_update` to declare an update event that
        modifies the discrete state at a recurring interval.
    """
    self.discrete_state_name = name
    if as_array:
        default_value = utils.make_array(default_value, dtype=dtype, shape=shape)

    # Tree-map the default value to ensure that it is an array-like with the
    # correct shape and dtype. This is necessary because the default value
    # may be a list, tuple, or other PyTree-structured object.
    default_value = tree_util.tree_map(npa.asarray, default_value)

    # T-027a: if Zeno protection is already installed, pack the user's
    # value alongside the existing Zeno tracker rather than overwriting it.
    if self._zeno_protected_events:
        self._zeno_user_default = default_value
        current = self._default_discrete_state
        zeno_xd = (
            current.zeno
            if isinstance(current, self._zeno_combined_type)
            else current
        )
        self._default_discrete_state = self._zeno_combined_type(
            user=default_value, zeno=zeno_xd
        )
    else:
        self._default_discrete_state = default_value

`declare_mode_output(name=None)`

Declare a mode output port in the system.

This method creates a new block-level output port which returns the component of the system's state corresponding to the discrete "mode" or "stage".

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the output port. Defaults to None.	`None`

Returns:

Name	Type	Description
`int`	`int`	The index of the declared mode output port.

Source code in jaxonomy/framework/leaf_system.py

def declare_mode_output(self, name: str = None) -> int:
    """Declare a mode output port in the system.

    This method creates a new block-level output port which returns the component
    of the system's state corresponding to the discrete "mode" or "stage".

    Args:
        name (str, optional):
            The name of the output port. Defaults to None.

    Returns:
        int:
            The index of the declared mode output port.
    """

    def _callback(time: Scalar, state: LeafState, *inputs, **parameters):
        return state.mode

    self._mode_output_port_idx = self.declare_output_port(
        _callback,
        name=name,
        prerequisites_of_calc=[DependencyTicket.mode],
        default_value=self._default_mode,
        requires_inputs=False,
    )

    return self._mode_output_port_idx

`declare_output_port(callback=None, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=None, units=None)`

Declare an output port in the LeafSystem.

This method accepts a callback function with the block-level signature callback(time, state, *inputs, **parameters) -> value and wraps it to the signature expected by SystemBase.declare_output_port: callback(context) -> value

Parameters:

Name	Type	Description	Default
`callback`	`Callable`	The callback function defining the output port.	`None`
`period`	`float`	If not None, the port will act as a "sample-and-hold", with the callback function used to define a periodic update event that refreshes the value that will be returned by the port. Typically this should match the update period of some associated update event in the system. Defaults to None.	`None`
`offset`	`float`	The offset of the periodic update event. Defaults to 0.0. Will be ignored unless `period` is not None.	`0.0`
`name`	`str`	The name of the output port. Defaults to None.	`None`
`default_value`	`Array`	The default value of the output port, if known. Defaults to None.	`None`
`requires_inputs`	`bool \| list[int] \| None`	Whether the callback reads input port values. Defaults to `None`. `None` resolves to `True` (collect all inputs) in every case except the unambiguous `prerequisites_of_calc=[DependencyTicket.nothing]` declaration, which resolves to `False` (T-A4-followup-requires-inputs-infer). The inference is deliberately conservative: `prerequisites_of_calc` may list upstream / transitive tickets (e.g. `xcdot` for a derivative output whose callback still reads `u`, or `xd` for a sample-and-hold port whose update event reads `u`), so a non-input prereq list does not imply the callback is input-free — only `[nothing]` does. Set this to `False` explicitly whenever the output does NOT depend on any input port (e.g. a ZOH output that returns a stored discrete state, or a CT output that only reads continuous state). This serves two purposes: 1. Eliminates false-positive algebraic-loop detection. The diagram-level algebraic-loop checker conservatively assumes every output with `requires_inputs=True` has direct feedthrough from all connected inputs. Declaring `requires_inputs=False` tells the checker there is no feedthrough from inputs to this output, which is required to break apparent cycles in discrete feedback topologies (A→B→A) that are valid because updates use x⁻. 2. Reduces compile time by avoiding unnecessary input collection. Can also be specified as a list of integer port indices to declare selective feedthrough (only the listed inputs feed through to this output). Defaults to `True` (collect all inputs, assume full feedthrough).	`None`
`prerequisites_of_calc`	`List[DependencyTicket]`	The dependency tickets for the output port computation. Defaults to None, in which case the assumption is a dependency on either (nothing) if `requires_inputs` is False otherwise (inputs).	`None`

Returns:

Name	Type	Description
`int`	`int`	The index of the declared output port.

Source code in jaxonomy/framework/leaf_system.py

def declare_output_port(
    self,
    callback: Callable = None,
    period: float = None,
    offset: float = 0.0,
    name: str = None,
    prerequisites_of_calc: List[DependencyTicket] = None,
    default_value: Array = None,
    requires_inputs: bool | list[int] | None = None,
    units=None,
) -> int:
    """Declare an output port in the LeafSystem.

    This method accepts a callback function with the block-level signature
        `callback(time, state, *inputs, **parameters) -> value`
    and wraps it to the signature expected by SystemBase.declare_output_port:
        `callback(context) -> value`

    Args:
        callback (Callable):
            The callback function defining the output port.
        period (float, optional):
            If not None, the port will act as a "sample-and-hold", with the
            callback function used to define a periodic update event that refreshes
            the value that will be returned by the port. Typically this should
            match the update period of some associated update event in the system.
            Defaults to None.
        offset (float, optional):
            The offset of the periodic update event. Defaults to 0.0.  Will be ignored
            unless `period` is not None.
        name (str, optional):
            The name of the output port. Defaults to None.
        default_value (Array, optional):
            The default value of the output port, if known. Defaults to None.
        requires_inputs (bool | list[int] | None, optional):
            Whether the callback reads input port values.

            **Defaults to ``None``.** ``None`` resolves to ``True``
            (collect all inputs) in every case except the unambiguous
            ``prerequisites_of_calc=[DependencyTicket.nothing]``
            declaration, which resolves to ``False``
            (T-A4-followup-requires-inputs-infer). The inference is
            deliberately conservative: ``prerequisites_of_calc`` may list
            *upstream / transitive* tickets (e.g. ``xcdot`` for a
            derivative output whose callback still reads ``u``, or ``xd``
            for a sample-and-hold port whose *update* event reads ``u``),
            so a non-input prereq list does **not** imply the callback is
            input-free — only ``[nothing]`` does.

            **Set this to ``False`` explicitly whenever the output does NOT
            depend on any input port** (e.g. a ZOH output that returns a
            stored discrete state, or a CT output that only reads continuous
            state).  This serves two purposes:
              1. **Eliminates false-positive algebraic-loop detection.**  The
                 diagram-level algebraic-loop checker conservatively assumes every
                 output with ``requires_inputs=True`` has direct feedthrough from
                 all connected inputs.  Declaring ``requires_inputs=False`` tells
                 the checker there is no feedthrough from inputs to this output,
                 which is required to break apparent cycles in discrete feedback
                 topologies (A→B→A) that are valid because updates use x⁻.
              2. **Reduces compile time** by avoiding unnecessary input collection.

            Can also be specified as a list of integer port indices to declare
            selective feedthrough (only the listed inputs feed through to this
            output).  Defaults to ``True`` (collect all inputs, assume full
            feedthrough).
        prerequisites_of_calc (List[DependencyTicket], optional):
            The dependency tickets for the output port computation.  Defaults to
            None, in which case the assumption is a dependency on either (nothing)
            if `requires_inputs` is False otherwise (inputs).

    Returns:
        int: The index of the declared output port.
    """

    # T-A4-followup-requires-inputs-infer: when the caller leaves
    # ``requires_inputs`` unset (None) but supplies ``prerequisites_of_calc``,
    # infer whether inputs are needed from the prerequisites rather than
    # forcing the user to keep the two arguments consistent by hand.
    requires_inputs = self._resolve_requires_inputs(
        requires_inputs, prerequisites_of_calc
    )

    if default_value is not None:
        default_value = npa.array(default_value)

    cache_index = None
    if period is not None:
        # The output port will be of "sample-and-hold" type, so we have to declare a
        # periodic event to update the value.  The callback will be used to define the
        # update event, and the output callback will simply return the stored value.

        # This is the index that this port value will have in state.cache
        cache_index = len(self._default_cache)
        self._default_cache.append(default_value)

    output_port_idx = super().declare_output_port(
        callback, name=name, cache_index=cache_index, units=units
    )

    self.configure_output_port(
        output_port_idx,
        callback,
        period=period,
        offset=offset,
        prerequisites_of_calc=prerequisites_of_calc,
        default_value=default_value,
        requires_inputs=requires_inputs,
    )

    return output_port_idx

`declare_zero_crossing(guard, reset_map=None, start_mode=None, end_mode=None, direction='crosses_zero', terminal=False, name=None, enable_tracing=None, zeno_tolerance=None, grad_guard=None)`

Declare an event triggered by a zero-crossing of a guard function.

Optionally, the system can also transition between discrete modes If start_mode and end_mode are specified, the system will transition from start_mode to end_mode when the event is triggered according to guard. This event will be active conditionally on state.mode == start_mode and when triggered will result in applying the reset map. In addition, the mode will be updated to end_mode.

If start_mode and end_mode are not specified, the event will always be active and will not result in a mode transition.

The guard function should have the signature

guard(time, state, *inputs, **parameters) -> float

and the reset map should have the signature of an unrestricted update

reset_map(time, state, *inputs, **parameters) -> state

Parameters:

Name	Type	Description	Default
`guard`	`Callable`	The guard function which triggers updates on zero crossing.	required
`reset_map`	`Callable`	The reset map which is applied when the event is triggered. If None (default), no reset is applied.	`None`
`start_mode`	`int`	The mode or stage of the system in which the guard will be actively monitored. If None (default), the event will always be active.	`None`
`end_mode`	`int`	The mode or stage of the system to which the system will transition when the event is triggered. If start_mode is None, this is ignored. Otherwise it must be specified, though it can be the same as start_mode.	`None`
`direction`	`str`	The direction of the zero crossing. Options are "crosses_zero" (default), "positive_then_non_positive", "negative_then_non_negative", and "edge_detection". All except edge detection operate on continuous signals; edge detection operates on boolean signals and looks for a jump from False to True or vice versa.	`'crosses_zero'`
`terminal`	`bool`	If True, the event will halt simulation if and when the zero-crossing occurs. If this event is triggered the reset map will still be applied as usual prior to termination. Defaults to False.	`False`
`name`	`str`	The name of the event. Defaults to None.	`None`
`enable_tracing`	`bool`	If True, enable tracing for this event. Defaults to None.	`None`

Notes

By default the system state does not have a "mode" component, so in order to declare "state transitions" with non-null start and end modes, the user must first call declare_default_mode to set the default mode to be some integer (initial condition for the system).

Source code in jaxonomy/framework/leaf_system.py

def declare_zero_crossing(
    self,
    guard: Callable,
    reset_map: Callable = None,
    start_mode: int = None,
    end_mode: int = None,
    direction: str = "crosses_zero",
    terminal: bool = False,
    name: str = None,
    enable_tracing: bool = None,
    zeno_tolerance: float | None = None,
    grad_guard: Callable = None,
):
    """Declare an event triggered by a zero-crossing of a guard function.

    Optionally, the system can also transition between discrete modes
    If `start_mode` and `end_mode` are specified, the system will transition
    from `start_mode` to `end_mode` when the event is triggered according to `guard`.
    This event will be active conditionally on `state.mode == start_mode` and when
    triggered will result in applying the reset map. In addition, the mode will be
    updated to `end_mode`.

    If `start_mode` and `end_mode` are not specified, the event will always be active
    and will not result in a mode transition.

    The guard function should have the signature:
        `guard(time, state, *inputs, **parameters) -> float`

    and the reset map should have the signature of an unrestricted update:
        `reset_map(time, state, *inputs, **parameters) -> state`

    Args:
        guard (Callable):
            The guard function which triggers updates on zero crossing.
        reset_map (Callable, optional):
            The reset map which is applied when the event is triggered. If None
            (default), no reset is applied.
        start_mode (int, optional):
            The mode or stage of the system in which the guard will be
            actively monitored. If None (default), the event will always be
            active.
        end_mode (int, optional):
            The mode or stage of the system to which the system will transition
            when the event is triggered. If start_mode is None, this is ignored.
            Otherwise it _must_ be specified, though it can be the same as
            start_mode.
        direction (str, optional):
            The direction of the zero crossing. Options are "crosses_zero"
            (default), "positive_then_non_positive", "negative_then_non_negative",
            and "edge_detection".  All except edge detection operate on continuous
            signals; edge detection operates on boolean signals and looks for a
            jump from False to True or vice versa.
        terminal (bool, optional):
            If True, the event will halt simulation if and when the zero-crossing
            occurs. If this event is triggered the reset map will still be applied
            as usual prior to termination. Defaults to False.
        name (str, optional):
            The name of the event. Defaults to None.
        enable_tracing (bool, optional):
            If True, enable tracing for this event. Defaults to None.

    Notes:
        By default the system state does not have a "mode" component, so in
        order to declare "state transitions" with non-null start and end modes,
        the user must first call `declare_default_mode` to set the default mode
        to be some integer (initial condition for the system).
    """

    logger.debug(
        f"Declaring transition for {self.name} with guard {guard} and reset map {reset_map}"
    )

    if enable_tracing is None:
        enable_tracing = True

    if start_mode is not None or end_mode is not None:
        assert (
            self._default_mode is not None
        ), "System has no mode: call `declare_default_mode` before transitions."
        assert isinstance(start_mode, int) and isinstance(end_mode, int)

    # T-027: optional Zeno-hold protection. If `zeno_tolerance` is set,
    # wrap the user's reset_map to flag a Zeno entry, declare a companion
    # exit event, and freeze the continuous-state ODE while held.
    if zeno_tolerance is not None:
        assert (
            isinstance(zeno_tolerance, (float, int)) and float(zeno_tolerance) > 0.0
        ), "zeno_tolerance must be a positive float"
        reset_map, _zeno_companion = self._install_zeno_protection(
            reset_map=reset_map,
            guard=guard,
            direction=direction,
            tol=float(zeno_tolerance),
            name=name,
        )
    else:
        _zeno_companion = None

    # Wrap the reset map with a mode update if necessary
    def _reset_and_update_mode(
        time: Scalar, state: LeafState, *inputs, **parameters
    ) -> LeafState:
        if reset_map is not None:
            state = reset_map(time, state, *inputs, **parameters)
        logger.debug(f"Updating mode from {state.mode} to {end_mode}")

        # If the start and end modes are declared, update the mode
        if start_mode is not None:
            logger.debug(f"Updating mode from {state.mode} to {end_mode}")
            state = state.with_mode(end_mode)

        return state

    _wrapped_guard = self.wrap_callback(guard)
    # Optional smooth guard residual for the event-time (saltation) gradient
    # only — wrapped the same way as the trigger guard.  ``None`` keeps the
    # legacy behaviour (the saltation paths fall back to ``guard``).
    _wrapped_grad_guard = (
        self.wrap_callback(grad_guard) if grad_guard is not None else None
    )
    _wrapped_reset = _wrap_reset_map(
        self, _reset_and_update_mode, _wrapped_guard, terminal,
        grad_guard=_wrapped_grad_guard,
    )

    event = ZeroCrossingEvent(
        system_id=self.system_id,
        guard=_wrapped_guard,
        grad_guard=_wrapped_grad_guard,
        reset_map=_wrapped_reset,
        passthrough=self._passthrough,
        direction=direction,
        is_terminal=terminal,
        name=name,
        event_data=ZeroCrossingEventData(active=True, triggered=False),
        enable_tracing=enable_tracing,
        active_mode=start_mode,
    )

    event_index = len(self._zero_crossing_events)
    self._zero_crossing_events.append(event)

    # T-115-followup-saturate-rate-classification: bump the
    # behavioral-ZC counter only when this event actually does
    # something on trigger — has a user-supplied reset map or
    # participates in a mode transition. Pure guard-only events
    # (Saturate / DeadZone clip boundaries) are solver hints with
    # no behavioral effect, so they should not flip the block's
    # rate-group classification to ``event_driven``.
    if (
        reset_map is not None
        or start_mode is not None
        or end_mode is not None
    ):
        self._n_behavioral_zc_events += 1

    # Record the transition in the transition map (for debugging or analysis)
    if start_mode is not None:
        if start_mode not in self.transition_map:
            self.transition_map[start_mode] = []
        self.transition_map[start_mode].append((event_index, event))

    # T-027: register the companion `_exit_zeno` event AFTER the main event
    # so the slot index is finalized first.
    if _zeno_companion is not None:
        _zeno_companion(event_index)

`initialize(**parameters)`

Hook for initializing a system. Called during context creation.

If the parameters are instances of Parameter, they will be resolved. If implemented, the function signature should contain all the declared parameters.

This function should not be called directly. It will be called implicitly after init with the resolved parameters.

Source code in jaxonomy/framework/leaf_system.py

def initialize(self, **parameters):
    """Hook for initializing a system. Called during context creation.

    If the parameters are instances of Parameter, they will be resolved.
    If implemented, the function signature should contain all the declared
    parameters.

    This function should not be called directly. It will be called implicitly
    after __init__ with the resolved parameters.
    """
    pass

`reset_default_values(**dynamic_parameters)`

This function is used to reset default values for continuous/discrete states, ports and mode based on dynamic parameters. It is called in create_state() and used to reset states in ensemble sims and optimization with the context method with_new_state().

Note that dtypes and shapes can't be changed after initialization because the diagram may already have been jax-compiled. Only values may change.

Source code in jaxonomy/framework/leaf_system.py

def reset_default_values(self, **dynamic_parameters):
    """This function is used to reset default values for
    continuous/discrete states, ports and mode based on dynamic parameters.
    It is called in `create_state()` and used to reset states in ensemble sims
    and optimization with the context method `with_new_state()`.

    Note that dtypes and shapes can't be changed after initialization because
    the diagram may already have been jax-compiled. Only values may change.
    """
    pass

`with_parameter(name, value)`

Return a copy of this system with one dynamic parameter replaced.

The returned system is a new instance. The original is unchanged.

Parameters:

Name	Type	Description	Default
`name`	`str`	Parameter name (must exist as a dynamic parameter).	required
`value`		New value (typically a JAX array for `jax.grad` / `jax.vmap`).	required

Raises:

Type	Description
`KeyError`	If `name` is not a dynamic parameter.
`TypeError`	If `name` is a static parameter.

Source code in jaxonomy/framework/leaf_system.py

def with_parameter(self, name: str, value) -> LeafSystem:
    """Return a copy of this system with one dynamic parameter replaced.

    The returned system is a new instance. The original is unchanged.

    Args:
        name: Parameter name (must exist as a dynamic parameter).
        value: New value (typically a JAX array for ``jax.grad`` / ``jax.vmap``).

    Raises:
        KeyError: If ``name`` is not a dynamic parameter.
        TypeError: If ``name`` is a static parameter.
    """
    if name in self._static_parameters:
        raise TypeError(
            f"Parameter {name!r} is static on {self.name!r}; static parameters "
            "cannot be replaced at runtime without recompilation."
        )
    if name not in self._dynamic_parameters:
        available = sorted(
            {*self._static_parameters.keys(), *self._dynamic_parameters.keys()}
        )
        raise KeyError(
            f"Parameter {name!r} is not a dynamic parameter on {self.name!r}. "
            f"Available: {available}"
        )

    old_param = self._dynamic_parameters[name]
    old_val = Parameter.unwrap(old_param)
    try:
        value = _check_values_compatible(old_val, value)
    except ValueError as e:
        raise ValueError(f"{e} (parameter {name!r} on {self.name!r})") from None

    new = copy.deepcopy(self)
    new.parent = None
    new._dependency_graph = None
    new.feedthrough_pairs = None
    new._cache_update_events = None
    new._cached_input_ports.clear()
    new._cached_output_ports.clear()

    if isinstance(old_param, Parameter):
        new._dynamic_parameters[name] = dataclasses.replace(
            old_param,
            value=value,
            name=name,
            system=new,
        )
    else:
        new._dynamic_parameters[name] = Parameter(
            value=value,
            name=name,
            system=new,
        )

    return new

`wrap_callback(callback, collect_inputs=True)`

Wrap an update function to unpack local variables and block inputs.

The callback should have the signature callback(time, state, *inputs, **params) -> result and will be wrapped to have the signature callback(context) -> result, as expected by the event handling logic.

This is used internally for declaration methods like declare_periodic_update so that users can write more intuitive block-level update functions without worrying about the "context", and have them automatically wrapped to have the right interface. It can also be called directly by users to wrap their own update functions, for example to create a callback function for declare_output_port.

The context and state are strictly immutable, so the callback should not attempt to change any values in the context or state. Even in cases where it is impossible to enforce this (e.g. a state component is a list, which is always mutable in Python), the callback should be careful to avoid direct modification of the context or state, which may lead to unexpected behavior or JAX tracer errors.

Parameters:

Name	Type	Description	Default
`callback`	`Callable`	The (pure) function to be wrapped. See above for expected signature.	required
`collect_inputs`	`bool`	If True, the callback will eval input ports to gather input values. Normally this should be True, but it can be set to False if the return value depends only on the state but not inputs, for instance. This helps reduce the number of expressions that need to be JIT compiled. Can also be specified as a list of integer port indices. Default is True (collect all inputs).	`True`

Returns:

Name	Type	Description
`Callable`	`Callable`	The wrapped function, with signature `callback(context) -> result`.

Source code in jaxonomy/framework/leaf_system.py

def wrap_callback(
    self, callback: Callable, collect_inputs: bool | list[int] = True
) -> Callable:
    """Wrap an update function to unpack local variables and block inputs.

    The callback should have the signature
    `callback(time, state, *inputs, **params) -> result`
    and will be wrapped to have the signature `callback(context) -> result`,
    as expected by the event handling logic.

    This is used internally for declaration methods like
    `declare_periodic_update` so that users can write more intuitive
    block-level update functions without worrying about the "context", and have
    them automatically wrapped to have the right interface.  It can also be
    called directly by users to wrap their own update functions, for example to
    create a callback function for `declare_output_port`.

    The context and state are strictly immutable, so the callback should not
    attempt to change any values in the context or state.  Even in cases where
    it is impossible to _enforce_ this (e.g. a state component is a list, which
    is always mutable in Python), the callback should be careful to avoid direct
    modification of the context or state, which may lead to unexpected behavior
    or JAX tracer errors.

    Args:
        callback (Callable):
            The (pure) function to be wrapped. See above for expected signature.
        collect_inputs (bool):
            If True, the callback will eval input ports to gather input values.
            Normally this should be True, but it can be set to False if the
            return value depends only on the state but not inputs, for
            instance. This helps reduce the number of expressions that need to
            be JIT compiled. Can also be specified as a list of integer port indices.
            Default is True (collect all inputs).

    Returns:
        Callable:
            The wrapped function, with signature `callback(context) -> result`.
    """
    return partial(
        _wrap_leaf_user_callback,
        owner=self,
        user_callback=callback,
        collect_inputs=collect_inputs,
    )

`Parameter` `dataclass`

Source code in jaxonomy/framework/parameter.py

@dataclasses.dataclass
class Parameter:
    value: Union[ParameterExpr, "Parameter", ArrayLike, str, tuple]

    # shape & dtype are set at init time when constructing the parameter,
    # they are not necessarily the actual value's shape and dtype
    dtype: DTypeLike = None
    shape: ShapeLike = None
    as_array: bool = False

    # name is used by reference submodels, model parameters and init script
    # variables so that they can be referred to in other fields
    # (we need this for serialization).
    name: str = None

    # For complex parameter values, we can specify a Python expression as string
    # This is useful for expressions like "np.eye(p)" where p is a parameter.
    is_python_expr: bool = False
    py_namespace: dict = None

    is_static: bool = False  # TODO: staticness should be propagated to dependents
    system: "SystemBase" = None

    # T-038a — opt-in per-parameter dtype hint.  When non-None and the value
    # is array-like (`np.ndarray` / `jax.Array`), the value is cast to this
    # dtype at construction time.  This is the metadata foundation for the
    # per-block dtype override mechanism; downstream block-side code may also
    # read ``_dtype_hint`` to allocate compatible buffers.  Default ``None``
    # is byte-equivalent to the pre-T-038a behavior.
    _dtype_hint: DTypeLike = None

    def get(self):
        value = ParameterCache.get(self)
        if self.as_array and not isinstance(value, Array):
            value = utils.make_array(value, self.dtype, self.shape)
        return value

    def set(self, value: Union["Parameter", ArrayLike, str, tuple]):
        ParameterCache.replace(self, value)

    @property
    def static_dependents(self):
        return ParameterCache.static_dependents(self)

    @property
    def is_dirty(self):
        return ParameterCache.__is_dirty__[self]

    @classmethod
    def unwrap(cls, value):
        """Get the underlying value of raw arrays and Parameter objects alike."""
        if value is None:
            return None
        if isinstance(value, (Array, bool, int, float, complex)):
            return value
        if isinstance(value, (np.ndarray, np.number)):
            if np.issubdtype(value.dtype, np.number):
                return value
            if value.shape == ():
                return Parameter.unwrap(value.item())
            return Parameter(value).get()
        if isinstance(value, Parameter):
            return value.get()
        if isinstance(value, list):
            return [cls.unwrap(val) for val in value]
        if isinstance(value, tuple):
            return tuple(cls.unwrap(val) for val in value)
        if isinstance(value, dict):
            return {key: cls.unwrap(val) for key, val in value.items()}
        # Fallback for unhandled types: forward to __compute__
        return Parameter(value).get()

    def __post_init__(self):
        ParameterCache._register(self)  # thread-safe registration

        # T-038a — apply optional dtype hint to array-like values.  Skipped
        # when the hint is None (the common path) so this remains a no-op
        # for every existing caller.  Only triggers on concrete array values
        # — ParameterExpr / strings / nested Parameters are left untouched
        # so the resolved value picks up its dtype downstream.
        if self._dtype_hint is not None and isinstance(
            self.value, (Array, np.ndarray)
        ):
            self.value = jnp.asarray(self.value, dtype=self._dtype_hint)

        if isinstance(self.value, Parameter):
            ParameterCache.add_dependent(self.value, self)
        if isinstance(self.value, ParameterExpr):
            for val in self.value:
                if isinstance(val, Parameter):
                    ParameterCache.add_dependent(val, self)
        if isinstance(self.value, (list, tuple)):
            _add_dependents(self.value, self)
        if self.is_python_expr and isinstance(self.value, str) and self.py_namespace:
            # A string expression like "k" or "np.eye(p)" depends on the
            # Parameter objects it names in its evaluation scope. Registering
            # them here (not only at deserialization time) means the links
            # survive deepcopy — __deepcopy__ re-runs __post_init__ — so
            # set() on a copied alias still invalidates copied referencing
            # blocks (T-141).
            for dep in _expr_parameter_refs(self.value, self.py_namespace):
                ParameterCache.add_dependent(dep, self)

        _record_parameter_creation(self)

    def __setstate__(self, state):
        self.__dict__.update(state)
        ParameterCache._register(self)
        if isinstance(self.value, Parameter):
            ParameterCache.add_dependent(self.value, self)
        if isinstance(self.value, ParameterExpr):
            for val in self.value:
                if isinstance(val, Parameter):
                    ParameterCache.add_dependent(val, self)
        if isinstance(self.value, (list, tuple)):
            _add_dependents(self.value, self)
        if self.is_python_expr and isinstance(self.value, str) and self.py_namespace:
            for dep in _expr_parameter_refs(self.value, self.py_namespace):
                ParameterCache.add_dependent(dep, self)

    def __deepcopy__(self, memo):
        """Copy fields and re-run post-init so :class:`ParameterCache` bookkeeping matches."""
        cls = type(self)
        result = cls.__new__(cls)
        memo[id(self)] = result
        for field in dataclasses.fields(cls):
            value = getattr(self, field.name)
            if field.name == "py_namespace" and value is not None:
                # py_namespace is the evaluation scope for a string-valued
                # parameter expression: {**globals, **locals}, which includes
                # imported modules. Modules are un-deep-copyable singletons
                # (deepcopy raises "cannot pickle 'module' object"), so the
                # scope dict itself is rebuilt shallow — sharing module/global
                # references. Parameter entries are the exception: they must
                # go through the memo so a copied expression evaluates against
                # the *copied* aliases (the ones with_parameters mutates), not
                # the originals (T-141).
                setattr(
                    result,
                    field.name,
                    {
                        k: copy.deepcopy(v, memo) if isinstance(v, Parameter) else v
                        for k, v in value.items()
                    },
                )
            else:
                setattr(result, field.name, copy.deepcopy(value, memo))
        result.__post_init__()
        return result

    def __add__(self, other):
        return _op(Ops.ADD, self, other)

    def __radd__(self, other):
        return _op(Ops.ADD, other, self)

    def __sub__(self, other):
        return _op(Ops.SUB, self, other)

    def __rsub__(self, other):
        return _op(Ops.SUB, other, self)

    def __mul__(self, other):
        return _op(Ops.MUL, self, other)

    def __rmul__(self, other):
        return _op(Ops.MUL, other, self)

    def __truediv__(self, other):
        return _op(Ops.DIV, self, other)

    def __rtruediv__(self, other):
        return _op(Ops.DIV, other, self)

    def __floordiv__(self, other):
        return _op(Ops.FLOORDIV, self, other)

    def __rfloordiv__(self, other):
        return _op(Ops.FLOORDIV, other, self)

    def __mod__(self, other):
        return _op(Ops.MOD, self, other)

    def __rmod__(self, other):
        return _op(Ops.MOD, other, self)

    def __pow__(self, other):
        return _op(Ops.POW, self, other)

    def __rpow__(self, other):
        return _op(Ops.POW, other, self)

    def __neg__(self):
        p = Parameter(value=ParameterExpr([Ops.NEG, self]))
        ParameterCache.add_dependent(self, p)
        return p

    def __pos__(self):
        p = Parameter(value=ParameterExpr([Ops.POS, self]))
        ParameterCache.add_dependent(self, p)
        return p

    def __abs__(self):
        p = Parameter(value=ParameterExpr([Ops.ABS, self]))
        ParameterCache.add_dependent(self, p)
        return p

    def __eq__(self, other):
        return _op(Ops.EQ, self, other)

    def __ne__(self, other):
        return _op(Ops.NE, self, other)

    def __lt__(self, other):
        return _op(Ops.LT, self, other)

    def __le__(self, other):
        return _op(Ops.LE, self, other)

    def __gt__(self, other):
        return _op(Ops.GT, self, other)

    def __ge__(self, other):
        return _op(Ops.GE, self, other)

    def __del__(self):
        ParameterCache.remove(self)

    def __hash__(self):
        return id(self)

    def __str__(self):
        # Calling str() on a Parameter object is confusing. What's the intent?
        # 1. Serializing to a valid Python expression?
        # 2. Is it for logs? For debugging?
        # 3. Is it part of building a wider expression (like a list of parameters)?
        # 4. Evaluating the actual value of a string parameter?
        # Here, we support 2 & 4. We'll likely have to change this when we want support
        # for non-literal string parameters in the UI.

        expr, _ = self.value_as_api_param(
            allow_param_name=True,
            allow_string_literal=True,
        )
        return expr

    def __matmul__(self, other):
        return _op(Ops.MATMUL, self, other)

    def __int__(self):
        if self.dtype is not None:
            return self.dtype(self.get())
        return int(self.get())

    def __float__(self):
        if self.dtype is not None:
            return self.dtype(self.get())
        return float(self.get())

    # NOTE: __bool__ is intentionally not defined. Adding bool(Parameter) ->
    # bool(self.get()) broke some tests (it forces concretization of the
    # wrapped value); keep numeric coercions only.

    def __complex__(self):
        return complex(self.get())

    def value_as_api_param(
        self, allow_param_name=True, allow_string_literal=True
    ) -> tuple[str, bool]:
        """Returns an API-compatible expression[1] that defines this parameter

        What we return depends on the caller's context, since it depends on
        whether we are serializing for a model, submodel or block parameter.

        The boolean is the value of 'is_string' (means "string literal" or
        "do not call eval").

        [1] The returned string can be serialized to JSON, but it is not an
            already escaped JSON string!

        Args:
            allow_param_name: Set to false for (sub)model parameters. Optional.
                If true, and the value is defined by a name, just the name will
                be returned.
            allow_string_literal: Set to false for (sub)model parameters. Optional.
                If true, and the value is a string, then the string will be
                returned and 'is_string' will be returned as True.
        """
        if self.name is not None and allow_param_name:
            return self.name, False

        if self.is_python_expr and isinstance(self.value, str):
            return self.value, False

        if allow_string_literal and isinstance(self.value, str):
            return self.value, True

        return _value_as_str(self.value), False

    def __repr__(self):
        # This must return a valid python expression since it is used for
        # serialization to Python.
        ex, _ = self.value_as_api_param(allow_string_literal=False)
        if len(ex) > 100:
            ex = ex[:50] + "..." + ex[-50:]

        return (
            "Parameter("
            f"name={self.name}, value={ex}, "
            f"is_python_expr={self.is_python_expr}, "
            f"system={self.system.name if self.system is not None else None}"
            ")"
        )

`deepcopy(memo)`

Copy fields and re-run post-init so :class:ParameterCache bookkeeping matches.

Source code in jaxonomy/framework/parameter.py

def __deepcopy__(self, memo):
    """Copy fields and re-run post-init so :class:`ParameterCache` bookkeeping matches."""
    cls = type(self)
    result = cls.__new__(cls)
    memo[id(self)] = result
    for field in dataclasses.fields(cls):
        value = getattr(self, field.name)
        if field.name == "py_namespace" and value is not None:
            # py_namespace is the evaluation scope for a string-valued
            # parameter expression: {**globals, **locals}, which includes
            # imported modules. Modules are un-deep-copyable singletons
            # (deepcopy raises "cannot pickle 'module' object"), so the
            # scope dict itself is rebuilt shallow — sharing module/global
            # references. Parameter entries are the exception: they must
            # go through the memo so a copied expression evaluates against
            # the *copied* aliases (the ones with_parameters mutates), not
            # the originals (T-141).
            setattr(
                result,
                field.name,
                {
                    k: copy.deepcopy(v, memo) if isinstance(v, Parameter) else v
                    for k, v in value.items()
                },
            )
        else:
            setattr(result, field.name, copy.deepcopy(value, memo))
    result.__post_init__()
    return result

`unwrap(value)` `classmethod`

Get the underlying value of raw arrays and Parameter objects alike.

Source code in jaxonomy/framework/parameter.py

@classmethod
def unwrap(cls, value):
    """Get the underlying value of raw arrays and Parameter objects alike."""
    if value is None:
        return None
    if isinstance(value, (Array, bool, int, float, complex)):
        return value
    if isinstance(value, (np.ndarray, np.number)):
        if np.issubdtype(value.dtype, np.number):
            return value
        if value.shape == ():
            return Parameter.unwrap(value.item())
        return Parameter(value).get()
    if isinstance(value, Parameter):
        return value.get()
    if isinstance(value, list):
        return [cls.unwrap(val) for val in value]
    if isinstance(value, tuple):
        return tuple(cls.unwrap(val) for val in value)
    if isinstance(value, dict):
        return {key: cls.unwrap(val) for key, val in value.items()}
    # Fallback for unhandled types: forward to __compute__
    return Parameter(value).get()

`value_as_api_param(allow_param_name=True, allow_string_literal=True)`

Returns an API-compatible expression[1] that defines this parameter

What we return depends on the caller's context, since it depends on whether we are serializing for a model, submodel or block parameter.

The boolean is the value of 'is_string' (means "string literal" or "do not call eval").

[1] The returned string can be serialized to JSON, but it is not an already escaped JSON string!

Parameters:

Name	Type	Description	Default
`allow_param_name`		Set to false for (sub)model parameters. Optional. If true, and the value is defined by a name, just the name will be returned.	`True`
`allow_string_literal`		Set to false for (sub)model parameters. Optional. If true, and the value is a string, then the string will be returned and 'is_string' will be returned as True.	`True`

Source code in jaxonomy/framework/parameter.py

def value_as_api_param(
    self, allow_param_name=True, allow_string_literal=True
) -> tuple[str, bool]:
    """Returns an API-compatible expression[1] that defines this parameter

    What we return depends on the caller's context, since it depends on
    whether we are serializing for a model, submodel or block parameter.

    The boolean is the value of 'is_string' (means "string literal" or
    "do not call eval").

    [1] The returned string can be serialized to JSON, but it is not an
        already escaped JSON string!

    Args:
        allow_param_name: Set to false for (sub)model parameters. Optional.
            If true, and the value is defined by a name, just the name will
            be returned.
        allow_string_literal: Set to false for (sub)model parameters. Optional.
            If true, and the value is a string, then the string will be
            returned and 'is_string' will be returned as True.
    """
    if self.name is not None and allow_param_name:
        return self.name, False

    if self.is_python_expr and isinstance(self.value, str):
        return self.value, False

    if allow_string_literal and isinstance(self.value, str):
        return self.value, True

    return _value_as_str(self.value), False

`ParameterCache`

Global parameter value cache used by all :class:Parameter instances.

Thread safety

All public methods are protected by a class-level reentrant lock (threading.RLock). Using an RLock rather than a plain Lock is necessary because __compute__ may call param.get() recursively (for compound parameter expressions), which would deadlock under a non-reentrant lock held by the outer get() call.

Concurrent simulations in separate threads sharing the same Parameter objects are serialised correctly. However, mutating a parameter from one thread while another thread is actively simulating with it is not recommended — the lock ensures the state remains consistent, but the simulation semantics of mid-run mutation are undefined.

Source code in jaxonomy/framework/parameter.py

class ParameterCache:
    """Global parameter value cache used by all :class:`Parameter` instances.

    Thread safety:
        All public methods are protected by a class-level reentrant lock
        (``threading.RLock``).  Using an ``RLock`` rather than a plain ``Lock``
        is necessary because ``__compute__`` may call ``param.get()`` recursively
        (for compound parameter expressions), which would deadlock under a
        non-reentrant lock held by the outer ``get()`` call.

        Concurrent simulations in separate threads sharing the same ``Parameter``
        objects are serialised correctly.  However, mutating a parameter from one
        thread while another thread is actively simulating with it is not
        recommended — the lock ensures the state remains consistent, but the
        simulation semantics of mid-run mutation are undefined.
    """

    __dependents__: dict["Parameter", set["Parameter"]] = {}
    __cache__: dict["Parameter", ArrayLike] = {}
    __is_dirty__ = defaultdict(lambda: True)
    _lock: threading.RLock = threading.RLock()

    @classmethod
    def _register(cls, param: "Parameter") -> None:
        """Register a newly created Parameter in the cache (called from __post_init__)."""
        with cls._lock:
            if param not in cls.__dependents__:
                cls.__dependents__[param] = set()

    @classmethod
    def get(cls, param: "Parameter") -> ArrayLike:
        with cls._lock:
            if cls.__is_dirty__[param]:
                cls.__cache__[param] = cls.__compute__(param)
                cls.__is_dirty__[param] = False
            return cls.__cache__[param]

    @classmethod
    def replace(cls, param: "Parameter", value: ArrayLike):
        with cls._lock:
            param.value = value
            # Invalidate this parameter and propagate dirty flag recursively to all dependents.
            cls.__invalidate__(param)

    @classmethod
    def remove(cls, param: "Parameter"):
        with cls._lock:
            # Remove param from every set it appears in as a dependent.
            # Use list() to snapshot values so dict iteration is safe if anything
            # changes (e.g. via __del__ called on another thread simultaneously).
            for dependents in list(cls.__dependents__.values()):
                dependents.discard(param)  # discard is safe if param not present

            cls.__dependents__.pop(param, None)
            cls.__cache__.pop(param, None)
            if param in cls.__is_dirty__:
                del cls.__is_dirty__[param]

    @classmethod
    def add_dependent(cls, param: "Parameter", dependent: "Parameter"):
        # Mark 'dependent' as having a dependency on 'param', that is,
        # 'param' is built as an expression that involves 'dependent'.
        with cls._lock:
            cls._register(param)
            cls.__dependents__[param].add(dependent)

    @classmethod
    def get_dependents(cls, param: "Parameter"):
        with cls._lock:
            cls._register(param)
            return cls.__dependents__[param]

    @classmethod
    def print_dependents(cls, param: "Parameter", indent=0):
        """Prints the dependents tree of a parameter"""
        with cls._lock:
            cls._register(param)
            indent_str = "|" + "--" * indent if indent > 0 else ""
            print(indent_str + repr(param))
            for dependent in list(cls.__dependents__[param]):
                cls.print_dependents(dependent, indent + 1)

    @classmethod
    def static_dependents(cls, param: "Parameter"):
        with cls._lock:
            cls._register(param)
            dependents = set()
            for dependent in list(cls.__dependents__[param]):
                if dependent.is_static:
                    dependents.add(dependent)
                dependents |= cls.static_dependents(dependent)
            return dependents

    @classmethod
    def __invalidate__(cls, param: "Parameter"):
        # Caller MUST already hold cls._lock (called from replace() or recursively).
        cls.__cache__[param] = None
        cls.__is_dirty__[param] = True
        # Snapshot the dependent set to avoid mutation-during-iteration if another
        # thread somehow inserts a new dependent while we traverse (belt-and-suspenders).
        for dependent in list(cls.__dependents__.get(param, ())):
            cls.__invalidate__(dependent)

    @classmethod
    def __compute__(cls, param: "Parameter"):
        if isinstance(param.value, ParameterExpr):
            acc = None
            right_value = None
            op = None
            i = 0

            while i < len(param.value):
                val = param.value[i]

                if isinstance(val, Parameter):
                    right_value = val.get()
                elif isinstance(val, ArrayLikeTypes):
                    right_value = val
                elif isinstance(val, Ops):
                    if val in (Ops.NEG, Ops.POS, Ops.ABS):
                        if i + 1 >= len(param.value):
                            raise ParameterError(
                                param, message="Invalid parameter value"
                            )
                        if isinstance(param.value[i + 1], Parameter):
                            right_value = __OPS_FN__[val](param.value[i + 1].get())
                        elif isinstance(param.value[i + 1], ArrayLikeTypes):
                            right_value = __OPS_FN__[val](param.value[i + 1])
                        else:
                            raise ParameterError(
                                param,
                                message=f"Invalid value in parameter list: {param.value[i + 1]} of type {type(param.value[i + 1])}",
                            )
                        i += 1
                    else:
                        op = val
                else:
                    raise ParameterError(
                        param,
                        message=f"Invalid value in parameter list: {val} of type {type(val)}",
                    )

                if acc is not None and right_value is not None and op is not None:
                    acc = __OPS_FN__[op](acc, right_value)
                    op = None
                    right_value = None
                elif right_value is not None:
                    acc = right_value
                    right_value = None
                i += 1

            if acc is not None:
                return acc
            if right_value is not None:
                return right_value
            raise ParameterError(param, message="Invalid parameter value")

        if isinstance(param.value, Parameter):
            return cls.__compute__(param.value)

        if isinstance(param.value, tuple):
            t = _compute_list(param.value, is_tuple=True)
            return t

        if isinstance(param.value, list):
            t = _compute_list(param.value, is_tuple=False)
            return t

        if isinstance(param.value, dict):
            return {key: Parameter.unwrap(val) for key, val in param.value.items()}

        if isinstance(param.value, np.ndarray):
            vals = _resolve_array_param_value(param)
            return np.array(vals, dtype=param.value.dtype)

        if isinstance(param.value, Array):
            vals = _resolve_array_param_value(param)
            if param.value.weak_type:
                return jnp.array(vals)
            return jnp.array(vals, dtype=param.value.dtype)

        if isinstance(param.value, np.number):
            if isinstance(param.value.item(), Parameter):
                return type(param.value)(cls.__compute__(param.value.item()))
            return param.value

        if isinstance(param.value, str) and param.is_python_expr:
            # T-002: produce a useful diagnosis when the expression references
            # an undefined symbol or when py_namespace was never populated.
            scope = param.py_namespace
            try:
                _, resolved_parameters = resolve_parameters(param.value, scope)
                return eval(
                    param.value,
                    scope,
                    {**scope, **resolved_parameters},
                )
            except (TypeError, NameError, KeyError) as err:
                available = sorted(scope.keys()) if scope else "none"
                raise ValueError(
                    f"Parameter {param.name!r} expression {param.value!r} "
                    f"references undefined symbol(s). Available symbols: "
                    f"{available}"
                ) from err

        return param.value

`print_dependents(param, indent=0)` `classmethod`

Prints the dependents tree of a parameter

Source code in jaxonomy/framework/parameter.py

@classmethod
def print_dependents(cls, param: "Parameter", indent=0):
    """Prints the dependents tree of a parameter"""
    with cls._lock:
        cls._register(param)
        indent_str = "|" + "--" * indent if indent > 0 else ""
        print(indent_str + repr(param))
        for dependent in list(cls.__dependents__[param]):
            cls.print_dependents(dependent, indent + 1)

`RuntimeVariantSubsystem`

Bases: LeafSystem

Switch between pre-built submodel choices via a discrete selector input.

This is the runtime counterpart to select_variant / Variant. Unlike the build-time selector (which never instantiates the unselected branches), RuntimeVariantSubsystem builds every choice and routes the selected branch's output through. The selector is a normal input port, so it can be driven by any discrete control signal in the diagram and the active branch follows at simulate time. This is the runtime-controlled variant pattern, as opposed to the label-mode build-time variant.

Implementation: the block stacks all branches' outputs along a new leading axis and picks out the selected slice with integer indexing. This is the same mechanism used by MultiPortSwitch (T-118), reused here at the framework level so it does not pull a library dependency.

Contract — "all branches integrated each step"

Because the underlying stack traces every branch, every choice's submodel runs on every step and sees the same input trajectory. The consequences:

Pure (memoryless) branches behave exactly as you'd expect: only the selected branch's output is exposed; gradients w.r.t. the active branch's parameters are non-zero, and gradients w.r.t. the others are zero (matching MultiPortSwitch's data-input semantics).
The selector is non-differentiable (round + clip zero out its gradient), as expected for a control signal.
If a branch holds internal discrete state (e.g. a hold latch) the caller is responsible for supplying that state. RuntimeVariantSubsystem itself is stateless; if you need stateful sub-Diagrams, hoist the state out, or build the runtime switch by composing MultiPortSwitch (T-118) with N pre-built sub-diagrams in a parent DiagramBuilder.

All branches must return outputs that are broadcast-compatible (the stack op requires a common shape/dtype after broadcasting).

Parameters:

Name	Type	Description	Default
`choices`		Either a sequence of submodel callables `[f0, f1, ..., f_{N-1}]` or a mapping `{int: callable}` (integer keys must be `0..N-1`). Each callable has signature `f(*inputs) -> output` and must be JAX-traceable.	required
`n_inputs`	`int`	Number of user inputs forwarded to every branch. Input port 0 is always the selector; ports `1..n_inputs` are the user inputs. Defaults to 1.	`1`
`default_choice`	`int`	Index of the choice used as the default. Stored for introspection / documentation; the runtime selector value still controls which branch is exposed each step. Defaults to 0.	`0`
`name`		Optional block name.	required

Input ports

(0) selector — scalar integer-valued signal in [0, N-1]. Floating values are rounded and clipped. (1..n_inputs) user inputs forwarded to every branch.

Output ports

(0) The selected branch's output.

Raises:

Type	Description
`VariantError`	If `choices` is empty / non-callable / has bad keys, or if `default_choice` is out of range.

Source code in jaxonomy/framework/variants.py

class RuntimeVariantSubsystem(LeafSystem):
    """Switch between pre-built submodel choices via a discrete selector input.

    This is the runtime counterpart to ``select_variant`` / ``Variant``. Unlike
    the build-time selector (which never instantiates the unselected branches),
    ``RuntimeVariantSubsystem`` builds *every* choice and routes the selected
    branch's output through. The selector is a normal input port, so it can be
    driven by any discrete control signal in the diagram and the active branch
    follows at simulate time. This is the runtime-controlled variant pattern,
    as opposed to the label-mode build-time variant.

    Implementation: the block stacks all branches' outputs along a new leading
    axis and picks out the selected slice with integer indexing. This is the
    same mechanism used by ``MultiPortSwitch`` (T-118), reused here at the
    framework level so it does not pull a library dependency.

    Contract — "all branches integrated each step"
    -----------------------------------------------
    Because the underlying ``stack`` traces every branch, every choice's
    submodel runs on every step and sees the same input trajectory. The
    consequences:

    - Pure (memoryless) branches behave exactly as you'd expect: only the
      selected branch's output is exposed; gradients w.r.t. the active
      branch's parameters are non-zero, and gradients w.r.t. the others
      are zero (matching ``MultiPortSwitch``'s data-input semantics).

    - The selector is non-differentiable (``round`` + ``clip`` zero out
      its gradient), as expected for a control signal.

    - If a branch holds internal discrete state (e.g. a hold latch) the
      caller is responsible for supplying that state. ``RuntimeVariantSubsystem``
      itself is stateless; if you need stateful sub-Diagrams, hoist the
      state out, or build the runtime switch by composing ``MultiPortSwitch``
      (T-118) with N pre-built sub-diagrams in a parent ``DiagramBuilder``.

    All branches must return outputs that are broadcast-compatible (the
    stack op requires a common shape/dtype after broadcasting).

    Args:
        choices:
            Either a sequence of submodel callables ``[f0, f1, ..., f_{N-1}]``
            *or* a mapping ``{int: callable}`` (integer keys
            must be ``0..N-1``). Each callable has signature
            ``f(*inputs) -> output`` and must be JAX-traceable.
        n_inputs:
            Number of user inputs forwarded to every branch. Input port 0
            is always the selector; ports ``1..n_inputs`` are the user
            inputs. Defaults to 1.
        default_choice:
            Index of the choice used as the default. Stored for
            introspection / documentation; the runtime selector value
            still controls which branch is exposed each step. Defaults
            to 0.
        name:
            Optional block name.

    Input ports:
        (0) selector  — scalar integer-valued signal in ``[0, N-1]``.
            Floating values are rounded and clipped.
        (1..n_inputs) user inputs forwarded to every branch.

    Output ports:
        (0) The selected branch's output.

    Raises:
        VariantError:
            If ``choices`` is empty / non-callable / has bad keys, or if
            ``default_choice`` is out of range.
    """

    def __init__(
        self,
        choices,
        n_inputs: int = 1,
        default_choice: int = 0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        # Normalize choices to an indexable tuple of callables.
        normalized = self._normalize_choices(choices, name=kwargs.get("name"))
        self._choices: tuple[Callable, ...] = normalized
        n_choices = len(normalized)

        if not (0 <= int(default_choice) < n_choices):
            raise VariantError(
                f"RuntimeVariantSubsystem {kwargs.get('name')!r}: "
                f"default_choice={default_choice} is out of range "
                f"[0, {n_choices - 1}]."
            )
        if int(n_inputs) < 0:
            raise VariantError(
                f"RuntimeVariantSubsystem {kwargs.get('name')!r}: "
                f"n_inputs must be >= 0, got {n_inputs}."
            )

        self._n_choices = n_choices
        self._n_inputs = int(n_inputs)
        self._default_choice = int(default_choice)

        # Port 0 = selector; ports 1..n_inputs forwarded to every branch.
        self.declare_input_port(name="selector")
        for i in range(self._n_inputs):
            self.declare_input_port(name=f"u_{i}")

        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── helpers ────────────────────────────────────────────────────────────

    @staticmethod
    def _normalize_choices(choices, name=None) -> tuple[Callable, ...]:
        if isinstance(choices, Mapping):
            if not choices:
                raise VariantError(
                    f"RuntimeVariantSubsystem {name!r}: choices mapping is empty."
                )
            try:
                keys = sorted(int(k) for k in choices.keys())
            except (TypeError, ValueError) as exc:
                raise VariantError(
                    f"RuntimeVariantSubsystem {name!r}: choices mapping keys "
                    f"must be integers; got {list(choices.keys())!r}."
                ) from exc
            if keys != list(range(len(keys))):
                raise VariantError(
                    f"RuntimeVariantSubsystem {name!r}: choices mapping keys "
                    f"must be a contiguous 0..N-1 range; got {keys!r}."
                )
            ordered = tuple(choices[k] for k in keys)
        else:
            ordered = tuple(choices)
            if not ordered:
                raise VariantError(
                    f"RuntimeVariantSubsystem {name!r}: choices sequence is empty."
                )
        for i, fn in enumerate(ordered):
            if not callable(fn):
                raise VariantError(
                    f"RuntimeVariantSubsystem {name!r}: choice [{i}] is not "
                    f"callable (got {type(fn).__name__}). Pass a submodel "
                    f"function f(*inputs) -> output."
                )
        return ordered

    # ── output computation ────────────────────────────────────────────────

    def _compute_output(self, _time, _state, *inputs, **_params):
        selector = inputs[0]
        user_inputs = inputs[1 : 1 + self._n_inputs]
        # Evaluate every branch — this is the "all branches integrated each
        # step" contract. ``jnp.stack`` requires a common shape/dtype across
        # branch outputs.
        branch_outputs = [
            jnp.asarray(fn(*user_inputs)) for fn in self._choices
        ]
        stacked = jnp.stack(branch_outputs, axis=0)
        idx = jnp.clip(
            jnp.round(selector).astype(jnp.int32), 0, self._n_choices - 1
        )
        return stacked[idx]

    # ── introspection ─────────────────────────────────────────────────────

    @property
    def n_choices(self) -> int:
        """Number of variant choices held by this block."""
        return self._n_choices

    @property
    def default_choice(self) -> int:
        """Default choice index (documentary; runtime selector still rules)."""
        return self._default_choice

`default_choice` `property`

Default choice index (documentary; runtime selector still rules).

`n_choices` `property`

Number of variant choices held by this block.

`ShapeMismatchError`

Bases: StaticError

Block parameters or input/outputs have mismatched shapes.

Source code in jaxonomy/framework/error.py

class ShapeMismatchError(StaticError):
    """Block parameters or input/outputs have mismatched shapes."""

    def __init__(self, expected_shape=None, actual_shape=None, **kwargs):
        super().__init__(**kwargs)
        self.expected_shape = expected_shape
        self.actual_shape = actual_shape

    def __str__(self):
        if self.expected_shape or self.actual_shape:
            return (
                f"Shape mismatch: expected {self.expected_shape}, "
                f"got {self.actual_shape}" + self._context_info()
            )
        return f"Shape mismatch{self._context_info()}"

`StaticError`

Bases: JaxonomyError

Wraps a Python exception to record the offending block id. The original exception is found in the 'cause' field.

See jaxonomy.framework.context_factory._check_types for use.

This is called 'static' (as opposed to say 'runtime') meaning this is for wrapping errors detected prior to running a simulation.

Source code in jaxonomy/framework/error.py

class StaticError(JaxonomyError):
    """Wraps a Python exception to record the offending block id. The original
    exception is found in the '__cause__' field.

    See jaxonomy.framework.context_factory._check_types for use.

    This is called 'static' (as opposed to say 'runtime') meaning this is for
    wrapping errors detected prior to running a simulation."""

    pass

`SystemBase` `dataclass`

Basic building block for simulation in jaxonomy.

NOTE: Type hints in SystemBase indicate the union between what would be returned by a LeafSystem and a Diagram. See type hints of the subclasses for the specific argument and return types.

Source code in jaxonomy/framework/system_base.py

@dataclasses.dataclass
class SystemBase:
    """Basic building block for simulation in jaxonomy.

    NOTE: Type hints in SystemBase indicate the union between what would be returned
    by a LeafSystem and a Diagram. See type hints of the subclasses for the specific
    argument and return types.
    """

    # Generated unique ID for this system
    system_id: Hashable = dataclasses.field(default_factory=next_system_id, init=False)
    name: Optional[str] = None  # Human-readable name for this system (optional but never None)
    ui_id: Optional[str] = None  # UUID of the block when loaded from JSON (optional, may be None)

    # Immediate parent of the current system (can only be a Diagram).
    # If None, _this_ is the root system.
    parent: Optional[Diagram] = None

    def __post_init__(self):
        if self.name is None:
            self.name = f"{type(self).__name__}_{self.system_id}_"

        # All "cache sources" for this system. Typically these will correspond to
        # input ports, output ports, time derivative calculations, and any custom
        # "cached data" declared by the user (e.g. see ModelicaFMU block).
        self.callbacks: List[SystemCallback] = []

        # Index into SystemCallbacks for each port. For instance, input port `i` can
        # be retrieved by `self.callbacks[self.input_port_indices[i]]`. The
        # `input_ports` and `output_ports` properties give more convenient access.
        self.input_port_indices: List[int] = []
        self.output_port_indices: List[int] = []

        # Override this or set manually to provide a custom characteristic time scale
        # for the system. At the moment this is only used for zero-crossing isolation
        # in the simulator.
        self.characteristic_time = 1.0

        # A dependency graph for the system, mapping prerequisites of each calculation.
        # `None` indicates that the dependency graph has not been constructed yet.  If
        # accessed via the `dependency_graph` property, it will be constructed
        # automatically as necessary.
        self._dependency_graph: DependencyGraph = None

        # Static parameters are not jax-traceable
        self._static_parameters: dict[str, Parameter] = {}

        # If not empty, this defines the shape and data type of the numeric parameters
        # in the LeafSystem. This value will be used to initialize the context, so it
        # will also serve as the initial value unless explicitly overridden. In the
        # simplest cases, parameters could be stored as attributes of the LeafSystem,
        # but declaring them has the advantage of moving the values to the context,
        # allowing them to be traced by JAX rather than stored as static data. This
        # means they can be differentiated, vmapped, or otherwise modified without
        # re-compiling the simulation.
        self._dynamic_parameters: dict[str, Array] = {}

        # Map from (input_port, output_port) if that pair is feedthrough
        # `None` indicates that the feedthrough is unknown for this system.
        # This will be computed automatically using the dependency graph
        # during algebraic loop detection unless it is set manually.
        # To manually set feedthrough, either declare this explicitly or
        # override `get_feedthrough`.
        self.feedthrough_pairs: List[Tuple[int, int]] = None

        # Pre-sorted list of all output update events for this system.  This will
        # be created when the associated property is first accessed.  This should
        # only need to be done for the root system.
        self._cache_update_events: EventCollection = None

        # Cached lists of i/o ports SystemCallbacks. Do not read this directly.
        # ``_PortAccessList`` gives an actionable IndexError when a user indexes
        # an empty / too-short port list (T-B1-followup-empty-output-ports).
        self._cached_input_ports: List[SystemCallback] = _PortAccessList(
            self.name, "in"
        )
        self._cached_output_ports: List[SystemCallback] = _PortAccessList(
            self.name, "out"
        )

        # Should the system use caching or re-evaluate the output every time?
        # By default this should be False, but it can be overridden, for instance
        # during the main simulation loop.  When not in simulation mode, it is
        # better to have the cache disabled to avoid stale data.
        self._cache_enabled = False

    def __deepcopy__(self, memo):
        """Deep-copy while keeping partially constructed copies hashable.

        Subsystems reference themselves via callbacks; the default deepcopy order can
        call :meth:`__hash__` (via dict/set operations) before ``system_id`` exists
        on the copy. Assign a new ``system_id`` immediately after memo registration.
        """
        cls = type(self)
        result = cls.__new__(cls)
        memo[id(self)] = result
        result.system_id = next_system_id()
        for key, value in self.__dict__.items():
            if key == "system_id":
                continue
            setattr(result, key, copy.deepcopy(value, memo))
        return result

    def __hash__(self) -> Hashable:
        return hash(self.system_id)

    def pprint(self, output=print, fancy=True) -> str:
        """Pretty-print the system and its hierarchy."""
        output(self._pprint_helper(fancy=fancy).strip())

    def _pprint_helper(self, prefix="", fancy=True) -> str:
        if fancy:
            return pprint_fancy(prefix, self)
        return f"{prefix}|-- {self.name}(id={self.system_id})\n"

    def post_simulation_finalize(self) -> None:
        """Finalize the system after simulation has completed.

        This is only intended for special blocks that need to clean up
        resources and close files."""

    @property
    def root(self) -> SystemBase:
        """Get the root system of the current system."""
        if self.parent is None:
            return self
        return self.parent.root

    @property
    def static_parameters(self) -> dict[str, Parameter]:
        return self._static_parameters

    @static_parameters.setter
    def static_parameters(self, value):
        self._static_parameters = value

    @property
    def dynamic_parameters(self) -> dict[str, Parameter]:
        return self._dynamic_parameters

    @property
    def parameters(self) -> dict[str, Parameter]:
        return {**self.static_parameters, **self.dynamic_parameters}

    #
    # Simulation interface
    #
    @property
    def cache_enabled(self) -> bool:
        if self.parent is None:
            return self._cache_enabled
        return self.root._cache_enabled

    @cache_enabled.setter
    def cache_enabled(self, value: bool):
        self.root._cache_enabled = value

    @property
    @abc.abstractmethod
    def has_feedthrough_side_effects(self) -> bool:
        """Check if the system includes any feedthrough calls to `io_callback`."""
        # This is a tricky one to explain and is almost always False except for a
        # PythonScript block that is not JAX traced.  Basically, if the output of
        # the system is used as an ODE right-hand-side, will it fail in the case where
        # the ODE solver defines a custom VJP?  This happens in diffrax, so for example
        # if a PythonScript block is used to compute the ODE right-hand-side, it will
        # fail with "Effects not supported in `custom_vjp`"
        pass

    @property
    @abc.abstractmethod
    def has_ode_side_effects(self) -> bool:
        """Check if the ODE RHS for the system includes any calls to `io_callback`."""
        # This flag indicates that the system `has_feedthrough_side_effects` AND that
        # signal is used as an ODE right-hand-side.  This is used to determine whether
        # a JAX ODE solver can be used to integrate the system.
        pass

    @property
    @abc.abstractmethod
    def has_continuous_state(self) -> bool:
        pass

    @property
    @abc.abstractmethod
    def has_discrete_state(self) -> bool:
        pass

    @property
    @abc.abstractmethod
    def has_zero_crossing_events(self) -> bool:
        pass

    def eval_time_derivatives(self, context: ContextBase) -> StateComponent:
        """Evaluate the continuous time derivatives for this system.

        Given the _root_ context, evaluate the continuous time derivatives,
        which must have the same PyTree structure as the continuous state.

        In principle, this can be overridden by custom implementations, but
        in general it is preferable to declare continuous states for LeafSystems
        using `declare_continuous_state`, which accepts a callback function
        that will be used to compute the derivatives. For Diagrams, the time
        derivatives are computed automatically using the callback functions for
        all child systems with continuous state.

        Args:
            context (ContextBase): root context of this system

        Returns:
            StateComponent:
                Continuous time derivatives for this system, or None if the system
                has no continuous state.
        """
        return None

    @property
    @abc.abstractmethod
    def mass_matrix(self) -> StateComponent:
        """Mass matrix for this system.

        Returns PyTree-structured data where each leaf is an (n, n) array.
        This is used for implicit integration methods (currently only BDF).
        """
        pass

    @property
    @abc.abstractmethod
    def has_mass_matrix(self) -> bool:
        """Returns True if any component of the system has a nontrivial mass matrix."""
        pass

    @abc.abstractmethod
    def eval_zero_crossing_updates(
        self,
        context: ContextBase,
        events: EventCollection,
    ) -> State:
        """Evaluate reset maps associated with zero-crossing events.

        Args:
            context (ContextBase):
                The context for the system, containing the current state and parameters.
            events (EventCollection):
                The collection of events to be evaluated (for example zero-crossing or
                periodic events for this system).

        Returns:
            State: The complete state with all updates applied.

        Notes:
            (1) Following the Drake definition, "unrestricted" updates are allowed to
            modify any component of the state: continuous, discrete, or mode.  These
            updates are evaluated in the order in which they were declared, so it is
            _possible_ (but should be strictly avoided) for multiple events to modify the
            same state component at the same time.

            Each update computes its results given the _current_ state of the system
            (the "minus" values) and returns the _updated_ state (the "plus" values).
            The update functions cannot access any information about the "plus" values of
            its own state or the state of any other block.  This could change in the future
            but for now it ensures consistency with Drake's discrete semantices:

            More specifically, since all unrestricted updates can modify the entire state,
            any time there are multiple unrestricted updates, the resulting states are
            ALWAYS in conflict.  For example, suppose a system has two unrestricted
            updates, `event1` and `event2`.  At time t_n, `event1` is active and `event2`
            is inactive.  First, `event1` is evaluated, and the state is updated.  Then
            `event2` is evaluated, but the state is not updated.  Which one is valid?
            Obviously, the `event1` return is valid, but how do we communicate this to JAX?
            The situation is more complicated if both `event1` and `event2` happen to be
            active.  In this case the states have to be "merged" somehow.  In the worst
            case, these two will modify the same components of the state in different ways.

            The implementation updates the state in a local copy of the context (since both
            are immutable).  This allows multiple unrestricted updates, but leaves open the
            possibility of multiple active updates modifying the state in conflicting ways.
            This should strictly be avoided by the implementer of the LeafSystem.  If it is
            at all unclear how to do this, it may be better to split the system into
            multiple blocks to be safe.

            (2) The events are evaluated conditionally on being marked "active"
            (indicating that their guard function triggered), so the entire event
            collection can be passed without filtering to active events. This is necessary
            to make the function calls work with JAX tracing, which do not allow for
            variable-sized arguments or returns.
        """
        pass

    def handle_discrete_update(
        self,
        events: EventCollection,
        context: ContextBase,
        *,
        topological_order: bool = False,
    ) -> ContextBase:
        """Compute and apply active discrete updates.

        Given the _root_ context, evaluate the discrete updates, which must have the
        same PyTree structure as the discrete states of this system. This should be
        a pure function, so that it does not modify any aspect of the context in-place
        (even though it is difficult to strictly prevent this in Python).

        This will evaluate the set of events that result from declaring state or output
        update events on systems using `LeafSystem.declare_periodic_update` and
        `LeafSystem.declare_output_port` with an associated periodic update rate.

        This is intended for internal use by the simulator and should not normally need
        to be invoked directly by users. Events are evaluated conditionally on being
        marked "active", so the entire event collection can be passed without filtering
        to active events. This is necessary to make the function calls work with JAX
        tracing, which do not allow for variable-sized arguments or returns.

        For a discrete system updating at a particular rate, the update rule for a
        particular block is:

        ```
        x[n+1] = f(t[n], x[n], u[n])
        y[n]   = g(t[n], x[n], u[n])
        ```

        Additionally, the value y[n] is held constant until the next update from the
        point of view of other continuous-time or asynchronous discrete-time blocks.

        Because each output `y[n]` may in general depend on the input `u[n]` evaluated
        _at the same time_, the composite discrete update function represents a
        system of equations.  However, since algebraic loops are prohibited, the events
        can be ordered and executed sequentially to ensure that the updates are applied
        in the correct order.  This is implemented in
        `SystemBase.sorted_callbacks`.

        Multirate systems work in the same way, except that the events are evaluated
        conditionally on whether the current time corresponds to an update time for each
        event.

        Args:
            events (EventCollection): collection of discrete update events
            context (ContextBase): root context for this system

        Returns:
            ContextBase:
                updated context with all active updates applied to the discrete state
        """
        logger.debug(
            f"Handling {events.num_events} discrete update events at t={context.time}"
        )
        if events.has_events:
            # Two-phase approach for correct x⁻ / x⁺ semantics:
            #
            # Phase 1 — OUTPUT CACHE UPDATES (is_state_update=False):
            #   Process in sorted dependency order without a snapshot.  Each cache
            #   update computes y[n] = g(x[n]) and stores it in state.cache.  These
            #   are safe to sequence without a snapshot because they only READ discrete
            #   state (never write it), so no block can see another block's x⁺ here.
            #
            # Phase 2 — DISCRETE STATE UPDATES (is_state_update=True):
            #   After Phase 1, ALL blocks have correct y[n] in their state.cache.
            #   Now take a snapshot.  Each state update is evaluated against a "blended"
            #   context where:
            #     • the current block uses its accumulated context (so it sees any
            #       intra-block cache updates that already fired in Phase 1)
            #     • every OTHER block uses the snapshot (x⁻ discrete states + y[n]
            #       caches) — preventing block B from reading block A's x⁺ state.
            #
            # This correctly handles both:
            #   (a) DiscreteClock → UnitDelay: UnitDelay state update reads clock's
            #       y[n]=1.0 (set in Phase 1), not the stale y[n-1]=0.0 snapshot.
            #   (b) Cross-block swap (A reads B, B reads A): each sees x⁻ of the other
            #       via the Phase 2 snapshot, not x⁺.

            # T-105-followup-priority-scheduler-hook — compute the
            # per-system priority/rate tiebreak key once.  ``None`` means
            # no leaf has an explicit ``priority`` attribute set, so the
            # downstream scheduler must use its legacy byte-equivalent
            # order (declaration order here, ``str(system_id)`` for
            # Kahn's in the topological branch below).  ``_phase1`` is
            # computed up-front so Phase 1 cache updates honour the
            # same per-block ordering as Phase 2.
            _phase1_priority_tiebreak = _build_priority_tiebreak(self)

            # Phase 1: cache/output update events. These arrive already in
            # execution (topological) order — ``cache_update_events`` sorts them
            # via ``sorted_callbacks`` → ``sort_trackers``, which respects
            # feedthrough dependencies, INCLUDING transitively through pure-
            # feedthrough intermediates (e.g. a Mux between two sample-and-holds).
            # The one missing piece was refreshing the port cache BETWEEN them:
            # a downstream feedthrough sample-and-hold (e.g. ZeroOrderHold) must
            # read its upstream's fresh current-tick output, not the stale pre-
            # tick value. Without the refresh a same-rate ZOH lags its source by
            # a step, violating Simulink sorted-execution semantics (a same-rate
            # ZOH is the identity; two in series are a no-op).
            phase1_events = [e for e in events if not e.is_state_update]
            _multi_p1 = len(phase1_events) > 1
            if _phase1_priority_tiebreak is not None:
                # Stable sort: events for the same system_id keep their
                # declared order.
                phase1_events = sorted(
                    phase1_events,
                    key=lambda e: _phase1_priority_tiebreak(e.system_id),
                )
            for event in phase1_events:
                system_id = event.system_id
                state = event.handle(context)
                local_context = context[system_id].with_state(state)
                context = context.with_subcontext(system_id, local_context)
                if _multi_p1:
                    # Refresh so the next (downstream) feedthrough cache update
                    # reads this block's fresh current-tick output.
                    context = context.refresh_port_cache()

            # Between phases: refresh the root port_cache so that Phase 2 state
            # update callbacks (which use OutputPort.eval → port_cache lookup) see
            # the updated y[n] values written in Phase 1, not stale values from
            # before Phase 1.
            snapshot = context.refresh_port_cache()

            # Phase 2: state update events.  Two ordering modes:
            #
            #   (a) Diagonal (default, `topological_order=False`): evaluate
            #       events in their declared order against the snapshot.
            #       Each block sees x⁻ for every other block — preserves the
            #       cross-block-swap atomicity documented above.
            #
            #   (b) Lower-triangular (T-022a, opt-in via
            #       `topological_order=True`): evaluate events in the
            #       topological order of the discrete dependency graph.
            #       Block B's update sees the post-update x⁺ of any block A
            #       upstream of B (A→B), AND the post-update y⁺[n+1] of
            #       any sample-and-hold output of A — both via the
            #       accumulated context.  Cycles raise
            #       ``DependencyCycleError``.
            state_events = [e for e in events if e.is_state_update]

            # Reuse the Phase-1 priority tiebreak (already accounts for
            # an "any leaf has explicit priority" check).  ``None`` →
            # legacy byte-equivalent order downstream.
            priority_tiebreak = _phase1_priority_tiebreak

            if topological_order:
                from .discrete_dependencies import (
                    DependencyCycleError,
                    build_discrete_dependency_graph,
                    topological_sort,
                )
                graph = build_discrete_dependency_graph(self)
                event_ids = {e.system_id for e in events}  # all events
                trimmed = {
                    n: deps & event_ids
                    for n, deps in graph.items() if n in event_ids
                }
                topological_sort(
                    trimmed, tiebreak_key=priority_tiebreak,
                )  # validate (raises on cycle)
                order = topological_sort(
                    trimmed, tiebreak_key=priority_tiebreak,
                )
                # Group all events by system; topological branch
                # interleaves state + cache updates per block so
                # downstream blocks see y[n+1] (cache reflects post-
                # update state).
                all_events_by_id: dict = {}
                for e in events:
                    all_events_by_id.setdefault(e.system_id, []).append(e)
                ancestors: dict = {n: set() for n in order}
                for n in order:
                    for d in trimmed.get(n, ()):
                        ancestors[n] |= ancestors.get(d, set()) | {d}
                for sid in order:
                    block_events = all_events_by_id.get(sid, [])
                    blended = snapshot
                    for use_acc in {sid, *ancestors[sid]}:
                        blended = blended.with_subcontext(
                            use_acc, context[use_acc],
                        )
                    # Empty port cache so downstream OutputPort.eval
                    # recomputes against the accumulated state.
                    blended = blended.with_port_cache({})
                    # State updates first (so cache updates can read x⁺).
                    for event in block_events:
                        if not event.is_state_update:
                            continue
                        state = event.handle(blended)
                        local_context = context[sid].with_state(state)
                        context = context.with_subcontext(sid, local_context)
                        blended = blended.with_subcontext(sid, local_context)
                    # Cache updates re-run against post-state context
                    # so that y[n+1] reflects x[n+1] for sample-and-hold
                    # output ports.
                    for event in block_events:
                        if event.is_state_update:
                            continue
                        state = event.handle(blended)
                        local_context = context[sid].with_state(state)
                        context = context.with_subcontext(sid, local_context)
                        blended = blended.with_subcontext(sid, local_context)
            else:
                # Diagonal (legacy) order.  When ANY leaf declared an
                # explicit ``priority`` the events are stable-sorted by
                # the (rate, priority, name) tiebreak key — block
                # authors who opt into priorities see deterministic
                # within-tick ordering even on the legacy snapshot path.
                # When no priorities are declared the list is returned
                # unchanged so byte-equivalence with the pre-followup
                # diagonal scheduler is preserved.
                ordered_state_events = state_events
                if priority_tiebreak is not None:
                    ordered_state_events = sorted(
                        state_events,
                        key=lambda e: priority_tiebreak(e.system_id),
                    )
                for event in ordered_state_events:
                    system_id = event.system_id
                    # Own block: use accumulated context (may have updated cache from Phase 1)
                    # Other blocks: use snapshot (x⁻ discrete states + y[n] caches)
                    blended = snapshot.with_subcontext(system_id, context[system_id])
                    state = event.handle(blended)
                    local_context = context[system_id].with_state(state)
                    context = context.with_subcontext(system_id, local_context)

        return context

    def handle_zero_crossings(
        self, events: EventCollection, context: ContextBase
    ) -> ContextBase:
        """Compute and apply active zero-crossing events.

        This is intended for internal use by the simulator and should not normally need
        to be invoked directly by users. Events are evaluated conditionally on being
        marked "active", so the entire event collection can be passed without filtering
        to active events. This is necessary to make the function calls work with JAX
        tracing, which do not allow for variable-sized arguments or returns.

        Args:
            events (EventCollection): collection of zero-crossing events
            context (ContextBase): root context for this system

        Returns:
            ContextBase: updated context with all active zero-crossing events applied
        """
        logger.debug(
            "Handling %d state transition events at t=%s",
            events.num_events,
            context.time,
        )
        if events.num_events > 0:
            # Compute the updated state for all active events
            context = context.refresh_port_cache()
            state = self.eval_zero_crossing_updates(context, events)

            # Apply the updates
            context = context.with_state(state)
        return context

    # Events that happen at some regular interval
    @property
    def periodic_events(self) -> FlatEventCollection:
        return self.cache_update_events + self.state_update_events

    # Events that update the discrete state of the system
    @property
    @abc.abstractmethod
    def state_update_events(self) -> FlatEventCollection:
        pass

    # Events that refresh sample-and-hold outputs of the system
    @property
    def cache_update_events(self) -> FlatEventCollection:
        if self._cache_update_events is None:
            # Sort the output update events and store in the private attribute.
            self._cache_update_events = [
                cb.event for cb in self.sorted_callbacks if cb.event is not None
            ]

        return FlatEventCollection(tuple(self._cache_update_events))

    @property
    @abc.abstractmethod
    def _flat_callbacks(self) -> List[SystemCallback]:
        """Get a flat list of callbacks for this and all child systems."""
        pass

    @property
    def sorted_callbacks(self) -> List[SystemCallback]:
        """Sort and return the callbacks for this system."""
        # Collect all of the callbacks associated with this system.  These are
        # SystemCallback objects, so they are all associated with trackers in the
        # dependency graph.  We can use these to sort the events in execution order.
        trackers = sort_trackers([cb.tracker for cb in self._flat_callbacks])

        # Retrieve the callback associated with each tracker.
        return [tracker.cache_source for tracker in trackers]

    def recompute_port_cache(self, context: ContextBase) -> dict[Hashable, Array]:
        """Recompute all output ports and return them in a dictionary."""
        # The callbacks are already sorted in the correct order, so we can just
        # evaluate them in sequence, assuming all upstream values are correct.
        for cb in self.sorted_callbacks:
            if not isinstance(cb, OutputPort):
                continue
            # Evaluate the callback and store the result in the cache
            val = cb.calc(context)
            context = context.with_port_cache_entry(hash(cb), val)

        return context

    # Events that are triggered by a "guard" function and may induce a "reset" map
    @property
    @abc.abstractmethod
    def zero_crossing_events(self) -> EventCollection:
        pass

    @abc.abstractmethod
    def determine_active_guards(self, context: ContextBase) -> EventCollection:
        """Determine active guards for zero-crossing events.

        This method is responsible for evaluating and determining which
        zero-crossing events are active based on the current system mode
        and other conditions.  This can be overridden to flag active/inactive
        guards on a block-specific basis, for instance in a StateMachine-type
        block. By default all guards are marked active at this point unless
        the zero-crossing event was declared with a non-default `start_mode`, in
        which case the guard is activated conditionally on the current mode.

        For example, in a system with finite state transitions, where a transition from
        mode A to mode B is triggered by a guard function g_AB and the inverse
        transition is triggered by a guard function g_BA, this function would activate
        g_AB if the system is in mode A and g_BA if the system is in mode B. The other
        guard function would be inactive.  If the zero-crossing event is not associated
        with a starting mode, it is considered to be always active.

        Args:
            context (ContextBase):
                The root context containing the overall state and parameters.

        Returns:
            EventCollection:
                A collection of zero-crossing events with active/inactive status
                updated based on the current system mode and other conditions.
        """
        pass

    #
    # I/O ports
    #
    @property
    def input_ports(self) -> List[InputPort]:
        if len(self._cached_input_ports) != len(self.input_port_indices):
            ports = list(self.callbacks[i] for i in self.input_port_indices)
            self._cached_input_ports.clear()
            self._cached_input_ports.extend(ports)
        return self._cached_input_ports

    def get_input_port(self, name: str) -> tuple[InputPort, int]:
        """Retrieve a specific input port by name."""
        for i, port in enumerate(self.input_ports):
            if port.name == name:
                return port, i
        raise ValueError(
            f"System {self.name} has no input port named {name}. "
            f"Available ports: {list(map(lambda x: x.name,self.input_ports))}"
        )

    @property
    def num_input_ports(self) -> int:
        return len(self.input_port_indices)

    @property
    def output_ports(self) -> List[OutputPort]:
        if len(self._cached_output_ports) != len(self.output_port_indices):
            ports = list(self.callbacks[i] for i in self.output_port_indices)
            self._cached_output_ports.clear()
            self._cached_output_ports.extend(ports)
        return self._cached_output_ports

    def get_output_port(self, name: str) -> OutputPort:
        """Retrieve a specific output port by name."""
        for port in self.output_ports:
            if port.name == name:
                return port
        raise ValueError(f"System {self.name} has no output port named {name}")

    @property
    def num_output_ports(self) -> int:
        return len(self.output_port_indices)

    def eval_input(self, context: ContextBase, port_index: int = 0) -> Array:
        """Get the input for a given port.

        This works by evaluating the callback function associated with the port, which
        will "pull" the upstream output port values.

        Args:
            context (ContextBase): root context for this system
            port_index (int, optional): index into `self.input_ports`, for example
                the value returned by `declare_input_port`. Defaults to 0.

        Returns:
            Array: current input values
        """
        return self.input_ports[port_index].eval(context)

    def collect_inputs(
        self, context: ContextBase, port_indices: list[int] = None
    ) -> List[Array]:
        """Collect all current inputs for this system.

        Args:
            context (ContextBase): root context for this system
            port_indices (List[int], optional): list of input port indices to collect.
                If None (default), will return values from all ports.  Otherwise will
                return a list of length(num_input_ports), where the values are None for
                ports not in the list.

        Returns:
            List[Array]: list of all current input values
        """
        if port_indices is None:
            port_indices = range(self.num_input_ports)

        # Some blocks are hard-coded to have no inputs, so we should make
        # sure that a list full of None is not returned in that case. This
        # happens if the callback signature is (time, state, **parameters)
        # instead of the more general (time, state, *inputs, **parameters)
        if port_indices == []:
            return []

        inputs = []
        for i in range(self.num_input_ports):
            u_i = self.eval_input(context, i) if i in port_indices else None
            inputs.append(u_i)

        return inputs

    def _eval_input_port(self, context: ContextBase, port_index: int) -> Array:
        """Evaluate an upstream input port given the _root_ context.

        Intended for internal use as a callback function. Users and developers
        should typically call `eval_input` in order to get this information. That
        method will call the callback function associated with the input port,
        which will have a reference to this method.

        Args:
            context (ContextBase): root context for this system
            port_index (int): index of the input port to evaluate on the target system

        Returns:
            Array: current input values
        """
        # A helper function to evaluate an upstream input port given the _root_ context.

        port_locator = self.input_ports[port_index].locator

        if self.parent is None:
            # This is currently the root system.  Typically root input ports should not be evaluated,
            #  but we can get here during subsystem construction (e.g. type inference).  In that case,
            #  we should just defer evaluation and rely on the graph analysis to determine that
            #  everything is connected correctly.
            # See https://jaxonomy.atlassian.net/browse/WC-51.
            # This should not happen during simulation or root context construction.
            logger.debug(
                "    ---> %s is the root system, deferring evaluation of %s[%s]",
                self.name,
                port_locator[0].name,
                port_locator[1],
            )
            raise UpstreamEvalError(port_locator=(self, "in", port_index))

        # The `eval_subsystem_input_port` method is only defined for Diagrams, but
        # the parent system is guaranteed to be a Diagram if this is not the root.
        # If it is the root, it should not have any (un-fixed) input ports.
        return self.parent.eval_subsystem_input_port(context, port_locator)

    #
    # Declaration utilities
    #
    def _next_input_port_name(self, name: str | None = None) -> str:
        """Automatically generate a unique name for the next input port."""
        if name is not None:
            assert name != ""
            return name
        return f"in_{self.num_input_ports}"

    def _next_output_port_name(self, name: str | None = None) -> str:
        """Automatically generate a unique name for the next output port."""
        if name is not None:
            assert name != ""
            return name
        return f"out_{self.num_output_ports}"

    def declare_input_port(
        self,
        name: str = None,
        prerequisites_of_calc: List[DependencyTicket] = None,
        units=None,
    ) -> int:
        """Add an input port to the system.

        Returns the corresponding index into the system input_port_indices list
        Note that this is different from the callbacks index - typically it
        will make more sense to retrieve via system.input_ports[port_index], but

        Args:
            name (str, optional): name of the new port. Defaults to None, which will
                use the default naming scheme for the system (e.g. "u_0")
            prerequisites_of_calc (List[DependencyTicket], optional): list of
                dependencies for the callback function. Defaults to None.
            units (Unit, optional): physical unit of the signal carried on this
                port (T-104 phase 1).  Default ``None`` is treated as
                ``dimensionless`` — existing diagrams that never declare a unit
                continue to connect to anything.

        Returns:
            int: port index of the newly created port in `input_ports`
        """
        port_index = self.num_input_ports
        port_name = self._next_input_port_name(name)

        for port in self.input_ports:
            assert (
                port.name != port_name
            ), f"System {self.name} already has an input port named {port.name}"

        _callback = partial(
            _input_port_eval_callback, owner=self, port_index=port_index
        )

        callback_index = len(self.callbacks)
        port = InputPort(
            callback=_callback,
            system=self,
            callback_index=callback_index,
            name=port_name,
            index=port_index,
            prerequisites_of_calc=prerequisites_of_calc,
        )
        # T-104 phase 1: stash the optional unit on the port instance.
        # Stored as a plain attribute (not a dataclass field) to keep the
        # zero-units default path byte-equivalent and to avoid disturbing
        # the existing PortBase __init__ argument ordering.
        port.units = units

        assert isinstance(port, InputPort)
        assert port.system is self
        assert port.name != ""

        # Check that name is unique
        for p in self.input_ports:
            assert (
                p.name != port.name
            ), f"System {self.name} already has an input port named {port.name}"

        self.input_port_indices.append(callback_index)
        self.callbacks.append(port)
        self._cached_input_ports.clear()

        return port_index

    def declare_output_port(
        self,
        callback: Callable,
        name: str = None,
        prerequisites_of_calc: List[DependencyTicket] = None,
        default_value: Array = None,
        event: DiscreteUpdateEvent = None,
        cache_index: int = None,
        units=None,
    ) -> int:
        """Add an output port to the system.

        This output port could represent any function of the context available to
        the system, so a callback function is required.  This function should have
        the form
            `callback(context: ContextBase) -> Array`
        SystemBase implementations have some specific convenience wrappers, e.g.:
            `LeafSystem.declare_continuous_state_output`
            `Diagram.export_output`

        Common cases are:
        - Feedthrough blocks: gather inputs and return some function of the
            inputs (e.g. a gain)
        - Stateful blocks: use LeafSystem.declare_(...)_state_output_port to
            return the value of a particular state
        - Diagrams: create and export a diagram-level port to the parent system using
            the callback function associated with the system-level port

        Returns the corresponding index into the system output_port_indices list
        Note that this is different from the callbacks index - typically it
        will make more sense to retrieve via system.output_ports[port_index].

        Args:
            callback (Callable, optional): computes the value of the output port given
                the root context.
            name (str, optional): name of the new port. Defaults to None, which will
                use the default naming scheme for the system (e.g. "y_0")
            prerequisites_of_calc (List[DependencyTicket], optional): list of
                dependencies for the callback function. Defaults to None, which will
                use the default dependencies for the system (all sources).  This may
                conservatively flag the system as having algebraic loops, so it is
                better to be specific here when possible.  This is done automatically
                in the wrapper functions like `LeafSystem.declare_(...)_output_port`
            default_value (Array, optional): A default array-like value used to seed
                the context and perform type inference, when this is known up front.
                Defaults to None, which will use information propagation through the
                graph along with type promotion to determine an appropriate value.
            event (DiscreteUpdateEvent, optional): A discrete update event associated
                with this output port that will periodically refresh the value that
                will be returned by the callback function. This makes the port act as
                a sample-and-hold rather than a direct function evaluation.
            cache_index (int, optional): Index into the cache state component
                corresponding to the output port result, if the output port is of
                periodically-updated sample-and-hold type.

        Returns:
            int: port index of the newly created port
        """
        port_index = self.num_output_ports
        port_name = self._next_output_port_name(name)

        for port in self.output_ports:
            assert (
                port.name != port_name
            ), f"System {self.name} already has an output port named {port.name}"

        if prerequisites_of_calc is None:
            prerequisites_of_calc = [DependencyTicket.all_sources]

        callback_index = len(self.callbacks)
        port = OutputPort(
            callback,
            system=self,
            callback_index=callback_index,
            name=port_name,
            index=port_index,
            prerequisites_of_calc=prerequisites_of_calc,
            default_value=default_value,
            event=event,
            cache_index=cache_index,
        )
        # T-104 phase 1: see declare_input_port for the rationale.
        port.units = units

        assert isinstance(port, OutputPort)
        assert port.system is self
        assert port.name != ""

        # Check that name is unique
        for p in self.output_ports:
            assert (
                p.name != port.name
            ), f"System {self.name} already has an output port named {port.name}"

        logger.debug("Adding output port %s to %s", port, self.name)
        self.output_port_indices.append(callback_index)
        self.callbacks.append(port)
        self._cached_output_ports.clear()

        logger.debug(
            "    ---> %s now has %s output ports: %s",
            self.name,
            len(self.output_ports),
            self.output_ports,
        )
        logger.debug(
            "    ---> %s now has %s cache sources: %s",
            self.name,
            len(self.callbacks),
            self.callbacks,
        )

        return port_index

    def configure_output_port(
        self,
        port_index: int,
        callback: Callable,
        prerequisites_of_calc: List[DependencyTicket] = None,
        default_value: Array = None,
        event: DiscreteUpdateEvent = None,
        cache_index: int = None,
    ):
        """Configure an output port of the system.

        See `declare_output_port` for a description of the arguments.

        Args:
            port_index (int): index of the output port to configure

        Returns:
            None
        """

        if prerequisites_of_calc is None:
            prerequisites_of_calc = [DependencyTicket.all_sources]

        port = self.output_ports[port_index]
        port.port_index = port_index
        port._callback = callback
        port.prerequisites_of_calc = prerequisites_of_calc
        port.default_value = default_value
        port.event = event
        port.cache_index = cache_index
        self.callbacks[port.callback_index] = port
        self._cached_output_ports.clear()

        logger.debug(
            "    ---> %s now has %s output ports: %s",
            self.name,
            len(self.output_ports),
            self.output_ports,
        )
        logger.debug(
            "    ---> %s now has %s cache sources: %s",
            self.name,
            len(self.callbacks),
            self.callbacks,
        )

    @abc.abstractmethod
    def get_feedthrough(self) -> List[Tuple[int, int]]:
        """Determine pairs of direct feedthrough ports for this system.

        By default, the algorithm relies on the dependency tracking system to determine
        feedthrough, but this can be overridden by implementing this method directly
        in a subclass, for instance if the automatic dependency tracking is too
        conservative in determining feedthrough.

        Returns:
            List[Tuple[int, int]]:
                A list of tuples (u, v) indicating that output port v has a direct
                dependency on input port u, resulting in a feedthrough path in the system.
                The indices u and v correspond to the indices of the input and output
                ports in the system's input and output port lists.
        """
        pass

    #
    # Initialization
    #
    def create_context(self, **kwargs) -> ContextBase:
        """Create a new context for this system.

        The context will contain all variable information used in
        simulation/analysis/optimization, such as state and parameters.

        Returns:
            ContextBase: new context for this system
        """
        return self.context_factory(**kwargs)

    def check_types(self, context: ContextBase, error_collector: ErrorCollector = None):
        """Perform any system-specific static analysis."""
        pass

    @abc.abstractproperty
    def context_factory(self) -> ContextFactory:
        """Factory object for creating contexts for this system.

        Should not be called directly - use `system.create_context` instead.
        """
        pass

    @property
    def dependency_graph(self) -> DependencyGraph:
        """Retrieve (or create if necessary) the dependency graph for this system."""
        return self._dependency_graph

    @abc.abstractproperty
    def dependency_graph_factory(self) -> DependencyGraphFactory:
        """Factory object for creating dependency graphs for this system.

        Should not be called directly - use `system.create_dependency_graph` instead.
        """
        pass

    def create_dependency_graph(self):
        """Create a dependency graph for this system."""
        self._dependency_graph = self.dependency_graph_factory()

    def initialize_static_data(self, context: ContextBase) -> ContextBase:
        """Initialize any context data that has to be done after context creation.

        Use this to define custom auxiliary data or type inference that doesn't
        get traced by JAX. See the `ZeroOrderHold` implementation for an example.
        Since this is only applied during context initialization, it is allowed to
        modify the context directly (or the system itself).

        Typically this should not be called outside of the ContextFactory.

        Args:
            context (ContextBase): partially initialized context for this system.
        """
        return context

    @property
    def ports(self) -> dict[str, PortBase]:
        """Dictionary of all ports in this system, indexed by name"""
        return {port.name: port for port in self.input_ports + self.output_ports}

    # Convenience functions for errors and UI logs

    @property
    def name_path(self) -> list[str]:
        """Get the human-readable path to this system. None if some names are not set."""
        if self.parent is None:
            return [self.name]  # Likely to be 'root'
        if self.parent.parent is None:
            return [self.name]  # top-level block
        return self.parent.name_path + [self.name]

    @property
    def name_path_str(self) -> str:
        """Get the human-readable path to this system as a string."""
        return ".".join(self.name_path)

    @property
    def ui_id_path(self) -> Union[list[str], None]:
        """Get the uuid node path to this system. None if some IDs are not set."""
        if self.ui_id is None:
            return None
        if self.parent is None:
            return [self.ui_id]
        if self.parent.parent is None:
            return [self.ui_id]  # top-level block
        parent_path = self.parent.ui_id_path
        if parent_path is None:
            return None
        return parent_path + [self.ui_id]

    def declare_static_parameters(self, **params):
        """Declare a set of static parameters for the system.

        These parameters are not JAX-traceable and therefore can't be optimized.

        Examples of static parameters include booleans, strings, parameters
        used in shapes, etc.

        The args should be a dict of name-value pairs, where the values are either
        strings, bool, arrays, or Parameters.

        Typical usage:

        ```python
        class MyBlock(LeafSystem):
            def __init__(self, param1=True, param2=1.0):
                super().__init__()
                self.declare_static_parameters(param1=param1, param2=param2)
        ```
        """
        for name, value in params.items():
            if name in self.dynamic_parameters:
                raise BlockParameterError(
                    "Parameter already declared as dynamic parameter",
                    system=self,
                    parameter_name=name,
                )
            if isinstance(value, list):
                self._static_parameters[name] = Parameter(
                    value=np.array(value),
                    system=self,
                    is_static=True,
                )
            else:
                self._static_parameters[name] = Parameter(
                    value=value, system=self, is_static=True
                )

    def declare_static_parameter(self, name, value):
        """Declare a single static parameter for the system.

        This is a convenience function for declaring a single static parameter.

        Args:
            name (str): name of the parameter
            value (Union[Array, Parameter]): value of the parameter
        """
        self.declare_static_parameters(**{name: value})

    def declare_dynamic_parameter(
        self,
        name: str,
        default_value: Array | Parameter = None,
        shape: ShapeLike = None,
        dtype: DTypeLike = None,
        as_array: bool = True,
    ):
        """Declare a numeric parameter for the system.

        Parameters are declared in the system and accessed through the context to
        maintain separation of data ownership. This method creates an entry in the
        system's dynamic_parameters, recording the name, default value, and dependency
        ticket for later reference.

        The default value will be used to initialize the context, so it
        will also serve as the initial value unless explicitly overridden. In the
        simplest cases, parameters could be stored as attributes of the LeafSystem,
        but declaring them has the advantage of moving the values to the context,
        allowing them to be traced by JAX rather than stored as static data. This
        means they can be differentiated, vmapped, or otherwise modified without
        re-compiling the simulation.

        Args:
            name (str):
                The name of the parameter.
            default_value (Union[Array, Parameter], optional):
                The default value of the parameter. Parameters are used
                primarily internally for serialization and should not normally need
                to be used directly when implementing LeafSystems. Defaults to None.
            shape (ShapeLike, optional):
                The shape of the parameter. Defaults to None.
            dtype (DTypeLike, optional):
                The data type of the parameter. Defaults to None.
            as_array (bool, optional):
                If True, treat the default_value as an array-like (cast if necessary).
                Otherwise, it will be stored as the default state without modification.

        Raises:
            AssertionError:
                If the parameter with the given name is already declared.

        Notes:
            (1) Only one of `shape` and `default_value` should be provided. If
            `default_value` is provided, it will be used as the initial value of the
            continuous state. If `shape` is provided, the initial value will be a
            zero vector of the given shape and specified dtype.
        """
        # assert (
        #     name not in self._dynamic_parameters
        # ), f"Parameter {name} already declared"

        if name in self.static_parameters:
            raise BlockParameterError(
                "Parameter already declared as static parameter",
                system=self,
                parameter_name=name,
            )

        try:
            if isinstance(default_value, Parameter):
                self._dynamic_parameters[name] = Parameter(
                    value=default_value,
                    dtype=dtype,
                    shape=shape,
                    system=self,
                    as_array=as_array,
                )
            else:
                if as_array:
                    default_value = utils.make_array(
                        default_value, dtype=dtype, shape=shape
                    )
                self._dynamic_parameters[name] = Parameter(
                    value=default_value,
                    dtype=dtype,
                    shape=shape,
                    system=self,
                )

            logger.debug(
                "Adding parameter %s to %s with default: %s",
                name,
                self.name,
                default_value,
            )

        except Exception as e:
            traceback.print_exc()
            raise BlockParameterError(
                "Error declaring parameter",
                system=self,
                parameter_name=name,
            ) from e

    def get_parameter(self, name: str):
        """Get a parameter value by name.

        Checks dynamic parameters first, then static parameters. Values are
        returned as concrete array-like / Python scalars via
        :meth:`Parameter.unwrap`.

        Args:
            name: Parameter name on this system (not a dotted path).

        Raises:
            KeyError: If ``name`` is not found; the message lists available names.
        """
        if name in self._dynamic_parameters:
            return Parameter.unwrap(self._dynamic_parameters[name])
        if name in self._static_parameters:
            return Parameter.unwrap(self._static_parameters[name])
        available = sorted(
            {*self._static_parameters.keys(), *self._dynamic_parameters.keys()}
        )
        raise KeyError(
            f"Parameter {name!r} not found on {self.name!r}. Available: {available}"
        )

    def list_parameters(self) -> dict:
        """Return all parameters as a flat ``{name: value}`` mapping.

        Dynamic parameters override static ones when names collide. Values are
        unwrapped the same way as :meth:`get_parameter`.
        """
        merged = {**self._static_parameters, **self._dynamic_parameters}
        return {k: Parameter.unwrap(v) for k, v in merged.items()}

    @property
    def has_dirty_static_parameters(self) -> bool:
        """Check if any static parameters have been modified."""
        return any(param.is_dirty for param in self.static_parameters.values())

`dependency_graph` `property`

Retrieve (or create if necessary) the dependency graph for this system.

`has_dirty_static_parameters` `property`

Check if any static parameters have been modified.

`has_feedthrough_side_effects` `abstractmethod` `property`

Check if the system includes any feedthrough calls to io_callback.

`has_mass_matrix` `abstractmethod` `property`

Returns True if any component of the system has a nontrivial mass matrix.

`has_ode_side_effects` `abstractmethod` `property`

Check if the ODE RHS for the system includes any calls to io_callback.

`mass_matrix` `abstractmethod` `property`

Mass matrix for this system.

Returns PyTree-structured data where each leaf is an (n, n) array. This is used for implicit integration methods (currently only BDF).

`name_path` `property`

Get the human-readable path to this system. None if some names are not set.

`name_path_str` `property`

Get the human-readable path to this system as a string.

`ports` `property`

Dictionary of all ports in this system, indexed by name

`root` `property`

Get the root system of the current system.

`sorted_callbacks` `property`

Sort and return the callbacks for this system.

`ui_id_path` `property`

Get the uuid node path to this system. None if some IDs are not set.

`deepcopy(memo)`

Deep-copy while keeping partially constructed copies hashable.

Subsystems reference themselves via callbacks; the default deepcopy order can call :meth:__hash__ (via dict/set operations) before system_id exists on the copy. Assign a new system_id immediately after memo registration.

Source code in jaxonomy/framework/system_base.py

def __deepcopy__(self, memo):
    """Deep-copy while keeping partially constructed copies hashable.

    Subsystems reference themselves via callbacks; the default deepcopy order can
    call :meth:`__hash__` (via dict/set operations) before ``system_id`` exists
    on the copy. Assign a new ``system_id`` immediately after memo registration.
    """
    cls = type(self)
    result = cls.__new__(cls)
    memo[id(self)] = result
    result.system_id = next_system_id()
    for key, value in self.__dict__.items():
        if key == "system_id":
            continue
        setattr(result, key, copy.deepcopy(value, memo))
    return result

`check_types(context, error_collector=None)`

Perform any system-specific static analysis.

Source code in jaxonomy/framework/system_base.py

def check_types(self, context: ContextBase, error_collector: ErrorCollector = None):
    """Perform any system-specific static analysis."""
    pass

`collect_inputs(context, port_indices=None)`

Collect all current inputs for this system.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	root context for this system	required
`port_indices`	`List[int]`	list of input port indices to collect. If None (default), will return values from all ports. Otherwise will return a list of length(num_input_ports), where the values are None for ports not in the list.	`None`

Returns:

Type	Description
`List[Array]`	List[Array]: list of all current input values

Source code in jaxonomy/framework/system_base.py

def collect_inputs(
    self, context: ContextBase, port_indices: list[int] = None
) -> List[Array]:
    """Collect all current inputs for this system.

    Args:
        context (ContextBase): root context for this system
        port_indices (List[int], optional): list of input port indices to collect.
            If None (default), will return values from all ports.  Otherwise will
            return a list of length(num_input_ports), where the values are None for
            ports not in the list.

    Returns:
        List[Array]: list of all current input values
    """
    if port_indices is None:
        port_indices = range(self.num_input_ports)

    # Some blocks are hard-coded to have no inputs, so we should make
    # sure that a list full of None is not returned in that case. This
    # happens if the callback signature is (time, state, **parameters)
    # instead of the more general (time, state, *inputs, **parameters)
    if port_indices == []:
        return []

    inputs = []
    for i in range(self.num_input_ports):
        u_i = self.eval_input(context, i) if i in port_indices else None
        inputs.append(u_i)

    return inputs

`configure_output_port(port_index, callback, prerequisites_of_calc=None, default_value=None, event=None, cache_index=None)`

Configure an output port of the system.

See declare_output_port for a description of the arguments.

Parameters:

Name	Type	Description	Default
`port_index`	`int`	index of the output port to configure	required

Returns:

Type	Description
	None

Source code in jaxonomy/framework/system_base.py

def configure_output_port(
    self,
    port_index: int,
    callback: Callable,
    prerequisites_of_calc: List[DependencyTicket] = None,
    default_value: Array = None,
    event: DiscreteUpdateEvent = None,
    cache_index: int = None,
):
    """Configure an output port of the system.

    See `declare_output_port` for a description of the arguments.

    Args:
        port_index (int): index of the output port to configure

    Returns:
        None
    """

    if prerequisites_of_calc is None:
        prerequisites_of_calc = [DependencyTicket.all_sources]

    port = self.output_ports[port_index]
    port.port_index = port_index
    port._callback = callback
    port.prerequisites_of_calc = prerequisites_of_calc
    port.default_value = default_value
    port.event = event
    port.cache_index = cache_index
    self.callbacks[port.callback_index] = port
    self._cached_output_ports.clear()

    logger.debug(
        "    ---> %s now has %s output ports: %s",
        self.name,
        len(self.output_ports),
        self.output_ports,
    )
    logger.debug(
        "    ---> %s now has %s cache sources: %s",
        self.name,
        len(self.callbacks),
        self.callbacks,
    )

`context_factory()`

Factory object for creating contexts for this system.

Should not be called directly - use system.create_context instead.

Source code in jaxonomy/framework/system_base.py

@abc.abstractproperty
def context_factory(self) -> ContextFactory:
    """Factory object for creating contexts for this system.

    Should not be called directly - use `system.create_context` instead.
    """
    pass

`create_context(**kwargs)`

Create a new context for this system.

The context will contain all variable information used in simulation/analysis/optimization, such as state and parameters.

Returns:

Name	Type	Description
`ContextBase`	`ContextBase`	new context for this system

Source code in jaxonomy/framework/system_base.py

def create_context(self, **kwargs) -> ContextBase:
    """Create a new context for this system.

    The context will contain all variable information used in
    simulation/analysis/optimization, such as state and parameters.

    Returns:
        ContextBase: new context for this system
    """
    return self.context_factory(**kwargs)

`create_dependency_graph()`

Create a dependency graph for this system.

Source code in jaxonomy/framework/system_base.py

def create_dependency_graph(self):
    """Create a dependency graph for this system."""
    self._dependency_graph = self.dependency_graph_factory()

`declare_dynamic_parameter(name, default_value=None, shape=None, dtype=None, as_array=True)`

Declare a numeric parameter for the system.

Parameters are declared in the system and accessed through the context to maintain separation of data ownership. This method creates an entry in the system's dynamic_parameters, recording the name, default value, and dependency ticket for later reference.

The default value will be used to initialize the context, so it will also serve as the initial value unless explicitly overridden. In the simplest cases, parameters could be stored as attributes of the LeafSystem, but declaring them has the advantage of moving the values to the context, allowing them to be traced by JAX rather than stored as static data. This means they can be differentiated, vmapped, or otherwise modified without re-compiling the simulation.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the parameter.	required
`default_value`	`Union[Array, Parameter]`	The default value of the parameter. Parameters are used primarily internally for serialization and should not normally need to be used directly when implementing LeafSystems. Defaults to None.	`None`
`shape`	`ShapeLike`	The shape of the parameter. Defaults to None.	`None`
`dtype`	`DTypeLike`	The data type of the parameter. Defaults to None.	`None`
`as_array`	`bool`	If True, treat the default_value as an array-like (cast if necessary). Otherwise, it will be stored as the default state without modification.	`True`

Raises:

Type	Description
`AssertionError`	If the parameter with the given name is already declared.

Notes

(1) Only one of shape and default_value should be provided. If default_value is provided, it will be used as the initial value of the continuous state. If shape is provided, the initial value will be a zero vector of the given shape and specified dtype.

Source code in jaxonomy/framework/system_base.py

def declare_dynamic_parameter(
    self,
    name: str,
    default_value: Array | Parameter = None,
    shape: ShapeLike = None,
    dtype: DTypeLike = None,
    as_array: bool = True,
):
    """Declare a numeric parameter for the system.

    Parameters are declared in the system and accessed through the context to
    maintain separation of data ownership. This method creates an entry in the
    system's dynamic_parameters, recording the name, default value, and dependency
    ticket for later reference.

    The default value will be used to initialize the context, so it
    will also serve as the initial value unless explicitly overridden. In the
    simplest cases, parameters could be stored as attributes of the LeafSystem,
    but declaring them has the advantage of moving the values to the context,
    allowing them to be traced by JAX rather than stored as static data. This
    means they can be differentiated, vmapped, or otherwise modified without
    re-compiling the simulation.

    Args:
        name (str):
            The name of the parameter.
        default_value (Union[Array, Parameter], optional):
            The default value of the parameter. Parameters are used
            primarily internally for serialization and should not normally need
            to be used directly when implementing LeafSystems. Defaults to None.
        shape (ShapeLike, optional):
            The shape of the parameter. Defaults to None.
        dtype (DTypeLike, optional):
            The data type of the parameter. Defaults to None.
        as_array (bool, optional):
            If True, treat the default_value as an array-like (cast if necessary).
            Otherwise, it will be stored as the default state without modification.

    Raises:
        AssertionError:
            If the parameter with the given name is already declared.

    Notes:
        (1) Only one of `shape` and `default_value` should be provided. If
        `default_value` is provided, it will be used as the initial value of the
        continuous state. If `shape` is provided, the initial value will be a
        zero vector of the given shape and specified dtype.
    """
    # assert (
    #     name not in self._dynamic_parameters
    # ), f"Parameter {name} already declared"

    if name in self.static_parameters:
        raise BlockParameterError(
            "Parameter already declared as static parameter",
            system=self,
            parameter_name=name,
        )

    try:
        if isinstance(default_value, Parameter):
            self._dynamic_parameters[name] = Parameter(
                value=default_value,
                dtype=dtype,
                shape=shape,
                system=self,
                as_array=as_array,
            )
        else:
            if as_array:
                default_value = utils.make_array(
                    default_value, dtype=dtype, shape=shape
                )
            self._dynamic_parameters[name] = Parameter(
                value=default_value,
                dtype=dtype,
                shape=shape,
                system=self,
            )

        logger.debug(
            "Adding parameter %s to %s with default: %s",
            name,
            self.name,
            default_value,
        )

    except Exception as e:
        traceback.print_exc()
        raise BlockParameterError(
            "Error declaring parameter",
            system=self,
            parameter_name=name,
        ) from e

`declare_input_port(name=None, prerequisites_of_calc=None, units=None)`

Add an input port to the system.

Returns the corresponding index into the system input_port_indices list Note that this is different from the callbacks index - typically it will make more sense to retrieve via system.input_ports[port_index], but

Parameters:

Name	Type	Description	Default
`name`	`str`	name of the new port. Defaults to None, which will use the default naming scheme for the system (e.g. "u_0")	`None`
`prerequisites_of_calc`	`List[DependencyTicket]`	list of dependencies for the callback function. Defaults to None.	`None`
`units`	`Unit`	physical unit of the signal carried on this port (T-104 phase 1). Default `None` is treated as `dimensionless` — existing diagrams that never declare a unit continue to connect to anything.	`None`

Returns:

Name	Type	Description
`int`	`int`	port index of the newly created port in `input_ports`

Source code in jaxonomy/framework/system_base.py

def declare_input_port(
    self,
    name: str = None,
    prerequisites_of_calc: List[DependencyTicket] = None,
    units=None,
) -> int:
    """Add an input port to the system.

    Returns the corresponding index into the system input_port_indices list
    Note that this is different from the callbacks index - typically it
    will make more sense to retrieve via system.input_ports[port_index], but

    Args:
        name (str, optional): name of the new port. Defaults to None, which will
            use the default naming scheme for the system (e.g. "u_0")
        prerequisites_of_calc (List[DependencyTicket], optional): list of
            dependencies for the callback function. Defaults to None.
        units (Unit, optional): physical unit of the signal carried on this
            port (T-104 phase 1).  Default ``None`` is treated as
            ``dimensionless`` — existing diagrams that never declare a unit
            continue to connect to anything.

    Returns:
        int: port index of the newly created port in `input_ports`
    """
    port_index = self.num_input_ports
    port_name = self._next_input_port_name(name)

    for port in self.input_ports:
        assert (
            port.name != port_name
        ), f"System {self.name} already has an input port named {port.name}"

    _callback = partial(
        _input_port_eval_callback, owner=self, port_index=port_index
    )

    callback_index = len(self.callbacks)
    port = InputPort(
        callback=_callback,
        system=self,
        callback_index=callback_index,
        name=port_name,
        index=port_index,
        prerequisites_of_calc=prerequisites_of_calc,
    )
    # T-104 phase 1: stash the optional unit on the port instance.
    # Stored as a plain attribute (not a dataclass field) to keep the
    # zero-units default path byte-equivalent and to avoid disturbing
    # the existing PortBase __init__ argument ordering.
    port.units = units

    assert isinstance(port, InputPort)
    assert port.system is self
    assert port.name != ""

    # Check that name is unique
    for p in self.input_ports:
        assert (
            p.name != port.name
        ), f"System {self.name} already has an input port named {port.name}"

    self.input_port_indices.append(callback_index)
    self.callbacks.append(port)
    self._cached_input_ports.clear()

    return port_index

`declare_output_port(callback, name=None, prerequisites_of_calc=None, default_value=None, event=None, cache_index=None, units=None)`

Add an output port to the system.

This output port could represent any function of the context available to the system, so a callback function is required. This function should have the form callback(context: ContextBase) -> Array SystemBase implementations have some specific convenience wrappers, e.g.: LeafSystem.declare_continuous_state_output Diagram.export_output

Common cases are: - Feedthrough blocks: gather inputs and return some function of the inputs (e.g. a gain) - Stateful blocks: use LeafSystem.declare_(...)_state_output_port to return the value of a particular state - Diagrams: create and export a diagram-level port to the parent system using the callback function associated with the system-level port

Returns the corresponding index into the system output_port_indices list Note that this is different from the callbacks index - typically it will make more sense to retrieve via system.output_ports[port_index].

Parameters:

Name	Type	Description	Default
`callback`	`Callable`	computes the value of the output port given the root context.	required
`name`	`str`	name of the new port. Defaults to None, which will use the default naming scheme for the system (e.g. "y_0")	`None`
`prerequisites_of_calc`	`List[DependencyTicket]`	list of dependencies for the callback function. Defaults to None, which will use the default dependencies for the system (all sources). This may conservatively flag the system as having algebraic loops, so it is better to be specific here when possible. This is done automatically in the wrapper functions like `LeafSystem.declare_(...)_output_port`	`None`
`default_value`	`Array`	A default array-like value used to seed the context and perform type inference, when this is known up front. Defaults to None, which will use information propagation through the graph along with type promotion to determine an appropriate value.	`None`
`event`	`DiscreteUpdateEvent`	A discrete update event associated with this output port that will periodically refresh the value that will be returned by the callback function. This makes the port act as a sample-and-hold rather than a direct function evaluation.	`None`
`cache_index`	`int`	Index into the cache state component corresponding to the output port result, if the output port is of periodically-updated sample-and-hold type.	`None`

Returns:

Name	Type	Description
`int`	`int`	port index of the newly created port

Source code in jaxonomy/framework/system_base.py

def declare_output_port(
    self,
    callback: Callable,
    name: str = None,
    prerequisites_of_calc: List[DependencyTicket] = None,
    default_value: Array = None,
    event: DiscreteUpdateEvent = None,
    cache_index: int = None,
    units=None,
) -> int:
    """Add an output port to the system.

    This output port could represent any function of the context available to
    the system, so a callback function is required.  This function should have
    the form
        `callback(context: ContextBase) -> Array`
    SystemBase implementations have some specific convenience wrappers, e.g.:
        `LeafSystem.declare_continuous_state_output`
        `Diagram.export_output`

    Common cases are:
    - Feedthrough blocks: gather inputs and return some function of the
        inputs (e.g. a gain)
    - Stateful blocks: use LeafSystem.declare_(...)_state_output_port to
        return the value of a particular state
    - Diagrams: create and export a diagram-level port to the parent system using
        the callback function associated with the system-level port

    Returns the corresponding index into the system output_port_indices list
    Note that this is different from the callbacks index - typically it
    will make more sense to retrieve via system.output_ports[port_index].

    Args:
        callback (Callable, optional): computes the value of the output port given
            the root context.
        name (str, optional): name of the new port. Defaults to None, which will
            use the default naming scheme for the system (e.g. "y_0")
        prerequisites_of_calc (List[DependencyTicket], optional): list of
            dependencies for the callback function. Defaults to None, which will
            use the default dependencies for the system (all sources).  This may
            conservatively flag the system as having algebraic loops, so it is
            better to be specific here when possible.  This is done automatically
            in the wrapper functions like `LeafSystem.declare_(...)_output_port`
        default_value (Array, optional): A default array-like value used to seed
            the context and perform type inference, when this is known up front.
            Defaults to None, which will use information propagation through the
            graph along with type promotion to determine an appropriate value.
        event (DiscreteUpdateEvent, optional): A discrete update event associated
            with this output port that will periodically refresh the value that
            will be returned by the callback function. This makes the port act as
            a sample-and-hold rather than a direct function evaluation.
        cache_index (int, optional): Index into the cache state component
            corresponding to the output port result, if the output port is of
            periodically-updated sample-and-hold type.

    Returns:
        int: port index of the newly created port
    """
    port_index = self.num_output_ports
    port_name = self._next_output_port_name(name)

    for port in self.output_ports:
        assert (
            port.name != port_name
        ), f"System {self.name} already has an output port named {port.name}"

    if prerequisites_of_calc is None:
        prerequisites_of_calc = [DependencyTicket.all_sources]

    callback_index = len(self.callbacks)
    port = OutputPort(
        callback,
        system=self,
        callback_index=callback_index,
        name=port_name,
        index=port_index,
        prerequisites_of_calc=prerequisites_of_calc,
        default_value=default_value,
        event=event,
        cache_index=cache_index,
    )
    # T-104 phase 1: see declare_input_port for the rationale.
    port.units = units

    assert isinstance(port, OutputPort)
    assert port.system is self
    assert port.name != ""

    # Check that name is unique
    for p in self.output_ports:
        assert (
            p.name != port.name
        ), f"System {self.name} already has an output port named {port.name}"

    logger.debug("Adding output port %s to %s", port, self.name)
    self.output_port_indices.append(callback_index)
    self.callbacks.append(port)
    self._cached_output_ports.clear()

    logger.debug(
        "    ---> %s now has %s output ports: %s",
        self.name,
        len(self.output_ports),
        self.output_ports,
    )
    logger.debug(
        "    ---> %s now has %s cache sources: %s",
        self.name,
        len(self.callbacks),
        self.callbacks,
    )

    return port_index

`declare_static_parameter(name, value)`

Declare a single static parameter for the system.

This is a convenience function for declaring a single static parameter.

Parameters:

Name	Type	Description	Default
`name`	`str`	name of the parameter	required
`value`	`Union[Array, Parameter]`	value of the parameter	required

Source code in jaxonomy/framework/system_base.py

def declare_static_parameter(self, name, value):
    """Declare a single static parameter for the system.

    This is a convenience function for declaring a single static parameter.

    Args:
        name (str): name of the parameter
        value (Union[Array, Parameter]): value of the parameter
    """
    self.declare_static_parameters(**{name: value})

`declare_static_parameters(**params)`

Declare a set of static parameters for the system.

These parameters are not JAX-traceable and therefore can't be optimized.

Examples of static parameters include booleans, strings, parameters used in shapes, etc.

The args should be a dict of name-value pairs, where the values are either strings, bool, arrays, or Parameters.

Typical usage:

class MyBlock(LeafSystem):
    def __init__(self, param1=True, param2=1.0):
        super().__init__()
        self.declare_static_parameters(param1=param1, param2=param2)

Source code in jaxonomy/framework/system_base.py

def declare_static_parameters(self, **params):
    """Declare a set of static parameters for the system.

    These parameters are not JAX-traceable and therefore can't be optimized.

    Examples of static parameters include booleans, strings, parameters
    used in shapes, etc.

    The args should be a dict of name-value pairs, where the values are either
    strings, bool, arrays, or Parameters.

    Typical usage:

    ```python
    class MyBlock(LeafSystem):
        def __init__(self, param1=True, param2=1.0):
            super().__init__()
            self.declare_static_parameters(param1=param1, param2=param2)
    ```
    """
    for name, value in params.items():
        if name in self.dynamic_parameters:
            raise BlockParameterError(
                "Parameter already declared as dynamic parameter",
                system=self,
                parameter_name=name,
            )
        if isinstance(value, list):
            self._static_parameters[name] = Parameter(
                value=np.array(value),
                system=self,
                is_static=True,
            )
        else:
            self._static_parameters[name] = Parameter(
                value=value, system=self, is_static=True
            )

`dependency_graph_factory()`

Factory object for creating dependency graphs for this system.

Should not be called directly - use system.create_dependency_graph instead.

Source code in jaxonomy/framework/system_base.py

@abc.abstractproperty
def dependency_graph_factory(self) -> DependencyGraphFactory:
    """Factory object for creating dependency graphs for this system.

    Should not be called directly - use `system.create_dependency_graph` instead.
    """
    pass

`determine_active_guards(context)` `abstractmethod`

Determine active guards for zero-crossing events.

This method is responsible for evaluating and determining which zero-crossing events are active based on the current system mode and other conditions. This can be overridden to flag active/inactive guards on a block-specific basis, for instance in a StateMachine-type block. By default all guards are marked active at this point unless the zero-crossing event was declared with a non-default start_mode, in which case the guard is activated conditionally on the current mode.

For example, in a system with finite state transitions, where a transition from mode A to mode B is triggered by a guard function g_AB and the inverse transition is triggered by a guard function g_BA, this function would activate g_AB if the system is in mode A and g_BA if the system is in mode B. The other guard function would be inactive. If the zero-crossing event is not associated with a starting mode, it is considered to be always active.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	The root context containing the overall state and parameters.	required

Returns:

Name	Type	Description
`EventCollection`	`EventCollection`	A collection of zero-crossing events with active/inactive status updated based on the current system mode and other conditions.

Source code in jaxonomy/framework/system_base.py

@abc.abstractmethod
def determine_active_guards(self, context: ContextBase) -> EventCollection:
    """Determine active guards for zero-crossing events.

    This method is responsible for evaluating and determining which
    zero-crossing events are active based on the current system mode
    and other conditions.  This can be overridden to flag active/inactive
    guards on a block-specific basis, for instance in a StateMachine-type
    block. By default all guards are marked active at this point unless
    the zero-crossing event was declared with a non-default `start_mode`, in
    which case the guard is activated conditionally on the current mode.

    For example, in a system with finite state transitions, where a transition from
    mode A to mode B is triggered by a guard function g_AB and the inverse
    transition is triggered by a guard function g_BA, this function would activate
    g_AB if the system is in mode A and g_BA if the system is in mode B. The other
    guard function would be inactive.  If the zero-crossing event is not associated
    with a starting mode, it is considered to be always active.

    Args:
        context (ContextBase):
            The root context containing the overall state and parameters.

    Returns:
        EventCollection:
            A collection of zero-crossing events with active/inactive status
            updated based on the current system mode and other conditions.
    """
    pass

`eval_input(context, port_index=0)`

Get the input for a given port.

This works by evaluating the callback function associated with the port, which will "pull" the upstream output port values.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	root context for this system	required
`port_index`	`int`	index into `self.input_ports`, for example the value returned by `declare_input_port`. Defaults to 0.	`0`

Returns:

Name	Type	Description
`Array`	`Array`	current input values

Source code in jaxonomy/framework/system_base.py

def eval_input(self, context: ContextBase, port_index: int = 0) -> Array:
    """Get the input for a given port.

    This works by evaluating the callback function associated with the port, which
    will "pull" the upstream output port values.

    Args:
        context (ContextBase): root context for this system
        port_index (int, optional): index into `self.input_ports`, for example
            the value returned by `declare_input_port`. Defaults to 0.

    Returns:
        Array: current input values
    """
    return self.input_ports[port_index].eval(context)

`eval_time_derivatives(context)`

Evaluate the continuous time derivatives for this system.

Given the root context, evaluate the continuous time derivatives, which must have the same PyTree structure as the continuous state.

In principle, this can be overridden by custom implementations, but in general it is preferable to declare continuous states for LeafSystems using declare_continuous_state, which accepts a callback function that will be used to compute the derivatives. For Diagrams, the time derivatives are computed automatically using the callback functions for all child systems with continuous state.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	root context of this system	required

Returns:

Name	Type	Description
`StateComponent`	`StateComponent`	Continuous time derivatives for this system, or None if the system has no continuous state.

Source code in jaxonomy/framework/system_base.py

def eval_time_derivatives(self, context: ContextBase) -> StateComponent:
    """Evaluate the continuous time derivatives for this system.

    Given the _root_ context, evaluate the continuous time derivatives,
    which must have the same PyTree structure as the continuous state.

    In principle, this can be overridden by custom implementations, but
    in general it is preferable to declare continuous states for LeafSystems
    using `declare_continuous_state`, which accepts a callback function
    that will be used to compute the derivatives. For Diagrams, the time
    derivatives are computed automatically using the callback functions for
    all child systems with continuous state.

    Args:
        context (ContextBase): root context of this system

    Returns:
        StateComponent:
            Continuous time derivatives for this system, or None if the system
            has no continuous state.
    """
    return None

`eval_zero_crossing_updates(context, events)` `abstractmethod`

Evaluate reset maps associated with zero-crossing events.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	The context for the system, containing the current state and parameters.	required
`events`	`EventCollection`	The collection of events to be evaluated (for example zero-crossing or periodic events for this system).	required

Returns:

Name	Type	Description
`State`	`State`	The complete state with all updates applied.

Notes

(1) Following the Drake definition, "unrestricted" updates are allowed to modify any component of the state: continuous, discrete, or mode. These updates are evaluated in the order in which they were declared, so it is possible (but should be strictly avoided) for multiple events to modify the same state component at the same time.

Each update computes its results given the current state of the system (the "minus" values) and returns the updated state (the "plus" values). The update functions cannot access any information about the "plus" values of its own state or the state of any other block. This could change in the future but for now it ensures consistency with Drake's discrete semantices:

More specifically, since all unrestricted updates can modify the entire state, any time there are multiple unrestricted updates, the resulting states are ALWAYS in conflict. For example, suppose a system has two unrestricted updates, event1 and event2. At time t_n, event1 is active and event2 is inactive. First, event1 is evaluated, and the state is updated. Then event2 is evaluated, but the state is not updated. Which one is valid? Obviously, the event1 return is valid, but how do we communicate this to JAX? The situation is more complicated if both event1 and event2 happen to be active. In this case the states have to be "merged" somehow. In the worst case, these two will modify the same components of the state in different ways.

The implementation updates the state in a local copy of the context (since both are immutable). This allows multiple unrestricted updates, but leaves open the possibility of multiple active updates modifying the state in conflicting ways. This should strictly be avoided by the implementer of the LeafSystem. If it is at all unclear how to do this, it may be better to split the system into multiple blocks to be safe.

(2) The events are evaluated conditionally on being marked "active" (indicating that their guard function triggered), so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns.

Source code in jaxonomy/framework/system_base.py

@abc.abstractmethod
def eval_zero_crossing_updates(
    self,
    context: ContextBase,
    events: EventCollection,
) -> State:
    """Evaluate reset maps associated with zero-crossing events.

    Args:
        context (ContextBase):
            The context for the system, containing the current state and parameters.
        events (EventCollection):
            The collection of events to be evaluated (for example zero-crossing or
            periodic events for this system).

    Returns:
        State: The complete state with all updates applied.

    Notes:
        (1) Following the Drake definition, "unrestricted" updates are allowed to
        modify any component of the state: continuous, discrete, or mode.  These
        updates are evaluated in the order in which they were declared, so it is
        _possible_ (but should be strictly avoided) for multiple events to modify the
        same state component at the same time.

        Each update computes its results given the _current_ state of the system
        (the "minus" values) and returns the _updated_ state (the "plus" values).
        The update functions cannot access any information about the "plus" values of
        its own state or the state of any other block.  This could change in the future
        but for now it ensures consistency with Drake's discrete semantices:

        More specifically, since all unrestricted updates can modify the entire state,
        any time there are multiple unrestricted updates, the resulting states are
        ALWAYS in conflict.  For example, suppose a system has two unrestricted
        updates, `event1` and `event2`.  At time t_n, `event1` is active and `event2`
        is inactive.  First, `event1` is evaluated, and the state is updated.  Then
        `event2` is evaluated, but the state is not updated.  Which one is valid?
        Obviously, the `event1` return is valid, but how do we communicate this to JAX?
        The situation is more complicated if both `event1` and `event2` happen to be
        active.  In this case the states have to be "merged" somehow.  In the worst
        case, these two will modify the same components of the state in different ways.

        The implementation updates the state in a local copy of the context (since both
        are immutable).  This allows multiple unrestricted updates, but leaves open the
        possibility of multiple active updates modifying the state in conflicting ways.
        This should strictly be avoided by the implementer of the LeafSystem.  If it is
        at all unclear how to do this, it may be better to split the system into
        multiple blocks to be safe.

        (2) The events are evaluated conditionally on being marked "active"
        (indicating that their guard function triggered), so the entire event
        collection can be passed without filtering to active events. This is necessary
        to make the function calls work with JAX tracing, which do not allow for
        variable-sized arguments or returns.
    """
    pass

`get_feedthrough()` `abstractmethod`

Determine pairs of direct feedthrough ports for this system.

By default, the algorithm relies on the dependency tracking system to determine feedthrough, but this can be overridden by implementing this method directly in a subclass, for instance if the automatic dependency tracking is too conservative in determining feedthrough.

Returns:

Type	Description
`List[Tuple[int, int]]`	List[Tuple[int, int]]: A list of tuples (u, v) indicating that output port v has a direct dependency on input port u, resulting in a feedthrough path in the system. The indices u and v correspond to the indices of the input and output ports in the system's input and output port lists.

Source code in jaxonomy/framework/system_base.py

@abc.abstractmethod
def get_feedthrough(self) -> List[Tuple[int, int]]:
    """Determine pairs of direct feedthrough ports for this system.

    By default, the algorithm relies on the dependency tracking system to determine
    feedthrough, but this can be overridden by implementing this method directly
    in a subclass, for instance if the automatic dependency tracking is too
    conservative in determining feedthrough.

    Returns:
        List[Tuple[int, int]]:
            A list of tuples (u, v) indicating that output port v has a direct
            dependency on input port u, resulting in a feedthrough path in the system.
            The indices u and v correspond to the indices of the input and output
            ports in the system's input and output port lists.
    """
    pass

`get_input_port(name)`

Retrieve a specific input port by name.

Source code in jaxonomy/framework/system_base.py

def get_input_port(self, name: str) -> tuple[InputPort, int]:
    """Retrieve a specific input port by name."""
    for i, port in enumerate(self.input_ports):
        if port.name == name:
            return port, i
    raise ValueError(
        f"System {self.name} has no input port named {name}. "
        f"Available ports: {list(map(lambda x: x.name,self.input_ports))}"
    )

`get_output_port(name)`

Retrieve a specific output port by name.

Source code in jaxonomy/framework/system_base.py

def get_output_port(self, name: str) -> OutputPort:
    """Retrieve a specific output port by name."""
    for port in self.output_ports:
        if port.name == name:
            return port
    raise ValueError(f"System {self.name} has no output port named {name}")

`get_parameter(name)`

Get a parameter value by name.

Checks dynamic parameters first, then static parameters. Values are returned as concrete array-like / Python scalars via :meth:Parameter.unwrap.

Parameters:

Name	Type	Description	Default
`name`	`str`	Parameter name on this system (not a dotted path).	required

Raises:

Type	Description
`KeyError`	If `name` is not found; the message lists available names.

Source code in jaxonomy/framework/system_base.py

def get_parameter(self, name: str):
    """Get a parameter value by name.

    Checks dynamic parameters first, then static parameters. Values are
    returned as concrete array-like / Python scalars via
    :meth:`Parameter.unwrap`.

    Args:
        name: Parameter name on this system (not a dotted path).

    Raises:
        KeyError: If ``name`` is not found; the message lists available names.
    """
    if name in self._dynamic_parameters:
        return Parameter.unwrap(self._dynamic_parameters[name])
    if name in self._static_parameters:
        return Parameter.unwrap(self._static_parameters[name])
    available = sorted(
        {*self._static_parameters.keys(), *self._dynamic_parameters.keys()}
    )
    raise KeyError(
        f"Parameter {name!r} not found on {self.name!r}. Available: {available}"
    )

`handle_discrete_update(events, context, *, topological_order=False)`

Compute and apply active discrete updates.

Given the root context, evaluate the discrete updates, which must have the same PyTree structure as the discrete states of this system. This should be a pure function, so that it does not modify any aspect of the context in-place (even though it is difficult to strictly prevent this in Python).

This will evaluate the set of events that result from declaring state or output update events on systems using LeafSystem.declare_periodic_update and LeafSystem.declare_output_port with an associated periodic update rate.

This is intended for internal use by the simulator and should not normally need to be invoked directly by users. Events are evaluated conditionally on being marked "active", so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns.

For a discrete system updating at a particular rate, the update rule for a particular block is:

x[n+1] = f(t[n], x[n], u[n])
y[n]   = g(t[n], x[n], u[n])

Additionally, the value y[n] is held constant until the next update from the point of view of other continuous-time or asynchronous discrete-time blocks.

Because each output y[n] may in general depend on the input u[n] evaluated at the same time, the composite discrete update function represents a system of equations. However, since algebraic loops are prohibited, the events can be ordered and executed sequentially to ensure that the updates are applied in the correct order. This is implemented in SystemBase.sorted_callbacks.

Multirate systems work in the same way, except that the events are evaluated conditionally on whether the current time corresponds to an update time for each event.

Parameters:

Name	Type	Description	Default
`events`	`EventCollection`	collection of discrete update events	required
`context`	`ContextBase`	root context for this system	required

Returns:

Name	Type	Description
`ContextBase`	`ContextBase`	updated context with all active updates applied to the discrete state

Source code in jaxonomy/framework/system_base.py

def handle_discrete_update(
    self,
    events: EventCollection,
    context: ContextBase,
    *,
    topological_order: bool = False,
) -> ContextBase:
    """Compute and apply active discrete updates.

    Given the _root_ context, evaluate the discrete updates, which must have the
    same PyTree structure as the discrete states of this system. This should be
    a pure function, so that it does not modify any aspect of the context in-place
    (even though it is difficult to strictly prevent this in Python).

    This will evaluate the set of events that result from declaring state or output
    update events on systems using `LeafSystem.declare_periodic_update` and
    `LeafSystem.declare_output_port` with an associated periodic update rate.

    This is intended for internal use by the simulator and should not normally need
    to be invoked directly by users. Events are evaluated conditionally on being
    marked "active", so the entire event collection can be passed without filtering
    to active events. This is necessary to make the function calls work with JAX
    tracing, which do not allow for variable-sized arguments or returns.

    For a discrete system updating at a particular rate, the update rule for a
    particular block is:

    ```
    x[n+1] = f(t[n], x[n], u[n])
    y[n]   = g(t[n], x[n], u[n])
    ```

    Additionally, the value y[n] is held constant until the next update from the
    point of view of other continuous-time or asynchronous discrete-time blocks.

    Because each output `y[n]` may in general depend on the input `u[n]` evaluated
    _at the same time_, the composite discrete update function represents a
    system of equations.  However, since algebraic loops are prohibited, the events
    can be ordered and executed sequentially to ensure that the updates are applied
    in the correct order.  This is implemented in
    `SystemBase.sorted_callbacks`.

    Multirate systems work in the same way, except that the events are evaluated
    conditionally on whether the current time corresponds to an update time for each
    event.

    Args:
        events (EventCollection): collection of discrete update events
        context (ContextBase): root context for this system

    Returns:
        ContextBase:
            updated context with all active updates applied to the discrete state
    """
    logger.debug(
        f"Handling {events.num_events} discrete update events at t={context.time}"
    )
    if events.has_events:
        # Two-phase approach for correct x⁻ / x⁺ semantics:
        #
        # Phase 1 — OUTPUT CACHE UPDATES (is_state_update=False):
        #   Process in sorted dependency order without a snapshot.  Each cache
        #   update computes y[n] = g(x[n]) and stores it in state.cache.  These
        #   are safe to sequence without a snapshot because they only READ discrete
        #   state (never write it), so no block can see another block's x⁺ here.
        #
        # Phase 2 — DISCRETE STATE UPDATES (is_state_update=True):
        #   After Phase 1, ALL blocks have correct y[n] in their state.cache.
        #   Now take a snapshot.  Each state update is evaluated against a "blended"
        #   context where:
        #     • the current block uses its accumulated context (so it sees any
        #       intra-block cache updates that already fired in Phase 1)
        #     • every OTHER block uses the snapshot (x⁻ discrete states + y[n]
        #       caches) — preventing block B from reading block A's x⁺ state.
        #
        # This correctly handles both:
        #   (a) DiscreteClock → UnitDelay: UnitDelay state update reads clock's
        #       y[n]=1.0 (set in Phase 1), not the stale y[n-1]=0.0 snapshot.
        #   (b) Cross-block swap (A reads B, B reads A): each sees x⁻ of the other
        #       via the Phase 2 snapshot, not x⁺.

        # T-105-followup-priority-scheduler-hook — compute the
        # per-system priority/rate tiebreak key once.  ``None`` means
        # no leaf has an explicit ``priority`` attribute set, so the
        # downstream scheduler must use its legacy byte-equivalent
        # order (declaration order here, ``str(system_id)`` for
        # Kahn's in the topological branch below).  ``_phase1`` is
        # computed up-front so Phase 1 cache updates honour the
        # same per-block ordering as Phase 2.
        _phase1_priority_tiebreak = _build_priority_tiebreak(self)

        # Phase 1: cache/output update events. These arrive already in
        # execution (topological) order — ``cache_update_events`` sorts them
        # via ``sorted_callbacks`` → ``sort_trackers``, which respects
        # feedthrough dependencies, INCLUDING transitively through pure-
        # feedthrough intermediates (e.g. a Mux between two sample-and-holds).
        # The one missing piece was refreshing the port cache BETWEEN them:
        # a downstream feedthrough sample-and-hold (e.g. ZeroOrderHold) must
        # read its upstream's fresh current-tick output, not the stale pre-
        # tick value. Without the refresh a same-rate ZOH lags its source by
        # a step, violating Simulink sorted-execution semantics (a same-rate
        # ZOH is the identity; two in series are a no-op).
        phase1_events = [e for e in events if not e.is_state_update]
        _multi_p1 = len(phase1_events) > 1
        if _phase1_priority_tiebreak is not None:
            # Stable sort: events for the same system_id keep their
            # declared order.
            phase1_events = sorted(
                phase1_events,
                key=lambda e: _phase1_priority_tiebreak(e.system_id),
            )
        for event in phase1_events:
            system_id = event.system_id
            state = event.handle(context)
            local_context = context[system_id].with_state(state)
            context = context.with_subcontext(system_id, local_context)
            if _multi_p1:
                # Refresh so the next (downstream) feedthrough cache update
                # reads this block's fresh current-tick output.
                context = context.refresh_port_cache()

        # Between phases: refresh the root port_cache so that Phase 2 state
        # update callbacks (which use OutputPort.eval → port_cache lookup) see
        # the updated y[n] values written in Phase 1, not stale values from
        # before Phase 1.
        snapshot = context.refresh_port_cache()

        # Phase 2: state update events.  Two ordering modes:
        #
        #   (a) Diagonal (default, `topological_order=False`): evaluate
        #       events in their declared order against the snapshot.
        #       Each block sees x⁻ for every other block — preserves the
        #       cross-block-swap atomicity documented above.
        #
        #   (b) Lower-triangular (T-022a, opt-in via
        #       `topological_order=True`): evaluate events in the
        #       topological order of the discrete dependency graph.
        #       Block B's update sees the post-update x⁺ of any block A
        #       upstream of B (A→B), AND the post-update y⁺[n+1] of
        #       any sample-and-hold output of A — both via the
        #       accumulated context.  Cycles raise
        #       ``DependencyCycleError``.
        state_events = [e for e in events if e.is_state_update]

        # Reuse the Phase-1 priority tiebreak (already accounts for
        # an "any leaf has explicit priority" check).  ``None`` →
        # legacy byte-equivalent order downstream.
        priority_tiebreak = _phase1_priority_tiebreak

        if topological_order:
            from .discrete_dependencies import (
                DependencyCycleError,
                build_discrete_dependency_graph,
                topological_sort,
            )
            graph = build_discrete_dependency_graph(self)
            event_ids = {e.system_id for e in events}  # all events
            trimmed = {
                n: deps & event_ids
                for n, deps in graph.items() if n in event_ids
            }
            topological_sort(
                trimmed, tiebreak_key=priority_tiebreak,
            )  # validate (raises on cycle)
            order = topological_sort(
                trimmed, tiebreak_key=priority_tiebreak,
            )
            # Group all events by system; topological branch
            # interleaves state + cache updates per block so
            # downstream blocks see y[n+1] (cache reflects post-
            # update state).
            all_events_by_id: dict = {}
            for e in events:
                all_events_by_id.setdefault(e.system_id, []).append(e)
            ancestors: dict = {n: set() for n in order}
            for n in order:
                for d in trimmed.get(n, ()):
                    ancestors[n] |= ancestors.get(d, set()) | {d}
            for sid in order:
                block_events = all_events_by_id.get(sid, [])
                blended = snapshot
                for use_acc in {sid, *ancestors[sid]}:
                    blended = blended.with_subcontext(
                        use_acc, context[use_acc],
                    )
                # Empty port cache so downstream OutputPort.eval
                # recomputes against the accumulated state.
                blended = blended.with_port_cache({})
                # State updates first (so cache updates can read x⁺).
                for event in block_events:
                    if not event.is_state_update:
                        continue
                    state = event.handle(blended)
                    local_context = context[sid].with_state(state)
                    context = context.with_subcontext(sid, local_context)
                    blended = blended.with_subcontext(sid, local_context)
                # Cache updates re-run against post-state context
                # so that y[n+1] reflects x[n+1] for sample-and-hold
                # output ports.
                for event in block_events:
                    if event.is_state_update:
                        continue
                    state = event.handle(blended)
                    local_context = context[sid].with_state(state)
                    context = context.with_subcontext(sid, local_context)
                    blended = blended.with_subcontext(sid, local_context)
        else:
            # Diagonal (legacy) order.  When ANY leaf declared an
            # explicit ``priority`` the events are stable-sorted by
            # the (rate, priority, name) tiebreak key — block
            # authors who opt into priorities see deterministic
            # within-tick ordering even on the legacy snapshot path.
            # When no priorities are declared the list is returned
            # unchanged so byte-equivalence with the pre-followup
            # diagonal scheduler is preserved.
            ordered_state_events = state_events
            if priority_tiebreak is not None:
                ordered_state_events = sorted(
                    state_events,
                    key=lambda e: priority_tiebreak(e.system_id),
                )
            for event in ordered_state_events:
                system_id = event.system_id
                # Own block: use accumulated context (may have updated cache from Phase 1)
                # Other blocks: use snapshot (x⁻ discrete states + y[n] caches)
                blended = snapshot.with_subcontext(system_id, context[system_id])
                state = event.handle(blended)
                local_context = context[system_id].with_state(state)
                context = context.with_subcontext(system_id, local_context)

    return context

`handle_zero_crossings(events, context)`

Compute and apply active zero-crossing events.

This is intended for internal use by the simulator and should not normally need to be invoked directly by users. Events are evaluated conditionally on being marked "active", so the entire event collection can be passed without filtering to active events. This is necessary to make the function calls work with JAX tracing, which do not allow for variable-sized arguments or returns.

Parameters:

Name	Type	Description	Default
`events`	`EventCollection`	collection of zero-crossing events	required
`context`	`ContextBase`	root context for this system	required

Returns:

Name	Type	Description
`ContextBase`	`ContextBase`	updated context with all active zero-crossing events applied

Source code in jaxonomy/framework/system_base.py

def handle_zero_crossings(
    self, events: EventCollection, context: ContextBase
) -> ContextBase:
    """Compute and apply active zero-crossing events.

    This is intended for internal use by the simulator and should not normally need
    to be invoked directly by users. Events are evaluated conditionally on being
    marked "active", so the entire event collection can be passed without filtering
    to active events. This is necessary to make the function calls work with JAX
    tracing, which do not allow for variable-sized arguments or returns.

    Args:
        events (EventCollection): collection of zero-crossing events
        context (ContextBase): root context for this system

    Returns:
        ContextBase: updated context with all active zero-crossing events applied
    """
    logger.debug(
        "Handling %d state transition events at t=%s",
        events.num_events,
        context.time,
    )
    if events.num_events > 0:
        # Compute the updated state for all active events
        context = context.refresh_port_cache()
        state = self.eval_zero_crossing_updates(context, events)

        # Apply the updates
        context = context.with_state(state)
    return context

`initialize_static_data(context)`

Initialize any context data that has to be done after context creation.

Use this to define custom auxiliary data or type inference that doesn't get traced by JAX. See the ZeroOrderHold implementation for an example. Since this is only applied during context initialization, it is allowed to modify the context directly (or the system itself).

Typically this should not be called outside of the ContextFactory.

Parameters:

Name	Type	Description	Default
`context`	`ContextBase`	partially initialized context for this system.	required

Source code in jaxonomy/framework/system_base.py

def initialize_static_data(self, context: ContextBase) -> ContextBase:
    """Initialize any context data that has to be done after context creation.

    Use this to define custom auxiliary data or type inference that doesn't
    get traced by JAX. See the `ZeroOrderHold` implementation for an example.
    Since this is only applied during context initialization, it is allowed to
    modify the context directly (or the system itself).

    Typically this should not be called outside of the ContextFactory.

    Args:
        context (ContextBase): partially initialized context for this system.
    """
    return context

`list_parameters()`

Return all parameters as a flat {name: value} mapping.

Dynamic parameters override static ones when names collide. Values are unwrapped the same way as :meth:get_parameter.

Source code in jaxonomy/framework/system_base.py

def list_parameters(self) -> dict:
    """Return all parameters as a flat ``{name: value}`` mapping.

    Dynamic parameters override static ones when names collide. Values are
    unwrapped the same way as :meth:`get_parameter`.
    """
    merged = {**self._static_parameters, **self._dynamic_parameters}
    return {k: Parameter.unwrap(v) for k, v in merged.items()}

`post_simulation_finalize()`

Finalize the system after simulation has completed.

This is only intended for special blocks that need to clean up resources and close files.

Source code in jaxonomy/framework/system_base.py

def post_simulation_finalize(self) -> None:
    """Finalize the system after simulation has completed.

    This is only intended for special blocks that need to clean up
    resources and close files."""

`pprint(output=print, fancy=True)`

Pretty-print the system and its hierarchy.

Source code in jaxonomy/framework/system_base.py

def pprint(self, output=print, fancy=True) -> str:
    """Pretty-print the system and its hierarchy."""
    output(self._pprint_helper(fancy=fancy).strip())

`recompute_port_cache(context)`

Recompute all output ports and return them in a dictionary.

Source code in jaxonomy/framework/system_base.py

def recompute_port_cache(self, context: ContextBase) -> dict[Hashable, Array]:
    """Recompute all output ports and return them in a dictionary."""
    # The callbacks are already sorted in the correct order, so we can just
    # evaluate them in sequence, assuming all upstream values are correct.
    for cb in self.sorted_callbacks:
        if not isinstance(cb, OutputPort):
            continue
        # Evaluate the callback and store the result in the cache
        val = cb.calc(context)
        context = context.with_port_cache_entry(hash(cb), val)

    return context

`SystemCallback` `dataclass`

A function associated with a system that has has specified dependencies.

This can include port update rules, discrete update functions, the right-hand-side of an ODE, etc. Storing these functions as SystemCallbacks allows the system, or a Diagram containing the system, to track dependencies across the system or diagram.

Attributes:

Name	Type	Description
`system`	`SystemBase`	The system that owns this callback.
`ticket`	`int`	The dependency ticket associated with this callback. See DependencyTicket for built-in tickets. If None, a new ticket will be generated.
`name`	`str`	A short description of this callback function.
`prerequisites_of_calc`	`List[DependencyTicket]`	Direct prerequisites of the computation, used for dependency tracking. These might be built-in tickets or tickets associated with other SystemCallbacks.
`default_value`	`Array`	A dummy value of the same shape/dtype as the result, if known. If None, any type checking will rely on propagating upstream information via the callback.
`callback_index`	`int`	The index of this function in the system's list of associated callbacks.
`event`	`Event`	Optionally, the callback function may be associated with an event. If so, the associated trackers can be used to sort event execution order in addition to the regular callback execution order. For example, if an OutputPort is of sample-and-hold type, then this will be the event that periodically updates the output value. Default is None.

Source code in jaxonomy/framework/cache.py

@dataclasses.dataclass
class SystemCallback:
    """A function associated with a system that has has specified dependencies.

    This can include port update rules, discrete update functions, the right-hand-side
    of an ODE, etc. Storing these functions as SystemCallbacks allows the system, or a
    Diagram containing the system, to track dependencies across the system or diagram.

    Attributes:
        system (SystemBase):
            The system that owns this callback.
        ticket (int):
            The dependency ticket associated with this callback.  See DependencyTicket
            for built-in tickets. If None, a new ticket will be generated.
        name (str):
            A short description of this callback function.
        prerequisites_of_calc (List[DependencyTicket]):
            Direct prerequisites of the computation, used for dependency tracking.
            These might be built-in tickets or tickets associated with other
            SystemCallbacks.
        default_value (Array):
            A dummy value of the same shape/dtype as the result, if known.  If None,
            any type checking will rely on propagating upstream information via the
            callback.
        callback_index (int):
            The index of this function in the system's list of associated callbacks.
        event (Event):
            Optionally, the callback function may be associated with an event.  If so,
            the associated trackers can be used to sort event execution order in addition
            to the regular callback execution order. For example, if an OutputPort is of
            sample-and-hold type, then this will be the event that periodically updates
            the output value. Default is None.
    """

    callback: dataclasses.InitVar[Callable[[ContextBase], Array]]
    system: SystemBase
    callback_index: int
    ticket: DependencyTicket = None
    name: str = None
    prerequisites_of_calc: List[DependencyTicket] = None
    default_value: Array = None
    event: Event = None

    # If the result is cached (e.g. an output port of "sample-and-hold" type),
    # this will be the index of the cache in the system's cache list.
    cache_index: int = None

    def __post_init__(self, callback):
        self._callback = callback  # Given root context, return calculated value

        if self.ticket is None:
            self.ticket = next_dependency_ticket()
        assert isinstance(self.ticket, int)

        if self.prerequisites_of_calc is None:
            self.prerequisites_of_calc = []

        logger.debug(
            "Initialized callback %s:%s with prereqs %s",
            self.system.name_path_str,
            self.name,
            self.prerequisites_of_calc,
        )


    def __hash__(self) -> int:
        locator = (self.system, self.callback_index)
        return hash(locator)

    def __repr__(self) -> str:
        return f"{self.name}(ticket = {self.ticket})"

    def calc(self, root_context: ContextBase) -> Array:
        """Unconditionally evaluate the callback function.

        This does not check the cache status, but will always recompute the value.
        Typically `eval` should be preferred to `calc` to take advantage of caching
        where possible.

        Args:
            root_context: The root context used for the evaluation.

        Returns:
            The calculated value from the callback, expected to be a Array.
        """
        return self._callback(root_context)

    def eval(self, root_context: ContextBase) -> Array:
        """Evaluate the callback function and return the calculated value.

        Within a single top-level call, repeated evaluations of the same
        callback against the same context are memoized (see ``_eval_memo``
        above) — this keeps eager evaluation of diagrams with fan-out /
        reconvergence linear in graph size instead of exponential in
        composition depth.  Nothing is cached across top-level calls.

        Args:
            root_context: The root context used for the evaluation.

        Returns:
            The calculated value from the callback, expected to be a Array.
        """
        if not root_context.is_initialized:
            if self.default_value is None:
                self.default_value = self.calc(root_context)
            return self.default_value

        memo = _eval_memo.get()
        if memo is None:
            # Outermost eval of this tree: install a fresh memo, always
            # clear it on exit so nothing leaks across top-level calls
            # (or across JAX traces).
            token = _eval_memo.set({})
            try:
                return self._eval_initialized(root_context)
            finally:
                _eval_memo.reset(token)

        key = (id(root_context), self)
        if key in memo:
            return memo[key]
        result = self._eval_initialized(root_context)
        memo[key] = result
        return result

    def _eval_initialized(self, root_context: ContextBase) -> Array:
        try:
            result = self.calc(root_context)
        except ValueError as e:
            # this error is raised if the callback is not differentiable
            if "do not support JVP." in str(e):
                raise CallbackIsNotDifferentiableError(
                    system=self.system,
                    port_name=self.name,
                )
            raise
        return result

    @property
    def tracker(self) -> DependencyTracker:
        return self.system.dependency_graph[self.ticket]

`calc(root_context)`

Unconditionally evaluate the callback function.

This does not check the cache status, but will always recompute the value. Typically eval should be preferred to calc to take advantage of caching where possible.

Parameters:

Name	Type	Description	Default
`root_context`	`ContextBase`	The root context used for the evaluation.	required

Returns:

Type	Description
`Array`	The calculated value from the callback, expected to be a Array.

Source code in jaxonomy/framework/cache.py

def calc(self, root_context: ContextBase) -> Array:
    """Unconditionally evaluate the callback function.

    This does not check the cache status, but will always recompute the value.
    Typically `eval` should be preferred to `calc` to take advantage of caching
    where possible.

    Args:
        root_context: The root context used for the evaluation.

    Returns:
        The calculated value from the callback, expected to be a Array.
    """
    return self._callback(root_context)

`eval(root_context)`

Evaluate the callback function and return the calculated value.

Within a single top-level call, repeated evaluations of the same callback against the same context are memoized (see _eval_memo above) — this keeps eager evaluation of diagrams with fan-out / reconvergence linear in graph size instead of exponential in composition depth. Nothing is cached across top-level calls.

Parameters:

Name	Type	Description	Default
`root_context`	`ContextBase`	The root context used for the evaluation.	required

Returns:

Type	Description
`Array`	The calculated value from the callback, expected to be a Array.

Source code in jaxonomy/framework/cache.py

def eval(self, root_context: ContextBase) -> Array:
    """Evaluate the callback function and return the calculated value.

    Within a single top-level call, repeated evaluations of the same
    callback against the same context are memoized (see ``_eval_memo``
    above) — this keeps eager evaluation of diagrams with fan-out /
    reconvergence linear in graph size instead of exponential in
    composition depth.  Nothing is cached across top-level calls.

    Args:
        root_context: The root context used for the evaluation.

    Returns:
        The calculated value from the callback, expected to be a Array.
    """
    if not root_context.is_initialized:
        if self.default_value is None:
            self.default_value = self.calc(root_context)
        return self.default_value

    memo = _eval_memo.get()
    if memo is None:
        # Outermost eval of this tree: install a fresh memo, always
        # clear it on exit so nothing leaks across top-level calls
        # (or across JAX traces).
        token = _eval_memo.set({})
        try:
            return self._eval_initialized(root_context)
        finally:
            _eval_memo.reset(token)

    key = (id(root_context), self)
    if key in memo:
        return memo[key]
    result = self._eval_initialized(root_context)
    memo[key] = result
    return result

`TriggerEdge`

Allowed string values for TriggeredSubsystem.edge.

Source code in jaxonomy/framework/containers.py

class TriggerEdge:
    """Allowed string values for ``TriggeredSubsystem.edge``."""

    RISING = "rising"
    FALLING = "falling"
    EITHER = "either"

    @classmethod
    def valid(cls) -> tuple[str, ...]:
        return (cls.RISING, cls.FALLING, cls.EITHER)

`TriggeredSubsystem`

Bases: LeafSystem

Container block: latch the submodel output on edge transitions (the child still RUNS every step — only the output is gated).

Important: this does not skip execution of the submodel on non-triggered steps. The submodel is evaluated on every step so its inputs participate in the JAX trace; the trigger only controls whether a fresh result is latched into the held output. If you need to actually skip computation between triggers, gate it yourself with jax.lax.cond at the application level.

Phase-1 implementation runs the submodel on every step (so the inputs participate in the trace) but only latches a new output on an edge transition of the trigger signal. Between transitions the output holds the most recently latched value.

The trigger signal is sampled at sample_period. Edges are detected by comparing the current trigger sample against the previously-stored sample held in discrete state.

This is not the eventual zero-crossing-driven TriggeredSubsystem described in the T-120 architecture notes (that requires hooking into the continuous-time event detector); but it is functionally correct for any sample-rate use case and matches the behaviour documented in the test fixtures.

Parameters:

Name	Type	Description	Default
`submodel`	`Callable`	Callable `f(*inputs) -> output` taking the non-trigger user inputs. Must be JAX-traceable.	required
`n_inputs`	`int`	Number of user inputs (NOT counting the trigger).	`1`
`edge`	`Literal['rising', 'falling', 'either']`	`"rising"` (low→high), `"falling"` (high→low) or `"either"`.	`RISING`
`sample_period`	`float`	Period (seconds) at which the trigger signal is sampled and the latch is updated. Must be positive.	`0.0`
`initial_value`		Latched output value before any edge has been detected. Defines output shape/dtype.	`0.0`
`name`		Optional block name.	required

Limitations (phase 1): - Trigger detection runs on the periodic sample grid, not on continuous-time zero crossings. Trigger pulses shorter than sample_period may be missed. - The latch is a single discrete state; the submodel must produce a single output array. - The submodel runs on every output evaluation; only the output is gated. Users who need to skip computation on non-triggered steps should use jax.lax.cond at the application level.

Source code in jaxonomy/framework/containers.py

class TriggeredSubsystem(LeafSystem):
    """Container block: latch the submodel output on edge transitions
    (the child still RUNS every step — only the *output* is gated).

    Important: this does **not** skip execution of the submodel on
    non-triggered steps. The submodel is evaluated on every step so its
    inputs participate in the JAX trace; the trigger only controls
    whether a fresh result is *latched* into the held output. If you
    need to actually skip computation between triggers, gate it yourself
    with ``jax.lax.cond`` at the application level.

    Phase-1 implementation runs the submodel on every step (so the
    inputs participate in the trace) but only *latches* a new output
    on an edge transition of the trigger signal. Between transitions
    the output holds the most recently latched value.

    The trigger signal is sampled at ``sample_period``. Edges are
    detected by comparing the current trigger sample against the
    previously-stored sample held in discrete state.

    This is *not* the eventual zero-crossing-driven ``TriggeredSubsystem``
    described in the T-120 architecture notes (that requires hooking
    into the continuous-time event detector); but it is functionally
    correct for any sample-rate use case and matches the behaviour
    documented in the test fixtures.

    Args:
        submodel: Callable ``f(*inputs) -> output`` taking the
            non-trigger user inputs. Must be JAX-traceable.
        n_inputs: Number of user inputs (NOT counting the trigger).
        edge: ``"rising"`` (low→high), ``"falling"`` (high→low) or
            ``"either"``.
        sample_period: Period (seconds) at which the trigger signal is
            sampled and the latch is updated. Must be positive.
        initial_value: Latched output value before any edge has been
            detected. Defines output shape/dtype.
        name: Optional block name.

    Limitations (phase 1):
        - Trigger detection runs on the periodic sample grid, not on
          continuous-time zero crossings. Trigger pulses shorter than
          ``sample_period`` may be missed.
        - The latch is a single discrete state; the submodel must
          produce a single output array.
        - The submodel runs on every output evaluation; only the
          *output* is gated. Users who need to skip computation on
          non-triggered steps should use ``jax.lax.cond`` at the
          application level.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        edge: Literal["rising", "falling", "either"] = TriggerEdge.RISING,
        sample_period: float = 0.0,
        initial_value=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if edge not in TriggerEdge.valid():
            raise ValueError(
                f"TriggeredSubsystem: edge must be one of "
                f"{TriggerEdge.valid()!r}, got {edge!r}"
            )
        if sample_period is None or float(sample_period) <= 0.0:
            raise ValueError(
                "TriggeredSubsystem requires a positive sample_period "
                "(seconds) for trigger sampling."
            )
        if n_inputs < 0:
            raise ValueError(
                f"TriggeredSubsystem: n_inputs must be >= 0, got {n_inputs}"
            )

        self._submodel = submodel
        self._edge = edge
        self._sample_period = float(sample_period)
        self._initial = jnp.asarray(initial_value)

        # Port 0 is the trigger signal; ports 1..n_inputs are the user
        # inputs forwarded to the submodel.
        self.declare_input_port(name="trigger")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        # Discrete state pair: (latched_output, previous_trigger).
        # Pack as a flat 1-D array so the existing scalar-friendly
        # discrete-state machinery handles them uniformly. The two
        # pieces have different shapes in general, so use a tuple-like
        # tree via two separate periodic updates? Simpler: pack as a
        # dict. ``LeafSystem.declare_discrete_state`` only accepts a
        # single default_value, so we encode (latch, prev_trigger) as
        # a flat concatenation when both are scalar. For phase 1 we
        # require a scalar trigger so this packing is safe.
        #
        # Layout: discrete_state[..., 0] holds the previous trigger
        # sample; discrete_state[..., 1:] holds the latched output
        # (flattened). For a scalar latch this collapses to length 2.
        flat_init = jnp.concatenate(
            [
                jnp.zeros((1,), dtype=self._initial.dtype),
                jnp.atleast_1d(self._initial).reshape(-1),
            ]
        )
        self._latch_shape = self._initial.shape
        self._latch_size = int(jnp.atleast_1d(self._initial).reshape(-1).size)
        self.declare_discrete_state(default_value=flat_init)
        self.declare_periodic_update(
            self._latch_update,
            period=self._sample_period,
            offset=0.0,
        )
        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── helpers ───────────────────────────────────────────────────────────

    def _unpack(self, ds):
        prev_trig = ds[0]
        latch_flat = ds[1 : 1 + self._latch_size]
        latch = latch_flat.reshape(self._latch_shape)
        return prev_trig, latch

    def _pack(self, prev_trig, latch):
        return jnp.concatenate(
            [
                jnp.atleast_1d(prev_trig).reshape(-1)[:1],
                jnp.atleast_1d(latch).reshape(-1),
            ]
        )

    def _edge_detected(self, prev_trig, cur_trig):
        prev_b = jnp.asarray(prev_trig).astype(bool)
        cur_b = jnp.asarray(cur_trig).astype(bool)
        if self._edge == TriggerEdge.RISING:
            return jnp.logical_and(jnp.logical_not(prev_b), cur_b)
        if self._edge == TriggerEdge.FALLING:
            return jnp.logical_and(prev_b, jnp.logical_not(cur_b))
        # either
        return jnp.not_equal(prev_b, cur_b)

    # ── callbacks ─────────────────────────────────────────────────────────

    def _submodel_output(self, inputs):
        user_inputs = inputs[1:]  # skip trigger
        return jnp.asarray(self._submodel(*user_inputs))

    def _latch_update(self, time, state, *inputs, **params):
        ds = state.discrete_state
        prev_trig, latch = self._unpack(ds)
        cur_trig = jnp.asarray(inputs[0])
        edge = self._edge_detected(prev_trig, cur_trig)
        y_sub = self._submodel_output(inputs)
        new_latch = jnp.where(edge, y_sub, latch)
        # Store the current trigger sample (cast to the same dtype as
        # the rest of the discrete-state vector).
        new_prev = jnp.asarray(cur_trig).astype(ds.dtype).reshape(())
        return self._pack(new_prev, new_latch)

    def _compute_output(self, time, state, *inputs, **params):
        # Output reads the latched value. The submodel is *also* invoked
        # via the latch_update path on the periodic sample grid; here we
        # additionally consult the current trigger so a same-step rising
        # edge surfaces immediately rather than one sample later.
        ds = state.discrete_state
        prev_trig, latch = self._unpack(ds)
        cur_trig = jnp.asarray(inputs[0])
        edge = self._edge_detected(prev_trig, cur_trig)
        y_sub = self._submodel_output(inputs)
        return jnp.where(edge, y_sub, latch)

`Unit` `dataclass`

Immutable SI dimensional value.

Units are compared by their dimension exponents and scale factor. The optional name is informational (used in error messages) and is not part of equality.

Source code in jaxonomy/framework/units.py

@dataclass(frozen=True, eq=False)
class Unit:
    """Immutable SI dimensional value.

    Units are compared by their dimension exponents and scale factor.
    The optional ``name`` is informational (used in error messages) and is
    not part of equality.
    """

    # Tuple of seven integers: exponents on (kg, m, s, A, K, mol, cd).
    dims: Tuple[int, int, int, int, int, int, int] = field(
        default=(0, 0, 0, 0, 0, 0, 0)
    )
    # Multiplicative scale on the base SI unit (1.0 == base).  Phase 1 only
    # supports unit / unit comparisons that are dimension-equal AND
    # scale-equal; conversions are deferred to Phase 2.
    scale: float = 1.0
    # Additive offset applied after the multiplicative ``scale`` when
    # mapping a raw value into base SI:
    #   physical_value_in_base_SI = scale * raw_value + offset
    # The vast majority of physical units (length, mass, time, etc.) have
    # offset == 0; the canonical exceptions are temperature scales
    # (Celsius, Fahrenheit). Unit multiplication / division for units
    # with a non-zero ``offset`` is not well-defined and raises.
    # (T-104 followup temperature conversion.)
    offset: float = 0.0
    # T-104-followup-currency-units: per-currency exponents.
    # Tuple of five integers (USD, EUR, GBP, JPY, CAD). Default all-zero,
    # which keeps every pre-existing :class:`Unit` instance byte-equal
    # to its previous form under the dataclass equality below. See the
    # module-level :data:`CURRENCY_CODES` for the canonical axis order.
    currency: Tuple[int, int, int, int, int] = field(
        default=_ZERO_CURRENCY
    )
    # Optional human-readable label, e.g. "m/s".  Not part of equality.
    name: str | None = None
    # T-104 phase 3: physical-quantity disambiguation.
    #
    # Two units with identical SI dimensions can describe different
    # physical quantities (the textbook example: ``N·m`` is both torque
    # AND energy; ``Pa·m^3`` is both work and pressure*volume). When set,
    # ``physical_quantity`` tags the *interpretation* so the
    # connect-time check refuses a torque-out → energy-in connection
    # even though their dims agree.
    #
    # Compatibility rule (in :func:`are_units_compatible`):
    #   * either side ``None`` → match (default-off byte-equivalence)
    #   * both set → must be string-equal
    #
    # Multiplicative algebra drops the tag (``torque * angular_velocity``
    # has no canonical interpretation; user must re-tag the result).
    # Default ``None`` keeps every pre-existing :class:`Unit` instance
    # byte-equal to its previous form.
    physical_quantity: str | None = None

    # ---- algebra -----------------------------------------------------

    def __post_init__(self):
        # Coerce dims into a 7-tuple of ints to make the dataclass robust
        # against being constructed from lists / numpy ints.
        if not isinstance(self.dims, tuple) or len(self.dims) != 7:
            object.__setattr__(self, "dims", tuple(int(d) for d in self.dims))
        # Same defensive coercion for the currency 5-tuple. We accept
        # any iterable of length 5 (lists, generators) so that callers
        # can spell ``currency=[1, 0, 0, 0, 0]`` without surprises.
        if (
            not isinstance(self.currency, tuple)
            or len(self.currency) != len(CURRENCY_CODES)
        ):
            coerced = tuple(int(c) for c in self.currency)
            if len(coerced) != len(CURRENCY_CODES):
                raise ValueError(
                    f"Unit.currency must have {len(CURRENCY_CODES)} entries "
                    f"(one per code in CURRENCY_CODES = {CURRENCY_CODES}); "
                    f"got {coerced!r}."
                )
            object.__setattr__(self, "currency", coerced)

    def _require_zero_offset_for_algebra(self, other: "Unit", op: str) -> None:
        """Raise a clear error when multiplying / dividing offsetted units.

        Offsetted (affine) units like Celsius / Fahrenheit do not form
        a multiplicative algebra: ``celsius * celsius`` is undefined,
        and ``celsius / second`` cannot be expressed as a scalar
        ``Unit`` either. We refuse the operation rather than silently
        returning a numerically-bogus result.
        """
        if self.offset != 0.0 or other.offset != 0.0:
            raise UnitMismatchError(
                f"Cannot {op} offsetted units: {self!r} {op} {other!r}. "
                "Affine units (e.g. Celsius, Fahrenheit) do not form a "
                "multiplicative algebra. Convert to a base unit "
                "(e.g. kelvin) via convert_offset_aware() first."
            )

    def __mul__(self, other: "Unit") -> "Unit":
        if not isinstance(other, Unit):
            return NotImplemented
        self._require_zero_offset_for_algebra(other, "multiply")
        new_dims = tuple(a + b for a, b in zip(self.dims, other.dims))
        new_currency = tuple(
            a + b for a, b in zip(self.currency, other.currency)
        )
        return Unit(
            dims=new_dims,
            scale=self.scale * other.scale,
            currency=new_currency,
        )

    def __truediv__(self, other: "Unit") -> "Unit":
        if not isinstance(other, Unit):
            return NotImplemented
        self._require_zero_offset_for_algebra(other, "divide")
        new_dims = tuple(a - b for a, b in zip(self.dims, other.dims))
        new_currency = tuple(
            a - b for a, b in zip(self.currency, other.currency)
        )
        return Unit(
            dims=new_dims,
            scale=self.scale / other.scale,
            currency=new_currency,
        )

    def __pow__(self, exponent: int) -> "Unit":
        if not isinstance(exponent, int):
            raise TypeError(
                f"Unit exponent must be int, got {type(exponent).__name__}"
            )
        if self.offset != 0.0 and exponent != 1:
            raise UnitMismatchError(
                f"Cannot raise offsetted unit {self!r} to a non-unity "
                "power. Affine units (e.g. Celsius, Fahrenheit) do not "
                "form a multiplicative algebra."
            )
        new_dims = tuple(d * exponent for d in self.dims)
        new_currency = tuple(c * exponent for c in self.currency)
        return Unit(
            dims=new_dims,
            scale=self.scale ** exponent,
            currency=new_currency,
        )

    # ---- equality / hashing -----------------------------------------

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Unit):
            return NotImplemented
        # T-104 phase 3: physical_quantity participates in equality so two
        # Units with the same dims but different physical interpretations
        # (e.g. N·m as torque vs. energy) compare distinct. ``name`` stays
        # informational and out of the equality contract.
        return (
            self.dims == other.dims
            and self.scale == other.scale
            and self.offset == other.offset
            and self.currency == other.currency
            and self.physical_quantity == other.physical_quantity
        )

    def __hash__(self) -> int:
        return hash((
            self.dims, self.scale, self.offset, self.currency,
            self.physical_quantity,
        ))

    # ---- dimensional helpers ----------------------------------------

    @property
    def is_dimensionless(self) -> bool:
        return (
            all(d == 0 for d in self.dims)
            and all(c == 0 for c in self.currency)
            and self.scale == 1.0
            and self.offset == 0.0
        )

    def same_dimension_as(self, other: "Unit") -> bool:
        """True if exponents match (ignoring scale).  Phase 1 doesn't use
        this for the connect check (which is strict-equal), but it's part
        of the public surface so Phase 2 can layer scalar-conversion
        warnings on top."""
        return (
            isinstance(other, Unit)
            and self.dims == other.dims
            and self.currency == other.currency
        )

    # ---- repr -------------------------------------------------------

    def __repr__(self) -> str:
        if self.name:
            return f"Unit({self.name!r})"
        if self.is_dimensionless:
            return "Unit(dimensionless)"
        parts = []
        for exp, label in zip(self.dims, _DIM_NAMES):
            if exp == 0:
                continue
            parts.append(label if exp == 1 else f"{label}^{exp}")
        # T-104-followup-currency-units: render non-zero currency
        # exponents alongside the SI dimensions so a $/m composite
        # prints as ``Unit(USD*m^-1)`` rather than collapsing.
        for exp, code in zip(self.currency, CURRENCY_CODES):
            if exp == 0:
                continue
            parts.append(code if exp == 1 else f"{code}^{exp}")
        body = "*".join(parts) if parts else "1"
        if self.scale != 1.0:
            body = f"{self.scale}*{body}"
        if self.offset != 0.0:
            sign = "+" if self.offset >= 0 else "-"
            body = f"{body}{sign}{abs(self.offset)}"
        if self.physical_quantity is not None:
            body = f"{body}@{self.physical_quantity}"
        return f"Unit({body})"

    # ---- T-104 phase 3: serialization + human-readable summary ----

    def to_dict(self) -> dict:
        """Return a JSON-friendly dict representation of this Unit.

        Round-trips losslessly via :meth:`from_dict`. Keys are stable
        across versions; new optional fields are always added with
        defaults so older serialised forms continue to load.

        The default value for any field is omitted from the output for
        compactness — every legacy ``Unit()`` instance serialises to
        ``{}``.
        """
        out: dict = {}
        if self.dims != (0, 0, 0, 0, 0, 0, 0):
            out["dims"] = list(self.dims)
        if self.scale != 1.0:
            out["scale"] = self.scale
        if self.offset != 0.0:
            out["offset"] = self.offset
        if self.currency != _ZERO_CURRENCY:
            out["currency"] = list(self.currency)
        if self.name is not None:
            out["name"] = self.name
        if self.physical_quantity is not None:
            out["physical_quantity"] = self.physical_quantity
        return out

    @classmethod
    def from_dict(cls, data: dict) -> "Unit":
        """Construct a :class:`Unit` from a dict produced by
        :meth:`to_dict`. Missing keys take their dataclass defaults so
        the empty dict ``{}`` round-trips to ``Unit()``.
        """
        return cls(
            dims=tuple(data.get("dims", (0, 0, 0, 0, 0, 0, 0))),
            scale=float(data.get("scale", 1.0)),
            offset=float(data.get("offset", 0.0)),
            currency=tuple(data.get("currency", _ZERO_CURRENCY)),
            name=data.get("name"),
            physical_quantity=data.get("physical_quantity"),
        )

    def to_json(self, *, indent: int | None = None) -> str:
        """Serialise :meth:`to_dict` via :func:`json.dumps`.

        Args:
            indent: Optional JSON indent (default ``None`` for compact
                form; pass an int for pretty-printed output).
        """
        import json

        return json.dumps(self.to_dict(), indent=indent, sort_keys=True)

    @classmethod
    def from_json(cls, json_str: str) -> "Unit":
        """Inverse of :meth:`to_json`.

        Raises:
            ValueError: If ``json_str`` is not a JSON object.
        """
        import json

        data = json.loads(json_str)
        if not isinstance(data, dict):
            raise ValueError(
                f"Unit.from_json: expected a JSON object at the top "
                f"level; got {type(data).__name__}."
            )
        return cls.from_dict(data)

    def summary(self) -> str:
        """Return a human-readable one-line summary of this Unit.

        Designed for ``print()`` / display contexts where ``repr(unit)``
        is too terse. Includes the dimension exponents (with SI labels),
        scale, offset, currency exponents, and the
        ``physical_quantity`` tag when set.
        """
        if self.is_dimensionless:
            base = "dimensionless"
        else:
            parts = []
            for exp, label in zip(self.dims, _DIM_NAMES):
                if exp != 0:
                    parts.append(label if exp == 1 else f"{label}^{exp}")
            for exp, code in zip(self.currency, CURRENCY_CODES):
                if exp != 0:
                    parts.append(code if exp == 1 else f"{code}^{exp}")
            base = " · ".join(parts) if parts else "1"
        bits = [base]
        if self.scale != 1.0:
            bits.append(f"scale={self.scale}")
        if self.offset != 0.0:
            bits.append(f"offset={self.offset}")
        if self.name is not None:
            bits.append(f'name="{self.name}"')
        if self.physical_quantity is not None:
            bits.append(f'physical_quantity="{self.physical_quantity}"')
        return "Unit(" + ", ".join(bits) + ")"

`from_dict(data)` `classmethod`

Construct a :class:Unit from a dict produced by :meth:to_dict. Missing keys take their dataclass defaults so the empty dict {} round-trips to Unit().

Source code in jaxonomy/framework/units.py

@classmethod
def from_dict(cls, data: dict) -> "Unit":
    """Construct a :class:`Unit` from a dict produced by
    :meth:`to_dict`. Missing keys take their dataclass defaults so
    the empty dict ``{}`` round-trips to ``Unit()``.
    """
    return cls(
        dims=tuple(data.get("dims", (0, 0, 0, 0, 0, 0, 0))),
        scale=float(data.get("scale", 1.0)),
        offset=float(data.get("offset", 0.0)),
        currency=tuple(data.get("currency", _ZERO_CURRENCY)),
        name=data.get("name"),
        physical_quantity=data.get("physical_quantity"),
    )

`from_json(json_str)` `classmethod`

Inverse of :meth:to_json.

Raises:

Type	Description
`ValueError`	If `json_str` is not a JSON object.

Source code in jaxonomy/framework/units.py

@classmethod
def from_json(cls, json_str: str) -> "Unit":
    """Inverse of :meth:`to_json`.

    Raises:
        ValueError: If ``json_str`` is not a JSON object.
    """
    import json

    data = json.loads(json_str)
    if not isinstance(data, dict):
        raise ValueError(
            f"Unit.from_json: expected a JSON object at the top "
            f"level; got {type(data).__name__}."
        )
    return cls.from_dict(data)

`same_dimension_as(other)`

True if exponents match (ignoring scale). Phase 1 doesn't use this for the connect check (which is strict-equal), but it's part of the public surface so Phase 2 can layer scalar-conversion warnings on top.

Source code in jaxonomy/framework/units.py

def same_dimension_as(self, other: "Unit") -> bool:
    """True if exponents match (ignoring scale).  Phase 1 doesn't use
    this for the connect check (which is strict-equal), but it's part
    of the public surface so Phase 2 can layer scalar-conversion
    warnings on top."""
    return (
        isinstance(other, Unit)
        and self.dims == other.dims
        and self.currency == other.currency
    )

`summary()`

Return a human-readable one-line summary of this Unit.

Designed for print() / display contexts where repr(unit) is too terse. Includes the dimension exponents (with SI labels), scale, offset, currency exponents, and the physical_quantity tag when set.

Source code in jaxonomy/framework/units.py

def summary(self) -> str:
    """Return a human-readable one-line summary of this Unit.

    Designed for ``print()`` / display contexts where ``repr(unit)``
    is too terse. Includes the dimension exponents (with SI labels),
    scale, offset, currency exponents, and the
    ``physical_quantity`` tag when set.
    """
    if self.is_dimensionless:
        base = "dimensionless"
    else:
        parts = []
        for exp, label in zip(self.dims, _DIM_NAMES):
            if exp != 0:
                parts.append(label if exp == 1 else f"{label}^{exp}")
        for exp, code in zip(self.currency, CURRENCY_CODES):
            if exp != 0:
                parts.append(code if exp == 1 else f"{code}^{exp}")
        base = " · ".join(parts) if parts else "1"
    bits = [base]
    if self.scale != 1.0:
        bits.append(f"scale={self.scale}")
    if self.offset != 0.0:
        bits.append(f"offset={self.offset}")
    if self.name is not None:
        bits.append(f'name="{self.name}"')
    if self.physical_quantity is not None:
        bits.append(f'physical_quantity="{self.physical_quantity}"')
    return "Unit(" + ", ".join(bits) + ")"

`to_dict()`

Return a JSON-friendly dict representation of this Unit.

Round-trips losslessly via :meth:from_dict. Keys are stable across versions; new optional fields are always added with defaults so older serialised forms continue to load.

The default value for any field is omitted from the output for compactness — every legacy Unit() instance serialises to {}.

Source code in jaxonomy/framework/units.py

def to_dict(self) -> dict:
    """Return a JSON-friendly dict representation of this Unit.

    Round-trips losslessly via :meth:`from_dict`. Keys are stable
    across versions; new optional fields are always added with
    defaults so older serialised forms continue to load.

    The default value for any field is omitted from the output for
    compactness — every legacy ``Unit()`` instance serialises to
    ``{}``.
    """
    out: dict = {}
    if self.dims != (0, 0, 0, 0, 0, 0, 0):
        out["dims"] = list(self.dims)
    if self.scale != 1.0:
        out["scale"] = self.scale
    if self.offset != 0.0:
        out["offset"] = self.offset
    if self.currency != _ZERO_CURRENCY:
        out["currency"] = list(self.currency)
    if self.name is not None:
        out["name"] = self.name
    if self.physical_quantity is not None:
        out["physical_quantity"] = self.physical_quantity
    return out

`to_json(*, indent=None)`

Serialise :meth:to_dict via :func:json.dumps.

Parameters:

Name	Type	Description	Default
`indent`	`int \| None`	Optional JSON indent (default `None` for compact form; pass an int for pretty-printed output).	`None`

Source code in jaxonomy/framework/units.py

def to_json(self, *, indent: int | None = None) -> str:
    """Serialise :meth:`to_dict` via :func:`json.dumps`.

    Args:
        indent: Optional JSON indent (default ``None`` for compact
            form; pass an int for pretty-printed output).
    """
    import json

    return json.dumps(self.to_dict(), indent=indent, sort_keys=True)

`UnitMismatchError`

Bases: StaticError

Raised at diagram build time when two connected ports have incompatible units.

Attributes are populated through :class:StaticError so the regular ErrorCollector / system-locator machinery still works.

Source code in jaxonomy/framework/units.py

class UnitMismatchError(StaticError):
    """Raised at diagram build time when two connected ports have
    incompatible units.

    Attributes are populated through :class:`StaticError` so the regular
    ``ErrorCollector`` / system-locator machinery still works.
    """

`Variant` `dataclass`

A frozen description of N variant choices for build-time selection.

Parameters:

Name	Type	Description	Default
`choices`	`Mapping[str, Callable[[], SystemBase]]`	Mapping from choice name to a zero-argument builder callable. Each callable, when invoked, must return a `SystemBase` (typically a fully-built `Diagram`). Unselected callables are never invoked.	required
`default`	`str`	Name of the choice to use when `select_variant` is called without an explicit `name`. Required (no implicit "first choice") so that adding a new variant later doesn't silently change the default. Must be a key in `choices`.	required
`name`	`Optional[str]`	Optional human-readable label for diagnostics / logging. Does not affect resolution.	`None`

Raises:

Type	Description
`VariantError`	If `choices` is empty, `default` is not in `choices`, or any choice is not callable.

Source code in jaxonomy/framework/variants.py

@dataclass(frozen=True)
class Variant:
    """A frozen description of N variant choices for build-time selection.

    Args:
        choices:
            Mapping from choice name to a zero-argument builder callable.
            Each callable, when invoked, must return a ``SystemBase``
            (typically a fully-built ``Diagram``). Unselected callables are
            never invoked.
        default:
            Name of the choice to use when ``select_variant`` is called
            without an explicit ``name``. Required (no implicit "first
            choice") so that adding a new variant later doesn't silently
            change the default. Must be a key in ``choices``.
        name:
            Optional human-readable label for diagnostics / logging. Does
            not affect resolution.

    Raises:
        VariantError: If ``choices`` is empty, ``default`` is not in
            ``choices``, or any choice is not callable.
    """

    choices: Mapping[str, Callable[[], SystemBase]]
    default: str
    name: Optional[str] = None

    def __post_init__(self):
        if not self.choices:
            raise VariantError(
                f"Variant {self.name!r}: choices must contain at least one entry"
            )
        for choice_name, builder in self.choices.items():
            if not callable(builder):
                raise VariantError(
                    f"Variant {self.name!r}: choice {choice_name!r} is not "
                    f"callable (got {type(builder).__name__}). Pass a "
                    f"zero-argument builder, e.g. ``lambda: build_pid()``."
                )
        if self.default not in self.choices:
            raise VariantError(
                f"Variant {self.name!r}: default {self.default!r} is not in "
                f"choices {list(self.choices)!r}"
            )

    @property
    def choice_names(self) -> tuple[str, ...]:
        """Stable tuple of available choice names (for introspection / CLI)."""
        return tuple(self.choices.keys())

`choice_names` `property`

Stable tuple of available choice names (for introspection / CLI).

`VariantError`

Bases: ValueError

Raised when a variant configuration is invalid or a selection is bad.

Source code in jaxonomy/framework/variants.py

class VariantError(ValueError):
    """Raised when a variant configuration is invalid or a selection is bad."""

`WhileLoop`

Bases: LeafSystem

Container block: run body_fn until cond_fn is False.

WhileLoop wraps :func:jax.lax.while_loop with a built-in iteration counter that caps execution at max_iter to guarantee termination under jit.

The block declares an input port for the initial carry value (port 0) plus n_inputs additional ports for upstream signals that the loop body / condition can consume. A single output port returns the carry after the loop exits (either because cond_fn returned False, or because max_iter was hit).

Parameters:

Name	Type	Description	Default
`body_fn`	`Callable`	Callable. Either `carry -> carry` (legacy) or `(carry, inputs) -> carry` when `n_inputs > 0`. Must be JAX-traceable. The signature is detected via :func:`inspect.signature`; functions that accept more than one positional argument (or `args`) receive all upstream input values. Callables that only need a subset should either accept the rest as throwaway positional args, or use `*args` and index into it.	required
`cond_fn`	`Callable`	Callable. Either `carry -> bool` (legacy) or `(carry, *inputs) -> bool`. Loop continues while this returns True. Signature detection mirrors `body_fn`; the same all-inputs-or-none rule applies.	required
`max_iter`	`int`	Positive integer cap on iterations. Required to keep traces bounded under jit. Defaults to 1000.	`1000`
`n_inputs`	`int`	Number of additional upstream input ports (default 0). When `n_inputs > 0` the block exposes ports `u_0..u_{n_inputs-1}` after the `carry_init` port. The current values of these inputs are passed to `cond_fn` / `body_fn` (if they accept them) on every iteration — so the condition can compare the carry against a live upstream signal (e.g. "iterate until the input exceeds a threshold").	`0`
`name`		Optional block name.	required

Differentiability

jax.grad flows through the carry as long as body_fn and cond_fn are pure. The number of iterations is data-dependent and not differentiable; lax.while_loop is itself non-differentiable in reverse mode (use jax.jvp for forward mode, or refactor with :func:jax.lax.scan if you need a reverse-mode-friendly bounded loop).

Notes

On hitting max_iter the loop exits silently. Users who want a runtime warning should jax.debug.callback from body_fn or test the post-loop carry.
The carry pytree structure must be invariant across iterations (a lax.while_loop requirement).
The condition is re-evaluated against the current upstream input values inside the loop trace — the inputs are captured once at output-evaluation time and held constant for the duration of the loop (the diagram doesn't re-tick during a single major step).

Source code in jaxonomy/framework/containers.py

class WhileLoop(LeafSystem):
    """Container block: run ``body_fn`` until ``cond_fn`` is False.

    ``WhileLoop`` wraps :func:`jax.lax.while_loop` with a built-in
    iteration counter that caps execution at ``max_iter`` to guarantee
    termination under jit.

    The block declares an input port for the *initial carry value*
    (port 0) plus ``n_inputs`` additional ports for upstream signals
    that the loop body / condition can consume. A single output port
    returns the carry after the loop exits (either because ``cond_fn``
    returned False, or because ``max_iter`` was hit).

    Args:
        body_fn: Callable. Either ``carry -> carry`` (legacy) or
            ``(carry, *inputs) -> carry`` when ``n_inputs > 0``. Must be
            JAX-traceable. The signature is detected via
            :func:`inspect.signature`; functions that accept more than
            one positional argument (or ``*args``) receive *all*
            upstream input values. Callables that only need a subset
            should either accept the rest as throwaway positional
            args, or use ``*args`` and index into it.
        cond_fn: Callable. Either ``carry -> bool`` (legacy) or
            ``(carry, *inputs) -> bool``. Loop continues while this
            returns True. Signature detection mirrors ``body_fn``; the
            same all-inputs-or-none rule applies.
        max_iter: Positive integer cap on iterations. Required to
            keep traces bounded under jit. Defaults to 1000.
        n_inputs: Number of additional upstream input ports (default 0).
            When ``n_inputs > 0`` the block exposes ports
            ``u_0..u_{n_inputs-1}`` after the ``carry_init`` port. The
            current values of these inputs are passed to ``cond_fn`` /
            ``body_fn`` (if they accept them) on every iteration — so
            the condition can compare the carry against a live upstream
            signal (e.g. "iterate until the input exceeds a threshold").
        name: Optional block name.

    Differentiability:
        ``jax.grad`` flows through the carry as long as ``body_fn`` and
        ``cond_fn`` are pure. The number of iterations is data-dependent
        and not differentiable; ``lax.while_loop`` is itself
        non-differentiable in reverse mode (use ``jax.jvp`` for forward
        mode, or refactor with :func:`jax.lax.scan` if you need a
        reverse-mode-friendly bounded loop).

    Notes:
        - On hitting ``max_iter`` the loop exits silently. Users who
          want a runtime warning should ``jax.debug.callback`` from
          ``body_fn`` or test the post-loop carry.
        - The carry pytree structure must be invariant across
          iterations (a ``lax.while_loop`` requirement).
        - The condition is re-evaluated against the *current* upstream
          input values inside the loop trace — the inputs are captured
          once at output-evaluation time and held constant for the
          duration of the loop (the diagram doesn't re-tick during a
          single major step).
    """

    def __init__(
        self,
        body_fn: Callable,
        cond_fn: Callable,
        max_iter: int = 1000,
        n_inputs: int = 0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if not isinstance(max_iter, int):
            raise TypeError(
                f"WhileLoop: max_iter must be a Python int (static), got "
                f"{type(max_iter).__name__}"
            )
        if max_iter <= 0:
            raise ValueError(
                f"WhileLoop: max_iter must be > 0 to keep traces bounded, "
                f"got {max_iter}"
            )
        if not isinstance(n_inputs, int):
            raise TypeError(
                f"WhileLoop: n_inputs must be a Python int, got "
                f"{type(n_inputs).__name__}"
            )
        if n_inputs < 0:
            raise ValueError(
                f"WhileLoop: n_inputs must be >= 0, got {n_inputs}"
            )

        self._body_fn = body_fn
        self._cond_fn = cond_fn
        self._max_iter = int(max_iter)
        self._n_inputs = int(n_inputs)

        # Signature inspection: detect whether the user-supplied
        # callables want the upstream inputs forwarded. We do this once
        # at construction so the JIT-compiled hot path doesn't pay any
        # introspection cost. The detection is intentionally permissive
        # — anything with more than one positional argument or a
        # ``*args`` gets the extra-args treatment. Legacy single-arg
        # callables are wrapped to ignore the inputs, preserving the
        # T-120-followup-loop-blocks contract byte-for-byte.
        self._body_takes_inputs = _accepts_extra_args(body_fn)
        self._cond_takes_inputs = _accepts_extra_args(cond_fn)

        # Port 0 is always the initial carry; ports 1..n_inputs carry
        # the upstream input signals forwarded to ``cond_fn`` /
        # ``body_fn``.
        self.declare_input_port(name="carry_init")
        for i in range(self._n_inputs):
            self.declare_input_port(name=f"u_{i}")

        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _compute_output(self, time, state, *inputs, **params):
        initial_carry = inputs[0]
        user_inputs = inputs[1:]  # the upstream signals (may be empty)
        max_iter = self._max_iter
        body_fn = self._body_fn
        cond_fn = self._cond_fn
        body_takes_inputs = self._body_takes_inputs
        cond_takes_inputs = self._cond_takes_inputs

        def safe_cond(loop_state):
            count, carry = loop_state
            if cond_takes_inputs:
                user_cond = cond_fn(carry, *user_inputs)
            else:
                user_cond = cond_fn(carry)
            user_cond = jnp.asarray(user_cond).astype(bool)
            return jnp.logical_and(count < max_iter, user_cond)

        def safe_body(loop_state):
            count, carry = loop_state
            if body_takes_inputs:
                new_carry = body_fn(carry, *user_inputs)
            else:
                new_carry = body_fn(carry)
            return (count + 1, new_carry)

        _final_count, final_carry = jax.lax.while_loop(
            safe_cond, safe_body, (jnp.asarray(0, dtype=jnp.int32), initial_carry)
        )
        return final_carry

`ZeroCrossingEvent` `dataclass`

Bases: Event

An event that triggers when a specified "guard" function crosses zero.

The event is triggered when the guard function crosses zero in the specified direction. In addition to the guard callback, the event also has a "reset map" which is called when the event is triggered. The reset map may update any state component in the system.

The event can also be defined as "terminal", which means that the simulation will terminate when the event is triggered. (TODO: Does the reset map still happen?)

The "direction" of the zero-crossing is one of the following: - "none": Never trigger the event (can be useful for debugging) - "positive_then_non_positive": Trigger when the guard goes from positive to non-positive - "negative_then_non_negative": Trigger when the guard goes from negative to non-negative - "crosses_zero": Trigger when the guard crosses zero in either direction - "edge_detection": Trigger when the guard changes value

Notes

This class should typically not need to be used directly by users. Instead, declare the guard function and reset map on a LeafSystem using the declare_zero_crossing method. The event will then be auto-generated for simulation.

Source code in jaxonomy/framework/event.py

@tree_util.register_pytree_node_class
@dataclasses.dataclass
class ZeroCrossingEvent(Event):
    """An event that triggers when a specified "guard" function crosses zero.

    The event is triggered when the guard function crosses zero in the specified
    direction. In addition to the guard callback, the event also has a "reset map"
    which is called when the event is triggered. The reset map may update any state
    component in the system.

    The event can also be defined as "terminal", which means that the simulation will
    terminate when the event is triggered. (TODO: Does the reset map still happen?)

    The "direction" of the zero-crossing is one of the following:
        - "none": Never trigger the event (can be useful for debugging)
        - "positive_then_non_positive": Trigger when the guard goes from positive to
            non-positive
        - "negative_then_non_negative": Trigger when the guard goes from negative to
            non-negative
        - "crosses_zero": Trigger when the guard crosses zero in either direction
        - "edge_detection": Trigger when the guard changes value

    Notes:
        This class should typically not need to be used directly by users. Instead,
        declare the guard function and reset map on a LeafSystem using the
        `declare_zero_crossing` method.  The event will then be auto-generated for
        simulation.
    """

    # Supersede type hints in Event with the specific signature for full-state updates
    callback: Callable[[ContextBase], LeafState] = None
    passthrough: Callable[[ContextBase], LeafState] = None

    guard: Callable[[ContextBase], Scalar] = None
    reset_map: dataclasses.InitVar[Callable[[ContextBase], LeafState]] = None
    direction: str = "crosses_zero"
    is_terminal: bool = False
    event_data: ZeroCrossingEventData = None

    # Optional *smooth* guard residual (``context -> scalar``) used ONLY by the
    # reverse-mode event-time (saltation) gradient machinery — never for
    # triggering or localization, which always use ``guard``.  When the trigger
    # ``guard`` is non-smooth (e.g. a boolean ``where(x>c, 1, -1)`` predicate as
    # a StateMachine emits), its gradient is identically zero and the
    # implicit-function event-time formula ``dt_e/dp = -∇g/D`` is unrecoverable;
    # supplying a smooth residual whose zero coincides with the trigger (e.g.
    # ``x - c``) lets ``∇g`` / ``D`` be taken from it instead.  ``None`` (the
    # default) means "use ``guard``" — byte-equivalent to the legacy path.
    # T-NEW-sm-smooth-guard.
    grad_guard: Callable[[ContextBase], Scalar] = None

    # If not none, only trigger when in this mode. This logic is handled by the owning
    # leaf system.
    active_mode: int = None

    def __post_init__(self, reset_map):  # pylint: disable=arguments-differ
        if self.callback is None:
            self.callback = reset_map

    def _should_trigger(self, w0: Scalar, w1: Scalar) -> bool:
        """Determine if the event should trigger.

        This will use the provided beginning/ending guard value (w0 and w1, resp.),
        as well as the direction of the zero-crossing event. Additionally, the event
        will only trigger if it has been marked as "active", indicating for example
        that the system is in the correct "mode" or "stage" from which the event might
        trigger.
        """
        active = self.event_data.active

        trigger_func = _zero_crossing_trigger_functions[self.direction]
        return active & trigger_func(w0, w1)

    def should_trigger(self) -> bool:
        """Determine if the event should trigger based on the stored guard values."""
        return self._should_trigger(self.event_data.w0, self.event_data.w1)

    def handle(self, context: ContextBase) -> LeafState:
        """Conditionally compute the result of the zero crossing callback

        If the zero crossing is marked "inactive" via its event data attribute, the passthrough
        callback will be called instead of the update callback. Otherwise, the update
        callback will be called. The return types of both callbacks must match, but the
        specific type will depend on the kind of event.
        """
        if self.enable_tracing:  # not driven by simulator.enable_tracing.
            return _handle_with_severed_discretes(
                self.event_data.active & self.event_data.triggered,
                self.callback,
                self.passthrough,
                context,
            )

        # No tracing: use standard control flow
        if self.event_data.active & self.event_data.triggered:
            return self.callback(context)
        return self.passthrough(context)

    #
    # PyTree registration
    #
    def tree_flatten(self):
        children = (self.event_data,)
        aux_data = (
            self.system_id,
            self.guard,
            self.callback,
            self.name,
            self.direction,
            self.is_terminal,
            self.passthrough,
            self.enable_tracing,
            self.active_mode,
            self.grad_guard,
        )
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        (event_data,) = children
        (
            system_id,
            guard,
            callback,
            name,
            direction,
            is_terminal,
            passthrough,
            enable_tracing,
            active_mode,
            grad_guard,
        ) = aux_data
        return cls(
            system_id=system_id,
            event_data=event_data,
            guard=guard,
            grad_guard=grad_guard,
            callback=callback,
            name=name,
            direction=direction,
            is_terminal=is_terminal,
            passthrough=passthrough,
            enable_tracing=enable_tracing,
            active_mode=active_mode,
        )

`handle(context)`

Conditionally compute the result of the zero crossing callback

If the zero crossing is marked "inactive" via its event data attribute, the passthrough callback will be called instead of the update callback. Otherwise, the update callback will be called. The return types of both callbacks must match, but the specific type will depend on the kind of event.

Source code in jaxonomy/framework/event.py

def handle(self, context: ContextBase) -> LeafState:
    """Conditionally compute the result of the zero crossing callback

    If the zero crossing is marked "inactive" via its event data attribute, the passthrough
    callback will be called instead of the update callback. Otherwise, the update
    callback will be called. The return types of both callbacks must match, but the
    specific type will depend on the kind of event.
    """
    if self.enable_tracing:  # not driven by simulator.enable_tracing.
        return _handle_with_severed_discretes(
            self.event_data.active & self.event_data.triggered,
            self.callback,
            self.passthrough,
            context,
        )

    # No tracing: use standard control flow
    if self.event_data.active & self.event_data.triggered:
        return self.callback(context)
    return self.passthrough(context)

`should_trigger()`

Determine if the event should trigger based on the stored guard values.

Source code in jaxonomy/framework/event.py

def should_trigger(self) -> bool:
    """Determine if the event should trigger based on the stored guard values."""
    return self._should_trigger(self.event_data.w0, self.event_data.w1)

`ZeroCrossingTriggeredSubsystem`

Bases: LeafSystem

Container block: latch the submodel output at zero-crossings.

Like :class:TriggeredSubsystem, but uses the framework's continuous zero-crossing detector rather than a periodic sample grid. The submodel fires exactly when the trigger signal crosses zero in the configured direction — this gives sub-sample-period precision for the latched event time, which is the property normally expected from a triggered subsystem driven by a continuous signal.

Wiring matches :class:TriggeredSubsystem:

Input port 0 is the trigger signal (a continuous scalar; the block monitors its sign).
Input ports 1..n_inputs are the submodel inputs.
The single output port returns the most recently latched submodel output (initialized to initial_value).

Parameters:

Name	Type	Description	Default
`submodel`	`Callable`	Callable `f(*inputs) -> output` taking the non-trigger user inputs. Must be JAX-traceable.	required
`n_inputs`	`int`	Number of user inputs (NOT counting the trigger).	`1`
`edge`	`Literal['rising', 'falling', 'either']`	`"rising"` (low→high zero crossing of the trigger signal), `"falling"` (high→low), or `"either"`.	`RISING`
`initial_value`		Latched output value before the first crossing fires. Also defines the output shape/dtype.	`0.0`
`name`		Optional block name.	required

Differentiability

jax.grad flows through the submodel inputs along the path through the latch (so when the latched value depends on a differentiable input, the gradient propagates). The trigger signal itself is consumed by the zero-crossing event detector; the gradient through the discontinuity at the firing instant is zero by design (the latched value is constant between crossings).

Notes

The framework localizes the zero crossing to within the integrator's tolerance, so the latched output reflects the submodel inputs at the crossing instant, not at the next periodic sample. Compare with the phase-1 :class:TriggeredSubsystem, which can only resolve the edge to the nearest sample_period.
The latch is a single discrete-state component; the submodel must produce a single output array of fixed shape.
This is a leaf block (no nested mode machinery), so the "hold between crossings" semantics fall out naturally: the output port simply returns state.discrete_state.

Source code in jaxonomy/framework/containers.py

class ZeroCrossingTriggeredSubsystem(LeafSystem):
    """Container block: latch the submodel output at zero-crossings.

    Like :class:`TriggeredSubsystem`, but uses the framework's continuous
    zero-crossing detector rather than a periodic sample grid. The
    submodel fires *exactly* when the trigger signal crosses zero in the
    configured direction — this gives sub-sample-period precision for
    the latched event time, which is the property normally expected
    from a triggered subsystem driven by a continuous signal.

    Wiring matches :class:`TriggeredSubsystem`:

    - Input port 0 is the trigger signal (a continuous scalar; the
      block monitors its sign).
    - Input ports 1..n_inputs are the submodel inputs.
    - The single output port returns the most recently latched
      submodel output (initialized to ``initial_value``).

    Args:
        submodel: Callable ``f(*inputs) -> output`` taking the
            non-trigger user inputs. Must be JAX-traceable.
        n_inputs: Number of user inputs (NOT counting the trigger).
        edge: ``"rising"`` (low→high zero crossing of the trigger
            signal), ``"falling"`` (high→low), or ``"either"``.
        initial_value: Latched output value before the first crossing
            fires. Also defines the output shape/dtype.
        name: Optional block name.

    Differentiability:
        ``jax.grad`` flows through the submodel inputs along the path
        through the latch (so when the latched value depends on a
        differentiable input, the gradient propagates). The trigger
        signal itself is consumed by the zero-crossing event detector;
        the gradient through the discontinuity at the firing instant
        is zero by design (the latched value is constant between
        crossings).

    Notes:
        - The framework localizes the zero crossing to within the
          integrator's tolerance, so the latched output reflects the
          submodel inputs *at the crossing instant*, not at the next
          periodic sample. Compare with the phase-1
          :class:`TriggeredSubsystem`, which can only resolve the edge
          to the nearest ``sample_period``.
        - The latch is a single discrete-state component; the submodel
          must produce a single output array of fixed shape.
        - This is a leaf block (no nested mode machinery), so the
          ``"hold between crossings"`` semantics fall out naturally:
          the output port simply returns ``state.discrete_state``.
    """

    def __init__(
        self,
        submodel: Callable,
        n_inputs: int = 1,
        edge: Literal["rising", "falling", "either"] = TriggerEdge.RISING,
        initial_value=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if edge not in TriggerEdge.valid():
            raise ValueError(
                f"ZeroCrossingTriggeredSubsystem: edge must be one of "
                f"{TriggerEdge.valid()!r}, got {edge!r}"
            )
        if n_inputs < 0:
            raise ValueError(
                f"ZeroCrossingTriggeredSubsystem: n_inputs must be >= 0, "
                f"got {n_inputs}"
            )

        self._submodel = submodel
        self._edge = edge
        self._initial = jnp.asarray(initial_value)

        # Port 0 is the trigger signal; ports 1..n_inputs are the user
        # inputs forwarded to the submodel.
        self.declare_input_port(name="trigger")
        for i in range(n_inputs):
            self.declare_input_port(name=f"u_{i}")

        # Discrete state: the latched submodel output.
        self.declare_discrete_state(default_value=self._initial)

        # Zero-crossing event: guard returns the trigger-signal value;
        # reset_map evaluates the submodel and writes the result to the
        # discrete state.
        self.declare_zero_crossing(
            guard=self._guard,
            reset_map=self._on_trigger,
            direction=_EDGE_TO_DIRECTION[edge],
            name="zc_trigger",
        )

        # Output reads the latched value.
        self.declare_output_port(
            self._compute_output,
            prerequisites_of_calc=[port.ticket for port in self.input_ports],
            default_value=self._initial,
        )

    # ── callbacks ─────────────────────────────────────────────────────────

    def _guard(self, time, state, *inputs, **params):
        # The trigger is input port 0. Cast to the framework's working
        # dtype so the zero-crossing solver sees a plain scalar.
        return jnp.asarray(inputs[0])

    def _on_trigger(self, time, state, *inputs, **params):
        # Run the submodel on the user inputs (everything past the
        # trigger) and latch the result into the discrete state.
        user_inputs = inputs[1:]
        new_value = jnp.asarray(self._submodel(*user_inputs))
        # Preserve dtype/shape of the discrete state; jnp.broadcast_to
        # is safe for the scalar-latch case and also handles the
        # rank-preserving case when initial_value supplied a shape.
        new_value = jnp.asarray(new_value, dtype=state.discrete_state.dtype)
        new_value = jnp.reshape(new_value, state.discrete_state.shape)
        return state.with_discrete_state(new_value)

    def _compute_output(self, time, state, *inputs, **params):
        return state.discrete_state

`ForEach(submodel, n, n_inputs=1, in_axes=None, name=None)`

Container block: evaluate a submodel n times in parallel.

ForEach is a block-diagram-vocabulary alias for the existing :class:jaxonomy.library.ReplicatedFunction (T-010). It exists so that users familiar with the ForEach block name can find it without paying a duplication tax: the implementation is exactly :class:ReplicatedFunction under the hood.

Parameters:

Name	Type	Description	Default
`submodel`	`Callable`	Callable `f(*inputs) -> output`. Must be JAX-traceable.	required
`n`	`int`	Number of replicas (the iteration count).	required
`n_inputs`	`int`	Number of input ports the block declares.	`1`
`in_axes`		As in :func:`jax.vmap` / `ReplicatedFunction`: a length-`n_inputs` tuple of `0` (input is batched along the leading axis) or `None` (input is broadcast). Default is all-batched.	`None`
`name`	`str \| None`	Optional block name.	`None`

Returns:

Type	Description
	A configured :class:`ReplicatedFunction` instance, ready to be
	wired into a :class:`DiagramBuilder`.

Source code in jaxonomy/framework/containers.py

def ForEach(
    submodel: Callable,
    n: int,
    n_inputs: int = 1,
    in_axes=None,
    name: str | None = None,
):
    """Container block: evaluate a submodel ``n`` times in parallel.

    ``ForEach`` is a block-diagram-vocabulary alias for the existing
    :class:`jaxonomy.library.ReplicatedFunction` (T-010). It exists so
    that users familiar with the ``ForEach`` block name can find it
    without paying a duplication tax: the implementation
    is exactly :class:`ReplicatedFunction` under the hood.

    Args:
        submodel: Callable ``f(*inputs) -> output``. Must be
            JAX-traceable.
        n: Number of replicas (the iteration count).
        n_inputs: Number of input ports the block declares.
        in_axes: As in :func:`jax.vmap` / ``ReplicatedFunction``: a
            length-``n_inputs`` tuple of ``0`` (input is batched along
            the leading axis) or ``None`` (input is broadcast). Default
            is all-batched.
        name: Optional block name.

    Returns:
        A configured :class:`ReplicatedFunction` instance, ready to be
        wired into a :class:`DiagramBuilder`.
    """
    # Lazy import: ReplicatedFunction lives in jaxonomy.library, which
    # imports the framework. Importing it eagerly here would create a
    # cycle. The lazy import is exercised only when a user actually
    # constructs a ForEach block.
    from ..library.replicated import ReplicatedFunction

    kwargs: dict = {}
    if name is not None:
        kwargs["name"] = name
    return ReplicatedFunction(
        submodel=submodel,
        n=n,
        n_inputs=n_inputs,
        in_axes=in_axes,
        **kwargs,
    )

`apply_variant_config(diagram, **overrides)`

Return a copy of diagram with named variants reconfigured.

Walks the diagram tree, finds every subsystem that was produced by :func:select_variant from a named Variant, and -- for each override_name=choice keyword -- replaces matching subsystems with a freshly-built copy from select_variant(variant, name=choice). All other diagram structure (non-variant blocks, connections, exported ports) is preserved.

The original diagram is not modified.

Example::

builder = DiagramBuilder()
ctrl = select_variant(controller_variant, name="pid")  # default
plant = select_variant(plant_variant, name="lti")
builder.add(ctrl)
builder.add(plant)
...
diagram = builder.build()

# Reconfigure post-build:
runtime_a = apply_variant_config(diagram, controller="pid", plant="lti")
runtime_b = apply_variant_config(diagram, controller="lqr", plant="lti")

Parameters:

Name	Type	Description	Default
`diagram`		A built `Diagram` (typically the output of `DiagramBuilder.build`).	required
`**overrides`		Map from a variant's `name` (the `name=` kwarg passed to :class:`Variant`) to the choice name to activate. Variants whose name is not mentioned in `overrides` keep their currently-active choice.	`{}`

Returns:

Type	Description
	A new `Diagram` with the requested variant choices resolved.
	If `overrides` is empty, returns a structurally identical deep
	copy of `diagram` (the same default-off semantics as
	meth:`Diagram.with_parameters` with no updates).

Raises:

Type	Description
`VariantError`	If an override name does not match any variant in the diagram, or if the requested choice is not in that variant's `choices`.

Source code in jaxonomy/framework/variants.py

def apply_variant_config(diagram, **overrides):
    """Return a copy of ``diagram`` with named variants reconfigured.

    Walks the diagram tree, finds every subsystem that was produced by
    :func:`select_variant` from a named ``Variant``, and -- for each
    ``override_name=choice`` keyword -- replaces matching subsystems with
    a freshly-built copy from ``select_variant(variant, name=choice)``.
    All other diagram structure (non-variant blocks, connections,
    exported ports) is preserved.

    The original diagram is not modified.

    Example::

        builder = DiagramBuilder()
        ctrl = select_variant(controller_variant, name="pid")  # default
        plant = select_variant(plant_variant, name="lti")
        builder.add(ctrl)
        builder.add(plant)
        ...
        diagram = builder.build()

        # Reconfigure post-build:
        runtime_a = apply_variant_config(diagram, controller="pid", plant="lti")
        runtime_b = apply_variant_config(diagram, controller="lqr", plant="lti")

    Args:
        diagram: A built ``Diagram`` (typically the output of
            ``DiagramBuilder.build``).
        **overrides: Map from a variant's ``name`` (the ``name=`` kwarg
            passed to :class:`Variant`) to the choice name to activate.
            Variants whose name is not mentioned in ``overrides`` keep
            their currently-active choice.

    Returns:
        A new ``Diagram`` with the requested variant choices resolved.
        If ``overrides`` is empty, returns a structurally identical deep
        copy of ``diagram`` (the same default-off semantics as
        :meth:`Diagram.with_parameters` with no updates).

    Raises:
        VariantError: If an override name does not match any variant in
            the diagram, or if the requested choice is not in that
            variant's ``choices``.
    """
    # Local imports to avoid an import cycle at module load time.
    import copy as _copy
    from .diagram import (
        Diagram,
        _diagram_rewrite_child_refs,
        _diagram_refresh_exported_outputs_for_child,
        _diagram_rebuild_leaf_systems,
    )
    from .system_base import next_system_id

    if not isinstance(diagram, Diagram):
        raise VariantError(
            f"apply_variant_config: expected a Diagram, got "
            f"{type(diagram).__name__}."
        )

    new = _copy.deepcopy(diagram)
    new.system_id = next_system_id()
    new.parent = None
    new._dependency_graph = None
    new.feedthrough_pairs = None
    new._cache_update_events = None

    if not overrides:
        # Default-off: no overrides → identical-equivalent diagram.
        _diagram_rebuild_leaf_systems(new)
        return new

    # Index every tagged subsystem in the (deep-copied) tree by variant name.
    found_by_name: dict[str, list[tuple[Diagram, int, _VariantMetadata]]] = {}
    for parent_d, idx, tag in _iter_tagged(new):
        vname = tag.variant.name
        if vname is None:
            # Anonymous variant — can't be addressed by name. Skip; the
            # error path below will still fire for any unmatched override.
            continue
        found_by_name.setdefault(vname, []).append((parent_d, idx, tag))

    # Validate: every override must match at least one tagged variant.
    unknown = [k for k in overrides if k not in found_by_name]
    if unknown:
        raise VariantError(
            f"apply_variant_config: no Variant with name in {unknown!r} "
            f"found in diagram {diagram.name!r}. Available variant "
            f"names: {sorted(found_by_name)!r}. (Anonymous Variants "
            f"-- those built without a ``name=`` kwarg -- cannot be "
            f"addressed by ``apply_variant_config``.)"
        )

    # Apply each override.
    for vname, choice in overrides.items():
        for parent_d, idx, tag in found_by_name[vname]:
            if choice not in tag.variant.choices:
                raise VariantError(
                    f"apply_variant_config: variant {vname!r}: unknown "
                    f"choice {choice!r}; available: "
                    f"{list(tag.variant.choices)!r}"
                )
            old_child = parent_d.nodes[idx]
            # Build a fresh subsystem from the requested choice. This
            # re-tags the result with updated metadata, so subsequent
            # apply_variant_config calls keep working.
            repl = select_variant(tag.variant, name=choice)
            parent_d.nodes[idx] = repl
            repl.parent = parent_d
            _diagram_rewrite_child_refs(parent_d, old_child, repl)
            _diagram_refresh_exported_outputs_for_child(parent_d, repl)

    _diagram_rebuild_leaf_systems(new)
    return new

`are_units_compatible(src, dst)`

Return True if a connection from src to dst should be allowed under the Phase-1 rules:

Either side being None (unset) is always OK.
If both sides are :class:BusUnit, compatible iff every shared field's :class:Unit is pair-wise compatible AND the field sets match. A BusUnit on one side and None on the other is always OK (default-off byte-equivalence with the no-units bus).
Otherwise, both sides must be plain :class:Unit; either being :data:dimensionless is OK, else units must be equal.

Scalar conversion (Phase 2) is layered on top by :func:assert_units_compatible_with_scale — see there.

Source code in jaxonomy/framework/units.py

def are_units_compatible(
    src: Unit | BusUnit | None,
    dst: Unit | BusUnit | None,
) -> bool:
    """Return True if a connection from ``src`` to ``dst`` should be
    allowed under the Phase-1 rules:

    * Either side being ``None`` (unset) is always OK.
    * If both sides are :class:`BusUnit`, compatible iff every shared
      field's :class:`Unit` is pair-wise compatible AND the field sets
      match. A ``BusUnit`` on one side and ``None`` on the other is
      always OK (default-off byte-equivalence with the no-units bus).
    * Otherwise, both sides must be plain :class:`Unit`; either being
      :data:`dimensionless` is OK, else units must be equal.

    Scalar conversion (Phase 2) is layered on top by
    :func:`assert_units_compatible_with_scale` — see there.
    """
    # T-117-followup-bus-units: handle compound bus-unit case first.
    if isinstance(src, BusUnit) or isinstance(dst, BusUnit):
        # ``None`` wildcards match anything (preserves default-off
        # byte-equivalence with the unit-less BusCreator).
        if src is None or dst is None:
            return True
        if not (isinstance(src, BusUnit) and isinstance(dst, BusUnit)):
            # BusUnit cannot connect to a scalar Unit and vice-versa.
            return False
        if set(src.fields.keys()) != set(dst.fields.keys()):
            return False
        return all(
            are_units_compatible(src.fields[k], dst.fields[k])
            for k in src.fields
        )

    src_u = resolve_unit(src)
    dst_u = resolve_unit(dst)
    if src_u.is_dimensionless or dst_u.is_dimensionless:
        return True
    # T-104 phase 3: dims/scale/offset/currency must match, AND the
    # physical_quantity disambiguation tag must agree when both sides
    # carry one. One side ``None`` is treated as a wildcard so
    # legacy callers that never set the tag stay byte-equivalent.
    if not (
        src_u.dims == dst_u.dims
        and src_u.scale == dst_u.scale
        and src_u.offset == dst_u.offset
        and src_u.currency == dst_u.currency
    ):
        return False
    if src_u.physical_quantity is None or dst_u.physical_quantity is None:
        return True
    return src_u.physical_quantity == dst_u.physical_quantity

`assert_unit_compatible(src, dst, *, src_label='source port', dst_label='destination port')`

Raise :class:UnitMismatchError if the two units are not Phase-1-compatible. See :func:are_units_compatible.

The labels are interpolated into the message so the caller (typically :meth:DiagramBuilder.connect) can name both ports.

Source code in jaxonomy/framework/units.py

def assert_unit_compatible(
    src: Unit | BusUnit | None,
    dst: Unit | BusUnit | None,
    *,
    src_label: str = "source port",
    dst_label: str = "destination port",
) -> None:
    """Raise :class:`UnitMismatchError` if the two units are not
    Phase-1-compatible.  See :func:`are_units_compatible`.

    The labels are interpolated into the message so the caller (typically
    :meth:`DiagramBuilder.connect`) can name both ports.
    """
    if are_units_compatible(src, dst):
        return
    # T-117-followup-bus-units: render BusUnit and Unit cases distinctly
    # so the error message points at the actual mismatched value rather
    # than passing a BusUnit through ``resolve_unit`` (which would
    # collapse it to the dimensionless sentinel).
    src_repr = src if isinstance(src, BusUnit) else resolve_unit(src)
    dst_repr = dst if isinstance(dst, BusUnit) else resolve_unit(dst)
    raise UnitMismatchError(
        f"Unit mismatch: {src_label} has units {src_repr!r} but "
        f"{dst_label} has units {dst_repr!r}."
    )

`clear_fx_rates()`

Empty the FX rate table. Tests use this to keep their state isolated; production code should rarely need to call it.

Source code in jaxonomy/framework/units.py

def clear_fx_rates() -> None:
    """Empty the FX rate table. Tests use this to keep their state
    isolated; production code should rarely need to call it.
    """
    _FX_RATES.clear()
    _FX_AUTO_REVERSE.clear()

`convert_currency(value, from_unit, to_unit)`

Convert a numeric value carried in from_unit to the equivalent value in to_unit using the current FX rate table.

Self-conversion (same currency on both sides) is a no-op and returns the value unchanged. Cross-currency conversion looks up the rate via :func:get_fx_rate and multiplies; a missing rate raises :class:KeyError.

Parameters:

Name	Type	Description	Default
`value`		Numeric value (Python scalar, NumPy array, JAX array). The helper only uses `*`, so it composes transparently through `jit` / `vmap` / `grad`.	required
`from_unit`	`'Unit \| str'`	Source currency, either a :class:`Unit` (such as :data:`usd`) or a string code (`"USD"`).	required
`to_unit`	`'Unit \| str'`	Destination currency, ditto.	required

Returns:

Type	Description
	`value * rate` where `rate = get_fx_rate(from_unit, to_unit)`.

Raises:

Type	Description
`UnitMismatchError`	if either argument carries non-currency dimensions (e.g. seconds), so the conversion is undefined.
`KeyError`	if the relevant FX rate has not been registered.

Example:

>>> set_fx_rate("USD", "EUR", 0.92)
>>> convert_currency(100.0, usd, eur)
92.0

Source code in jaxonomy/framework/units.py

def convert_currency(
    value,
    from_unit: "Unit | str",
    to_unit: "Unit | str",
):
    """Convert a numeric ``value`` carried in ``from_unit`` to the
    equivalent value in ``to_unit`` using the current FX rate table.

    Self-conversion (same currency on both sides) is a no-op and
    returns the value unchanged. Cross-currency conversion looks up
    the rate via :func:`get_fx_rate` and multiplies; a missing rate
    raises :class:`KeyError`.

    Args:
        value: Numeric value (Python scalar, NumPy array, JAX array).
            The helper only uses ``*``, so it composes transparently
            through ``jit`` / ``vmap`` / ``grad``.
        from_unit: Source currency, either a :class:`Unit` (such as
            :data:`usd`) or a string code (``"USD"``).
        to_unit: Destination currency, ditto.

    Returns:
        ``value * rate`` where ``rate = get_fx_rate(from_unit, to_unit)``.

    Raises:
        UnitMismatchError: if either argument carries non-currency
            dimensions (e.g. seconds), so the conversion is undefined.
        KeyError: if the relevant FX rate has not been registered.

    Example:

        >>> set_fx_rate("USD", "EUR", 0.92)
        >>> convert_currency(100.0, usd, eur)
        92.0
    """
    # Both sides must be pure currency units — otherwise it's a
    # mistake (a USD value cannot be converted to seconds). Reuse
    # the canonical-currency resolver so the error message names the
    # actual offending unit.
    src = _canonical_currency_code(from_unit)
    dst = _canonical_currency_code(to_unit)
    if src == dst:
        return value
    rate = _FX_RATES.get((src, dst))
    if rate is None:
        raise KeyError(
            f"No FX rate registered for {src}->{dst}. "
            f"Call set_fx_rate({src!r}, {dst!r}, ...) first."
        )
    return value * rate

`derived_unit(name, symbol=None, components=None)`

Define a new derived :class:Unit from existing components.

This is a convenience constructor for users who want to spell a composite unit once (with a friendly name) rather than recomposing its base components at every port declaration site. The returned :class:Unit has the same (dims, scale, offset) as components — it is therefore equal (under Unit.__eq__) to any other unit with matching dimensions and scale — but carries a custom name for friendlier error messages and pprint output.

Parameters:

Name	Type	Description	Default
`name`	`str`	Long-form descriptive name (e.g. `"my_torque"`). Used only when `symbol` is omitted.	required
`symbol`	`str \| None`	Short-form printable symbol (e.g. `"τ"`). When provided, it overrides `name` as the Unit's display label.	`None`
`components`	`'Unit \| None'`	A :class:`Unit` expression describing the dimensions of the new unit (e.g. `meter * newton`). Must not be `None` and must have `offset == 0` — affine units cannot be re-aliased this way.	`None`

Returns:

Type	Description
`'Unit'`	A fresh :class:`Unit` with `components.dims` /
`'Unit'`	`components.scale` / `components.offset` and a `name`
`'Unit'`	set to `symbol` (when provided) or `name`.

Raises:

Type	Description
`TypeError`	if `components` is not a :class:`Unit`.
`UnitMismatchError`	if `components` has a non-zero offset (affine units cannot be aliased).

Example:

>>> from jaxonomy.framework.units import (
...     derived_unit, meter, newton,
... )
>>> torque = derived_unit("torque", "N·m", meter * newton)
>>> torque.dims == (1, 2, -2, 0, 0, 0, 0)
True
>>> torque == meter * newton
True

Source code in jaxonomy/framework/units.py

def derived_unit(
    name: str,
    symbol: str | None = None,
    components: "Unit | None" = None,
) -> "Unit":
    """Define a new derived :class:`Unit` from existing components.

    This is a convenience constructor for users who want to spell a
    composite unit once (with a friendly name) rather than recomposing
    its base components at every port declaration site. The returned
    :class:`Unit` has the same ``(dims, scale, offset)`` as
    ``components`` — it is therefore equal (under ``Unit.__eq__``) to
    any other unit with matching dimensions and scale — but carries
    a custom ``name`` for friendlier error messages and pprint output.

    Args:
        name: Long-form descriptive name (e.g. ``"my_torque"``).
            Used only when ``symbol`` is omitted.
        symbol: Short-form printable symbol (e.g. ``"τ"``). When
            provided, it overrides ``name`` as the Unit's display label.
        components: A :class:`Unit` expression describing the
            dimensions of the new unit (e.g. ``meter * newton``).
            Must not be ``None`` and must have ``offset == 0`` — affine
            units cannot be re-aliased this way.

    Returns:
        A fresh :class:`Unit` with ``components.dims`` /
        ``components.scale`` / ``components.offset`` and a ``name``
        set to ``symbol`` (when provided) or ``name``.

    Raises:
        TypeError: if ``components`` is not a :class:`Unit`.
        UnitMismatchError: if ``components`` has a non-zero offset
            (affine units cannot be aliased).

    Example:

        >>> from jaxonomy.framework.units import (
        ...     derived_unit, meter, newton,
        ... )
        >>> torque = derived_unit("torque", "N·m", meter * newton)
        >>> torque.dims == (1, 2, -2, 0, 0, 0, 0)
        True
        >>> torque == meter * newton
        True
    """
    if components is None or not isinstance(components, Unit):
        raise TypeError(
            "derived_unit(...) requires a `components` Unit expression, "
            f"got {type(components).__name__}: {components!r}"
        )
    if components.offset != 0.0:
        raise UnitMismatchError(
            f"derived_unit({name!r}, ...) cannot alias an offsetted "
            f"(affine) unit {components!r}; affine units do not form a "
            "multiplicative algebra and cannot be re-aliased."
        )
    label = symbol if symbol is not None else name
    return Unit(
        dims=components.dims,
        scale=components.scale,
        offset=components.offset,
        name=label,
    )

`flatten_diagram(diagram)`

Flatten a nested Diagram into a single-depth Diagram.

All intermediate sub-Diagrams are dissolved. The resulting Diagram has: - nodes: all LeafSystem instances from the original tree - connection_map: remapped to only reference leaf-to-leaf connections - exported inputs/outputs: preserved (still reference the same leaf ports)

Parameters:

Name	Type	Description	Default
`diagram`	`Diagram`	The (possibly nested) Diagram to flatten.	required

Returns:

Type	Description
`Diagram`	A new single-depth Diagram with all original LeafSystems as direct
`Diagram`	children and all connections resolved to the leaf level.

Source code in jaxonomy/framework/flatten.py

def flatten_diagram(diagram: Diagram) -> Diagram:
    """Flatten a nested Diagram into a single-depth Diagram.

    All intermediate sub-Diagrams are dissolved. The resulting Diagram has:
    - nodes: all LeafSystem instances from the original tree
    - connection_map: remapped to only reference leaf-to-leaf connections
    - exported inputs/outputs: preserved (still reference the same leaf ports)

    Args:
        diagram: The (possibly nested) Diagram to flatten.

    Returns:
        A new single-depth Diagram with all original LeafSystems as direct
        children and all connections resolved to the leaf level.
    """
    # If there are no sub-Diagrams, nothing to flatten
    if not any(isinstance(node, Diagram) for node in diagram.nodes):
        return diagram

    # Gather all leaf systems
    leaf_systems = list(diagram.leaf_systems)

    # Collect all leaf-to-leaf connections
    connections = _collect_connections(diagram)

    # Collect diagram-level port exports
    exported_inputs = _collect_exported_inputs(diagram)
    exported_outputs = _collect_exported_outputs(diagram)

    # Build the new flat diagram
    builder = DiagramBuilder()
    for leaf in leaf_systems:
        # Reset parent so the builder doesn't complain about re-registration
        leaf.parent = None
        builder.add(leaf)

    # Wire up all connections
    for input_loc, output_loc in connections:
        input_sys, input_idx = input_loc
        output_sys, output_idx = output_loc
        builder.connect(
            output_sys.output_ports[output_idx],
            input_sys.input_ports[input_idx],
        )

    # Re-export diagram-level ports
    for input_loc, port_name in exported_inputs:
        input_sys, input_idx = input_loc
        builder.export_input(input_sys.input_ports[input_idx], name=port_name)

    for output_loc, port_name in exported_outputs:
        output_sys, output_idx = output_loc
        builder.export_output(output_sys.output_ports[output_idx], name=port_name)

    return builder.build(name=diagram.name)

`get_active_variant(diagram, variant_name)`

Return the currently-selected choice for the named variant.

Parameters:

Name	Type	Description	Default
`diagram`		A built `Diagram`.	required
`variant_name`	`str`	The human-readable label of the variant to look up (the `name=` kwarg passed to :class:`Variant`).	required

Returns:

Type	Description
	The name of the active choice (a string in
	`Variant.choice_names`), or `None` if no variant with the
	given name is found in the diagram. The `None` sentinel lets
	CLI / introspection code treat "no such variant" as a soft
	miss; use :func:`get_variant_choices` if you want a hard error
	for unknown names.

Source code in jaxonomy/framework/variants.py

def get_active_variant(diagram, variant_name: str):
    """Return the currently-selected choice for the named variant.

    Args:
        diagram: A built ``Diagram``.
        variant_name: The human-readable label of the variant to look
            up (the ``name=`` kwarg passed to :class:`Variant`).

    Returns:
        The name of the active choice (a string in
        ``Variant.choice_names``), or ``None`` if no variant with the
        given name is found in the diagram. The ``None`` sentinel lets
        CLI / introspection code treat "no such variant" as a soft
        miss; use :func:`get_variant_choices` if you want a hard error
        for unknown names.
    """
    tag = _find_first_tag_by_name(diagram, variant_name)
    if tag is None:
        return None
    return tag.active_choice

`get_fx_rate(from_currency, to_currency)`

Return the previously-set FX rate from from_currency to to_currency. Self-rates are always 1.0 even when unset.

Raises:

Type	Description
`KeyError`	if no rate has been set for the requested pair AND the two codes differ.
`UnitMismatchError`	if either argument is not a pure currency.

Source code in jaxonomy/framework/units.py

def get_fx_rate(
    from_currency: "Unit | str",
    to_currency: "Unit | str",
) -> float:
    """Return the previously-set FX rate from ``from_currency`` to
    ``to_currency``. Self-rates are always ``1.0`` even when unset.

    Raises:
        KeyError: if no rate has been set for the requested pair AND
            the two codes differ.
        UnitMismatchError: if either argument is not a pure currency.
    """
    src = _canonical_currency_code(from_currency)
    dst = _canonical_currency_code(to_currency)
    if src == dst:
        return 1.0
    try:
        return _FX_RATES[(src, dst)]
    except KeyError as e:
        raise KeyError(
            f"No FX rate registered for {src}->{dst}. "
            f"Call set_fx_rate({src!r}, {dst!r}, ...) first."
        ) from e

`get_variant_choices(diagram, variant_name)`

Return the choice names of the named variant in diagram.

Parameters:

Name	Type	Description	Default
`diagram`		A built `Diagram`.	required
`variant_name`	`str`	The human-readable label of the variant to look up (the `name=` kwarg passed to :class:`Variant`).	required

Returns:

Type	Description
`tuple`	A tuple of choice names (`Variant.choice_names`), in
`tuple`	insertion order.

Raises:

Type	Description
`VariantError`	If no variant with the given name is found in the diagram. Anonymous Variants (built without `name=`) are never matched and so cannot be queried via this helper.

Source code in jaxonomy/framework/variants.py

def get_variant_choices(diagram, variant_name: str) -> tuple:
    """Return the choice names of the named variant in ``diagram``.

    Args:
        diagram: A built ``Diagram``.
        variant_name: The human-readable label of the variant to look
            up (the ``name=`` kwarg passed to :class:`Variant`).

    Returns:
        A tuple of choice names (``Variant.choice_names``), in
        insertion order.

    Raises:
        VariantError: If no variant with the given name is found in the
            diagram. Anonymous Variants (built without ``name=``) are
            never matched and so cannot be queried via this helper.
    """
    tag = _find_first_tag_by_name(diagram, variant_name)
    if tag is None:
        available = sorted(
            {n for n, _, _ in list_variants(diagram) if n is not None}
        )
        raise VariantError(
            f"get_variant_choices: no Variant with name {variant_name!r} "
            f"found in diagram. Available variant names: {available!r}. "
            f"(Anonymous Variants -- those built without a ``name=`` "
            f"kwarg -- cannot be addressed by name.)"
        )
    return tag.variant.choice_names

`list_variants(diagram)`

List every variant point found in a (possibly nested) diagram.

Walks diagram recursively and returns a metadata triple for every subsystem that was produced by :func:select_variant. Each triple has the shape (name, choice_names, active_choice):

name is the variant's human-readable label (Variant.name). None for anonymous Variants.
choice_names is the stable tuple of available choice names (Variant.choice_names).
active_choice is the name of the choice currently bound at this point in the diagram.

Iteration order follows the diagram's tree-traversal order (parent before children, siblings in registration order). If the same Variant instance is reused at multiple points in the diagram, each occurrence yields its own entry — callers that want a deduped view should collapse on name.

Parameters:

Name	Type	Description	Default
`diagram`		A built `Diagram` (typically the output of `DiagramBuilder.build`) or any `SystemBase`. Passing a non-Diagram subsystem returns `[]` (no children to walk; a tagged-leaf-as-root case is not produced by the current API surface, but the helper degrades gracefully).	required

Returns:

Type	Description
`list[tuple]`	A list of `(name, choice_names, active_choice)` tuples. Empty
`list[tuple]`	if `diagram` contains no variant points (default-off path).

Source code in jaxonomy/framework/variants.py

def list_variants(diagram) -> list[tuple]:
    """List every variant point found in a (possibly nested) diagram.

    Walks ``diagram`` recursively and returns a metadata triple for
    every subsystem that was produced by :func:`select_variant`. Each
    triple has the shape ``(name, choice_names, active_choice)``:

    - ``name`` is the variant's human-readable label (``Variant.name``).
      ``None`` for anonymous Variants.
    - ``choice_names`` is the stable tuple of available choice names
      (``Variant.choice_names``).
    - ``active_choice`` is the name of the choice currently bound at
      this point in the diagram.

    Iteration order follows the diagram's tree-traversal order (parent
    before children, siblings in registration order). If the same
    ``Variant`` instance is reused at multiple points in the diagram,
    each occurrence yields its own entry — callers that want a deduped
    view should collapse on ``name``.

    Args:
        diagram: A built ``Diagram`` (typically the output of
            ``DiagramBuilder.build``) or any ``SystemBase``. Passing a
            non-Diagram subsystem returns ``[]`` (no children to walk;
            a tagged-leaf-as-root case is not produced by the current
            API surface, but the helper degrades gracefully).

    Returns:
        A list of ``(name, choice_names, active_choice)`` tuples. Empty
        if ``diagram`` contains no variant points (default-off path).
    """
    # Avoid an import cycle (diagram imports leaf_system, which is
    # imported here at module load time).
    from .diagram import Diagram

    if not isinstance(diagram, Diagram):
        return []

    out: list[tuple] = []
    for _parent, _idx, tag in _iter_tagged(diagram):
        out.append((tag.variant.name, tag.variant.choice_names, tag.active_choice))
    return out

`next_dependency_ticket()`

Create a new unique dependency ticket using the next available value.

Source code in jaxonomy/framework/dependency_graph.py

def next_dependency_ticket():
    """Create a new unique dependency ticket using the next available value."""
    return DependencyTicket.next_available_ticket()

`parameters(static=None, dynamic=None)`

Decorator to apply to a system class to declare static or dynamic parameters.

Source code in jaxonomy/framework/system_decorators.py

def parameters(static: list[str] = None, dynamic: list[str] = None):
    """Decorator to apply to a system class to declare
    static or dynamic parameters."""

    if static is None:
        static = []

    if dynamic is None:
        dynamic = []

    static_param_names = set(static)
    dynamic_param_names = set(dynamic)

    def decorator(entity: Union[Callable[P, T], type]) -> Callable[P, T]:
        if isinstance(entity, type):
            init_func = entity.__init__
            # Useful for class introspection like parsing custom leaf system in
            # the frontend to configure the UI block.
            entity.__parameters__ = static + dynamic
        elif callable(entity):
            init_func = entity

        @wraps(init_func)
        def wrapped_init(self, *args, **kwargs):
            resolved_args = [
                arg.get() if isinstance(arg, Parameter) else arg for arg in args
            ]
            resolved_kwargs = {
                k: kwarg.get() if isinstance(kwarg, Parameter) else kwarg
                for k, kwarg in kwargs.items()
            }

            init_func(self, *resolved_args, **resolved_kwargs)

            # TODO: Prevent parameters from being inherited from parent systems.
            # This is necessary to avoid unknown behaviors when a child parameter
            # is used to define a parent parameter, eg. what we used to do in
            # PID continuous block where gains were used to calculate A, B, C, D
            # matrices.
            # This will force the implementor of the block to implement jitted
            # callbacks in such a way that they only depend on the current system's
            # parameters.
            # We should also allow inheritance of params with a flag or annotation.
            # https://github.com/machinavitalis/jaxonomy/pull/6790
            # self._static_parameters = {}
            # self._dynamic_parameters = {}

            static_params = _get_params(static_param_names, init_func, args, kwargs)
            for param_name, value in static_params.items():
                self.declare_static_parameter(param_name, value)

            dyn_params = _get_params(dynamic_param_names, init_func, args, kwargs)
            for param_name, value in dyn_params.items():
                if value is not None:
                    self.declare_dynamic_parameter(param_name, value)

        if isinstance(entity, type):
            entity.__init__ = wrapped_init
            return entity
        elif callable(entity):
            return wrapped_init

    return decorator

`select_variant(variant, name=None)`

Resolve a Variant at build time and return the active sub-system.

Only the chosen builder is invoked; the others are never called. This matches the "active variant only" code-generation behavior familiar from established block-diagram tools -- nothing about the unselected branches enters the JIT trace, the parameter pytree, or the diagram's registered-systems list.

Parameters:

Name	Type	Description	Default
`variant`	`Variant`	The `Variant` to resolve.	required
`name`	`Optional[str]`	Name of the choice to activate. If `None`, `variant.default` is used.	`None`

Returns:

Type	Description
`SystemBase`	The `SystemBase` returned by the chosen builder.

Raises:

Type	Description
`VariantError`	If `name` is not one of `variant.choices`, or if the chosen builder returns something that isn't a `SystemBase`.

Source code in jaxonomy/framework/variants.py

def select_variant(
    variant: Variant,
    name: Optional[str] = None,
) -> SystemBase:
    """Resolve a ``Variant`` at build time and return the active sub-system.

    Only the chosen builder is invoked; the others are never called. This
    matches the "active variant only" code-generation behavior familiar
    from established block-diagram tools -- nothing about the unselected
    branches enters the JIT trace, the parameter pytree, or the diagram's
    registered-systems list.

    Args:
        variant: The ``Variant`` to resolve.
        name:
            Name of the choice to activate. If ``None``, ``variant.default``
            is used.

    Returns:
        The ``SystemBase`` returned by the chosen builder.

    Raises:
        VariantError: If ``name`` is not one of ``variant.choices``, or if
            the chosen builder returns something that isn't a ``SystemBase``.
    """
    chosen = variant.default if name is None else name
    if chosen not in variant.choices:
        raise VariantError(
            f"Variant {variant.name!r}: unknown choice {chosen!r}; "
            f"available: {list(variant.choices)!r}"
        )
    builder = variant.choices[chosen]
    result = builder()
    if not isinstance(result, SystemBase):
        raise VariantError(
            f"Variant {variant.name!r}: choice {chosen!r} builder returned "
            f"{type(result).__name__}, expected a SystemBase (Diagram or LeafSystem)."
        )
    # T-111-followup-with-config: tag the resolved subsystem with its
    # originating variant so post-build configurators (apply_variant_config /
    # Diagram.with_config) can locate and swap it later. Tagging is a
    # no-op for the legacy / phase-1 path: nothing in the simulator,
    # context factory, or pytree machinery reads VARIANT_METADATA_ATTR;
    # it's purely a hint for the configurator walker.
    setattr(
        result,
        VARIANT_METADATA_ATTR,
        _VariantMetadata(variant=variant, active_choice=chosen),
    )
    return result

`set_fx_rate(from_currency, to_currency, rate)`

Record an FX rate so that one unit of from_currency equals rate units of to_currency.

Both directions are written: setting USD→EUR at 0.92 simultaneously sets EUR→USD at 1.0 / 0.92 so round-trips are exact under the floating-point reciprocal. A zero or non-finite rate is rejected (FX rates must be positive finite numbers).

Parameters:

Name	Type	Description	Default
`from_currency`	`'Unit \| str'`	Source currency, either a :class:`Unit` (such as :data:`usd`) or a string code (`"USD"`).	required
`to_currency`	`'Unit \| str'`	Destination currency, ditto.	required
`rate`	`float`	Strictly positive multiplicative conversion factor.	required

Raises:

Type	Description
`ValueError`	if `rate` is non-positive or non-finite.
`UnitMismatchError`	if either argument is not a pure currency.

Example:

>>> set_fx_rate("USD", "EUR", 0.92)
>>> get_fx_rate(usd, eur)
0.92
>>> # Self-rate is always 1.0 and need not be set explicitly.
>>> get_fx_rate(usd, usd)
1.0

Source code in jaxonomy/framework/units.py

def set_fx_rate(
    from_currency: "Unit | str",
    to_currency: "Unit | str",
    rate: float,
) -> None:
    """Record an FX rate so that one unit of ``from_currency`` equals
    ``rate`` units of ``to_currency``.

    Both directions are written: setting USD→EUR at 0.92 simultaneously
    sets EUR→USD at ``1.0 / 0.92`` so round-trips are exact under the
    floating-point reciprocal. A zero or non-finite ``rate`` is rejected
    (FX rates must be positive finite numbers).

    Args:
        from_currency: Source currency, either a :class:`Unit` (such as
            :data:`usd`) or a string code (``"USD"``).
        to_currency: Destination currency, ditto.
        rate: Strictly positive multiplicative conversion factor.

    Raises:
        ValueError: if ``rate`` is non-positive or non-finite.
        UnitMismatchError: if either argument is not a pure currency.

    Example:

        >>> set_fx_rate("USD", "EUR", 0.92)
        >>> get_fx_rate(usd, eur)
        0.92
        >>> # Self-rate is always 1.0 and need not be set explicitly.
        >>> get_fx_rate(usd, usd)
        1.0
    """
    src = _canonical_currency_code(from_currency)
    dst = _canonical_currency_code(to_currency)
    rate_f = float(rate)
    if not (rate_f > 0.0):
        raise ValueError(
            f"set_fx_rate({from_currency!r}, {to_currency!r}, {rate!r}): "
            "FX rate must be a strictly positive finite number."
        )
    # math.isinf check without importing math: rate_f != rate_f for NaN,
    # rate_f - rate_f == 0 is False for ±inf.
    if rate_f != rate_f or (rate_f - rate_f) != 0.0:
        raise ValueError(
            f"set_fx_rate({from_currency!r}, {to_currency!r}, {rate!r}): "
            "FX rate must be a finite number (got NaN or inf)."
        )
    _FX_RATES[(src, dst)] = rate_f
    # This direction is now explicitly user-set, so it is no longer a
    # candidate for auto-refresh from its own reverse.
    _FX_AUTO_REVERSE.discard((src, dst))
    # Auto-populate (or refresh) the reverse direction unless the user
    # already set an explicit (potentially asymmetric) reverse — common in
    # real markets when bid/ask spreads matter. Re-setting an existing pair
    # (e.g. a daily-snapshot refresh) must update the derived reverse too,
    # otherwise it goes stale and round-trips stop being exact.
    reverse = (dst, src)
    if reverse not in _FX_RATES or reverse in _FX_AUTO_REVERSE:
        _FX_RATES[reverse] = 1.0 / rate_f
        _FX_AUTO_REVERSE.add(reverse)

`submodel_function(system, output_ports=None, input_ports=None, auto_seed=True)`

Wrap system's ports as a pure function of (context, *inputs).

Parameters:

Name	Type	Description	Default
`system`	`'SystemBase'`	The `LeafSystem` or `Diagram` to wrap.	required
`output_ports`	`'Sequence[OutputPort] \| None'`	Output ports whose values to return. Defaults to all of `system.output_ports`.	`None`
`input_ports`	`'Sequence[InputPort] \| None'`	Input ports that the closure will feed. Defaults to all of `system.input_ports`. Inputs not listed here are assumed already connected or fixed.	`None`
`auto_seed`	`bool`	If True (default), any input port in `input_ports` that is not already fixed or connected is pre-fixed to a zero placeholder so `create_context` succeeds on systems with dangling exported inputs. Set to False if you have seeded placeholders yourself.	`True`

Returns:

Type	Description
`Callable`	`f(context, *inputs) -> outputs`. When a single output port
`Callable`	is selected the return is a scalar / array; otherwise a tuple
`Callable`	in `output_ports` declaration order.

Example::

bld = jaxonomy.DiagramBuilder()
plant = bld.add(MyPlant())
bld.export_input(plant.input_ports[0], name="u")
bld.export_output(plant.output_ports[0], name="y")
diagram = bld.build()

f = jaxonomy.submodel_function(diagram)
ctx = diagram.create_context()    # auto-seeded placeholders
y = f(ctx, u)
dy_du = jax.grad(lambda u: f(ctx, u))(u0)
y_batch = jax.vmap(f, in_axes=(None, 0))(ctx, u_batch)

Performance envelope (T-008, follow-up finding 2026-05-16): Each call invokes the diagram's full evaluation machinery — port-fix context managers, dependency-tracked output evaluation, cache invalidation. That overhead is fine for one-shot rollouts, batched evaluation (where the cost amortises across the batch via jax.vmap), and gradient computation via jax.grad (the closure is traced once, then the compiled XLA program runs at native speed).

It is **not** fine for tight Python-side loops that call ``f``
thousands of times per simulated second — typical MPC inner
loops where the prediction model is re-evaluated at every
sample of a ``jax.lax.scan``-style rollout. There the per-call
Python overhead dominates and the wall-clock blows up by 100×
or more relative to closing over the underlying primitive
directly (e.g. ``interp_2d``, ``lookup_table_nd``, or a
hand-rolled JAX function). The canonical workaround in that
case is to skip ``submodel_function`` entirely for the inner
loop and call the primitive directly inside the scan body. See
``docs/examples/engine_map_fitting_to_mpc.ipynb`` for an
example of the hand-rolled-scan pattern.

Rule of thumb: if the closure will be invoked from a
Python-level loop more than ~100 times per simulation, profile
first. ``jax.jit(f)`` + ``jax.vmap`` over the entire batch
usually beats a Python loop by orders of magnitude.

Source code in jaxonomy/framework/submodel.py

def submodel_function(
    system: "SystemBase",
    output_ports: "Sequence[OutputPort] | None" = None,
    input_ports: "Sequence[InputPort] | None" = None,
    auto_seed: bool = True,
) -> Callable:
    """Wrap ``system``'s ports as a pure function of (context, *inputs).

    Args:
        system: The ``LeafSystem`` or ``Diagram`` to wrap.
        output_ports: Output ports whose values to return.  Defaults to
            all of ``system.output_ports``.
        input_ports: Input ports that the closure will feed.  Defaults to
            all of ``system.input_ports``.  Inputs not listed here are
            assumed already connected or fixed.
        auto_seed: If True (default), any input port in ``input_ports``
            that is not already fixed or connected is pre-fixed to a
            zero placeholder so ``create_context`` succeeds on systems
            with dangling exported inputs.  Set to False if you have
            seeded placeholders yourself.

    Returns:
        ``f(context, *inputs) -> outputs``.  When a single output port
        is selected the return is a scalar / array; otherwise a tuple
        in ``output_ports`` declaration order.

    Example::

        bld = jaxonomy.DiagramBuilder()
        plant = bld.add(MyPlant())
        bld.export_input(plant.input_ports[0], name="u")
        bld.export_output(plant.output_ports[0], name="y")
        diagram = bld.build()

        f = jaxonomy.submodel_function(diagram)
        ctx = diagram.create_context()    # auto-seeded placeholders
        y = f(ctx, u)
        dy_du = jax.grad(lambda u: f(ctx, u))(u0)
        y_batch = jax.vmap(f, in_axes=(None, 0))(ctx, u_batch)

    Performance envelope (T-008, follow-up finding 2026-05-16):
        Each call invokes the diagram's full evaluation machinery —
        port-fix context managers, dependency-tracked output evaluation,
        cache invalidation. That overhead is fine for **one-shot
        rollouts**, **batched evaluation** (where the cost amortises
        across the batch via ``jax.vmap``), and **gradient computation
        via ``jax.grad``** (the closure is traced once, then the
        compiled XLA program runs at native speed).

        It is **not** fine for tight Python-side loops that call ``f``
        thousands of times per simulated second — typical MPC inner
        loops where the prediction model is re-evaluated at every
        sample of a ``jax.lax.scan``-style rollout. There the per-call
        Python overhead dominates and the wall-clock blows up by 100×
        or more relative to closing over the underlying primitive
        directly (e.g. ``interp_2d``, ``lookup_table_nd``, or a
        hand-rolled JAX function). The canonical workaround in that
        case is to skip ``submodel_function`` entirely for the inner
        loop and call the primitive directly inside the scan body. See
        ``docs/examples/engine_map_fitting_to_mpc.ipynb`` for an
        example of the hand-rolled-scan pattern.

        Rule of thumb: if the closure will be invoked from a
        Python-level loop more than ~100 times per simulation, profile
        first. ``jax.jit(f)`` + ``jax.vmap`` over the entire batch
        usually beats a Python loop by orders of magnitude.
    """
    out_ports = tuple(output_ports) if output_ports is not None else tuple(system.output_ports)
    in_ports = tuple(input_ports) if input_ports is not None else tuple(system.input_ports)

    if not out_ports:
        raise ValueError(
            f"submodel_function({system.name!r}): the system has no output ports "
            "to evaluate.  Supply output_ports= explicitly if you want to evaluate "
            "intermediate ports."
        )

    if auto_seed:
        for p in in_ports:
            _seed_placeholder(p)

    def _call(context: "ContextBase", *inputs):
        if len(inputs) != len(in_ports):
            raise TypeError(
                f"submodel_function({system.name!r}) expected {len(in_ports)} "
                f"input values (one per input port), got {len(inputs)}."
            )
        with ExitStack() as stack:
            for port, value in zip(in_ports, inputs):
                stack.enter_context(port.fixed(value))
            ys = tuple(p.eval(context) for p in out_ports)
        return ys[0] if len(ys) == 1 else ys

    _call.__name__ = f"submodel_{system.name}"
    _call.__doc__ = (
        f"Evaluate {system.name!r} as a pure function of inputs "
        f"{[p.name for p in in_ports]} → outputs "
        f"{[p.name for p in out_ports]}.\n\n"
        "Signature: f(context, *inputs) -> outputs.  See "
        "jaxonomy.submodel_function for details."
    )
    return _call

`variant_subsystem(choices, name=None, default=None)`

Build a resolver closure for a one-shot variant point.

Convenience wrapper around Variant + select_variant for the common case::

controller = variant_subsystem(
    choices={
        "pid": lambda: build_pid(),
        "lqr": lambda: build_lqr(),
    },
    default="pid",
)

# Later, at "configure" time:
active = controller(name="lqr")   # returns the lqr Diagram
active = controller()             # returns the pid Diagram (default)

Parameters:

Name	Type	Description	Default
`choices`	`Mapping[str, Callable[[], SystemBase]]`	See `Variant.choices`.	required
`name`	`Optional[str]`	See `Variant.name`.	`None`
`default`	`Optional[str]`	Choice name to use when the returned closure is called without an argument. If `None`, the first key of `choices` is used (insertion order, which is guaranteed in Python 3.7+).	`None`

Returns:

Type	Description
`Callable[..., SystemBase]`	A closure `select(name=None) -> SystemBase` that, on each call,
`Callable[..., SystemBase]`	resolves to a freshly-built sub-system for the named choice.

Source code in jaxonomy/framework/variants.py

def variant_subsystem(
    choices: Mapping[str, Callable[[], SystemBase]],
    name: Optional[str] = None,
    default: Optional[str] = None,
) -> Callable[..., SystemBase]:
    """Build a resolver closure for a one-shot variant point.

    Convenience wrapper around ``Variant`` + ``select_variant`` for the
    common case::

        controller = variant_subsystem(
            choices={
                "pid": lambda: build_pid(),
                "lqr": lambda: build_lqr(),
            },
            default="pid",
        )

        # Later, at "configure" time:
        active = controller(name="lqr")   # returns the lqr Diagram
        active = controller()             # returns the pid Diagram (default)

    Args:
        choices: See ``Variant.choices``.
        name: See ``Variant.name``.
        default:
            Choice name to use when the returned closure is called without
            an argument. If ``None``, the *first* key of ``choices`` is
            used (insertion order, which is guaranteed in Python 3.7+).

    Returns:
        A closure ``select(name=None) -> SystemBase`` that, on each call,
        resolves to a freshly-built sub-system for the named choice.
    """
    if not choices:
        raise VariantError(
            f"variant_subsystem {name!r}: choices must contain at least one entry"
        )
    chosen_default = default if default is not None else next(iter(choices))
    variant = Variant(choices=dict(choices), default=chosen_default, name=name)

    def _select(name: Optional[str] = None) -> SystemBase:
        return select_variant(variant, name=name)

    # Surface the underlying Variant for introspection / tests.
    _select.variant = variant  # type: ignore[attr-defined]
    return _select

Framework

jaxonomy.framework

BlockInitializationError

BlockParameterError

BlockRuntimeError

BusUnit dataclass

field_unit(name)

ContextBase dataclass

__getitem__(key) abstractmethod

with_continuous_state(value) abstractmethod

with_discrete_state(value) abstractmethod

with_mode(value) abstractmethod

with_new_state() abstractmethod

with_parameter(name, value)

with_parameters(new_parameters) abstractmethod

with_state(state) abstractmethod

with_subcontext(key, ctx) abstractmethod

with_time(value)

with_updated_parameters() abstractmethod

DependencyTicket

Diagram dataclass

continuous_substep_vector property

has_dirty_static_parameters property

check_no_algebraic_loops()

check_types(context, error_collector=None)

declare_dynamic_parameter(name, parameter)

eval_subsystem_input_port(context, port_locator)

eval_subsystem_output_port(context, port_locator)

export_input(locator, port_name)

export_output(locator, port_name)

get_parameter(path)

initialize_static_data(context)

list_parameters(prefix='')

post_simulation_finalize()

print_schedule(*, format='text', file=None, ensure_initialized=True)

with_parameters(updates)

DiagramBuilder

__init__(*, validate_rates_at_connect=None, unit_conversion='auto', auto_insert_rate_transitions=False)

add(*systems)

build(name='root', ui_id=None, parameters=None)

connect(src, dest)

export_input(port, name=None)

export_output(port, name=None)

DiagramContext dataclass

with_parameters(new_parameters)

DiscreteUpdateEvent dataclass

DtypeMismatchError

EnabledMode

EnabledStateMode

EnabledSubsystem

ErrorCollector

add_error(error)

context(parent=None) classmethod

EventCollection

ForLoop

IntegerTime

as_decimal(time) classmethod

from_decimal(time) classmethod

JaxonomyError

__init__(message=None, *, system=None, system_id=None, name_path=None, ui_id_path=None, port_index=None, port_name=None, port_direction=None, parameter_name=None, loop=None)

caused_by(exc_type)

LeafContext dataclass

__getitem__(key)

with_parameters(new_parameters)

with_subcontext(key, ctx)

LeafState dataclass

with_cached_value(index, value)

with_continuous_state(value)

with_discrete_state(value)

with_mode(value)

LeafSystem dataclass

continuous_state_default property

continuous_substep_vector property

has_multirate_substeps property

configure_continuous_state(callback_idx, shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None)

configure_output_port(port_index, callback, period=None, offset=0.0, prerequisites_of_calc=None, default_value=None, requires_inputs=None)

configure_periodic_update(event_index, callback, period, offset, enable_tracing=None)

declare_cache(callback, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=True)

declare_continuous_state(shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None, substeps=1, project=None)

declare_continuous_state_output(name=None)

`jaxonomy.framework`

`BlockInitializationError`

`BlockParameterError`

`BlockRuntimeError`

`BusUnit` `dataclass`

`field_unit(name)`

`ContextBase` `dataclass`

`getitem(key)` `abstractmethod`

`with_continuous_state(value)` `abstractmethod`

`with_discrete_state(value)` `abstractmethod`

`with_mode(value)` `abstractmethod`

`with_new_state()` `abstractmethod`

`with_parameter(name, value)`

`with_parameters(new_parameters)` `abstractmethod`

`with_state(state)` `abstractmethod`

`with_subcontext(key, ctx)` `abstractmethod`

`with_time(value)`

`with_updated_parameters()` `abstractmethod`

`DependencyTicket`

`Diagram` `dataclass`

`continuous_substep_vector` `property`

`has_dirty_static_parameters` `property`

`check_no_algebraic_loops()`

`check_types(context, error_collector=None)`

`declare_dynamic_parameter(name, parameter)`

`eval_subsystem_input_port(context, port_locator)`

`eval_subsystem_output_port(context, port_locator)`

`export_input(locator, port_name)`

`export_output(locator, port_name)`

`get_parameter(path)`

`initialize_static_data(context)`

`list_parameters(prefix='')`

`post_simulation_finalize()`

`print_schedule(*, format='text', file=None, ensure_initialized=True)`

`with_parameters(updates)`

`DiagramBuilder`

`init(*, validate_rates_at_connect=None, unit_conversion='auto', auto_insert_rate_transitions=False)`

`add(*systems)`

`build(name='root', ui_id=None, parameters=None)`

`connect(src, dest)`

`export_input(port, name=None)`

`export_output(port, name=None)`

`DiagramContext` `dataclass`

`with_parameters(new_parameters)`

`DiscreteUpdateEvent` `dataclass`

`DtypeMismatchError`

`EnabledMode`

`EnabledStateMode`

`EnabledSubsystem`

`ErrorCollector`

`add_error(error)`

`context(parent=None)` `classmethod`

`EventCollection`

`ForLoop`

`IntegerTime`

`as_decimal(time)` `classmethod`

`from_decimal(time)` `classmethod`

`JaxonomyError`

`init(message=None, *, system=None, system_id=None, name_path=None, ui_id_path=None, port_index=None, port_name=None, port_direction=None, parameter_name=None, loop=None)`

`caused_by(exc_type)`

`LeafContext` `dataclass`

`getitem(key)`

`with_parameters(new_parameters)`

`with_subcontext(key, ctx)`

`LeafState` `dataclass`

`with_cached_value(index, value)`

`with_continuous_state(value)`

`with_discrete_state(value)`

`with_mode(value)`

`LeafSystem` `dataclass`

`continuous_state_default` `property`

`continuous_substep_vector` `property`

`has_multirate_substeps` `property`

`configure_continuous_state(callback_idx, shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None)`

`configure_output_port(port_index, callback, period=None, offset=0.0, prerequisites_of_calc=None, default_value=None, requires_inputs=None)`

`configure_periodic_update(event_index, callback, period, offset, enable_tracing=None)`

`declare_cache(callback, period=None, offset=0.0, name=None, prerequisites_of_calc=None, default_value=None, requires_inputs=True)`

`declare_continuous_state(shape=None, default_value=None, dtype=None, ode=None, mass_matrix=None, as_array=True, requires_inputs=True, prerequisites_of_calc=None, substeps=1, project=None)`

`declare_continuous_state_output(name=None)`

`declare_discrete_state(shape=None, default_value=None, dtype=None, as_array=True, name=None)`