feat: EqualData, bin edges with equal data counts computed from the data when filling

README.md

@@ -170,6 +170,11 @@ combine: {cmssw_dir: /work/USER/CMSSW_14_1_9, scram_arch: el9_amd64_gcc12}
 
 - Expressions (cuts, weights, variables) are pandas `eval` syntax on ntuple columns, evaluated
   with numexpr: `(pt_1 > 25) & (abs(eta_1) < 2.1)`, `id_wgt_tau_1 * trg_wgt`.
+- The edges of a variable are numbers or `EqualData(n_bins, low, high)`: `n_bins` bins with equal
+  data counts in (low, high), as smhtt_ul `gof/build_binning.py` (`binning.py`). Every fill
+  (`hist`, a measurement) computes them from the data of the category in the nominal region, uses
+  them for every region and variation, and records them in `binning.json` next to the histogram
+  file.
 - Histograms are the small numpy `shapesmith.histogram.Histogram`; ROOT files are written and
   read through uproot (`boost-histogram`/`hist` are broken in LCG_108, their axis edges come out
   constant).

src/shapesmith/binning.py

@@ -0,0 +1,85 @@
+"""Equal-data binning: the edges of an `EqualData` variable from the data of the run, as smhtt_ul gof/build_binning.py
+(get_data_selection and get_1d_binning).
+
+Per category: the data process of the channel in the nominal region, its events in the category, the values of the
+variable in (low, high); the unweighted percentiles 0, 100 / n_bins, ..., 100 (linear interpolation), duplicates
+removed, every edge moved by -1e-4 and the last one by +2e-4 on top, so that the smallest and the largest value are
+inside. `run_hist` resolves the edges whenever it fills, so every region and variation of a category shares them, and
+records them in `binning.json` next to its output. smhtt_ul's mask of default values (-11, -999, -10, -1) is not
+applied.
+"""
+from __future__ import annotations
+
+import dataclasses
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from shapesmith.config import RunConfig
+from shapesmith.events import Query, load
+from shapesmith.expressions import columns_of, evaluate, mask
+from shapesmith.model import NOMINAL, Analysis, EqualData
+
+if TYPE_CHECKING:
+    from shapesmith.fill import Target
+
+EPSILON = 1e-4
+RECORD = "binning.json"
+
+
+def inside(values: np.ndarray, rule: EqualData) -> np.ndarray:
+    """The finite values in (low, high), as float64."""
+    x = np.asarray(values, dtype=np.float64)
+    return x[np.isfinite(x) & (x > rule.low) & (x < rule.high)]
+
+
+def equal_data_edges(values: np.ndarray, rule: EqualData) -> tuple[float, ...]:
+    x = inside(values, rule)
+    if not x.size:
+        raise ValueError(f"no data in ({rule.low}, {rule.high})")
+    edges = sorted({float(edge) for edge in np.percentile(x, np.linspace(0.0, 100.0, rule.n_bins + 1))})
+    if len(edges) < 2:
+        raise ValueError(f"fewer than two distinct edges from {x.size} values")
+    edges = [edge - EPSILON for edge in edges]
+    edges[-1] += 2 * EPSILON
+    return tuple(edges)
+
+
+def resolve(config: RunConfig, analysis: Analysis, channel_name: str, targets: list[Target]) -> tuple[list[Target], dict]:
+    """The targets with the edges of their EqualData variables, and the record of those edges
+    ({category: {variable: {"rule", "n_data", "edges"}}}). The data are read only if a target needs them."""
+    pending = [t for t in targets if isinstance(t.variable.edges, EqualData)]
+    if not pending:
+        return list(targets), {}
+    channel = analysis.channel(channel_name)
+    exprs = [t.variable.expr for t in pending] + [t.cut for t in pending if t.cut]
+    frame = load(config, analysis, Query(channel_name, channel.data(), NOMINAL), columns_of(exprs)).frame
+    resolved, record = [], {}
+    for target in targets:
+        rule = target.variable.edges
+        if not isinstance(rule, EqualData):
+            resolved.append(target)
+            continue
+        values = evaluate(frame[mask(frame, [target.cut])] if target.cut else frame, target.variable.expr)
+        try:
+            edges = equal_data_edges(values, rule)
+        except ValueError as error:
+            raise ValueError(f"{channel_name}/{target.category}/{target.variable.name}: equal-data binning: {error}") from None
+        resolved.append(dataclasses.replace(target, variable=dataclasses.replace(target.variable, edges=edges)))
+        record.setdefault(target.category, {})[target.variable.name] = {"rule": dataclasses.asdict(rule), "n_data": int(inside(values, rule).size), "edges": list(edges)}
+    return resolved, record
+
+
+def write_record(path: Path, record: dict) -> Path:
+    """Merge `record` ({channel: {category: {variable: entry}}}) into the JSON at `path`: its entries replace the stored
+    ones, the others are kept (as the histograms of other scopes in the shapes file)."""
+    path = Path(path)
+    stored = json.loads(path.read_text()) if path.exists() else {}
+    for channel, categories in record.items():
+        for category, variables in categories.items():
+            stored.setdefault(channel, {}).setdefault(category, {}).update(variables)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(stored, indent=1))
+    return path

src/shapesmith/fill.py

@@ -14,11 +14,12 @@
 
 import numpy as np
 
+from shapesmith.binning import RECORD, resolve, write_record
 from shapesmith.config import RunConfig
 from shapesmith.events import event_weights, lumi, select
 from shapesmith.expressions import columns_of, evaluate, mask, shift
 from shapesmith.histogram import INCLUSIVE, NOMINAL_VARIATION, HistKey, Histogram, HistogramSet
-from shapesmith.model import ABCD, NOMINAL, Analysis, Channel, ColumnVariation, DataMinus, Process, Variable, Variation, applies
+from shapesmith.model import ABCD, NOMINAL, Analysis, Channel, ColumnVariation, DataMinus, EqualData, Process, Variable, Variation, applies
 from shapesmith.parallel import run_jobs
 from shapesmith.store import SKIM_COLUMNS, read_skims, schema
 
@@ -126,6 +127,9 @@ def plan_booking(channel: Channel, booking: Booking, targets_: list[Target], ava
 
 def fill_booking(config: RunConfig, analysis: Analysis, channel_name: str, booking: Booking, targets_: list[Target]) -> dict[HistKey, Histogram]:
     """All histograms of one booking, from one frame of the process's skims."""
+    unresolved = [t.variable.name for t in targets_ if isinstance(t.variable.edges, EqualData)]
+    if unresolved:
+        raise ValueError(f"{channel_name}: the equal-data binning of {', '.join(unresolved)} is not resolved (fill through run_hist)")
     channel = analysis.channel(channel_name)
     process = booking.process
     nicks = [s.nick for s in channel.samples_of(process.group)]
@@ -159,12 +163,13 @@ def run_hist(config: RunConfig, analysis: Analysis, channels: list[str] | None,
     """Fill every booked histogram of the requested channels into `output` (+ .json index).
 
     The requested scopes (channel, process, region, category, variable) are always filled; histograms of other
-    scopes already in `output` are kept."""
+    scopes already in `output` are kept. The edges of EqualData variables are computed from the data first and
+    recorded in binning.json next to `output`."""
     output = Path(output)
-    jobs, scopes = [], set()
+    jobs, scopes, record = [], set(), {}
     for channel_name in channels or config.channels:
         channel = analysis.channel(channel_name)
-        targets_ = targets(channel, control, variables)
+        targets_, record[channel_name] = resolve(config, analysis, channel_name, targets(channel, control, variables))
         for booking in bookings(channel, systematics, regions):
             if processes and booking.process.name not in processes:
                 continue
@@ -181,4 +186,7 @@ def run_hist(config: RunConfig, analysis: Analysis, channels: list[str] | None,
             raise RuntimeError(f"{job[2]}/{job[3].process.name}: {error}") from error
         hset.update(result)
     hset.save(output)
+    record = {channel: entries for channel, entries in record.items() if entries}
+    if record:
+        write_record(output.parent / RECORD, record)
     return hset

src/shapesmith/measurements/tau_id_es/plots.py

@@ -89,7 +89,9 @@ def get(process: str, variation: str = NOMINAL_VARIATION):
         return hset.get(HistKey(CHANNEL, category, process, NOMINAL, variation, variable.name))
 
     backgrounds = [(name, h) for name in channel.backgrounds() if (h := get(name)) is not None]
-    edges, widths = np.asarray(variable.edges), np.diff(variable.edges)
+    data = get(channel.data())
+    edges = data.edges  # those of the filled histograms: an EqualData variable of the analysis carries no numbers
+    widths = np.diff(edges)
     total = np.sum([h.values for _, h in backgrounds], axis=0)
     signal = get(SIGNAL)
     fig, ax = plt.subplots(figsize=(10, 8))
@@ -98,7 +100,6 @@ def get(process: str, variation: str = NOMINAL_VARIATION):
         shifted = get(SIGNAL, grid_name(shift))
         if shifted is not None:
             hep.histplot((total - signal.values + shifted.values) / widths, bins=edges, histtype="step", linewidth=2, label=f"ES {shift / 10:+.1f} %", ax=ax)
-    data = get(channel.data())
     hep.histplot(data.values / widths, bins=edges, yerr=np.sqrt(data.variances) / widths, histtype="errorbar", color="black", label="Data", ax=ax)
     ax.set_ylabel("dN / dm$_{vis}$ [1/GeV]")
     ax.set_xlabel(r"$m_{vis}$ [GeV]")

src/shapesmith/model.py

@@ -86,11 +86,21 @@ def __post_init__(self):
 NOMINAL_REGION = Region(NOMINAL)
 
 
+@dataclass(frozen=True)
+class EqualData:
+    """Edges from the data of the run (binning.py): `n_bins` bins with equal data counts in (low, high), computed from
+    the data of the category in the nominal region whenever histograms are filled."""
+
+    n_bins: int
+    low: float
+    high: float
+
+
 @dataclass(frozen=True)
 class Variable:
     name: str
     expr: str
-    edges: tuple[float, ...]
+    edges: tuple[float, ...] | EqualData
 
 
 @dataclass(frozen=True)

src/shapesmith/validate.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from shapesmith.histogram import part_of
-from shapesmith.model import NOMINAL, ROLES, SAMPLE_KINDS, Analysis, AnalysisError, Channel, ColumnVariation, DataMinus, TemplateShift, VariationSum, WeightVariation
+from shapesmith.model import NOMINAL, ROLES, SAMPLE_KINDS, Analysis, AnalysisError, Channel, ColumnVariation, DataMinus, EqualData, TemplateShift, VariationSum, WeightVariation
 
 
 def _duplicates(what: str, names: list[str]) -> list[str]:
@@ -93,11 +93,26 @@ def _estimator_problems(channel: Channel, region_names: set[str]) -> list[str]:
     return problems
 
 
+def _binning_problems(channel: Channel) -> list[str]:
+    problems = []
+    for variable in [c.variable for c in channel.categories] + list(channel.variables.values()):
+        rule = variable.edges
+        if not isinstance(rule, EqualData):
+            continue
+        if rule.n_bins < 1:
+            problems.append(f"variable {variable.name}: equal-data binning needs n_bins >= 1")
+        if not rule.low < rule.high:
+            problems.append(f"variable {variable.name}: equal-data binning needs low < high")
+        if not any(p.role == "data" for p in channel.processes):
+            problems.append(f"variable {variable.name}: equal-data binning needs a data process")
+    return problems
+
+
 def channel_problems(channel: Channel) -> list[str]:
     region_names = {NOMINAL, *(r.name for r in channel.regions)}
     problems = _process_problems(channel) + _region_problems(channel)
     problems += _duplicates("category", [c.name for c in channel.categories])
-    problems += _variation_problems(channel, region_names) + _estimator_problems(channel, region_names)
+    problems += _variation_problems(channel, region_names) + _estimator_problems(channel, region_names) + _binning_problems(channel)
     return [f"channel {channel.name}: {problem}" for problem in problems]
 
 

tests/test_binning.py

@@ -0,0 +1,105 @@
+import dataclasses
+import json
+
+import numpy as np
+import pytest
+
+from shapesmith.binning import equal_data_edges
+from shapesmith.fill import Target, bookings, fill_booking, run_hist
+from shapesmith.histogram import INCLUSIVE, HistKey
+from shapesmith.model import Category, EqualData, Variable
+from shapesmith.store import read_skims
+
+RULE = EqualData(4, 20.0, 180.0)
+
+
+def smhtt_ul_edges(values, percentiles, low=30.0, high=160.0):
+    """smhtt_ul gof/build_binning.py (tauID_SFs_dev): the m_vis cut of get_data_selection, then get_1d_binning."""
+    arr = values[(low < values) & (values < high)]
+    edges = np.percentile(arr, percentiles)
+    edges = sorted(set(float(x) for x in edges))
+    edges = [e - 1e-4 for e in edges]
+    edges[-1] += 2e-4
+    return edges
+
+
+@pytest.mark.parametrize("percentiles", [
+    [0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0],  # the three lists of build_binning.py
+    [0.0, 12.5, 25.0, 37.5, 50.0, 62.5, 75.0, 87.5, 100.0],
+    [0.0, 20.0, 40.0, 60.0, 80.0, 100.0],
+])
+def test_edges_are_those_of_smhtt_ul(percentiles):
+    values = np.random.default_rng(7).gamma(9.0, 9.0, 20000)
+    assert list(equal_data_edges(values, EqualData(len(percentiles) - 1, 30.0, 160.0))) == smhtt_ul_edges(values, percentiles)
+
+
+def test_edges_keep_the_range_open_and_drop_duplicates():
+    assert equal_data_edges(np.array([30.0, 40.0, 50.0, 160.0, np.nan]), EqualData(1, 30.0, 160.0)) == pytest.approx((40.0 - 1e-4, 50.0 + 1e-4), abs=1e-12)
+    assert equal_data_edges(np.repeat([1.0, 2.0, 3.0], 100), EqualData(10, 0.0, 5.0)) == pytest.approx((1.0 - 1e-4, 2.0 - 1e-4, 3.0 + 1e-4), abs=1e-12)
+    with pytest.raises(ValueError, match=r"no data in \(30.0, 160.0\)"):
+        equal_data_edges(np.array([10.0, 200.0]), EqualData(10, 30.0, 160.0))
+    with pytest.raises(ValueError, match="fewer than two distinct edges from 3 values"):
+        equal_data_edges(np.array([50.0, 50.0, 50.0]), EqualData(10, 30.0, 160.0))
+
+
+def _with_categories(analysis, categories, variables=None):
+    channel = analysis.channel("mt")
+    channel = dataclasses.replace(channel, categories=categories, variables=channel.variables if variables is None else variables)
+    return dataclasses.replace(analysis, channels={"mt": channel})
+
+
+def _equal_data(analysis, rule=RULE):
+    m_vis = Variable("m_vis", "m_vis", rule)
+    return _with_categories(analysis, (Category("sig", "cls == 0", m_vis), Category("bkg", "cls == 1", m_vis)))
+
+
+def test_run_hist_fills_every_region_and_variation_into_the_edges_of_the_data(mini_run):
+    config, analysis = mini_run
+    hset = run_hist(config, _equal_data(analysis), ["mt"], control=False, variables=None, systematics=True, processes=None, output=config.output_dir / "shapes.root")
+
+    data = read_skims(config.skim_dir, "mt", ["DATA_A"], None)
+    nominal = (data["veto"] < 0.5) & ((data["q_1"] * data["q_2"]) < 0) & (data["id_tight"] > 0.5)
+    record = json.loads((config.output_dir / "binning.json").read_text())
+    for cls, category in enumerate(("sig", "bkg")):
+        values = data.loc[nominal & (data["cls"] == cls), "m_vis"].to_numpy()
+        expected = equal_data_edges(values, RULE)
+        in_range = int(((values > RULE.low) & (values < RULE.high)).sum())
+        keys = hset.select(category=category)
+        assert {(k.region, k.variation) for k in keys} >= {("nominal", "Nominal"), ("anti_iso", "Nominal"), ("nominal", "CMS_puUp")}
+        assert all(hset[k].edges.tolist() == list(expected) for k in keys)
+        assert hset[HistKey("mt", category, "data", "nominal", "Nominal", "m_vis")].sum() == pytest.approx(in_range)  # every data value is inside
+        assert record["mt"][category]["m_vis"] == {"rule": {"n_bins": 4, "low": 20.0, "high": 180.0}, "n_data": in_range, "edges": list(expected)}
+
+
+def test_the_record_keeps_the_entries_of_other_runs(mini_run):
+    config, analysis = mini_run
+    analysis = _equal_data(analysis)
+    run_hist(config, analysis, ["mt"], False, None, False, ["data"], config.output_dir / "shapes.root")
+    control = dataclasses.replace(analysis.channel("mt"), variables={"m_vis": Variable("m_vis", "m_vis", EqualData(3, 0.0, 200.0))})
+    run_hist(config, dataclasses.replace(analysis, channels={"mt": control}), ["mt"], True, None, False, ["data"], config.output_dir / "control_shapes.root")
+
+    record = json.loads((config.output_dir / "binning.json").read_text())
+    assert sorted(record["mt"]) == sorted(["sig", "bkg", INCLUSIVE])
+    assert record["mt"][INCLUSIVE]["m_vis"]["rule"] == {"n_bins": 3, "low": 0.0, "high": 200.0}
+
+
+def test_fixed_edges_write_no_record(mini_run):
+    config, analysis = mini_run
+    run_hist(config, analysis, ["mt"], False, None, False, ["data"], config.output_dir / "plain" / "shapes.root")
+    assert not (config.output_dir / "plain" / "binning.json").exists()
+
+
+def test_a_category_without_data_names_itself(mini_run):
+    config, analysis = mini_run
+    empty = _with_categories(analysis, (Category("none", "cls == 7", Variable("m_vis", "m_vis", RULE)),))
+    with pytest.raises(ValueError, match=r"mt/none/m_vis: equal-data binning: no data in \(20.0, 180.0\)"):
+        run_hist(config, empty, ["mt"], False, None, False, None, config.output_dir / "shapes.root")
+
+
+def test_fill_booking_refuses_unresolved_edges(mini_run):
+    config, analysis = mini_run
+    analysis = _equal_data(analysis)
+    channel = analysis.channel("mt")
+    booking = bookings(channel, False)[0]
+    with pytest.raises(ValueError, match="mt: the equal-data binning of m_vis, m_vis is not resolved"):
+        fill_booking(config, analysis, "mt", booking, [Target(c.name, c.cut, c.variable) for c in channel.categories])

tests/test_tau_id_es.py

@@ -14,7 +14,7 @@
 from shapesmith.measurements.tau_id_es.grid import CATEGORIES, factor, grid, grid_name, mass, shift_of
 from shapesmith.measurements.tau_id_es.payload import correction_set
 from shapesmith.measurements.tau_id_es.synced import shapes_path, signal_name, write_synced
-from shapesmith.model import Analysis, Category, Channel, ColumnVariation, DataMinus, Process, Sample, Variable
+from shapesmith.model import Analysis, Category, Channel, ColumnVariation, DataMinus, EqualData, Process, Sample, Variable
 
 GRID = grid(-200, 200, 2)  # the predecessor's grid: +-20 % in steps of 0.2 %
 M_VIS = Variable("m_vis", "m_vis", (30.0, 60.0, 90.0))
@@ -204,6 +204,10 @@ def test_plots_are_written(tmp_path):
     for process, variation, value in (("data", "Nominal", 3.0), ("EMB", "Nominal", 2.0), ("EMB", "es-2", 1.8), ("ZL", "Nominal", 1.0)):
         hset[HistKey("mt", "DM0", process, "nominal", variation, "m_vis")] = _h(value)
     assert plots.plot_control(hset, _analysis(), "DM0", tmp_path / "control")[0].exists()
+    analysis = _analysis()
+    equal_data = Variable("m_vis", "m_vis", EqualData(2, 30.0, 90.0))  # the edges come from the histograms
+    mt = dataclasses.replace(analysis.channel("mt"), categories=(Category("DM0", "tau_decaymode_2 == 0", equal_data),))
+    assert plots.plot_control(hset, dataclasses.replace(analysis, channels={**analysis.channels, "mt": mt}), "DM0", tmp_path / "equal_data")[0].exists()
     with uproot.recreate(tmp_path / "postfit_shapes.root") as f:
         for name, value in (("data_obs", 3.0), ("EMB_DM0", 2.0), ("ZL", 1.0), ("TotalProcs", 3.0)):
             f[f"htt_mt_7_Run2018_postfit/{name}"] = _h(value).to_root(name)

tests/test_validate.py

@@ -122,3 +122,16 @@ def test_validate_checks_variation_sums():
         with pytest.raises(model.AnalysisError, match=problem):
             validate(_with_channel(analysis, variations=channel.variations + parts, estimators=(estimator,)))
 
+
+def test_validate_checks_equal_data_binning():
+    analysis = build()
+    channel = analysis.channel("mt")
+    good = model.Variable("m_vis", "m_vis", model.EqualData(10, 30.0, 160.0))
+    validate(_with_channel(analysis, categories=(model.Category("sig", "cls == 0", good),)))
+    bad = model.Variable("m_vis", "m_vis", model.EqualData(0, 160.0, 30.0))
+    with pytest.raises(model.AnalysisError) as error:
+        validate(_with_channel(analysis, categories=(model.Category("sig", "cls == 0", bad),)))
+    assert "equal-data binning needs n_bins >= 1" in str(error.value) and "equal-data binning needs low < high" in str(error.value)
+    no_data = tuple(p for p in channel.processes if p.role != "data")
+    with pytest.raises(model.AnalysisError, match="variable m_vis: equal-data binning needs a data process"):
+        validate(_with_channel(analysis, processes=no_data, estimators=(), variables={"m_vis": good}))