OP System

op_system

op_system.

Domain-agnostic RHS specification + compilation utilities.

Public API (v1)

Primary user entrypoints: - compile_spec: Validate, normalize, and compile a RHS specification in one step. - normalize_rhs: Validate and normalize a YAML-friendly RHS specification. - compile_rhs: Compile a NormalizedRhs into an efficient callable RHS.

Core data structures: - NormalizedRhs - CompiledRhs - OperatorDescriptor

Design guarantees: - No dependency on provider/adapters (eg flepimop2). - Stable interface for downstream engines. - Forward-compatible with multiphysics extensions.

IdentifierString = Annotated[str, AfterValidator(_validate_identifier_string)] module-attribute

Custom pydantic type for validated identifier strings used in op_system.

Identifier strings are used for state names, dimension names, and other keys in the system. They must be non-empty, contain only alphanumeric characters, and start with a letter. Leading and trailing whitespace is stripped before validation.

Examples:

>>> from pydantic import BaseModel
>>> from op_system import IdentifierString
>>> class ExampleModel(BaseModel):
...     identifier: IdentifierString
...
>>> ExampleModel(identifier="S")
ExampleModel(identifier='S')
>>> ExampleModel(identifier="  Foobar  ")
ExampleModel(identifier='Foobar')
>>> ExampleModel(identifier="123abc")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must contain only alphanumerical characters and start with a letter. [...]
    For further information visit ...
>>> ExampleModel(identifier="")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must not be empty. [...]
    For further information visit ...

Array

Bases: Protocol

Structural Array-API protocol.

Any object whose runtime type implements shape, dtype, __array_namespace__ and item satisfies this protocol. NumPy

= 2.0 ndarrays, JAX arrays (concrete and traced), and PyTorch tensors (via the array-api compat layer) all qualify.

The namespace returned by __array_namespace__ is the only gate op_system uses to dispatch operations: input → namespace → output in that same namespace. No conversion, no coercion, no compile-time backend selector.

BlockAxisInfo(name, size, state_axis_pos, param_axis_pos) dataclass

Metadata for one block-diagonal (factorizable) axis.

A BlockAxisInfo is attached to :class:~op_system.compile.CompiledRhs for each axis declared in spec["factorize_axes"] that passes the IR separability check in :func:analyze_block_axes. Engines (e.g. the diffrax plugin) consume this to partition ODE solves with jax.vmap over the block axis.

Attributes:

Name	Type	Description
name	str	Axis name string (e.g. "loc").
size	int	Number of elements along the axis.
state_axis_pos	dict[str, int]	Maps each state-template base name to the integer position of this axis within that template's shape tuple. Only templates that carry this axis appear in the dict.
param_axis_pos	dict[str, int | None]	Maps each shaped parameter name to the integer position of this axis within the actual runtime array that the engine passes to the eval function, or None if the parameter does not carry this axis (broadcast). For non-time-varying shaped parameters the position is the index in the parameter's axis tuple. For time-varying parameters the runtime array has time prepended at index 0, so the position equals the index of the block axis in the full (time, *spatial_axes) tuple. Parameters that are entirely scalar (not shaped) do not appear in this dict and are always broadcast.

Note

BlockAxisInfo uses dict fields and therefore cannot be used as a hash key. It is stored on :class:~op_system.compile.CompiledRhs with hash=False and is fully pickle-stable.

Examples:

>>> info = BlockAxisInfo(
...     name="loc",
...     size=3,
...     state_axis_pos={"S": 1, "I": 1, "R": 1},
...     param_axis_pos={"rho": 0, "beta": None},
... )
>>> info.name
'loc'
>>> info.size
3
>>> info.state_axis_pos["S"]
1
>>> info.param_axis_pos["rho"]
0
>>> info.param_axis_pos["beta"] is None
True

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.

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, _rhs=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, "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, "_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

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.

ExprRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="expr" specs (explicit d(state)/dt equations).

Produced by :func:normalize_expr_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

ExpressionString(source) dataclass

Validated expression wrapper with cached AST and symbol names.

as_ir()

Return the typed IR representation for this expression.

Returns:

Type	Description
Expr	Parsed typed IR tree.

Source code in src/op_system/_symbols.py

def as_ir(self) -> Expr:
    """Return the typed IR representation for this expression.

    Returns:
        Parsed typed IR tree.
    """
    return to_ir(self.ast)

as_lowered_ir()

Return helper-lowered typed IR for this expression.

Returns:

Type	Description
Expr	Typed IR tree with helper calls lowered to Reduce nodes.

Source code in src/op_system/_symbols.py

def as_lowered_ir(self) -> Expr:
    """Return helper-lowered typed IR for this expression.

    Returns:
        Typed IR tree with helper calls lowered to ``Reduce`` nodes.
    """
    return parse_expr_to_ir(self.source, lower_helpers=True)

OperatorDescriptor(axis, kind=None, bc=None, velocity=None, rate=None, kernel=None) dataclass

Typed description of a spatial operator declared in an op_system RHS spec.

Captures the model-level description of an operator that is known at compile time: which axis it acts on, its kind (advection, diffusion, etc.), optional boundary condition, and the names of any runtime parameters (velocity, rate) it consumes. Grid geometry, CN matrix construction, and solver staging remain downstream concerns handled by the engine.

Attributes:

Name	Type	Description
axis	str	Name of the axis the operator acts on (e.g. "loc").
kind	str | None	Operator type string, e.g. "advection", "diffusion", "transport", "jump_integral", or None if unspecified.
bc	str | None	Boundary condition, e.g. "absorbing", "periodic", "neumann", "reflecting", or None if unspecified.
velocity	str | None	Parameter name for an advection velocity, or None.
rate	str | None	Parameter name for a diffusion rate/coefficient, or None.
kernel	Mapping[str, Any] | None	Mixing-kernel sub-specification, or None.

Examples:

>>> od = OperatorDescriptor(axis="loc")
>>> od.axis
'loc'
>>> od.kind is None
True
>>> od.bc is None
True
>>> od.velocity is None
True
>>> od2 = OperatorDescriptor(
...     axis="loc",
...     kind="advection",
...     bc="absorbing",
...     velocity="v_advec",
...     rate="diff_r",
... )
>>> od2.kind
'advection'
>>> od2.bc
'absorbing'
>>> od2.velocity
'v_advec'
>>> od2.rate
'diff_r'

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.

StateString

Bases: BaseModel

Structured representation of a state string.

A state string is either a bare state name like "S" or a state name followed immediately by bracketed dimensions like "R[age,vax]".

Examples:

>>> StateString.model_validate("S")
StateString(name='S', dims=())
>>> recovery = StateString.model_validate("R[age,vax]")
>>> recovery
StateString(name='R', dims=('age', 'vax'))
>>> print(recovery)
R[age,vax]
>>> recovery.model_dump()
'R[age,vax]'
>>> StateString.model_validate("Foobar[ age , vax ]")
StateString(name='Foobar', dims=('age', 'vax'))

__str__()

Return the compact string form.

Returns:

Type	Description
str	Compact state string.

Examples:

>>> str(StateString(name="S", dims=()))
'S'
>>> str(StateString(name="R", dims=("age",)))
'R[age]'
>>> str(StateString(name="R", dims=("age", "vax")))
'R[age,vax]'
>>> str(StateString(name="lambda", dims=("age", "vax", "state")))
'lambda[age,vax,state]'

Source code in src/op_system/_state_string.py

def __str__(self) -> str:
    """
    Return the compact string form.

    Returns:
        Compact state string.

    Examples:
        >>> str(StateString(name="S", dims=()))
        'S'
        >>> str(StateString(name="R", dims=("age",)))
        'R[age]'
        >>> str(StateString(name="R", dims=("age", "vax")))
        'R[age,vax]'
        >>> str(StateString(name="lambda", dims=("age", "vax", "state")))
        'lambda[age,vax,state]'
    """
    if not self.dims:
        return self.name
    dims = ",".join(self.dims)
    return f"{self.name}[{dims}]"

TransitionsRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="transitions" specs (per-capita hazard diagram).

Produced by :func:normalize_transitions_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

compile_rhs(rhs, *, xp=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 (y.__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
xp	object | None	Deprecated. Formerly the compile-time array backend 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(rhs: NormalizedRhs, *, xp: object | 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
    (``y.__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`.
        xp: **Deprecated.** Formerly the compile-time array backend
            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)

    raw_history_requirements = _history_requirements_from_ir(
        aliases_ir=rhs.aliases_ir,
        equations_ir=rhs.equations_ir,
    )
    _validate_history_kinds(raw_history_requirements)

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

    return CompiledRhs(
        state_names=rhs.state_names,
        param_names=tuple(rhs.param_names),
        eval_fn=eval_fn,
        meta=rhs.meta,
        operators=_parse_operator_descriptors(rhs.meta),
        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,
        _rhs=rhs,
    )

compile_spec(spec, *, xp=None, backend=DEFAULT_ARRAY_BACKEND)

Validate, normalize, and compile a RHS specification in one call.

This is the recommended public entrypoint for most users and adapters.

The compiled eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (y.__array_namespace__()), so a single compiled callable handles NumPy, JAX (concrete and traced), and any other Array-API backend natively. No compile-time backend selection is required.

Parameters:

Name	Type	Description	Default
spec	dict[str, object]	Raw RHS specification mapping (YAML/JSON friendly).	required
xp	object | None	Deprecated. Formerly the compile-time array backend namespace. Now ignored — see compile_rhs for details. Will be removed in a future release.	None
backend	Literal['numpy', 'jax']	Deprecated. Formerly selected the compile-time backend ("numpy" or "jax"). Now ignored. Will be removed in a future release.	DEFAULT_ARRAY_BACKEND

Returns:

Name	Type	Description
CompiledRhs	CompiledRhs	Runnable RHS callable container.

Source code in src/op_system/__init__.py

def compile_spec(  # noqa: RUF067
    spec: dict[str, object],
    *,
    xp: object | None = None,
    backend: Literal["numpy", "jax"] = DEFAULT_ARRAY_BACKEND,
) -> CompiledRhs:
    """
    Validate, normalize, and compile a RHS specification in one call.

    This is the recommended public entrypoint for most users and adapters.

    The compiled ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (``y.__array_namespace__()``), so a single compiled callable handles
    NumPy, JAX (concrete and traced), and any other Array-API backend
    natively. No compile-time backend selection is required.

    Args:
        spec: Raw RHS specification mapping (YAML/JSON friendly).
        xp: **Deprecated.** Formerly the compile-time array backend
            namespace. Now ignored — see ``compile_rhs`` for details.
            Will be removed in a future release.
        backend: **Deprecated.** Formerly selected the compile-time
            backend (``"numpy"`` or ``"jax"``). Now ignored. Will be
            removed in a future release.

    Returns:
        CompiledRhs: Runnable RHS callable container.
    """
    if xp is not None or backend != DEFAULT_ARRAY_BACKEND:
        warnings.warn(
            "compile_spec(xp=..., backend=...) is deprecated and ignored. "
            "The compiled eval_fn now infers its array namespace from the "
            "input `y` at call time via __array_namespace__(); pass JAX "
            "arrays for a JAX-native call, NumPy arrays for a NumPy call. "
            "These kwargs will be removed in a future release.",
            DeprecationWarning,
            stacklevel=2,
        )

    rhs = normalize_rhs(spec)
    return compile_rhs(rhs)

normalize_expr_rhs(spec)

Normalize an expression-based RHS specification.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any]	Raw RHS specification mapping.	required

Returns:

Type	Description
ExprRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_expr_rhs(spec: Mapping[str, Any]) -> ExprRhs:  # noqa: C901, PLR0914, PLR0915
    """Normalize an expression-based RHS specification.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    equations_map = spec.get("equations")
    if not isinstance(equations_map, dict):
        raise InvalidRhsSpecError(detail="equations must be a mapping of state->expr")

    equations_map = {_normalize_bracket_key(k): v for k, v in equations_map.items()}

    axes_meta = _normalize_axes(spec.get("axes"))
    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=set(state_raw),
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "axes": axes_meta,
        "state_axes": meta_parts[1],
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    for reserved_key in ("sources", "couplings", "constraints"):
        if reserved_key in spec:
            meta[reserved_key] = spec.get(reserved_key)

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    # Pre-scan raw aliases + equations for shaped-parameter references before
    # alias/template expansion mangles bracketed bases into per-coord names.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    raw_expressions.extend(v for v in equations_map.values() if isinstance(v, str))
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    time_axis_name = _resolve_time_axis_name(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    if time_varying_full:
        _strip_time_axis_in_mapping(
            equations_map,
            tv_full_axes=time_varying_full,
            time_axis_name=time_axis_name,
        )
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            meta_parts[0],
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_expr_chains(
            chains=chain_block,
            state_expanded=state_expanded,
            equations_map=equations_map,
        )

    unknown_keys = [
        k
        for k in equations_map
        if k not in state_expanded and k not in template_map_all
    ]
    if unknown_keys:
        raise InvalidRhsSpecError(
            detail=f"unknown equation key(s): {sorted(unknown_keys)}"
        )

    equations_ir_built, equations_ir_reduce, all_syms = _build_equations_ir_from_raw(
        state_expanded=state_expanded,
        equations_map=equations_map,
        template_map=template_map_all,
        axes=axes_meta,
        shaped_params=shaped_params,
        axis_lookup=axis_lookup_dict,
        aliases_ir=aliases_ir_map,
    )
    # Collect free symbols from alias bodies (alias bodies may reference
    # params not appearing in any equation directly). Walk the *reduce*
    # map (Reduce nodes still folded) rather than the fully expanded
    # map: alias inlining is identical between the two, but the reduce
    # form is orders of magnitude smaller for continuum specs. Share a
    # single id-keyed memo across the per-cell entries so common
    # subtrees are visited at most once (issue #145).
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_ir_val in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_ir_val, memo=fs_memo)

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations (e.g. coord-pinned copies of the same
    # alias template body) are rendered only once (issue #145).
    unparse_memo: dict[tuple[int, int, bool], str] = {}
    equations_strings = _derive_equation_strings(
        equations_ir_built, _unparse_memo=unparse_memo
    )
    aliases_strings = _derive_alias_strings(
        aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo
    )

    return ExprRhs(
        state_names=tuple(state_expanded),
        equations=equations_strings,
        aliases=aliases_strings,
        param_names=_sorted_unique(
            sym
            for sym in all_syms
            if sym not in set(state_expanded)
            and sym not in aliases_ir_map
            and sym not in shaped_set
            and sym not in time_varying_set
            and sym not in axis_name_set
            and sym not in template_base_set
            and sym not in _SHAPED_PARAM_BUILTIN_NAMES
        ),
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta=meta,
        state_templates=_build_state_templates(
            state_raw,
            axes=axes_meta,
            state_template_map=state_template_map,
            state_expanded=state_expanded,
        ),
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )

normalize_rhs(spec)

Normalize a RHS specification dict into a backend-facing representation.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any] | None	Raw RHS specification mapping.	required

Returns:

Type	Description
NormalizedRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.
UnsupportedFeatureError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_rhs(spec: Mapping[str, Any] | None) -> NormalizedRhs:
    """Normalize a RHS specification dict into a backend-facing representation.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
        UnsupportedFeatureError: If validation fails.
    """
    if spec is None:
        raise InvalidRhsSpecError(detail="rhs specification is required")

    kind = str(spec.get("kind", "expr")).strip().lower()

    if kind == "expr":
        return normalize_expr_rhs(spec)

    if kind == "transitions":
        return normalize_transitions_rhs(spec)

    raise UnsupportedFeatureError(
        feature=f"rhs.kind={kind}",
        detail="Only 'expr' and 'transitions' are supported in v1.",
    )

normalize_transitions_rhs(spec)

Normalize a transition-based RHS specification (diagram/hazard semantics).

Returns:

Type	Description
TransitionsRhs	Backend-facing normalized RHS representation for the transitions kind.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_transitions_rhs(  # noqa: C901, PLR0912, PLR0914, PLR0915
    spec: Mapping[str, Any],
) -> TransitionsRhs:
    """Normalize a transition-based RHS specification (diagram/hazard semantics).

    Returns:
        Backend-facing normalized RHS representation for the transitions kind.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    transitions_raw = spec.get("transitions")
    if transitions_raw is None:
        transitions_raw = []
    elif isinstance(transitions_raw, list):
        transitions_raw = list(transitions_raw)
    else:
        raise InvalidRhsSpecError(detail="transitions must be a list")

    axes_meta = _normalize_axes(spec.get("axes"))

    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=None,
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "transitions": transitions_raw,
        "axes": axes_meta,
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    meta.update({
        k: spec[k] for k in ("sources", "couplings", "constraints") if k in spec
    })

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_transition_chains(
            chains=chain_block,
            state_raw=state_raw,
            transitions_raw=transitions_raw,
            state_set=set(state_raw),
        )

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    # Pre-scan raw aliases + transition rates for shaped-parameter references.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    for tr in transitions_raw:
        if isinstance(tr, _MappingABC):
            r = tr.get("rate")
            if isinstance(r, str):
                raw_expressions.append(r)
            n = tr.get("name")
            if isinstance(n, str):
                raw_expressions.append(n)
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    time_axis_name = _resolve_time_axis_name(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    if time_varying_full:
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )
        for tr in transitions_raw:
            if not isinstance(tr, _MappingABC):
                continue
            r = tr.get("rate")
            if isinstance(r, str):
                tr["rate"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    r,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )
            n = tr.get("name")
            if isinstance(n, str):
                tr["name"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    n,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            aliases_raw_map,
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    state_set = set(state_expanded)
    # Collect alias symbols from IR. Walk the Reduce-folded map (same
    # inlined symbol set as the full-expansion map but vastly smaller
    # on continuum specs) and share an id-keyed memo across per-cell
    # entries that share subtrees by identity (issue #145).
    all_syms: set[str] = set()
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_expr in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_expr, memo=fs_memo)

    _apply_coord_shifts(
        transitions_raw=transitions_raw,
        state_expanded=state_expanded,
        axes=axes_meta,
        state_template_map=state_template_map,
    )

    if not transitions_raw:
        raise InvalidRhsSpecError(
            detail="transitions must be non-empty after applying chain expansion"
        )

    # Build equations and collect expanded transitions in one IR-native pass
    pinned_mask_names, pinned_mask_values = _discover_pinned_token_masks(
        transitions_raw, axis_lookup=axis_lookup_dict
    )
    if pinned_mask_values:
        # Register one-hot masks as shaped params so the vectorizer's
        # extra-param-buffers plumbing assembles them at eval time; stash
        # the actual values under ``meta`` so ``compile_rhs`` can inject
        # them into the eval_fn's ``params`` automatically.
        for (axis, _coord), mask_name in pinned_mask_names.items():
            shaped_params[mask_name] = (axis,)
        meta["op_system_synth_constants"] = dict(pinned_mask_values)

    equations_ir_pre_inline, equations_ir_reduce, transitions_expanded, rate_syms = (
        _build_transition_equations_ir(
            transitions_raw,
            state_set=state_set,
            state_expanded=state_expanded,
            axes=axes_meta,
            axis_lookup=axis_lookup_dict,
            template_map=template_map_all,
            shaped_params=shaped_params,
            mask_names=pinned_mask_names,
            time_axis_name=time_axis_name,
            alias_bases={
                parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map
            },
        )
    )
    all_syms |= rate_syms

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations are rendered only once (issue #145).
    unparse_memo_final: dict[tuple[int, int, bool], str] = {}
    eqs_tuple = _derive_equation_strings(
        equations_ir_pre_inline, _unparse_memo=unparse_memo_final
    )
    # Inline aliases directly into the per-cell IR we already built in
    # ``_build_transition_equations_ir`` rather than re-parsing the
    # round-tripped equation strings. Avoids 73k x ``parse_expr_to_ir``
    # plus a redundant IR rebuild on large continuum specs (issue #145).
    # ``aliases_ir_map`` is already fully alias-inlined inside
    # ``_build_aliases_ir_from_raw`` so cycle detection is redundant here.
    #
    # The dominant cost on large specs is the per-cell ``inline_aliases``
    # call. Because synthesized transitions install the SAME ``synth_to`` /
    # ``synth_neg`` IR object into every cell of a template, the *terms*
    # of the per-cell sum are shared across many cells even though the
    # outer ``Apply(op="+", ...)`` wrapper is unique per cell. Inlining
    # term-by-term with a shared ``result_memo`` keyed on ``id(term)``
    # collapses the alias-substitution work from O(n_state) to
    # O(n_unique_terms) (issue #145).
    alias_inline_memo: dict[int, frozenset[str]] = {}
    alias_inline_result_memo: dict[int, Expr] = {}

    def _inline_one(expr: Expr) -> Expr:
        return inline_aliases(
            expr,
            aliases_ir_map,
            memo=alias_inline_memo,
            skip_cycle_check=True,
            result_memo=alias_inline_result_memo,
        )

    # Dedup the post-inline outer Apply by tuple of arg ids so cells
    # whose inlined-term tuple is identical share one wrapper, letting
    # downstream identity-keyed memos (e.g. unparse_ir) cache the
    # rendered string once per unique equation. Issue #145.
    apply_plus_dedup: dict[tuple[int, ...], Expr] = {}
    # The upstream ``sum_dedup_full`` collapses ``equations_ir_pre_inline``
    # to a handful of unique outer Apply objects shared across many
    # cells (e.g. 7 unique exprs across 72,828 cells on the COVID19_USA
    # continuum spec). Cache the per-expr inlined result by ``id(expr)``
    # so we skip the per-term ``_inline_one`` calls entirely on cache
    # hits -- this collapses ~1.9M memo-hit calls to ~7 real inlines
    # plus 72k dict lookups (issue #145).
    outer_inline_memo: dict[int, Expr] = {}

    def _inline_outer(expr: Expr) -> Expr:
        """Run alias-inlining for one outer equation expression.

        Returns:
            The inlined expression (possibly the original ``expr`` when no
            inlining was needed).
        """
        if isinstance(expr, Apply) and expr.op == "+":
            new_args = tuple(_inline_one(a) for a in expr.args)
            if all(n is o for n, o in zip(new_args, expr.args, strict=True)):
                return expr
            dedup_key = tuple(id(a) for a in new_args)
            cached_apply = apply_plus_dedup.get(dedup_key)
            if cached_apply is None:
                cached_apply = Apply(op="+", args=new_args)
                apply_plus_dedup[dedup_key] = cached_apply
            return cached_apply
        return _inline_one(expr)

    equations_ir_built_list: list[Expr | None] = []
    for expr in equations_ir_pre_inline:
        if expr is None:
            equations_ir_built_list.append(None)
            continue
        if not aliases_ir_map:
            equations_ir_built_list.append(expr)
            continue
        cached_outer = outer_inline_memo.get(id(expr))
        if cached_outer is not None:
            equations_ir_built_list.append(cached_outer)
            continue
        try:
            result_expr = _inline_outer(expr)
        except (ValueError, RecursionError):
            result_expr = expr
        outer_inline_memo[id(expr)] = result_expr
        equations_ir_built_list.append(result_expr)
    equations_ir_built = tuple(equations_ir_built_list)
    return TransitionsRhs(
        state_names=tuple(state_expanded),
        equations=eqs_tuple,
        aliases=_derive_alias_strings(
            aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo_final
        ),
        param_names=_sorted_unique(
            sym
            for sym in all_syms
            if sym not in state_set
            and sym not in aliases_ir_map
            and sym not in shaped_set
            and sym not in time_varying_set
            and sym not in axis_name_set
            and sym not in template_base_set
            and sym not in _SHAPED_PARAM_BUILTIN_NAMES
        ),
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta={**meta, "transitions": transitions_expanded},
        state_templates=_build_state_templates(
            state_raw,
            axes=axes_meta,
            state_template_map=state_template_map,
            state_expanded=state_expanded,
        ),
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )