Metadata injection, physical-weight based reweighting, and optional full-ntuple output for centrally produced inputs

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)

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 `<tbase_dir>/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` &rarr; `--split-components` &rarr; `--reweight` (`--rw`) &rarr; `--rw-merge` (`--rwm`) &rarr; `--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:

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",
 ]

tests/integration/test_run.py

@@ -9,15 +9,28 @@
 import numpy as np
 import pytest
 from ftag import get_mock_file
+from numpy.lib import recfunctions as rfn
 
 from upp.main import main
 
 this_dir = Path(__file__).parent
 
 
 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):

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):

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

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

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."""