rootfig 0.1.0__py3-none-any.whl

rootfig/errors.py

@@ -0,0 +1,73 @@
+"""Exception and warning types raised by rootfig.
+
+All exceptions derive from :class:`RootfigError` so callers can catch
+everything from the library with a single ``except`` clause. The more
+specific subclasses also inherit from the matching built-in exception
+(``ValueError``, ``KeyError``-like ``LookupError``) so generic handlers
+keep working.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class RootfigError(Exception):
+    """Base class for all errors raised by rootfig."""
+
+
+class ExpressionError(RootfigError, ValueError):
+    """An expression string could not be parsed, is not allowed, or failed to evaluate."""
+
+
+class MissingBranchError(ExpressionError, LookupError):
+    """An expression refers to a name that is not a branch, alias, function, or constant."""
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        available: Sequence[str] = (),
+        suggestions: Sequence[str] = (),
+        context: str | None = None,
+    ) -> None:
+        self.name = name
+        self.available = tuple(available)
+        self.suggestions = tuple(suggestions)
+        message = f"unknown name {name!r}"
+        if context:
+            message += f" in {context}"
+        if self.suggestions:
+            message += ". Did you mean " + ", ".join(repr(s) for s in self.suggestions) + "?"
+        elif self.available:
+            shown = ", ".join(repr(a) for a in self.available[:12])
+            more = "" if len(self.available) <= 12 else f", ... ({len(self.available)} total)"
+            message += f". Available names: {shown}{more}"
+        super().__init__(message)
+
+
+class SelectionError(RootfigError, ValueError):
+    """Variable, selection, and weight arrays have incompatible structures."""
+
+
+class IncompatibleWeightError(SelectionError):
+    """A weight array cannot be broadcast to the variable being histogrammed."""
+
+
+class BinningError(RootfigError, ValueError):
+    """A binning or range specification is invalid or cannot be inferred."""
+
+
+class SourceError(RootfigError):
+    """A file, tree, or in-memory data source cannot be opened or interpreted."""
+
+
+class LuminosityError(RootfigError, ValueError):
+    """A sample cannot be scaled to a luminosity (no luminosity, cross section or event count)."""
+
+
+class RootfigWarning(UserWarning):
+    """Base class for warnings emitted by rootfig (dropped values, empty selections, ...)."""

rootfig/expressions/__init__.py

@@ -0,0 +1,34 @@
+"""Expression strings over branches.
+
+rootfig expressions use ordinary Python syntax and are evaluated with NumPy
+and Awkward Array semantics, so ``Muon_pt > 20`` produces a boolean array
+with the same (possibly jagged) structure as ``Muon_pt``.
+
+Supported syntax
+----------------
+* Names refer to branches (``Muon_pt``). Branch names that are not valid
+  Python identifiers are written in backticks: ```jet1_b-tag` > 0.5``.
+* Arithmetic ``+ - * / // % **`` and unary ``-``.
+* Comparisons ``== != < <= > >=`` including chained comparisons
+  (``20 < Muon_pt < 100``), which are rewritten element-wise.
+* Boolean combinations with ``&``, ``|``, ``~`` **or** with ``and``, ``or``,
+  ``not``; the keyword forms are rewritten to the element-wise operators.
+* Indexing and slicing (``Muon_pt[:, 0]``).
+* A fixed set of functions (see :data:`FUNCTIONS`): element-wise NumPy
+  functions such as ``abs``, ``sqrt``, ``log``, ``where``, and per-event
+  reductions over jagged branches: ``count``, ``sum``, ``min``, ``max``,
+  ``mean``, ``any``, ``all``, ``first``.
+* Constants ``pi``, ``e``, ``inf``, ``nan``, ``True``, ``False``.
+
+Everything else (attribute access, lambdas, comprehensions, string literals,
+calls to unknown functions) is rejected at parse time with
+:class:`~rootfig.errors.ExpressionError`.
+
+The public entry points are :func:`parse`, :class:`Expression`, and
+:func:`evaluate`.
+"""
+
+from rootfig.expressions.functions import CONSTANTS, FUNCTIONS
+from rootfig.expressions.parser import Expression, ExpressionLike, evaluate, parse
+
+__all__ = ["CONSTANTS", "FUNCTIONS", "Expression", "ExpressionLike", "evaluate", "parse"]

rootfig/expressions/functions.py

@@ -0,0 +1,151 @@
+"""Functions and constants available inside expressions."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Mapping
+from types import MappingProxyType
+from typing import Any, Final
+
+import awkward as ak
+import numpy as np
+
+from rootfig.errors import ExpressionError
+
+__all__ = ["CONSTANTS", "FUNCTIONS"]
+
+# --------------------------------------------------------------------------------------
+
+
+def _require_jagged(name: str, array: Any) -> Any:
+    if not isinstance(array, ak.Array):
+        array = ak.Array(array)
+    if array.layout.purelist_depth < 2:
+        msg = (
+            f"{name}() reduces over the objects in each event and needs a jagged "
+            f"(per-object) array, but got a flat (per-event) array"
+        )
+        raise ExpressionError(msg)
+    return array
+
+
+def _reduction(name: str, fn: Callable[..., Any]) -> Callable[[Any], Any]:
+    def reduce(array: Any) -> Any:
+        return fn(_require_jagged(name, array), axis=-1)
+
+    reduce.__name__ = name
+    return reduce
+
+
+def _clip(x: Any, lo: Any, hi: Any) -> Any:
+    return np.minimum(np.maximum(x, lo), hi)
+
+
+def _first(array: Any) -> Any:
+    return ak.firsts(_require_jagged("first", array), axis=-1)
+
+
+# -- kinematics from Cartesian components (EDM4hep stores momentum.x/y/z and energy) ----
+
+
+def _pt(px: Any, py: Any) -> Any:
+    return np.hypot(px, py)
+
+
+def _p(px: Any, py: Any, pz: Any) -> Any:
+    return np.sqrt(px * px + py * py + pz * pz)
+
+
+def _theta(px: Any, py: Any, pz: Any) -> Any:
+    return np.arctan2(np.hypot(px, py), pz)
+
+
+def _costheta(px: Any, py: Any, pz: Any) -> Any:
+    with np.errstate(divide="ignore", invalid="ignore"):  # nan for a zero momentum
+        return pz / _p(px, py, pz)
+
+
+def _eta(px: Any, py: Any, pz: Any) -> Any:
+    with np.errstate(divide="ignore"):  # +-inf along the beam axis
+        return np.arcsinh(pz / np.hypot(px, py))
+
+
+def _phi(px: Any, py: Any) -> Any:
+    return np.arctan2(py, px)
+
+
+def _mass(energy: Any, px: Any, py: Any, pz: Any) -> Any:
+    return np.sqrt(np.maximum(energy * energy - (px * px + py * py + pz * pz), 0.0))
+
+
+FUNCTIONS: Final[Mapping[str, Callable[..., Any]]] = MappingProxyType(
+    {
+        # element-wise (NumPy ufuncs dispatch to Awkward automatically)
+        "abs": np.abs,
+        "sqrt": np.sqrt,
+        "cbrt": np.cbrt,
+        "exp": np.exp,
+        "expm1": np.expm1,
+        "log": np.log,
+        "log10": np.log10,
+        "log2": np.log2,
+        "log1p": np.log1p,
+        "power": np.power,
+        "hypot": np.hypot,
+        "sin": np.sin,
+        "cos": np.cos,
+        "tan": np.tan,
+        "arcsin": np.arcsin,
+        "arccos": np.arccos,
+        "arctan": np.arctan,
+        "arctan2": np.arctan2,
+        "sinh": np.sinh,
+        "cosh": np.cosh,
+        "tanh": np.tanh,
+        "arcsinh": np.arcsinh,
+        "arccosh": np.arccosh,
+        "arctanh": np.arctanh,
+        "deg2rad": np.deg2rad,
+        "rad2deg": np.rad2deg,
+        "floor": np.floor,
+        "ceil": np.ceil,
+        "round": np.rint,
+        "trunc": np.trunc,
+        "sign": np.sign,
+        "minimum": np.minimum,
+        "maximum": np.maximum,
+        "clip": _clip,
+        "isnan": np.isnan,
+        "isinf": np.isinf,
+        "isfinite": np.isfinite,
+        "where": ak.where,
+        # per-event reductions over jagged branches
+        "count": _reduction("count", ak.num),
+        "len": _reduction("len", ak.num),
+        "sum": _reduction("sum", ak.sum),
+        "prod": _reduction("prod", ak.prod),
+        "min": _reduction("min", ak.min),
+        "max": _reduction("max", ak.max),
+        "mean": _reduction("mean", ak.mean),
+        "std": _reduction("std", ak.std),
+        "any": _reduction("any", ak.any),
+        "all": _reduction("all", ak.all),
+        "argmin": _reduction("argmin", ak.argmin),
+        "argmax": _reduction("argmax", ak.argmax),
+        "first": _first,
+        # kinematics from Cartesian components
+        "pt": _pt,
+        "p": _p,
+        "theta": _theta,
+        "costheta": _costheta,
+        "eta": _eta,
+        "phi": _phi,
+        "mass": _mass,
+    }
+)
+"""Functions callable from expressions, keyed by the name used in the expression."""
+
+CONSTANTS: Final[Mapping[str, float]] = MappingProxyType(
+    {"pi": math.pi, "e": math.e, "inf": math.inf, "nan": math.nan}
+)
+"""Named constants available in expressions when no branch of that name exists."""