diff --git a/src/qsp_codegen/observables.py b/src/qsp_codegen/observables.py index 8e21516..cb97b6f 100644 --- a/src/qsp_codegen/observables.py +++ b/src/qsp_codegen/observables.py @@ -33,15 +33,63 @@ 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. +#: Anything the C++ infix can carry that is not ``+ - * /``, a literal, or one of +#: the calls :func:`_rewrite_calls` normalises away before this ever runs. _NON_ARITHMETIC = re.compile(r"std::\w+|&&|\|\||[<>!]=?|==") +#: The two C++ calls the rules use that have array-agnostic Python equivalents. +_CALL = re.compile(r"std::(pow|max)\s*\(") + +#: Names the emitted module defines itself, so the dependency walk must not chase +#: them into the model and report them missing. +_HELPERS = frozenset({"_maximum"}) + + +def _rewrite_calls(expression: str) -> str: + """``std::pow(a,b)`` -> ``((a) ** (b))``, ``std::max(a,b)`` -> ``_maximum(a,b)``. + + Applied once when the rules are read, so everything downstream sees infix + arithmetic and never has to know these existed. Both forms keep the emitted + module free of imports: ``**`` and ``abs`` are defined on floats, numpy arrays + and JAX tracers alike, which is what lets one generated file serve all three. + + Arguments are split on the top-level comma rather than by regex, since a Hill + exponent is routinely itself a ``std::pow``. + """ + while True: + match = _CALL.search(expression) + if match is None: + return expression + depth, i, comma = 1, match.end(), None + while True: + if i >= len(expression): + raise ValueError(f"unbalanced parentheses in {expression!r}") + char = expression[i] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + break + elif char == "," and depth == 1: + comma = i + i += 1 + if comma is None: + raise ValueError( + f"std::{match.group(1)} takes two arguments in {expression!r}" + ) + left = expression[match.end():comma].strip() + right = expression[comma + 1:i].strip() + rewritten = (f"(({left}) ** ({right}))" if match.group(1) == "pow" + else f"_maximum({left}, {right})") + expression = expression[:match.start()] + rewritten + expression[i + 1:] + 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. + A comparison would otherwise reach the dependency walk and report its operands + as symbols missing from the model. """ found = sorted(set(_NON_ARITHMETIC.findall(expression))) if found: @@ -73,7 +121,8 @@ def _classify(sbml) -> Tuple[Dict[str, str], Set[str], Dict[str, float]]: 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} + rules = {r["variable_name"]: _rewrite_calls(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: @@ -106,7 +155,8 @@ def visit(name: str) -> None: if name in rules: _reject_non_arithmetic(name, rules[name]) for dep in _TOKEN.findall(rules[name]): - visit(dep) + if dep not in _HELPERS: + visit(dep) order.append(name) elif name in species: states.add(name) @@ -175,9 +225,21 @@ def emit(closure: Closure, *, source: str, function_name: str = "observables") - _reject_non_arithmetic(name, closure.expressions[name]) rename = _rename_map(closure) + body = {name: _translate(closure.expressions[name], rename) + for name in closure.order} lines = [ f'"""Observables generated by qsp-codegen from {source}. Do not edit."""', "", + ] + if any("_maximum(" in line for line in body.values()): + lines += [ + "", + "def _maximum(a, b):", + ' """Elementwise ``max`` without an array library. Kinked at a == b."""', + " return (a + b + abs(a - b)) / 2", + "", + ] + lines += [ "STATES = (", *(f" {name!r}," for name in closure.states), ")", @@ -200,7 +262,7 @@ def emit(closure: Closure, *, source: str, function_name: str = "observables") - 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.append(f" {rename[name]} = {body[name]}") lines += [ "", " return {", diff --git a/tests/test_observables.py b/tests/test_observables.py index 5ce0fc9..30533bb 100644 --- a/tests/test_observables.py +++ b/tests/test_observables.py @@ -139,10 +139,34 @@ def test_it_works_on_arrays(self): 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"]) + m = _Model(rules=[("x", "a > b")], species=["a", "b"]) with pytest.raises(ValueError, match="non-arithmetic"): emit(resolve(m, ["x"]), source="t.sbml") + def test_pow_becomes_an_exponent(self): + m = _Model(rules=[("x", "std::pow(a, b)")], species=["a", "b"]) + assert _run(emit(resolve(m, ["x"]), source="t.sbml"), + {"a": 2.0, "b": 3.0})["x"] == 8.0 + + def test_max_is_translated_and_stays_elementwise(self): + np = pytest.importorskip("numpy") + m = _Model(rules=[("x", "std::max(a, b)")], species=["a", "b"]) + out = _run(emit(resolve(m, ["x"]), source="t.sbml"), + {"a": np.array([1.0, 5.0]), "b": np.array([4.0, 2.0])}) + assert np.allclose(out["x"], [4.0, 5.0]) + + def test_a_nested_call_resolves_from_the_inside(self): + """A Hill exponent is routinely itself a power, so one pass is not enough.""" + m = _Model(rules=[("x", "std::pow(std::max(a, b), 2.0)")], + species=["a", "b"]) + assert _run(emit(resolve(m, ["x"]), source="t.sbml"), + {"a": 1.0, "b": 3.0})["x"] == 9.0 + + def test_a_translated_call_does_not_leak_into_the_closure(self): + """``_maximum`` is the emitter's own name; chasing it would report it missing.""" + m = _Model(rules=[("x", "std::max(a, b)")], species=["a", "b"]) + assert set(resolve(m, ["x"]).states) == {"a", "b"} + 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"):