Generalise the framework beyond jets to arbitrary global objects by afroch · Pull Request #156 · umami-hep/umami-preprocessing

afroch · 2026-06-22T14:04:52Z

Summary

This pull request introduces the following changes

Honour the configured global object name in every stage — split, resample, hist, reweight, merge and the prepare/download helpers previously assumed the literal jets group and silently broke with a custom name
Rename all jet-named config options and code identifiers to global_object names (jets_name→global_name, num_jets→num_global_objects, equal_jets→equal_global_objects, add_jet_vars→add_global_vars, add_jet_flavour_label→add_global_object_label, …)
Remap deprecated config keys on load with a single deprecation warning, so existing configs keep working unchanged
Bump to atlas-ftag-tools 0.3.5 (and therefore also to puma 0.5.4), whose H5Reader/H5Writer API is now object-agnostic (global_objects_name, num_global_objects, equal_global_objects)
Derive output h5 metadata attribute names from global_name: unique_<name>, num_<name>, <name>_counts — so global_name: jets still writes unique_jets/num_jets, keeping downstream readers working
Keep physics flavour names (bjets, cjets, …), the default object name (jets) and the default Jet … plot labels unchanged

Relates to the following issues

N/A — addresses review feedback on this PR

Conformity

Several stages assumed the literal "jets" group instead of the configured jets_name, breaking object-agnostic preprocessing for non-"jets" object names. - split_containers: use config.jets_name for batch access; pass jets_name to the H5Reader/H5Writer calls that previously defaulted to "jets". - rw_merge: use config.jets_name for attr_to_write and variables keys; thread jets_name through do_merge_with_weights into its H5Reader/H5Writer. - reweight: pass jets_name to the input H5Reader and the fallback all_vars key. - download_and_prepare: pass jets_name to the metadata H5Reader. Add an end-to-end integration test (and fixture) that renames the mock object group to "objects" and runs split -> weights -> merge with jets_name: objects, asserting the output uses "objects" and no hardcoded "jets" group remains. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

github-actions · 2026-06-22T14:05:22Z

Docs preview: https://umami-hep.github.io/umami-preprocessing/pr-156/

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

codecov · 2026-06-22T14:09:29Z

Codecov Report

❌ Patch coverage is 95.66929% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.34%. Comparing base (446a567) to head (cc31690).

Files with missing lines	Patch %	Lines
upp/classes/components.py	93.54%	2 Missing ⚠️
upp/stages/split_containers.py	86.66%	2 Missing ⚠️
upp/classes/preprocessing_config.py	97.67%	1 Missing ⚠️
upp/classes/reweight_config.py	66.66%	1 Missing ⚠️
upp/stages/merging.py	97.95%	1 Missing ⚠️
upp/stages/resampling.py	97.14%	1 Missing ⚠️
upp/stages/reweight.py	93.33%	1 Missing ⚠️
upp/stages/rw_merge.py	90.90%	1 Missing ⚠️
upp/utils/check_input_samples.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #156      +/-   ##
==========================================
- Coverage   95.38%   95.34%   -0.05%     
==========================================
  Files          24       24              
  Lines        2298     2318      +20     
==========================================
+ Hits         2192     2210      +18     
- Misses        106      108       +2

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

dkobylianskii · 2026-06-22T14:20:47Z

If we assume it work with non-jets object, maybe we could change jets_name to a global_name or so?

afroch · 2026-06-22T14:28:12Z

If we assume it work with non-jets object, maybe we could change jets_name to a global_name or so?

I would love to, but we need to make the changes upstream in atlas-ftag-tools first. I can change it here already, but it would still be jets_name=global_name

Edit: I mis-read what you actually ask for. I implemented it as good as possible, but this will break backward compatibility.

Generalise the framework beyond jets: the config key that names the main per-object dataset is renamed from jets_name to global_name, along with all internal attributes/parameters and the {jets_name} ylabel placeholder. The external ftag H5Reader/H5Writer keyword (jets_name=) and reader.jets_name attribute are kept as-is, since they belong to the ftag library API; call sites now pass jets_name=<config>.global_name. Breaking change: existing configs must rename the global: jets_name key to global_name. All in-repo configs, docs, tests and the {global_name} ylabel placeholder are updated accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

dkobylianskii

If we want to make it general, it's better to clean all jet references in the code now (except atlas-ftag-tools dependence)

dkobylianskii · 2026-06-23T07:58:52Z

Edit: I mis-read what you actually ask for. I implemented it as good as possible, but this will break backward compatibility.

For backward compatibility with old configs we could add a small helper in the config parsing, that will map jets_name to global_name on the fly.

Rename every jet-named config option and code identifier in UPP to object-agnostic global_object names (e.g. num_jets -> num_global_objects, equal_jets -> equal_global_objects, .jets -> .global_objects, add_jet_vars -> add_global_vars, add_jet_flavour_label -> add_global_object_label, unique_jets -> unique_global_objects, jet_counts -> global_object_counts), addressing the PR #156 review. Old configs keep working: PreprocessingConfig.from_file recursively remaps deprecated keys to their new names on load and warns. The atlas-ftag-tools boundary is left untouched: ftag H5Reader/H5Writer keyword args (jets_name=, equal_jets=, num_jets=), reader/writer attributes (reader.num_jets, estimate_available_jets, get_attr("unique_jets")), the h5 data-format attribute strings, and physics flavour names (bjets/cjets/taujets/ghost*). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

afroch · 2026-06-23T09:39:49Z

If we want to make it general, it's better to clean all jet references in the code now (except atlas-ftag-tools dependence)

Done in 98d956a. Every jet-named config option and code identifier is now object-agnostic (num_jets→num_global_objects, equal_jets→equal_global_objects, .jets→.global_objects, unique_jets→unique_global_objects, jet_counts→global_object_counts, etc.), including docstrings, comments and plot/log text. The only jet names left are the atlas-ftag-tools boundary you flagged: the ftag H5Reader/H5Writer kwargs (jets_name=, equal_jets=, num_jets=), reader attributes (reader.num_jets, estimate_available_jets, get_attr("unique_jets")), the h5 data-format attribute strings, and physics flavour names (bjets/cjets/taujets/ghost*). All 112 tests pass.

afroch · 2026-06-23T09:40:08Z

For backward compatibility with old configs we could add a small helper in the config parsing, that will map jets_name to global_name on the fly.

Added. PreprocessingConfig.from_file now runs a recursive _rename_legacy_keys pass that remaps all deprecated keys (jets_name, num_jets*, equal_jets, show_num_jets, …) to their new names on load and emits one deprecation warning, so existing configs keep working unchanged. Covered by a unit test (test_rename_legacy_keys_remaps_nested_and_records).

Align with the object-agnostic ftag H5Reader/H5Writer API and rename the remaining jet-named touchpoints: ftag kwargs/attributes (jets_name, num_jets, equal_jets, estimate_available_jets) and the output h5 metadata attribute strings (num_jets, unique_jets, jet_counts) to their global_object equivalents.

Restore the jet-flavoured default variable labels in _default_variable_labels (pt/eta/mass), and rename the remaining N_Jets sample-check log label to N_Objects.

Derive the output h5 metadata attribute names from global_name instead of hardcoding them, so a config with global_name=jets again writes unique_jets / num_jets / jets_counts (restoring backward compatibility for jet datasets), while a custom name yields unique_<name> / num_<name> / <name>_counts.

afroch · 2026-06-30T21:37:51Z

Hi @dkobylianskii

I updated atlas-ftag-tools and puma in here. Now the removal of jet* should be done

dkobylianskii

Thanks @afroch , it's a nice step towards generalization.

However, we are still depend on flavours in configs, so it won't work with anything except jets

Rename the user-facing settings flavour_config -> class_config, flavour_category -> class_category and the per-component flavours -> classes across the config schema, in-repo configs and tests. Old keys are remapped on load via LEGACY_KEY_MAP so existing configs keep working with a deprecation warning. Internal flavour identifiers and the flavour_label output attribute are unchanged.

class_config is annotated Path | None, so it was skipped by the Path resolution loop in __post_init__ and, unlike other path fields, resolved against the CWD instead of base_dir. Resolve it against base_dir like vds_dir so a relative class_config next to the config is found.

Add a Custom classes section to the configuration docs describing the class definition yaml schema and how class_config generalises the framework beyond jets. Add an end-to-end reweighting test that renames the mock object group and drives the pipeline with a custom class_config (relative path) defining non-standard classes, asserting the output carries the custom class labels.

afroch · 2026-07-02T15:34:04Z

@dkobylianskii you're right that the config previously spoke only in jet-flavour terms. This is now addressed:

Custom classes for any object type. The class definitions are no longer tied to the atlas-ftag-tools bundled jet flavours — point class_config at your own class-definition YAML (a list of name/label/cuts/colour/category entries) and the framework classifies whatever object you like. A relative class_config now resolves against the config's base_dir, and it takes precedence over class_category (standard/extended, the bundled labels).

Renamed the settings to match. To drop the jet-flavour framing, the three user-facing options are now class_config, class_category, and the per-component classes (was flavour_config/flavour_category/flavours) — aligning with the class_var/class_target terms already used in the reweighting config. Old keys are remapped on load via LEGACY_KEY_MAP with a deprecation warning, so existing configs keep working unchanged.

Documented + proven end-to-end. There's a new Custom classes section in the config docs, and a test_rw_custom_flavours integration test that renames the mock object group and drives split → reweight → merge with a custom class_config (non-standard heavy/light classes), asserting the output carries the custom class labels and no hardcoded jets group.

The one thing intentionally left as-is is the flavour_label output attribute (and the reweighting class_var: flavour_label), since it's a data-format contract read by downstream training. I would not change this in this PR, because it is already huge. If you are fine with this, I would open an issue and fix it another, separate PR

nikitapond

This seems like a pretty simple MR - just renaming jets to global_object to make it more easily understandable when used outside of flavour/jet tagging work. I've looked through the code briefly, but as its mainly just a renaming and tests pass so I'm approving.

afroch self-assigned this Jun 22, 2026

afroch added enhancement New feature or request AI-assisted AI coding assistants were used in the development of the changes in this merge request. labels Jun 22, 2026

afroch linked an issue Jun 22, 2026 that may be closed by this pull request

Remove references to "jets" #56

Closed

afroch marked this pull request as ready for review June 22, 2026 14:06

afroch requested review from WeiShengL and dkobylianskii June 22, 2026 14:06

Add changelog entry for jets_name fix

6633708

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

afroch changed the title ~~Honour configured jets_name in remaining stages~~ Generalise framework: honour & rename the global object name (jets_name -> global_name) Jun 22, 2026

Update changelog.md

4040d01

afroch added the No backward compatibility This PR will introduce changes that are not backward compatible label Jun 22, 2026

dkobylianskii reviewed Jun 23, 2026

View reviewed changes

afroch changed the title ~~Generalise framework: honour & rename the global object name (jets_name -> global_name)~~ Generalise framework beyond jets: honour & rename the global object name Jun 24, 2026

Update atlas-ftag-tools and puma

4c2675c

afroch changed the title ~~Generalise framework beyond jets: honour & rename the global object name~~ Generalise the framework beyond jets to arbitrary global objects Jun 30, 2026

afroch added 3 commits June 30, 2026 23:10

Keep default Jet plot labels and rename N_Jets log label

9b50c48

Restore the jet-flavoured default variable labels in _default_variable_labels (pt/eta/mass), and rename the remaining N_Jets sample-check log label to N_Objects.

dkobylianskii reviewed Jul 2, 2026

View reviewed changes

Comment thread tests/unit/stages/test_plotting.py Outdated

Comment thread upp/stages/plot.py Outdated

afroch added 4 commits July 2, 2026 16:43

Label the ATLAS second tag with the configured global object name

4595d52

Merge branch 'main' into remove-hardcoded-jets

cc31690

nikitapond approved these changes Aug 3, 2026

View reviewed changes

nikitapond merged commit 8058141 into main Aug 3, 2026
15 checks passed

nikitapond deleted the remove-hardcoded-jets branch August 3, 2026 10:15

github-actions Bot added a commit that referenced this pull request Aug 3, 2026

Remove docs preview for PR #156

a6f4b61

afroch mentioned this pull request Aug 4, 2026

Rename the flavour_label output field to an object-agnostic name #161

Open

Uh oh!

Conversation

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

Uh oh!

Summary

Conformity

Uh oh!

github-actions Bot commented Jun 22, 2026

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

dkobylianskii commented Jun 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

dkobylianskii left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dkobylianskii commented Jun 23, 2026

Uh oh!

afroch commented Jun 23, 2026

Uh oh!

afroch commented Jun 23, 2026

Uh oh!

afroch commented Jun 30, 2026

Uh oh!

dkobylianskii left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

afroch commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nikitapond left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

afroch commented Jun 22, 2026 •

edited

Loading

codecov Bot commented Jun 22, 2026 •

edited

Loading

afroch commented Jun 22, 2026 •

edited

Loading

afroch commented Jul 2, 2026 •

edited

Loading