rootfig 0.1.0__py3-none-any.whl

rootfig/histograms/build.py

@@ -0,0 +1,197 @@
+"""Filling ``Hist`` objects and the :class:`Histogram` wrapper."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any
+
+import hist
+import numpy as np
+
+from rootfig._typing import FloatArray, Hist
+from rootfig.histograms.stats import Summary
+
+if TYPE_CHECKING:
+    from rootfig.model.binning import Axis
+    from rootfig.model.samples import HistType, Sample
+    from rootfig.selection import Columns
+
+__all__ = ["Histogram", "as_weight_storage", "fill"]
+
+
+def fill(axes: Sequence[Axis], columns: Columns) -> Hist:
+    """Create a ``Hist`` with ``Weight`` storage and fill it from ``columns``.
+
+    One axis is expected per column in ``columns.arrays``. Unweighted data is
+    filled without weights, so bin variances equal bin counts.
+    """
+    if len(axes) != len(columns.arrays):
+        msg = f"got {len(axes)} axes for {len(columns.arrays)} columns"
+        raise ValueError(msg)
+    histogram = hist.Hist(*axes, storage=hist.storage.Weight())
+    if columns.n_entries:
+        if columns.weights is None:
+            histogram.fill(*columns.arrays)
+        else:
+            histogram.fill(*columns.arrays, weight=columns.weights)
+    return histogram
+
+
+def as_weight_storage(histogram: Hist) -> Hist:
+    """Return ``histogram`` with ``Weight`` storage (a copy if it had another storage).
+
+    Plain count storages (``Double``, ``Int64``, ...) carry no sum of squared
+    weights; their variances are taken as ``hist`` reports them, i.e. the
+    counts (Poisson) for unweighted fills. If a storage reports no variances at
+    all (a ``Double`` histogram filled with weights) the contents are used as
+    variances, the best available guess, and the caller is not told otherwise
+    because ``hist`` itself makes the same assumption when drawing.
+    """
+    if histogram.storage_type is hist.storage.Weight:
+        return histogram
+    if histogram.ndim and histogram.storage_type not in (
+        hist.storage.Double,
+        hist.storage.Int64,
+        hist.storage.AtomicInt64,
+        hist.storage.Unlimited,
+    ):
+        msg = (
+            f"histograms with {histogram.storage_type.__name__} storage are not supported; "
+            "use Weight (or a plain count) storage"
+        )
+        raise TypeError(msg)
+    values = np.asarray(histogram.values(flow=True), dtype=float)
+    reported = histogram.variances(flow=True)
+    variances = values if reported is None else np.asarray(reported, dtype=float)
+    result = hist.Hist(*histogram.axes, storage=hist.storage.Weight())
+    view: Any = result.view(flow=True)
+    view.value = values
+    view.variance = variances
+    return result
+
+
+@dataclass(frozen=True)
+class Histogram:
+    """A filled histogram together with its provenance and drawing hints.
+
+    The underlying :class:`Hist` (``Weight`` storage, so bin variances
+    are the sums of squared weights) is available as :attr:`hist`; everything
+    else is metadata used for legends, ratios and statistics.
+
+    Attributes
+    ----------
+    hist
+        The histogram itself (1D or 2D).
+    label
+        Legend label.
+    sample
+        The :class:`~rootfig.model.Sample` this was filled from, if any.
+    stats
+        Unbinned :class:`~rootfig.histograms.Summary` statistics of the filled
+        values, if available.
+    is_data
+        Whether this represents observed data.
+    color, histtype
+        Drawing hints, ``None`` for style defaults.
+    normalization
+        Description of the normalisation applied (``None`` for raw counts).
+    """
+
+    hist: Hist
+    label: str
+    sample: Sample | None = None
+    stats: Summary | None = None
+    is_data: bool = False
+    color: str | None = None
+    histtype: HistType | None = None
+    normalization: str | None = None
+
+    # -- convenience accessors -----------------------------------------------------------
+
+    @property
+    def ndim(self) -> int:
+        """Number of axes."""
+        return int(self.hist.ndim)
+
+    @property
+    def axis(self) -> Any:
+        """The first axis."""
+        return self.hist.axes[0]
+
+    @property
+    def edges(self) -> FloatArray:
+        """Bin edges of the first axis."""
+        return np.asarray(self.hist.axes[0].edges, dtype=float)
+
+    @property
+    def centers(self) -> FloatArray:
+        """Bin centres of the first axis."""
+        return np.asarray(self.hist.axes[0].centers, dtype=float)
+
+    @property
+    def widths(self) -> FloatArray:
+        """Bin widths of the first axis."""
+        return np.asarray(self.hist.axes[0].widths, dtype=float)
+
+    def values(self, *, flow: bool = False) -> FloatArray:
+        """Bin contents (sum of weights)."""
+        return np.asarray(self.hist.values(flow=flow), dtype=float)
+
+    def variances(self, *, flow: bool = False) -> FloatArray:
+        """Bin variances (sum of squared weights)."""
+        return np.asarray(self.hist.variances(flow=flow), dtype=float)
+
+    def errors(self, *, flow: bool = False) -> FloatArray:
+        """Bin uncertainties, ``sqrt(variances)``."""
+        return np.asarray(np.sqrt(self.variances(flow=flow)), dtype=float)
+
+    @property
+    def integral(self) -> float:
+        """Sum of the visible bin contents (flow bins excluded)."""
+        return float(self.values().sum())
+
+    @property
+    def underflow(self) -> float:
+        """Content of the underflow bin (first axis, 1D only; ``0`` if the axis has none)."""
+        return self._flow_cell("underflow", self.hist.values(flow=True))
+
+    @property
+    def overflow(self) -> float:
+        """Content of the overflow bin (first axis, 1D only; ``0`` if the axis has none)."""
+        return self._flow_cell("overflow", self.hist.values(flow=True))
+
+    @property
+    def underflow_variance(self) -> float:
+        """Variance of the underflow bin (first axis, 1D only; ``0`` if the axis has none)."""
+        return self._flow_cell("underflow", self.hist.variances(flow=True))
+
+    @property
+    def overflow_variance(self) -> float:
+        """Variance of the overflow bin (first axis, 1D only; ``0`` if the axis has none)."""
+        return self._flow_cell("overflow", self.hist.variances(flow=True))
+
+    def _flow_cell(self, side: str, cells: Any) -> float:
+        if self.ndim != 1:
+            return float("nan")
+        traits = self.axis.traits
+        if not getattr(traits, side):
+            return 0.0
+        return float(cells[0 if side == "underflow" else -1])
+
+    @property
+    def entries(self) -> int:
+        """Number of filled entries, from the statistics if available."""
+        if self.stats is not None:
+            return self.stats.entries
+        if self.normalization is not None:
+            return 0
+        return int(np.rint(float(np.sum(self.hist.values(flow=True)))))
+
+    def with_(self, **changes: Any) -> Histogram:
+        """Return a copy with the given fields replaced."""
+        return replace(self, **changes)
+
+    def scaled(self, factor: float) -> Histogram:
+        """Return a copy multiplied by ``factor`` (variances scale with ``factor**2``)."""
+        return replace(self, hist=self.hist * factor)

rootfig/histograms/cutflow.py

@@ -0,0 +1,195 @@
+"""Cut flows: event counts and weighted yields after successive selections."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from rootfig.expressions import parse
+from rootfig.histograms.pipeline import combined_weight, read_arrays
+from rootfig.model.cuts import Cut, CutLike, as_cut
+from rootfig.model.samples import Sample
+from rootfig.selection import NonFinitePolicy, event_mask, event_weights
+
+__all__ = ["Cutflow", "CutflowStep", "CutflowTable", "cutflow"]
+
+
+@dataclass(frozen=True)
+class CutflowStep:
+    """Yields after one more cut has been applied.
+
+    Attributes
+    ----------
+    label
+        The cut's label (or expression), ``"All"`` for the first step.
+    expression
+        The cut expression, empty for the first step.
+    events
+        Raw number of events passing all cuts so far.
+    yield_
+        Weighted yield (sum of weights, including the sample scale and luminosity).
+    error
+        Statistical uncertainty on ``yield_``, ``sqrt(sum w^2)``.
+    """
+
+    label: str
+    expression: str
+    events: int
+    yield_: float
+    error: float
+
+
+@dataclass(frozen=True)
+class Cutflow:
+    """The cut flow of one sample: a sequence of :class:`CutflowStep`."""
+
+    sample: str
+    steps: tuple[CutflowStep, ...]
+
+    @property
+    def labels(self) -> list[str]:
+        """Step labels."""
+        return [step.label for step in self.steps]
+
+    @property
+    def yields(self) -> np.ndarray:
+        """Weighted yield per step."""
+        return np.array([step.yield_ for step in self.steps], dtype=float)
+
+    @property
+    def events(self) -> np.ndarray:
+        """Raw event count per step."""
+        return np.array([step.events for step in self.steps], dtype=int)
+
+    @property
+    def efficiencies(self) -> np.ndarray:
+        """Weighted efficiency of each step relative to the previous one (1 for the first)."""
+        y = self.yields
+        with np.errstate(divide="ignore", invalid="ignore"):
+            rel = np.where(y[:-1] > 0, y[1:] / y[:-1], np.nan)
+        return np.concatenate([[1.0], rel])
+
+    @property
+    def absolute_efficiencies(self) -> np.ndarray:
+        """Weighted efficiency of each step relative to the first."""
+        y = self.yields
+        with np.errstate(divide="ignore", invalid="ignore"):
+            return np.where(y[0] > 0, y / y[0], np.nan)
+
+
+@dataclass(frozen=True)
+class CutflowTable:
+    """Cut flows of several samples with the same cuts; ``str(table)`` is an aligned text table."""
+
+    rows: tuple[Cutflow, ...]
+
+    def get(self, sample: str) -> Cutflow:
+        """Return the cut flow of the sample labelled ``sample``."""
+        for flow in self.rows:
+            if flow.sample == sample:
+                return flow
+        msg = f"no cut flow for sample {sample!r}; have {[f.sample for f in self.rows]}"
+        raise KeyError(msg)
+
+    @property
+    def samples(self) -> list[str]:
+        """Sample labels in order."""
+        return [flow.sample for flow in self.rows]
+
+    @property
+    def labels(self) -> list[str]:
+        """Step labels (identical for all samples)."""
+        return self.rows[0].labels if self.rows else []
+
+    def __str__(self) -> str:
+        if not self.rows:
+            return ""
+        header = ["", *(f"{flow.sample}" for flow in self.rows)]
+        lines = [header]
+        for index, label in enumerate(self.labels):
+            cells = [label]
+            for flow in self.rows:
+                step = flow.steps[index]
+                eff = flow.efficiencies[index]
+                if step.events == step.yield_ == int(step.yield_):
+                    cell = f"{step.events}"
+                else:
+                    cell = f"{_fmt(step.yield_)} ± {_fmt(step.error)} ({step.events})"
+                if index > 0:
+                    cell += f"  {100 * eff:5.1f}%"
+                cells.append(cell)
+            lines.append(cells)
+        widths = [max(len(row[i]) for row in lines) for i in range(len(header))]
+        return "\n".join(
+            "  ".join(
+                cell.ljust(widths[i]) if i == 0 else cell.rjust(widths[i])
+                for i, cell in enumerate(row)
+            )
+            for row in lines
+        )
+
+
+def _fmt(value: float) -> str:
+    if value == 0:
+        return "0"
+    if abs(value) >= 100:
+        return f"{value:,.0f}"
+    return f"{value:.3g}"
+
+
+def cutflow(
+    sample: Sample,
+    cuts: Sequence[CutLike],
+    *,
+    weight: str | None = None,
+    lumi: float | str | None = None,
+    nonfinite: NonFinitePolicy = "drop",
+) -> Cutflow:
+    """Apply ``cuts`` one after another to ``sample`` and count events and yields.
+
+    The sample's own selection (if any) defines the first step; the plot-level
+    ``weight`` is multiplied with the sample's weight, scale and luminosity
+    factor. Per-object cuts pass an event when any object passes. Events with a
+    missing or non-finite weight are excluded from every step (counts and
+    yields); non-finite ones are reported following ``nonfinite`` (a
+    :class:`~rootfig.errors.RootfigWarning`, or a
+    :class:`~rootfig.errors.SelectionError` for ``"error"``).
+    """
+    steps: list[Cut] = []
+    for item in cuts:
+        cut = as_cut(item)
+        if cut is not None:
+            steps.append(cut)
+    weight_expr = combined_weight(sample, weight)
+    expressions = [parse(c.expression) for c in steps]
+    if sample.selection is not None:
+        expressions.append(sample.selection.parsed())
+    if weight_expr is not None:
+        expressions.append(parse(weight_expr))
+    arrays, n_events = read_arrays(sample, expressions)
+    weights = event_weights(
+        weight_expr, arrays, n_events, nonfinite=nonfinite, context=sample.label
+    )
+    passing = np.isfinite(weights)
+    weights = np.where(passing, weights, 0.0) * sample.scale * sample.lumi_scale(lumi)
+    if sample.selection is not None:
+        passing &= event_mask(sample.selection.expression, arrays, length=n_events)
+
+    def step(label: str, expression: str) -> CutflowStep:
+        selected = weights[passing]
+        return CutflowStep(
+            label=label,
+            expression=expression,
+            events=int(np.count_nonzero(passing)),
+            yield_=float(selected.sum()),
+            error=float(np.sqrt(np.sum(selected**2))),
+        )
+
+    first_label = sample.selection.label or "All" if sample.selection is not None else "All"
+    result = [step(first_label, sample.selection.expression if sample.selection else "")]
+    for cut in steps:
+        passing &= event_mask(cut.expression, arrays, length=n_events)
+        result.append(step(cut.label or cut.expression, cut.expression))
+    return Cutflow(sample=sample.label, steps=tuple(result))

rootfig/histograms/efficiency.py

@@ -0,0 +1,214 @@
+"""Efficiencies (pass / total with binomial intervals) and profiles (a statistic of y per x bin)."""
+
+from __future__ import annotations
+
+import math
+import warnings
+from dataclasses import dataclass
+from typing import Literal, TypeAlias
+
+import numpy as np
+
+from rootfig._typing import FloatArray, Hist
+from rootfig.errors import BinningError, RootfigWarning
+from rootfig.histograms.ratio import compatible_binning
+
+__all__ = ["Efficiency", "Profile", "ProfileStatistic", "efficiency", "profile"]
+
+ProfileStatistic: TypeAlias = Literal["mean", "std"]
+"""What a profile shows per bin: the weighted mean of ``y`` or its standard deviation."""
+
+
+@dataclass(frozen=True)
+class Efficiency:
+    """Bin-by-bin efficiency ``passed / total`` with a binomial confidence interval.
+
+    Attributes
+    ----------
+    values
+        The efficiency; ``nan`` where ``total`` is empty.
+    lower, upper
+        Bounds of the Wilson score interval (``z`` standard deviations; ``z = 1``
+        is the usual 68 % band), computed with the effective number of entries
+        so weighted samples get sensible intervals. The interval always contains
+        the value; it is ``nan`` where the value is outside ``[0, 1]`` (negative
+        weights), since a binomial interval is undefined there.
+    edges
+        Bin edges.
+    label
+        Legend label.
+    """
+
+    values: FloatArray
+    lower: FloatArray
+    upper: FloatArray
+    edges: FloatArray
+    label: str = ""
+
+    @property
+    def centers(self) -> FloatArray:
+        """Bin centres."""
+        return np.asarray(0.5 * (self.edges[1:] + self.edges[:-1]), dtype=float)
+
+    @property
+    def half_widths(self) -> FloatArray:
+        """Half bin widths."""
+        return np.asarray(0.5 * np.diff(self.edges), dtype=float)
+
+    @property
+    def errors(self) -> tuple[FloatArray, FloatArray]:
+        """``(values - lower, upper - values)``, ready for ``yerr``."""
+        return (
+            np.asarray(self.values - self.lower, dtype=float),
+            np.asarray(self.upper - self.values, dtype=float),
+        )
+
+
+def efficiency(passed: Hist, total: Hist, *, z: float = 1.0, label: str = "") -> Efficiency:
+    """Compute ``passed / total`` per bin with a Wilson score interval.
+
+    Both histograms must share their binning; ``passed`` should be a subset of
+    ``total``. With weights the effective counts ``(sum w)^2 / sum w^2`` of the
+    total replace the raw counts in the interval. The interval is clipped to
+    ``[0, 1]`` and always contains the efficiency (at 0 % and 100 % the
+    respective bound coincides with the value). Bins whose efficiency falls
+    outside ``[0, 1]``, which can only happen with negative weights, get ``nan``
+    bounds and a :class:`~rootfig.errors.RootfigWarning`: no binomial interval
+    describes them.
+
+    Raises
+    ------
+    BinningError
+        If the binnings differ or ``z`` is not a positive finite number.
+    """
+    if not compatible_binning(passed, total):
+        msg = "efficiency requires two one-dimensional histograms with identical bin edges"
+        raise BinningError(msg)
+    if not (math.isfinite(z) and z > 0):
+        msg = f"z must be a positive finite number of standard deviations, got {z!r}"
+        raise BinningError(msg)
+    k = np.asarray(passed.values(), dtype=float)
+    n = np.asarray(total.values(), dtype=float)
+    vn = np.asarray(total.variances(), dtype=float)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        ok = n > 0
+        p = np.where(ok, k / n, np.nan)
+        n_eff = np.where(vn > 0, n**2 / vn, n)
+        z2 = z * z
+        denominator = 1.0 + z2 / n_eff
+        centre = (p + z2 / (2.0 * n_eff)) / denominator
+        half = (z / denominator) * np.sqrt(p * (1.0 - p) / n_eff + z2 / (4.0 * n_eff**2))
+        valid = ok & (p >= 0.0) & (p <= 1.0)
+        # the Wilson interval contains p by construction; guard against round-off at 0 and 1
+        lower = np.where(valid, np.minimum(np.clip(centre - half, 0.0, 1.0), p), np.nan)
+        upper = np.where(valid, np.maximum(np.clip(centre + half, 0.0, 1.0), p), np.nan)
+    undefined = int(np.count_nonzero(ok & ~valid))
+    if undefined:
+        warnings.warn(
+            f"{label + ': ' if label else ''}{undefined} bin(s) have an efficiency outside "
+            "[0, 1] (negative weights?); no confidence interval is drawn for them",
+            RootfigWarning,
+            stacklevel=2,
+        )
+    return Efficiency(
+        values=np.asarray(p, dtype=float),
+        lower=np.asarray(lower, dtype=float),
+        upper=np.asarray(upper, dtype=float),
+        edges=np.asarray(total.axes[0].edges, dtype=float),
+        label=label,
+    )
+
+
+@dataclass(frozen=True)
+class Profile:
+    """A statistic of ``y`` in bins of ``x``.
+
+    Attributes
+    ----------
+    values
+        The weighted mean (``statistic="mean"``) or standard deviation
+        (``"std"``) of ``y`` per bin; ``nan`` for empty bins, and ``nan`` for
+        the standard deviation where negative weights make the weighted
+        variance negative.
+    errors
+        Standard error: ``std / sqrt(n_eff)`` for the mean, ``std / sqrt(2 n_eff)``
+        for the standard deviation, with the effective entries ``n_eff``.
+    counts
+        Sum of weights per bin.
+    edges
+        Bin edges.
+    statistic, label
+        What is shown and the legend label.
+    """
+
+    values: FloatArray
+    errors: FloatArray
+    counts: FloatArray
+    edges: FloatArray
+    statistic: ProfileStatistic = "mean"
+    label: str = ""
+
+    @property
+    def centers(self) -> FloatArray:
+        """Bin centres."""
+        return np.asarray(0.5 * (self.edges[1:] + self.edges[:-1]), dtype=float)
+
+    @property
+    def half_widths(self) -> FloatArray:
+        """Half bin widths."""
+        return np.asarray(0.5 * np.diff(self.edges), dtype=float)
+
+
+def profile(
+    x: FloatArray,
+    y: FloatArray,
+    edges: FloatArray,
+    *,
+    weights: FloatArray | None = None,
+    statistic: ProfileStatistic = "mean",
+    label: str = "",
+) -> Profile:
+    """Bin ``x`` with ``edges`` and compute the weighted mean or standard deviation of ``y``.
+
+    The variance is computed from the deviations from the bin mean (two passes),
+    so a narrow spread at a large offset, e.g. ``std([1e9, 1e9 + 1]) = 0.5``, is
+    not lost to cancellation.
+    """
+    x = np.asarray(x, dtype=float)
+    y = np.asarray(y, dtype=float)
+    edges = np.asarray(edges, dtype=float)
+    w = np.ones_like(x) if weights is None else np.asarray(weights, dtype=float)
+    if not (x.shape == y.shape == w.shape):
+        msg = "x, y and weights must have the same length"
+        raise BinningError(msg)
+    if statistic not in ("mean", "std"):
+        msg = f"statistic must be 'mean' or 'std', got {statistic!r}"
+        raise BinningError(msg)
+    n_bins = len(edges) - 1
+    index = np.digitize(x, edges) - 1
+    inside = (index >= 0) & (index < n_bins)
+    index, y, w = index[inside], y[inside], w[inside]
+    sw = np.bincount(index, weights=w, minlength=n_bins)
+    sw2 = np.bincount(index, weights=w * w, minlength=n_bins)
+    swy = np.bincount(index, weights=w * y, minlength=n_bins)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        ok = sw > 0
+        mean = np.where(ok, swy / sw, np.nan)
+        residual = y - np.where(ok, mean, 0.0)[index]
+        swr2 = np.bincount(index, weights=w * residual * residual, minlength=n_bins)
+        variance = np.where(ok, swr2 / sw, np.nan)
+        variance = np.where(variance < 0, np.nan, variance)  # negative weights: undefined
+        std = np.sqrt(variance)
+        n_eff = np.where(sw2 > 0, sw**2 / sw2, 0.0)
+        if statistic == "mean":
+            values, errors = mean, np.where(n_eff > 0, std / np.sqrt(n_eff), np.nan)
+        else:
+            values, errors = std, np.where(n_eff > 0, std / np.sqrt(2.0 * n_eff), np.nan)
+    return Profile(
+        values=np.asarray(values, dtype=float),
+        errors=np.asarray(errors, dtype=float),
+        counts=np.asarray(sw, dtype=float),
+        edges=edges,
+        statistic=statistic,
+        label=label,
+    )