Metadata injection, physical-weight based reweighting, and optional full-ntuple output for centrally produced inputs by Zhuyuanda · Pull Request #142 · umami-hep/umami-preprocessing

Zhuyuanda · 2026-05-27T14:16:32Z

Summary

This PR extends UPP to support a metadata → split → reweight → rw-merge workflow for centrally produced ntuples from TDD that lack pre-computed cross-section metadata and physicalWeight. It also adds an opt-in keep_all_variables flag so train/val/test VDS outputs preserve all top-level HDF5 datasets (tracks, towers, flow, …) while still appending reweight columns on jets.

Motivation

Centrally produced .h5 ntuples by TDD often need:

Metadata (cross section, k-factor, gen filter efficiency) to calculate physics weights.
physicalWeight on jets for histogram-based reweighting that respects MC event weights and sample normalization.
Full ntuple content in final VDS (not just the variables.yaml subset), e.g. for downstream training or analysis that needs tracks/towers/flow.

Changes

1) New `--metadata` stage

Adds a metadata injection stage via MetadataInjector.
Introduces a CLI flag --metadata in upp/main.py.
For each input file:
- creates a .bak backup
- injects metadata using ftag.find_metadata.MetadataFinder
- computes and appends physicalWeight to the jets structured array using:
  $$(\frac{\text{XS} \times \text{eff} \times \text{kfactor}}{\text{SOW}}) \times \text{mcEventWeight}$$
- restores original jets dataset attributes after dataset recreation
In-place modification: designed to run on workspace copies, not central originals.

Files:

upp/stages/metadata_injector.py
upp/main.py

2) Physical-weight based reweighting

Reweighting uses physicalWeight if present; otherwise falls back to uniform weights.
Caps reweight factor via RW_CAP = 1e4 and uses HIST_FLOOR = 1e-6.

Files:

upp/stages/reweight.py

3) Weighted histogram support

bin_jets(..., weights=...) supports weighted histograms (count when weights is None, sum otherwise).

Files:

upp/stages/hist.py

4) `global.keep_all_variables` (full top-level dataset passthrough)

New config flag: global.keep_all_variables: true (default false)
Wired through:
- PreprocessingConfig → VariableConfig(keep_all=...)
Effects:
- Split outputs (train/val/test): write full top-level dataset fields for all splits (not just a subset).
- RW columns: reweight columns still appended only under jets.

Files:

upp/classes/preprocessing_config.py
upp/stages/split_containers.py
upp/stages/rw_merge.py
upp/stages/normalisation.py (keeps norm computation consistent with declared inputs/labels)

Config example:

upp/configs/metadata.yaml
upp/configs/flavours-r10.yaml

Conformity

Resolve hist.py and reweight.py conflicts: keep PW-weighted histograms, HIST_FLOOR/RW_CAP workflow, and upstream float64/np.divide warning fixes. Co-authored-by: Cursor <cursoragent@cursor.com>

codecov · 2026-06-03T09:45:11Z

Codecov Report

❌ Patch coverage is 98.73418% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.35%. Comparing base (239ff9e) to head (20548f2).

Files with missing lines	Patch %	Lines
upp/stages/metadata_injector.py	98.52%	1 Missing ⚠️
upp/stages/normalisation.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #142      +/-   ##
==========================================
+ Coverage   95.38%   96.35%   +0.97%     
==========================================
  Files          24       25       +1     
  Lines        2298     2444     +146     
==========================================
+ Hits         2192     2355     +163     
+ Misses        106       89      -17

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

afroch

A changelog entry still seems to be missing — could you add one? It would also be good to document the usage of the new --metadata stage and the keep_all_variables flag in the docs.

afroch · 2026-06-18T07:40:52Z

+                safe_hist = np.maximum(hist, HIST_FLOOR)
                weights = np.zeros_like(hist, dtype=float)
-                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
-                else:
-                    idx_below_min |= this_idx_below_min
-            # If we have any bins where we have 0 of a given flavour, we set all the
-            # weights to 0
-            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
+                np.divide(target, safe_hist, out=weights, where=hist > HIST_FLOOR)
+                output_weights[rw_group][rw_rep]["weights"][cls] = np.clip(weights, 0, RW_CAP)


This changes reweighting for all configs, not just the physicalWeight path: the old "zero a bin across all classes when any class is empty" logic is gone and weights are now capped at RW_CAP=1e4. Could you confirm this is intended and validate against current outputs? Also safe_hist is effectively dead — where=hist > HIST_FLOOR means it's only used where safe_hist == hist, so the floor never applies. And PW_CAP (L181) / RW_CAP are the same value defined twice; consider a single named constant.

afroch · 2026-06-18T07:41:00Z

        self.variables = VariableConfig(
-            self.config["variables"], self.jets_name, self.is_test, selectors
+            self.config["variables"],
+            self.jets_name,


This swaps the 3rd VariableConfig arg from is_test to keep_all_variables, so test splits no longer get "keep all vars" for free via combined(). It's patched back in rw_merge/split_containers, but could you double-check every other combined()/.keep_all consumer (normalisation, resampling, merging) still behaves for the test split?

- Relocate metadata/flavours/variables-r10 configs under upp/configs/MetadataRW/; revert default variables/xbb/single-b configs to main - split_containers: gate physicalWeight append on column presence / keep_all - reweight: restore idx_below_min zeroing, unify WEIGHT_CAP applied only on the physicalWeight path, drop dead safe_hist and debug prints - metadata_injector: guard sow==0 / non-finite metadata, collect per-file failures and re-raise a summary instead of log-and-continue - preprocessing_config: keep_all = keep_all_variables or is_test (restore test split) - Condense inline block comments; add changelog entry and document --metadata stage and keep_all_variables flag Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # upp/classes/preprocessing_config.py

Weight --plot histograms by physicalWeight (capped) and reweight columns when present, and label plots "Pre/Post Reweighting" instead of "Resampling" on the reweight route. Default resampling route stays unweighted. Adds unit tests and updates changelog/docs. Co-authored-by: Cursor <cursoragent@cursor.com>

afroch

A few non-blocking follow-ups on the new metadata/reweight/plot route (inline). The earlier review threads (cap gating, restored empty-bin zeroing, single WEIGHT_CAP, keep_all_variables or is_test, sow/failure guards) all look addressed — thanks

afroch · 2026-06-24T10:20:25Z

+                output_weights[rw_group][rw_rep]["weights"][cls] = np.where(
+                    hist > 0, all_targets[rw_group][rw_rep] / hist, 0
                )


all_targets[...] / hist is evaluated over the whole array before np.where selects, so the hist == 0 bins divide by zero and re-emit the RuntimeWarning: divide by zero that #140 specifically removed. The masked np.divide keeps it warning-free:

weights = np.zeros_like(hist, dtype=float) np.divide(all_targets[rw_group][rw_rep], hist, out=weights, where=hist > 0) output_weights[rw_group][rw_rep]["weights"][cls] = weights

Non-blocking, but it reintroduces a warning we'd already silenced.

Switched back to the masked np.divide (out=weights, where=hist > 0), so empty bins no longer trigger the divide-by-zero warning

afroch · 2026-06-24T10:20:39Z

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


Since the plot route keys off the presence of physicalWeight, injecting it into every mock means the default resampling integration tests now traverse the weighted / "Pre Reweighting"-labelled plotting branch too. The unit test covers _reweight_weight_fields directly, so this isn't blocking — but could you either gate this on the reweight test(s) only, or add a brief comment on why all mocks need the column? Otherwise the genuinely-unweighted default route is no longer exercised end-to-end.

generate_mock now only injects physicalWeight when with_physical_weight=True (default off), so the default resampling route stays exercised unweighted end-to-end; the weighted route is covered by test_run_rw.py.

afroch · 2026-06-24T10:20:52Z

+
+                    # 6. Delete and recreate the dataset while restoring attributes
+                    del f["jets"]
+                    new_ds = f.create_dataset("jets", data=updated_jets, compression="gzip")


The recreated jets dataset is hard-coded to compression="gzip" even if the original was uncompressed (or used different chunking/shuffle), so the metadata stage silently changes the storage layout. Consider preserving the source dataset's compression/compression_opts/chunks/shuffle from old_jets_ds alongside the attrs you already restore.

Now preserving the source dataset's compression/compression_opts/chunks/shuffle on recreation instead of hard-coding gzip.

- reweight: use masked np.divide so empty bins no longer re-emit the divide-by-zero RuntimeWarning that umami-hep#140 removed - configs: restore mcCampaignYear resampling bins in xbb / xbb-gn3x - test_run: only inject physicalWeight on demand so the default (unweighted) resampling route stays exercised end-to-end - metadata_injector: preserve source compression/chunks/shuffle instead of hard-coding gzip on the recreated jets dataset Co-authored-by: Cursor <cursoragent@cursor.com>

yuanda and others added 6 commits May 27, 2026 15:46

metadata injector reweighting workflow

bdbddab

Merge remote-tracking branch 'upstream/main' into physics

2d09a5c

Resolve hist.py and reweight.py conflicts: keep PW-weighted histograms, HIST_FLOOR/RW_CAP workflow, and upstream float64/np.divide warning fixes. Co-authored-by: Cursor <cursoragent@cursor.com>

precommit fix

94b4b94

test logic added

76bb44c

precommit and pytest fix

29ef221

Merge branch 'main' into physics

88a84ed

yuanda added 2 commits June 8, 2026 13:47

precommit fix and improve coverage

3cd4a57

3.13 and 3.14 fix

917cf3c

afroch assigned Zhuyuanda Jun 17, 2026

afroch added the enhancement New feature or request label Jun 17, 2026

afroch requested review from WeiShengL and dkobylianskii June 17, 2026 09:36

afroch requested changes Jun 18, 2026

View reviewed changes

Comment thread upp/configs/MetadataRW/flavours-r10.yaml

Comment thread upp/configs/MetadataRW/metadata.yaml

Comment thread upp/configs/variables.yaml Outdated

Comment thread upp/configs/variables.yaml Outdated

afroch requested changes Jun 18, 2026

View reviewed changes

Comment thread upp/stages/split_containers.py Outdated

afroch reviewed Jun 18, 2026

View reviewed changes

yuanda and others added 3 commits June 22, 2026 12:10

Merge remote-tracking branch 'upstream/main' into physics

a1d9d1b

Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # upp/classes/preprocessing_config.py

afroch reviewed Jun 24, 2026

View reviewed changes

Comment thread upp/configs/xbb-gn3x.yaml

Comment thread upp/configs/xbb.yaml

afroch reviewed Jun 24, 2026

View reviewed changes

Uh oh!

Conversation

Zhuyuanda commented May 27, 2026

Summary

Motivation

Changes

1) New --metadata stage

Files:

2) Physical-weight based reweighting

Files:

3) Weighted histogram support

Files:

4) global.keep_all_variables (full top-level dataset passthrough)

Files:

Config example:

Conformity

Uh oh!

codecov Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

afroch left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

afroch left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

1) New `--metadata` stage

4) `global.keep_all_variables` (full top-level dataset passthrough)

codecov Bot commented Jun 3, 2026 •

edited

Loading

afroch left a comment •

edited

Loading