Emit Python observables from the model's assignment rules

src/qsp_codegen/codegen.py

@@ -2600,6 +2600,9 @@ def main(argv: Optional[List[str]] = None) -> int:
     from .verify import add_subparser as _add_verify
     _add_verify(sub)
 
+    from .observables import add_subparser as _add_observables
+    _add_observables(sub)
+
     args = parser.parse_args(argv)
     if not getattr(args, "command", None):
         parser.print_help()

src/qsp_codegen/observables.py

@@ -0,0 +1,243 @@
+"""Python observables from an SBML model's assignment rules.
+
+A consumer that scores a model against data has to recompute the derived
+quantities the data names: totals, lineages, volume fractions. Those are
+``repeatedAssignment`` rules in the model, and the C++ emitter already reads
+them, so emitting Python from the same SBML keeps the two in step by
+construction. Re-deriving them by hand is a second definition that drifts.
+
+:func:`resolve` answers the other half: which raw species a set of requested
+symbols actually needs. That closure is the output contract for anything
+standing in for the simulator.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, List, Sequence, Set, Tuple
+
+__all__ = ["Closure", "resolve", "emit", "UnresolvedSymbol", "NameCollision"]
+
+#: Identifiers as the C++ emitter writes them, dots and all: ``V_T.CD8_exh``.
+#: The lookbehind keeps the ``e`` of a literal like ``1e-06`` from reading as one.
+_TOKEN = re.compile(r"(?<![A-Za-z0-9_.])[A-Za-z_][A-Za-z0-9_.]*")
+
+
+class UnresolvedSymbol(LookupError):
+    """A symbol that is neither a rule, a species, a compartment nor a parameter."""
+
+
+class NameCollision(ValueError):
+    """Two model symbols that sanitise to the same Python identifier."""
+
+
+#: Anything the C++ infix can carry that is not ``+ - * /`` or a literal.
+_NON_ARITHMETIC = re.compile(r"std::\w+|&&|\|\||[<>!]=?|==")
+
+
+def _reject_non_arithmetic(name: str, expression: str) -> None:
+    """Fail on the rule, not later on its tokens.
+
+    ``std::max(a, b)`` would otherwise reach the dependency walk and report
+    ``std`` and ``max`` as symbols missing from the model.
+    """
+    found = sorted(set(_NON_ARITHMETIC.findall(expression)))
+    if found:
+        raise ValueError(
+            f"rule {name!r} uses non-arithmetic operators this emitter cannot "
+            f"translate: " + ", ".join(found)
+        )
+
+
+@dataclass(frozen=True)
+class Closure:
+    """What a set of requested symbols needs, and in what order to compute it."""
+
+    requested: Tuple[str, ...]
+    states: Tuple[str, ...]        # raw species: the simulator's output contract
+    constants: Dict[str, float]    # fixed parameters and compartment sizes
+    order: Tuple[str, ...]         # rules to evaluate, dependencies first
+    expressions: Dict[str, str]    # rule name -> infix, as the C++ emitter sees it
+
+    @property
+    def derived(self) -> Tuple[str, ...]:
+        """Requested symbols that are rules rather than raw states."""
+        return tuple(s for s in self.requested if s in self.expressions)
+
+
+def _classify(sbml) -> Tuple[Dict[str, str], Set[str], Dict[str, float]]:
+    """Rules, species, and everything with a fixed value.
+
+    A compartment carrying a rule is derived, not a constant: ``V_T`` is computed
+    from the species it holds, and its listed size is only an initial value.
+    """
+    rules = {r["variable_name"]: r["expression"] for r in sbml.assignment_rules}
+    species = {s["name"] for s in sbml.species}
+    fixed: Dict[str, float] = {p["name"]: p["value"] for p in sbml.parameters}
+    for c in sbml.compartments:
+        if c["name"] not in rules:
+            fixed[c["name"]] = c["size"]
+    return rules, species, fixed
+
+
+def resolve(sbml, symbols: Iterable[str]) -> Closure:
+    """Close ``symbols`` over the model's assignment rules.
+
+    Raises :class:`UnresolvedSymbol` rather than dropping anything: a symbol that
+    silently vanishes is a readout that silently stops being computed.
+    """
+    rules, species, fixed = _classify(sbml)
+    requested = tuple(sorted(set(symbols)))
+
+    states: Set[str] = set()
+    constants: Dict[str, float] = {}
+    order: List[str] = []
+    state: Dict[str, int] = {}     # 0 = visiting, 1 = done
+    unresolved: Set[str] = set()
+
+    def visit(name: str) -> None:
+        if state.get(name) == 1:
+            return
+        if state.get(name) == 0:
+            raise UnresolvedSymbol(f"assignment rules are cyclic through {name!r}")
+        state[name] = 0
+        if name in rules:
+            _reject_non_arithmetic(name, rules[name])
+            for dep in _TOKEN.findall(rules[name]):
+                visit(dep)
+            order.append(name)
+        elif name in species:
+            states.add(name)
+        elif name in fixed:
+            constants[name] = fixed[name]
+        else:
+            unresolved.add(name)
+        state[name] = 1
+
+    for symbol in requested:
+        visit(symbol)
+
+    if unresolved:
+        raise UnresolvedSymbol(
+            f"{len(unresolved)} symbol(s) are not in the model: "
+            + ", ".join(sorted(unresolved))
+        )
+
+    return Closure(
+        requested=requested,
+        states=tuple(sorted(states)),
+        constants=dict(sorted(constants.items())),
+        order=tuple(order),
+        expressions={n: rules[n] for n in order},
+    )
+
+
+def _sanitise(name: str) -> str:
+    return name.replace(".", "_")
+
+
+def _rename_map(closure: Closure) -> Dict[str, str]:
+    """Model symbol -> Python identifier, refusing to collapse two into one."""
+    out: Dict[str, str] = {}
+    seen: Dict[str, str] = {}
+    for name in (*closure.states, *closure.constants, *closure.order):
+        ident = _sanitise(name)
+        if ident in seen and seen[ident] != name:
+            raise NameCollision(
+                f"{name!r} and {seen[ident]!r} both sanitise to {ident!r}"
+            )
+        if not ident.isidentifier():
+            raise NameCollision(f"{name!r} sanitises to {ident!r}, not an identifier")
+        seen[ident] = name
+        out[name] = ident
+    return out
+
+
+def _translate(expression: str, rename: Dict[str, str]) -> str:
+    """Rewrite one infix expression into Python identifiers.
+
+    A single pass over the original string, so a substituted name can never be
+    rewritten again by a later one.
+    """
+    return _TOKEN.sub(lambda m: rename.get(m.group(0), m.group(0)), expression)
+
+
+def emit(closure: Closure, *, source: str, function_name: str = "observables") -> str:
+    """Python source for a module computing ``closure.requested`` from raw states.
+
+    Pure arithmetic, so the result works on floats, numpy arrays or JAX tracers
+    without importing an array library. Rechecked here so a hand-built closure
+    cannot smuggle C++ into a .py file.
+    """
+    for name in closure.order:
+        _reject_non_arithmetic(name, closure.expressions[name])
+
+    rename = _rename_map(closure)
+    lines = [
+        f'"""Observables generated by qsp-codegen from {source}. Do not edit."""',
+        "",
+        "STATES = (",
+        *(f"    {name!r}," for name in closure.states),
+        ")",
+        "",
+        "CONSTANTS = {",
+        *(f"    {name!r}: {value!r}," for name, value in closure.constants.items()),
+        "}",
+        "",
+        "OBSERVABLES = (",
+        *(f"    {name!r}," for name in closure.requested),
+        ")",
+        "",
+        "",
+        f"def {function_name}(states, constants=CONSTANTS):",
+        f'    """Every symbol in OBSERVABLES, from a mapping over STATES."""',
+    ]
+    for name in closure.states:
+        lines.append(f"    {rename[name]} = states[{name!r}]")
+    for name in closure.constants:
+        lines.append(f"    {rename[name]} = constants[{name!r}]")
+    lines.append("")
+    for name in closure.order:
+        lines.append(f"    {rename[name]} = {_translate(closure.expressions[name], rename)}")
+    lines += [
+        "",
+        "    return {",
+        *(f"        {name!r}: {rename[name]}," for name in closure.requested),
+        "    }",
+        "",
+    ]
+    return "\n".join(lines)
+
+
+def add_subparser(subparsers) -> None:
+    """Register the ``emit-observables`` subcommand."""
+    p = subparsers.add_parser(
+        "emit-observables",
+        help="Emit Python for the model's derived quantities, from its SBML rules.",
+    )
+    p.add_argument("--sbml", required=True, type=Path, help="SBML model file.")
+    p.add_argument("--symbols", required=True, type=Path,
+                   help="File of symbol names to emit, one per line; # comments ok.")
+    p.add_argument("--out", required=True, type=Path, help="Python file to write.")
+    p.add_argument("--function-name", default="observables")
+    p.set_defaults(_handler=_handle)
+
+
+def _read_symbols(path: Path) -> List[str]:
+    lines = (line.split("#", 1)[0].strip() for line in path.read_text().splitlines())
+    return [line for line in lines if line]
+
+
+def _handle(args) -> int:
+    from .codegen import SBMLModel
+
+    closure = resolve(SBMLModel(str(args.sbml)), _read_symbols(args.symbols))
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(emit(closure, source=args.sbml.name,
+                             function_name=args.function_name))
+    print(f"{args.out}: {len(closure.requested)} observables, "
+          f"{len(closure.states)} states, {len(closure.constants)} constants, "
+          f"{len(closure.order)} rules")
+    return 0