Specs

specs

op_system.specs.

RHS specification models and normalization utilities for op_system.

Design goals

• Domain-agnostic core: no imports from flepimop2 or other adapters.
• YAML-friendly RHS specifications that compile into a normalized representation consumable by op_engine (and other numerical backends).
• Minimal v1 implementation that demonstrates the idea without blocking future multiphysics extensions (IMEX operators, PDE terms, sources, etc.).

Current supported RHS kinds

1) kind: "expr" - User provides explicit expressions for d(state)/dt per state variable.

2) kind: "transitions" - User provides diagram-style transitions and per-capita hazard expressions. - Each transition contributes a flow: flow = hazard_expr * from_state and updates derivatives: d(from)/dt -= flow d(to)/dt += flow

Future-facing (not implemented, but reserved)

• kind: "multiphysics" or additional top-level keys such as:
  • sources: explicit additive per-state terms (births/imports/forcing)
  • operators: implicit operator specs/factories for IMEX (diffusion, transport)
  • couplings: structured couplings across axes (space/age/traits) and fields The normalization outputs include placeholders to carry these blocks forward, so adapters/backends can extend without changing the fundamental contract.

ConstraintRule

Bases: NamedTuple

Validated constraint rule produced by _normalize_constraints.

NormalizedRhs(kind, state_names, equations, aliases, param_names, all_symbols, meta) dataclass

Normalized RHS representation suitable for compilation/execution.

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
NormalizedRhs	Backend-facing normalized RHS representation.

Source code in src/op_system/specs.py

def normalize_expr_rhs(spec: Mapping[str, Any]) -> NormalizedRhs:
    """
    Normalize an expression-based RHS specification.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        _raise_invalid_rhs_spec(detail="state contains duplicate names")

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

    # Normalize equation keys so that e.g. "u[x, y]" matches template "u[x,y]"
    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),
    )

    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_invalid_rhs_spec(detail="expanded state contains duplicates")

    aliases, alias_template_map = _expand_alias_templates(
        meta_parts[0], axes=axes_meta, template_map_seed=state_template_map
    )
    template_map_all = {**state_template_map, **alias_template_map}

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

    # Validate equation keys: allow either concrete states or template keys
    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_invalid_rhs_spec(
            detail=f"unknown equation key(s): {sorted(unknown_keys)}"
        )

    all_syms = _collect_alias_symbols(aliases, axes=axes_meta)
    eqs = _gather_equations(
        state_expanded,
        equations_map,
        all_syms,
        axes=axes_meta,
        template_map=template_map_all,
    )

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

    return NormalizedRhs(
        kind="expr",
        state_names=tuple(state_expanded),
        equations=tuple(eqs),
        aliases=aliases,
        param_names=_sorted_unique(
            sym
            for sym in all_syms
            if sym not in set(state_expanded) and sym not in aliases
        ),
        all_symbols=frozenset(all_syms | set(aliases.keys())),
        meta=meta,
    )

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.

Source code in src/op_system/specs.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.
    """
    if spec is None:
        _raise_invalid_rhs_spec(detail="rhs specification is required")

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

    if kind == "expr":  # lowest-level escape hatch
        return normalize_expr_rhs(spec)

    if kind == "transitions":  # diagram-style hazards
        return normalize_transitions_rhs(spec)

    _raise_unsupported_feature(
        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
NormalizedRhs	Backend-facing normalized RHS representation for transitions kind.

Source code in src/op_system/specs.py

def normalize_transitions_rhs(
    spec: Mapping[str, Any],
) -> NormalizedRhs:
    """Normalize a transition-based RHS specification (diagram/hazard semantics).

    Returns:
        Backend-facing normalized RHS representation for transitions kind.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        _raise_invalid_rhs_spec(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_invalid_rhs_spec(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
    )

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

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

    aliases, alias_template_map = _expand_alias_templates(
        meta_parts[0], axes=axes_meta, template_map_seed=state_template_map
    )
    template_map_all = {**state_template_map, **alias_template_map}

    state_set = set(state_expanded)
    d_terms: dict[str, list[str]] = {s: [] for s in state_expanded}
    all_syms = _collect_alias_symbols(aliases, axes=axes_meta)

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

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

    for state_name in state_expanded:
        d_terms.setdefault(state_name, [])

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

    transitions_expanded = _expand_transition_templates(
        transitions_raw,
        axes=axes_meta,
        template_map=template_map_all,
    )

    for idx, tr_map in enumerate(transitions_expanded):
        _apply_transition(
            idx=idx,
            tr=tr_map,
            state_set=state_set,
            all_syms=all_syms,
            d_terms=d_terms,
        )

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

    return NormalizedRhs(
        kind="transitions",
        state_names=tuple(state_expanded),
        equations=tuple(_build_transition_equations(state_expanded, d_terms)),
        aliases=aliases,
        param_names=_sorted_unique(
            sym for sym in all_syms if sym not in state_set and sym not in aliases
        ),
        all_symbols=frozenset(all_syms | set(aliases.keys())),
        meta={**meta, "transitions": transitions_expanded},
    )