Skip to content

Specs

specs

op_system.specs.

Public façade for RHS specification models and normalization.

The implementation lives in op_system._normalize. This module re-exports the public types and functions under the stable op_system.specs namespace for backward compatibility.

Supported RHS kinds

1) kind: "expr" - explicit d(state)/dt equations per state variable. 2) kind: "transitions" - diagram-style per-capita hazard transitions.

See op_system._normalize for the full implementation.

AxisDeclaration

Bases: NamedTuple

Physical interpretation of a named axis.

Attributes:

Name Type Description
units str | None

Coordinate units, independent of measure units.
measure AxisMeasure | None

Explicit measure, or None when undeclared.

AxisMeasure

Bases: NamedTuple

Explicit physical measure in the axis coordinate order.

Attributes:

Name Type Description
kind _MeasureKind

Counting, cell widths/volumes, or point quadrature weights.
weights tuple[float, ...]

Finite strictly positive weights, one per coordinate.
units str | None

Optional measure units, without dimensional verification.

from_mapping(value, *, size) classmethod

Validate a measure without inferring geometry from coordinates.

Parameters:

Name Type Description Default
value object

Mapping with kind, optional units, and weights for cell/quadrature.

 required
size int

Number of axis coordinates.

 required

Returns:

Type Description
AxisMeasure

Immutable measure in coordinate order.

Raises:

Type Description
InvalidRhsSpecError

If the kind, units, weights or length are invalid.
Source code in src/op_system/_scientific_meta.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
@classmethod
def from_mapping(cls, value: object, *, size: int) -> AxisMeasure:
    """Validate a measure without inferring geometry from coordinates.

    Args:
        value: Mapping with kind, optional units, and weights for cell/quadrature.
        size: Number of axis coordinates.

    Returns:
        Immutable measure in coordinate order.

    Raises:
        InvalidRhsSpecError: If the kind, units, weights or length are invalid.
    """
    if not isinstance(size, int) or isinstance(size, bool) or size < 1:
        raise InvalidRhsSpecError(
            detail="axis measure size must be a positive integer"
        )
    if not isinstance(value, Mapping) or set(value) - {"kind", "weights", "units"}:
        raise InvalidRhsSpecError(
            detail="axis measure must contain kind, weights and/or units"
        )
    kind = value.get("kind")
    if not isinstance(kind, str) or kind not in {"counting", "cell", "quadrature"}:
        raise InvalidRhsSpecError(
            detail="axis measure kind must be counting, cell or quadrature"
        )
    if kind == "counting":
        if "weights" in value:
            raise InvalidRhsSpecError(
                detail="counting measure has unit weights; omit weights"
            )
        weights = (1.0,) * size
    else:
        raw_weights = value.get("weights")
        if not isinstance(raw_weights, (list, tuple)) or len(raw_weights) != size:
            raise InvalidRhsSpecError(
                detail="axis measure weights must match the coordinate count"
            )
        weights = tuple(_positive_weight(weight) for weight in raw_weights)
    return cls(cast("_MeasureKind", kind), weights, _optional_text(value, "units"))

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 
43
44
45
46
47
48
49
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 
51
52
53
54
55
56
57
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)

ObservableDefinition

Bases: NamedTuple

Backend-neutral definition of one instantaneous latent observable.

Attributes:

Name Type Description
name str

Unique output name within a compiled model.
expression str

Restricted pointwise state/parameter/time expression.
axes tuple[str, ...]

Output axes, after any explicit reduction.
reduction ObservableReduction | None

Optional unweighted sum or physical integral.
units str | None

Optional declared output units, without dimensional verification.
description str | None

Optional human-readable interpretation.

from_mapping(name, value) classmethod

Validate the bounded observable schema.

Parameters:

Name Type Description Default
name str

Output identifier.

 required
value object

Mapping with expr, optional axes/reduction/units/description.

 required

Returns:

Type Description
ObservableDefinition

Detached immutable definition.

Raises:

Type Description
InvalidRhsSpecError

If identifiers, fields or reductions are invalid.
Source code in src/op_system/_observables.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@classmethod
def from_mapping(cls, name: str, value: object) -> ObservableDefinition:
    """Validate the bounded observable schema.

    Args:
        name: Output identifier.
        value: Mapping with expr, optional axes/reduction/units/description.

    Returns:
        Detached immutable definition.

    Raises:
        InvalidRhsSpecError: If identifiers, fields or reductions are invalid.
    """
    if not _identifier(name):
        raise InvalidRhsSpecError(detail="observable names must be identifiers")
    if not isinstance(value, Mapping) or set(value) - {
        "expr",
        "axes",
        "reduction",
        "units",
        "description",
    }:
        raise InvalidRhsSpecError(detail=f"observable {name!r} has invalid fields")
    expression = _text(value, "expr", required=True)
    axes = _axes(value.get("axes", ()))
    reduction = None
    if "reduction" in value:
        raw = value["reduction"]
        if not isinstance(raw, Mapping) or set(raw) != {"kind", "axes"}:
            raise InvalidRhsSpecError(
                detail="observable reduction needs kind and axes"
            )
        kind = raw["kind"]
        reduction_axes = _axes(raw["axes"])
        if (
            not isinstance(kind, str)
            or kind not in {"sum", "integral"}
            or not reduction_axes
        ):
            raise InvalidRhsSpecError(
                detail="observable reduction must be a nonempty sum or integral"
            )
        reduction = ObservableReduction(
            cast('Literal["sum", "integral"]', kind), reduction_axes
        )
        if set(axes) & set(reduction_axes):
            raise InvalidRhsSpecError(
                detail="observable output and reduction axes must be disjoint"
            )
    return cls(
        name,
        cast("str", expression),
        axes,
        reduction,
        _text(value, "units"),
        _text(value, "description"),
    )

ObservableReduction

Bases: NamedTuple

Explicit reduction independent of legacy equation quadrature.

Attributes:

Name Type Description
kind Literal['sum', 'integral']

Unweighted sum or integral with declared physical measures.
axes tuple[str, ...]

Named axes to remove, in the declared reduction order.

ParameterDeclaration

Bases: NamedTuple

Optional scientific metadata for an inferred runtime parameter.

Attributes:

Name Type Description
support ParameterSupport | None

Declared domain, or None when it was not specified.
units str | None

Unit annotation, without dimensional-analysis verification.
description str | None

Optional human-readable parameter description.
fixed bool

Whether inference should treat the parameter as fixed. This does not bind a value or remove a runtime parameter.

from_mapping(value) classmethod

Validate and copy one parameter declaration.

Parameters:

Name Type Description Default
value object

Mapping with optional support, units, description and fixed.

 required

Returns:

Type Description
ParameterDeclaration

Immutable parameter declaration.

Raises:

Type Description
InvalidRhsSpecError

If an annotation is malformed or unknown.
Source code in src/op_system/_parameter_meta.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@classmethod
def from_mapping(cls, value: object) -> ParameterDeclaration:
    """Validate and copy one parameter declaration.

    Args:
        value: Mapping with optional support, units, description and fixed.

    Returns:
        Immutable parameter declaration.

    Raises:
        InvalidRhsSpecError: If an annotation is malformed or unknown.
    """
    if not isinstance(value, Mapping) or any(
        key not in {"support", "units", "description", "fixed"} for key in value
    ):
        raise InvalidRhsSpecError(
            detail=(
                "parameter declaration must contain only support, units, "
                "description and fixed"
            )
        )
    fixed = value.get("fixed", False)
    if not isinstance(fixed, bool):
        raise InvalidRhsSpecError(detail="parameter fixed must be a boolean")
    support = (
        ParameterSupport.from_mapping(value["support"])
        if "support" in value
        else None
    )
    return cls(
        support,
        _optional_text(value, "units"),
        _optional_text(value, "description"),
        fixed,
    )

ParameterSupport

Bases: NamedTuple

Declared domain; bounds are inclusive and positivity is strict.

Attributes:

Name Type Description
kind _SupportKind

Domain identifier; unit_interval includes zero and one.
lower float | None

Finite lower bound for bounded support, otherwise None.
upper float | None

Finite upper bound for bounded support, otherwise None.

from_mapping(value) classmethod

Parse a domain name or bounded-domain mapping.

Parameters:

Name Type Description Default
value object

Domain string or mapping with kind, lower and upper.

 required

Returns:

Type Description
ParameterSupport

Immutable support declaration.

Raises:

Type Description
InvalidRhsSpecError

If the domain or bounds are invalid.
Source code in src/op_system/_parameter_meta.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@classmethod
def from_mapping(cls, value: object) -> ParameterSupport:
    """Parse a domain name or bounded-domain mapping.

    Args:
        value: Domain string or mapping with kind, lower and upper.

    Returns:
        Immutable support declaration.

    Raises:
        InvalidRhsSpecError: If the domain or bounds are invalid.
    """
    if isinstance(value, str) and value in {"real", "positive", "unit_interval"}:
        return cls(cast("_SupportKind", value))
    if not isinstance(value, Mapping) or set(value) != {"kind", "lower", "upper"}:
        raise InvalidRhsSpecError(
            detail=(
                "parameter support must be real, positive, unit_interval, "
                "or a bounded mapping with kind, lower and upper"
            )
        )
    lower = _finite_bound(value["lower"])
    upper = _finite_bound(value["upper"])
    if value["kind"] != "bounded" or lower >= upper:
        raise InvalidRhsSpecError(
            detail="bounded parameter support requires finite numeric lower < upper"
        )
    return cls("bounded", float(lower), float(upper))

PinnedToken(axis, coord) dataclass

Axis token pinned to a specific coordinate.

StateDeclaration

Bases: NamedTuple

Optional representation and physical reduction requirements of a state.

Attributes:

Name Type Description
representation _Representation | None

Extensive counts/mass/counter or intensive density/ concentration; None when unknown. Cumulative means an extensive accumulated counter, without asserting monotonicity of its equations.
units str | None

State value units, without dimensional verification.
description str | None

Human-readable annotation.
measure_axes tuple[str, ...] | None

Axes whose measures are required for an intensive integral. None means undeclared; an empty tuple explicitly needs none.
scale Literal['identity', 'log']

Requested identity/log coordinate for an explicit coordinate_view. Primary compiled equations and reactions retain natural coordinates.

from_mapping(value) classmethod

Validate a state annotation without changing its equations.

Parameters:

Name Type Description Default
value object

Mapping of representation, units, description and measure_axes.

 required

Returns:

Type Description
StateDeclaration

Immutable declaration; omitted information remains unknown.

Raises:

Type Description
InvalidRhsSpecError

If fields or their values are invalid.
Source code in src/op_system/_scientific_meta.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
@classmethod
def from_mapping(cls, value: object) -> StateDeclaration:
    """Validate a state annotation without changing its equations.

    Args:
        value: Mapping of representation, units, description and measure_axes.

    Returns:
        Immutable declaration; omitted information remains unknown.

    Raises:
        InvalidRhsSpecError: If fields or their values are invalid.
    """
    if not isinstance(value, Mapping) or set(value) - {
        "representation",
        "units",
        "description",
        "measure_axes",
        "scale",
    }:
        raise InvalidRhsSpecError(
            detail="state metadata contains unknown fields or is not a mapping"
        )
    scale = value.get("scale", "identity")
    if not isinstance(scale, str) or scale not in {"identity", "log"}:
        raise InvalidRhsSpecError(detail="state scale must be identity or log")
    representation = value.get("representation")
    if "representation" in value and (
        not isinstance(representation, str)
        or representation
        not in {"counts", "cell_mass", "density", "concentration", "cumulative"}
    ):
        raise InvalidRhsSpecError(
            detail=(
                "state representation must be counts, cell_mass, density, "
                "concentration or cumulative"
            )
        )
    measure_axes = None
    if "measure_axes" in value:
        axes = value["measure_axes"]
        if (
            not isinstance(axes, (list, tuple))
            or any(not isinstance(axis, str) or not axis.strip() for axis in axes)
            or len(set(axes)) != len(axes)
        ):
            raise InvalidRhsSpecError(
                detail="state measure_axes must list unique nonempty axis names"
            )
        measure_axes = tuple(axes)
        if representation in {"counts", "cell_mass", "cumulative"} and measure_axes:
            raise InvalidRhsSpecError(
                detail="extensive states do not require measure weights"
            )
    return cls(
        cast("_Representation | None", representation),
        _optional_text(value, "units"),
        _optional_text(value, "description"),
        measure_axes,
        cast('Literal["identity", "log"]', scale),
    )

StateTemplate(base, axes, shape, expanded_names, coord_assignments, offset) dataclass

Structural record for a state template prior to scalar expansion.

A spec like state: ["S[age, vax]"] over axes age (size 4) and vax (size 2) produces a single StateTemplate with base="S", axes=("age", "vax"), shape=(4, 2), and expanded_names listing the eight scalar state names in cartesian-product order (age outer, vax inner — matching itertools.product and the order used to expand the flat state vector).

Non-templated entries (e.g. a bare "D" in state) are also reported as StateTemplate records with axes=(), shape=(), and a single expanded_names = ("D",) so consumers can iterate templates uniformly.

Attributes:

Name Type Description
base str

Compartment name without selector brackets (e.g. "S").
axes tuple[str, ...]

Wildcard axes in declaration order. Empty for scalar templates.
shape tuple[int, ...]

Per-axis sizes in axes order. Empty for scalar templates.
expanded_names tuple[str, ...]

Flat scalar state names, in cartesian-product order over axes (consistent with state_names).
coord_assignments tuple[Mapping[str, str], ...]

For each entry in expanded_names, the axis -> coord mapping. Empty dict for scalar templates.
offset int

Index of expanded_names[0] within the parent NormalizedRhs.state_names tuple (i.e. where this template's slice starts in the flat state vector).

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=(), reactions_ir=(), reaction_gaps_ir=()) 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.

WildcardToken(axis) dataclass

Axis token that expands over all coordinates of that axis.

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 
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
def normalize_expr_rhs(spec: Mapping[str, Any]) -> ExprRhs:  # ruff: ignore[complex-structure, too-many-locals, too-many-statements]
    """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()}

    time_axis_name = _resolve_time_axis_name(spec)
    axes_meta = _normalize_axes(spec.get("axes"), time_axis_name=time_axis_name)
    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],
    }
    meta["op_system_may_have_history"] = _may_have_history(spec)
    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")

    _validate_bare_template_references(
        state_raw,
        state_expanded,
        chain(
            (
                (f"aliases[{key!r}]", value)
                for key, value in (meta_parts[0] or {}).items()
            ),
            ((f"equations[{key!r}]", value) for key, value in equations_map.items()),
        ),
    )

    # 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)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    meta.update(
        _normalize_time_metadata(
            spec,
            time_axis_name=time_axis_name,
            axes_meta=axes_meta,
            time_varying_params=time_varying_full,
        )
    )
    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
        ]

    state_set = set(state_expanded)
    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_lazy(
        equations_ir_built, _unparse_memo=unparse_memo
    )
    aliases_strings = _derive_alias_strings(
        aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo
    )

    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
    )

    state_templates = _build_state_templates(
        state_raw,
        axes=axes_meta,
        state_template_map=state_template_map,
        state_expanded=state_expanded,
    )
    meta.update(
        _state_declaration_metadata(
            spec.get("state_metadata"),
            templates=state_templates,
            axes=axes_meta,
            state_axes=meta_parts[1],
        )
    )
    normalized = ExprRhs(
        state_names=tuple(state_expanded),
        equations=equations_strings,
        aliases=aliases_strings,
        param_names=param_names,
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta=meta,
        state_templates=state_templates,
        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,
        ),
    )
    return _finish_observables(normalized, spec)

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 
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
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 
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
def normalize_transitions_rhs(  # ruff: ignore[complex-structure, too-many-branches, too-many-locals, too-many-statements]
    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):
        # Copy each entry: time-axis stripping rewrites rates in place.
        # Stamp each with its spec position for reaction coverage records.
        transitions_raw = [
            {**tr, ORIGIN_KEY: f"transitions[{idx}]"}
            if isinstance(tr, _MappingABC)
            else tr
            for idx, tr in enumerate(transitions_raw)
        ]
        _hoist_coord_shift_rates(transitions_raw)
    else:
        raise InvalidRhsSpecError(detail="transitions must be a list")

    time_axis_name = _resolve_time_axis_name(spec)
    axes_meta = _normalize_axes(spec.get("axes"), time_axis_name=time_axis_name)

    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["op_system_may_have_history"] = _may_have_history(spec)
    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")

    _validate_bare_template_references(
        state_raw,
        state_expanded,
        chain(
            (
                (f"aliases[{key!r}]", value)
                for key, value in (meta_parts[0] or {}).items()
            ),
            (
                (f"{tr.get(ORIGIN_KEY, f'transitions[{index}]')}.rate", tr.get("rate"))
                for index, tr in enumerate(transitions_raw)
                if isinstance(tr, _MappingABC)
            ),
        ),
    )

    # 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)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    meta.update(
        _normalize_time_metadata(
            spec,
            time_axis_name=time_axis_name,
            axes_meta=axes_meta,
            time_varying_params=time_varying_full,
        )
    )
    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
    )
    shift_axes, shift_values = _discover_coord_shift_constants(
        transitions_raw, axis_lookup=axis_lookup_dict
    )
    # Routing reactions mask their no-op diagonal (see op_system._reactions).
    for tr in transitions_raw:
        offdiag_axis = routing_offdiag_axis(tr) if isinstance(tr, dict) else None
        if offdiag_axis is None:
            continue
        n_coords = len(axis_lookup_dict[offdiag_axis])
        offdiag_name = offdiag_constant_name(offdiag_axis)
        shift_axes[offdiag_name] = (offdiag_axis, offdiag_axis)
        shift_values[offdiag_name] = tuple(
            tuple(0.0 if i == j else 1.0 for j in range(n_coords))
            for i in range(n_coords)
        )
    if pinned_mask_values or shift_values:
        # Register one-hot masks and shift matrices 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,)
        shaped_params.update(shift_axes)
        meta["op_system_synth_constants"] = {**pinned_mask_values, **shift_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

    reactions_ir, reaction_gaps_ir = build_reaction_artifacts_ir(
        transitions_raw,
        axes=axes_meta,
        axis_lookup=axis_lookup_dict,
        shaped_params=shaped_params,
        time_axis_name=time_axis_name,
        aliases_raw=aliases_raw_map,
        state_axes={
            base: tuple(tok.axis for tok in tokens)
            for state_s in state_raw
            for base, tokens in (parse_selector(state_s),)
        },
    )

    _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).
    # Both memos key ``id()`` of terms owned by ``equations_ir_pre_inline``,
    # which outlives them -- the invariant that makes an identity key sound
    # (issue #200). Passing a freshly-built expression to either would
    # silently return another expression's inlined result.
    alias_inline_memo = _InlineMemo()
    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)
    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
    )

    state_templates = _build_state_templates(
        state_raw,
        axes=axes_meta,
        state_template_map=state_template_map,
        state_expanded=state_expanded,
    )
    meta.update(
        _state_declaration_metadata(
            spec.get("state_metadata"),
            templates=state_templates,
            axes=axes_meta,
            state_axes=meta_parts[1],
        )
    )
    normalized = TransitionsRhs(
        reactions_ir=reactions_ir,
        reaction_gaps_ir=reaction_gaps_ir,
        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=param_names,
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta={**meta, "transitions": transitions_expanded},
        state_templates=state_templates,
        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,
        ),
    )
    return _finish_observables(normalized, spec)