Reactions

reactions

Compile op_system reaction artifacts into a flat stochastic network.

op_system publishes one typed :class:ReactionArtifact per named transition (CompiledRhs.reactions). :func:compile_reaction_network expands each artifact into one channel per source cell and builds the stoichiometry, reactant orders, and propensity callback that :class:~op_engine.DirectSSASolver, :class:~op_engine.TauLeapingSolver, and :class:~op_engine.AdaptiveTauLeapingSolver consume. :func:from_compiled_rhs does the same from a compiled op_system RHS.

The module reads artifacts structurally and imports neither op_system nor flepimop2.

CompiledReactionNetwork(stoichiometry, reactant_stoichiometry, channel_reactions, channel_names, _blocks, _reactions, _event_shapes, _params, reactants_complete=False, incomplete_reactions=(), dependency_incidence=None, propensity_orders=None) dataclass

Flat stochastic network compiled from typed reaction artifacts.

channel_reactions retains the parent reaction name for every expanded source-cell channel. channel_names adds integer coordinates and is intended for diagnostics only.

n_channels property

Return the number of expanded reaction channels.

n_state property

Return the number of flattened state cells.

mean_drift(time, state)

Return deterministic mean drift from the compiled channels.

Returns:

Type	Description
Array	stoichiometry @ propensity in the state's namespace.

Source code in src/op_engine/reactions.py

def mean_drift(self, time: float, state: Array) -> Array:
    """Return deterministic mean drift from the compiled channels.

    Returns:
        ``stoichiometry @ propensity`` in the state's namespace.
    """
    xp = _namespace_of(state)
    propensity = self.propensity(time, state)
    stoichiometry = cast(
        "Array",
        xp.asarray(self.stoichiometry, dtype=state.dtype),
    )
    return cast("Array", xp.matmul(stoichiometry, propensity))

propensity(time, state)

Evaluate every reaction and concatenate its source-cell rates.

The stochastic core stores the state as (n_state, 1). A flat state is accepted as well so the same callback can construct hybrid mean drift. Returned shape mirrors the input convention.

Returns:

Type	Description
Array	Propensities in the evolving state's array namespace.

Source code in src/op_engine/reactions.py

def propensity(self, time: float, state: Array) -> Array:
    """Evaluate every reaction and concatenate its source-cell rates.

    The stochastic core stores the state as ``(n_state, 1)``. A flat
    state is accepted as well so the same callback can construct hybrid
    mean drift. Returned shape mirrors the input convention.

    Returns:
        Propensities in the evolving state's array namespace.
    """
    xp = _namespace_of(state)
    batched = state.shape == (self.n_state, 1)
    if not batched and state.shape != (self.n_state,):
        msg = (
            f"Reaction propensity received state shape {state.shape}; expected "
            f"{(self.n_state,)} or {(self.n_state, 1)}."
        )
        raise ValueError(msg)
    flat_state = cast("Array", xp.reshape(state, (self.n_state,)))
    indexable_state = cast("_IndexableArray", flat_state)
    state_dict = {
        block.base: cast(
            "Array",
            xp.reshape(
                indexable_state[block.offset : block.offset + block.size],
                block.shape,
            ),
        )
        for block in self._blocks
    }

    values: list[Array] = []
    for reaction, expected_shape in zip(
        self._reactions,
        self._event_shapes,
        strict=True,
    ):
        raw = reaction.propensity_fn(time, state_dict, **self._params)
        raw_namespace = getattr(raw, "__array_namespace__", None)
        if raw_namespace is not None and raw_namespace() is not xp:
            msg = f"Reaction {reaction.name!r} changed the state array namespace."
            raise TypeError(msg)
        value = cast("Array", xp.asarray(raw, dtype=state.dtype))
        if value.shape != expected_shape:
            msg = (
                f"Reaction {reaction.name!r} returned propensity shape "
                f"{value.shape}; expected {expected_shape}."
            )
            raise ValueError(msg)
        values.append(cast("Array", xp.reshape(value, (-1,))))

    combined = cast("Array", xp.concat(tuple(values), axis=0))
    if batched:
        return cast("Array", xp.reshape(combined, (self.n_channels, 1)))
    return combined

CompiledRhsLike

Bases: Protocol

Structural surface of op_system.CompiledRhs read here.

ReactionArtifact

Bases: Protocol

Structural surface published by op_system.CompiledReaction.

ReactionReactantArtifact

Bases: Protocol

Structural molecular-reactant surface published by op_system.

compile_reaction_network(reactions, *, template_shapes, axis_sizes, params, n_state=None, reaction_names=None)

Compile op_system reaction artifacts into flat channels.

Each source-cell propensity becomes one channel. Consequently, a collapsed axis is represented by distinct columns that share one destination row; summing simultaneous firings is then exactly the stoichiometric matrix multiplication performed by the core stochastic solvers. An offset axis moves each channel to its shifted coordinate; when that leaves the axis, the column only removes the donor.

Parameters:

Name	Type	Description	Default
reactions	Sequence[ReactionArtifact]	Typed reaction artifacts, such as CompiledRhs.reactions.	required
template_shapes	Mapping[str, Sequence[int]]	Shape of each state template, in the flat state's order, such as CompiledRhs.template_shapes.	required
axis_sizes	Mapping[str, int]	Number of coordinates on each axis.	required
params	Mapping[str, object]	Parameter values passed to every propensity.	required
n_state	int | None	Number of cells in the flat state. Defaults to the total that template_shapes describes; when given, the two must agree.	None
reaction_names	Sequence[str] | None	Optional parent reaction names to select for a hybrid jump partition. None selects every artifact.	None

Returns:

Type	Description
CompiledReactionNetwork	Validated flat reaction network.

Raises:

Type	Description
TypeError	If an argument or artifact field has the wrong type.
ValueError	If the artifacts are inconsistent with each other or with the state layout.

Source code in src/op_engine/reactions.py

def compile_reaction_network(  # noqa: PLR0913
    reactions: Sequence[ReactionArtifact],
    *,
    template_shapes: Mapping[str, Sequence[int]],
    axis_sizes: Mapping[str, int],
    params: Mapping[str, object],
    n_state: int | None = None,
    reaction_names: Sequence[str] | None = None,
) -> CompiledReactionNetwork:
    """Compile op_system reaction artifacts into flat channels.

    Each source-cell propensity becomes one channel. Consequently, a collapsed
    axis is represented by distinct columns that share one destination row;
    summing simultaneous firings is then exactly the stoichiometric matrix
    multiplication performed by the core stochastic solvers. An offset axis
    moves each channel to its shifted coordinate; when that leaves the axis,
    the column only removes the donor.

    Args:
        reactions: Typed reaction artifacts, such as ``CompiledRhs.reactions``.
        template_shapes: Shape of each state template, in the flat state's
            order, such as ``CompiledRhs.template_shapes``.
        axis_sizes: Number of coordinates on each axis.
        params: Parameter values passed to every propensity.
        n_state: Number of cells in the flat state. Defaults to the total
            that ``template_shapes`` describes; when given, the two must
            agree.
        reaction_names: Optional parent reaction names to select for a hybrid
            jump partition. ``None`` selects every artifact.

    Returns:
        Validated flat reaction network.

    Raises:
        TypeError: If an argument or artifact field has the wrong type.
        ValueError: If the artifacts are inconsistent with each other or
            with the state layout.
    """
    if not isinstance(reactions, tuple | list):
        msg = "'reactions' must be a sequence of reaction artifacts."
        raise TypeError(msg)
    selected = tuple(cast("ReactionArtifact", item) for item in reactions)
    if reaction_names is not None:
        requested = tuple(reaction_names)
        if len(requested) != len(set(requested)):
            msg = "stochastic_reactions must not contain duplicates."
            raise ValueError(msg)
        available = {reaction.name for reaction in selected}
        missing = sorted(set(requested) - available)
        if missing:
            msg = f"Unknown stochastic reaction names: {', '.join(missing)}."
            raise ValueError(msg)
        wanted = set(requested)
        selected = tuple(reaction for reaction in selected if reaction.name in wanted)
    if not selected:
        msg = "No typed reactions are available for stochastic execution."
        raise ValueError(msg)
    reactions = selected

    if not isinstance(template_shapes, Mapping):
        msg = "'template_shapes' must map state template names to shapes."
        raise TypeError(msg)
    if n_state is None:
        n_state = sum(
            int(np.prod(tuple(shape), dtype=np.int64)) if shape else 1
            for shape in template_shapes.values()
            if isinstance(shape, tuple | list)
        )
    blocks_tuple = _state_blocks(
        cast("Mapping[object, object]", template_shapes), n_state=n_state
    )
    blocks = {block.base: block for block in blocks_tuple}

    if not isinstance(axis_sizes, Mapping):
        msg = "'axis_sizes' must be a mapping."
        raise TypeError(msg)
    raw_axis_sizes = axis_sizes
    axis_sizes = {}
    for axis, size in raw_axis_sizes.items():
        if (
            not isinstance(axis, str)
            or not isinstance(size, int)
            or isinstance(size, bool)
            or size < 1
        ):
            msg = "'axis_sizes' contains an invalid entry."
            raise TypeError(msg)
        axis_sizes[axis] = size
    base_axes = _reaction_axes(reactions, blocks, axis_sizes)

    columns: list[np.ndarray] = []
    reactant_columns: list[np.ndarray] = []
    dependency_columns: list[np.ndarray] = []
    channel_orders: list[int] = []
    channel_reactions: list[str] = []
    channel_names: list[str] = []
    event_shapes: list[tuple[int, ...]] = []
    reactants_complete = True
    incomplete: list[str] = []

    for reaction in reactions:
        complete = getattr(reaction, "reactants_complete", False)
        if not isinstance(complete, bool):
            msg = f"Reaction {reaction.name!r} reactants_complete must be boolean."
            raise TypeError(msg)
        if complete and getattr(reaction, "reactants", None) is None:
            msg = (
                f"Reaction {reaction.name!r} claims complete reactants without "
                "publishing reactant metadata."
            )
            raise ValueError(msg)
        # Complete reactants take precedence; otherwise dependencies and a
        # propensity order can describe the reaction instead.
        order = 0 if complete else _dependency_order(reaction)
        covered = complete or order > 0
        reactants_complete = reactants_complete and covered
        if not covered:
            incomplete.append(reaction.name)
        from_axes = _require_string_tuple(reaction.from_axes, field="from_axes")
        target_axes = base_axes[reaction.to_base]
        pinned, from_pinned, offsets, routed = _reaction_axis_metadata(
            reaction,
            from_axes=from_axes,
            source_axes=(
                target_axes
                if reaction.from_base is None
                else base_axes[reaction.from_base]
            ),
            target_axes=target_axes,
        )

        # A routed target coordinate is a trailing channel dimension.
        event_shape = tuple(axis_sizes[axis] for axis in from_axes + routed)
        event_shapes.append(event_shape)
        coordinates = np.ndindex(event_shape) if event_shape else iter(((),))
        for coordinate in coordinates:
            varying = dict(zip(from_axes, coordinate[: len(from_axes)], strict=True))
            routed_to = dict(zip(routed, coordinate[len(from_axes) :], strict=True))
            column = np.zeros(n_state, dtype=np.int64)
            if reaction.from_base is not None:
                source_coordinates = {**varying, **from_pinned}
                source = _flat_cell(
                    blocks[reaction.from_base],
                    base_axes[reaction.from_base],
                    source_coordinates,
                )
                column[source] -= 1

            reactants = _expanded_reactants(
                reaction,
                varying=varying,
                from_axes=from_axes,
                base_axes=base_axes,
                blocks=blocks,
                n_state=n_state,
            )

            read = np.zeros(n_state, dtype=np.int64)
            if order > 0:
                read = _expanded_entries(
                    reaction,
                    getattr(reaction, "dependencies", ()),
                    label="dependency",
                    varying=varying,
                    from_axes=from_axes,
                    base_axes=base_axes,
                    blocks=blocks,
                    n_state=n_state,
                )
                # A consumed species is read too, whether or not listed.
                read = ((read > 0) | (reactants > 0)).astype(np.int64)

            destination = _destination_cell(
                blocks[reaction.to_base],
                base_axes[reaction.to_base],
                {**varying, **pinned, **routed_to},
                offsets=offsets,
            )
            if destination is not None:
                column[destination] += 1
            columns.append(column)
            reactant_columns.append(reactants)
            dependency_columns.append(read)
            channel_orders.append(order)
            channel_reactions.append(reaction.name)
            coordinate_text = ",".join(
                [f"{axis}={index}" for axis, index in varying.items()]
                + [f"to:{axis}={index}" for axis, index in routed_to.items()]
            )
            channel_names.append(
                reaction.name
                if not coordinate_text
                else f"{reaction.name}[{coordinate_text}]"
            )

    return CompiledReactionNetwork(
        stoichiometry=np.stack(columns, axis=1),
        reactant_stoichiometry=np.stack(reactant_columns, axis=1),
        channel_reactions=tuple(channel_reactions),
        channel_names=tuple(channel_names),
        _blocks=blocks_tuple,
        _reactions=reactions,
        _event_shapes=tuple(event_shapes),
        _params=dict(params),
        reactants_complete=reactants_complete,
        incomplete_reactions=tuple(incomplete),
        dependency_incidence=(
            np.stack(dependency_columns, axis=1) if any(channel_orders) else None
        ),
        propensity_orders=(
            np.asarray(channel_orders, dtype=np.int64) if any(channel_orders) else None
        ),
    )

compiled_rhs_axis_sizes(compiled)

Return each axis's coordinate count from a compiled RHS's metadata.

Returns:

Type	Description
dict[str, int]	Mapping from axis name to size, from compiled.meta["axes"].

Raises:

Type	Description
TypeError	If an axis entry has no name.

Source code in src/op_engine/reactions.py

def compiled_rhs_axis_sizes(compiled: CompiledRhsLike) -> dict[str, int]:
    """Return each axis's coordinate count from a compiled RHS's metadata.

    Returns:
        Mapping from axis name to size, from ``compiled.meta["axes"]``.

    Raises:
        TypeError: If an axis entry has no name.
    """
    sizes: dict[str, int] = {}
    for axis in compiled.meta.get("axes", ()) or ():
        name = axis.get("name") if isinstance(axis, Mapping) else None
        if not isinstance(name, str):
            msg = "compiled.meta['axes'] entries must name their axis."
            raise TypeError(msg)
        size = axis.get("size") or len(axis.get("coords") or ())
        sizes[name] = int(size)
    return sizes

from_compiled_rhs(compiled, params, *, reaction_names=None, allow_gaps=False)

Compile the reaction network of a compiled op_system RHS.

The flat state follows compiled.template_shapes: each template's cells in C order, templates in declaration order. That is also the layout of compiled.eval_fn's state vector when the RHS is vectorized, so network.mean_drift(t, y) should equal compiled.eval_fn(t, y, **params) for every y. Check that identity once for a new model: a mismatch means some dynamics have no reaction.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhsLike	A compiled RHS, such as op_system.compile_spec(spec).	required
params	Mapping[str, object]	Parameter values passed to every propensity.	required
reaction_names	Sequence[str] | None	Optional reaction names to compile, for a hybrid jump partition. None compiles every reaction.	None
allow_gaps	bool	Compile even when compiled.reaction_gaps lists dynamics without a reaction. Selecting reaction_names also skips the check, since a partition is partial by design.	False

Returns:

Type	Description
CompiledReactionNetwork	Validated flat reaction network.

Raises:

Type	Description
TypeError	If the RHS has no vectorized state layout.
ValueError	If the RHS has reaction gaps and neither allow_gaps nor reaction_names is given.

Source code in src/op_engine/reactions.py

def from_compiled_rhs(
    compiled: CompiledRhsLike,
    params: Mapping[str, object],
    *,
    reaction_names: Sequence[str] | None = None,
    allow_gaps: bool = False,
) -> CompiledReactionNetwork:
    """Compile the reaction network of a compiled ``op_system`` RHS.

    The flat state follows ``compiled.template_shapes``: each template's
    cells in C order, templates in declaration order. That is also the
    layout of ``compiled.eval_fn``'s state vector when the RHS is
    vectorized, so ``network.mean_drift(t, y)`` should equal
    ``compiled.eval_fn(t, y, **params)`` for every ``y``. Check that
    identity once for a new model: a mismatch means some dynamics have no
    reaction.

    Args:
        compiled: A compiled RHS, such as ``op_system.compile_spec(spec)``.
        params: Parameter values passed to every propensity.
        reaction_names: Optional reaction names to compile, for a hybrid
            jump partition. ``None`` compiles every reaction.
        allow_gaps: Compile even when ``compiled.reaction_gaps`` lists
            dynamics without a reaction. Selecting ``reaction_names`` also
            skips the check, since a partition is partial by design.

    Returns:
        Validated flat reaction network.

    Raises:
        TypeError: If the RHS has no vectorized state layout.
        ValueError: If the RHS has reaction gaps and neither ``allow_gaps``
            nor ``reaction_names`` is given.
    """
    template_shapes = compiled.template_shapes
    if template_shapes is None:
        msg = (
            "The compiled RHS has no vectorized state layout "
            "(template_shapes is None), so its reactions cannot be indexed."
        )
        raise TypeError(msg)
    gaps = tuple(compiled.reaction_gaps)
    if gaps and not allow_gaps and reaction_names is None:
        described = ", ".join(
            f"{getattr(gap, 'name', None) or getattr(gap, 'origin', '?')} "
            f"({getattr(gap, 'reason', 'unknown')})"
            for gap in gaps
        )
        msg = (
            "The compiled RHS has dynamics without a reaction artifact: "
            f"{described}. A network of its reactions alone would drop them; "
            "pass allow_gaps=True to compile it anyway."
        )
        raise ValueError(msg)
    return compile_reaction_network(
        tuple(compiled.reactions),
        template_shapes=template_shapes,
        axis_sizes=compiled_rhs_axis_sizes(compiled),
        params=params,
        reaction_names=reaction_names,
    )