Compile

compile

op_system.compile.

Compile normalized RHS specifications into efficient, runnable callables.

This module is domain-agnostic and intentionally does not import flepimop2.

Contract

• Accepts NormalizedRhs (from op_system.specs) and produces a CompiledRhs object containing an eval_fn suitable for numerical backends.
• Raises built-in exceptions with standardized messages.

Security note

Expressions are parsed and validated with a conservative AST whitelist. RHS and latent-observable expressions use cached functions with explicit runtime bindings. Their constructor evaluates a generated lambda once; separate reaction callbacks retain their existing restricted code-object evaluation.

BodyEvalFn

Bases: Protocol

Callable that evaluates history signal bodies at (t, y, **params).

Returns a mapping from signal_id to the evaluated body array/value.

CompiledReactant(state_base, state_axes, full_axes, pinned, order) dataclass

Molecular reactant metadata aligned to an expanded reaction channel.

The record is numerical-array neutral. Providers combine state_axes with the enclosing reaction's channel coordinates and pinned indices to locate a state cell, then place order in their reactant-order matrix. Catalysts appear here even when their net stoichiometric change is zero.

CompiledReaction(name, from_base, from_axes, full_axes, to_base, to_axes, sum_axes, pinned, from_pinned, propensity_fn, reactants=(), reactants_complete=False, offsets=(), to_full_axes=None, routed_axes=(), dependencies=(), propensity_order=None, dependencies_complete=False, resolved_propensity_fn=None, time_varying_params=(), independent_axes=(), uses_time=False) dataclass

Compiled propensity + axis bookkeeping for one named transition.

Produced from :class:op_system._reactions.ReactionArtifactIR at compile time. A firing event at from-cell index i (in the natural N-D shape given by from_axes) depletes 1 from from_base at that cell and adds 1 to to_base at the cell obtained by: keeping every axis in to_axes at i's value for that axis, and using the fixed coordinate index from pinned for every axis in pinned (this includes both axes that are wildcard on from_axes and pinned on to -- a "collapse to a fixed target" transition, where multiple source cells share one destination and the axis also appears in sum_axes -- and axes pinned on BOTH from and to, a "point-to-point" shift between two specific coordinates on the same axis, e.g. a vaccination-dose-progression transition, which do NOT appear in sum_axes since there's only ever one source value on that axis). Consumers that need a single combined scatter/reduce target (e.g. an engine applying many simultaneous firings) sum over sum_axes when depositing into to_base -- op_system does not do that summation itself, since whether/how to combine simultaneous firings across a summed axis is an execution-semantics decision for the consumer, not a compile-time one.

from_base is None for a SOURCE-ONLY (from: null) reaction -- an exogenous hazard with no compartment to deplete (e.g. cross-district case importation). A consumer must special-case this: skip every from-side depletion/clamp step and apply only the to_base scatter, using from_axes/full_axes (which are the to-side template's own axes in this case, not a nonexistent from-side one -- see ReactionArtifactIR).

An axis-wide coord_shift publishes one reaction per state template with a non-empty offsets: the destination index on each offset axis is the source index plus its step (see offsets). The destination cell is then built from to_axes (copied), pinned (fixed), and offsets (shifted), which together cover full_axes.

Examples:

>>> import numpy as np
>>> from op_system import compile_spec
>>> compiled = compile_spec({
...     "kind": "transitions",
...     "state": ["S", "I"],
...     "transitions": [
...         {"name": "infect", "from": "S", "to": "I", "rate": "k"},
...     ],
... })
>>> reaction = compiled.reactions[0]
>>> float(
...     reaction.propensity_fn(
...         0.0, {"S": np.array(10.0), "I": np.array(0.0)}, k=0.2
...     )
... )
2.0
>>> reaction.time_varying_params, reaction.uses_time
((), False)

CompiledRhs(state_names, param_names, eval_fn, meta=(lambda: MappingProxyType({}))(), operators=tuple(), factorize_axes=tuple(), block_axes=tuple(), pytree_eval_fn=None, template_shapes=None, block_pytree_eval_fn=None, block_template_shapes=None, history_requirements=tuple(), history_eval_fn=None, body_eval_fn=None, block_history_eval_fn=None, block_body_eval_fn=None, reactions=tuple(), reaction_gaps=tuple(), _rhs=None, param_meta=(lambda: MappingProxyType({}))(), state_meta=(lambda: MappingProxyType({}))(), axis_meta=(lambda: MappingProxyType({}))(), observables=(lambda: MappingProxyType({}))(), gauges=None, gauge_fix=None) dataclass

Container for a compiled RHS evaluation function.

Instances produced by :func:compile_rhs retain a private reference to their source :class:NormalizedRhs so the container can be pickled and re-hydrated by re-running the compile pipeline on load. eval_fn itself is a closure (and on the vectorized path captures compiled code objects), so it is dropped from the pickle and rebuilt by :func:compile_rhs in :meth:__setstate__. Round-tripping a CompiledRhs therefore costs one compile on load and yields a functionally equivalent instance whose eval_fn produces identical outputs for identical inputs.

__getstate__()

Return picklable state.

The compiled eval_fn is a closure (and on the vectorized path captures compiled :class:types.CodeType objects), which is not portably picklable. Instead we serialize just the source :class:NormalizedRhs and let :meth:__setstate__ recompile.

Raises:

Type	Description
TypeError	If the source NormalizedRhs was not retained (i.e. the instance was constructed directly rather than via :func:compile_rhs).

Source code in src/op_system/compile.py

def __getstate__(self) -> dict[str, Any]:
    """Return picklable state.

    The compiled ``eval_fn`` is a closure (and on the vectorized path
    captures compiled :class:`types.CodeType` objects), which is not
    portably picklable. Instead we serialize just the source
    :class:`NormalizedRhs` and let :meth:`__setstate__` recompile.

    Raises:
        TypeError: If the source ``NormalizedRhs`` was not retained
            (i.e. the instance was constructed directly rather than
            via :func:`compile_rhs`).
    """
    if self._rhs is None:
        msg = (
            "CompiledRhs is not picklable: the source NormalizedRhs was "
            "not retained. Construct via compile_rhs() to produce a "
            "picklable CompiledRhs."
        )
        raise TypeError(msg)
    return {"_rhs": self._rhs}

__setstate__(state)

Restore by recompiling from the pickled :class:NormalizedRhs.

Source code in src/op_system/compile.py

def __setstate__(self, state: Mapping[str, Any]) -> None:
    """Restore by recompiling from the pickled :class:`NormalizedRhs`."""
    rhs = state["_rhs"]
    rebuilt = compile_rhs(rhs)
    # frozen+slots dataclass: bypass __setattr__ via object.__setattr__
    object.__setattr__(self, "state_names", rebuilt.state_names)
    object.__setattr__(self, "param_names", rebuilt.param_names)
    object.__setattr__(self, "param_meta", rebuilt.param_meta)
    object.__setattr__(self, "state_meta", rebuilt.state_meta)
    object.__setattr__(self, "axis_meta", rebuilt.axis_meta)
    object.__setattr__(self, "observables", rebuilt.observables)
    object.__setattr__(self, "gauges", rebuilt.gauges)
    object.__setattr__(self, "gauge_fix", rebuilt.gauge_fix)
    object.__setattr__(self, "eval_fn", rebuilt.eval_fn)
    object.__setattr__(self, "meta", rebuilt.meta)
    object.__setattr__(self, "operators", rebuilt.operators)
    object.__setattr__(self, "factorize_axes", rebuilt.factorize_axes)
    object.__setattr__(self, "block_axes", rebuilt.block_axes)
    object.__setattr__(self, "pytree_eval_fn", rebuilt.pytree_eval_fn)
    object.__setattr__(self, "template_shapes", rebuilt.template_shapes)
    object.__setattr__(self, "block_pytree_eval_fn", rebuilt.block_pytree_eval_fn)
    object.__setattr__(self, "block_template_shapes", rebuilt.block_template_shapes)
    object.__setattr__(self, "history_requirements", rebuilt.history_requirements)
    object.__setattr__(self, "history_eval_fn", rebuilt.history_eval_fn)
    object.__setattr__(self, "body_eval_fn", rebuilt.body_eval_fn)
    object.__setattr__(self, "block_history_eval_fn", rebuilt.block_history_eval_fn)
    object.__setattr__(self, "block_body_eval_fn", rebuilt.block_body_eval_fn)
    object.__setattr__(self, "reactions", rebuilt.reactions)
    object.__setattr__(self, "reaction_gaps", rebuilt.reaction_gaps)
    object.__setattr__(self, "_rhs", rhs)

bind(params)

Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

Parameters:

Name	Type	Description	Default
params	Mapping[str, object]	Mapping of parameter names to values.	required

Returns:

Type	Description
Callable[[object, object], Float64Array]	A callable rhs(t, y) that evaluates the RHS with params fixed.

Source code in src/op_system/compile.py

def bind(
    self, params: Mapping[str, object]
) -> Callable[[object, object], Float64Array]:
    """Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

    Args:
        params: Mapping of parameter names to values.

    Returns:
        A callable `rhs(t, y)` that evaluates the RHS with `params` fixed.
    """
    params_dict = dict(params)

    def rhs(t: object, y: object) -> Float64Array:
        return self.eval_fn(t, y, **params_dict)

    return rhs

coordinate_view()

Build deterministic evaluators in the declared state coordinates.

Returns:

Type	Description
StateCoordinateView	An explicit view with initial-state encoders and natural-scale
StateCoordinateView	observations. The original evaluator/reactions remain natural.

Source code in src/op_system/compile.py

def coordinate_view(self) -> StateCoordinateView:
    """Build deterministic evaluators in the declared state coordinates.

    Returns:
        An explicit view with initial-state encoders and natural-scale
        observations. The original evaluator/reactions remain natural.
    """
    from .coordinates import _build_state_coordinates  # ruff: ignore[import-outside-top-level]

    return _build_state_coordinates(self)

EvalFn

Bases: Protocol

Callable RHS evaluator supporting runtime parameter kwargs.

Accepts a flat (n_state,) state array and returns a flat (n_state,) derivative array in the same array namespace.

HistoryEvalFn

Bases: Protocol

Callable RHS evaluator with history provider support.

Accepts y as a StateDict and a history_provider object implementing a query(signal_id, body, **options) method. Returns a StateDict of derivatives.

PytreeEvalFn

Bases: Protocol

Callable RHS evaluator operating on shaped PyTree state dicts.

Accepts y as a StateDict (mapping from state-template base name to a shaped array with the template's natural N-D shape) and returns a StateDict of the same structure containing the derivative. Enables the engine to skip the flatten/unflatten step entirely and expose the full tensor structure to JAX/XLA.

ReactionPropensityFn

Bases: Protocol

Callable propensity evaluator for one named reaction/transition.

Accepts y as a StateDict and returns an array shaped like the reaction's from_axes (one independent rate per source cell) -- see :class:CompiledReaction.

compile_rhs(rhs, *, xp=None, detect_gauges=False, gauge_fix=None)

Compile a normalized RHS into a runnable evaluation function.

Always uses the vectorized eval path that operates on shaped buffers (one tensor expression per state template) for specs that declare axes. Specs without axes (genuinely scalar models) fall back to the scalar path. Raising :class:UnsupportedFeatureError if an axis-indexed spec cannot be vectorized, rather than silently falling back to the catastrophically slow scalar path.

The returned eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (through array_api_compat.array_namespace), and returns arrays in that same namespace. Calling it with JAX arrays (or tracers) yields a JAX-native computation suitable for jax.jit / jax.vmap without any correctness wrapping.

Parameters:

Name	Type	Description	Default
rhs	NormalizedRhs	Normalized RHS produced by op_system.specs.normalize_rhs.	required
detect_gauges	bool	Opt in to conservative structural scaling detection.	False
gauge_fix	Mapping[str, float] | None	Positive scalar anchors for the reported scaling action.	None
xp	object | None	Deprecated. Formerly the compile-time array namespace. Now ignored — the namespace is resolved per call from the input y. Will be removed in a future release.	None

Returns:

Type	Description
CompiledRhs	A CompiledRhs containing an eval_fn(t, y, **params) -> dydt.
CompiledRhs	For axis-indexed specs the returned object also carries
CompiledRhs	pytree_eval_fn and template_shapes. If the spec declares
CompiledRhs	axes but the vectorizer cannot build a plan an
CompiledRhs	UnsupportedFeatureError is raised (see bail reason in the detail
CompiledRhs	message) rather than silently degrading to the scalar path.

Source code in src/op_system/compile.py

def compile_rhs(  # ruff: ignore[too-many-locals]
    rhs: NormalizedRhs,
    *,
    xp: object | None = None,
    detect_gauges: bool = False,
    gauge_fix: Mapping[str, float] | None = None,
) -> CompiledRhs:
    """Compile a normalized RHS into a runnable evaluation function.

    Always uses the vectorized eval path that operates on shaped buffers
    (one tensor expression per state template) for specs that declare axes.
    Specs without axes (genuinely scalar models) fall back to the scalar path.
    Raising :class:`UnsupportedFeatureError` if an axis-indexed spec cannot be
    vectorized, rather than silently falling back to the catastrophically slow
    scalar path.

    The returned ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (through ``array_api_compat.array_namespace``), and returns arrays in that same
    namespace. Calling it with JAX arrays (or tracers) yields a JAX-native
    computation suitable for ``jax.jit`` / ``jax.vmap`` without any
    correctness wrapping.

    Args:
        rhs: Normalized RHS produced by `op_system.specs.normalize_rhs`.
        detect_gauges: Opt in to conservative structural scaling detection.
        gauge_fix: Positive scalar anchors for the reported scaling action.
        xp: **Deprecated.** Formerly the compile-time array namespace.
            Now ignored — the namespace is resolved per call
            from the input ``y``. Will be removed in a future release.

    Returns:
        A `CompiledRhs` containing an `eval_fn(t, y, **params) -> dydt`.
        For axis-indexed specs the returned object also carries
        ``pytree_eval_fn`` and ``template_shapes``.  If the spec declares
        axes but the vectorizer cannot build a plan an
        ``UnsupportedFeatureError`` is raised (see bail reason in the detail
        message) rather than silently degrading to the scalar path.
    """
    _warn_on_deprecated_xp(xp)
    _validate_rhs_type(rhs)

    _validate_history_kinds(
        _history_requirements_from_ir(
            aliases_ir=rhs.aliases_ir,
            equations_ir=rhs.equations_ir,
        )
        if _reduce_ir_has_history(rhs)
        else ()
    )

    vec, plan, eval_fn, pytree_eval_fn, template_shapes = _build_primary_eval_artifacts(
        rhs
    )

    # Apply both time-varying and synth-const wrappers to ``pytree_eval_fn``
    # BEFORE ``_build_history_artifacts`` consumes it.  history_eval_fn /
    # body_eval_fn capture the pytree_eval_fn reference at construction
    # time, so any later re-wrapping would not propagate into them; that
    # would leave runtime history-body calls missing time-varying param
    # slicing and synthesized constants (e.g. __op_system_mask__* one-hot
    # arrays for pinned transition selectors).
    synth_consts: Mapping[str, object] | None = None
    if eval_fn is not None:
        eval_fn, pytree_eval_fn = _wrap_time_varying_artifacts(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )
        eval_fn, pytree_eval_fn, synth_consts = _apply_synth_const_wrappers(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )

    history_requirements, history_eval_fn, body_eval_fn = _build_history_artifacts(
        rhs=rhs,
        plan=plan,
        pytree_eval_fn=pytree_eval_fn,
    )

    eval_fn = _resolve_eval_fn(
        rhs=rhs,
        eval_fn=eval_fn,
        history_requirements=history_requirements,
    )

    block_axes = analyze_block_axes(rhs)

    # ------------------------------------------------------------------
    # Block-stripped compile: produce a per-block-coord pytree_eval_fn by
    # stripping the first factorize axis from the RHS and re-running the
    # vectorizer.  Engines can jax.vmap this over the block axis instead
    # of baking literal axis indices that break under vmap.
    # ------------------------------------------------------------------
    block_pytree_eval_fn, block_template_shapes = _build_block_pytree_artifacts(
        rhs=rhs,
        vec=vec,
        pytree_eval_fn=pytree_eval_fn,
        block_axes=block_axes,
        synth_consts=synth_consts,
    )

    # Build per-block history / body eval fns when both the block compile and
    # history path succeeded.  These wrap ``block_pytree_eval_fn`` exactly as
    # ``history_eval_fn`` / ``body_eval_fn`` wrap ``pytree_eval_fn``.
    block_history_eval_fn, block_body_eval_fn = _build_block_history_artifacts(
        block_pytree_eval_fn=block_pytree_eval_fn,
        history_requirements=history_requirements,
    )

    reaction_artifacts = _build_reaction_artifacts(rhs=rhs, plan=plan, vec=vec)
    compiled = CompiledRhs(
        state_names=rhs.state_names,
        param_names=tuple(rhs.param_names),
        eval_fn=eval_fn,
        meta=rhs.meta,
        param_meta=MappingProxyType(
            dict(
                cast(
                    "Mapping[str, ParameterDeclaration]",
                    rhs.meta.get("parameter_declarations", {}),
                )
            )
        ),
        state_meta=MappingProxyType(
            dict(
                cast(
                    "Mapping[str, StateDeclaration]",
                    rhs.meta.get("state_declarations", {}),
                )
            )
        ),
        axis_meta=MappingProxyType(
            _axis_declarations(
                cast("Sequence[Mapping[str, object]]", rhs.meta.get("axes", ()))
            )
        ),
        operators=_parse_operator_descriptors(rhs.meta),
        observables=MappingProxyType(_compile_observables(rhs)),
        factorize_axes=_parse_factorize_axes(rhs.meta),
        block_axes=block_axes,
        pytree_eval_fn=pytree_eval_fn,
        template_shapes=template_shapes,
        block_pytree_eval_fn=block_pytree_eval_fn,
        block_template_shapes=block_template_shapes,
        history_requirements=history_requirements,
        history_eval_fn=history_eval_fn,
        body_eval_fn=body_eval_fn,
        block_history_eval_fn=block_history_eval_fn,
        block_body_eval_fn=block_body_eval_fn,
        reactions=reaction_artifacts.reactions,
        reaction_gaps=reaction_artifacts.gaps,
        _rhs=rhs,
    )

    return _gauge_options(compiled, detect=detect_gauges, anchors=gauge_fix)