Translate std::pow and std::max in emitted observables

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 {",

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"):