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