diff --git a/changelog.md b/changelog.md index f084858..e40fbfd 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,7 @@ ### [Latest] +- Add `--metadata` stage (metadata injection + physicalWeight) and `global.keep_all_variables` full-ntuple passthrough for centrally produced inputs; `--plot` weights distributions by `physicalWeight` (capped) × reweight columns when present so plots reflect the reweight route [#142](https://github.com/umami-hep/umami-preprocessing/pull/142) ### [v0.3.1](https://github.com/umami-hep/umami-preprocessing/releases/tag/v0.3.1) (19.06.2026) - Make skip-resampling work end-to-end; support `num_jets: -1` to write all jets passing cuts, and record the resampling method in the output metadata [#153](https://github.com/umami-hep/umami-preprocessing/pull/153) diff --git a/docs/run.md b/docs/run.md index b78d987..204d0ba 100644 --- a/docs/run.md +++ b/docs/run.md @@ -134,6 +134,42 @@ The normalize stage (`--norm`) calculates scaling and shifting values for all va The plotting stage (`--plot`) produces histograms of resampled variables to verify the resampling quality. You can find these plots in `/plots/`. +### Centrally produced inputs: metadata + reweighting workflow + +Centrally produced `.h5` ntuples from TDD often lack pre-computed cross-section metadata and a `physicalWeight` column. For these inputs UPP provides an alternative chain that replaces resample/merge with a metadata injection and a histogram-based reweighting: + +`--metadata` → `--split-components` → `--reweight` (`--rw`) → `--rw-merge` (`--rwm`) → `--norm` + +An example config is provided in [`upp/configs/MetadataRW/metadata.yaml`](https://github.com/umami-hep/umami-preprocessing/blob/main/upp/configs/MetadataRW/metadata.yaml). + +#### Metadata injection (`--metadata`) + +The metadata stage injects cross section, k-factor and generator filter efficiency (via `ftag.find_metadata.MetadataFinder`) and appends a `physicalWeight` column to the `jets` dataset: + +$$\text{physicalWeight} = \frac{\text{XS} \times \text{eff} \times \text{kfactor}}{\text{SOW}} \times \text{mcEventWeight}$$ + +```bash +preprocess --config upp/configs/MetadataRW/metadata.yaml --metadata +``` + +!!!warning "In-place modification" + The metadata stage edits the input files in place (a `.bak` backup is created per file and removed on success). Run it on **workspace copies**, not on central originals. If the sum-of-weights is zero or any metadata value is non-finite the stage raises, and any per-file failures are reported as a summary at the end of the run. + +The downstream reweighting stage automatically uses the `physicalWeight` column when present, and otherwise falls back to uniform weights. + +On this route the plotting stage (`--plot`) also detects `physicalWeight` and the rw-merge weight columns and weights its histograms by their product (`physicalWeight` is capped at the reweight `WEIGHT_CAP` for consistency with the reweight histograms). When neither column is present (the default resampling route) plots remain unweighted. + +#### Full-ntuple passthrough (`global.keep_all_variables`) + +By default the split and rw-merge outputs only keep the variables declared in your `variables` config. Set `keep_all_variables: true` under `global` to instead preserve **all** top-level HDF5 datasets (e.g. `tracks`, `towers`, `flow`) and full `jets` fields in the train/val/test outputs, while still appending the reweight columns on `jets`: + +```yaml +global: + keep_all_variables: true +``` + +The `test` split always keeps all variables regardless of this flag. + ### Additional Scripts: Initial Sample Check The check for the initial samples from the prepare stage can also be run stand-alone. This is important if you plan to run in parallel mode. To do so, you can simply use the following command: diff --git a/pyproject.toml b/pyproject.toml index 7c65c07..4b1c772 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "puma-hep==0.5.3", "pyyaml-include==1.3", "PyYAML>=6.0.2", + # Needed by ftag.find_metadata (missing from atlas-ftag-tools' own deps) + "requests>=2.32", "rich>=14.1.0", "scipy>=1.15.3", ] diff --git a/tests/integration/test_run.py b/tests/integration/test_run.py index d86a1d6..5737b14 100644 --- a/tests/integration/test_run.py +++ b/tests/integration/test_run.py @@ -9,6 +9,7 @@ import numpy as np import pytest from ftag import get_mock_file +from numpy.lib import recfunctions as rfn from upp.main import main @@ -16,8 +17,20 @@ class TestClass: - def generate_mock(self, out_file, N=100_000): + def generate_mock(self, out_file, N=100_000, with_physical_weight=False): + # Inject physicalWeight only on demand; the weighted route is covered in test_run_rw.py. _, f = get_mock_file(num_jets=N, fname=out_file) + if with_physical_weight: + jets = f["jets"][:] + if "physicalWeight" not in jets.dtype.names: + jets2 = rfn.append_fields( + jets, + "physicalWeight", + np.ones(jets.shape[0], dtype="f4"), + usemask=False, + ) + del f["jets"] + f.create_dataset("jets", data=jets2) f.close() def setup_method(self, method): diff --git a/tests/integration/test_run_rw.py b/tests/integration/test_run_rw.py index 1c07c8c..8b9ec9a 100644 --- a/tests/integration/test_run_rw.py +++ b/tests/integration/test_run_rw.py @@ -7,6 +7,7 @@ import h5py import numpy as np from ftag.mock import JET_VARS, get_mock_file +from numpy.lib import recfunctions as rfn from upp.main import main @@ -18,7 +19,25 @@ class TestRunRW: def generate_mock(self, out_file, N=1_000): _, f = get_mock_file(num_jets=N, fname=out_file) - f["jets"]["eventNumber"] = np.arange(N, dtype="i4") + jets = f["jets"][:] + if "eventNumber" in jets.dtype.names: + jets["eventNumber"] = np.arange(N, dtype="i4") + else: + jets = rfn.append_fields( + jets, + "eventNumber", + np.arange(N, dtype="i4"), + usemask=False, + ) + if "physicalWeight" not in jets.dtype.names: + jets = rfn.append_fields( + jets, + "physicalWeight", + np.ones(jets.shape[0], dtype="f4"), + usemask=False, + ) + del f["jets"] + f.create_dataset("jets", data=jets) f.close() def setup_method(self, method): diff --git a/tests/unit/stages/test_metadata_injector.py b/tests/unit/stages/test_metadata_injector.py new file mode 100644 index 0000000..b42cbfd --- /dev/null +++ b/tests/unit/stages/test_metadata_injector.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from pathlib import Path + +import h5py +import numpy as np +import pytest + +import upp.stages.metadata_injector as mi + + +def _make_injector_file(tmp_path, jets_dtype, set_attrs=False): + """Create a minimal HDF5 file for MetadataInjector tests.""" + fpath = tmp_path / "in.h5" + n = 10 + jets = np.zeros(n, dtype=jets_dtype) + if "mcEventWeight" in jets.dtype.names: + jets["mcEventWeight"] = np.arange(n, dtype="f4") + 1.0 + sow = float(np.sum(jets["mcEventWeight"])) if "mcEventWeight" in jets.dtype.names else 55.0 + with h5py.File(fpath, "w") as f: + ds = f.create_dataset("jets", data=jets) + if set_attrs: + ds.attrs["description"] = "test_attr" + ds.attrs["version"] = 42 + cb = f.create_group("cutBookkeeper").create_group("nominal") + cb.create_dataset("counts", data=np.array([sow], dtype="f8")) + md = f.create_group("metadata").create_group("dummy_dsid") + md.create_dataset("cross_section_pb", data=np.array(2.0)) + md.create_dataset("genFiltEff", data=np.array(0.5)) + md.create_dataset("kfactor", data=np.array(1.0)) + return fpath + + +class _Cfg: + def __init__(self, path: Path) -> None: + self.config = {"inputs": {"train": {"input_files": [str(path)]}}} + + +def _stub_finder(monkeypatch): + class _Finder: + def __init__(self, *_a, **_k) -> None: + pass + + def inject_metadata(self) -> None: + return None + + monkeypatch.setattr(mi, "MetadataFinder", _Finder, raising=False) + + +def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): + """Smoke-test MetadataInjector: physicalWeight is appended and fields are preserved.""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")]) + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + out = f["jets"][:] + assert "physicalWeight" in out.dtype.names + assert "mcEventWeight" in out.dtype.names + + +def test_metadata_injector_missing_mcEventWeight_recovers(tmp_path, monkeypatch): + """Exception path: missing mcEventWeight restores the backup and a summary is raised.""" + _stub_finder(monkeypatch) + # jets without mcEventWeight triggers the per-file failure + restore path + fpath = _make_injector_file(tmp_path, [("someField", "f4")]) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + injector = mi.MetadataInjector(_Cfg(fpath)) + # The file is restored from backup, but failures are surfaced as a summary error. + with pytest.raises(RuntimeError, match="Metadata injection failed"): + injector.run() + + assert fpath.exists() + assert not backup_path.exists() + + +def test_metadata_injector_zero_sow_raises(tmp_path, monkeypatch): + """A zero sum-of-weights is rejected instead of producing inf/nan weights.""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")]) + # Force the stored sum-of-weights to zero. + with h5py.File(fpath, "a") as f: + del f["cutBookkeeper/nominal/counts"] + f["cutBookkeeper/nominal"].create_dataset("counts", data=np.array([0.0], dtype="f8")) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + injector = mi.MetadataInjector(_Cfg(fpath)) + with pytest.raises(RuntimeError, match="Metadata injection failed"): + injector.run() + + assert fpath.exists() + assert not backup_path.exists() + + +def test_metadata_injector_drops_existing_physicalWeight(tmp_path, monkeypatch): + """Existing physicalWeight field is dropped before recomputing (line 73).""" + _stub_finder(monkeypatch) + dtype = [("mcEventWeight", "f4"), ("physicalWeight", "f4")] + fpath = _make_injector_file(tmp_path, dtype) + # Overwrite physicalWeight with sentinel value to verify it gets replaced + with h5py.File(fpath, "a") as f: + data = f["jets"][:] + data["physicalWeight"] = 999.0 + del f["jets"] + f.create_dataset("jets", data=data) + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + out = f["jets"][:] + assert "physicalWeight" in out.dtype.names + assert not np.all(out["physicalWeight"] == 999.0) + + +def test_metadata_injector_preserves_dataset_attrs(tmp_path, monkeypatch): + """Original dataset attrs are restored after jets dataset recreation (lines 84-85).""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")], set_attrs=True) + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + assert f["jets"].attrs["description"] == "test_attr" + assert f["jets"].attrs["version"] == 42 diff --git a/tests/unit/stages/test_normalisation.py b/tests/unit/stages/test_normalisation.py index f2bbaee..450cff2 100644 --- a/tests/unit/stages/test_normalisation.py +++ b/tests/unit/stages/test_normalisation.py @@ -4,6 +4,7 @@ import subprocess from pathlib import Path +import numpy as np import pytest from ftag import get_mock_file @@ -105,6 +106,30 @@ def test_combine_mean_std(): assert combined_mean_ref == combined_mean assert combined_std_ref == combined_std + def test_get_class_dict_integer_label(self): + """Integer-typed label vars are counted in get_class_dict (lines 160-161).""" + norm = Normalisation( + config=PreprocessingConfig.from_file( + Path(CFG_DIR / "test_config_pdf_auto_umami.yaml"), "train" + ) + ) + # config labels for jets: [pt, eta, mass] + resampling vars [pt, abs_eta] + # → effective labels: [pt, eta, mass, abs_eta]; make "pt" integer → lines 160-161 + jets_dtype = np.dtype( + [ + ("pt_btagJes", "f4"), + ("eta_btagJes", "f4"), + ("flavour_label", "i4"), + ("pt", "i4"), + ("eta", "f4"), + ("mass", "f4"), + ("abs_eta", "f4"), + ] + ) + batch = {"jets": np.zeros(10, dtype=jets_dtype)} + result = norm.get_class_dict(batch) + assert "pt" in result["jets"] + @staticmethod def test_combine_norm_dict(): # Test combination of mean and std diff --git a/tests/unit/stages/test_plotting.py b/tests/unit/stages/test_plotting.py index 54420d9..68bf160 100644 --- a/tests/unit/stages/test_plotting.py +++ b/tests/unit/stages/test_plotting.py @@ -210,6 +210,85 @@ def test_post_resampling_paths_split_mode(tmp_path): assert paths == [tmp_path / "test" / "pp_output_test*.h5"] +def test_reweight_weight_fields_detects_present_columns(): + """Check reweight-route weight columns are picked up only when present.""" + from upp.classes.reweight_config import SingleReweightConfig + + rw = SingleReweightConfig( + group="jets", + reweight_vars=["pt"], + bins={"pt": [[20_000, 250_000, 5]]}, + class_var="flavour_label", + class_target="uniform", + ) + rw_name = repr(rw) + assert rw_name == "weight_jets_pt_target_uniform_flavour_label" + config = SimpleNamespace(rw_config=SimpleNamespace(reweights=[rw])) + + # both physicalWeight and the rw-merge column present + assert plot_mod._reweight_weight_fields(config, {"physicalWeight", rw_name, "pt"}) == [ + "physicalWeight", + rw_name, + ] + # rw column absent (e.g. raw ntuple after --metadata) -> only physicalWeight + assert plot_mod._reweight_weight_fields(config, {"physicalWeight", "pt"}) == ["physicalWeight"] + # default resampling route: no rw_config and no physicalWeight -> unweighted + assert plot_mod._reweight_weight_fields(SimpleNamespace(), {"pt"}) == [] + + +def test_make_hist_applies_capped_weights(monkeypatch, tmp_path): + """Check make_hist multiplies weight columns and caps physicalWeight.""" + import numpy as np + + from upp.stages.reweight import WEIGHT_CAP + + dtype = [("pt", "f4"), ("HadronConeExclTruthLabelID", "i4"), ("physicalWeight", "f4")] + arr = np.zeros(5, dtype=dtype) + arr["pt"] = [10_000, 20_000, 30_000, 40_000, 50_000] + arr["HadronConeExclTruthLabelID"] = 5 # bjets + arr["physicalWeight"] = [1.0, 2.0, 3.0 * WEIGHT_CAP, 4.0, 5.0] + + captured = {} + + class FakeHist: + def __init__(self, **kwargs): + captured["weights"] = kwargs.get("weights") + self.bin_edges = np.array([0.0, 100.0]) + + class FakePlot: + def __init__(self, **kwargs): + pass + + def add(self, **kwargs): + pass + + def draw(self): + pass + + def make_linestyle_legend(self, **kwargs): + pass + + def savefig(self, _path): + pass + + monkeypatch.setattr(plot_mod, "Histogram", FakeHist) + monkeypatch.setattr(plot_mod, "HistogramPlot", FakePlot) + + make_hist( + stage="initial", + values_dict={"": arr}, + flavours=[Flavours["bjets"]], + variable="pt", + out_dir=tmp_path, + out_format_list=["png"], + weight_fields=["physicalWeight"], + ) + + expected = np.clip(arr["physicalWeight"].astype(float), 0.0, WEIGHT_CAP) + assert captured["weights"] is not None + assert np.allclose(captured["weights"], expected) + + def test_plot_initial_uses_split_suffix_and_plotting_jet_count(monkeypatch, tmp_path): """Check initial plot calls include split-specific suffixes and plotting counts.""" diff --git a/tests/unit/stages/test_rw_merge.py b/tests/unit/stages/test_rw_merge.py index eb60cc4..6cd8df4 100644 --- a/tests/unit/stages/test_rw_merge.py +++ b/tests/unit/stages/test_rw_merge.py @@ -2,9 +2,12 @@ from __future__ import annotations +from unittest.mock import MagicMock, patch + import numpy as np import pytest +import upp.stages.rw_merge as rw_merge_module from upp.stages.rw_merge import RWMerge @@ -54,3 +57,78 @@ def test_empty_input(self): result = RWMerge._assign_weights(rw, bins, classes) assert result.shape == (0,) + + +class TestRWMergeInit: + """Cover RWMerge.__init__ assertion for outfile_idx_range (line 26).""" + + def test_non_tuple_idx_range_raises(self): + config = MagicMock() + with pytest.raises(AssertionError): + RWMerge(config, outfile_idx_range=[0, 1]) + + def test_wrong_length_tuple_raises(self): + config = MagicMock() + with pytest.raises(AssertionError): + RWMerge(config, outfile_idx_range=(0, 1, 2)) + + +class TestStartMp: + """Cover start_mp multiprocess branch (lines 314-316).""" + + def test_start_mp_multiprocess(self): + called_with: list = [] + + class FakePool: + def __init__(self, n: int) -> None: + pass + + def __enter__(self) -> FakePool: + return self + + def __exit__(self, *_: object) -> None: + pass + + def starmap(self, _fn: object, args_list: list) -> None: + for args in args_list: + called_with.append(args) + + def fn(x: int) -> int: + return x + + with patch.object(rw_merge_module, "Pool", FakePool): + RWMerge.start_mp(fn, [(1,), (2,), (3,)], n_threads=2) + + assert called_with == [(1,), (2,), (3,)] + + +class TestGetSampleWeightsException: + """Cover the except-and-reraise path in get_sample_weights (lines 185-187).""" + + def test_exception_propagates(self): + n = 5 + jets = np.zeros(n, dtype=[("x", "f4"), ("class_var", "i4")]) + batch = {"jets": jets} + weights = { + "jets": { + "rw1": { + "rw_vars": ["x"], + "class_var": "class_var", + "bins": [np.linspace(0.0, 1.0, 6)], + "weights": {"0": np.ones(5)}, + } + } + } + + def fake_bin_jets(*_a: object, **_k: object) -> tuple: + return (np.zeros(5), np.zeros((1, n), dtype=int)) + + def bad_assign(*_a: object, **_k: object) -> None: + raise ValueError("injected error") + + with ( + patch.object(rw_merge_module, "bin_jets", fake_bin_jets), + patch.object(RWMerge, "_assign_weights", staticmethod(bad_assign)), + pytest.raises(ValueError, match="injected error"), + ): + RWMerge.get_sample_weights(batch, weights) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 896652f..bcd1e98 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1,10 +1,12 @@ from __future__ import annotations +import argparse from argparse import Namespace +from unittest.mock import MagicMock, patch from pytest import fixture -from upp.main import parse_args +from upp.main import main, parse_args, run_pp @fixture @@ -25,6 +27,7 @@ def test_parse_args_with_config(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=None, resample=True, merge=None, @@ -51,6 +54,7 @@ def test_parse_args_flags_not_given(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -88,6 +92,7 @@ def test_parse_args_flags_negative(config_file): # Check if the parsed arguments match the expected values expected_args = Namespace( config=config_file, + metadata=False, prep=False, resample=False, merge=False, @@ -123,6 +128,7 @@ def test_parse_args_flags_positive(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -160,6 +166,7 @@ def test_parse_args_component(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -197,6 +204,7 @@ def test_parse_args_region(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -216,3 +224,81 @@ def test_parse_args_region(config_file): ) assert parsed_args == expected_args + + +def test_parse_args_metadata_flag(config_file): + args = ["--config", str(config_file), "--metadata"] + parsed_args = parse_args(args) + assert parsed_args.metadata is True + + +def _base_args(config_file: object, **overrides: object) -> argparse.Namespace: + defaults: dict = dict( + config=config_file, + metadata=False, + prep=False, + resample=False, + merge=False, + norm=False, + plot=False, + split="train", + component=None, + region=None, + container=None, + grid=False, + split_components=False, + reweight=False, + rw_merge=False, + rw_merge_idx=None, + files=None, + skip_sample_check=False, + ) + defaults.update(overrides) + return argparse.Namespace(**defaults) + + +def test_run_pp_metadata_injection(tmp_path): + """run_pp with metadata=True constructs and runs MetadataInjector (lines 222-224).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + args = _base_args(config_file, metadata=True) + + mock_injector = MagicMock() + with ( + patch("upp.main.PreprocessingConfig.from_file", return_value=MagicMock()), + patch("upp.main.MetadataInjector", return_value=mock_injector), + ): + run_pp(args) + + mock_injector.run.assert_called_once() + + +def test_run_pp_rw_merge_with_idx(tmp_path): + """run_pp with rw_merge_idx parses the comma-separated pair (lines 269-272).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + args = _base_args(config_file, rw_merge=True, rw_merge_idx="0,1") + + mock_rw = MagicMock() + with ( + patch("upp.main.PreprocessingConfig.from_file", return_value=MagicMock()), + patch("upp.main.RWMerge", return_value=mock_rw), + ): + run_pp(args) + + mock_rw.run.assert_called_once() + + +def test_main_split_all(tmp_path): + """main() with split='all' calls run_pp three times (lines 303-310).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + mock_args = _base_args(config_file, split="all") + + with ( + patch("upp.main.parse_args", return_value=mock_args), + patch("upp.main.run_pp") as mock_run_pp, + ): + main() + + assert mock_run_pp.call_count == 3 diff --git a/upp/classes/preprocessing_config.py b/upp/classes/preprocessing_config.py index 1a7b0a9..9004278 100644 --- a/upp/classes/preprocessing_config.py +++ b/upp/classes/preprocessing_config.py @@ -121,6 +121,9 @@ class PreprocessingConfig: Skip checks for the input files. This is used for grid submission skip_config_copy : bool, optional Decide, if the config copying is skipped or not. By default False + keep_all_variables : bool, optional + When true: split-containers and rw-merge keep full jets fields and all top-level + HDF5 datasets (PW and RW columns still added on jets). vds_dir : Path | None, optional Directory name for creation of virtual datasets. By default None If none is given, virtual datasets is created next to input ntuples @@ -147,6 +150,8 @@ class PreprocessingConfig: num_jets_per_output_file: int | None = None skip_checks: bool = False skip_config_copy: bool = False + # Keep all top-level datasets (not just the variables.yaml subset) through split/rw-merge. + keep_all_variables: bool = False vds_dir: Path | None = None def __post_init__(self): @@ -207,9 +212,12 @@ def __post_init__(self): if selection := groups.get("selection", None): selectors[name] = TrackSelector(Cuts.from_list(selection)) - # configure variables + # test split always keeps all variables; keep_all_variables also forces it for train/val. self.variables = VariableConfig( - self.config["variables"], self.jets_name, self.is_test, selectors + self.config["variables"], + self.jets_name, + self.keep_all_variables or self.is_test, + selectors, ) if self.sampl_cfg is not None and self.sampl_cfg.variables: self.variables = self.variables.add_jet_vars( diff --git a/upp/classes/reweight_config.py b/upp/classes/reweight_config.py index a488621..42f41b1 100644 --- a/upp/classes/reweight_config.py +++ b/upp/classes/reweight_config.py @@ -11,12 +11,18 @@ class ReweightConfig: # Number of jets to estimate, if None, use the global num jets estimate num_jets_estimate: None | int = None merge_num_proc: int = 1 # Number of processes to use for merging + # Upper bound for physical weights and reweight factors, + # which can span many orders of magnitude + weight_cap: float = 1e4 reweights: list[SingleReweightConfig] = field(default_factory=list) def __post_init__(self): if self.num_jets_estimate is not None and self.num_jets_estimate <= 0: raise ValueError("num_jets_estimate must be a positive integer or None") + if self.weight_cap <= 0: + raise ValueError("weight_cap must be a positive number") + parsed_reweights = [] for rw in self.reweights: parsed_reweights.append(SingleReweightConfig(**rw)) diff --git a/upp/configs/MetadataRW/flavours-r10.yaml b/upp/configs/MetadataRW/flavours-r10.yaml new file mode 100644 index 0000000..0003bea --- /dev/null +++ b/upp/configs/MetadataRW/flavours-r10.yaml @@ -0,0 +1,44 @@ +# Flavour definitions for the metadata/reweight workflow on centrally produced +# large-R ntuples. Labels are contiguous integers 0-5. + +- name: djets + category: jets + label: "d-jets" + colour: blue + cuts: + - [PartonTruthLabelID, "==", 1] + +- name: ujets + category: jets + label: "u-jets" + colour: cyan + cuts: + - [PartonTruthLabelID, "==", 2] + +- name: sjets + category: jets + label: "s-jets" + colour: purple + cuts: + - [PartonTruthLabelID, "==", 3] + +- name: cjets + category: jets + label: "c-jets" + colour: red + cuts: + - [PartonTruthLabelID, "==", 4] + +- name: bjets + category: jets + label: "b-jets" + colour: orange + cuts: + - [PartonTruthLabelID, "==", 5] + +- name: gluon + category: jets + label: "gluon" + colour: pink + cuts: + - [PartonTruthLabelID, "==", 21] diff --git a/upp/configs/MetadataRW/metadata.yaml b/upp/configs/MetadataRW/metadata.yaml new file mode 100644 index 0000000..0aef604 --- /dev/null +++ b/upp/configs/MetadataRW/metadata.yaml @@ -0,0 +1,88 @@ +global: + base_dir: /eos/home-y/yuanda/dev + # Metadata injection (--metadata) modifies files in-place; use copies under central_working, not origin. + ntuple_dir: /eos/home-y/yuanda/dev/sample/central_working + # When true: split + rw-merge preserve all top-level HDF5 datasets / full jet fields (PW + RW columns still added on jets). + keep_all_variables: true + out_dir: output_metadata + jets_name: jets + batch_size: 40000 + num_jets_estimate: 700000 + num_jets_per_output_file: 1_000_000 + # test split: plot.py uses merged VDS path (glob) instead of per-sample pp_output_test_*.h5 + merge_test_samples: true + flavour_config: /eos/home-y/yuanda/dev/umami-preprocessing/upp/configs/MetadataRW/flavours-r10.yaml + +inputs: + train: + input_files: + - /eos/home-y/yuanda/dev/sample/central_working/*.h5 + +global_cuts: + common: [] + train: [[eventNumber, "%10<=", 7]] + val: [[eventNumber, "%10==", 8]] + test: [[eventNumber, "%10==", 9]] + +sample_all: &sample_all + name: sample_all + pattern: ["*.h5"] + +components: + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 1]] } + sample: { <<: *sample_all } + flavours: [djets] + num_jets: 60000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 2]] } + sample: { <<: *sample_all } + flavours: [ujets] + num_jets: 120000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 3]] } + sample: { <<: *sample_all } + flavours: [sjets] + num_jets: 15000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 4]] } + sample: { <<: *sample_all } + flavours: [cjets] + num_jets: 10000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 5]] } + sample: { <<: *sample_all } + flavours: [bjets] + num_jets: 6000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 21]] } + sample: { <<: *sample_all } + flavours: [gluon] + num_jets: 160000 + +variables: + jets: + inputs: [pt, eta, eventNumber] + labels: [PartonTruthLabelID, physicalWeight] + +resampling: + method: none + target: ujets + variables: + pt: + bins: [[20000, 2480000, 100]] + +reweighting: + num_jets_estimate: 700000 + reweights: + - group: jets + reweight_vars: [pt] + bins: + pt: [[20000, 2480000, 100]] + class_var: flavour_label + class_target: uniform + +flavour_categories: + jets: [djets, ujets, sjets, cjets, bjets, gluon] + +custom_flavours: + djets: { label: 0, cuts: [[PartonTruthLabelID, "==", 1]] } + ujets: { label: 1, cuts: [[PartonTruthLabelID, "==", 2]] } + sjets: { label: 2, cuts: [[PartonTruthLabelID, "==", 3]] } + cjets: { label: 3, cuts: [[PartonTruthLabelID, "==", 4]] } + bjets: { label: 4, cuts: [[PartonTruthLabelID, "==", 5]] } + gluon: { label: 5, cuts: [[PartonTruthLabelID, "==", 21]] } diff --git a/upp/configs/MetadataRW/variables-r10.yaml b/upp/configs/MetadataRW/variables-r10.yaml new file mode 100644 index 0000000..534e559 --- /dev/null +++ b/upp/configs/MetadataRW/variables-r10.yaml @@ -0,0 +1,53 @@ +jets: + inputs: + - pt_btagJes + - eta_btagJes + - R10TruthLabel_R22v1 + - PartonTruthLabelID + - physicalWeight + labels: + - HadronConeExclTruthLabelID + - HadronConeExclExtendedTruthLabelID + - HadronConeExclTruthLabelPt + - HadronConeExclTruthLabelLxy + - HadronConeExclTruthLabelDR + - HadronGhostTruthLabelID + - HadronGhostExtendedTruthLabelID + - HadronGhostTruthLabelPt + - HadronGhostTruthLabelLxy + - HadronGhostTruthLabelDR + - pt + - eta + - mass + - n_tracks + - n_truth_promptLepton + - eventNumber + - jetFoldHash + - physicalWeight + +tracks: + inputs: + - d0 + - z0SinTheta + - dphi + - deta + - qOverP + - lifetimeSignedD0Significance + - lifetimeSignedZ0SinThetaSignificance + - phiUncertainty + - thetaUncertainty + - qOverPUncertainty + - numberOfPixelHits + - numberOfSCTHits + - numberOfInnermostPixelLayerHits + - numberOfNextToInnermostPixelLayerHits + - numberOfInnermostPixelLayerSharedHits + - numberOfInnermostPixelLayerSplitHits + - numberOfPixelSharedHits + - numberOfPixelSplitHits + - numberOfSCTSharedHits + - leptonID + labels: + - ftagTruthOriginLabel + - ftagTruthTypeLabel + - ftagTruthVertexIndex diff --git a/upp/main.py b/upp/main.py index 1dd23b9..9d7a746 100644 --- a/upp/main.py +++ b/upp/main.py @@ -22,6 +22,7 @@ from upp.classes.preprocessing_config import PreprocessingConfig from upp.stages.hist import create_histograms from upp.stages.merging import Merging +from upp.stages.metadata_injector import MetadataInjector from upp.stages.normalisation import Normalisation from upp.stages.plot import plot_resampling_dists from upp.stages.resampling import Resampling @@ -55,6 +56,12 @@ def parse_args(args: Any) -> argparse.Namespace: type=valid_path, help="Path to config file", ) + parser.add_argument( + "--metadata", + action="store_true", + default=False, + help="Run metadata injection stage before reweighting", + ) parser.add_argument( "--prep", action="store_true", @@ -212,6 +219,12 @@ def run_pp(args: argparse.Namespace) -> None: # load config config = PreprocessingConfig.from_file(args.config, args.split, skip_checks=args.grid) + # run metadata injection + if getattr(args, "metadata", False): + log.info("Running metadata injection...") + injector = MetadataInjector(config) + injector.run() + if args.split_components: log.info("Splitting containers...") split = SplitContainers(args.config) @@ -223,6 +236,7 @@ def run_pp(args: argparse.Namespace) -> None: # If we aren't running on the grid, we create the metadata after splitting if not args.grid: split.create_meta_data() + if args.reweight: log.info("Running reweighting...") reweight = Reweight(config) diff --git a/upp/stages/hist.py b/upp/stages/hist.py index 6b37c02..63a3611 100644 --- a/upp/stages/hist.py +++ b/upp/stages/hist.py @@ -18,7 +18,9 @@ from upp.classes.preprocessing_config import PreprocessingConfig -def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: +def bin_jets( + array: dict, bins: list, weights: np.ndarray | None = None +) -> tuple[np.ndarray, np.ndarray]: """Create the histogram and bins for the given resampling variables. Parameters @@ -28,6 +30,9 @@ def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: variables. bins : list Flat list with the bins which are to be used. + weights : np.ndarray | None, optional + Per-jet weights for the histogram (e.g. physicalWeight). When None, bins are filled + by count; otherwise by sum of weights. Returns ------- @@ -42,10 +47,11 @@ def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: `expand_binnumbers` argument. See `Notes` for details. """ sample = s2u(array).astype(np.float64, copy=False) + statistic_mode = "count" if weights is None else "sum" hist, _, out_bins = binned_statistic_dd( sample=sample, - values=None, - statistic="count", + values=weights, + statistic=statistic_mode, bins=bins, expand_binnumbers=True, ) diff --git a/upp/stages/metadata_injector.py b/upp/stages/metadata_injector.py new file mode 100644 index 0000000..4df2d5d --- /dev/null +++ b/upp/stages/metadata_injector.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import glob +import math +import shutil +from pathlib import Path + +import h5py +from ftag.find_metadata import MetadataFinder +from numpy.lib import recfunctions as rfn + +from upp.utils.logger import setup_logger + + +class MetadataInjector: + def __init__(self, config): + self.config = config + self.log = setup_logger() + + def run(self): + # Initialize input file list + input_files = [] + if hasattr(self.config, "config") and "inputs" in self.config.config: + raw_inputs = self.config.config["inputs"] + if "train" in raw_inputs and "input_files" in raw_inputs["train"]: + input_files.extend(raw_inputs["train"]["input_files"]) + + # Expand wildcards/glob patterns + expanded_files = [] + for f in input_files: + matched = glob.glob(f) + expanded_files.extend(matched) + + failures: list[tuple[str, str]] = [] + for fpath_str in expanded_files: + fpath = Path(fpath_str) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + try: + # 0. Create a physical backup + shutil.copy(fpath, backup_path) + + # 1. Inject metadata (this adds a new 'metadata' group) + finder = MetadataFinder(fpath_str) + finder.inject_metadata() + + with h5py.File(fpath_str, "a") as f: + # 2. Calculate parameters required for physical weights + dsid = next(iter(f["metadata"].keys())) + xs = float(f[f"metadata/{dsid}/cross_section_pb"][()]) + eff = float(f[f"metadata/{dsid}/genFiltEff"][()]) + kf = float(f[f"metadata/{dsid}/kfactor"][()]) + + # Read Sum of Weights (SOW) with support for different formats + sow_ds = f["cutBookkeeper/nominal/counts"] + sow = ( + float(sow_ds["sumOfWeights"][0]) if sow_ds.dtype.names else float(sow_ds[0]) + ) + + # Guard against div-by-zero / non-finite metadata (silent inf/nan weights). + if sow == 0 or not math.isfinite(sow): + raise ValueError(f"Invalid sum-of-weights (sow={sow}) in {fpath_str}") + if not all(math.isfinite(v) for v in (xs, eff, kf)): + raise ValueError( + f"Non-finite metadata in {fpath_str}: xs={xs}, eff={eff}, kfactor={kf}" + ) + + # 3. Read original jets and their attributes + old_jets_ds = f["jets"] + original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata + # Preserve source storage layout so rewrite keeps compression/chunking. + create_kwargs = { + "compression": old_jets_ds.compression, + "compression_opts": old_jets_ds.compression_opts, + "chunks": old_jets_ds.chunks, + "shuffle": old_jets_ds.shuffle, + } + jets_data = old_jets_ds[:] + + # 4. Calculate Physical Weight + if "mcEventWeight" not in jets_data.dtype.names: + raise KeyError(f"mcEventWeight missing in {fpath_str}") + + mcw = jets_data["mcEventWeight"].astype("f8") + # Formula: Weight = (XS * Efficiency * k-factor / SOW) * MC_Weight + physical_w = (xs * eff * kf / sow) * mcw + + # 5. Construct the updated structured array + if "physicalWeight" in jets_data.dtype.names: + jets_data = rfn.drop_fields(jets_data, "physicalWeight") + + updated_jets = rfn.append_fields( + jets_data, "physicalWeight", physical_w.astype("f4"), usemask=False + ) + + # 6. Delete and recreate the dataset while restoring attributes + del f["jets"] + new_ds = f.create_dataset("jets", data=updated_jets, **create_kwargs) + + # Restore all original attributes (e.g., descriptions for eventNumber, etc.) + for k, v in original_attrs.items(): + new_ds.attrs[k] = v + + # 7. Processing successful: remove the backup + backup_path.unlink() + self.log.info( + f"Successfully updated {fpath.name} (Attributes and original fields preserved)" + ) + + except Exception as e: + self.log.error(f"Failed for {fpath_str}: {e}") + if backup_path.exists(): + self.log.warning(f"Restoring {fpath.name} from backup...") + shutil.move(backup_path, fpath) + failures.append((fpath_str, str(e))) + + # Re-raise a summary so partial failures are not silently missed in a long run. + if failures: + summary = "\n".join(f" - {path}: {err}" for path, err in failures) + raise RuntimeError( + f"Metadata injection failed for {len(failures)}/{len(expanded_files)} file(s):\n" + f"{summary}" + ) diff --git a/upp/stages/normalisation.py b/upp/stages/normalisation.py index 6ef5446..516cdf6 100644 --- a/upp/stages/normalisation.py +++ b/upp/stages/normalisation.py @@ -259,9 +259,19 @@ def run(self): norm_dict = None class_dict = None total = None - vars = self.variables.combined() + combined_partial = self.variables.combined() + vars: dict[str, list[str] | None] = {} + for name in self.variables: + cols = combined_partial[name] + if cols is None: + # keep_all_variables: explicit stream columns; norm still on yaml inputs/labels + cols = self.variables[name]["inputs"] + self.variables[name].get("labels", []) + vars[name] = list(cols) with h5py.File(reader.files[0]) as f: - if "flavour_label" in f[self.jets_name].dtype.names: + if ( + "flavour_label" in f[self.jets_name].dtype.names + and "flavour_label" not in vars[self.jets_name] + ): vars[self.jets_name].append("flavour_label") stream = reader.stream(vars, self.num_jets) diff --git a/upp/stages/plot.py b/upp/stages/plot.py index ee0c1a1..c8aad12 100644 --- a/upp/stages/plot.py +++ b/upp/stages/plot.py @@ -1,5 +1,6 @@ from __future__ import annotations +import glob import logging as log import re from dataclasses import dataclass @@ -7,6 +8,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +import h5py +import numpy as np from ftag import Cuts from ftag.hdf5 import H5Reader from puma import Histogram, HistogramPlot @@ -418,6 +421,62 @@ def _stitching_regions(regions: list[PlotRegion], pt_variable: str | None) -> li return stitching_regions +def _available_jet_fields(config: PreprocessingConfig, in_paths: Any) -> set[str]: + """Return the jet field names present in the first matching input file. + + Parameters + ---------- + config : PreprocessingConfig + Active preprocessing configuration. + in_paths : Any + Input HDF5 file path, glob, or list of paths. + + Returns + ------- + set[str] + Jet dtype field names of the first existing file, or an empty set. + """ + paths = in_paths if isinstance(in_paths, list) else [in_paths] + for pattern in paths: + for fpath in sorted(glob.glob(str(pattern))): + with h5py.File(fpath, "r") as f: + if config.jets_name in f: + return set(f[config.jets_name].dtype.names or ()) + return set() + + +def _reweight_weight_fields(config: PreprocessingConfig, available: set[str]) -> list[str]: + """Return per-jet weight columns to apply on the reweight route. + + Combines ``physicalWeight`` (physics weight from --metadata) with the + reweight columns written by --rw-merge. Only columns present in ``available`` + are returned, so the default resampling route (which has neither) stays + unweighted. + + Parameters + ---------- + config : PreprocessingConfig + Active preprocessing configuration. + available : set[str] + Jet field names present in the input being plotted. + + Returns + ------- + list[str] + Weight column names to multiply together when histogramming. + """ + fields = [] + if "physicalWeight" in available: + fields.append("physicalWeight") + rw_config = getattr(config, "rw_config", None) + if rw_config is not None: + for rw in rw_config.reweights: + name = repr(rw) + if name in available: + fields.append(name) + return fields + + def _load_jets(config: PreprocessingConfig, in_paths: Any, vars_to_load: list[str]) -> Any: """Load jet variables for plotting. @@ -461,6 +520,8 @@ def make_hist( selection_cuts: Cuts | None = None, atlas_second_tag: str | None = None, plotting: PlottingConfig | None = None, + weight_fields: list[str] | None = None, + weight_cap: float | None = None, ) -> None: """Make a flavour-split histogram for one variable. @@ -502,11 +563,22 @@ def make_hist( centre-of-mass energy label is shown. plotting : PlottingConfig | None, optional Plot labels and style settings. If ``None``, use the defaults. + weight_fields : list[str] | None, optional + Jet weight columns multiplied per jet on the reweight route (e.g. + ``physicalWeight`` and the rw-merge weight column). ``physicalWeight`` is + capped at ``weight_cap`` to match the reweight histograms. If ``None`` or + empty, histograms are unweighted (default resampling behaviour). + weight_cap : float | None, optional + Upper bound applied to ``physicalWeight``. If ``None``, the default + ``WEIGHT_CAP`` from the reweight stage is used. """ from upp.classes.plotting_config import PlottingConfig + from upp.stages.reweight import WEIGHT_CAP selection_cuts = selection_cuts or Cuts.empty() plotting = plotting or PlottingConfig() + if weight_cap is None: + weight_cap = WEIGHT_CAP # Setup the histogram plot = HistogramPlot( @@ -541,9 +613,20 @@ def make_hist( selected_values = cuts(values).values histo_values = _display_values(variable, selected_values[variable]) + # Reweight route: weight per jet by capped physicalWeight x rw-merge columns. + histo_weights = None + if weight_fields: + histo_weights = np.ones(len(selected_values), dtype=np.float64) + for field in weight_fields: + column = selected_values[field].astype(np.float64) + if field == "physicalWeight": + column = np.clip(column, 0.0, weight_cap) + histo_weights *= column + # Get the histogram object histo = Histogram( values=histo_values, + weights=histo_weights, bins=plotting.bins, bins_range=bins_range, norm=plotting.norm, @@ -600,8 +683,12 @@ def _plot_initial(config: PreprocessingConfig) -> None: for flavour in config.components.flavours: vars_to_load += flavour.cuts.variables + in_paths = list(sample.path) + weight_fields = _reweight_weight_fields(config, _available_jet_fields(config, in_paths)) + vars_to_load += weight_fields + stage_status = "Pre Reweighting" if weight_fields else "Pre Resampling" values_dict = { - sample.name: _load_jets(config, list(sample.path), vars_to_load), + sample.name: _load_jets(config, in_paths, vars_to_load), } pt_range = _pt_bounds_from_cuts(selection_cuts, pt_var) if pt_var else None @@ -627,10 +714,11 @@ def _plot_initial(config: PreprocessingConfig) -> None: num_jets=_plotting_num_jets(config, region_components.num_jets) if config.plotting.show_num_jets else None, - resampling_status="Pre Resampling", + resampling_status=stage_status, ), plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, ) @@ -680,6 +768,9 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: if full_region := _full_region(base_regions, pt_var): plot_regions.append(full_region) + in_paths = _post_resampling_paths(config, stage) + weight_fields = _reweight_weight_fields(config, _available_jet_fields(config, in_paths)) + sample_names = [sample.name for sample in config.components.samples] atlas_second_tag = _atlas_second_tag( *sample_names, @@ -687,15 +778,16 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: num_jets=_plotting_num_jets(config, config.components.num_jets) if config.plotting.show_num_jets else None, - resampling_status="Post Resampling", + resampling_status="Post Reweighting" if weight_fields else "Post Resampling", ) vars_to_load = list(config.sampl_cfg.vars) + ["flavour_label"] for region in plot_regions: vars_to_load += region.cuts.variables + vars_to_load += weight_fields values_dict = { - "": _load_jets(config, _post_resampling_paths(config, stage), vars_to_load), + "": _load_jets(config, in_paths, vars_to_load), } for variable in config.sampl_cfg.vars: @@ -718,6 +810,8 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: atlas_second_tag=atlas_second_tag, plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, + weight_cap=config.rw_config.weight_cap if config.rw_config else None, ) if _is_pt_variable(variable): @@ -735,6 +829,8 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: atlas_second_tag=atlas_second_tag, plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, + weight_cap=config.rw_config.weight_cap if config.rw_config else None, ) diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index 99ff3b2..4b49e3b 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -14,6 +14,10 @@ from upp.classes.preprocessing_config import PreprocessingConfig from upp.stages.hist import bin_jets +# Default upper bound for physical weights and reweight factors, which can span +# many orders of magnitude. Configurable via `reweighting.weight_cap`. +WEIGHT_CAP = 1e4 + class Reweight: def __init__(self, config: PreprocessingConfig): @@ -23,6 +27,7 @@ def __init__(self, config: PreprocessingConfig): assert self.rw_config is not None, ( "Reweighting configuration is not set in the preprocessing config" ) + self.weight_cap = self.rw_config.weight_cap self.organised_components_config = ( Path(config.base_dir) / "split-components/organised-components.yaml" ) @@ -107,6 +112,7 @@ def calculate_weights( print("N per file : ", self.num_jets_estimate) # Get the variables we need to reweight + use_physical_weight = False for rw in reweights: rw_group = rw.group if rw_group not in all_vars: @@ -116,6 +122,9 @@ def calculate_weights( all_vars[rw_group].extend(rw.reweight_vars) if "valid" in existing_vars[rw_group]: all_vars[rw_group] += ["valid"] + if "physicalWeight" in existing_vars.get(rw_group, []): + all_vars[rw_group] += ["physicalWeight"] + use_physical_weight = True if "jets" not in all_vars: all_vars["jets"] = ["pt"] all_vars = {k: list(set(v)) for k, v in all_vars.items()} @@ -171,10 +180,23 @@ def calculate_weights( assert "valid" in data.dtype.names data = data[data["valid"]] classes = np.unique(data[rw.class_var]) if rw.class_var is not None else [None] - for cls in classes: mask = data[rw.class_var] == cls - hist, _outbins = bin_jets(data[mask][rw.reweight_vars], rw.flat_bins) + data_masked = data[mask] + + # Use physicalWeight if present, else uniform weights. + if "physicalWeight" in data_masked.dtype.names: + w = np.asarray(data_masked["physicalWeight"], dtype=np.float64) + w = np.clip(w, 0, self.weight_cap) + else: + w = np.ones(mask.sum(), dtype=float) + + hist, _outbins = bin_jets( + data_masked[rw.reweight_vars], + rw.flat_bins, + weights=w, + ) + if rw.class_var is not None: cls = str(cls) if rw_group not in all_histograms: @@ -274,12 +296,7 @@ def calculate_weights( for cls, hist in all_histograms[rw_group][rw_rep]["histograms"].items(): this_idx_below_min = hist == 0 # | (all_targets[rw_group][rw_rep] == 0) weights = np.zeros_like(hist, dtype=float) - np.divide( - all_targets[rw_group][rw_rep], - hist, - out=weights, - where=hist > 0, - ) + np.divide(all_targets[rw_group][rw_rep], hist, out=weights, where=hist > 0) output_weights[rw_group][rw_rep]["weights"][cls] = weights if idx_below_min is None: idx_below_min = this_idx_below_min @@ -290,6 +307,15 @@ def calculate_weights( if np.any(idx_below_min): for cls in all_histograms[rw_group][rw_rep]["histograms"]: output_weights[rw_group][rw_rep]["weights"][cls][idx_below_min] = 0 + # Cap final factors only on the physicalWeight path; default path unchanged. + if use_physical_weight: + for cls in output_weights[rw_group][rw_rep]["weights"]: + np.clip( + output_weights[rw_group][rw_rep]["weights"][cls], + 0, + self.weight_cap, + out=output_weights[rw_group][rw_rep]["weights"][cls], + ) return output_weights diff --git a/upp/stages/rw_merge.py b/upp/stages/rw_merge.py index 36f9576..9cbd232 100644 --- a/upp/stages/rw_merge.py +++ b/upp/stages/rw_merge.py @@ -5,6 +5,7 @@ from multiprocessing import Pool from pathlib import Path +import h5py import numpy as np import yaml from ftag.hdf5 import H5Reader, H5Writer, join_structured_arrays @@ -76,13 +77,23 @@ def run(self): num_jets_per_file = self.config.num_jets_per_output_file or total_jets batches_per_file = num_jets_per_file // batch_size or 1 - num_batches = ( - total_jets // batch_size + (1 if total_jets % num_jets_per_file != 0 else 0) - ) or 1 - - variables = self.config.variables.combined() if self.config.split != "test" else None - if variables and "flavour_label" not in variables: - variables["jets"] += ["flavour_label"] + num_batches = (total_jets // batch_size + (1 if total_jets % batch_size != 0 else 0)) or 1 + + # stream(None) only loads jets; for full ntuple merge use {dataset: None} per + # top-level HDF5 dataset. + if self.config.variables.keep_all: + p0 = Path(all_files[0]) + with h5py.File(p0, "r") as hf: + variables = {k: None for k in hf if isinstance(hf[k], h5py.Dataset)} + else: + variables = self.config.variables.combined() if self.config.split != "test" else None + if ( + variables is not None + and not self.config.variables.keep_all + and isinstance(variables.get(self.config.jets_name), list) + and "flavour_label" not in variables[self.config.jets_name] + ): + variables[self.config.jets_name] += ["flavour_label"] args_list = [] for i, bi in enumerate(range(0, num_batches, batches_per_file)): args_list.append( diff --git a/upp/stages/split_containers.py b/upp/stages/split_containers.py index efff193..318d925 100644 --- a/upp/stages/split_containers.py +++ b/upp/stages/split_containers.py @@ -86,10 +86,9 @@ def get_all_fp_vars(file: Path | str) -> list[str]: all_vars = get_all_vars(file) # combine the values in this dict into a single list - fp_vars = [ - v for v in all_vars if ("pt" in v.lower() or "energy" in v.lower() or "mass" in v.lower()) - ] + target_keywords = ["pt", "energy", "mass", "weight", "physicalWeight"] + fp_vars = [v for v in all_vars if any(key in v.lower() for key in target_keywords)] return fp_vars @@ -127,18 +126,25 @@ def split_file( output_name=None, variables: dict[str, dict[str, list[str]]] | None = None, flavour_label_list: list[str] | None = None, + keep_all_variables: bool = False, ): if isinstance(input_file, str): input_file = Path(input_file) add_flavour_label = flavour_label_list is not None # All variables for test file all_variables = get_all_datasets(input_file) - print("All variables: ", all_variables, flush=True) # Subset of variables for train/val files parsed_variables: dict[str, list[str]] | dict[str, None] = ( parse_variables(variables) if variables is not None else all_variables ) - print("parsed variables: ", parsed_variables, flush=True) + # Request physicalWeight only if present (added by --metadata), else reader raises. + if isinstance(parsed_variables, dict): + jets_cols = parsed_variables.get("jets") + if isinstance(jets_cols, list) and "physicalWeight" not in jets_cols: + with h5py.File(input_file, "r") as f: + has_pw = "jets" in f and "physicalWeight" in (f["jets"].dtype.names or ()) + if has_pw: + jets_cols.append("physicalWeight") start = time.time() reader = H5Reader(input_file, batch_size=batch_size, shuffle=False) if output_name is None: @@ -163,6 +169,7 @@ def split_file( print(f"At least 1 output file exists for {input_file}. Skipping it", flush=True) return + use_all_cols = keep_all_variables or ("test" in split) writers_by_sample_components[split] = H5Writer.from_file( input_file, num_jets=None, @@ -170,7 +177,7 @@ def split_file( precision="half", full_precision_vars=fp_vars, shuffle=False, - variables=all_variables if "test" in split else parsed_variables, + variables=all_variables if use_all_cols else parsed_variables, compression="gzip", add_flavour_label=add_flavour_label, ) @@ -308,8 +315,10 @@ def run( assert container is not None, "Can only specify files if a container is specified" for container, cuts_by_component in containers_with_split_cuts.items(): + # Sanitize '*' and '/' so the subdir isn't treated as a glob by H5Reader. + container_dir_name = container.replace("*", "all").replace("/", "_") or "default" this_out_dir = ( - Path(self.config.base_dir) / "split-components" / container + Path(self.config.base_dir) / "split-components" / container_dir_name if output_dir is None else Path(".") ) @@ -340,6 +349,7 @@ def run( variables=self.config.config["variables"], flavour_label_list=all_flavours, output_name="output", + keep_all_variables=self.config.keep_all_variables, ) def create_meta_data(self):