diff --git a/README.md b/README.md index 4844118..b8f2b58 100644 --- a/README.md +++ b/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). diff --git a/src/shapesmith/binning.py b/src/shapesmith/binning.py new file mode 100644 index 0000000..b44b155 --- /dev/null +++ b/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 diff --git a/src/shapesmith/fill.py b/src/shapesmith/fill.py index 0568538..29b1871 100644 --- a/src/shapesmith/fill.py +++ b/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 diff --git a/src/shapesmith/measurements/tau_id_es/plots.py b/src/shapesmith/measurements/tau_id_es/plots.py index 8f6d578..a9a71e3 100644 --- a/src/shapesmith/measurements/tau_id_es/plots.py +++ b/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]") diff --git a/src/shapesmith/model.py b/src/shapesmith/model.py index 2992370..d0b18b0 100644 --- a/src/shapesmith/model.py +++ b/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) diff --git a/src/shapesmith/validate.py b/src/shapesmith/validate.py index 9b0ff44..0a6e8f1 100644 --- a/src/shapesmith/validate.py +++ b/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] diff --git a/tests/test_binning.py b/tests/test_binning.py new file mode 100644 index 0000000..dd50a30 --- /dev/null +++ b/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]) diff --git a/tests/test_tau_id_es.py b/tests/test_tau_id_es.py index 40ea60b..80b41c0 100644 --- a/tests/test_tau_id_es.py +++ b/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) diff --git a/tests/test_validate.py b/tests/test_validate.py index cee72fe..df48776 100644 --- a/tests/test_validate.py +++ b/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}))