From 194f781a9fa302f3d9cd283bb478de78d0f651a8 Mon Sep 17 00:00:00 2001 From: Joel Eliason Date: Mon, 3 Aug 2026 14:39:26 -0600 Subject: [PATCH] feat: emit Python observables from the model's assignment rules A consumer scoring the model against data has to recompute the derived quantities the data names: totals, lineages, volume fractions. Those are repeatedAssignment rules the C++ emitter already reads, so emitting Python from the same SBML keeps the two in step. Re-deriving them by hand is a second definition that drifts, and on the PDAC model a hand expansion of the aggregate module had already missed phi_collagen, which lives in a different module and is a nonlinear ratio. resolve() closes a requested symbol set over the rules and returns the raw species it needs, which is the output contract for anything standing in for the simulator. It raises on unknown symbols and on cycles rather than dropping them. emit() writes pure arithmetic, so the module needs no array library and runs on floats, numpy arrays or JAX tracers alike. Non-arithmetic operators are refused where the rule is first read, so the error names the rule instead of reporting 'std' and 'max' as missing symbols. On PDAC: 35 symbols close to 30 states, 13 constants and 20 rules, and every aggregate identity reproduces exactly on random states. --- src/qsp_codegen/codegen.py | 3 + src/qsp_codegen/observables.py | 243 +++++++++++++++++++++++++++++++++ tests/test_observables.py | 197 ++++++++++++++++++++++++++ 3 files changed, 443 insertions(+) create mode 100644 src/qsp_codegen/observables.py create mode 100644 tests/test_observables.py diff --git a/src/qsp_codegen/codegen.py b/src/qsp_codegen/codegen.py index 4d1c96c..573ad2d 100644 --- a/src/qsp_codegen/codegen.py +++ b/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() diff --git a/src/qsp_codegen/observables.py b/src/qsp_codegen/observables.py new file mode 100644 index 0000000..8e21516 --- /dev/null +++ b/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"(?!]=?|==") + + +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 diff --git a/tests/test_observables.py b/tests/test_observables.py new file mode 100644 index 0000000..5ce0fc9 --- /dev/null +++ b/tests/test_observables.py @@ -0,0 +1,197 @@ +"""Tests for the Python observables emitter. + +The claim is that a consumer recomputing a model's derived quantities gets the +same numbers the C++ side would, so the tests check the closure (which raw +states a symbol needs), the topological order, and that the emitted module +evaluates to the arithmetic the rules state. +""" +import textwrap +import xml.etree.ElementTree as ET + +import pytest + +from qsp_codegen.codegen import MATH_NS, SBML_NS, SBMLModel +from qsp_codegen.observables import ( + NameCollision, + UnresolvedSymbol, + emit, + resolve, +) + + +class _Model: + """The four attributes resolve() reads. SBML parsing is covered elsewhere.""" + + def __init__(self, rules=(), species=(), parameters=(), compartments=()): + self.assignment_rules = [{"variable_name": n, "expression": e} + for n, e in rules] + self.species = [{"name": n} for n in species] + self.parameters = [{"name": n, "value": v} for n, v in parameters] + self.compartments = [{"name": n, "size": s} for n, s in compartments] + + +MODEL = _Model( + rules=[ + ("total", "a + b + scaled"), + ("scaled", "c * k"), + ("frac", "total / (total + V_box)"), + ("unused", "a + 999.0"), + ], + species=["a", "b", "c", "d"], + parameters=[("k", 2.0), ("spare", 7.0)], + compartments=[("V_box", 5.0)], +) + + +class TestResolve: + def test_closure_keeps_only_what_is_reachable(self): + c = resolve(MODEL, ["total"]) + assert c.states == ("a", "b", "c") # not d + assert c.constants == {"k": 2.0} # not spare + assert "unused" not in c.order + + def test_dependencies_come_before_dependents(self): + c = resolve(MODEL, ["total"]) + assert c.order.index("scaled") < c.order.index("total") + + def test_a_compartment_without_a_rule_is_a_constant(self): + c = resolve(MODEL, ["frac"]) + assert c.constants["V_box"] == 5.0 + assert "V_box" not in c.states + + def test_a_compartment_with_a_rule_is_derived(self): + m = _Model(rules=[("V_box", "a + b")], species=["a", "b"], + compartments=[("V_box", 5.0)]) + c = resolve(m, ["V_box"]) + assert c.states == ("a", "b") + assert "V_box" not in c.constants and "V_box" in c.order + + def test_a_raw_species_needs_no_rule(self): + c = resolve(MODEL, ["d"]) + assert c.states == ("d",) and c.order == () + + def test_derived_names_the_requested_rules(self): + c = resolve(MODEL, ["total", "d"]) + assert c.derived == ("total",) + + def test_an_unknown_symbol_raises_rather_than_vanishing(self): + with pytest.raises(UnresolvedSymbol, match="nope"): + resolve(MODEL, ["total", "nope"]) + + def test_a_cycle_raises(self): + m = _Model(rules=[("x", "y + 1.0"), ("y", "x + 1.0")], species=["a"]) + with pytest.raises(UnresolvedSymbol, match="cyclic"): + resolve(m, ["x"]) + + def test_scientific_notation_is_not_read_as_an_identifier(self): + m = _Model(rules=[("x", "a * 1e-06")], species=["a"], parameters=[("e", 1.0)]) + assert resolve(m, ["x"]).constants == {} + + +def _run(source, states, function="observables"): + ns = {} + exec(compile(source, "", "exec"), ns) + return ns[function](states) + + +class TestEmit: + def test_the_module_computes_the_rules(self): + src = emit(resolve(MODEL, ["total"]), source="t.sbml") + out = _run(src, {"a": 1.0, "b": 2.0, "c": 3.0}) + assert out["total"] == 1.0 + 2.0 + 3.0 * 2.0 + + def test_a_nonlinear_rule_survives_translation(self): + src = emit(resolve(MODEL, ["frac"]), source="t.sbml") + out = _run(src, {"a": 1.0, "b": 2.0, "c": 3.0}) + total = 9.0 + assert out["frac"] == pytest.approx(total / (total + 5.0)) + + def test_dotted_names_sanitise_without_eating_each_other(self): + m = _Model(rules=[("s", "V_T + V_T.C1 + V_T.C1_x")], + species=["V_T.C1", "V_T.C1_x"], compartments=[("V_T", 4.0)]) + src = emit(resolve(m, ["s"]), source="t.sbml") + assert _run(src, {"V_T.C1": 1.0, "V_T.C1_x": 2.0})["s"] == 7.0 + + def test_the_generated_module_declares_its_contract(self): + src = emit(resolve(MODEL, ["total", "d"]), source="t.sbml") + ns = {} + exec(compile(src, "", "exec"), ns) + assert ns["STATES"] == ("a", "b", "c", "d") + assert ns["OBSERVABLES"] == ("d", "total") + assert ns["CONSTANTS"] == {"k": 2.0} + + def test_constants_can_be_overridden_at_the_call(self): + src = emit(resolve(MODEL, ["total"]), source="t.sbml") + ns = {} + exec(compile(src, "", "exec"), ns) + out = ns["observables"]({"a": 0.0, "b": 0.0, "c": 1.0}, {"k": 10.0}) + assert out["total"] == 10.0 + + def test_the_function_can_be_renamed(self): + src = emit(resolve(MODEL, ["d"]), source="t.sbml", function_name="h_r") + assert _run(src, {"d": 1.0}, function="h_r") == {"d": 1.0} + + def test_it_works_on_arrays(self): + np = pytest.importorskip("numpy") + src = emit(resolve(MODEL, ["total"]), source="t.sbml") + out = _run(src, {k: np.full(4, v) for k, v in + (("a", 1.0), ("b", 2.0), ("c", 3.0))}) + assert np.allclose(out["total"], 9.0) + + def test_a_non_arithmetic_rule_is_refused_not_mistranslated(self): + m = _Model(rules=[("x", "std::max(a, b)")], species=["a", "b"]) + with pytest.raises(ValueError, match="non-arithmetic"): + emit(resolve(m, ["x"]), source="t.sbml") + + def test_colliding_names_raise(self): + m = _Model(rules=[("s", "V_T.C + V_T_C")], species=["V_T.C", "V_T_C"]) + with pytest.raises(NameCollision, match="sanitise"): + emit(resolve(m, ["s"]), source="t.sbml") + + +SBML_DOC = textwrap.dedent("""\ + + + + + + + + + + + + + + + + s1 + s2p1 + + + + +""").format(sns=SBML_NS.strip("{}"), mns=MATH_NS.strip("{}")) + + +def test_round_trip_from_real_sbml(tmp_path): + """Parser to closure to module, so the pieces agree on names and shapes. + + A rule's variable is a non-constant parameter, and the parser qualifies a + species by its compartment, so ``CD8`` in the file is ``V_T.CD8`` downstream. + """ + path = tmp_path / "m.sbml" + doc = SBML_DOC.replace( + '', + '' + '') + path.write_text(doc) + + closure = resolve(SBMLModel(str(path)), ["CD8_total"]) + assert closure.states == ("V_T.CD8", "V_T.CD8_exh") + assert closure.constants == {"w": 10.0} + out = _run(emit(closure, source="m.sbml"), + {"V_T.CD8": 3.0, "V_T.CD8_exh": 4.0}) + assert out["CD8_total"] == 3.0 + 4.0 * 10.0