Generalise the framework beyond jets to arbitrary global objects

changelog.md

@@ -2,6 +2,8 @@
 
 ### [Latest]
 
+- Generalise the framework beyond jets to arbitrary global objects [#156](https://github.com/umami-hep/umami-preprocessing/pull/156)
+
 ### [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/configuration.md

@@ -10,7 +10,7 @@ Each aspect of the configuration is described in detail below.
 
 Here we define the input h5 samples which are to be preprocessed.
 Each sample is defined using one or more DSIDs, which generally come from the [training-dataset-dumper](https://gitlab.cern.ch/atlas-flavor-tagging-tools/training-dataset-dumper).
-If a list of DSIDs is provided, jets from each DSID will be merged according to the `equal_jets` flag (see below).
+If a list of DSIDs is provided, jets from each DSID will be merged according to the `equal_global_objects` flag (see below).
 The samples are used to define components later on in configs and so one should define them with [anchors](https://support.atlassian.com/bitbucket-cloud/docs/yaml-anchors/).
 
 Below is an example and a table explaining each setting.
@@ -28,7 +28,7 @@ Below is an example and a table explaining each setting.
     ```yaml
     ttbar: &ttbar
       name: ttbar
-      equal_jets: False
+      equal_global_objects: False
       pattern:
         - name1.*.410470.*/*.h5
         - name2.*.410470.*/*.h5
@@ -38,7 +38,7 @@ Below is an example and a table explaining each setting.
 | ------- | ---- | ----------- | ------- |
 |`name`   |`str`| The name of the sample, used in output filenames.| *Required* |
 |`pattern`|`str` or `list[str]`| A single pattern or a list of pattern that match h5 files in a downloaded dataset. H5 files matching each pattern will be transparently merged using virtual datasets. | *Required* |
-|`equal_jets`|`bool`| Only relevant when providing a list of patterns. If `True`, the same number of jets from each DSID are selected. This is required for e.g. in Xbb QCD where each DSID belongs to a different slice, and the resampling would break if you tried to resample with one or more slices missing. If `False` this is not enforced, allowing for larger numbers of available jets. | `True` |
+|`equal_global_objects`|`bool`| Only relevant when providing a list of patterns. If `True`, the same number of jets from each DSID are selected. This is required for e.g. in Xbb QCD where each DSID belongs to a different slice, and the resampling would break if you tried to resample with one or more slices missing. If `False` this is not enforced, allowing for larger numbers of available jets. | `True` |
 
 The virtual dataset files created from wildcard patterns are by default stored alongside the input ntuples.
 If you have no write access to the input ntuples directory and would like to collect all VDS files in an accessible directory instead, set `vds_dir` in the global config (see [Global Config](#global-config)).
@@ -123,15 +123,15 @@ components:
     <<: *lowpt
     sample:
     <<: *ttbar
-    flavours: [bjets, cjets, ujets]
-    num_jets: 10_000_000
+    classes: [bjets, cjets, ujets]
+    num_global_objects: 10_000_000
 
 - region:
     <<: *highpt
     sample:
     <<: *zprime
-    flavours: [bjets, cjets, ujets]
-    num_jets: 5_000_000
+    classes: [bjets, cjets, ujets]
+    num_global_objects: 5_000_000
 ```
 
 Notice that we use `<<*` insertion tool to insert already defined regions and samples.
@@ -140,13 +140,47 @@ Notice that we use `<<*` insertion tool to insert already defined regions and sa
 | ------- | ---- | ----------- |
 | `region`| anchor | The pre-defined kinematic region anchor, e.g. `lowpt` or `highpt`, or `inclusive` if not splitting in $p_T$ |
 | `sample`| anchor | The pre-defined sample anchor, e.g. $t\bar{t}$ or $Z'$ |
-| `flavours` | `list[str]` | One or more jet flavours, e.g. `[bjets]` or `[ujets]`. The list syntax is pure syntactic sugar. If more then one is provided, separate components are created for each flavour.|
-|`num_jets`|`int`| The number of jets to be sampled from this component in the training split. When resampling is skipped, `-1` writes all jets of this component passing the cuts.|
-|`num_jets_val`|`int`| **Optional** (default: `num_jets//10`) number of jets of this component in validation set.|
-|`num_jets_test`|`int`| **Optional** (default: `num_jets//10`) number of jets of this component in a test set.|
+| `classes` | `list[str]` | One or more object classes (flavours), e.g. `[bjets]` or `[ujets]`. Each name must exist in the active class container (the atlas-ftag-tools bundled flavours by default, or your own file via `class_config` — see [Custom classes](#custom-classes)). The list syntax is pure syntactic sugar. If more then one is provided, separate components are created for each class.|
+|`num_global_objects`|`int`| The number of jets to be sampled from this component in the training split. When resampling is skipped, `-1` writes all jets of this component passing the cuts.|
+|`num_global_objects_val`|`int`| **Optional** (default: `num_global_objects//10`) number of jets of this component in validation set.|
+|`num_global_objects_test`|`int`| **Optional** (default: `num_global_objects//10`) number of jets of this component in a test set.|
 
 
 
+### Custom classes
+
+By default the class definitions come from the flavour labels bundled with
+`atlas-ftag-tools`, selected with `class_category` (`standard` or `extended`).
+These are jet flavours, but the framework itself is object-agnostic: to classify
+any other object type, point `class_config` at your own classes yaml. It is a
+list of class definitions, each with a `name`, plotting `label`, selection
+`cuts`, a `colour`, a `category`, and an optional `_px` probability name:
+
+```yaml
+- name: heavy
+  label: Heavy objects
+  cuts: ["HadronConeExclTruthLabelID == 5"]
+  colour: tab:red
+  category: custom
+- name: light
+  label: Light objects
+  cuts: ["HadronConeExclTruthLabelID == 0"]
+  colour: tab:blue
+  category: custom
+```
+
+Reference it from the global config; a relative path is resolved against
+`base_dir`, and `class_config` takes precedence over `class_category`:
+
+```yaml
+global:
+  global_name: objects
+  class_config: custom_flavours.yaml
+```
+
+The `classes` listed for each component then refer to the `name` entries in this
+file (e.g. `classes: [heavy, light]`).
+
 ### Variables
 
 The next thing you need is to provide the variables that are taken from the TDD files and written in the resampled dataset.
@@ -251,7 +285,7 @@ Plot labels and styles can be configured under the optional `plotting:` key. Any
 
 ```yaml
 plotting:
-  num_jets_plotting: 10_000_000
+  num_global_objects_plotting: 10_000_000
   variable_labels:
     pt: "Jet $p_\\mathrm{T}$ [GeV]"
     eta: "Jet $|\\eta|$"
@@ -273,7 +307,7 @@ plotting:
   output_directory: plots
 ```
 
-The `ylabel` setting supports a `{jets_name}` placeholder. Histogram normalisation and overflow handling can be controlled with `norm` and `underoverflow`.
+The `ylabel` setting supports a `{global_name}` placeholder. Histogram normalisation and overflow handling can be controlled with `norm` and `underoverflow`.
 
 ::: upp.classes.plotting_config.PlottingConfig
 

docs/reweighting.md

@@ -69,7 +69,7 @@ Once all the samples are prepared, we can calculate the weights. An example conf
 ```yaml
 
 reweighting:
-  num_jets_estimate: 1_500_000
+  num_global_objects_estimate: 1_500_000
   merge_num_proc: 20
   reweights:
     - group: jets
@@ -87,7 +87,7 @@ reweighting:
 
 ```
 
-`num_jets_estimate` represents the number of each jet flavour used to generate the reweighting histograms. The `merge_num_proc` variable will be relevant in the next section of these docs.
+`num_global_objects_estimate` represents the number of each jet flavour used to generate the reweighting histograms. The `merge_num_proc` variable will be relevant in the next section of these docs.
 Then, you have the `reweights` section, which includes a list of reweight configurations. In this example, we have the first reweight calculated over the jets group. It reweights based on the flavour-label, over the pt and eta distributions. The bins follow the same logic as in resampling.
 The class target can then either be chosen as a single label (e.g, if 0 then the reweighting would target the distribution for `flavour_label==0`), or one of `mean, min, max` which will instead target either the mean distribution, or always take the maximum/minimum bin counts as the target.
 The reweighting can also be performed over track variables, for example
@@ -122,7 +122,7 @@ Finally, we can merge all the relevant jets with their weights. This is done by
 preprocess --config {config} --rwm --split {train/test/val}
 ```
 
-This can either work in series to create 1 single large file, or we can produce multiple files with multi-processing. To do this, ensure the `global` section of the pre-processing config includes `num_jets_per_output_file` and the `reweighting` section has `merge_num_proc>1`.
-This will then launch `merge_num_proc` processes, with approximately `num_jets_per_output_file` per file*.
+This can either work in series to create 1 single large file, or we can produce multiple files with multi-processing. To do this, ensure the `global` section of the pre-processing config includes `num_global_objects_per_output_file` and the `reweighting` section has `merge_num_proc>1`.
+This will then launch `merge_num_proc` processes, with approximately `num_global_objects_per_output_file` per file*.
 
 * Due to the nature of the H5Reader, the actual number of jets per file will be slightly smaller than what is requested, on the order of 0.1%.

docs/run.md

@@ -51,7 +51,7 @@ The stages are described below.
 #### 1. Prepare
 The prepare stage (`--prep`) checks first the number of initial jets that are available per group/sample. For each of the entries in the `pattern` of the group, it checks how many jets are in total available. If this differs too much between the entries in `pattern`, an error is thrown because it indicates that you will might introduce biases in the training. For example, usually entries in `pattern` are different MC campaigns and by using drastically different numbers of initial jets, a campaign dependency can be introduced. If you manually checked it and you expect large differences, you can skip this by adding the command line argument `--skip-sample-check`. If you run the script the first time and you want to run the prepare stage in parallel, please let this script run first! It creates virtual datasets for each entry in `pattern` which could become corrupted if you do run this script in parallel multiple times! Instructions on how to run this check stand-alone can be found in [here](#additional-scripts-initial-sample-check).
 
-Afterwards, the prepare stage reads a specified number of jets (`num_jets_estimate_hist`) for each flavor and constructs histograms of the resampling variables. These histograms are stored in `<base_dir>/hists`.
+Afterwards, the prepare stage reads a specified number of jets (`num_global_objects_estimate_hist`) for each flavor and constructs histograms of the resampling variables. These histograms are stored in `<base_dir>/hists`.
 
 ???info "Paralellisation"
     This step can be parallelized to speed up the histogram creation. To do so, you need to provide the additional `--component` flag. The argument for the flag is the name of the component, which is to be processed. The argument can be constructed when looking closer at the different blocks in the `components` part of the config file. As an example, we take the `ghost-highstat.yaml` config file from the `gn3` folder in `configs/`:
@@ -61,9 +61,9 @@ Afterwards, the prepare stage reads a specified number of jets (`num_jets_estima
         <<: *lowpt
         sample:
         <<: *ttbar
-        flavours: [ghostsplitbjets]
-        num_jets: 22_000_000
-        num_jets_test: 2_000_000
+        classes: [ghostsplitbjets]
+        num_global_objects: 22_000_000
+        num_global_objects_test: 2_000_000
     ```
 
     The argument for the component flag can be constructed by taking the name of the region (this is defined in the definition of `lowpt`)
@@ -81,7 +81,7 @@ Afterwards, the prepare stage reads a specified number of jets (`num_jets_estima
     ```yaml
     ttbar: &ttbar
     name: ttbar
-    equal_jets: False
+    equal_global_objects: False
     pattern:
         - "user.svanstro.601589.e8547_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-09-17_v00_output.h5/*.h5" # mc20d
         - "user.svanstro.601589.e8549_s4159_r14799_p6368.tdd.GN3_dev.25_2_27.24-09-17_v00_output.h5/*.h5" # mc23a
@@ -127,7 +127,7 @@ The merge stage (`--merge`) combines the resampled samples into a single file na
 It also handles shuffling.
 
 #### 4. Normalise 
-The normalize stage (`--norm`) calculates scaling and shifting values for all variables intended for training based on (`num_jets_estimate_norm`). The results are stored in` <tbase_dir>/<out_dir>/norm_dict.yaml`.
+The normalize stage (`--norm`) calculates scaling and shifting values for all variables intended for training based on (`num_global_objects_estimate_norm`). The results are stored in` <tbase_dir>/<out_dir>/norm_dict.yaml`.
 
 #### 5. Plotting 
 

docs/sampling.md

@@ -5,7 +5,7 @@ For resampling, UPP has two different methods implemented. The desired method (`
 
 ### Skipping resampling
 
-Resampling can be disabled entirely by either omitting the `resampling` block from the config or setting `method: none`. In this case no `target`, resampling `variables`, or histogram (`--prep`) step are required. The jets passing the cuts are written directly, capped at each component's `num_jets`. Setting `num_jets: -1` (also valid for `num_jets_val` / `num_jets_test`) writes **all** jets of that component passing the cuts.
+Resampling can be disabled entirely by either omitting the `resampling` block from the config or setting `method: none`. In this case no `target`, resampling `variables`, or histogram (`--prep`) step are required. The jets passing the cuts are written directly, capped at each component's `num_global_objects`. Setting `num_global_objects: -1` (also valid for `num_global_objects_val` / `num_global_objects_test`) writes **all** jets of that component passing the cuts.
 
 Note that the `--no-resample` command line flag is different: it only skips the resampling *stage* (for example to re-run the merge/norm/plot stages on existing component files) and does not disable resampling.
 
@@ -15,7 +15,7 @@ This is an implementation of an [importance sampling](https://en.wikipedia.org/w
 
 The resampling is done using the following steps:
 
-1. A `num_jets_estimate` number of jets are binned for each flavour using the configurations for resampling variable bins. This histogram, `pdf_resampled_flavour`, is the initial estimate of the pdf of jets of each flavour.
+1. A `num_global_objects_estimate` number of jets are binned for each flavour using the configurations for resampling variable bins. This histogram, `pdf_resampled_flavour`, is the initial estimate of the pdf of jets of each flavour.
 2. The importance function is estimated by using the ratio of the histograms for each flavor to that of the target flavour, `pdf_target_flavour/pdf_resampled_flavour`. Safe division is used, which ensures that if for a bin in `pdf_resampled_flavour` is 0, we skip that bin. This ensures that we do not divide by 0. If a bin in `pdf_target_flavour` is 0, we also skip the bin. 
 3. Optionally, the importance function is upscaled. This means that it is interpolated using cubic spline interpolation to a finer grid of bins. The centres of bins are used as nodes for the splines. The new bins are created by splitting the old bins into `upscale_pdf` number of bins of equal width. The function is evaluated in the centers of the new bins. This way, the edge bins of each binning region are actually extrapolated rather than interpolated.
 4. The new batch of jets is being read and after the cuts are applied `n_batch` jets remain. The jets are binned with the the binning from step 1 (if upscaling is not used) or upscaled binning defined by 3 (if upscaling is used) and the reference number of the bin for each jet is saved.
@@ -28,7 +28,7 @@ This algorithm is used for all the flavours except the target flavour for which
 
 Countup resampling tries to select as many unique jets from each bin as possible before selecting the duplicates.
 
-1. `num_jets_estimate` jets are binned for each flavour using the configurations for resampling variable bins. This histogram is the initial estimate of the pdf of jets of each flavour.
+1. `num_global_objects_estimate` jets are binned for each flavour using the configurations for resampling variable bins. This histogram is the initial estimate of the pdf of jets of each flavour.
 2. The new batch of jets is being read and after the cuts are applied `n_batch` jets remain. The jets are binned with the the binning from step 1 (if upscaling is not used) or upscaled binning defined by 3 (if upscaling is used) and the reference number of the bin for each jet is saved.
 3. The number of **requested** jets in each bin are calculated as `floor(n_batch*pdf_target_flavour+uniform([0, 1]))` so that if `n_batch*flavour.sampling_fraction*pdf_target_flavour=1.2` it has a 80% chance to be rounded up to 1 and 20% chance to be rounded up to 2 so that for each bin we get an integer number that on average corresponds to the expected value. 
 4. From each bin we select consecutively (without replacement) the required number of jets. If the bin holds less jets than the **requested** number the rest of jets in this bin is chosen at random from this bin with replacement. This way only few jets in each bin are repeated for `flavour.sampling_fraction=1` and rarely any are repeated for smaller sampling fractions

pyproject.toml

@@ -19,10 +19,10 @@ classifiers = [
 ]
 
 dependencies = [
-  "atlas-ftag-tools==0.3.3",
+  "atlas-ftag-tools==0.3.5",
   "dotmap>=1.3.30",
   "numpy>=2.2.6",
-  "puma-hep==0.5.3",
+  "puma-hep==0.5.4",
   "pyyaml-include==1.3",
   "PyYAML>=6.0.2",
   "rich>=14.1.0",

tests/integration/fixtures/custom_flavours.yaml

@@ -0,0 +1,13 @@
+# Custom, non-jet flavour definitions used by test_rw_custom_flavours.
+# Names deliberately differ from the atlas-ftag-tools bundled flavours to prove
+# the container is loaded from this file. Cuts select on the mock truth variable.
+- name: heavy
+  label: Heavy objects
+  cuts: ["HadronConeExclTruthLabelID == 5"]
+  colour: tab:red
+  category: custom
+- name: light
+  label: Light objects
+  cuts: ["HadronConeExclTruthLabelID == 0"]
+  colour: tab:blue
+  category: custom

tests/integration/fixtures/test_config_countup.yaml

@@ -32,43 +32,43 @@ components:
       <<: *lowpt
     sample:
       <<: *ttbar
-    flavours: [bjets]
-    num_jets: 7_000
+    classes: [bjets]
+    num_global_objects: 7_000
 
   - region:
       <<: *lowpt
     sample:
       <<: *ttbar
-    flavours: [cjets]
-    num_jets: 7_000
+    classes: [cjets]
+    num_global_objects: 7_000
 
   - region:
       <<: *lowpt
     sample:
       <<: *ttbar
-    flavours: [ujets]
-    num_jets: 7_000
+    classes: [ujets]
+    num_global_objects: 7_000
 
   - region:
       <<: *highpt
     sample:
       <<: *zprime
-    flavours: [bjets]
-    num_jets: 3_000
+    classes: [bjets]
+    num_global_objects: 3_000
 
   - region:
       <<: *highpt
     sample:
       <<: *zprime
-    flavours: [cjets]
-    num_jets: 3_000
+    classes: [cjets]
+    num_global_objects: 3_000
 
   - region:
       <<: *highpt
     sample:
       <<: *zprime
-    flavours: [ujets]
-    num_jets: 3_000
+    classes: [ujets]
+    num_global_objects: 3_000
 
 resampling:
   target: bjets
@@ -82,10 +82,10 @@ resampling:
 
 global:
   batch_size: 10_000
-  num_jets_estimate: 5_000
-  num_jets_estimate_norm: 100
-  num_jets_estimate_available: -1
-  num_jets_per_output_file: 15_000
+  num_global_objects_estimate: 5_000
+  num_global_objects_estimate_norm: 100
+  num_global_objects_estimate_available: -1
+  num_global_objects_per_output_file: 15_000
   base_dir: tmp/upp-tests/integration/temp_workspace/
   out_dir: test_out
   ntuple_dir: ntuples