rootfig 0.1.0__py3-none-any.whl

rootfig/model/binning.py

@@ -0,0 +1,237 @@
+"""Binning specifications and their resolution into ``hist`` axes."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias
+
+import hist
+import numpy as np
+
+from rootfig.errors import BinningError
+
+if TYPE_CHECKING:
+    from rootfig.model.variables import Variable
+
+__all__ = [
+    "ROBUST_THRESHOLD",
+    "Axis",
+    "Bins",
+    "RangeSpec",
+    "auto_range",
+    "log_bins",
+    "resolve_axis",
+    "validate_bins",
+]
+
+Axis: TypeAlias = hist.axis.Regular | hist.axis.Variable
+"""The histogram axis types rootfig produces."""
+
+Bins: TypeAlias = int | tuple[int, float, float] | Sequence[float] | np.ndarray | Axis
+"""Accepted binning specifications.
+
+* ``int`` - number of equal-width bins; the range is taken from ``range`` or
+  inferred from the data.
+* ``(n, low, high)`` - ``n`` equal-width bins between ``low`` and ``high``.
+* a sequence of edges - variable-width bins.
+* a :class:`hist.axis.Regular` or :class:`hist.axis.Variable` - used as is.
+"""
+
+RangeSpec: TypeAlias = tuple[float, float] | Literal["auto", "robust"] | None
+"""How to choose the histogram range when ``bins`` is an integer.
+
+* ``(low, high)`` - explicit.
+* ``"auto"`` (default) - the finite minimum and maximum over all samples.
+* ``"robust"`` - like ``"auto"`` but ignoring outliers far from the bulk of
+  the data (see :func:`auto_range`), padded by 5 percent; useful when
+  sentinel values such as ``-999`` would otherwise dominate the range.
+"""
+
+
+# --------------------------------------------------------------------------------------
+
+
+def log_bins(n: int, low: float, high: float) -> np.ndarray:
+    """Return ``n + 1`` logarithmically spaced bin edges between ``low`` and ``high``.
+
+    Examples
+    --------
+    >>> log_bins(3, 1, 1000).tolist()
+    [1.0, 10.0, 100.0, 1000.0]
+    """
+    if n < 1:
+        msg = f"number of bins must be positive, got {n}"
+        raise BinningError(msg)
+    if not (0 < low < high):
+        msg = f"log bins need 0 < low < high, got low={low}, high={high}"
+        raise BinningError(msg)
+    return np.geomspace(low, high, n + 1)
+
+
+def validate_bins(bins: Bins, range_: RangeSpec) -> None:
+    """Raise :class:`BinningError` if ``bins``/``range_`` are not a valid specification."""
+    if isinstance(bins, hist.axis.Regular | hist.axis.Variable):
+        return
+    if isinstance(bins, bool):
+        msg = "bins must be an int, (n, low, high), edges, or a hist axis, got a bool"
+        raise BinningError(msg)
+    if isinstance(bins, int):
+        if bins < 1:
+            msg = f"number of bins must be positive, got {bins}"
+            raise BinningError(msg)
+        if isinstance(range_, tuple):
+            _validate_range(range_)
+        elif range_ not in ("auto", "robust", None):
+            msg = f"range must be (low, high), 'auto', or 'robust', got {range_!r}"
+            raise BinningError(msg)
+        return
+    if isinstance(bins, tuple) and len(bins) == 3 and isinstance(bins[0], int):
+        if bins[0] < 1:
+            msg = f"number of bins must be positive, got {bins[0]}"
+            raise BinningError(msg)
+        _validate_range((bins[1], bins[2]))
+        return
+    edges = _edges_from(bins)
+    if len(edges) < 2:
+        msg = "bin edges must contain at least two values"
+        raise BinningError(msg)
+    if not np.all(np.diff(edges) > 0):
+        msg = "bin edges must be strictly increasing"
+        raise BinningError(msg)
+    if not np.all(np.isfinite(edges)):
+        msg = "bin edges must be finite"
+        raise BinningError(msg)
+
+
+def _validate_range(range_: tuple[float, float]) -> None:
+    if len(range_) != 2:  # runtime guard for untyped callers
+        msg = f"range must have two values (low, high), got {range_!r}"  # type: ignore[unreachable]
+        raise BinningError(msg)
+    low, high = range_
+    if not (math.isfinite(low) and math.isfinite(high)):
+        msg = f"range must be finite, got {range_!r}"
+        raise BinningError(msg)
+    if not low < high:
+        msg = f"range must satisfy low < high, got {range_!r}"
+        raise BinningError(msg)
+
+
+def _edges_from(bins: Any) -> np.ndarray:
+    try:
+        edges = np.asarray(bins, dtype=float)
+    except (TypeError, ValueError) as exc:
+        msg = f"cannot interpret {bins!r} as bin edges"
+        raise BinningError(msg) from exc
+    if edges.ndim != 1:
+        msg = f"bin edges must be one-dimensional, got shape {edges.shape}"
+        raise BinningError(msg)
+    return edges
+
+
+ROBUST_THRESHOLD = 10.0
+"""Modified z-score (in units of the median absolute deviation) beyond which values are
+ignored by ``range="robust"``. Generous enough to keep physical tails, strict enough to
+drop sentinel values such as ``-999``."""
+
+
+def auto_range(
+    arrays: Sequence[np.ndarray],
+    *,
+    mode: Literal["auto", "robust"] = "auto",
+) -> tuple[float, float]:
+    """Choose a histogram range covering the finite values of all ``arrays``.
+
+    ``mode="auto"`` uses the overall minimum and maximum, with the upper edge
+    nudged up so the maximum value lands inside the last bin. ``mode="robust"``
+    first discards outliers whose modified z-score (``0.6745 * |x - median| /
+    MAD``) exceeds :data:`ROBUST_THRESHOLD`, which removes sentinel values such
+    as ``-999`` without trimming ordinary tails, and pads the result by 5
+    percent of its span. A degenerate range (all values equal) is widened
+    symmetrically; if there are no finite values at all, ``(0.0, 1.0)`` is
+    returned.
+    """
+    values = [np.asarray(a, dtype=float).ravel() for a in arrays if len(a)]
+    finite = [v[np.isfinite(v)] for v in values]
+    finite = [v for v in finite if v.size]
+    if not finite:
+        return (0.0, 1.0)
+    combined = np.concatenate(finite)
+    if mode == "robust":
+        kept = _reject_outliers(combined)
+        low, high = float(kept.min()), float(kept.max())
+        pad = 0.05 * (high - low)
+        low, high = low - pad, high + pad
+    else:
+        low, high = float(combined.min()), float(combined.max())
+        span = high - low
+        high = high + (span * 1e-3 if span > 0 else 0.0)
+    if not high > low:
+        width = abs(low) * 0.1 if low != 0 else 0.5
+        low, high = low - width, high + width
+    return (low, high)
+
+
+def _reject_outliers(values: np.ndarray, threshold: float = ROBUST_THRESHOLD) -> np.ndarray:
+    median = np.median(values)
+    deviation = np.abs(values - median)
+    mad = np.median(deviation)
+    if mad <= 0:
+        # More than half the values are identical (e.g. all zero): fall back to the
+        # mean absolute deviation, and if that is zero too, keep everything.
+        mad = float(deviation.mean())
+        if mad <= 0:
+            return values
+    score = 0.6745 * deviation / mad
+    kept = values[score <= threshold]
+    return kept if kept.size else values
+
+
+def resolve_axis(
+    variable: Variable,
+    data: Sequence[np.ndarray] = (),
+    *,
+    name: str = "x",
+) -> Axis:
+    """Turn a :class:`Variable`'s binning into a concrete ``hist`` axis.
+
+    Parameters
+    ----------
+    variable
+        The variable whose ``bins`` and ``range`` are interpreted.
+    data
+        Flat arrays of values (one per sample) used to infer the range when
+        ``bins`` is an integer and ``range`` is ``"auto"``/``"robust"``.
+    name
+        Axis name stored in the histogram.
+
+    Raises
+    ------
+    BinningError
+        If the range must be inferred but no data was given.
+    """
+    bins = variable.bins
+    label = variable.axis_label
+    if isinstance(bins, hist.axis.Regular | hist.axis.Variable):
+        axis = bins
+        if not axis.label:
+            axis.label = label
+        return axis
+    if isinstance(bins, int) and not isinstance(bins, bool):
+        if isinstance(variable.range, tuple):
+            low, high = variable.range
+        else:
+            if not data:
+                msg = (
+                    f"cannot infer a range for {variable.expression!r} without data; "
+                    "pass range=(low, high) or bins=(n, low, high)"
+                )
+                raise BinningError(msg)
+            mode: Literal["auto", "robust"] = "robust" if variable.range == "robust" else "auto"
+            low, high = auto_range(data, mode=mode)
+        return hist.axis.Regular(bins, low, high, name=name, label=label)
+    if isinstance(bins, tuple) and len(bins) == 3 and isinstance(bins[0], int):
+        n, low, high = bins
+        return hist.axis.Regular(int(n), float(low), float(high), name=name, label=label)
+    edges = _edges_from(bins)
+    return hist.axis.Variable(edges, name=name, label=label)

rootfig/model/cuts.py

@@ -0,0 +1,85 @@
+"""Composable selection expressions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TypeAlias
+
+from rootfig.expressions import Expression, parse
+
+__all__ = ["Cut", "CutLike", "as_cut"]
+
+# --------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Cut:
+    """A selection expression that composes with ``&``, ``|`` and ``~``.
+
+    Plain strings are accepted wherever a ``Cut`` is; the class exists so
+    selections can be built up and named in analysis scripts::
+
+        base = rf.Cut("nMuon >= 2", label="2 muons")
+        signal_region = base & "abs(Muon_eta) < 2.4" & ~rf.Cut("has_bjet")
+
+    Parameters
+    ----------
+    expression
+        A rootfig expression producing booleans (see :mod:`rootfig.expressions`).
+    label
+        Optional short description for legends or logging.
+    """
+
+    expression: str
+    label: str | None = None
+
+    def __post_init__(self) -> None:
+        parse(self.expression)  # validate early
+
+    def __str__(self) -> str:
+        return self.expression
+
+    def parsed(self) -> Expression:
+        """Return the parsed :class:`~rootfig.expressions.Expression`."""
+        return parse(self.expression)
+
+    def _combine(self, other: CutLike, op: str) -> Cut:
+        rhs = Cut(other) if isinstance(other, str) else other
+        label = None
+        if self.label and rhs.label:
+            label = f"{self.label} {op} {rhs.label}"
+        return Cut(f"({self.expression}) {op} ({rhs.expression})", label=label)
+
+    def __and__(self, other: CutLike) -> Cut:
+        return self._combine(other, "&")
+
+    def __rand__(self, other: CutLike) -> Cut:
+        return (
+            Cut(other)._combine(self, "&") if isinstance(other, str) else other._combine(self, "&")
+        )
+
+    def __or__(self, other: CutLike) -> Cut:
+        return self._combine(other, "|")
+
+    def __ror__(self, other: CutLike) -> Cut:
+        return (
+            Cut(other)._combine(self, "|") if isinstance(other, str) else other._combine(self, "|")
+        )
+
+    def __invert__(self) -> Cut:
+        label = f"not {self.label}" if self.label else None
+        return Cut(f"~({self.expression})", label=label)
+
+
+CutLike: TypeAlias = str | Cut
+"""Anything accepted as a selection."""
+
+
+def as_cut(selection: CutLike | None) -> Cut | None:
+    """Coerce a string or ``Cut`` to a ``Cut`` (``None`` passes through)."""
+    if selection is None or isinstance(selection, Cut):
+        return selection
+    if isinstance(selection, str):
+        return Cut(selection)
+    msg = f"selection must be a string or Cut, got {type(selection).__name__}"  # type: ignore[unreachable]
+    raise TypeError(msg)

rootfig/model/samples.py

@@ -0,0 +1,284 @@
+"""The :class:`Sample` description of a dataset."""
+
+from __future__ import annotations
+
+import copy
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, fields
+from os import PathLike
+from typing import Any, Literal, TypeAlias
+
+import awkward as ak
+
+from rootfig.errors import LuminosityError, SourceError
+from rootfig.io import FileSource, Source, as_source
+from rootfig.model.cuts import Cut, CutLike, as_cut
+from rootfig.model.units import cross_section_pb, luminosity_fb
+
+__all__ = ["HistType", "Sample", "SampleLike", "as_samples"]
+
+HistType: TypeAlias = Literal["step", "fill", "errorbar", "band"]
+"""How a histogram is drawn: outline, filled area, points with error bars, or a band."""
+
+
+@dataclass(frozen=True, init=False)
+class Sample:
+    """A dataset: where it comes from, how it is labelled, and how it is drawn.
+
+    Parameters
+    ----------
+    data
+        File path(s), glob pattern(s), ``"path:tree"`` strings, a mapping of
+        arrays, an Awkward record array, or an existing
+        :class:`~rootfig.io.Source`.
+    tree
+        Tree name for file inputs. Auto-detected when the file holds exactly one
+        ``TTree``/``RNTuple``.
+    label
+        Legend label. Defaults to the file stem.
+    selection
+        A selection applied to this sample only, combined with the selection
+        given to :func:`rootfig.plot` with ``&``.
+    weight
+        Expression for per-event (or per-object) weights, e.g. ``"mc_weight"``
+        or ``"mc_weight * pileup_sf"``. Combined multiplicatively with the
+        weight given to :func:`rootfig.plot`.
+    is_data
+        Mark as observed data: drawn as black points with error bars, excluded
+        from stacks, and used as the numerator of data/MC ratios.
+    color
+        Matplotlib colour. Defaults to the style's colour cycle.
+    histtype
+        How to draw this sample; see :data:`HistType`. Defaults depend on the
+        plot type.
+    scale
+        Constant multiplied into every weight.
+    xsec
+        Cross section of the process, in pb or as a string with a unit
+        (``"1.2 fb"``). With a luminosity given to :func:`rootfig.plot`
+        (``lumi=``), every weight is multiplied by ``xsec * lumi / ngen`` so
+        that the histogram is the expected yield.
+    ngen
+        Number of generated events the cross section refers to: a number, the
+        name of an object in the file(s) holding it (a sum-of-weights histogram
+        or a ``TParameter`` such as FCCAnalyses' ``"eventsProcessed"``), or
+        ``None`` to use the number of entries in the source.
+    entry_start, entry_stop
+        Read only this range of entries (counted across files).
+
+    Examples
+    --------
+    >>> import rootfig as rf
+    >>> sig = rf.Sample("sig.root", tree="events", label="Signal", weight="mc_w")
+    >>> bkg = rf.Sample(["bkg_*.root"], tree="events", label="Background", color="gray")
+    >>> data = rf.Sample("data.root", tree="events", label="Data", is_data=True)
+    """
+
+    source: Source
+    label: str
+    selection: Cut | None = None
+    weight: str | None = None
+    is_data: bool = False
+    color: str | None = None
+    histtype: HistType | None = None
+    scale: float = 1.0
+    xsec: float | str | None = None
+    ngen: float | str | None = None
+
+    def __init__(
+        self,
+        data: Any,
+        *,
+        tree: str | None = None,
+        label: str | None = None,
+        selection: CutLike | None = None,
+        weight: str | None = None,
+        is_data: bool = False,
+        color: str | None = None,
+        histtype: HistType | None = None,
+        scale: float = 1.0,
+        xsec: float | str | None = None,
+        ngen: float | str | None = None,
+        entry_start: int | None = None,
+        entry_stop: int | None = None,
+    ) -> None:
+        source = as_source(data, tree=tree, entry_start=entry_start, entry_stop=entry_stop)
+        if label is None:
+            label = getattr(source, "default_label", None) or source.describe()
+        if weight is not None and not weight.strip():
+            weight = None
+        object.__setattr__(self, "source", source)
+        object.__setattr__(self, "label", label)
+        object.__setattr__(self, "selection", as_cut(selection))
+        object.__setattr__(self, "weight", weight)
+        object.__setattr__(self, "is_data", is_data)
+        object.__setattr__(self, "color", color)
+        object.__setattr__(self, "histtype", histtype)
+        if not math.isfinite(scale):
+            msg = f"sample {label!r}: scale must be a finite number, got {scale!r}"
+            raise ValueError(msg)
+        object.__setattr__(self, "scale", float(scale))
+        if xsec is not None:
+            cross_section_pb(xsec)  # validate early
+        if isinstance(ngen, int | float) and not math.isfinite(ngen):
+            msg = f"sample {label!r}: ngen must be a finite number, got {ngen!r}"
+            raise LuminosityError(msg)
+        object.__setattr__(self, "xsec", xsec)
+        object.__setattr__(self, "ngen", ngen)
+
+    def __repr__(self) -> str:
+        parts = [repr(self.source.describe()), f"label={self.label!r}"]
+        if self.selection is not None:
+            parts.append(f"selection={self.selection.expression!r}")
+        if self.weight is not None:
+            parts.append(f"weight={self.weight!r}")
+        if self.is_data:
+            parts.append("is_data=True")
+        if self.xsec is not None:
+            parts.append(f"xsec={self.xsec!r}")
+        return f"Sample({', '.join(parts)})"
+
+    def generated_events(self) -> float:
+        """Return the number of generated events ``ngen`` refers to (see the class docstring)."""
+        if isinstance(self.ngen, int | float):
+            return float(self.ngen)
+        if isinstance(self.ngen, str):
+            reader = getattr(self.source, "read_scalar", None)
+            if not callable(reader):
+                msg = f"sample {self.label!r}: ngen={self.ngen!r} needs a file source"
+                raise LuminosityError(msg)
+            return float(reader(self.ngen))
+        counter = getattr(self.source, "num_entries", None)
+        if callable(counter):
+            return float(counter())
+        first = self.source.branches()[:1]
+        return float(len(next(iter(self.source.arrays(first).values()))))
+
+    def lumi_scale(self, lumi: float | str | None) -> float:
+        """Factor turning weights into expected yields at ``lumi``; 1 without a cross section.
+
+        ``xsec`` (pb) times ``lumi`` (fb^-1, or a string with a unit) divided by
+        :meth:`generated_events`.
+        """
+        if self.xsec is None:
+            return 1.0
+        if lumi is None:
+            msg = (
+                f"sample {self.label!r} has a cross section ({self.xsec!r}) but no luminosity "
+                "was given; pass lumi=... (in fb^-1, or a string such as '10.8 ab^-1')"
+            )
+            raise LuminosityError(msg)
+        ngen = self.generated_events()
+        if not (math.isfinite(ngen) and ngen > 0):
+            msg = f"sample {self.label!r}: the number of generated events is {ngen}"
+            raise LuminosityError(msg)
+        factor = cross_section_pb(self.xsec) * 1e3 * luminosity_fb(lumi) / ngen
+        if not math.isfinite(factor):
+            msg = (
+                f"sample {self.label!r}: xsec * lumi / ngen = {factor} is not finite "
+                f"(xsec={self.xsec!r}, lumi={lumi!r}, ngen={ngen})"
+            )
+            raise LuminosityError(msg)
+        return factor
+
+    @property
+    def files(self) -> tuple[str, ...]:
+        """The resolved input files, or an empty tuple for in-memory sources."""
+        return self.source.files if isinstance(self.source, FileSource) else ()
+
+    def with_(self, **changes: Any) -> Sample:
+        """Return a copy with the given fields replaced, e.g. ``sample.with_(label="B")``."""
+        known = {f.name for f in fields(self)}
+        unknown = set(changes) - known
+        if unknown:
+            msg = f"unknown Sample field(s) {sorted(unknown)}; valid fields: {sorted(known)}"
+            raise TypeError(msg)
+        if "selection" in changes:
+            changes["selection"] = as_cut(changes["selection"])
+        clone = copy.copy(self)
+        for name, value in changes.items():
+            object.__setattr__(clone, name, value)
+        return clone
+
+
+SampleLike: TypeAlias = Sample | str | Sequence[Any] | Any
+"""A :class:`Sample`, file specification(s), or in-memory arrays."""
+
+
+def as_samples(
+    data: Any,
+    *,
+    tree: str | None = None,
+    labels: str | Sequence[str] | None = None,
+    entry_start: int | None = None,
+    entry_stop: int | None = None,
+) -> list[Sample]:
+    """Normalise the ``data`` argument of :func:`rootfig.plot` into a list of samples.
+
+    * A :class:`Sample` gives one sample.
+    * A single path/glob/mapping/array gives one sample.
+    * A list gives one sample per element. Elements may mix ``Sample`` objects
+      with raw file specifications; each raw element becomes its own sample,
+      so ``["sig.root", "bkg.root"]`` produces two samples. To combine several
+      files into *one* sample wrap them in a ``Sample`` or pass a glob.
+    * A mapping from label to file specification, in-memory arrays or
+      ``Sample`` gives one labelled sample per entry (a ``Sample`` keeps its
+      selection, weight and cross section and only takes the label). A mapping
+      whose values are columns (arrays or lists of numbers) is one in-memory
+      sample instead.
+
+    Raises
+    ------
+    SourceError
+        If ``labels`` does not match the number of samples.
+    """
+    if isinstance(data, Sample):
+        samples = [data]
+    elif isinstance(data, Mapping) and _looks_like_label_map(data):
+        samples = [
+            value.with_(label=key)
+            if isinstance(value, Sample)
+            else Sample(value, tree=tree, label=key, entry_start=entry_start, entry_stop=entry_stop)
+            for key, value in data.items()
+        ]
+    elif isinstance(data, list | tuple):
+        if not data:
+            msg = "no samples given"
+            raise SourceError(msg)
+        samples = [
+            item
+            if isinstance(item, Sample)
+            else Sample(item, tree=tree, entry_start=entry_start, entry_stop=entry_stop)
+            for item in data
+        ]
+    else:
+        samples = [Sample(data, tree=tree, entry_start=entry_start, entry_stop=entry_stop)]
+
+    if labels is not None:
+        label_list = [labels] if isinstance(labels, str) else list(labels)
+        if len(label_list) != len(samples):
+            msg = f"got {len(label_list)} labels for {len(samples)} samples"
+            raise SourceError(msg)
+        samples = [s.with_(label=lab) for s, lab in zip(samples, label_list, strict=True)]
+    return samples
+
+
+def _looks_like_label_map(data: Mapping[Any, Any]) -> bool:
+    """Distinguish ``{"Signal": "sig.root"}`` from a mapping of column arrays.
+
+    Values that describe a dataset (paths, lists of paths, ``Sample`` objects,
+    mappings of arrays, Awkward record arrays) make a label map; anything else
+    (NumPy arrays, lists of numbers, flat Awkward arrays) is a column.
+    """
+    return bool(data) and all(isinstance(k, str) and _is_dataset_spec(v) for k, v in data.items())
+
+
+def _is_dataset_spec(value: Any) -> bool:
+    if isinstance(value, str | PathLike | Sample | Mapping):
+        return True
+    if isinstance(value, ak.Array):
+        return bool(value.fields)
+    if isinstance(value, list | tuple):
+        return bool(value) and all(isinstance(item, str | PathLike) for item in value)
+    return False