ampzip 1.0.0__py3-none-any.whl

ampzip/__init__.py

@@ -0,0 +1,7 @@
+"""ampzip: training-free simplification of spinor-helicity amplitudes.
+
+``simplifier`` holds the CLI, the input check and the backends, ``search`` the best-first
+search, ``verify`` the finite-field test and ``webservice`` the HTTP service. The moves are
+executed by the extension module ``ampzip_engine`` built from ``ampzip/engine``.
+"""
+from ._version import __version__  # noqa: F401

ampzip/__main__.py

@@ -0,0 +1,7 @@
+"""`python -m ampzip` runs the simplifier CLI; `python -m ampzip.webservice` runs the web service."""
+import sys
+
+from .simplifier.cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

ampzip/_version.py

@@ -0,0 +1,26 @@
+"""The package version, from one source.
+
+An installed distribution reports the version in its metadata, taken from `pyproject.toml`
+at build time. A source checkout run without installing reads the same field from the
+`pyproject.toml` next to the package.
+"""
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version as _dist_version
+from pathlib import Path
+
+
+def _read() -> str:
+    try:
+        return _dist_version("ampzip")
+    except PackageNotFoundError:
+        pass
+    try:
+        import tomllib
+        data = tomllib.loads((Path(__file__).resolve().parents[1] / "pyproject.toml").read_text())
+        return str(data["project"]["version"]) + "+source"
+    except Exception:  # noqa: BLE001  no pyproject next to the package
+        return "0+unknown"
+
+
+__version__ = _read()

ampzip/atoms.py

@@ -0,0 +1,74 @@
+"""The bracket atoms and their canonical ordering.
+
+`ab(i, j)` is the angle bracket ⟨i j⟩ and `sb(i, j)` the square bracket [i j], both as
+uninterpreted sympy functions of two leg labels. Every expression the package handles is
+built from these two atoms and `opaque_den`, the engine's marker for a polynomial
+denominator factor, which is created where it is needed.
+
+`canonical_order(expr)` rewrites every bracket to `i < j`, pushing the antisymmetry sign into
+the coefficient: `ab(2, 1) -> -ab(1, 2)`. The engine, the atomizer and the display layer all
+rely on this spelling, so a Mandelstam `s_ij = ⟨ij⟩[ji]` is stored as `-ab(i,j)*sb(i,j)`.
+
+`ATOMS` is the registry of atom classes by name that the input check consults. A new atom
+kind is added by registering it here together with its canonicalisation rule.
+"""
+from __future__ import annotations
+
+from typing import Dict, Type
+
+import sympy as sp
+
+
+def _leg_labels(printer, args):
+    """The leg labels for display. An integer prints as itself, and a symbol such as `p_{3}`
+    prints as its bare label."""
+    return [printer._print(a).replace("p_{", "").replace("}", "") for a in args]
+
+
+class _Bracket(sp.Function):
+    """A two-leg bracket. Subclasses only choose the delimiters."""
+
+    _delims = ("", "")
+
+    # sympy would infer is_number from the integer leg labels and then evaluate bracket
+    # products numerically while printing, which takes seconds on a power of a sum.
+    is_number = False
+
+    def _latex(self, printer, exp=1):
+        i, j = _leg_labels(printer, self.args)
+        left, right = self._delims
+        body = f"{left}{i} {j}{right}"
+        return body if exp == 1 else f"{body}^{{{exp}}}"
+
+
+class ab(_Bracket):
+    """Angle bracket ⟨i j⟩ of legs i and j."""
+
+    _delims = (r"\langle ", r" \rangle")
+
+
+class sb(_Bracket):
+    """Square bracket [i j] of legs i and j."""
+
+    _delims = (r"\left[ ", r" \right]")
+
+
+ATOMS: Dict[str, Type[sp.Function]] = {"ab": ab, "sb": sb}
+
+
+def bracket_classes():
+    """The two bracket classes in the order every caller unpacks them, `ab, sb`."""
+    return ab, sb
+
+
+def canonical_order(expr: sp.Expr) -> sp.Expr:
+    """Rewrite each two-leg bracket with its legs ascending, times -1 when they were swapped.
+
+    It applies to every two-argument function atom in `expr`, which is what the bracket
+    antisymmetry needs, and leaves one-argument atoms such as `opaque_den` alone.
+    """
+    replace = {}
+    for fun in expr.atoms(sp.Function):
+        if len(fun.args) == 2 and fun.args[0] > fun.args[1]:
+            replace[fun] = fun.func(fun.args[1], fun.args[0]) * (-1)
+    return expr.subs(replace)

ampzip/search/__init__.py

@@ -0,0 +1,5 @@
+"""Best-first search over the engine `ampzip_engine`.
+
+`actions` enumerates legal moves on an `Expr` handle and `gbfs` runs a best-first search
+ranked by bracket count. The heuristic is `Expr.factored_bracket_count`.
+"""

ampzip/search/action_space.py

@@ -0,0 +1,101 @@
+"""A fixed-size action encoding and legality mask for a learned policy.
+
+It follows the SFT action spaces of `RL_amplitudes_claude/src/action_space_{4,5,6}pt.py`.
+
+    action = bracket_idx * (n_term_slots * identities_per_bracket)
+           + term_slot   * identities_per_bracket
+           + identity_idx
+
+The sizes n_brackets x n_term_slots x identities_per_bracket are 12x11x11 = 1452 at 4
+points, 20x26x19 = 9880 at 5 and 30x31x32 = 29760 at 6. The per-n `max_terms` defaults
+below set n_term_slots. `term_slot=0` is the whole expression and `term_slot=k>0` the
+engine's k-1'th numerator term. The identity order is that of `actions.identity_menu`,
+Schouten, then momentum, then momentum_sq.
+
+The mask marks the legal action indices, so a policy can exclude illegal actions when
+selecting from its logits. `decode_to_move` maps a selected index to an engine move for
+`Expr.apply`. The sequential search enumerates moves through `actions.legal_moves` and the
+parallel expansion through `Expr.legal_move_list`. Neither uses this fixed encoding.
+
+The default term-slot limits are 10, 25 and 30 for 4, 5 and 6-point inputs. Single-term
+actions exist only for the first `max_terms` numerator terms.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .actions import identity_menu, _term_brackets
+
+# The per-n term-slot cap, with n_term_slots = 1 + max_terms, matching the action-space sizes
+# above. It bounds how many per-term slots the fixed encoding can represent and is the
+# search's own choice, like the term_slot ordering.
+_DEFAULT_MAX_TERMS = {4: 10, 5: 25, 6: 30}
+
+
+def default_max_terms(n: int) -> int:
+    return _DEFAULT_MAX_TERMS.get(n, 10)
+
+
+class ActionSpace:
+    def __init__(self, n: int = 4, max_terms: int | None = None):
+        self.n = n
+        if max_terms is None:
+            max_terms = default_max_terms(n)
+        self.max_terms = max_terms
+        self.n_term_slots = 1 + max_terms
+        self.brackets = self._canonical_brackets(n)          # [(square,i,j)], ab before sb
+        self._bidx = {b: i for i, b in enumerate(self.brackets)}
+        # The menu has the same size for every bracket.
+        self.identities_per_bracket = len(identity_menu(False, 1, 2, n))
+        self.n_actions = len(self.brackets) * self.n_term_slots * self.identities_per_bracket
+
+    @staticmethod
+    def _canonical_brackets(n: int):
+        out = []
+        for square in (False, True):     # ab, then sb
+            for i in range(1, n + 1):
+                for j in range(i + 1, n + 1):
+                    out.append((square, i, j))
+        return out
+
+    def encode(self, bracket_idx: int, term_slot: int, identity_idx: int) -> int:
+        return (bracket_idx * (self.n_term_slots * self.identities_per_bracket)
+                + term_slot * self.identities_per_bracket + identity_idx)
+
+    def decode(self, action: int):
+        idy = action % self.identities_per_bracket
+        rest = action // self.identities_per_bracket
+        term_slot = rest % self.n_term_slots
+        bracket_idx = rest // self.n_term_slots
+        return bracket_idx, term_slot, idy
+
+    def decode_to_move(self, action: int):
+        """The move (square,i,j,term_slot,identity,params) of an action index, for Expr.apply."""
+        bracket_idx, term_slot, idy = self.decode(action)
+        square, i, j = self.brackets[bracket_idx]
+        ident, params = identity_menu(square, i, j, self.n)[idy]
+        return (square, i, j, term_slot, ident, params)
+
+    def mask(self, handle) -> np.ndarray:
+        """The 0/1 legality vector of length n_actions for the current expression.
+
+        Every identity of a bracket is legal once the bracket is present, as in
+        `build_action_mask_v3`.
+        """
+        m = np.zeros(self.n_actions, dtype=np.float32)
+        n_id = self.identities_per_bracket
+        present = handle.distinct_brackets()
+        per_term = [_term_brackets(s) for s in handle.term_monos()]
+        for b in present:
+            b_idx = self._bidx.get(tuple(b))
+            if b_idx is None:
+                continue
+            base = self.encode(b_idx, 0, 0)        # the whole-expression block
+            m[base:base + n_id] = 1.0
+            for k, bset in enumerate(per_term):
+                if k >= self.max_terms:
+                    break
+                if tuple(b) in bset:
+                    base = self.encode(b_idx, k + 1, 0)
+                    m[base:base + n_id] = 1.0
+        return m

ampzip/search/actions.py

@@ -0,0 +1,135 @@
+"""Legal-move enumeration on an `ampzip_engine.Expr` handle.
+
+A move is `(square, i, j, term_slot, identity, params)`. It applies the identity to the
+bracket `bk(i,j)`, square False being ab and True sb, either to the whole expression with
+`term_slot=0` or to one numerator term with `term_slot=k>0`, the engine's k-1'th term.
+
+The per-bracket identity menu reproduces the SFT action spaces of
+`RL_amplitudes_claude/src/action_space_{4,5,6}pt.py::get_identity_types_for_bracket`. Per
+bracket there are C(n-2,2) Schouten, 2(n-1) momentum and 2^(n-2) momentum_sq identities,
+11 at 4 points, 19 at 5 and 32 at 6.
+
+Legality follows `eval_sft_4pt.py::build_action_mask_v3`. A bracket move is legal at
+`term_slot=0` when the bracket is in the numerator and at `term_slot=k>0` when the bracket
+is in that term. There is no identity-specific masking. The full identity menu is offered
+once the bracket is present, and degenerate identities are filtered when the engine applies
+them.
+"""
+from __future__ import annotations
+
+import ast
+import re
+from itertools import combinations
+
+_BRACKET_RE = re.compile(r"\b(ab|sb)\((\d+),\s*(\d+)\)")
+
+
+def identity_menu(square: bool, i: int, j: int, n: int):
+    """The (identity, params) menu for bracket `bk(i,j)`, 11, 19 or 32 entries at n = 4, 5 or 6.
+
+    The order matches the SFT action spaces, Schouten, then momentum, then momentum_sq. The
+    params encoding matches `ampzip_engine`, [pk, pl] for schouten, [pout1, pout2] for
+    momentum and the mom_add list for momentum_sq.
+    """
+    other = [m for m in range(1, n + 1) if m != i and m != j]
+    menu = []
+    # Schouten, every one of the C(n-2, 2) pairs of the other legs in combinations order.
+    for pk, pl in combinations(other, 2):
+        menu.append(("schouten", [pk, pl]))
+    # Momentum, pout1 among the bracket's own legs and pout2 among pout1 and the other legs, 2(n-1) in all.
+    for pout1 in (i, j):
+        for pout2 in [pout1] + other:
+            menu.append(("momentum", [pout1, pout2]))
+    # Momentum_sq, every subset of the other legs as mom_add, 2^(n-2) in all.
+    for size in range(len(other) + 1):
+        for mom_add in combinations(other, size):
+            menu.append(("momentum_sq", list(mom_add)))
+    return menu
+
+
+def _term_brackets(mono_str: str):
+    """The set of (square, i, j) brackets present in a canonical monomial string."""
+    return {
+        (typ == "sb", int(a), int(b))
+        for typ, a, b in _BRACKET_RE.findall(mono_str)
+    }
+
+
+def legal_moves(handle, n: int):
+    """All legal `(square, i, j, term_slot, identity, params)` moves for `handle`.
+
+    It yields `term_slot=0`, the whole expression, for every present numerator bracket, plus
+    `term_slot=k+1` for each numerator term that contains the bracket.
+
+    The engine also enumerates moves through `Expr.legal_moves`, but returning a
+    Python tuple, string and parameter list for every move can make it slower than
+    this function, which reuses each identity's `params` list across term slots.
+    The threaded search avoids these allocations by keeping the move list in Rust
+    with `Expr.legal_move_list` and applying slices through `Expr.apply_slice_coded`,
+    falling back to `Expr.apply_slice` on older engines.
+    """
+    present = handle.distinct_brackets()  # the [(square, i, j)] of the numerator
+    if not present:
+        return []
+    per_term = [_term_brackets(m) for m in handle.term_monos()]
+    moves = []
+    for square, i, j in present:
+        menu = identity_menu(square, i, j, n)
+        # the whole expression
+        for ident, params in menu:
+            moves.append((square, i, j, 0, ident, params))
+        # a single term
+        for k, bset in enumerate(per_term):
+            if (square, i, j) in bset:
+                for ident, params in menu:
+                    moves.append((square, i, j, k + 1, ident, params))
+    return moves
+
+
+# The explicit name of the Python move enumerator, kept for compatibility and for comparison
+# with the Rust implementation.
+legal_moves_py = legal_moves
+
+
+def move_rule(square, i, j, identity, params, n):
+    """The local rewrite a move performs on `bk(i,j)`, as `(lhs_str, rhs_str)` in ab/sb
+    notation. The right-hand side comes from applying the identity to the isolated bracket,
+    so it is the literal substitution the engine makes, signs and all, with no re-derivation
+    of the algebra here. Returns `(lhs, None)` when the engine cannot build it, as with
+    degenerate params.
+    """
+    import ampzip_engine as _E
+    lhs = f"{'sb' if square else 'ab'}({i},{j})"
+    try:
+        rhs = _E.Expr(lhs).apply(square, i, j, identity, list(params), n, term_slot=0).to_str()
+    except Exception:
+        rhs = None
+    return lhs, rhs
+
+
+def annotate_trajectory_rules(traj, n):
+    """Add the ab/sb strings `rule_lhs` and `rule_rhs` to each trajectory step in place,
+    reconstructed from its `bracket`, `identity` and `params`. A step missing fields or whose
+    rule cannot be built is left as it is. It accepts the raw rollout trajectory, where
+    params is a list, and a serialised one, where params is a string.
+    """
+    if not traj:
+        return traj
+    for step in traj:
+        if not isinstance(step, dict) or step.get("rule_lhs"):
+            continue
+        m = _BRACKET_RE.search(str(step.get("bracket", "")))
+        ident, params = step.get("identity"), step.get("params")
+        if not m or ident is None or params is None:
+            continue
+        if not isinstance(params, (list, tuple)):
+            try:
+                params = ast.literal_eval(str(params))
+            except Exception:
+                continue
+        square, ii, jj = (m.group(1) == "sb"), int(m.group(2)), int(m.group(3))
+        lhs, rhs = move_rule(square, ii, jj, ident, list(params), n)
+        step["rule_lhs"] = lhs
+        if rhs is not None:
+            step["rule_rhs"] = rhs
+    return traj