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

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 ...

CompiledRhs(state_names, param_names, eval_fn, meta=(lambda: MappingProxyType({}))()) dataclass

Container for a compiled RHS evaluation function.

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[[float64, Float64Array], 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[[np.float64, Float64Array], 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: np.float64, y: Float64Array) -> Float64Array:
        return self.eval_fn(t, y, **params_dict)

    return rhs

EvalFn

Bases: Protocol

Callable RHS evaluator supporting runtime parameter kwargs.

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

Normalized RHS representation suitable for compilation/execution.

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}]"

compile_rhs(rhs)

Compile a normalized RHS into a runnable evaluation function.

Parameters:

Name	Type	Description	Default
rhs	NormalizedRhs	Normalized RHS produced by op_system.specs.normalize_rhs.	required

Returns:

Type	Description
CompiledRhs	A CompiledRhs containing an eval_fn(t, y, **params) -> dydt.

Source code in src/op_system/compile.py

def compile_rhs(rhs: NormalizedRhs) -> CompiledRhs:
    """Compile a normalized RHS into a runnable evaluation function.

    Args:
        rhs: Normalized RHS produced by `op_system.specs.normalize_rhs`.

    Returns:
        A `CompiledRhs` containing an `eval_fn(t, y, **params) -> dydt`.
    """
    if rhs.kind not in {"expr", "transitions"}:
        _raise_unsupported_feature(
            feature=f"rhs.kind={rhs.kind}",
            detail="Only 'expr' and 'transitions' are supported in v1.",
        )

    eval_fn = _make_eval_fn(
        state_names=rhs.state_names,
        aliases=rhs.aliases,
        equations=rhs.equations,
    )

    return CompiledRhs(
        state_names=rhs.state_names,
        param_names=rhs.param_names,
        eval_fn=eval_fn,
        meta=rhs.meta,
    )

compile_spec(spec)

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

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

Parameters:

Name	Type	Description	Default
spec	dict[str, object]	Raw RHS specification mapping (YAML/JSON friendly).	required

Returns:

Name	Type	Description
CompiledRhs	CompiledRhs	Runnable RHS callable container.

Source code in src/op_system/__init__.py

def compile_spec(spec: dict[str, object]) -> CompiledRhs:  # noqa: RUF067
    """
    Validate, normalize, and compile a RHS specification in one call.

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

    Args:
        spec: Raw RHS specification mapping (YAML/JSON friendly).

    Returns:
        CompiledRhs: Runnable RHS callable container.
    """
    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
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},
    )