diff --git a/changelog.md b/changelog.md index f084858a..0eceae3f 100644 --- a/changelog.md +++ b/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) diff --git a/docs/configuration.md b/docs/configuration.md index 2be9d80a..26694e2e 100644 --- a/docs/configuration.md +++ b/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 diff --git a/docs/reweighting.md b/docs/reweighting.md index aea08c59..013dc7e5 100644 --- a/docs/reweighting.md +++ b/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%. diff --git a/docs/run.md b/docs/run.md index b78d9877..bfb10376 100644 --- a/docs/run.md +++ b/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 `/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 `/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` //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` //norm_dict.yaml`. #### 5. Plotting diff --git a/docs/sampling.md b/docs/sampling.md index 031e2494..24eea175 100644 --- a/docs/sampling.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 7c65c075..93d975e6 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/integration/fixtures/custom_flavours.yaml b/tests/integration/fixtures/custom_flavours.yaml new file mode 100644 index 00000000..27e8f01f --- /dev/null +++ b/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 diff --git a/tests/integration/fixtures/test_config_countup.yaml b/tests/integration/fixtures/test_config_countup.yaml index 9fb46ef5..78055a10 100644 --- a/tests/integration/fixtures/test_config_countup.yaml +++ b/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 diff --git a/tests/integration/fixtures/test_config_countup_upscaled.yaml b/tests/integration/fixtures/test_config_countup_upscaled.yaml index 1fa14f9d..dc25c840 100644 --- a/tests/integration/fixtures/test_config_countup_upscaled.yaml +++ b/tests/integration/fixtures/test_config_countup_upscaled.yaml @@ -32,43 +32,43 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets] - num_jets: 70_000 + classes: [bjets] + num_global_objects: 70_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [cjets] - num_jets: 70_000 + classes: [cjets] + num_global_objects: 70_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 70_000 + classes: [ujets] + num_global_objects: 70_000 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets] - num_jets: 30_000 + classes: [bjets] + num_global_objects: 30_000 - region: <<: *highpt sample: <<: *zprime - flavours: [cjets] - num_jets: 30_000 + classes: [cjets] + num_global_objects: 30_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ujets] - num_jets: 30_000 + classes: [ujets] + num_global_objects: 30_000 resampling: target: cjets @@ -82,9 +82,9 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5000 - num_jets_estimate_norm: 100 - num_jets_estimate_available: -1 + num_global_objects_estimate: 5000 + num_global_objects_estimate_norm: 100 + num_global_objects_estimate_available: -1 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_method_none.yaml b/tests/integration/fixtures/test_config_method_none.yaml index f94f71de..ccc407e3 100644 --- a/tests/integration/fixtures/test_config_method_none.yaml +++ b/tests/integration/fixtures/test_config_method_none.yaml @@ -27,22 +27,22 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets] - num_jets: 1_000 + classes: [bjets] + num_global_objects: 1_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [cjets] - num_jets: 2_000 + classes: [cjets] + num_global_objects: 2_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 3_000 + classes: [ujets] + num_global_objects: 3_000 # Resampling explicitly disabled - no target or variables required. resampling: @@ -50,9 +50,9 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5_000 - num_jets_estimate_norm: 100 - num_jets_estimate_available: -1 + num_global_objects_estimate: 5_000 + num_global_objects_estimate_norm: 100 + num_global_objects_estimate_available: -1 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_no_resample.yaml b/tests/integration/fixtures/test_config_no_resample.yaml index 331dd2b0..3d7186a6 100644 --- a/tests/integration/fixtures/test_config_no_resample.yaml +++ b/tests/integration/fixtures/test_config_no_resample.yaml @@ -29,28 +29,28 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets] - num_jets: 1_000 + classes: [bjets] + num_global_objects: 1_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [cjets] - num_jets: 1_000 + classes: [cjets] + num_global_objects: 1_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: -1 + classes: [ujets] + num_global_objects: -1 global: batch_size: 10_000 - num_jets_estimate: 5_000 - num_jets_estimate_norm: 100 - num_jets_estimate_available: -1 + num_global_objects_estimate: 5_000 + num_global_objects_estimate_norm: 100 + num_global_objects_estimate_available: -1 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_pdf_auto.yaml b/tests/integration/fixtures/test_config_pdf_auto.yaml index 234590c6..eaa2520a 100644 --- a/tests/integration/fixtures/test_config_pdf_auto.yaml +++ b/tests/integration/fixtures/test_config_pdf_auto.yaml @@ -31,43 +31,43 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets] - num_jets: 12_000 + classes: [bjets] + num_global_objects: 12_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [cjets] - num_jets: 12_000 + classes: [cjets] + num_global_objects: 12_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 12_000 + classes: [ujets] + num_global_objects: 12_000 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets] - num_jets: 6_000 + classes: [bjets] + num_global_objects: 6_000 - region: <<: *highpt sample: <<: *zprime - flavours: [cjets] - num_jets: 6_000 + classes: [cjets] + num_global_objects: 6_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ujets] - num_jets: 6_000 + classes: [ujets] + num_global_objects: 6_000 resampling: target: cjets @@ -81,14 +81,14 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5000 - num_jets_estimate_available: -1 + num_global_objects_estimate: 5000 + num_global_objects_estimate_available: -1 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples plotting: - num_jets_plotting: 100 + num_global_objects_plotting: 100 variable_labels: pt: "$p_\\mathrm{T}$ [GeV]" sample_labels: diff --git a/tests/integration/fixtures/test_config_pdf_upscaled.yaml b/tests/integration/fixtures/test_config_pdf_upscaled.yaml index d90f2960..3bc4428a 100644 --- a/tests/integration/fixtures/test_config_pdf_upscaled.yaml +++ b/tests/integration/fixtures/test_config_pdf_upscaled.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 @@ -83,7 +83,7 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5000 + num_global_objects_estimate: 5000 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_rw.yaml b/tests/integration/fixtures/test_config_rw.yaml index 6b3a334a..ca7dd71e 100644 --- a/tests/integration/fixtures/test_config_rw.yaml +++ b/tests/integration/fixtures/test_config_rw.yaml @@ -19,14 +19,14 @@ global_cuts: !include GN3V01/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "data1.h5" - "data2.h5" zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "data3.h5" @@ -51,18 +51,18 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets, ujets, taujets] - num_jets: -1 + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets, ujets, taujets] - num_jets: -1 + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 reweighting: - num_jets_estimate: 200 + num_global_objects_estimate: 200 merge_num_proc: 1 reweights: - group: jets @@ -146,9 +146,9 @@ reweighting: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_rw_custom_flavours.yaml b/tests/integration/fixtures/test_config_rw_custom_flavours.yaml new file mode 100644 index 00000000..a39c9a95 --- /dev/null +++ b/tests/integration/fixtures/test_config_rw_custom_flavours.yaml @@ -0,0 +1,96 @@ +variables: + objects: + inputs: + - pt + - eta + labels: + - mass + - eventNumber + + tracks: + inputs: + - dphi + - deta + - qOverP + labels: + - qOverP + - leptonID +global_cuts: !include GN3V01/simple-split.yaml + +ttbar: &ttbar + name: ttbar + equal_global_objects: False + pattern: + - "data1.h5" + - "data2.h5" + +zprime: &zprime + name: zprime + equal_global_objects: False + pattern: + - "data3.h5" + +lowpt: &lowpt + name: lowpt + cuts: + - [pt, ">", 20_000] + - [pt, "<", 250_000] + - [eta, "<", 2.5] + - [eta, ">", -2.5] + +highpt: &highpt + name: highpt + cuts: + - [pt, ">", 250_000] + - [pt, "<", 6_000_000] + - [eta, "<", 2.5] + - [eta, ">", -2.5] + +components: + - region: + <<: *lowpt + sample: + <<: *ttbar + classes: [heavy, light] + num_global_objects: -1 + + - region: + <<: *highpt + sample: + <<: *zprime + classes: [heavy, light] + num_global_objects: -1 + +reweighting: + num_global_objects_estimate: 200 + merge_num_proc: 1 + reweights: + - group: objects + reweight_vars: [pt, eta] + bins: + pt: + [ + [20_000, 250_000, 50], + [250_000, 1_000_000, 50], + [1_000_000, 6_000_000, 50], + ] + eta: [[-2.5, 2.5, 40]] + class_var: flavour_label + class_target: mean + - group: tracks + reweight_vars: [deta] + bins: + deta: [[-1.0, 1.0, 50]] + class_var: leptonID + class_target: mean + +# note: sensible defaults are defined in the PreprocessingConfig constructor +# class_config is a relative path, resolved against base_dir +global: + global_name: objects + class_config: custom_flavours.yaml + batch_size: 1_000_000 + num_global_objects_estimate: 25_000_000 + base_dir: tmp/upp-tests/integration/temp_workspace/ + out_dir: test_out + ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_rw_custom_name.yaml b/tests/integration/fixtures/test_config_rw_custom_name.yaml new file mode 100644 index 00000000..1c222888 --- /dev/null +++ b/tests/integration/fixtures/test_config_rw_custom_name.yaml @@ -0,0 +1,94 @@ +variables: + objects: + inputs: + - pt + - eta + labels: + - mass + - eventNumber + + tracks: + inputs: + - dphi + - deta + - qOverP + labels: + - qOverP + - leptonID +global_cuts: !include GN3V01/simple-split.yaml + +ttbar: &ttbar + name: ttbar + equal_global_objects: False + pattern: + - "data1.h5" + - "data2.h5" + +zprime: &zprime + name: zprime + equal_global_objects: False + pattern: + - "data3.h5" + +lowpt: &lowpt + name: lowpt + cuts: + - [pt, ">", 20_000] + - [pt, "<", 250_000] + - [eta, "<", 2.5] + - [eta, ">", -2.5] + +highpt: &highpt + name: highpt + cuts: + - [pt, ">", 250_000] + - [pt, "<", 6_000_000] + - [eta, "<", 2.5] + - [eta, ">", -2.5] + +components: + - region: + <<: *lowpt + sample: + <<: *ttbar + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 + + - region: + <<: *highpt + sample: + <<: *zprime + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 + +reweighting: + num_global_objects_estimate: 200 + merge_num_proc: 1 + reweights: + - group: objects + reweight_vars: [pt, eta] + bins: + pt: + [ + [20_000, 250_000, 50], + [250_000, 1_000_000, 50], + [1_000_000, 6_000_000, 50], + ] + eta: [[-2.5, 2.5, 40]] + class_var: flavour_label + class_target: mean + - group: tracks + reweight_vars: [deta] + bins: + deta: [[-1.0, 1.0, 50]] + class_var: leptonID + class_target: mean + +# note: sensible defaults are defined in the PreprocessingConfig constructor +global: + global_name: objects + batch_size: 1_000_000 + num_global_objects_estimate: 25_000_000 + base_dir: tmp/upp-tests/integration/temp_workspace/ + out_dir: test_out + ntuple_dir: ntuples diff --git a/tests/integration/fixtures/test_config_track_selection.yaml b/tests/integration/fixtures/test_config_track_selection.yaml index 3aca725c..727a14c1 100644 --- a/tests/integration/fixtures/test_config_track_selection.yaml +++ b/tests/integration/fixtures/test_config_track_selection.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,9 +82,9 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5000 - num_jets_estimate_norm: 100 - num_jets_estimate_available: -1 + num_global_objects_estimate: 5000 + num_global_objects_estimate_norm: 100 + num_global_objects_estimate_available: -1 base_dir: tmp/upp-tests/integration/temp_workspace/ out_dir: test_out ntuple_dir: ntuples diff --git a/tests/integration/test_run.py b/tests/integration/test_run.py index d86a1d6c..64e4a167 100644 --- a/tests/integration/test_run.py +++ b/tests/integration/test_run.py @@ -119,14 +119,14 @@ def test_run_no_resample(self): assert os.path.exists(fname) with h5py.File(fname, "r") as f: jets = f["jets"][:] - jet_counts = json.loads(f.attrs["jet_counts"]) + jets_counts = json.loads(f.attrs["jets_counts"]) assert f.attrs["resampling_method"] == "none" - # capped components write exactly num_jets; the -1 component writes all its jets - assert jet_counts["lowpt_ttbar_bjets"]["num_jets"] == 1_000 - assert jet_counts["lowpt_ttbar_cjets"]["num_jets"] == 1_000 - assert jet_counts["lowpt_ttbar_ujets"]["num_jets"] > 1_000 - assert jet_counts["total"]["num_jets"] == len(jets) + # capped components write exactly num_jets; -1 writes all its objects + assert jets_counts["lowpt_ttbar_bjets"]["num_jets"] == 1_000 + assert jets_counts["lowpt_ttbar_cjets"]["num_jets"] == 1_000 + assert jets_counts["lowpt_ttbar_ujets"]["num_jets"] > 1_000 + assert jets_counts["total"]["num_jets"] == len(jets) def test_run_method_none(self): args = [ diff --git a/tests/integration/test_run_rw.py b/tests/integration/test_run_rw.py index 1c07c8c7..b2f33328 100644 --- a/tests/integration/test_run_rw.py +++ b/tests/integration/test_run_rw.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import shutil import subprocess from pathlib import Path @@ -127,15 +128,94 @@ def test_rw(self): self._calculate_weights() self._rw_merge() - def test_rw_unequal_jets(self): - """Test reweighting when a file has fewer jets than num_jets_estimate. + def _rename_mock_group(self, fname, old="jets", new="objects"): + """Rename the main object group in a mock file (jets -> objects).""" + with h5py.File(fname, "a") as f: + f.move(old, new) + + def test_rw_custom_object_name(self): + """End-to-end reweighting with a non-"jets" object group name. + + Guards against hardcoded "jets" references: the input group is renamed + to "objects" and the config sets global_name: objects. Every stage must + honour the configured name rather than assuming "jets". + """ + for container in ["data1.h5", "data2.h5", "data3.h5"]: + self._rename_mock_group(f"tmp/upp-tests/integration/temp_workspace/ntuples/{container}") + + config = str(Path(this_dir / "fixtures/test_config_rw_custom_name.yaml")) + + # Split + main(["--config", config, "--split", "train", "--split-components", *self.no]) + + # Calculate weights + main(["--config", config, "--rw", *self.no]) + hist_file = Path("tmp/upp-tests/integration/temp_workspace/test_out/histograms.h5") + with h5py.File(hist_file, "r") as f: + assert f.keys() == {"objects", "tracks"}, ( + f"Expected 'objects'/'tracks' groups, found {f.keys()}" + ) + + # Merge with weights + for split in ["train", "val", "test"]: + main(["--config", config, "--rwm", "--split", split, *self.no]) + outfile = Path( + f"tmp/upp-tests/integration/temp_workspace/test_out/pp_output_{split}_vds.h5" + ) + assert outfile.exists() + with h5py.File(outfile, "r") as f: + assert "objects" in f, "Expected 'objects' group in output file" + assert "jets" not in f, "Output must not contain a hardcoded 'jets' group" + assert "flavour_label" in f["objects"].attrs + assert "flavour_label" in f["objects"].dtype.names + + def test_rw_custom_flavours(self): + """End-to-end reweighting with a user-supplied custom class_config. + + The mock "jets" group is renamed to "objects" and the config points + class_config at a custom classes yaml (given as a path relative to + base_dir) defining non-standard classes (heavy/light). Proves custom + classes work end-to-end for a non-jet object and that a relative + class_config path resolves against base_dir. + """ + for container in ["data1.h5", "data2.h5", "data3.h5"]: + self._rename_mock_group(f"tmp/upp-tests/integration/temp_workspace/ntuples/{container}") + + # Copy the custom classes file into base_dir so the relative + # class_config path in the config resolves against base_dir. + base_dir = Path("tmp/upp-tests/integration/temp_workspace") + shutil.copy(this_dir / "fixtures/custom_flavours.yaml", base_dir / "custom_flavours.yaml") + + config = str(Path(this_dir / "fixtures/test_config_rw_custom_flavours.yaml")) + + main(["--config", config, "--split", "train", "--split-components", *self.no]) + main(["--config", config, "--rw", *self.no]) + + for split in ["train", "val", "test"]: + main(["--config", config, "--rwm", "--split", split, *self.no]) + outfile = Path( + f"tmp/upp-tests/integration/temp_workspace/test_out/pp_output_{split}_vds.h5" + ) + assert outfile.exists() + with h5py.File(outfile, "r") as f: + assert "objects" in f, "Expected 'objects' group in output file" + assert "jets" not in f, "Output must not contain a hardcoded 'jets' group" + assert "flavour_label" in f["objects"].dtype.names + labels = [ + x.decode() if isinstance(x, bytes) else str(x) + for x in f["objects"].attrs["flavour_label"] + ] + assert labels == ["heavy", "light"], f"Expected custom flavours, found {labels}" + + def test_rw_unequal_objects(self): + """Test reweighting when a file has fewer objects than num_global_objects_estimate. Previously this would crash with an assertion error. After the fix, - get_input_readers() caps per-reader jet counts at available jets and + get_input_readers() caps per-reader object counts at available objects and the batch loop handles StopIteration from shorter readers. """ # Overwrite data1.h5 with a smaller mock — after per-flavour splitting - # some flavour files will have fewer jets than num_jets_estimate + # some flavour files will have fewer objects than num_global_objects_estimate self.generate_mock("tmp/upp-tests/integration/temp_workspace/ntuples/data1.h5", N=500) self._run_split() self._calculate_weights() diff --git a/tests/unit/classes/test_components.py b/tests/unit/classes/test_components.py index cd885f2f..693e276d 100644 --- a/tests/unit/classes/test_components.py +++ b/tests/unit/classes/test_components.py @@ -20,9 +20,9 @@ def make_component(tmp_path: Path, vds_dir: Path | None) -> Component: flavour=Flavours.bjets, global_cuts=Cuts.empty(), dirname=tmp_path / "components" / "sub", - num_jets=100, - num_jets_estimate_available=0, - equal_jets=True, + num_global_objects=100, + num_global_objects_estimate_available=0, + equal_global_objects=True, ) diff --git a/tests/unit/classes/test_plotting_config.py b/tests/unit/classes/test_plotting_config.py index 6c263e4b..1e642a09 100644 --- a/tests/unit/classes/test_plotting_config.py +++ b/tests/unit/classes/test_plotting_config.py @@ -28,7 +28,7 @@ def test_plotting_config_default_mass_label(): @pytest.mark.parametrize( ("kwargs", "message"), [ - ({"num_jets_plotting": 0}, "plotting.num_jets_plotting"), + ({"num_global_objects_plotting": 0}, "plotting.num_global_objects_plotting"), ({"output_formats": []}, "plotting.output_formats"), ({"linestyles": []}, "plotting.linestyles"), ], diff --git a/tests/unit/classes/test_preprocessing_config.py b/tests/unit/classes/test_preprocessing_config.py index 0ab771fc..27b6217f 100644 --- a/tests/unit/classes/test_preprocessing_config.py +++ b/tests/unit/classes/test_preprocessing_config.py @@ -11,10 +11,50 @@ from ftag import Extended_Flavours, Flavours, LabelContainer, get_mock_file from upp import __version__ -from upp.classes.preprocessing_config import PreprocessingConfig +from upp.classes.preprocessing_config import ( + PreprocessingConfig, + _rename_legacy_keys, +) from upp.classes.resampling_config import ResamplingConfig +def test_rename_legacy_keys_remaps_nested_and_records(): + """Deprecated jet-named keys are remapped everywhere, others untouched.""" + raw = { + "global": { + "jets_name": "muons", + "num_jets_estimate": 5, + "flavour_config": "my_flavours.yaml", + "flavour_category": "extended", + }, + "components": [{"num_jets": 10, "sample": {"equal_jets": True}, "flavours": ["bjets"]}], + "plotting": {"show_num_jets": False, "kept": 1}, + } + found: set[str] = set() + out = _rename_legacy_keys(raw, found) + + assert out["global"] == { + "global_name": "muons", + "num_global_objects_estimate": 5, + "class_config": "my_flavours.yaml", + "class_category": "extended", + } + assert out["components"][0]["num_global_objects"] == 10 + assert out["components"][0]["sample"]["equal_global_objects"] is True + assert out["components"][0]["classes"] == ["bjets"] # class names (values) untouched + assert out["plotting"] == {"show_num_global_objects": False, "kept": 1} + assert found == { + "jets_name", + "num_jets_estimate", + "num_jets", + "equal_jets", + "show_num_jets", + "flavours", + "flavour_config", + "flavour_category", + } + + class TestPreprocessingConfig(unittest.TestCase): """unittest-based rewrite of the original pytest suite.""" @@ -57,7 +97,9 @@ def test_legacy_plotting_jet_count(self) -> None: skip_config_copy=True, ) - self.assertEqual(config.plotting.num_jets_plotting, config.num_jets_estimate_plotting) + self.assertEqual( + config.plotting.num_global_objects_plotting, config.num_global_objects_estimate_plotting + ) def test_plotting_config(self) -> None: config = PreprocessingConfig.from_file( @@ -67,7 +109,7 @@ def test_plotting_config(self) -> None: skip_config_copy=True, ) - self.assertEqual(config.plotting.num_jets_plotting, 100) + self.assertEqual(config.plotting.num_global_objects_plotting, 100) self.assertEqual(config.plotting.variable_label("pt"), "$p_\\mathrm{T}$ [GeV]") self.assertEqual(config.plotting.variable_label("mass"), "Jet Mass [GeV]") self.assertEqual(config.plotting.sample_label("ttbar"), "$t\\bar{t}$") @@ -189,7 +231,7 @@ def test_standard_flavour_config(self) -> None: "variables": {"jets": {"labels": ["test"]}}, }, base_dir=Path("/tmp/upp-tests/integration/temp_workspace/"), - flavour_category="standard", + class_category="standard", ) self.assertEqual(config.flavour_cont, Flavours) @@ -203,7 +245,7 @@ def test_extended_flavour_config(self) -> None: "variables": {"jets": {"labels": ["test"]}}, }, base_dir=Path("/tmp/upp-tests/integration/temp_workspace/"), - flavour_category="extended", + class_category="extended", ) self.assertEqual(config.flavour_cont, Extended_Flavours) @@ -247,13 +289,13 @@ def test_unsupported_flavour_config(self) -> None: "variables": {"jets": {"labels": ["test"]}}, }, base_dir=Path("/tmp/upp-tests/integration/temp_workspace/"), - flavour_category="error", + class_category="error", ) self.assertEqual( - "flavour_category error is not supported in the default " - + "flavours! If you want to use your own flavour config yaml file, please " - + "provide flavour_config!", + "class_category error is not supported in the default " + + "flavours! If you want to use your own class config yaml file, please " + + "provide class_config!", str(ctx.exception), ) @@ -267,9 +309,26 @@ def test_separate_flavour_config(self) -> None: "variables": {"jets": {"labels": ["test"]}}, }, base_dir=Path("/tmp/upp-tests/integration/temp_workspace/"), - flavour_config=self.CFG_DIR / "test_flavour_config.yaml", + class_config=self.CFG_DIR / "test_flavour_config.yaml", ) self.assertEqual( config.flavour_cont, LabelContainer.from_yaml(yaml_path=self.CFG_DIR / "test_flavour_config.yaml"), ) + + def test_relative_class_config_resolved_against_base_dir(self) -> None: + config = PreprocessingConfig( + config_path=self.CFG_DIR / "test.yaml", + split="train", + config={ + "resampling": {"variables": {"jets": {"labels": ["test"]}}, "target": "bjets"}, + "components": [], + "variables": {"jets": {"labels": ["test"]}}, + }, + base_dir=self.CFG_DIR, + class_config=Path("test_flavour_config.yaml"), + skip_checks=True, + ) + expected = (self.CFG_DIR / "test_flavour_config.yaml").absolute() + self.assertEqual(config.class_config, expected) + self.assertEqual(config.flavour_cont, LabelContainer.from_yaml(yaml_path=expected)) diff --git a/tests/unit/fixtures/test_config_pdf_auto_umami.yaml b/tests/unit/fixtures/test_config_pdf_auto_umami.yaml index 1b3c5499..4e11559e 100644 --- a/tests/unit/fixtures/test_config_pdf_auto_umami.yaml +++ b/tests/unit/fixtures/test_config_pdf_auto_umami.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: 5000 + num_global_objects_estimate: 5000 base_dir: /tmp/upp-tests/integration/temp_workspace/ out_dir: test_out - flavour_category: standard + class_category: standard variables: jets: diff --git a/tests/unit/fixtures/test_config_pdf_auto_umami_required.yaml b/tests/unit/fixtures/test_config_pdf_auto_umami_required.yaml index 805ba8b4..7b76eeca 100644 --- a/tests/unit/fixtures/test_config_pdf_auto_umami_required.yaml +++ b/tests/unit/fixtures/test_config_pdf_auto_umami_required.yaml @@ -32,43 +32,43 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [singlebjets] - num_jets: 7_000 + classes: [singlebjets] + num_global_objects: 7_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [singlecjets] - num_jets: 7_000 + classes: [singlecjets] + 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: [singlebjets] - num_jets: 3_000 + classes: [singlebjets] + num_global_objects: 3_000 - region: <<: *highpt sample: <<: *zprime - flavours: [singlecjets] - num_jets: 3_000 + classes: [singlecjets] + num_global_objects: 3_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ujets] - num_jets: 3_000 + classes: [ujets] + num_global_objects: 3_000 resampling: target: singlebjets @@ -82,10 +82,10 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 5000 + num_global_objects_estimate: 5000 base_dir: /tmp/upp-tests/integration/temp_workspace/ out_dir: test_out - flavour_category: extended + class_category: extended variables: jets: diff --git a/tests/unit/fixtures/test_config_rw.yaml b/tests/unit/fixtures/test_config_rw.yaml index 217ee590..7c507549 100644 --- a/tests/unit/fixtures/test_config_rw.yaml +++ b/tests/unit/fixtures/test_config_rw.yaml @@ -3,14 +3,14 @@ global_cuts: !include GN3V01/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "data1.h5" - "data3.h5" zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "data2.h5" @@ -30,15 +30,15 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets, ujets, taujets] - num_jets: -1 + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets, ujets, taujets] - num_jets: -1 + classes: [bjets, cjets, ujets, taujets] + num_global_objects: -1 resampling: target: ghostcjets @@ -56,7 +56,7 @@ resampling: bins: [[0, 2.5, 10]] reweighting: - num_jets_estimate: 100 + num_global_objects_estimate: 100 merge_num_proc: 1 reweights: - group: jets @@ -69,8 +69,8 @@ reweighting: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: /tmp/upp-tests/integration/temp_workspace/ out_dir: test_out diff --git a/tests/unit/stages/test_merging.py b/tests/unit/stages/test_merging.py index 4687f511..e9c2c052 100644 --- a/tests/unit/stages/test_merging.py +++ b/tests/unit/stages/test_merging.py @@ -22,21 +22,21 @@ # Stub H5Writer that just records calls, does NO real IO class MemWriter: - def __init__(self, dst, dtypes, shapes, jets_name, **_): + def __init__(self, dst, dtypes, shapes, global_objects_name, **_): self.dst = Path(dst) self.dtypes = dtypes self.shapes = shapes - self.num_jets = next(iter(shapes.values()))[0] - self.jets_name = jets_name + self.num_global_objects = next(iter(shapes.values()))[0] + self.global_objects_name = global_objects_name self.num_written = 0 self.attrs = {} # API that Merging calls def write(self, data: dict): - self.num_written += len(data[self.jets_name]) + self.num_written += len(data[self.global_objects_name]) def close(self): - assert self.num_written == self.num_jets + assert self.num_written == self.num_global_objects def add_attr(self, name, value, *_): self.attrs[name] = value @@ -59,7 +59,7 @@ class DummyComponent(SimpleNamespace): Simulates the subset of the real Component interface that Merging uses: - flavour - stream (generator of batches) - - num_jets + - num_global_objects - complete flag """ @@ -67,7 +67,7 @@ def __init__(self, flavour, jet_batches): super().__init__() self.flavour = Flavours[flavour] self._batches = list(jet_batches) - self.num_jets = sum(len(b["jets"]) for b in self._batches) + self.num_global_objects = sum(len(b["jets"]) for b in self._batches) self.complete = False self.out_path = Path("/dev/null") @@ -104,8 +104,8 @@ def _minimal_merging(monkeypatch, jets_per_file=10) -> merging_mod.Merging: components=SimpleNamespace(flavours=[Flavours["bjets"]]), variables=variables, batch_size=100, - jets_name="jets", - num_jets_per_output_file=jets_per_file, + global_name="jets", + num_global_objects_per_output_file=jets_per_file, file_tag="split", out_fname=Path("/tmp/merged.h5"), split="train", @@ -126,13 +126,13 @@ def test_add_jet_flavour_label(monkeypatch): jets = _jets_struct(5) comp = SimpleNamespace(flavour=Flavours["bjets"]) - tagged = merge.add_jet_flavour_label(jets, comp) + tagged = merge.add_global_object_label(jets, comp) assert "flavour_label" in tagged.dtype.names assert np.all(tagged["flavour_label"] == 0) # Calling again must not duplicate the column - same = merge.add_jet_flavour_label(tagged, comp) + same = merge.add_global_object_label(tagged, comp) assert same.dtype == tagged.dtype @@ -146,14 +146,16 @@ def test_open_writer_names_and_shapes(monkeypatch): merge._open_writer( sample=None, - jets_in_file=7, + global_objects_in_file=7, file_idx=0, - components=SimpleNamespace(unique_jets=True, jet_counts={}, dsids=[]), + components=SimpleNamespace( + unique_global_objects=True, global_object_counts=lambda _gn: {}, dsids=[] + ), ) writer = merge.writer assert isinstance(writer, MemWriter) - assert writer.num_jets == 7 + assert writer.num_global_objects == 7 assert writer.dst.name.startswith("merged_split_000") assert writer.dst.parent.name == "train" assert "flavour_label" in writer.attrs @@ -162,7 +164,7 @@ def test_open_writer_names_and_shapes(monkeypatch): def test_write_chunk_splits(monkeypatch): """Test split writing of write_chunk. - With num_jets_per_output_file=5 and a batch of 8 jets, write_chunk must + With num_global_objects_per_output_file=5 and a batch of 8 jets, write_chunk must create two MemWriter instances and write all jets. """ merge = _minimal_merging(monkeypatch, jets_per_file=5) @@ -175,10 +177,12 @@ def test_write_chunk_splits(monkeypatch): # Fake dtypes / shapes and open first writer merge.dtypes = {"jets": jets1["jets"].dtype} merge.base_shapes = {"jets": (8,)} - merge.total_jets = 8 + merge.total_global_objects = 8 merge._file_idx = 0 - merge.jets_written = 0 - merge.current_components = SimpleNamespace(unique_jets=True, jet_counts={}, dsids=[]) + merge.global_objects_written = 0 + merge.current_components = SimpleNamespace( + unique_global_objects=True, global_object_counts=lambda _gn: {}, dsids=[] + ) merge._sample = None merge._open_writer(None, 5, 0, merge.current_components) @@ -203,7 +207,7 @@ def test_write_chunk_rollover(monkeypatch): 1. close that writer, 2. open a fresh one (file_idx increments), 3. write the incoming batch into the new file, - 4. update jets_written. + 4. update global_objects_written. """ # Create a merger with 5 jets / file merge = _minimal_merging(monkeypatch, jets_per_file=5) @@ -216,11 +220,13 @@ def test_write_chunk_rollover(monkeypatch): # Fake the bookkeeping exactly as write_components() would merge.dtypes = {"jets": batch["jets"].dtype} merge.base_shapes = {"jets": (8,)} - merge.total_jets = 8 - merge.jets_written = 5 + merge.total_global_objects = 8 + merge.global_objects_written = 5 merge._file_idx = 0 merge._sample = None - merge.current_components = SimpleNamespace(unique_jets=True, jet_counts={}, dsids=[]) + merge.current_components = SimpleNamespace( + unique_global_objects=True, global_object_counts=lambda _gn: {}, dsids=[] + ) # Open the first writer with capacity 5 and mark it as "full" merge._open_writer(None, 5, 0, merge.current_components) @@ -232,7 +238,7 @@ def test_write_chunk_rollover(monkeypatch): assert n == 3 assert merge._file_idx == 1 assert merge.writer.num_written == 3 - assert merge.jets_written == 8 + assert merge.global_objects_written == 8 def test_write_chunk_returns_zero_when_no_space_left(monkeypatch): @@ -245,19 +251,21 @@ def test_write_chunk_returns_zero_when_no_space_left(monkeypatch): comps = [comp] # Mimic state after everything has already been written - merge.total_jets = 4 - merge.jets_written = 4 + merge.total_global_objects = 4 + merge.global_objects_written = 4 merge._file_idx = 0 - merge.current_components = SimpleNamespace(unique_jets=True, jet_counts={}, dsids=[]) + merge.current_components = SimpleNamespace( + unique_global_objects=True, global_object_counts=lambda _gn: {}, dsids=[] + ) merge._sample = None # We still need valid dtypes / shapes for _open_writer merge.dtypes = {"jets": jets_batch["jets"].dtype} merge.base_shapes = {"jets": (4,)} - # Open a writer that is already full (num_written == num_jets) + # Open a writer that is already full (num_written == num_global_objects) merge._open_writer(None, 4, 0, merge.current_components) - merge.writer.num_written = merge.writer.num_jets + merge.writer.num_written = merge.writer.num_global_objects n_written = merge.write_chunk(comps) @@ -266,7 +274,7 @@ def test_write_chunk_returns_zero_when_no_space_left(monkeypatch): # _file_idx is incremented once, but _open_writer was not called again assert merge._file_idx == 1 - assert merge.writer.num_written == merge.writer.num_jets + assert merge.writer.num_written == merge.writer.num_global_objects class ReaderStub: @@ -274,17 +282,17 @@ class ReaderStub: def __init__(self, batches: list[dict[str, np.ndarray]]): self._batches = batches - self.num_jets = sum(len(b["jets"]) for b in batches) + self.num_global_objects = sum(len(b["jets"]) for b in batches) def dtypes(self, _vars): # Assume at least one batch exists return {"jets": self._batches[0]["jets"].dtype} - def shapes(self, total_jets: int, _keys): + def shapes(self, total_global_objects: int, _keys): # Base-shapes are used only for dataset names and leading dim - return {"jets": (total_jets,)} + return {"jets": (total_global_objects,)} - def stream(self, _vars, _num_jets): + def stream(self, _vars, _num_global_objects): def _gen(): yield from self._batches @@ -297,7 +305,7 @@ class ComponentStub: def __init__(self, flavour: str, batches: list[dict[str, np.ndarray]]): self.flavour = Flavours[flavour] self._batches = batches - self.num_jets = sum(len(b["jets"]) for b in batches) + self.num_global_objects = sum(len(b["jets"]) for b in batches) self.complete = False self.out_path = Path("/dev/null") self.reader = None @@ -311,11 +319,13 @@ class ComponentsStub: def __init__(self, comps: list[ComponentStub]): self._comps = comps - self.num_jets = sum(c.num_jets for c in comps) - self.unique_jets = True - self.jet_counts: dict[str, int] = {} + self.num_global_objects = sum(c.num_global_objects for c in comps) + self.unique_global_objects = True self.dsids: list[int] = [] + def global_object_counts(self, _global_name="jets"): + return {} + def __iter__(self): return iter(self._comps) @@ -338,8 +348,8 @@ def _mk_merge_for_path(monkeypatch, out_path: Path, jets_per_file=5): components=SimpleNamespace(flavours=[Flavours["bjets"]]), variables=variables, batch_size=100, - jets_name="jets", - num_jets_per_output_file=jets_per_file, + global_name="jets", + num_global_objects_per_output_file=jets_per_file, file_tag="split", out_fname=out_path, split="train", @@ -390,7 +400,7 @@ def test_part_fname_formatting(monkeypatch, tmp_path): def test_detect_and_clean_completed_parts_empty_dir(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 10 + merge.total_global_objects = 10 merge.base_shapes = {"jets": (10,)} # no files exist idx = merge._detect_and_clean_completed_parts(None) @@ -399,7 +409,7 @@ def test_detect_and_clean_completed_parts_empty_dir(monkeypatch, tmp_path): def test_expected_rows_for_part_middle_and_tail(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 12 # 5 + 5 + 2 + merge.total_global_objects = 12 # 5 + 5 + 2 assert merge._expected_rows_for_part(0) == 5 assert merge._expected_rows_for_part(1) == 5 assert merge._expected_rows_for_part(2) == 2 @@ -408,7 +418,7 @@ def test_expected_rows_for_part_middle_and_tail(monkeypatch, tmp_path): def test_is_part_valid_happy_path(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 12 + merge.total_global_objects = 12 merge.base_shapes = {"jets": (12,)} # dataset names only; length unused here f0 = merge._part_fname(None, 0) @@ -419,7 +429,7 @@ def test_is_part_valid_happy_path(monkeypatch, tmp_path): def test_is_part_valid_multiple_datasets(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 5 + merge.total_global_objects = 5 merge.base_shapes = {"jets": (5,), "tracks": (5,)} f0 = merge._part_fname(None, 0) @@ -432,7 +442,7 @@ def test_is_part_valid_multiple_datasets(monkeypatch, tmp_path): def test_is_part_valid_wrong_length(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 12 + merge.total_global_objects = 12 merge.base_shapes = {"jets": (12,)} f1 = merge._part_fname(None, 1) @@ -442,7 +452,7 @@ def test_is_part_valid_wrong_length(monkeypatch, tmp_path): def test_is_part_valid_missing_jets_dataset(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 5 + merge.total_global_objects = 5 merge.base_shapes = {"jets": (5,)} f0 = merge._part_fname(None, 0) @@ -454,7 +464,7 @@ def test_is_part_valid_missing_jets_dataset(monkeypatch, tmp_path): def test_is_part_valid_mismatched_lengths(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 5 + merge.total_global_objects = 5 # Declare two expected datasets; only those present are compared merge.base_shapes = {"jets": (5,), "tracks": (5,)} @@ -468,7 +478,7 @@ def test_is_part_valid_mismatched_lengths(monkeypatch, tmp_path): def test_is_part_valid_corrupt_file(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 5 + merge.total_global_objects = 5 merge.base_shapes = {"jets": (5,)} f0 = merge._part_fname(None, 0) @@ -478,7 +488,7 @@ def test_is_part_valid_corrupt_file(monkeypatch, tmp_path): def test_detect_and_clean_completed_parts_counts_and_deletes(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 13 # parts: [5, 5, 3] + merge.total_global_objects = 13 # parts: [5, 5, 3] merge.base_shapes = {"jets": (13,)} # Create two valid parts 0 and 1 @@ -503,7 +513,7 @@ def test_detect_and_clean_completed_parts_counts_and_deletes(monkeypatch, tmp_pa def test_detect_and_clean_completed_parts_no_delete_when_disabled(monkeypatch, tmp_path): merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 10 + merge.total_global_objects = 10 merge.base_shapes = {"jets": (10,)} _write_valid_part(merge._part_fname(None, 0), rows=5) @@ -519,7 +529,7 @@ def test_detect_and_clean_completed_parts_no_delete_when_disabled(monkeypatch, t def test_detect_and_clean_handles_unlink_error(monkeypatch, tmp_path): """Simulate OSError during unlink to cover error logging branch.""" merge = _mk_merge_for_path(monkeypatch, tmp_path / "merged.h5", jets_per_file=5) - merge.total_jets = 10 + merge.total_global_objects = 10 merge.base_shapes = {"jets": (10,)} # Valid part 0 @@ -564,11 +574,11 @@ def test_resume_skips_completed_parts_and_opens_next(monkeypatch, tmp_path): open_calls = [] _orig_open = merging_mod.Merging._open_writer - def _wrapped_open(self, sample, jets_in_file, file_idx, components): + def _wrapped_open(self, sample, global_objects_in_file, file_idx, components): # Fast-forward must be finished when we open a real writer assert self._fast_forwarding is False - open_calls.append((sample, jets_in_file, file_idx)) - return _orig_open(self, sample, jets_in_file, file_idx, components) + open_calls.append((sample, global_objects_in_file, file_idx)) + return _orig_open(self, sample, global_objects_in_file, file_idx, components) monkeypatch.setattr(merging_mod.Merging, "_open_writer", _wrapped_open) @@ -577,13 +587,13 @@ def _wrapped_open(self, sample, jets_in_file, file_idx, components): # We expect the first (and only) open to be for part index 2 with capacity 3 assert len(open_calls) >= 1 - _, jets_in_file, file_idx = open_calls[0] + _, global_objects_in_file, file_idx = open_calls[0] assert file_idx == 2 - assert jets_in_file == 3 + assert global_objects_in_file == 3 # The MemWriter should have written exactly 3 jets in that last file assert isinstance(merge.writer, MemWriter) - assert merge.writer.num_jets == 3 + assert merge.writer.num_global_objects == 3 assert merge.writer.num_written == 3 @@ -606,11 +616,11 @@ def test_fast_forward_does_not_open_real_writer(monkeypatch, tmp_path): call_times = {"count": 0} _orig_open = merging_mod.Merging._open_writer - def _wrapped_open(self, sample, jets_in_file, file_idx, components): + def _wrapped_open(self, sample, global_objects_in_file, file_idx, components): # Must never be invoked while fast-forwarding assert self._fast_forwarding is False call_times["count"] += 1 - return _orig_open(self, sample, jets_in_file, file_idx, components) + return _orig_open(self, sample, global_objects_in_file, file_idx, components) monkeypatch.setattr(merging_mod.Merging, "_open_writer", _wrapped_open) @@ -648,10 +658,12 @@ def test_write_chunk_all_components_complete_early_return(monkeypatch): jets0 = _jets_struct(0) merge.dtypes = {"jets": jets0.dtype} merge.base_shapes = {"jets": (0,)} - merge.total_jets = 0 + merge.total_global_objects = 0 merge._file_idx = 0 - merge.jets_written = 0 - merge.current_components = SimpleNamespace(unique_jets=True, jet_counts={}, dsids=[]) + merge.global_objects_written = 0 + merge.current_components = SimpleNamespace( + unique_global_objects=True, global_object_counts=lambda _gn: {}, dsids=[] + ) merge._sample = None merge._open_writer(None, 0, 0, merge.current_components) @@ -660,10 +672,10 @@ def test_write_chunk_all_components_complete_early_return(monkeypatch): def test_write_components_single_file_mode(monkeypatch, tmp_path): - """When num_jets_per_output_file is None, no split suffix is added.""" + """When num_global_objects_per_output_file is None, no split suffix is added.""" out = tmp_path / "merged.h5" merge = _mk_merge_for_path(monkeypatch, out, jets_per_file=None) - merge.num_jets_per_output_file = None # ensure single-file mode + merge.num_global_objects_per_output_file = None # ensure single-file mode jets = {"jets": _jets_struct(7)} comp = ComponentStub("bjets", [jets]) @@ -672,7 +684,7 @@ def test_write_components_single_file_mode(monkeypatch, tmp_path): merge.write_components(sample=None, components=comps) assert isinstance(merge.writer, MemWriter) - assert merge.writer.num_jets == 7 + assert merge.writer.num_global_objects == 7 assert merge.writer.num_written == 7 assert merge.writer.dst.name == "merged.h5" @@ -704,8 +716,8 @@ def groupby_sample(self): components=FakeComponents(), variables=variables, batch_size=100, - jets_name="jets", - num_jets_per_output_file=10, + global_name="jets", + num_global_objects_per_output_file=10, file_tag="split", out_fname=tmp_path / "merged.h5", split="train", @@ -720,7 +732,7 @@ def groupby_sample(self): called = [] def _wc(self, sample, components): # noqa: ARG001 - called.append((sample, components.num_jets)) + called.append((sample, components.num_global_objects)) monkeypatch.setattr(merging_mod.Merging, "write_components", _wc) diff --git a/tests/unit/stages/test_plotting.py b/tests/unit/stages/test_plotting.py index 54420d95..45840068 100644 --- a/tests/unit/stages/test_plotting.py +++ b/tests/unit/stages/test_plotting.py @@ -38,19 +38,19 @@ def setup_method(self, method): "test": H5Reader( fname=self.fname1, batch_size=self.config.batch_size, - jets_name=self.config.jets_name, + global_objects_name=self.config.global_name, shuffle=False, - equal_jets=True, + equal_global_objects=True, ).load( { - self.config.jets_name: [ + self.config.global_name: [ "pt", "abs_eta", "mass", "HadronConeExclTruthLabelID", ] } - )[self.config.jets_name] + )[self.config.global_name] } print(f"setup_method, method: {method.__name__}") @@ -85,20 +85,32 @@ def test_make_hist_initial_no_pt(self): def test_plot_helpers_format_labels_and_ranges(): """Check compact labels and GeV unit conversion helpers.""" - assert plot_mod._format_num_jets(999) == "999" - assert plot_mod._format_num_jets(100_000) == "100k" - assert plot_mod._format_num_jets(10_000_000) == "10M" + assert plot_mod._format_num_global_objects(999) == "999" + assert plot_mod._format_num_global_objects(100_000) == "100k" + assert plot_mod._format_num_global_objects(10_000_000) == "10M" assert ( plot_mod._atlas_second_tag( "ttbar", "zprime", plotting=PlottingConfig(), - num_jets=100_000, + num_global_objects=100_000, resampling_status="Pre Resampling", ) == "$\\sqrt{s} = 13/13.6\\,\\mathrm{TeV}$, $t\\bar{t}$ + $Z'$ jets" "\nPre Resampling\n100k jets" ) + assert ( + plot_mod._atlas_second_tag( + "ttbar", + "zprime", + plotting=PlottingConfig(), + global_name="tracks", + num_global_objects=100_000, + resampling_status="Pre Resampling", + ) + == "$\\sqrt{s} = 13/13.6\\,\\mathrm{TeV}$, $t\\bar{t}$ + $Z'$ tracks" + "\nPre Resampling\n100k tracks" + ) assert plot_mod._display_range("pt_btagJes", (20_000, 250_000)) == (20, 250) assert plot_mod._display_range("JetFitterSecondaryVertex_mass", (0, 25_000)) == (0, 25) assert plot_mod._display_range("absEta_btagJes", (0, 2.5)) == (0, 2.5) @@ -169,7 +181,11 @@ def test_make_hist_adds_sample_linestyle_legend(monkeypatch, tmp_path): fname, file = get_mock_file(num_jets=100, fname=tmp_path / "sample.h5") file.close() values = H5Reader( - fname=fname, batch_size=100, jets_name="jets", shuffle=False, equal_jets=True + fname=fname, + batch_size=100, + global_objects_name="jets", + shuffle=False, + equal_global_objects=True, ).load({"jets": ["pt", "HadronConeExclTruthLabelID"]})["jets"] calls = [] monkeypatch.setattr( @@ -198,7 +214,7 @@ def test_post_resampling_paths_split_mode(tmp_path): out_fname=tmp_path / "pp_output_test.h5", split="test", merge_test_samples=False, - num_jets_per_output_file=10, + num_global_objects_per_output_file=10, components=components, ) @@ -229,19 +245,19 @@ def groupby_region(self): ) region_components = SimpleNamespace( flavours=[Flavours["bjets"]], - num_jets=100_000, + num_global_objects=100_000, ) return [(region, region_components)] config = SimpleNamespace( split="val", - plotting=PlottingConfig(num_jets_plotting=10_000), + plotting=PlottingConfig(num_global_objects_plotting=10_000), sampl_cfg=SimpleNamespace( vars=["pt"], bins={"pt": [[20_000, 250_000, 5]]}, ), components=FakeComponents(), - jets_name="jets", + global_name="jets", batch_size=100, out_dir=tmp_path, ) @@ -254,7 +270,7 @@ def fake_load_jets(_config, _in_paths, _vars_to_load): def fake_make_hist(**kwargs): calls.append(kwargs) - monkeypatch.setattr(plot_mod, "_load_jets", fake_load_jets) + monkeypatch.setattr(plot_mod, "_load_global_objects", fake_load_jets) monkeypatch.setattr(plot_mod, "make_hist", fake_make_hist) plot_mod._plot_initial(config) diff --git a/tests/unit/stages/test_reweight.py b/tests/unit/stages/test_reweight.py index 59b33f36..37622487 100644 --- a/tests/unit/stages/test_reweight.py +++ b/tests/unit/stages/test_reweight.py @@ -25,22 +25,25 @@ def _make_organised_components(tmpdir, jets_per_flavour): f.create_dataset("jets", data=data) files[flav] = [str(fpath)] - config = {"files": {"train": files}, "num_jets": {"train": jets_per_flavour}} + config = {"files": {"train": files}, "num_global_objects": {"train": jets_per_flavour}} config_path = tmpdir / "organised-components.yaml" with open(config_path, "w") as f: yaml.dump(config, f) return config_path -def _make_reweight_obj(tmpdir, jets_per_flavour, num_jets_estimate, batch_size=1000): +def _make_reweight_obj(tmpdir, jets_per_flavour, num_global_objects_estimate, batch_size=1000): """Create a Reweight instance with mocked config.""" config_path = _make_organised_components(tmpdir, jets_per_flavour) config = MagicMock() config.batch_size = batch_size config.base_dir = str(tmpdir) + config.global_name = "jets" - rw_config = SimpleNamespace(num_jets_estimate=num_jets_estimate, reweights=[]) + rw_config = SimpleNamespace( + num_global_objects_estimate=num_global_objects_estimate, reweights=[] + ) rw = object.__new__(Reweight) rw.config = config @@ -52,40 +55,40 @@ def _make_reweight_obj(tmpdir, jets_per_flavour, num_jets_estimate, batch_size=1 class TestGetInputReaders: def test_caps_at_available_jets(self, tmp_path): - """When a reader has fewer jets than num_jets_estimate, cap to available.""" + """When a reader has fewer jets than num_global_objects_estimate, cap to available.""" rw = _make_reweight_obj( tmp_path, jets_per_flavour={"bjets": 50, "cjets": 200}, - num_jets_estimate=100, + num_global_objects_estimate=100, ) - readers, per_reader_num_jets = rw.get_input_readers() + readers, per_reader_num_global_objects = rw.get_input_readers() assert len(readers) == 2 - assert len(per_reader_num_jets) == 2 + assert len(per_reader_num_global_objects) == 2 # bjets has 50 < 100, should be capped at 50 - assert per_reader_num_jets[0] == 50 + assert per_reader_num_global_objects[0] == 50 # cjets has 200 >= 100, should use estimate - assert per_reader_num_jets[1] == 100 + assert per_reader_num_global_objects[1] == 100 def test_all_above_estimate(self, tmp_path): - """When all readers have enough jets, use num_jets_estimate for all.""" + """When all readers have enough jets, use num_global_objects_estimate for all.""" rw = _make_reweight_obj( tmp_path, jets_per_flavour={"bjets": 500, "cjets": 300}, - num_jets_estimate=100, + num_global_objects_estimate=100, ) - _, per_reader_num_jets = rw.get_input_readers() - assert per_reader_num_jets == [100, 100] + _, per_reader_num_global_objects = rw.get_input_readers() + assert per_reader_num_global_objects == [100, 100] class TestCalculateWeightsStopIteration: """Test that the batch loop handles StopIteration from shorter readers.""" def test_unequal_reader_lengths(self, tmp_path): - """Readers with different num_jets don't crash the batch loop.""" + """Readers with different num_global_objects don't crash the batch loop.""" rw = _make_reweight_obj( tmp_path, jets_per_flavour={"bjets": 50, "cjets": 200}, - num_jets_estimate=200, + num_global_objects_estimate=200, batch_size=100, ) diff --git a/tests/unit/utils/test_check_input_samples.py b/tests/unit/utils/test_check_input_samples.py index 983d3594..ebc6b9cd 100644 --- a/tests/unit/utils/test_check_input_samples.py +++ b/tests/unit/utils/test_check_input_samples.py @@ -35,7 +35,7 @@ def error(self, *_a, **_k): # H5Reader stub: we don't care about values here, just that it's called class _H5: def __init__(self, **_kwargs): - self.num_jets = 123 + self.num_global_objects = 123 monkeypatch.setattr(cis, "H5Reader", _H5) @@ -52,7 +52,7 @@ def __init__(self, ntuple_dir: Path): self.config = {"blockX": {"pattern": [missing_dsid, missing_rtag, missing_both]}} self.ntuple_dir = ntuple_dir self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None cfg = _Cfg(tmp_path) @@ -89,7 +89,7 @@ def __init__(self, ntuple_dir: Path): self.config = {"blockY": {"pattern": {"oops": "dict-not-supported"}}} self.ntuple_dir = ntuple_dir self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None cfg = _Cfg(tmp_path) @@ -112,7 +112,7 @@ def __init__(self): # minimal attributes used by run_input_sample_check stub self.config = {} self.ntuple_dir = tmp_path self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None called = {"run": False, "cfg": None, "df": None, "v": None} @@ -160,7 +160,7 @@ def error(self, *_a, **_k): # H5Reader stub class _H5: def __init__(self, **_kwargs): - self.num_jets = 10 + self.num_global_objects = 10 monkeypatch.setattr(cis, "H5Reader", _H5) @@ -178,7 +178,7 @@ def __init__(self, ntuple_dir: Path): self.config = {"block": {"pattern": ["x.123456.y_r13167.z.h5"]}} self.ntuple_dir = ntuple_dir self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None cfg = _Cfg(tmp_path) @@ -205,7 +205,7 @@ def error(self, *_a, **_k): class _H5: def __init__(self, **_kwargs): - self.num_jets = 10 + self.num_global_objects = 10 monkeypatch.setattr(cis, "H5Reader", _H5) @@ -224,7 +224,7 @@ def __init__(self, ntuple_dir: Path): self.config = {"block": {"pattern": [sample]}} self.ntuple_dir = ntuple_dir self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None cfg = _Cfg(tmp_path) @@ -247,7 +247,7 @@ def test_script_entry_point_executes_main(tmp_path): class _H5: def __init__(self, **_kwargs): - self.num_jets = 1 + self.num_global_objects = 1 fake_hdf5.H5Reader = _H5 @@ -258,7 +258,7 @@ def __init__(self): self.config = {} self.ntuple_dir = tmp_path self.batch_size = 1 - self.jets_name = "jets" + self.global_name = "jets" self.vds_dir = None @classmethod diff --git a/upp/classes/components.py b/upp/classes/components.py index 62d60218..3bd2af03 100644 --- a/upp/classes/components.py +++ b/upp/classes/components.py @@ -37,12 +37,12 @@ class Component: Global cuts that should be applied for this component dirname : Path Directory of where this component is/will be stored - num_jets : int - Number of jets that are to be used from this component - num_jets_estimate_available : int - Estimated available jets for this component - equal_jets : bool - If the same number of jets should be used from the different samples + num_global_objects : int + Number of objects that are to be used from this component + num_global_objects_estimate_available : int + Estimated available objects for this component + equal_global_objects : bool + If the same number of objects should be used from the different samples """ region: Region @@ -50,14 +50,14 @@ class Component: flavour: Label global_cuts: Cuts dirname: Path - num_jets: int - num_jets_estimate_available: int - equal_jets: bool + num_global_objects: int + num_global_objects_estimate_available: int + equal_global_objects: bool def __post_init__(self): """Post init setup of internal variables.""" self.hist = Hist(self.dirname.parent.parent / "hists" / f"hist_{self.name}.h5") - self._unique_jets = -1 + self._unique_global_objects = -1 self._complete = None self._ups_ratio = None self._ups_max = None @@ -66,20 +66,20 @@ def __post_init__(self): def setup_reader( self, batch_size: int, - jets_name: str = "jets", + global_name: str = "jets", fname: Path | str | list[Path | str] | None = None, **kwargs, ) -> None: - """Set up the reader of the jets to load them from file. + """Set up the reader of the objects to load them from file. Parameters ---------- batch_size : int Batch size that is used for loading from file - jets_name : str, optional - Name of the group in which the jets are stored, by default "jets" + global_name : str, optional + Name of the group in which the objects are stored, by default "jets" fname : Path | str | list[Path | str] | None, optional - Filename of the file(s) from which the jets are loaded, by default None + Filename of the file(s) from which the objects are loaded, by default None **kwargs Additional kwargs passed to the H5Reader """ @@ -92,26 +92,26 @@ def setup_reader( self.reader = H5Reader( fname=fname, batch_size=batch_size, - jets_name=jets_name, - equal_jets=self.equal_jets, + global_objects_name=global_name, + equal_global_objects=self.equal_global_objects, **kwargs, ) log.debug(f"Setup component reader at: {fname}") - def setup_writer(self, variables: VariableConfig, jets_name: str = "jets") -> None: - """Set up the writer of the jets to file. + def setup_writer(self, variables: VariableConfig, global_name: str = "jets") -> None: + """Set up the writer of the objects to file. Parameters ---------- variables : VariableConfig Instance of VariableConfig in which the variables are stored. - jets_name : str, optional - Name of the group in which the jets are stored, by default "jets" + global_name : str, optional + Name of the group in which the objects are stored, by default "jets" """ dtypes = self.reader.dtypes(variables.combined()) - # num_jets == -1 ("write all") -> 0 leading dim so the writer grows dynamically - shapes = self.reader.shapes(max(self.num_jets, 0), variables.keys()) - self.writer = H5Writer(self.out_path, dtypes, shapes, jets_name=jets_name) + # num_global_objects == -1 ("write all") -> 0 leading dim so the writer grows dynamically + shapes = self.reader.shapes(max(self.num_global_objects, 0), variables.keys()) + self.writer = H5Writer(self.out_path, dtypes, shapes, global_objects_name=global_name) log.debug(f"Setup component writer at: {self.out_path}") @property @@ -162,27 +162,29 @@ def is_target(self, target_str: str) -> bool: """ return self.flavour.name == target_str - def get_jets(self, variables: list, num_jets: int, cuts: Cuts | None = None) -> dict: - """Load jets from file. + def get_global_objects( + self, variables: list, num_global_objects: int, cuts: Cuts | None = None + ) -> dict: + """Load objects from file. Parameters ---------- variables : list Variables that are to be loaded - num_jets : int - Number of jets that are to be loaded + num_global_objects : int + Number of objects that are to be loaded cuts : Cuts | None, optional - Cuts instance of the cuts that should be applied on the jets, by default None + Cuts instance of the cuts that should be applied on the objects, by default None Returns ------- dict - Dict with the loaded jets + Dict with the loaded objects """ - jn = self.reader.jets_name - return self.reader.load({jn: variables}, num_jets, cuts)[jn] + jn = self.reader.global_objects_name + return self.reader.load({jn: variables}, num_global_objects, cuts)[jn] - def check_num_jets( + def check_num_global_objects( self, num_req: int, sampling_fraction: float | None = None, @@ -190,38 +192,40 @@ def check_num_jets( silent: bool = False, raise_error: bool = True, ) -> None: - """Check the number of available jets. + """Check the number of available objects. - If more jets are requested than available, throw an Error. + If more objects are requested than available, throw an Error. Parameters ---------- num_req : int - Number of requested jets + Number of requested objects sampling_fraction : float | None, optional Sampling , by default None cuts : Cuts | None, optional - Cuts instance of the cuts that are to be applied on the jets, by default None + Cuts instance of the cuts that are to be applied on the objects, by default None silent : bool, optional Decide, if the debug and info log statements are printed, by default False raise_error : bool, optional - Decide if the error should be raised if not enough jets are available, + Decide if the error should be raised if not enough objects are available, by default True Raises ------ ValueError - If more jets are requsted than available + If more objects are requsted than available """ - # num_req < 0 means "use all available jets" - nothing to check + # num_req < 0 means "use all available objects" - nothing to check if num_req < 0: return - # Check if num_jets jets are aviailable after the cuts and sampling fraction + # Check if num_global_objects objects are aviailable after the cuts and sampling fraction num_est = ( - None if self.num_jets_estimate_available <= 0 else self.num_jets_estimate_available + None + if self.num_global_objects_estimate_available <= 0 + else self.num_global_objects_estimate_available ) - total = self.reader.estimate_available_jets(cuts, num_est) + total = self.reader.estimate_available_global_objects(cuts, num_est) available = total if sampling_fraction: available = int(total * sampling_fraction) @@ -229,22 +233,22 @@ def check_num_jets( # check with tolerance to avoid failure midway through preprocessing if available < num_req and raise_error: raise ValueError( - f"{num_req:,} jets requested, but only {total:,} are estimated to be" + f"{num_req:,} objects requested, but only {total:,} are estimated to be" f" in {self}. With a sampling fraction of {sampling_fraction}, at most" f" {available:,} of these are available. You can either reduce the" - " number of requested jets or increase the sampling fraction." + " number of requested objects or increase the sampling fraction." ) if not silent: log.debug(f"Sampling fraction {sampling_fraction}") log.info( - f"Estimated {available:,} {self} jets available - {num_req:,} requested" - f"({self.reader.num_jets:,} in {self.sample})" + f"Estimated {available:,} {self} objects available - {num_req:,} requested" + f"({self.reader.num_global_objects:,} in {self.sample})" ) def get_auto_sampling_fraction( self, - num_jets: int, + num_global_objects: int, cuts: Cuts | None = None, silent: bool = False, ) -> float: @@ -252,10 +256,10 @@ def get_auto_sampling_fraction( Parameters ---------- - num_jets : int - Number of jets available + num_global_objects : int + Number of objects available cuts : Cuts | None, optional - Cuts instance of the cuts that should be applied on the jets, by default None + Cuts instance of the cuts that should be applied on the objects, by default None silent : bool, optional Decide, if the debug and info log statements are printed, by default False @@ -265,10 +269,12 @@ def get_auto_sampling_fraction( Automatically estimated sampling fraction """ num_est = ( - None if self.num_jets_estimate_available <= 0 else self.num_jets_estimate_available + None + if self.num_global_objects_estimate_available <= 0 + else self.num_global_objects_estimate_available ) - total = self.reader.estimate_available_jets(cuts, num_est) - auto_sampling_frac = round(1.1 * num_jets / total, 3) # 1.1 is a tolerance factor + total = self.reader.estimate_available_global_objects(cuts, num_est) + auto_sampling_frac = round(1.1 * num_global_objects / total, 3) # 1.1 is a tolerance factor if not silent: log.debug(f"optimal sampling fraction {auto_sampling_frac:.3f}") return auto_sampling_frac @@ -284,18 +290,19 @@ def __str__(self) -> str: return self.name @property - def unique_jets(self) -> int: - """Return the number of unique jets for this component. + def unique_global_objects(self) -> int: + """Return the number of unique objects for this component. Returns ------- int - Number of unique jets for this component + Number of unique objects for this component """ - if self._unique_jets == -1: - self._unique_jets = sum([r.get_attr("unique_jets") for r in self.reader.readers]) + if self._unique_global_objects == -1: + attr = f"unique_{self.reader.global_objects_name}" + self._unique_global_objects = sum([r.get_attr(attr) for r in self.reader.readers]) - return self._unique_jets + return self._unique_global_objects class Components: @@ -320,9 +327,9 @@ def from_config(cls, config: PreprocessingConfig) -> Components: """ component_list = [] for component in config.config["components"]: - # Ensure equal_jets flag is correctly set - assert "equal_jets" not in component, ( - "equal_jets flag should be set in the sample config" + # Ensure equal_global_objects flag is correctly set + assert "equal_global_objects" not in component, ( + "equal_global_objects flag should be set in the sample config" ) # Get the region cuts @@ -333,9 +340,9 @@ def from_config(cls, config: PreprocessingConfig) -> Components: # Get the region and apply the region cuts region = Region(component["region"]["name"], region_cuts + config.global_cuts) - # Load the pattern and the equal_jets settings + # Load the pattern and the equal_global_objects settings pattern = component["sample"]["pattern"] - equal_jets = component["sample"].get("equal_jets", True) + equal_global_objects = component["sample"].get("equal_global_objects", True) if isinstance(pattern, list): pattern = tuple(pattern) @@ -349,12 +356,16 @@ def from_config(cls, config: PreprocessingConfig) -> Components: ) # Create the Component instances for the different flavours - for name in component["flavours"]: - num_jets = component["num_jets"] + for name in component["classes"]: + num_global_objects = component["num_global_objects"] if config.split == "val": - num_jets = component.get("num_jets_val", num_jets // 10) + num_global_objects = component.get( + "num_global_objects_val", num_global_objects // 10 + ) elif config.split == "test": - num_jets = component.get("num_jets_test", num_jets // 10) + num_global_objects = component.get( + "num_global_objects_test", num_global_objects // 10 + ) component_list.append( Component( region=region, @@ -362,9 +373,9 @@ def from_config(cls, config: PreprocessingConfig) -> Components: flavour=config.flavour_cont[name], global_cuts=config.global_cuts, dirname=config.components_dir, - num_jets=num_jets, - num_jets_estimate_available=config.num_jets_estimate_available, # type: ignore - equal_jets=equal_jets, + num_global_objects=num_global_objects, + num_global_objects_estimate_available=config.num_global_objects_estimate_available, # type: ignore + equal_global_objects=equal_global_objects, ) ) components = cls(component_list) @@ -388,7 +399,9 @@ def check_flavour_ratios(self) -> None: for region, components in self.groupby_region(): this_ratios = {} for f in flavours: - this_ratios[f.name] = components[f].num_jets / components.num_jets + this_ratios[f.name] = ( + components[f].num_global_objects / components.num_global_objects + ) ratios[region] = this_ratios ref = next(iter(ratios.values())) @@ -445,26 +458,26 @@ def cuts(self) -> Cuts: return sum((c.cuts for c in self), Cuts.from_list([])) @property - def num_jets(self) -> int: - """Return the number of jets available. + def num_global_objects(self) -> int: + """Return the number of objects available. Returns ------- int - Number of available jets + Number of available objects """ - return sum(c.num_jets for c in self) + return sum(c.num_global_objects for c in self) @property - def unique_jets(self) -> int: - """Return the number of unique jets available. + def unique_global_objects(self) -> int: + """Return the number of unique objects available. Returns ------- int - Number of available unique jets + Number of available unique objects """ - return sum(c.unique_jets for c in self) + return sum(c.unique_global_objects for c in self) @property def out_dir(self): @@ -472,14 +485,29 @@ def out_dir(self): assert len(out_dir) == 1 return next(iter(out_dir)) - @property - def jet_counts(self): + def global_object_counts(self, global_name: str = "jets") -> dict: + """Return per-component and total object counts. + + Parameters + ---------- + global_name : str, optional + Name of the global object, used to key the counts, by default "jets". + + Returns + ------- + dict + Counts keyed by ``num_{global_name}`` and ``unique_{global_name}``. + """ num_dict = { - c.name: {"num_jets": int(c.num_jets), "unique_jets": int(c.unique_jets)} for c in self + c.name: { + f"num_{global_name}": int(c.num_global_objects), + f"unique_{global_name}": int(c.unique_global_objects), + } + for c in self } num_dict["total"] = { - "num_jets": int(self.num_jets), - "unique_jets": int(self.unique_jets), + f"num_{global_name}": int(self.num_global_objects), + f"unique_{global_name}": int(self.unique_global_objects), } return num_dict diff --git a/upp/classes/plotting_config.py b/upp/classes/plotting_config.py index d87ed361..28022e3c 100644 --- a/upp/classes/plotting_config.py +++ b/upp/classes/plotting_config.py @@ -28,9 +28,9 @@ class PlottingConfig: Attributes ---------- - num_jets_plotting : int | None, optional - Number of jets loaded for plotting. If not set, use the global - `num_jets_estimate_plotting` value. By default None. + num_global_objects_plotting : int | None, optional + Number of objects loaded for plotting. If not set, use the global + `num_global_objects_estimate_plotting` value. By default None. variable_labels : dict[str, str], optional Display labels for plotted variables. Keys are matched case-insensitively against variable names, with the longest matching key taking precedence. @@ -39,14 +39,14 @@ class PlottingConfig: Display labels for input samples. User-provided labels are merged with the default ttbar and zprime labels. ylabel : str, optional - Label for the y-axis. The `{jets_name}` placeholder is replaced with the - configured jet dataset name. By default "Normalised Number of {jets_name}". + Label for the y-axis. The `{global_name}` placeholder is replaced with the + configured global object name. By default "Normalised Number of {global_name}". atlas_first_tag : str, optional First ATLAS plot label. By default "Simulation Internal". atlas_second_tag : str, optional Second ATLAS plot label. By default "$\\sqrt{s} = 13/13.6\\,\\mathrm{TeV}$". - show_num_jets : bool, optional - Decide, if the number of jets is shown in the ATLAS second tag + show_num_global_objects : bool, optional + Decide, if the number of objects is shown in the ATLAS second tag output_formats : list[str], optional File formats in which each plot is saved. By default `["pdf", "png"]`. linestyles : list[str], optional @@ -76,13 +76,13 @@ class PlottingConfig: By default "plots". """ - num_jets_plotting: int | None = None + num_global_objects_plotting: int | None = None variable_labels: dict[str, str] = field(default_factory=_default_variable_labels) sample_labels: dict[str, str] = field(default_factory=_default_sample_labels) - ylabel: str = "Normalised Number of {jets_name}" + ylabel: str = "Normalised Number of {global_name}" atlas_first_tag: str = "Simulation Internal" atlas_second_tag: str = "$\\sqrt{s} = 13/13.6\\,\\mathrm{TeV}$" - show_num_jets: bool = True + show_num_global_objects: bool = True output_formats: list[str] = field(default_factory=lambda: ["pdf", "png"]) linestyles: list[str] = field(default_factory=lambda: ["-", "--", "-.", ":"]) bins: int = 50 @@ -99,8 +99,10 @@ class PlottingConfig: def __post_init__(self) -> None: self.variable_labels = {**_default_variable_labels(), **self.variable_labels} self.sample_labels = {**_default_sample_labels(), **self.sample_labels} - if self.num_jets_plotting is not None and self.num_jets_plotting <= 0: - raise ValueError("plotting.num_jets_plotting must be a positive integer or None") + if self.num_global_objects_plotting is not None and self.num_global_objects_plotting <= 0: + raise ValueError( + "plotting.num_global_objects_plotting must be a positive integer or None" + ) if not self.output_formats: raise ValueError("plotting.output_formats must contain at least one format") if not self.linestyles: diff --git a/upp/classes/preprocessing_config.py b/upp/classes/preprocessing_config.py index 1a7b0a97..bc8ccb96 100644 --- a/upp/classes/preprocessing_config.py +++ b/upp/classes/preprocessing_config.py @@ -34,6 +34,42 @@ Split = Literal["train", "val", "test"] +# Deprecated config keys mapped to their generalised (object-agnostic) names. +# Old configs keep working: the keys are remapped on load (see from_file). +LEGACY_KEY_MAP = { + "jets_name": "global_name", + "num_jets": "num_global_objects", + "num_jets_val": "num_global_objects_val", + "num_jets_test": "num_global_objects_test", + "num_jets_estimate": "num_global_objects_estimate", + "num_jets_estimate_available": "num_global_objects_estimate_available", + "num_jets_estimate_hist": "num_global_objects_estimate_hist", + "num_jets_estimate_norm": "num_global_objects_estimate_norm", + "num_jets_estimate_plotting": "num_global_objects_estimate_plotting", + "num_jets_per_output_file": "num_global_objects_per_output_file", + "num_jets_plotting": "num_global_objects_plotting", + "show_num_jets": "show_num_global_objects", + "equal_jets": "equal_global_objects", + "flavours": "classes", + "flavour_config": "class_config", + "flavour_category": "class_category", +} + + +def _rename_legacy_keys(obj, found: set[str]): + """Recursively rename deprecated object-named config keys to their new names.""" + if isinstance(obj, dict): + out = {} + for key, value in obj.items(): + if key in LEGACY_KEY_MAP: + found.add(key) + key = LEGACY_KEY_MAP[key] + out[key] = _rename_legacy_keys(value, found) + return out + if isinstance(obj, list): + return [_rename_legacy_keys(value, found) for value in obj] + return obj + @dataclass class PreprocessingConfig: @@ -47,9 +83,9 @@ class PreprocessingConfig: For example: ```yaml global: - jets_name: jets + global_name: objects batch_size: 1_000_000 - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 base_dir: /my/stuff/ ntuple_dir: h5-inputs # resolved path: /my/stuff/h5-inputs/ ``` @@ -82,40 +118,49 @@ class PreprocessingConfig: `sampling_fraction*batch_size_after_cuts`. It is recommended to choose high batch sizes especially to the `countup` method to achive best agreement of target and resampled distributions. By default 100_000 - num_jets_estimate : int, optional + num_global_objects_estimate : int, optional Any of the further three arguments that are not specified will default to this value Is equal to 1_000_000 by default. - num_jets_estimate_available : int | None, optional - A sabsample taken from the whole sample to estimate the number of jets after the cuts. + num_global_objects_estimate_available : int | None, optional + A sabsample taken from the whole sample to estimate the number of objects after the cuts. Please keep this number high in order to not get poisson error of more then 5%. - If time allows you can use -1 to get a precise number of jets and not just an estimate - although it will be slow for large datasets. Is equal to num_jets_estimate by default. - num_jets_estimate_hist : int | None, optional - Number of jets of each flavour that are used to construct histograms for probability + If time allows you can use -1 to get a precise number of objects and not just an estimate + although it will be slow for large datasets. + Is equal to num_global_objects_estimate by default. + num_global_objects_estimate_hist : int | None, optional + Number of objects of each flavour that are used to construct histograms for probability density function estimation. Larger numbers give a better quality estmate of the pdfs. - Is equal to num_jets_estimate by default. - num_jets_estimate_norm : int | None, optional - Number of jets of each flavour that are used to estimate shifting and scaling during + Is equal to num_global_objects_estimate by default. + num_global_objects_estimate_norm : int | None, optional + Number of objects of each flavour that are used to estimate shifting and scaling during normalisation step. Larger numbers give a better quality estmates. - Is equal to num_jets_estimate by default. - num_jets_estimate_plotting : int | None, optional - Number of jets of each flavour used for plotting the initial and the final resampling + Is equal to num_global_objects_estimate by default. + num_global_objects_estimate_plotting : int | None, optional + Number of objects of each flavour used for plotting the initial and the final resampling variable distributions. Larger numbers give a better estimate of the full distributions. - Is equal to num_jets_estimate by default. + Is equal to num_global_objects_estimate by default. merge_test_samples : bool, optional Merge the test samples of the different processes into one file. By default False. - jets_name : str, optional - Name of the jets dataset in the input file. By default "jets". - flavour_config : Path | None, optional - Flavour config yaml file which is to be used. By default None - flavour_category : str, optional - Flavour categories that are to be used. By default, the "standard" (non-extended) - labels are loaded. The extended labels can be used by setting this value to "extended". - By default "standard". To use this option, flavour_config must be None. - num_jets_per_output_file : int | None, optional - Number of jets per final output file. If the number of total jets is larger + global_name : str, optional + Name of the global (per-object) dataset in the input file, e.g. the objects. + By default "jets". + class_config : Path | None, optional + Path to a custom class-definition yaml, used instead of the flavour labels + bundled with atlas-ftag-tools. This lets the framework classify any object type, + not just jets. The file is a list of class dicts with keys ``name``, ``label``, + ``cuts``, ``colour`` and ``category`` (plus an optional ``_px`` probability name); + the ``classes`` listed per component must match the ``name`` entries here. A + relative path is resolved against ``base_dir``. Takes precedence over + ``class_category``. By default None + class_category : str, optional + Class categories that are to be used from the atlas-ftag-tools bundled labels. + By default, the "standard" (non-extended) labels are loaded. The extended labels + can be used by setting this value to "extended". By default "standard". To use + this option, class_config must be None. + num_global_objects_per_output_file : int | None, optional + Number of objects per final output file. If the number of total objects is larger than this number, the final h5 output files are splitted in multiple smaller - files with this number of jets per file. By default None which produces one + files with this number of objects per file. By default None which produces one huge output file. skip_checks : bool, optional Skip checks for the input files. This is used for grid submission @@ -135,31 +180,33 @@ class PreprocessingConfig: out_dir: Path = Path("output") out_fname: Path = Path("pp_output.h5") batch_size: int = 100_000 - num_jets_estimate: int = 1_000_000 - num_jets_estimate_available: int | None = None - num_jets_estimate_hist: int | None = None - num_jets_estimate_norm: int | None = None - num_jets_estimate_plotting: int | None = None + num_global_objects_estimate: int = 1_000_000 + num_global_objects_estimate_available: int | None = None + num_global_objects_estimate_hist: int | None = None + num_global_objects_estimate_norm: int | None = None + num_global_objects_estimate_plotting: int | None = None merge_test_samples: bool = False - jets_name: str = "jets" - flavour_config: Path | None = None - flavour_category: str = "standard" - num_jets_per_output_file: int | None = None + global_name: str = "jets" + class_config: Path | None = None + class_category: str = "standard" + num_global_objects_per_output_file: int | None = None skip_checks: bool = False skip_config_copy: bool = False vds_dir: Path | None = None def __post_init__(self): # postprocess paths - if self.num_jets_estimate: - if self.num_jets_estimate_available is None: - self.num_jets_estimate_available = max(self.num_jets_estimate, int(1e6)) - if self.num_jets_estimate_hist is None: - self.num_jets_estimate_hist = self.num_jets_estimate - if self.num_jets_estimate_norm is None: - self.num_jets_estimate_norm = self.num_jets_estimate - if self.num_jets_estimate_plotting is None: - self.num_jets_estimate_plotting = self.num_jets_estimate + if self.num_global_objects_estimate: + if self.num_global_objects_estimate_available is None: + self.num_global_objects_estimate_available = max( + self.num_global_objects_estimate, int(1e6) + ) + if self.num_global_objects_estimate_hist is None: + self.num_global_objects_estimate_hist = self.num_global_objects_estimate + if self.num_global_objects_estimate_norm is None: + self.num_global_objects_estimate_norm = self.num_global_objects_estimate + if self.num_global_objects_estimate_plotting is None: + self.num_global_objects_estimate_plotting = self.num_global_objects_estimate for field in dataclasses.fields(self): if field.type == "Path" and field.name != "out_fname" and field.name != "base_dir": @@ -167,27 +214,30 @@ def __post_init__(self): # vds_dir is optional (Path | None), so the loop above skips it; resolve it here if self.vds_dir is not None: self.vds_dir = self.get_path(Path(self.vds_dir)) + # class_config is optional (Path | None) and also skipped above; resolve it too + if self.class_config is not None: + self.class_config = self.get_path(Path(self.class_config)) if not self.ntuple_dir.exists() and not self.skip_checks: raise FileNotFoundError(f"Path {self.ntuple_dir} does not exist") self.components_dir = self.components_dir / self.split self.out_fname = self.out_dir / path_append(self.out_fname, self.split) # Define the content of the flavour label container - if self.flavour_config: + if self.class_config: self.flavour_cont = LabelContainer.from_yaml( - yaml_path=self.flavour_config, + yaml_path=self.class_config, ) - elif self.flavour_category == "standard": + elif self.class_category == "standard": self.flavour_cont = Flavours - elif self.flavour_category == "extended": + elif self.class_category == "extended": self.flavour_cont = Extended_Flavours else: raise ValueError( - f"flavour_category {self.flavour_category} is not supported in the default " - "flavours! If you want to use your own flavour config yaml file, please " - "provide flavour_config!" + f"class_category {self.class_category} is not supported in the default " + "flavours! If you want to use your own class config yaml file, please " + "provide class_config!" ) # configure classes if sampl_cfg := self.config.get("resampling", None): @@ -209,10 +259,10 @@ def __post_init__(self): # configure variables self.variables = VariableConfig( - self.config["variables"], self.jets_name, self.is_test, selectors + self.config["variables"], self.global_name, self.is_test, selectors ) if self.sampl_cfg is not None and self.sampl_cfg.variables: - self.variables = self.variables.add_jet_vars( + self.variables = self.variables.add_global_vars( list(self.config["resampling"]["variables"].keys()), "labels" ) self.transform = ( @@ -227,8 +277,8 @@ def __post_init__(self): else None ) self.plotting = PlottingConfig(**self.config.get("plotting", {})) - if self.plotting.num_jets_plotting is None: - self.plotting.num_jets_plotting = self.num_jets_estimate_plotting + if self.plotting.num_global_objects_plotting is None: + self.plotting.num_global_objects_plotting = self.num_global_objects_estimate_plotting # reproducibility try: @@ -259,6 +309,14 @@ def from_file( raise FileNotFoundError(f"{config_path} does not exist - check your --config arg") with open(config_path) as file: config = yaml.safe_load(file) + legacy: set[str] = set() + config = _rename_legacy_keys(config, legacy) + if legacy: + log.warning( + "Deprecated object-named config keys %s were remapped to their " + "global-object names; please update your config.", + sorted(legacy), + ) return cls( config_path=config_path, split=split, diff --git a/upp/classes/reweight_config.py b/upp/classes/reweight_config.py index a488621c..8049bd75 100644 --- a/upp/classes/reweight_config.py +++ b/upp/classes/reweight_config.py @@ -8,14 +8,14 @@ @dataclass class ReweightConfig: - # Number of jets to estimate, if None, use the global num jets estimate - num_jets_estimate: None | int = None + # Number of objects to estimate, if None, use the global num objects estimate + num_global_objects_estimate: None | int = None merge_num_proc: int = 1 # Number of processes to use for merging 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.num_global_objects_estimate is not None and self.num_global_objects_estimate <= 0: + raise ValueError("num_global_objects_estimate must be a positive integer or None") parsed_reweights = [] for rw in self.reweights: diff --git a/upp/classes/variable_config.py b/upp/classes/variable_config.py index 023ebab6..1dad96bd 100644 --- a/upp/classes/variable_config.py +++ b/upp/classes/variable_config.py @@ -9,7 +9,7 @@ @dataclass(frozen=True) class VariableConfig: variables: dict[str, dict[str, list[str]]] - jets_name: str = "jets" + global_name: str = "jets" keep_all: bool = False selectors: dict[str, TrackSelector] | None = None @@ -27,17 +27,19 @@ def combined(self): return combined @property - def jets(self): - return self[self.jets_name] + def global_objects(self): + return self[self.global_name] @property def tracks(self): - return {name: var for name, var in self.variables.items() if name != self.jets_name} + return {name: var for name, var in self.variables.items() if name != self.global_name} - def add_jet_vars(self, variables: list[str], kind: str = "inputs") -> VariableConfig: + def add_global_vars(self, variables: list[str], kind: str = "inputs") -> VariableConfig: """Return a new VariableConfig instance.""" - vc = VariableConfig(deepcopy(self.variables), self.jets_name, self.keep_all, self.selectors) - vc.jets[kind] = list(dict.fromkeys(vc.jets[kind] + variables)) + vc = VariableConfig( + deepcopy(self.variables), self.global_name, self.keep_all, self.selectors + ) + vc.global_objects[kind] = list(dict.fromkeys(vc.global_objects[kind] + variables)) return vc def items(self): diff --git a/upp/configs/GN3EPCMV01/GN3EPCMV01.yaml b/upp/configs/GN3EPCMV01/GN3EPCMV01.yaml index 2be01a9b..0b63b560 100644 --- a/upp/configs/GN3EPCMV01/GN3EPCMV01.yaml +++ b/upp/configs/GN3EPCMV01/GN3EPCMV01.yaml @@ -3,7 +3,7 @@ global_cuts: !include /home/users/r/reisch/production/slurm-scripts/umami-prepro ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.treisch.601589.e8547_s3797_r13144_p7085.tdd.GN3_dev.25_2_76.26-01-07_CentralDump_p7085_output.h5/*.h5" # MC20d - "user.treisch.601589.e8549_s4159_r15530_p7085.tdd.GN3_dev.25_2_76.26-01-07_CentralDump_p7085_output.h5/*.h5" # MC23d @@ -13,7 +13,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.treisch.802818.e8599_s3797_r13144_p7085.tdd.GN3_dev.25_2_76.26-01-07_CentralDump_p7085_output.h5/*.h5" # MC20d - "user.treisch.802818.e8599_s4159_r15530_p7085.tdd.GN3_dev.25_2_76.26-01-07_CentralDump_p7085_output.h5/*.h5" # MC23d @@ -37,109 +37,109 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [ghostbjets] - num_jets: 168_968_784 - num_jets_test: 2_000_000 - num_jets_val: 2_000_000 + classes: [ghostbjets] + num_global_objects: 168_968_784 + num_global_objects_test: 2_000_000 + num_global_objects_val: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostcjets] - num_jets: 38_872_067 - num_jets_test: 2_000_000 - num_jets_val: 2_000_000 + classes: [ghostcjets] + num_global_objects: 38_872_067 + num_global_objects_test: 2_000_000 + num_global_objects_val: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsjets] - num_jets: 30_714_894 - num_jets_test: 1_706_567 - num_jets_val: 1_706_567 + classes: [ghostsjets] + num_global_objects: 30_714_894 + num_global_objects_test: 1_706_567 + num_global_objects_val: 1_706_567 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostudjets] - num_jets: 78_352_864 - num_jets_test: 2_000_000 - num_jets_val: 2_000_000 + classes: [ghostudjets] + num_global_objects: 78_352_864 + num_global_objects_test: 2_000_000 + num_global_objects_val: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostgjets] - num_jets: 92_319_320 - num_jets_test: 2_000_000 - num_jets_val: 2_000_000 + classes: [ghostgjets] + num_global_objects: 92_319_320 + num_global_objects_test: 2_000_000 + num_global_objects_val: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghosttaujets] - num_jets: 16_261_627 - num_jets_test: 896_611 - num_jets_val: 896_611 + classes: [ghosttaujets] + num_global_objects: 16_261_627 + num_global_objects_test: 896_611 + num_global_objects_val: 896_611 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostbjets] - num_jets: 40_230_662 #40_696_238 - num_jets_test: 476_190 - num_jets_val: 476_190 + classes: [ghostbjets] + num_global_objects: 40_230_662 #40_696_238 + num_global_objects_test: 476_190 + num_global_objects_val: 476_190 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostcjets] - num_jets: 9_255_254 #39_928_286 - num_jets_test: 476_190 - num_jets_val: 476_190 + classes: [ghostcjets] + num_global_objects: 9_255_254 #39_928_286 + num_global_objects_test: 476_190 + num_global_objects_val: 476_190 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsjets] - num_jets: 7_313_070 #25_814_913 - num_jets_test: 406_325 - num_jets_val: 406_325 + classes: [ghostsjets] + num_global_objects: 7_313_070 #25_814_913 + num_global_objects_test: 406_325 + num_global_objects_val: 406_325 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostudjets] - num_jets: 18_655_443 #28_206_672 - num_jets_test: 476_190 - num_jets_val: 476_190 + classes: [ghostudjets] + num_global_objects: 18_655_443 #28_206_672 + num_global_objects_test: 476_190 + num_global_objects_val: 476_190 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostgjets] - num_jets: 21_980_790 #55_036_052 - num_jets_test: 476_190 - num_jets_val: 476_190 + classes: [ghostgjets] + num_global_objects: 21_980_790 #55_036_052 + num_global_objects_test: 476_190 + num_global_objects_val: 476_190 - region: <<: *highpt sample: <<: *zprime - flavours: [ghosttaujets] - num_jets: 3_871_815 #25_623_679 - num_jets_test: 213_478 - num_jets_val: 213_478 + classes: [ghosttaujets] + num_global_objects: 3_871_815 #25_623_679 + num_global_objects_test: 213_478 + num_global_objects_val: 213_478 resampling: target: ghostcjets @@ -153,11 +153,11 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 50_000_000 - num_jets_estimate_plotting: 10_000_000 - num_jets_per_output_file: 5_000_000 + num_global_objects_estimate: 50_000_000 + num_global_objects_estimate_plotting: 10_000_000 + num_global_objects_per_output_file: 5_000_000 base_dir: /srv/beegfs/scratch/groups/rodem/DAOD_FTAG/ out_dir: p7085/upp/ ntuple_dir: p7085/ diff --git a/upp/configs/GN3V00/dr.yaml b/upp/configs/GN3V00/dr.yaml index a1ec2836..2ceb1d21 100644 --- a/upp/configs/GN3V00/dr.yaml +++ b/upp/configs/GN3V00/dr.yaml @@ -3,14 +3,14 @@ global_cuts: !include splits/simple-split.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 zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.svanstro.800030.e7954_s3681_r13144_p6368.tdd.GN3_dev.25_2_27.24-09-17_v00_output.h5/*.h5" # mc20d - "user.svanstro.800030.e7954_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-09-17_v00_output.h5/*.h5" # mc20d @@ -32,33 +32,33 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets, ujets] - num_jets: 6_000_000 - num_jets_test: 2_000_000 + classes: [bjets, cjets, ujets] + num_global_objects: 6_000_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [taujets] - num_jets: 2_000_000 - num_jets_test: 500_000 + classes: [taujets] + num_global_objects: 2_000_000 + num_global_objects_test: 500_000 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets, ujets] - num_jets: 3_000_000 - num_jets_test: 500_000 + classes: [bjets, cjets, ujets] + num_global_objects: 3_000_000 + num_global_objects_test: 500_000 - region: <<: *highpt sample: <<: *zprime - flavours: [taujets] - num_jets: 1_000_000 - num_jets_test: 200_000 + classes: [taujets] + num_global_objects: 1_000_000 + num_global_objects_test: 200_000 resampling: target: cjets @@ -72,8 +72,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 2_000_000 + num_global_objects_estimate: 2_000_000 base_dir: /unix/atlastracking/samples/gn3v00/dr/ ntuple_dir: /unix/atlastracking/samples/gn3v00/ntuples/ diff --git a/upp/configs/GN3V00/ghost-highstat.yaml b/upp/configs/GN3V00/ghost-highstat.yaml index d630d018..d39e0c28 100644 --- a/upp/configs/GN3V00/ghost-highstat.yaml +++ b/upp/configs/GN3V00/ghost-highstat.yaml @@ -3,7 +3,7 @@ global_cuts: !include splits/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.wlai.601589.e8547_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d - "user.wlai.601589.e8549_s4159_r14799_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc23a @@ -11,7 +11,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.wlai.800030.e7954_s3681_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d - "user.wlai.800030.e7954_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d @@ -33,97 +33,97 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplitbjets] - num_jets: 79_000_000 - num_jets_test: 2_000_000 + classes: [ghostsplitbjets] + num_global_objects: 79_000_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplitcjets] - num_jets: 26_500_000 - num_jets_test: 2_000_000 + classes: [ghostsplitcjets] + num_global_objects: 26_500_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplitsjets] - num_jets: 21_000_000 - num_jets_test: 1_000_000 + classes: [ghostsplitsjets] + num_global_objects: 21_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplitudjets] - num_jets: 54_000_000 - num_jets_test: 1_000_000 + classes: [ghostsplitudjets] + num_global_objects: 54_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplitgjets] - num_jets: 46_000_000 - num_jets_test: 1_000_000 + classes: [ghostsplitgjets] + num_global_objects: 46_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsplittaujets] - num_jets: 9_000_000 - num_jets_test: 500_000 + classes: [ghostsplittaujets] + num_global_objects: 9_000_000 + num_global_objects_test: 500_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplitbjets] - num_jets: 39_500_000 - num_jets_test: 2_000_000 + classes: [ghostsplitbjets] + num_global_objects: 39_500_000 + num_global_objects_test: 2_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplitcjets] - num_jets: 13_250_000 - num_jets_test: 2_000_000 + classes: [ghostsplitcjets] + num_global_objects: 13_250_000 + num_global_objects_test: 2_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplitsjets] - num_jets: 10_500_000 - num_jets_test: 1_000_000 + classes: [ghostsplitsjets] + num_global_objects: 10_500_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplitudjets] - num_jets: 27_000_000 - num_jets_test: 1_000_000 + classes: [ghostsplitudjets] + num_global_objects: 27_000_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplitgjets] - num_jets: 23_000_000 - num_jets_test: 1_000_000 + classes: [ghostsplitgjets] + num_global_objects: 23_000_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsplittaujets] - num_jets: 4_500_000 - num_jets_test: 200_000 + classes: [ghostsplittaujets] + num_global_objects: 4_500_000 + num_global_objects_test: 200_000 resampling: target: ghostsplitcjets @@ -137,8 +137,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: /unix/atlas2/weilai/datasets/atlas/upp_outs/ghost_high_stat_full ntuple_dir: /unix/atlas2/weilai/datasets/atlas/ntuples/gn3v00 diff --git a/upp/configs/GN3V00/ghost.yaml b/upp/configs/GN3V00/ghost.yaml index e7583a16..04cd5f36 100644 --- a/upp/configs/GN3V00/ghost.yaml +++ b/upp/configs/GN3V00/ghost.yaml @@ -3,7 +3,7 @@ global_cuts: !include splits/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.wlai.601589.e8547_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d - "user.wlai.601589.e8549_s4159_r14799_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc23a @@ -11,7 +11,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.wlai.800030.e7954_s3681_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d - "user.wlai.800030.e7954_s3797_r13144_p6368.tdd.GN3_dev.25_2_27.24-11-10_GN3v00_output.h5/*.h5" # mc20d @@ -33,33 +33,33 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [ghostbjets, ghostcjets, ghostujets] - num_jets: 6_000_000 - num_jets_test: 2_000_000 + classes: [ghostbjets, ghostcjets, ghostujets] + num_global_objects: 6_000_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghosttaujets] - num_jets: 2_000_000 - num_jets_test: 500_000 + classes: [ghosttaujets] + num_global_objects: 2_000_000 + num_global_objects_test: 500_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostbjets, ghostcjets, ghostujets] - num_jets: 3_000_000 - num_jets_test: 2_000_000 + classes: [ghostbjets, ghostcjets, ghostujets] + num_global_objects: 3_000_000 + num_global_objects_test: 2_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghosttaujets] - num_jets: 1_000_000 - num_jets_test: 200_000 + classes: [ghosttaujets] + num_global_objects: 1_000_000 + num_global_objects_test: 200_000 resampling: target: ghostcjets @@ -73,8 +73,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 2_000_000 + num_global_objects_estimate: 2_000_000 base_dir: /unix/atlas2/weilai/datasets/atlas/upp_outs/ghost_correct ntuple_dir: /unix/atlas2/weilai/datasets/atlas/ntuples/gn3v00 diff --git a/upp/configs/GN3V01/GN3V01-RW.yaml b/upp/configs/GN3V01/GN3V01-RW.yaml index 216a6774..0be49b3f 100644 --- a/upp/configs/GN3V01/GN3V01-RW.yaml +++ b/upp/configs/GN3V01/GN3V01-RW.yaml @@ -3,7 +3,7 @@ global_cuts: !include GN3V01/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.alfroch.601589.e8547_s3797_r13144_p6859.tdd.GN3_dev.25_2_56.250729_Central_Dump_output.h5" - "user.alfroch.601589.e8549_s4162_r14622_p6859.tdd.GN3_dev.25_2_56.250729_Central_Dump_output.h5" @@ -12,7 +12,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.alfroch.800030.e7954_s3797_r13144_p6859.tdd.GN3_dev.25_2_56.250729_Central_Dump_output.h5" - "user.alfroch.800030.e8514_s4162_r14622_p6859.tdd.GN3_dev.25_2_56.250729_Central_Dump_output.h5" @@ -37,7 +37,7 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: + classes: [ ghostbjets, ghostcjets, @@ -46,13 +46,13 @@ components: ghostgjets, ghosttaujets, ] - num_jets: -1 + num_global_objects: -1 - region: <<: *highpt sample: <<: *zprime - flavours: + classes: [ ghostbjets, ghostcjets, @@ -61,19 +61,19 @@ components: ghostgjets, ghosttaujets, ] - num_jets: -1 + num_global_objects: -1 # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 - num_jets_per_output_file: 25_000_000 + num_global_objects_estimate: 25_000_000 + num_global_objects_per_output_file: 25_000_000 base_dir: /home/xzcappon/phd/datasets/flavour_tagging/gn3/high-stats/ ntuple_dir: /home/xzcappon/phd/datasets/flavour_tagging/gn3/high-stats/ntuples reweighting: - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 merge_num_proc: 20 reweights: - group: jets diff --git a/upp/configs/GN3V01/GN3V01.yaml b/upp/configs/GN3V01/GN3V01.yaml index eccd90f2..f14dd9d7 100644 --- a/upp/configs/GN3V01/GN3V01.yaml +++ b/upp/configs/GN3V01/GN3V01.yaml @@ -3,7 +3,7 @@ global_cuts: !include GN3V01/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.alfroch.601589.e8547_s3797_r13144_p6698.tdd.GN3_dev.25_2_48.GN3V01-Training_output.h5/*.h5" # MC20d - "user.alfroch.601589.e8549_s4162_r14622_p6698.tdd.GN3_dev.25_2_48.GN3V01-Training_output.h5/*.h5" # MC23a @@ -11,7 +11,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.alfroch.800030.e7954_s3797_r13144_p6698.tdd.GN3_dev.25_2_48.GN3V01-Training_output.h5/*.h5" # MC20d - "user.alfroch.800030.e8514_s4162_r14622_p6698.tdd.GN3_dev.25_2_48.GN3V01-Training_output.h5/*.h5" # MC23a @@ -33,97 +33,97 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [ghostbjets] - num_jets: 100_000_000 - num_jets_test: 2_000_000 + classes: [ghostbjets] + num_global_objects: 100_000_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostcjets] - num_jets: 27_500_000 - num_jets_test: 2_000_000 + classes: [ghostcjets] + num_global_objects: 27_500_000 + num_global_objects_test: 2_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostsjets] - num_jets: 20_000_000 - num_jets_test: 1_000_000 + classes: [ghostsjets] + num_global_objects: 20_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostudjets] - num_jets: 55_000_000 - num_jets_test: 1_000_000 + classes: [ghostudjets] + num_global_objects: 55_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghostgjets] - num_jets: 65_000_000 - num_jets_test: 1_000_000 + classes: [ghostgjets] + num_global_objects: 65_000_000 + num_global_objects_test: 1_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ghosttaujets] - num_jets: 11_000_000 - num_jets_test: 500_000 + classes: [ghosttaujets] + num_global_objects: 11_000_000 + num_global_objects_test: 500_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostbjets] - num_jets: 50_000_000 - num_jets_test: 2_000_000 + classes: [ghostbjets] + num_global_objects: 50_000_000 + num_global_objects_test: 2_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostcjets] - num_jets: 13_750_000 - num_jets_test: 2_000_000 + classes: [ghostcjets] + num_global_objects: 13_750_000 + num_global_objects_test: 2_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostsjets] - num_jets: 10_000_000 - num_jets_test: 1_000_000 + classes: [ghostsjets] + num_global_objects: 10_000_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostudjets] - num_jets: 27_500_000 - num_jets_test: 1_000_000 + classes: [ghostudjets] + num_global_objects: 27_500_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghostgjets] - num_jets: 32_500_000 - num_jets_test: 1_000_000 + classes: [ghostgjets] + num_global_objects: 32_500_000 + num_global_objects_test: 1_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ghosttaujets] - num_jets: 5_500_000 - num_jets_test: 200_000 + classes: [ghosttaujets] + num_global_objects: 5_500_000 + num_global_objects_test: 200_000 resampling: target: ghostcjets @@ -137,8 +137,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: GN3V01_Training_preprocessed/ ntuple_dir: GN3V01_Training_h5/ diff --git a/upp/configs/extended_labels.yaml b/upp/configs/extended_labels.yaml index f97a8faa..0c717b8f 100644 --- a/upp/configs/extended_labels.yaml +++ b/upp/configs/extended_labels.yaml @@ -6,7 +6,7 @@ global_cuts: !include splits/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.svanstro.410470.e6337_s3681_r13144_p5922.tdd.EMPFlow_kfold.24_2_27.23-11-10_kfoldtrain_output.h5/*.h5" @@ -21,36 +21,36 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets] - num_jets: 25_000_000 + classes: [bjets] + num_global_objects: 25_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [D0meson] - num_jets: 12_500_000 + classes: [D0meson] + num_global_objects: 12_500_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [nonD0meson] - num_jets: 12_500_000 + classes: [nonD0meson] + num_global_objects: 12_500_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 50_000_000 + classes: [ujets] + num_global_objects: 50_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [taujets] - num_jets: 4_000_000 + classes: [taujets] + num_global_objects: 4_000_000 resampling: target: bjets @@ -64,8 +64,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: nominal_Loose + global_name: nominal_Loose batch_size: 1_000_000 - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 base_dir: /nfs/dust/atlas/user/nkumari/UPP_latest/umami-preprocessing/upp/configs/prep ntuple_dir: /nfs/dust/atlas/user/nkumari/FTAG_tdd diff --git a/upp/configs/open-dataset.yaml b/upp/configs/open-dataset.yaml index 3198004b..830ca364 100644 --- a/upp/configs/open-dataset.yaml +++ b/upp/configs/open-dataset.yaml @@ -47,7 +47,7 @@ global_cuts: ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.wlai.601589.e8549_s4159_r15530_p6698.tdd.OpenDataset.25_2_46.25-04-08_open-rc3_output.h5/*.h5" - "user.wlai.601589.e8549_s4162_r14622_p6698.tdd.OpenDataset.25_2_46.25-04-08_open-rc3_output.h5/*.h5" @@ -63,22 +63,22 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets] - num_jets: 13_000_000 + classes: [bjets, cjets] + num_global_objects: 13_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 26_000_000 + classes: [ujets] + num_global_objects: 26_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [taujets] - num_jets: 1_500_000 + classes: [taujets] + num_global_objects: 1_500_000 resampling: target: cjets @@ -92,8 +92,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: /unix/atlas2/weilai/datasets/atlas/upp_outs/opendata_rc3 ntuple_dir: /unix/atlas2/weilai/datasets/atlas/ntuples/opendata/rc3 diff --git a/upp/configs/plit_electron.yaml b/upp/configs/plit_electron.yaml index 828ebf85..afe84110 100644 --- a/upp/configs/plit_electron.yaml +++ b/upp/configs/plit_electron.yaml @@ -3,7 +3,7 @@ variables: !include plit_electron_variables.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.pgadow.LD_2023_11_28.601589.PhPy8EG_A14_ttbar_hdamp258p75_nonallhadron.e8547_s3797_r13144_p5934_TREE/*.h5" - "user.pgadow.LD_2023_11_28.601589.PhPy8EG_A14_ttbar_hdamp258p75_nonallhadron.e8549_s4159_r14799_p5934_TREE/*.h5" @@ -27,15 +27,15 @@ components: <<: *electron sample: <<: *ttbar - flavours: [elxprompt] - num_jets: 25_000_000 + classes: [elxprompt] + num_global_objects: 25_000_000 - region: <<: *electron sample: <<: *ttbar - flavours: [npxall] - num_jets: 13_000_000 + classes: [npxall] + num_global_objects: 13_000_000 resampling: @@ -50,9 +50,9 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: electrons + global_name: electrons batch_size: 1_000_000 - num_jets_estimate: 5_000 + num_global_objects_estimate: 5_000 base_dir: /nfs/dust/atlas/user/pgadow/plit/data/preprocessed/electrons_38M ntuple_dir: /nfs/dust/atlas/user/pgadow/plit/data/ntuples out_dir: /nfs/dust/atlas/user/pgadow/plit/data/preprocessed/electrons_38M diff --git a/upp/configs/plit_muon.yaml b/upp/configs/plit_muon.yaml index 5427c221..663c51f9 100644 --- a/upp/configs/plit_muon.yaml +++ b/upp/configs/plit_muon.yaml @@ -3,7 +3,7 @@ variables: !include plit_muon_variables.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.pgadow.LD_2023_11_28.601589.PhPy8EG_A14_ttbar_hdamp258p75_nonallhadron.e8547_s3797_r13144_p5934_TREE/*.h5" - "user.pgadow.LD_2023_11_28.601589.PhPy8EG_A14_ttbar_hdamp258p75_nonallhadron.e8549_s4159_r14799_p5934_TREE/*.h5" @@ -27,15 +27,15 @@ components: <<: *muon sample: <<: *ttbar - flavours: [muxprompt] - num_jets: 30_000_000 + classes: [muxprompt] + num_global_objects: 30_000_000 - region: <<: *muon sample: <<: *ttbar - flavours: [npxall] - num_jets: 10_000_000 + classes: [npxall] + num_global_objects: 10_000_000 resampling: @@ -50,9 +50,9 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: muons + global_name: muons batch_size: 1_000_000 - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 base_dir: /nfs/dust/atlas/user/pgadow/plit/data/preprocessed/muons_40M ntuple_dir: /nfs/dust/atlas/user/pgadow/plit/data/ntuples out_dir: /nfs/dust/atlas/user/pgadow/plit/data/preprocessed/muons_40M diff --git a/upp/configs/single-b-upgrade.yaml b/upp/configs/single-b-upgrade.yaml index 48648905..533a8cf4 100644 --- a/upp/configs/single-b-upgrade.yaml +++ b/upp/configs/single-b-upgrade.yaml @@ -6,7 +6,7 @@ global_cuts: !include splits/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "/atlas_cloud/triglion/data/TDD_ttbar_DiLep/user.tstreble.601230.e8481_s4446_r16176_p6677.tdd.upgrade.25_0_29.25-05-06_Run4_May6_output.h5/user.tstreble.44562989._*.output.h5" - "/atlas_cloud/triglion/data/TDD_ttbar_SingleLep/user.tstreble.601229.e8481_s4446_r16176_p6677.tdd.upgrade.25_0_29.25-05-06_Run4_May6_output.h5/user.tstreble.44562988._*.output.h5" @@ -14,7 +14,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "/atlas_cloud/triglion/data/TDD_Zprime/user.tstreble.800030.e8481_s4446_r16176_p6677.tdd.upgrade.25_0_29.25-05-06_Run4_May6_output.h5/user.tstreble.44562987._*.output.h5" @@ -36,43 +36,43 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets] - num_jets: 14_500_000 + classes: [bjets, cjets] + num_global_objects: 14_500_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 29_000_000 + classes: [ujets] + num_global_objects: 29_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [taujets] - num_jets: 2_013_889 + classes: [taujets] + num_global_objects: 2_013_889 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets] - num_jets: 5_800_000 + classes: [bjets, cjets] + num_global_objects: 5_800_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ujets] - num_jets: 11_600_000 + classes: [ujets] + num_global_objects: 11_600_000 - region: <<: *highpt sample: <<: *zprime - flavours: [taujets] - num_jets: 805_555 + classes: [taujets] + num_global_objects: 805_555 @@ -88,8 +88,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: /atlas_cloud/triglion/preprocessing ntuple_dir: /atlas_cloud/triglion/preprocessing diff --git a/upp/configs/single-b.yaml b/upp/configs/single-b.yaml index fa216f79..19f715bf 100644 --- a/upp/configs/single-b.yaml +++ b/upp/configs/single-b.yaml @@ -6,7 +6,7 @@ global_cuts: !include splits/simple-split.yaml ttbar: &ttbar name: ttbar - equal_jets: False + equal_global_objects: False pattern: - "user.svanstro.410470.e6337_s3681_r13144_p5922.tdd.EMPFlow_kfold.24_2_27.23-11-10_kfoldtrain_output.h5/*.h5" - "user.svanstro.601229.e8514_s4162_r14622_p5922.tdd.EMPFlow_kfold.24_2_27.23-11-10_kfoldtrain_output.h5/*.h5" @@ -14,7 +14,7 @@ ttbar: &ttbar zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: - "user.svanstro.800030.e7954_s3681_r13144_p5922.tdd.EMPFlow_kfold.24_2_27.23-11-10_kfoldtrain_output.h5/*.h5" - "user.svanstro.800030.e7954_s3797_r13144_p5922.tdd.EMPFlow_kfold.24_2_27.23-11-19_kfoldtrain_output.h5/*.h5" @@ -37,43 +37,43 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets] - num_jets: 45_000_000 + classes: [bjets, cjets] + num_global_objects: 45_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [ujets] - num_jets: 90_000_000 + classes: [ujets] + num_global_objects: 90_000_000 - region: <<: *lowpt sample: <<: *ttbar - flavours: [taujets] - num_jets: 6_250_000 + classes: [taujets] + num_global_objects: 6_250_000 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets] - num_jets: 18_000_000 + classes: [bjets, cjets] + num_global_objects: 18_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [ujets] - num_jets: 36_000_000 + classes: [ujets] + num_global_objects: 36_000_000 - region: <<: *highpt sample: <<: *zprime - flavours: [taujets] - num_jets: 2_500_000 + classes: [taujets] + num_global_objects: 2_500_000 resampling: @@ -88,8 +88,8 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 25_000_000 + num_global_objects_estimate: 25_000_000 base_dir: /home/xzcappon/phd/datasets/combined_run2_run3/p5922/high_stats/fold0 ntuple_dir: /home/xucapsva/data/ftag/h5/gn2_v01/ntuples diff --git a/upp/configs/test.yaml b/upp/configs/test.yaml index bf72dfcb..9653ae08 100644 --- a/upp/configs/test.yaml +++ b/upp/configs/test.yaml @@ -29,15 +29,15 @@ components: <<: *lowpt sample: <<: *ttbar - flavours: [bjets, cjets, ujets] - num_jets: 10_000 + classes: [bjets, cjets, ujets] + num_global_objects: 10_000 - region: <<: *highpt sample: <<: *zprime - flavours: [bjets, cjets, ujets] - num_jets: 10_000 + classes: [bjets, cjets, ujets] + num_global_objects: 10_000 resampling: target: cjets @@ -51,7 +51,7 @@ resampling: global: batch_size: 10_000 - num_jets_estimate: 10_000 + num_global_objects_estimate: 10_000 base_dir: /unix/atlastracking/samples/gnn_training/martino/v1/ out_dir: test_out ntuple_dir: source_data diff --git a/upp/configs/xbb-gn3x.yaml b/upp/configs/xbb-gn3x.yaml index 50fd9ff3..d3a1c54d 100644 --- a/upp/configs/xbb-gn3x.yaml +++ b/upp/configs/xbb-gn3x.yaml @@ -3,35 +3,35 @@ global_cuts: !include splits/simple-split.yaml htautauhad: &htautauhad name: htautauhad - equal_jets: False + equal_global_objects: False pattern: - user.edcritel.802168.e8558_s3797_r13144_p6453.tdd.*/*.h5 - user.edcritel.802168.e8558_s4159_r15530_p6453.tdd.*/*.h5 hbb: &hbb name: hbb - equal_jets: False + equal_global_objects: False pattern: - group.perf-flavtag.801471.e8441_s3681_r13144_p6453.tdd.*/*.h5 - group.perf-flavtag.801972.e8514_s4159_r15224_p6453.tdd.*/*.h5 hcc: &hcc name: hcc - equal_jets: False + equal_global_objects: False pattern: - group.perf-flavtag.801472.e8441_s3681_r13144_p6453.tdd.*/*.h5 - group.perf-flavtag.801973.e8514_s4159_r15224_p6453.tdd.*/*.h5 top: &Zprime name: Zprime - equal_jets: False + equal_global_objects: False pattern: - group.perf-flavtag.426345.e6880_s3681_r13144_p6453.tdd.*/*.h5 - group.perf-flavtag.802423.e8514_s4159_r15530_p6453.tdd.*/*.h5 qcd: &qcd name: qcd - equal_jets: False + equal_global_objects: False pattern: - group.perf-flavtag.364703.e7142_s3681_r13144_p6453.tdd.*/*.h5 - group.perf-flavtag.364704.e7142_s3681_r13144_p6453.tdd.*/*.h5 @@ -46,7 +46,7 @@ qcd: &qcd Wqq: &Wqq name: Wqq - equal_jets: False + equal_global_objects: False pattern: - group.perf-flavtag.802017.e8482_s3797_r13144_p6453.tdd.*/*.h5 - group.perf-flavtag.802017.e8557_s4159_r15224_p6453.tdd.*/*.h5 @@ -61,43 +61,43 @@ components: <<: *inclusive sample: <<: *htautauhad - flavours: [htautauhad] - num_jets: 10_000_000 + classes: [htautauhad] + num_global_objects: 10_000_000 - region: <<: *inclusive sample: <<: *hbb - flavours: [hbb] - num_jets: 40_000_000 + classes: [hbb] + num_global_objects: 40_000_000 - region: <<: *inclusive sample: <<: *hcc - flavours: [hcc] - num_jets: 40_000_000 + classes: [hcc] + num_global_objects: 40_000_000 - region: <<: *inclusive sample: <<: *Zprime - flavours: [top] - num_jets: 35_000_000 + classes: [top] + num_global_objects: 35_000_000 - region: <<: *inclusive sample: <<: *qcd - flavours: [qcd] - num_jets: 80_000_000 + classes: [qcd] + num_global_objects: 80_000_000 - region: <<: *inclusive sample: <<: *Wqq - flavours: [Wqq] - num_jets: 5_000_000 + classes: [Wqq] + num_global_objects: 5_000_000 resampling: @@ -117,7 +117,7 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: batch_size: 4_000_000 - num_jets_estimate: 10_000_000 + num_global_objects_estimate: 10_000_000 base_dir: /share/lustre/ecritelli/xbb_upp_Zprimeonly_withRegr/base_dir ntuple_dir: /share/lustre/ecritelli/ntuples/gn3xv00 out_dir: /share/lustre/ecritelli/xbb_upp_Zprimeonly_withRegr/output diff --git a/upp/configs/xbb-rw.yaml b/upp/configs/xbb-rw.yaml index 8faa56fa..73843ae8 100644 --- a/upp/configs/xbb-rw.yaml +++ b/upp/configs/xbb-rw.yaml @@ -4,7 +4,7 @@ global_cuts: !include splits/simple-split.yaml hbb: &hbb name: hbb - equal_jets: False + equal_global_objects: False pattern: # mc20 - "user.jabarr.801471.e8441_s3681_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -13,7 +13,7 @@ hbb: &hbb hcc: &hcc name: hcc - equal_jets: False + equal_global_objects: False pattern: # mc20 - "user.jabarr.801472.e8441_s3681_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -22,7 +22,7 @@ hcc: &hcc htautauhad: &htautauhad name: htautauhad - equal_jets: False + equal_global_objects: False pattern: # mc20 - "user.jabarr.802168.e8558_s3797_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -31,7 +31,7 @@ htautauhad: &htautauhad zprime: &zprime name: zprime - equal_jets: False + equal_global_objects: False pattern: # mc20 - "user.jabarr.426345.e6880_s3681_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -40,7 +40,7 @@ zprime: &zprime qcd: &qcd name: qcd - equal_jets: False + equal_global_objects: False pattern: # mc20 dijet - "user.jabarr.364703.e7142_s3681_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -57,7 +57,7 @@ qcd: &qcd wqq: &wqq name: wqq - equal_jets: False + equal_global_objects: False pattern: # mc20 - "user.jabarr.802017.e8482_s3797_r13144_p6453.tdd.FatJets.25_2_76.26-01-05_cb_output.h5" @@ -73,56 +73,56 @@ components: <<: *inclusive sample: <<: *hbb - flavours: [hbb] - num_jets: -1 + classes: [hbb] + num_global_objects: -1 - region: <<: *inclusive sample: <<: *hcc - flavours: [hcc] - num_jets: -1 + classes: [hcc] + num_global_objects: -1 - region: <<: *inclusive sample: <<: *htautauhad - flavours: [htautauhad] - num_jets: -1 + classes: [htautauhad] + num_global_objects: -1 - region: <<: *inclusive sample: <<: *zprime - flavours: [top] - num_jets: -1 + classes: [top] + num_global_objects: -1 - region: <<: *inclusive sample: <<: *qcd - flavours: [qcd] - num_jets: -1 + classes: [qcd] + num_global_objects: -1 - region: <<: *inclusive sample: <<: *wqq - flavours: [Wqq] - num_jets: -1 + classes: [Wqq] + num_global_objects: -1 global: - jets_name: jets + global_name: jets batch_size: 1_000_000 - num_jets_estimate: 10_000_000 - num_jets_per_output_file: 25_000_000 + num_global_objects_estimate: 10_000_000 + num_global_objects_per_output_file: 25_000_000 base_dir: /share/lustre/jbarr/training-files/2026-01-01-cb/ merge_test_samples: False reweighting: - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 merge_num_proc: 10 reweights: # 4D reweighting over pt, eta, mass, and mcCampaignYear diff --git a/upp/configs/xbb.yaml b/upp/configs/xbb.yaml index 866e41d7..0555a164 100644 --- a/upp/configs/xbb.yaml +++ b/upp/configs/xbb.yaml @@ -9,13 +9,13 @@ hcc: &hcc pattern: user.svanstro.801472.e8441_e7400_s3681_r13144_r13146_p5488.tdd.*/*.h5 top: &top name: top - equal_jets: true + equal_global_objects: true pattern: - user.svanstro.426345.e6880_s3681_r13144_p5488.tdd.*/*.h5 - user.svanstro.426345.e6880_s3681_r13145_p5488.tdd.*/*.h5 qcd: &qcd name: qcd - equal_jets: true + equal_global_objects: true pattern: - user.svanstro.364703.e7142_s3681_r13144_p5488.tdd.*/*.h5 - user.svanstro.364704.e7142_s3681_r13144_p5488.tdd.*/*.h5 @@ -30,29 +30,29 @@ components: <<: *inclusive sample: <<: *hbb - flavours: [hbb] - num_jets: 30_000_000 + classes: [hbb] + num_global_objects: 30_000_000 - region: <<: *inclusive sample: <<: *hcc - flavours: [hcc] - num_jets: 30_000_000 + classes: [hcc] + num_global_objects: 30_000_000 - region: <<: *inclusive sample: <<: *top - flavours: [top] - num_jets: 30_000_000 + classes: [top] + num_global_objects: 30_000_000 - region: <<: *inclusive sample: <<: *qcd - flavours: [qcd] - num_jets: 50_000_000 + classes: [qcd] + num_global_objects: 50_000_000 resampling: target: hbb @@ -71,6 +71,6 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: batch_size: 4_000_000 - num_jets_estimate: 10_000_000 + num_global_objects_estimate: 10_000_000 base_dir: /unix/atlastracking/samples/xbb/dumps/p5488_ext/ merge_test_samples: True diff --git a/upp/configs/xtautau.yaml b/upp/configs/xtautau.yaml index 46afab2d..454d3bf2 100644 --- a/upp/configs/xtautau.yaml +++ b/upp/configs/xtautau.yaml @@ -15,7 +15,7 @@ top: &top pattern: /atlas_cloud/fujimoto/umamiMCd/user.mfujimot.802423.e8514_s4159_r15530_p6453.tdd.FatJets.25_2_34.prod_131224_output.h5/*.h5 qcd: &qcd name: qcd - equal_jets: true + equal_global_objects: true pattern: - /atlas_cloud/fujimoto/umamiMCd/user.mfujimot.801168.e8514_s4159_r15224_p6453.tdd.FatJets.25_2_34.prod_131224_output.h5/*.h5 - /atlas_cloud/fujimoto/umamiMCd/user.mfujimot.801169.e8514_s4159_r15224_p6453.tdd.FatJets.25_2_34.prod_131224_output.h5/*.h5 @@ -30,36 +30,36 @@ components: <<: *inclusive sample: <<: *htauhad - flavours: [htauhad] - num_jets: 5_000_000 + classes: [htauhad] + num_global_objects: 5_000_000 - region: <<: *inclusive sample: <<: *hbb - flavours: [hbb] - num_jets: 14_500_000 + classes: [hbb] + num_global_objects: 14_500_000 - region: <<: *inclusive sample: <<: *hcc - flavours: [hcc] - num_jets: 14_500_000 + classes: [hcc] + num_global_objects: 14_500_000 - region: <<: *inclusive sample: <<: *top - flavours: [top] - num_jets: 8_000_000 + classes: [top] + num_global_objects: 8_000_000 - region: <<: *inclusive sample: <<: *qcd - flavours: [qcd] - num_jets: 22_000_000 + classes: [qcd] + num_global_objects: 22_000_000 resampling: target: hbb @@ -76,6 +76,6 @@ resampling: # note: sensible defaults are defined in the PreprocessingConfig constructor global: batch_size: 4_000_000 - num_jets_estimate: 5_000_000 + num_global_objects_estimate: 5_000_000 base_dir: /atlas_cloud/fujimoto/samplesMC23d/ merge_test_samples: True diff --git a/upp/grid/download_and_prepare.py b/upp/grid/download_and_prepare.py index 9abe1b33..9b116e51 100644 --- a/upp/grid/download_and_prepare.py +++ b/upp/grid/download_and_prepare.py @@ -129,12 +129,13 @@ def create_meta_data( output_file = output_dir / "organised-components.yaml" - # Create a reader for each components to get num jets - num_jets = { + # Create a reader for each components to get num objects + num_global_objects = { split: { flavour: H5Reader( files_by_component[split][flavour], - ).num_jets + global_objects_name=pp_config.global_name, + ).num_global_objects for flavour in files_by_component[split] } for split in files_by_component @@ -144,7 +145,7 @@ def create_meta_data( yaml.dump( { "files": files_by_component, - "num_jets": num_jets, + "num_global_objects": num_global_objects, }, f, default_flow_style=False, diff --git a/upp/main.py b/upp/main.py index 1dd23b99..98f45f5e 100644 --- a/upp/main.py +++ b/upp/main.py @@ -1,13 +1,14 @@ """ -Preprocessing pipeline for jet tagging. +Preprocessing pipeline for object tagging. By default all stages for the training split are run. To run with only specific stages enabled, include the flag for the required stages. To run without certain stages, include the corresponding negative flag. -To disable resampling, omit the `resampling` block or set `method: none`. The jets passing -the cuts are then written directly, capped at each component's `num_jets` (use `num_jets: -1` -to keep all of them). The `--no-resample` flag only skips the resampling *stage* (e.g. to +To disable resampling, omit the `resampling` block or set `method: none`. The objects passing +the cuts are then written directly, capped at each component's `num_global_objects` +(use `num_global_objects: -1` to keep all of them). The `--no-resample` flag only skips +the resampling *stage* (e.g. to re-run later stages); it does not disable resampling. """ diff --git a/upp/stages/__init__.py b/upp/stages/__init__.py index 84a76196..dcf1fa67 100644 --- a/upp/stages/__init__.py +++ b/upp/stages/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations -from upp.stages.hist import Hist, bin_jets, create_histograms +from upp.stages.hist import Hist, bin_global_objects, create_histograms from upp.stages.interpolation import subdivide_bins, upscale_array, upscale_array_regionally from upp.stages.merging import Merging from upp.stages.normalisation import Normalisation @@ -14,7 +14,7 @@ "Merging", "Normalisation", "Resampling", - "bin_jets", + "bin_global_objects", "create_histograms", "make_hist", "plot_resampling_dists", diff --git a/upp/stages/hist.py b/upp/stages/hist.py index 6b37c02b..7ce8e9fc 100644 --- a/upp/stages/hist.py +++ b/upp/stages/hist.py @@ -18,13 +18,13 @@ from upp.classes.preprocessing_config import PreprocessingConfig -def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: +def bin_global_objects(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: """Create the histogram and bins for the given resampling variables. Parameters ---------- array : dict - Dict with the loaded jets and the resampling + Dict with the loaded objects and the resampling variables. bins : list Flat list with the bins which are to be used. @@ -61,21 +61,24 @@ class Hist: def write_hist( self, - jets: dict, + global_objects: dict, resampling_vars: list, bins: list, + global_name: str = "jets", ) -> None: """ Write the histogram to file. Parameters ---------- - jets : dict - Dict with the loaded jets. + global_objects : dict + Dict with the loaded objects. resampling_vars : list List of the resampling variables. bins : list Flat list with the bins. + global_name : str, optional + Name of the global object, used for the count attribute, by default "jets". Raises ------ @@ -85,16 +88,16 @@ def write_hist( # make parent dir self.path.parent.mkdir(parents=True, exist_ok=True) - # bin jets - hist = bin_jets(jets[resampling_vars], bins)[0] - pbin = hist / len(jets) # probability (rate) of each bin + # bin objects + hist = bin_global_objects(global_objects[resampling_vars], bins)[0] + pbin = hist / len(global_objects) # probability (rate) of each bin if not math.isclose(pbin.sum(), 1, rel_tol=1e-4, abs_tol=1e-4): raise ValueError(f"{pbin.sum()} != 1, check cuts and binning") with h5py.File(self.path, "w") as f: f.create_dataset("pbin", data=pbin) f.create_dataset("hist", data=hist) - f.attrs.create("num_jets", len(jets)) + f.attrs.create(f"num_{global_name}", len(global_objects)) f.attrs.create("resampling_vars", resampling_vars) for i, v in enumerate(resampling_vars): f.attrs.create(f"bins_{v}", bins[i]) @@ -153,7 +156,9 @@ def create_histograms( title = " Writing PDFs " log.info(f"[bold green]{title:-^100}") - log.info(f"[bold green]Estimating PDFs using {config.num_jets_estimate_hist:,} jets...") + log.info( + f"[bold green]Estimating PDFs using {config.num_global_objects_estimate_hist:,} objects..." + ) # Create check variable to ensure at least one component was processed component_processed = not component_to_run @@ -164,32 +169,36 @@ def create_histograms( if isinstance(component_to_run, str) and component_to_run != component.name: continue - log.info(f"Estimating {component} PDF using {config.num_jets_estimate_hist:,} samples...") - component.setup_reader(batch_size=config.batch_size, jets_name=config.jets_name) + log.info( + f"Estimating {component} PDF using " + f"{config.num_global_objects_estimate_hist:,} samples..." + ) + component.setup_reader(batch_size=config.batch_size, global_name=config.global_name) cuts_no_split = component.cuts.ignore(["eventNumber"]) ### - # TODO: return the number of jets here and pass to the next function to get started + # TODO: return the number of objects here and pass to the next function to get started ### - component.check_num_jets( - config.num_jets_estimate_hist, + component.check_num_global_objects( + config.num_global_objects_estimate_hist, cuts=cuts_no_split, silent=False, raise_error=False, ) - # Load the jets from file used for resampling - jets = component.get_jets( + # Load the objects from file used for resampling + global_objects = component.get_global_objects( variables=sampl_vars, - num_jets=config.num_jets_estimate_hist, + num_global_objects=config.num_global_objects_estimate_hist, cuts=cuts_no_split, ) # Write out the hist used for resampling component.hist.write_hist( - jets=jets, + global_objects=global_objects, resampling_vars=sampl_vars, bins=config.sampl_cfg.flat_bins, + global_name=config.global_name, ) # Set the check variable to true diff --git a/upp/stages/merging.py b/upp/stages/merging.py index a87a7948..c5ac0433 100644 --- a/upp/stages/merging.py +++ b/upp/stages/merging.py @@ -26,10 +26,10 @@ def __init__(self, config: PreprocessingConfig): self.components = config.components self.variables = config.variables self.batch_size = config.batch_size - self.jets_name = config.jets_name + self.global_name = config.global_name self.rng = np.random.default_rng(42) self.flavours = self.components.flavours - self.num_jets_per_output_file = config.num_jets_per_output_file + self.num_global_objects_per_output_file = config.num_global_objects_per_output_file self.file_tag = "split" # Auto-resume toggle (make configurable if you prefer opt-in) @@ -48,10 +48,10 @@ def __init__(self, config: PreprocessingConfig): self.dtypes: dict[str, np.dtype] = {} self.base_shapes: dict[str, tuple[int, ...]] = {} - # Setup all the jet counters + # Setup all the object counters self._file_idx: int = 0 - self.total_jets: int = 0 - self.jets_written: int = 0 + self.total_global_objects: int = 0 + self.global_objects_written: int = 0 # Setup the sample string self._sample: str | None = None @@ -60,30 +60,32 @@ def __init__(self, config: PreprocessingConfig): self.current_components = cast("Components", None) self.writer = cast(H5Writer, None) - def add_jet_flavour_label(self, jets: np.ndarray, component: Component) -> np.ndarray: - """Add the jet flavour label to the jets. + def add_global_object_label( + self, global_objects: np.ndarray, component: Component + ) -> np.ndarray: + """Add the object flavour label to the objects. - If already present, jets will be returned without any changes. + If already present, objects will be returned without any changes. Parameters ---------- - jets : np.ndarray - Structured array of with the jets and their variables + global_objects : np.ndarray + Structured array of with the objects and their variables component : Component Component instance of the Returns ------- np.ndarray - Structured array of the jets and their variables with the + Structured array of the objects and their variables with the "flavour_label" added. """ - if "flavour_label" in jets.dtype.names: - return jets + if "flavour_label" in global_objects.dtype.names: + return global_objects int_label = self.flavours.index(component.flavour) - label_array = np.full(len(jets), int_label, dtype=[("flavour_label", "i4")]) + label_array = np.full(len(global_objects), int_label, dtype=[("flavour_label", "i4")]) - return join_structured_arrays([jets, label_array]) + return join_structured_arrays([global_objects, label_array]) def _part_fname(self, sample: str | None, file_idx: int) -> Path: """Construct the exact output filename for a given part index. @@ -115,7 +117,9 @@ def _part_fname(self, sample: str | None, file_idx: int) -> Path: return fname.parent / self.config.split / fname.name def _expected_rows_for_part(self, part_idx: int) -> int: - """Return the expected number of rows for part `part_idx` given total_jets and split size. + """Return the expected number of rows for part `part_idx`. + + Uses ``total_global_objects`` and the configured split size. Parameters ---------- @@ -128,13 +132,13 @@ def _expected_rows_for_part(self, part_idx: int) -> int: Expected number of rows for the given partial file """ # Assert that the final output file will be splitted - assert self.num_jets_per_output_file is not None + assert self.num_global_objects_per_output_file is not None - # Remaining jets starting at this part - start = part_idx * int(self.num_jets_per_output_file) - remaining = max(0, self.total_jets - start) + # Remaining objects starting at this part + start = part_idx * int(self.num_global_objects_per_output_file) + remaining = max(0, self.total_global_objects - start) - return min(int(self.num_jets_per_output_file), remaining) + return min(int(self.num_global_objects_per_output_file), remaining) def _is_part_valid(self, sample: str | None, part_idx: int) -> bool: """Heuristically validate that a part file is complete and consistent. @@ -166,15 +170,15 @@ def _is_part_valid(self, sample: str | None, part_idx: int) -> bool: # Collect expected dataset names from base_shapes (already computed) expected_names = list(self.base_shapes.keys()) - # Tolerate missing optional groups, but require the jet dataset at least - if self.jets_name not in f: - log.warning(f"Missing dataset '{self.jets_name}' in {fname}") + # Tolerate missing optional groups, but require the object dataset at least + if self.global_name not in f: + log.warning(f"Missing dataset '{self.global_name}' in {fname}") return False - # Determine observed length from anchor (jets_name) or first dataset - anchor = self.jets_name if self.jets_name in f else expected_names[0] + # Determine observed length from anchor (global_name) or first dataset + anchor = self.global_name if self.global_name in f else expected_names[0] if anchor not in f: - # if jets_name wasn't found, try any expected dataset that exists + # if global_name wasn't found, try any expected dataset that exists for nm in expected_names: if nm in f: anchor = nm @@ -193,7 +197,7 @@ def _is_part_valid(self, sample: str | None, part_idx: int) -> bool: return False # Compare with expected rows for this part (if split mode) - if self.num_jets_per_output_file is not None: + if self.num_global_objects_per_output_file is not None: exp_len = self._expected_rows_for_part(part_idx) if obs_len != exp_len: log.warning( @@ -227,7 +231,7 @@ def _detect_and_clean_completed_parts(self, sample: str | None) -> int: The index of the first missing/invalid part. """ # Check that multiple output files should be created - if self.num_jets_per_output_file is None: + if self.num_global_objects_per_output_file is None: return 0 # Define a counter @@ -267,11 +271,12 @@ class _NullWriter: """A minimal writer that discards data while tracking how much would be written.""" def __init__(self, capacity: int): - self.num_jets = capacity + # Mirrors the ftag H5Writer API (assigned to self.writer) + self.num_global_objects = capacity self.num_written = 0 def write(self, batch: dict[str, np.ndarray]) -> None: - """Count the number of jets that would be written. + """Count the number of objects that would be written. Parameters ---------- @@ -283,7 +288,7 @@ def write(self, batch: dict[str, np.ndarray]) -> None: return any_arr = next(iter(batch.values())) k = len(any_arr) - self.num_written = min(self.num_written + k, self.num_jets) + self.num_written = min(self.num_written + k, self.num_global_objects) def add_attr(self, *args, **kwargs): """Skip the attribute addition.""" @@ -296,7 +301,7 @@ def close(self): def _open_writer( self, sample: str | None, - jets_in_file: int, + global_objects_in_file: int, file_idx: int, components: Components, ) -> None: @@ -306,15 +311,15 @@ def _open_writer( ---------- sample : str | None Sample name (``None`` for the "train/val test" merge). - jets_in_file : int + global_objects_in_file : int Capacity of the new file (= leading dimension of every dataset). file_idx : int Running part index (0, 1, 2, …); used only for the filename suffix. components : Components - The `Components` object we are currently merging needed for `jet_counts`, etc. + The `Components` object we are currently merging needed for `global_object_counts`, etc. """ # Construct the filename - if self.num_jets_per_output_file is not None: + if self.num_global_objects_per_output_file is not None: fname = self._part_fname(sample, file_idx) else: fname = Path(self.config.out_fname) @@ -322,7 +327,9 @@ def _open_writer( fname = path_append(fname, sample) # Adjust shapes to the capacity of this file - shapes = {name: (jets_in_file,) + shape[1:] for name, shape in self.base_shapes.items()} + shapes = { + name: (global_objects_in_file,) + shape[1:] for name, shape in self.base_shapes.items() + } # Ensure the output directory exists before opening the writer fname.parent.mkdir(parents=True, exist_ok=True) @@ -332,19 +339,22 @@ def _open_writer( fname, self.dtypes, shapes, - add_flavour_label=self.jets_name, - jets_name=self.jets_name, - num_jets=jets_in_file, + add_flavour_label=self.global_name, + global_objects_name=self.global_name, + num_global_objects=global_objects_in_file, ) # Copy the metadata attributes self.writer.add_attr( "flavour_label", [f.name for f in self.flavours], - self.jets_name, + self.global_name, + ) + self.writer.add_attr(f"unique_{self.global_name}", components.unique_global_objects) + self.writer.add_attr( + f"{self.global_name}_counts", + json.dumps(components.global_object_counts(self.global_name)), ) - self.writer.add_attr("unique_jets", components.unique_jets) - self.writer.add_attr("jet_counts", json.dumps(components.jet_counts)) self.writer.add_attr("dsids", str(components.dsids)) self.writer.add_attr("config", json.dumps(self.config.config)) self.writer.add_attr("upp_hash", self.config.git_hash) @@ -368,7 +378,7 @@ def write_chunk(self, components: Components) -> int: Returns ------- int - The number of jets that were consumed from the components + The number of objects that were consumed from the components (== written to disk). When all components are exhausted the function returns 0 so that the caller can stop its loop. """ @@ -385,8 +395,8 @@ def write_chunk(self, components: Components) -> int: try: # shallow copy because we will add a field batch = copy(next(component.stream)) - batch[self.jets_name] = self.add_jet_flavour_label( - jets=batch[self.jets_name], component=component + batch[self.global_name] = self.add_global_object_label( + global_objects=batch[self.global_name], component=component ) except StopIteration: component.complete = True @@ -407,28 +417,28 @@ def write_chunk(self, components: Components) -> int: # Apply track selections for name in self.variables.variables: - if name == self.jets_name: + if name == self.global_name: continue if selector := self.variables.selectors.get(name): merged[name] = selector(merged[name]) - # Get the total length of jets from the batch and how much + # Get the total length of objects from the batch and how much # capacity is left in the file - merged_len = len(merged[self.jets_name]) - capacity_left = self.writer.num_jets - self.writer.num_written + merged_len = len(merged[self.global_name]) + capacity_left = self.writer.num_global_objects - self.writer.num_written if self._fast_forwarding: # Limit consumption to the remaining discard quota if merged_len <= capacity_left: self.writer.write(merged) - self.jets_written += merged_len + self.global_objects_written += merged_len return merged_len else: head = {n: a[:capacity_left] for n, a in merged.items()} tail = {n: a[capacity_left:] for n, a in merged.items()} self.writer.write(head) self._ff_pending = tail # keep remainder for next iteration - self.jets_written += capacity_left + self.global_objects_written += capacity_left return capacity_left # If current file is full (and not fast-forwarding), roll to next file @@ -438,15 +448,15 @@ def write_chunk(self, components: Components) -> int: # open the next one self._file_idx += 1 - remaining_total = self.total_jets - self.jets_written + remaining_total = self.total_global_objects - self.global_objects_written - # Quit writing when no jets are left to write + # Quit writing when no objects are left to write if remaining_total == 0: return 0 next_file_size = ( - min(self.num_jets_per_output_file, remaining_total) - if self.num_jets_per_output_file + min(self.num_global_objects_per_output_file, remaining_total) + if self.num_global_objects_per_output_file else remaining_total ) self._open_writer( @@ -457,7 +467,7 @@ def write_chunk(self, components: Components) -> int: ) # Recompute free space in the freshly-opened file - capacity_left = self.writer.num_jets - self.writer.num_written + capacity_left = self.writer.num_global_objects - self.writer.num_written # Write (or discard) the batch if merged_len <= capacity_left or self._fast_forwarding: @@ -470,12 +480,14 @@ def write_chunk(self, components: Components) -> int: self.writer.write(head) self.writer.close() - # Open a fresh file sized for the remaining jets + # Open a fresh file sized for the remaining objects self._file_idx += 1 - remaining_total = self.total_jets - (self.jets_written + capacity_left) + remaining_total = self.total_global_objects - ( + self.global_objects_written + capacity_left + ) next_file_size = ( - min(self.num_jets_per_output_file, remaining_total) - if self.num_jets_per_output_file + min(self.num_global_objects_per_output_file, remaining_total) + if self.num_global_objects_per_output_file else remaining_total ) self._open_writer(self._sample, next_file_size, self._file_idx, self.current_components) @@ -485,15 +497,15 @@ def write_chunk(self, components: Components) -> int: self.writer.write(tail) # Updating the progress-bar - self.jets_written += merged_len + self.global_objects_written += merged_len return merged_len def write_components(self, sample: str | None, components: Components) -> None: """Merge *components* into one or more HDF5 files. - If ``self.num_jets_per_output_file`` is ``None`` the behaviour is identical to the + If ``self.num_global_objects_per_output_file`` is ``None`` the behaviour is identical to the original implementation (exactly one output file). Otherwise the function - keeps opening new `H5Writer`s whenever the current file reaches that jet + keeps opening new `H5Writer`s whenever the current file reaches that object limit. All heavy work (splitting batches, rolling files) is handled in ``self.write_chunk``. @@ -504,54 +516,58 @@ def write_components(self, sample: str | None, components: Components) -> None: components : Components Components that are to be written """ - # Resolve "write all" (num_jets < 0) to the actual number of jets on disk + # Resolve "write all" (num_global_objects < 0) to the actual number of objects on disk for component in components: - if component.num_jets < 0: + if component.num_global_objects < 0: component.setup_reader( - self.batch_size, fname=component.out_path, jets_name=self.jets_name + self.batch_size, fname=component.out_path, global_name=self.global_name ) - component.num_jets = component.reader.num_jets + component.num_global_objects = component.reader.num_global_objects # Prepare every Component's reader for component in components: - batch_size = self.batch_size * component.num_jets // components.num_jets + 1 + batch_size = ( + self.batch_size * component.num_global_objects // components.num_global_objects + 1 + ) component.setup_reader( batch_size, fname=component.out_path, - jets_name=self.jets_name, + global_name=self.global_name, ) component.stream = component.reader.stream( self.variables.combined(), - component.reader.num_jets, + component.reader.num_global_objects, ) component.complete = False # Cache dtype / base shapes once (re-used for every new file) self.dtypes = components[0].reader.dtypes(self.variables.combined()) - self.base_shapes = components[0].reader.shapes(components.num_jets, self.variables.keys()) + self.base_shapes = components[0].reader.shapes( + components.num_global_objects, self.variables.keys() + ) # Bookkeeping shared with write_chunk - self.total_jets = components.num_jets - self.jets_written = 0 + self.total_global_objects = components.num_global_objects + self.global_objects_written = 0 self._file_idx = 0 self._sample = sample self.current_components = components # Auto-resume: detect contiguous valid parts; delete a corrupt last part if found resume_parts = 0 - if self.resume and isinstance(self.num_jets_per_output_file, int): + if self.resume and isinstance(self.num_global_objects_per_output_file, int): resume_parts = self._detect_and_clean_completed_parts(sample) - if resume_parts and isinstance(self.num_jets_per_output_file, int): - to_discard = resume_parts * int(self.num_jets_per_output_file) + if resume_parts and isinstance(self.num_global_objects_per_output_file, int): + to_discard = resume_parts * int(self.num_global_objects_per_output_file) log.info( f"[bold yellow]Resuming merge: found {resume_parts} completed part(s); " - f"skipping first {to_discard:,} jets." + f"skipping first {to_discard:,} objects." ) # Use a NullWriter to pre-consume data via the exact same logic self._fast_forwarding = True self.writer = self._NullWriter(to_discard) - while self.jets_written < to_discard: + while self.global_objects_written < to_discard: consumed = self.write_chunk(components) if consumed == 0: break @@ -560,13 +576,13 @@ def write_components(self, sample: str | None, components: Components) -> None: # Align counters with the next missing part self._file_idx = resume_parts - self.jets_written = to_discard + self.global_objects_written = to_discard # Decide capacity of the first real file - remaining_total = self.total_jets - self.jets_written + remaining_total = self.total_global_objects - self.global_objects_written first_file_size = ( - min(self.num_jets_per_output_file, remaining_total) - if self.num_jets_per_output_file + min(self.num_global_objects_per_output_file, remaining_total) + if self.num_global_objects_per_output_file else remaining_total ) @@ -576,11 +592,11 @@ def write_components(self, sample: str | None, components: Components) -> None: # Main merge loop with progress with ProgressBar() as progress: task = progress.add_task( - f"[green]Merging {components.num_jets:,} jets...", - total=components.num_jets, + f"[green]Merging {components.num_global_objects:,} objects...", + total=components.num_global_objects, ) - if self.jets_written: - progress.update(task, advance=self.jets_written) + if self.global_objects_written: + progress.update(task, advance=self.global_objects_written) while True: n = self.write_chunk(components) @@ -591,7 +607,7 @@ def write_components(self, sample: str | None, components: Components) -> None: # Close Writer self.writer.close() label = "merged" if sample is None else sample - log.info(f"[bold green]Finished merging {components.num_jets:,} {label} jets!") + log.info(f"[bold green]Finished merging {components.num_global_objects:,} {label} objects!") def run(self): """Run merging of the components.""" diff --git a/upp/stages/normalisation.py b/upp/stages/normalisation.py index 6ef5446c..c2bc3617 100644 --- a/upp/stages/normalisation.py +++ b/upp/stages/normalisation.py @@ -21,8 +21,8 @@ def __init__(self, config: PreprocessingConfig): self.config = config self.components = config.components self.variables = config.variables - self.jets_name = self.config.jets_name - self.num_jets = config.num_jets_estimate_norm + self.global_name = self.config.global_name + self.num_global_objects = config.num_global_objects_estimate_norm self.norm_fname = config.out_dir / config.config.get("norm_fname", "norm_dict.yaml") self.class_fname = config.out_dir / config.config.get("class_fname", "class_dict.yaml") @@ -48,9 +48,9 @@ def combine_mean_std( std_B : float Standard deviation of the variable B num_A : int - Number of jets for variable A + Number of objects for variable A num_B : int - Number of jets for variable B + Number of objects for variable B Returns ------- @@ -64,21 +64,21 @@ def combine_mean_std( return float(combined_mean), float(combined_std) def get_norm_dict(self, batch: dict) -> tuple[dict, int]: - """Get the normalisation dict with the mean and standard deviation for the given jets. + """Get the normalisation dict with the mean and standard deviation for the given objects. Parameters ---------- batch : dict - Dict with the jets and all variables + Dict with the objects and all variables Returns ------- tuple[dict, int] - Normalisation dict and the number of jets used to calculate it + Normalisation dict and the number of objects used to calculate it """ norm_dict: dict[str, dict] = {k: {} for k in self.variables} for name, array in batch.items(): - if name != self.variables.jets_name: + if name != self.variables.global_name: array = array[array["valid"]] for var in self.variables[name]["inputs"]: if var in ["valid"]: @@ -86,7 +86,7 @@ def get_norm_dict(self, batch: dict) -> tuple[dict, int]: mean = float(np.nanmean(array[var])) std = float(np.nanstd(array[var])) norm_dict[name][var] = {"mean": mean, "std": std} - return norm_dict, len(batch[self.variables.jets_name]) + return norm_dict, len(batch[self.variables.global_name]) def combine_norm_dict(self, norm_A: dict, norm_B: dict, num_A: int, num_B: int) -> dict: """Combine two normalisation dicts into one. @@ -98,9 +98,9 @@ def combine_norm_dict(self, norm_A: dict, norm_B: dict, num_A: int, num_B: int) norm_B : dict Normalisation dict B num_A : int - Number of jets used to calculate normalisation dict A + Number of objects used to calculate normalisation dict A num_B : int - Number of jets used to calculate normalisation dict B + Number of objects used to calculate normalisation dict B Returns ------- @@ -127,12 +127,12 @@ def combine_norm_dict(self, norm_A: dict, norm_B: dict, num_A: int, num_B: int) return combined def get_class_dict(self, batch: dict) -> dict: - """Get the class dict for the given jets. + """Get the class dict for the given objects. Parameters ---------- batch : dict - Dict with the jets and their variables + Dict with the objects and their variables Returns ------- @@ -148,10 +148,10 @@ def get_class_dict(self, batch: dict) -> dict: ] class_dict: dict[str, dict] = {k: {} for k in self.variables} for name, array in batch.items(): - if name != self.variables.jets_name: + if name != self.variables.global_name: array = array[array["valid"]] # separate case for flavour_label - if name == self.variables.jets_name and "flavour_label" in array.dtype.names: + if name == self.variables.global_name and "flavour_label" in array.dtype.names: counts = np.unique(array["flavour_label"], return_counts=True) class_dict[name]["flavour_label"] = counts for var in self.variables[name].get("labels", []): @@ -237,7 +237,7 @@ def run(self): fname = str(self.config.out_fname).replace(".h5", "_vds.h5") # Get the correct output names if multiple output files were written - elif self.config.num_jets_per_output_file is not None: + elif self.config.num_global_objects_per_output_file is not None: fname = ( self.config.out_fname.parent / self.config.split @@ -252,7 +252,7 @@ def run(self): fname, self.config.batch_size, precision="full", - jets_name=self.jets_name, + global_objects_name=self.global_name, ) log.debug(f"Setup reader at: {fname}") @@ -261,14 +261,14 @@ def run(self): total = None vars = self.variables.combined() with h5py.File(reader.files[0]) as f: - if "flavour_label" in f[self.jets_name].dtype.names: - vars[self.jets_name].append("flavour_label") - stream = reader.stream(vars, self.num_jets) + if "flavour_label" in f[self.global_name].dtype.names: + vars[self.global_name].append("flavour_label") + stream = reader.stream(vars, self.num_global_objects) with ProgressBar() as progress: task = progress.add_task( - f"[green]Computing normalisations using {self.num_jets:,} jets...", - total=self.num_jets, + f"[green]Computing normalisations using {self.num_global_objects:,} objects...", + total=self.num_global_objects, ) for i, batch in enumerate(stream): @@ -283,9 +283,12 @@ def run(self): norm_dict = self.combine_norm_dict(norm_dict, this_norm_dict, total, num) total += num - progress.update(task, advance=len(batch[self.variables.jets_name])) + progress.update(task, advance=len(batch[self.variables.global_name])) - log.info(f"[bold green]Finished computing normalisation params on {self.num_jets:,} jets!") + log.info( + f"[bold green]Finished computing normalisation params on " + f"{self.num_global_objects:,} objects!" + ) self.write_norm_dict(norm_dict) self.write_class_dict(class_dict) log.info(f"[bold green]Saved norm dict to {self.norm_fname}") diff --git a/upp/stages/plot.py b/upp/stages/plot.py index ee0c1a1b..1c63f6ca 100644 --- a/upp/stages/plot.py +++ b/upp/stages/plot.py @@ -27,7 +27,7 @@ class PlotRegion: name : str Name used in output file suffixes and log messages. cuts : Cuts - Jet-level selection defining the region. + Object-level selection defining the region. pt_range : tuple[float, float] | None Raw pT range defining the region. Values follow the input ntuple units, usually MeV. If ``None``, the region has no explicit pT range. @@ -157,35 +157,36 @@ def _sample_label(sample_name: str, plotting: PlottingConfig) -> str: return plotting.sample_label(sample_name) -def _format_num_jets(num_jets: int) -> str: - """Format a jet count with compact suffixes. +def _format_num_global_objects(num_global_objects: int) -> str: + """Format a object count with compact suffixes. Parameters ---------- - num_jets : int - Number of jets to format. + num_global_objects : int + Number of objects to format. Returns ------- str Compact count using ``k`` for thousands and ``M`` for millions. """ - if num_jets >= 1_000_000: - value = num_jets / 1_000_000 + if num_global_objects >= 1_000_000: + value = num_global_objects / 1_000_000 return f"{value:g}M" - if num_jets >= 1_000: - value = num_jets / 1_000 + if num_global_objects >= 1_000: + value = num_global_objects / 1_000 return f"{value:g}k" - return str(num_jets) + return str(num_global_objects) def _atlas_second_tag( *sample_names: str, plotting: PlottingConfig, - num_jets: int | None = None, + global_name: str = "jets", + num_global_objects: int | None = None, resampling_status: str | None = None, ) -> str: - """Build the second ATLAS tag with energy, sample labels, status, and jet count. + """Build the second ATLAS tag with energy, sample labels, status, and object count. Parameters ---------- @@ -193,8 +194,11 @@ def _atlas_second_tag( Sample names to include after the centre-of-mass energy. plotting : PlottingConfig Active plotting configuration. - num_jets : int | None, optional - Number of jets requested for plotting. If provided, it is added as an + global_name : str, optional + Name of the global object used to label the sample and count lines, + by default "jets". + num_global_objects : int | None, optional + Number of objects requested for plotting. If provided, it is added as an extra line using compact formatting. resampling_status : str | None, optional Resampling status added as an extra line. @@ -204,37 +208,37 @@ def _atlas_second_tag( str Multiline ATLAS second tag. Sample labels share the first line with the centre-of-mass energy, followed by the optional resampling status and - jet count. + object count. """ labels = [_sample_label(name, plotting) for name in dict.fromkeys(sample_names) if name] first_line = plotting.atlas_second_tag if labels: - first_line = f"{first_line}, {' + '.join(labels)} jets" + first_line = f"{first_line}, {' + '.join(labels)} {global_name}" lines = [first_line] if resampling_status is not None: lines.append(resampling_status) - if num_jets is not None: - lines.append(f"{_format_num_jets(num_jets)} jets") + if num_global_objects is not None: + lines.append(f"{_format_num_global_objects(num_global_objects)} {global_name}") return "\n".join(lines) -def _plotting_num_jets(config: PreprocessingConfig, available_jets: int) -> int: - """Return the number of jets requested for plotting. +def _plotting_num_global_objects(config: PreprocessingConfig, available_global_objects: int) -> int: + """Return the number of objects requested for plotting. Parameters ---------- config : PreprocessingConfig Active preprocessing configuration. - available_jets : int - Number of jets available for the plotted selection. + available_global_objects : int + Number of objects available for the plotted selection. Returns ------- int - Minimum of the available jets and ``plotting.num_jets_plotting``. + Minimum of the available objects and ``plotting.num_global_objects_plotting``. """ - assert config.plotting.num_jets_plotting is not None - return min(available_jets, config.plotting.num_jets_plotting) + assert config.plotting.num_global_objects_plotting is not None + return min(available_global_objects, config.plotting.num_global_objects_plotting) def _pt_bounds_from_cuts(cuts: Cuts, pt_variable: str) -> tuple[float, float] | None: @@ -243,7 +247,7 @@ def _pt_bounds_from_cuts(cuts: Cuts, pt_variable: str) -> tuple[float, float] | Parameters ---------- cuts : Cuts - Jet-level cuts associated with a component region. + Object-level cuts associated with a component region. pt_variable : str Name of the pT variable used for resampling. @@ -418,8 +422,10 @@ def _stitching_regions(regions: list[PlotRegion], pt_variable: str | None) -> li return stitching_regions -def _load_jets(config: PreprocessingConfig, in_paths: Any, vars_to_load: list[str]) -> Any: - """Load jet variables for plotting. +def _load_global_objects( + config: PreprocessingConfig, in_paths: Any, vars_to_load: list[str] +) -> Any: + """Load object variables for plotting. Parameters ---------- @@ -428,24 +434,24 @@ def _load_jets(config: PreprocessingConfig, in_paths: Any, vars_to_load: list[st in_paths : Any Input HDF5 file path, glob, or list of paths passed to ``H5Reader``. vars_to_load : list[str] - Jet variables needed for plotting and selections. + Object variables needed for plotting and selections. Returns ------- Any - Structured jet array loaded from the input files. + Structured object array loaded from the input files. """ return H5Reader( fname=in_paths, batch_size=config.batch_size, - jets_name=config.jets_name, + global_objects_name=config.global_name, shuffle=False, - equal_jets=True, + equal_global_objects=True, vds_dir=config.vds_dir, ).load( - {config.jets_name: list(dict.fromkeys(vars_to_load))}, - num_jets=config.plotting.num_jets_plotting, - )[config.jets_name] + {config.global_name: list(dict.fromkeys(vars_to_load))}, + num_global_objects=config.plotting.num_global_objects_plotting, + )[config.global_name] def make_hist( @@ -454,7 +460,7 @@ def make_hist( flavours: list, variable: str, out_dir: Path, - jets_name: str = "jets", + global_name: str = "jets", bins_range: tuple | None = None, suffix: str = "", out_format_list: tuple[str, ...] | list[str] | None = None, @@ -484,8 +490,8 @@ def make_hist( Variable that is to be histogrammed and plotted. out_dir : Path Output directory to which the plots are written. - jets_name: str, optional - Name of the jet dataset / the global objects + global_name: str, optional + Name of the object dataset / the global objects by default "jets" bins_range : tuple | None, optional bins_range argument from from puma.HistogramPlot, @@ -510,7 +516,7 @@ def make_hist( # Setup the histogram plot = HistogramPlot( - ylabel=plotting.ylabel.replace("{jets_name}", jets_name), + ylabel=plotting.ylabel.replace("{global_name}", global_name), xlabel=plotting.variable_label(variable), y_scale=plotting.y_scale, figsize=plotting.figsize, @@ -601,7 +607,7 @@ def _plot_initial(config: PreprocessingConfig) -> None: vars_to_load += flavour.cuts.variables values_dict = { - sample.name: _load_jets(config, list(sample.path), vars_to_load), + sample.name: _load_global_objects(config, list(sample.path), vars_to_load), } pt_range = _pt_bounds_from_cuts(selection_cuts, pt_var) if pt_var else None @@ -615,7 +621,7 @@ def _plot_initial(config: PreprocessingConfig) -> None: make_hist( stage="initial", values_dict=values_dict, - jets_name=config.jets_name, + global_name=config.global_name, flavours=region_components.flavours, variable=variable, bins_range=bins_range, @@ -624,8 +630,11 @@ def _plot_initial(config: PreprocessingConfig) -> None: atlas_second_tag=_atlas_second_tag( sample.name, plotting=config.plotting, - num_jets=_plotting_num_jets(config, region_components.num_jets) - if config.plotting.show_num_jets + global_name=config.global_name, + num_global_objects=_plotting_num_global_objects( + config, region_components.num_global_objects + ) + if config.plotting.show_num_global_objects else None, resampling_status="Pre Resampling", ), @@ -653,13 +662,13 @@ def _post_resampling_paths(config: PreprocessingConfig, stage: str) -> list[Path return [ ( config.out_fname.parent / config.split / f"{config.out_fname.stem}*.h5" - if config.num_jets_per_output_file is not None + if config.num_global_objects_per_output_file is not None else config.out_fname ) ] paths = [path_append(config.out_fname, sample.name) for sample in config.components.samples] - if config.num_jets_per_output_file is not None: + if config.num_global_objects_per_output_file is not None: return [path.parent / config.split / f"{path.stem}*.h5" for path in paths] return paths @@ -684,8 +693,11 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: atlas_second_tag = _atlas_second_tag( *sample_names, plotting=config.plotting, - num_jets=_plotting_num_jets(config, config.components.num_jets) - if config.plotting.show_num_jets + global_name=config.global_name, + num_global_objects=_plotting_num_global_objects( + config, config.components.num_global_objects + ) + if config.plotting.show_num_global_objects else None, resampling_status="Post Resampling", ) @@ -695,7 +707,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: vars_to_load += region.cuts.variables values_dict = { - "": _load_jets(config, _post_resampling_paths(config, stage), vars_to_load), + "": _load_global_objects(config, _post_resampling_paths(config, stage), vars_to_load), } for variable in config.sampl_cfg.vars: @@ -709,7 +721,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: make_hist( stage=stage, values_dict=values_dict, - jets_name=config.jets_name, + global_name=config.global_name, flavours=config.components.flavours, variable=variable, bins_range=bins_range, @@ -726,7 +738,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: make_hist( stage=stage, values_dict=values_dict, - jets_name=config.jets_name, + global_name=config.global_name, flavours=config.components.flavours, variable=variable, bins_range=_display_range(variable, stitching_region.pt_range), diff --git a/upp/stages/resampling.py b/upp/stages/resampling.py index dc9888b3..7180d958 100644 --- a/upp/stages/resampling.py +++ b/upp/stages/resampling.py @@ -10,7 +10,7 @@ from ftag.hdf5 import H5Reader from yamlinclude import YamlIncludeConstructor -from upp.stages.hist import bin_jets +from upp.stages.hist import bin_global_objects from upp.stages.interpolation import subdivide_bins, upscale_array_regionally from upp.utils.logger import ProgressBar @@ -52,14 +52,14 @@ def __init__(self, config: PreprocessingConfig): self.components = config.components self.variables = config.variables self.batch_size = config.batch_size - self.jets_name = config.jets_name + self.global_name = config.global_name self.transform = config.transform self.rng = np.random.default_rng(42) # Define what type self.target will be self.target = cast("Component", None) - # In skip mode no selection is applied; jets are written as-is (see sample()) + # In skip mode no selection is applied; objects are written as-is (see sample()) if self.skip: self.method = None self.upscale_pdf = 1 @@ -81,20 +81,20 @@ def __init__(self, config: PreprocessingConfig): ) self.select_func = self.methods_map[self.method] - def countup_select_func(self, jets: dict, component: Component) -> np.ndarray: + def countup_select_func(self, global_objects: dict, component: Component) -> np.ndarray: """Countup resampling function. Parameters ---------- - jets : dict - Dict with the jets which are to be resampled. + global_objects : dict + Dict with the objects which are to be resampled. component : Component Component instance for a given flavour/class. Returns ------- np.ndarray - Numpy array with the index numbers of the jets that are to be used. + Numpy array with the index numbers of the objects that are to be used. Raises ------ @@ -106,22 +106,22 @@ def countup_select_func(self, jets: dict, component: Component) -> np.ndarray: if self.upscale_pdf != 1: raise ValueError("Upscaling of histograms is not supported for countup method") - # Get the target number of jets and target PDF values - num_jets = int(len(jets) * component.sampling_fraction) + # Get the target number of objects and target PDF values + num_global_objects = int(len(global_objects) * component.sampling_fraction) target_pdf = self.target.hist.pbin # Get the target histograms - target_hist = target_pdf * num_jets + target_hist = target_pdf * num_global_objects target_hist = (np.floor(target_hist + self.rng.random(target_pdf.shape))).astype(int) # Create histogram and bins for the given resampling variables - _hist, binnumbers = bin_jets( - array=jets[self.config.vars], + _hist, binnumbers = bin_global_objects( + array=global_objects[self.config.vars], bins=self.config.flat_bins, ) assert target_pdf.shape == _hist.shape - # Loop over bins and select relevant jets (indicies) + # Loop over bins and select relevant objects (indicies) all_idx = [] for bin_id in np.ndindex(*target_hist.shape): idx = np.where((bin_id == binnumbers.T).all(axis=-1))[0][: target_hist[bin_id]] @@ -130,25 +130,25 @@ def countup_select_func(self, jets: dict, component: Component) -> np.ndarray: all_idx.append(idx) idx = np.concatenate(all_idx).astype(int) - # If not enough jets are found, re-use randomly some of them - if len(idx) < num_jets: - idx = np.concatenate([idx, self.rng.choice(idx, num_jets - len(idx))]) + # If not enough objects are found, re-use randomly some of them + if len(idx) < num_global_objects: + idx = np.concatenate([idx, self.rng.choice(idx, num_global_objects - len(idx))]) - # Shuffle the jet indicies + # Shuffle the object indicies self.rng.shuffle(idx) return idx - def pdf_select_func(self, jets: dict, component: Component) -> np.ndarray: - # bin jets + def pdf_select_func(self, global_objects: dict, component: Component) -> np.ndarray: + # bin objects if self.upscale_pdf > 1: bins = [subdivide_bins(bins, self.upscale_pdf) for bins in self.config.flat_bins] else: bins = self.config.flat_bins # Create histogram and bins for the given resampling variables - _hist, binnumbers = bin_jets( - array=jets[self.config.vars], + _hist, binnumbers = bin_global_objects( + array=global_objects[self.config.vars], bins=bins, ) @@ -157,7 +157,7 @@ def pdf_select_func(self, jets: dict, component: Component) -> np.ndarray: binnumbers = tuple(binnumbers[i] for i in range(len(binnumbers))) # importance sample with replacement - num_samples = int(len(jets) * component.sampling_fraction) + num_samples = int(len(global_objects) * component.sampling_fraction) # Calculate the ratios between the target and the to-be-resampled distribution ratios = safe_divide(a=self.target.hist.pbin, b=component.hist.pbin) @@ -173,13 +173,13 @@ def pdf_select_func(self, jets: dict, component: Component) -> np.ndarray: # Get the probabilities for the resampling from the ratios probs = ratios[binnumbers] - # Select the jets (indicies) for the resampled final output - idx = random.choices(np.arange(len(jets)), weights=probs, k=num_samples) + # Select the objects (indicies) for the resampled final output + idx = random.choices(np.arange(len(global_objects)), weights=probs, k=num_samples) return idx def track_upsampling_stats(self, idx: np.ndarray, component: Component) -> None: - """Tracking the upsampling ratio and update the number of unique jets. + """Tracking the upsampling ratio and update the number of unique objects. Parameters ---------- @@ -189,16 +189,16 @@ def track_upsampling_stats(self, idx: np.ndarray, component: Component) -> None: Component instance for a given flavour/class. """ unique, ups_counts = np.unique(idx, return_counts=True) - component._unique_jets += len(unique) + component._unique_global_objects += len(unique) max_ups = ups_counts.max() component._ups_max = max_ups if max_ups > component._ups_max else component._ups_max def _finalise_component(self, component: Component) -> None: """Write metadata attrs and close the writer for a completed component.""" - unique = component._unique_jets + unique = component._unique_global_objects component._ups_ratio = component.writer.num_written / unique if unique else 0.0 component.writer.add_attr("upsampling_ratio", component._ups_ratio) - component.writer.add_attr("unique_jets", component._unique_jets) + component.writer.add_attr(f"unique_{self.global_name}", component._unique_global_objects) component.writer.add_attr("dsid", str(component.sample.dsid)) component.writer.close() @@ -209,14 +209,14 @@ def sample( progress: Progress, selected_component: str | None = None, ) -> None: - """Sample the jets by the given selected indicies from the resampling function. + """Sample the objects by the given selected indicies from the resampling function. Parameters ---------- components : Components Components instance of the components which are to be resampled. stream : Generator[Any, None, None] - Generator of the jets which are to be resampled. + Generator of the objects which are to be resampled. progress : Progress Progress bar instance for updating the shown progress bar. selected_component : str | None, optional @@ -226,7 +226,7 @@ def sample( Raises ------ ValueError - If not enough jets for a given component are present. + If not enough objects for a given component are present. """ # Loop through input file for batch in stream: @@ -241,34 +241,34 @@ def sample( continue # Apply selections - comp_idx, _ = component.flavour.cuts(batch[self.variables.jets_name]) + comp_idx, _ = component.flavour.cuts(batch[self.variables.global_name]) if len(comp_idx) == 0: continue - # Get the batch of jets as a separate dict + # Get the batch of objects as a separate dict batch_out = select_batch(batch, comp_idx) # Apply sampling - idx = np.arange(len(batch_out[self.variables.jets_name])) + idx = np.arange(len(batch_out[self.variables.global_name])) # Check that the component is not the target and a resampling # function is set. if component != self.target and self.select_func: # Apply the resampling idx = self.select_func( - jets=batch_out[self.variables.jets_name], + global_objects=batch_out[self.variables.global_name], component=component, ) if len(idx) == 0: continue - # Get the resampled jets in the dict + # Get the resampled objects in the dict batch_out = select_batch(batch_out, idx) - # Check for completion (only when a finite num_jets is requested; - # num_jets < 0 means write all jets passing the cuts) - if 0 <= component.num_jets <= component.writer.num_written + len(idx): - keep = component.num_jets - component.writer.num_written + # Check for completion (only when a finite num_global_objects is requested; + # num_global_objects < 0 means write all objects passing the cuts) + if 0 <= component.num_global_objects <= component.writer.num_written + len(idx): + keep = component.num_global_objects - component.writer.num_written idx = idx[:keep] for name, array in batch_out.items(): batch_out[name] = array[:keep] @@ -277,7 +277,7 @@ def sample( # Track upsampling stats self.track_upsampling_stats(idx=idx, component=component) - # Write the resampled jets to file + # Write the resampled objects to file component.writer.write(batch_out) # Update the progress bar @@ -291,17 +291,17 @@ def sample( if all(component._complete for component in components): break - # Finalise components: unbounded ones (num_jets < 0) are done once the input is - # exhausted; bounded ones that didn't reach their target ran out of jets. + # Finalise components: unbounded ones (num_global_objects < 0) are done once the input is + # exhausted; bounded ones that didn't reach their target ran out of objects. for component in components: if component._complete: continue - if component.num_jets < 0: + if component.num_global_objects < 0: component._complete = True self._finalise_component(component) else: raise ValueError( - f"Ran out of {component} jets after writing {component.writer.num_written:,}" + f"Ran out of {component} objects after writing {component.writer.num_written:,}" ) def run_on_region( @@ -325,7 +325,7 @@ def run_on_region( Raises ------ ValueError - If the equal_jets flag is not the same for all components. + If the equal_global_objects flag is not the same for all components. """ # Get the target component (not needed when resampling is skipped) if not self.skip: @@ -337,27 +337,31 @@ def run_on_region( grouped_samples = components.groupby_sample() for sample, components in grouped_samples: - # Ensure all components have the same equal_jets flag - equal_jets_flags = [component.equal_jets for component in components] - if len(set(equal_jets_flags)) != 1: - raise ValueError("equal_jets must be the same for all components in a sample") - equal_jets_flag = equal_jets_flags[0] + # Ensure all components have the same equal_global_objects flag + equal_global_objects_flags = [ + component.equal_global_objects for component in components + ] + if len(set(equal_global_objects_flags)) != 1: + raise ValueError( + "equal_global_objects must be the same for all components in a sample" + ) + equal_global_objects_flag = equal_global_objects_flags[0] # Get the variables which are to be used - variables = self.variables.add_jet_vars(components.cuts.variables) + variables = self.variables.add_global_vars(components.cuts.variables) - # Setup the Reader for reading the jets + # Setup the Reader for reading the objects reader = H5Reader( sample.path, self.batch_size, - jets_name=self.jets_name, - equal_jets=equal_jets_flag, + global_objects_name=self.global_name, + equal_global_objects=equal_global_objects_flag, transform=self.transform, vds_dir=sample.vds_dir, ) - # Define a stream of jets with the cuts for the region and the variables used - stream = reader.stream(variables.combined(), reader.num_jets, region.cuts) + # Define a stream of objects with the cuts for the region and the variables used + stream = reader.stream(variables.combined(), reader.num_global_objects, region.cuts) # Run with progress bar with ProgressBar() as progress: @@ -368,12 +372,12 @@ def run_on_region( if selected_component and selected_component != component.name: component._complete = True component._ups_max = 0.0 - component._unique_jets = 0 + component._unique_global_objects = 0 else: component._complete = False component._ups_max = 1.0 - component._unique_jets = 0 + component._unique_global_objects = 0 # Add each component to the progress bar for component in components: @@ -381,11 +385,11 @@ def run_on_region( if selected_component and selected_component != component.name: continue - unbounded = component.num_jets < 0 - label = "all" if unbounded else f"{component.num_jets:,}" + unbounded = component.num_global_objects < 0 + label = "all" if unbounded else f"{component.num_global_objects:,}" component.pbar = progress.add_task( - f"[green]Sampling {label} jets from {component}...", - total=None if unbounded else component.num_jets, + f"[green]Sampling {label} objects from {component}...", + total=None if unbounded else component.num_global_objects, ) # Run the actual resampling sampling @@ -406,8 +410,8 @@ def run_on_region( log.info( f"{component} usampling ratio is {np.mean(component._ups_ratio):.3f}, with" f" {written / np.mean(component._ups_ratio):,.0f}/" - f"{written:,} written jets." - f" Jets are upsampled at most {np.max(component._ups_max):.0f} times" + f"{written:,} written objects." + f" Objects are upsampled at most {np.max(component._ups_max):.0f} times" ) def set_component_sampling_fractions(self, component: Component) -> None: @@ -441,7 +445,7 @@ def set_component_sampling_fractions(self, component: Component) -> None: else: sam_frac = component.get_auto_sampling_fraction( - num_jets=component.num_jets, + num_global_objects=component.num_global_objects, cuts=component.cuts, ) @@ -513,7 +517,7 @@ def run(self, region: str | None = None, component: str | None = None): # Setup the reader for the components iter_component.setup_reader( - self.batch_size, jets_name=self.jets_name, transform=self.transform + self.batch_size, global_name=self.global_name, transform=self.transform ) # If only one component is run, stop here for the target that needs to be @@ -522,20 +526,20 @@ def run(self, region: str | None = None, component: str | None = None): continue # Setup the writer for the component - iter_component.setup_writer(self.variables, jets_name=self.jets_name) + iter_component.setup_writer(self.variables, global_name=self.global_name) # Set sampling fraction self.set_component_sampling_fractions(component=iter_component) - # Check that enough jets are available + # Check that enough objects are available sampling_fraction = 1 if self.skip else self.config.sampling_fraction log.info( - "[bold green]Checking requested num_jets based on a sampling fraction of" + "[bold green]Checking requested num_global_objects based on a sampling fraction of" f" {sampling_fraction}..." ) frac = iter_component.sampling_fraction if self.select_func else 1 - iter_component.check_num_jets( - iter_component.num_jets, + iter_component.check_num_global_objects( + iter_component.num_global_objects, sampling_fraction=frac, cuts=iter_component.cuts, ) @@ -573,22 +577,24 @@ def run(self, region: str | None = None, component: str | None = None): # If a component is given, skip all components that are not selected if component and iter_component.name != component: continue - unique += iter_component.writer.get_attr("unique_jets") + unique += iter_component.writer.get_attr(f"unique_{self.global_name}") log.info( f"[bold green]Finished resampling of region {region}. " - f"A total of {self.components.num_jets:,} jets!" + f"A total of {self.components.num_global_objects:,} objects!" ) - log.info(f"[bold green]Estimated unique jets: {unique:,.0f}") + log.info(f"[bold green]Estimated unique objects: {unique:,.0f}") log.info(f"[bold green]Saved to {self.components.out_dir}/") else: unique = sum( - iter_component.writer.get_attr("unique_jets") for iter_component in self.components + iter_component.writer.get_attr(f"unique_{self.global_name}") + for iter_component in self.components ) log.info( - f"[bold green]Finished resampling a total of {self.components.num_jets:,} jets!" + f"[bold green]Finished resampling a total of " + f"{self.components.num_global_objects:,} objects!" ) - log.info(f"[bold green]Estimated unique jets: {unique:,.0f}") + log.info(f"[bold green]Estimated unique objects: {unique:,.0f}") log.info(f"[bold green]Saved to {self.components.out_dir}/") def get_num_bins_from_config(self) -> list[list[int]]: diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index 99ff3b20..bd266ff7 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -12,7 +12,7 @@ from puma import Histogram, HistogramPlot from upp.classes.preprocessing_config import PreprocessingConfig -from upp.stages.hist import bin_jets +from upp.stages.hist import bin_global_objects class Reweight: @@ -35,8 +35,8 @@ def hists_path(self): return self.config.out_dir / "histograms.h5" @property - def num_jets_estimate(self): - return self.rw_config.num_jets_estimate or self.config.num_jets_estimate + def num_global_objects_estimate(self): + return self.rw_config.num_global_objects_estimate or self.config.num_global_objects_estimate def get_input_readers(self): components_config = self.organised_components_config @@ -48,23 +48,24 @@ def get_input_readers(self): f: H5Reader( files_by_flavour[f], batch_size=self.config.batch_size, + global_objects_name=self.config.global_name, ) for f in files_by_flavour } - per_reader_num_jets = [] + per_reader_num_global_objects = [] for f, r in input_readers.items(): - n = min(self.num_jets_estimate, r.num_jets) - if r.num_jets < self.num_jets_estimate: + n = min(self.num_global_objects_estimate, r.num_global_objects) + if r.num_global_objects < self.num_global_objects_estimate: print( - f"WARNING: Requested {self.num_jets_estimate} jets for {f}, " - f"but only {r.num_jets} available. Using {r.num_jets}." + f"WARNING: Requested {self.num_global_objects_estimate} objects for {f}, " + f"but only {r.num_global_objects} available. Using {r.num_global_objects}." ) print( - f"Flavour {f} has {r.num_jets} jets, using {n}, " + f"Flavour {f} has {r.num_global_objects} objects, using {n}, " f"reading in batches of {self.config.batch_size}" ) - per_reader_num_jets.append(n) - return list(input_readers.values()), per_reader_num_jets + per_reader_num_global_objects.append(n) + return list(input_readers.values()), per_reader_num_global_objects def calculate_weights( self, @@ -88,7 +89,7 @@ def calculate_weights( """ reweights = self.rw_config.reweights print(f"Calculating weights for {len(reweights)} reweights") - readers, per_reader_num_jets = self.get_input_readers() + readers, per_reader_num_global_objects = self.get_input_readers() for reader in readers: assert reader.batch_size == readers[0].batch_size, ( "All readers must have the same batch size" @@ -104,7 +105,7 @@ def calculate_weights( rw_groups = list(set([rw.group for rw in reweights])) print("Found rw groups : ", rw_groups) print("Batch size : ", batch_size_per_file) - print("N per file : ", self.num_jets_estimate) + print("N per file : ", self.num_global_objects_estimate) # Get the variables we need to reweight for rw in reweights: @@ -116,19 +117,19 @@ def calculate_weights( all_vars[rw_group].extend(rw.reweight_vars) if "valid" in existing_vars[rw_group]: all_vars[rw_group] += ["valid"] - if "jets" not in all_vars: - all_vars["jets"] = ["pt"] + if self.config.global_name not in all_vars: + all_vars[self.config.global_name] = ["pt"] all_vars = {k: list(set(v)) for k, v in all_vars.items()} num_in_hists = {} all_histograms = {} print("Setting up streams with vars: ", all_vars, flush=True) reader_streams = [ - r.stream(all_vars, num_jets=n) - for r, n in zip(readers, per_reader_num_jets, strict=False) + r.stream(all_vars, num_global_objects=n) + for r, n in zip(readers, per_reader_num_global_objects, strict=False) ] - max_num_jets = max(per_reader_num_jets) - num_batches = max_num_jets // batch_size_per_file + ( - 1 if max_num_jets % batch_size_per_file != 0 else 0 + max_num_global_objects = max(per_reader_num_global_objects) + num_batches = max_num_global_objects // batch_size_per_file + ( + 1 if max_num_global_objects % batch_size_per_file != 0 else 0 ) start_time = time.time() for i in range(num_batches): @@ -174,7 +175,7 @@ def calculate_weights( for cls in classes: mask = data[rw.class_var] == cls - hist, _outbins = bin_jets(data[mask][rw.reweight_vars], rw.flat_bins) + hist, _outbins = bin_global_objects(data[mask][rw.reweight_vars], rw.flat_bins) if rw.class_var is not None: cls = str(cls) if rw_group not in all_histograms: diff --git a/upp/stages/rw_merge.py b/upp/stages/rw_merge.py index 36f95768..41c88f35 100644 --- a/upp/stages/rw_merge.py +++ b/upp/stages/rw_merge.py @@ -10,7 +10,7 @@ from ftag.hdf5 import H5Reader, H5Writer, join_structured_arrays from ftag.vds import create_virtual_file -from upp.stages.hist import bin_jets +from upp.stages.hist import bin_global_objects # from ftag_rw.weights.rw_utils import get_sample_weights from upp.stages.reweight import Reweight @@ -37,14 +37,16 @@ def __init__(self, config, outfile_idx_range=None): with open(self.organised_components_config) as f: organised_components = yaml.safe_load(f) - num_jets = sum(organised_components["num_jets"][self.config.split].values()) + num_global_objects = sum( + organised_components["num_global_objects"][self.config.split].values() + ) self.attr_to_write = { - "jets": { + self.config.global_name: { "flavour_label": [f.name for f in self.config.components.flavours], }, None: { - "unique_jets": num_jets, - "jet_counts": num_jets, + f"unique_{self.config.global_name}": num_global_objects, + f"{self.config.global_name}_counts": num_global_objects, "dsids": str(self.config.components.dsids), "config": json.dumps(self.config.config), "upp_hash": self.config.git_hash, @@ -57,32 +59,38 @@ def run(self): with open(self.organised_components_config) as f: organised_components = yaml.safe_load(f) - # Get the number of jets per flavour + # Get the number of objects per flavour files_by_flavour = organised_components["files"][self.config.split] - num_jets_per_flavours = organised_components["num_jets"][self.config.split] + num_global_objects_per_flavours = organised_components["num_global_objects"][ + self.config.split + ] all_files = [] for f in files_by_flavour: all_files.extend(files_by_flavour[f]) - total_jets = sum(num_jets_per_flavours.values()) + total_global_objects = sum(num_global_objects_per_flavours.values()) batch_size = 250_000 reader_kwargs = { "fname": all_files, "batch_size": batch_size, "shuffle": False, + "global_objects_name": self.config.global_name, } output_dir = self.config.out_dir / self.config.split output_dir.mkdir(parents=True, exist_ok=True) - num_jets_per_file = self.config.num_jets_per_output_file or total_jets + num_global_objects_per_file = ( + self.config.num_global_objects_per_output_file or total_global_objects + ) - batches_per_file = num_jets_per_file // batch_size or 1 + batches_per_file = num_global_objects_per_file // batch_size or 1 num_batches = ( - total_jets // batch_size + (1 if total_jets % num_jets_per_file != 0 else 0) + total_global_objects // batch_size + + (1 if total_global_objects % num_global_objects_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"] + variables[self.config.global_name] += ["flavour_label"] args_list = [] for i, bi in enumerate(range(0, num_batches, batches_per_file)): args_list.append( @@ -90,7 +98,7 @@ def run(self): reader_kwargs, output_dir / f"pp_output_{self.config.split}-full_{i}.h5", weights, - num_jets_per_flavours, + num_global_objects_per_flavours, variables, bi, i, @@ -98,6 +106,7 @@ def run(self): if (bi + batches_per_file) < num_batches else (num_batches - bi), self.attr_to_write, + self.config.global_name, ) ) print("Running with ", self.rw_config.merge_num_proc, "processes") @@ -164,7 +173,7 @@ def get_sample_weights( rw_vars = rw["rw_vars"] class_var = rw["class_var"] - _, bins = bin_jets(to_dump[rw_vars], rw["bins"]) + _, bins = bin_global_objects(to_dump[rw_vars], rw["bins"]) # Enforce that bins are of shape (nvars, num_objects) if len(rw["bins"]) == 1: bins = np.expand_dims(bins, axis=0) @@ -200,12 +209,13 @@ def do_merge_with_weights( reader_kwargs: dict[str, dict], output_file: str, weights: dict, - num_jets_per_flavour: dict, + num_global_objects_per_flavour: dict, variables: dict[str, list[str]] | None = None, skip_batches: int = 0, writer_id=0, limit_batches=False, attrs=None, + global_name="jets", ): """Take a series of input files and merge them into a single final output file. @@ -221,7 +231,7 @@ def do_merge_with_weights( ) batch_size = reader.batch_size # reader = H5Reader(input_file) - # num_jets = reader.num_jets if N == -1 else N + # num_global_objects = reader.num_global_objects if N == -1 else N writer: H5Writer = None additional_vars = {} @@ -238,12 +248,12 @@ def do_merge_with_weights( for group in dtypes: dtypes[group] = np.dtype(dtypes[group]) - num_jets_total = sum(num_jets_per_flavour.values()) + num_global_objects_total = sum(num_global_objects_per_flavour.values()) if limit_batches: num_batches = limit_batches else: - num_batches = num_jets_total // batch_size + ( - 1 if num_jets_total % batch_size != 0 else 0 + num_batches = num_global_objects_total // batch_size + ( + 1 if num_global_objects_total % batch_size != 0 else 0 ) print(f"Writer {writer_id} has batches: {num_batches} ", flush=True) @@ -252,7 +262,8 @@ def do_merge_with_weights( for i, batch in enumerate(reader.stream(variables, skip_batches=skip_batches)): print( - f"Writer {writer_id} Combined batch {i} has {len(batch['jets'])} jets", flush=True + f"Writer {writer_id} Combined batch {i} has {len(batch[global_name])} objects", + flush=True, ) all_sample_weights = RWMerge.get_sample_weights(batch, weights) to_write = {} @@ -263,7 +274,14 @@ def do_merge_with_weights( to_write[key] = batch[key] if writer is None: shapes = {k: (None,) + v.shape[1:] for k, v in to_write.items()} - writer = H5Writer(output_file, dtypes, shapes, shuffle=True, compression="gzip") + writer = H5Writer( + output_file, + dtypes, + shapes, + shuffle=True, + compression="gzip", + global_objects_name=global_name, + ) for group, g_attrs in attrs.items(): for attr, value in g_attrs.items(): writer.add_attr(attr, value, group) @@ -286,7 +304,7 @@ def do_merge_with_weights( print(f"Writer {writer_id} batches complete - closing writer", flush=True) writer.close() message = ( - f"Writer {writer_id} finished writing {writer.num_written} jets in " + f"Writer {writer_id} finished writing {writer.num_written} objects in " f"{time.time() - start_time:.2f}s" ) print(message, flush=True) diff --git a/upp/stages/split_containers.py b/upp/stages/split_containers.py index efff1938..988c25fa 100644 --- a/upp/stages/split_containers.py +++ b/upp/stages/split_containers.py @@ -130,6 +130,7 @@ def split_file( ): if isinstance(input_file, str): input_file = Path(input_file) + global_name = self.config.global_name add_flavour_label = flavour_label_list is not None # All variables for test file all_variables = get_all_datasets(input_file) @@ -140,18 +141,25 @@ def split_file( ) print("parsed variables: ", parsed_variables, flush=True) start = time.time() - reader = H5Reader(input_file, batch_size=batch_size, shuffle=False) + reader = H5Reader( + input_file, + batch_size=batch_size, + shuffle=False, + global_objects_name=self.config.global_name, + ) if output_name is None: output_name = input_file.name - num_jets = reader.num_jets - num_batches = num_jets // batch_size + (1 if num_jets % batch_size != 0 else 0) - if num_jets == 0: - print(f"File {input_file} has no jets. Skipping it", flush=True) + num_global_objects = reader.num_global_objects + num_batches = num_global_objects // batch_size + ( + 1 if num_global_objects % batch_size != 0 else 0 + ) + if num_global_objects == 0: + print(f"File {input_file} has no objects. Skipping it", flush=True) return sample_components = [] writers_by_sample_components = {} cuts_by_sample_components = {} - print(f"Input file: {input_file} has {num_jets} jets", flush=True) + print(f"Input file: {input_file} has {num_global_objects} objects", flush=True) print(f"Running {num_batches} batches of size {batch_size}", flush=True) print(f"Output directory: {output_dir}", flush=True) @@ -165,7 +173,7 @@ def split_file( writers_by_sample_components[split] = H5Writer.from_file( input_file, - num_jets=None, + num_global_objects=None, dst=output_file, precision="half", full_precision_vars=fp_vars, @@ -173,6 +181,7 @@ def split_file( variables=all_variables if "test" in split else parsed_variables, compression="gzip", add_flavour_label=add_flavour_label, + global_objects_name=global_name, ) cuts_by_sample_components[split] = component_cuts print(f"Creating writer for {split} saved to {output_file}", flush=True) @@ -198,28 +207,29 @@ def split_file( for sample_component in sample_components: writer = writers_by_sample_components[sample_component] cuts = cuts_by_sample_components[sample_component] - sel_idx = cuts(batch["jets"]).idx + sel_idx = cuts(batch[global_name]).idx sel_batch = {k: v[sel_idx] for k, v in batch.items()} if add_flavour_label: this_flavour_label = flavour_label_by_component[sample_component] tfl_arr = ( - np.ones(sel_batch["jets"].shape[0], dtype=np.int32) * this_flavour_label + np.ones(sel_batch[global_name].shape[0], dtype=np.int32) + * this_flavour_label ) # I think this is only going to happen during tests, as the mock file has # flavour_label in, but we wont - if "flavour_label" in sel_batch["jets"].dtype.names: + if "flavour_label" in sel_batch[global_name].dtype.names: if i == 0: print( f"Warning: {sample_component} already has a flavour label. " "We will overwrite it now.", flush=True, ) - sel_batch["jets"]["flavour_label"] = tfl_arr + sel_batch[global_name]["flavour_label"] = tfl_arr else: # Get the - sel_batch["jets"] = rfn.append_fields( - sel_batch["jets"], + sel_batch[global_name] = rfn.append_fields( + sel_batch[global_name], "flavour_label", tfl_arr, usemask=False, @@ -247,7 +257,7 @@ def split_file( message += "\n" for sample_component, writer in writers_by_sample_components.items(): perc = 100 * (num_written_by_sample_components[sample_component] / sum_written) - message += f"{sample_component}: {writer.num_written} jets ({perc:.2f}%)\n" + message += f"{sample_component}: {writer.num_written} objects ({perc:.2f}%)\n" writer.close() @@ -255,9 +265,9 @@ def split_file( message += "-" * 20 + "\n" message += "Summary\n" message += "-" * 20 + "\n" - message += f"Total number of jets: {num_jets}\n" + message += f"Total number of objects: {num_global_objects}\n" message += f"Total number of batches: {num_batches}\n" - message += f"Total number of written jets: {sum_written}\n" + message += f"Total number of written objects: {sum_written}\n" message += f"Total time taken: {time_taken:.2f} seconds\n" message += "-" * 20 print(message, flush=True) @@ -279,9 +289,11 @@ def _make_tmp_vds(self, files: list[str] | str | Path) -> Generator[Path, None, create_virtual_file(str(tmp_dir / "*.h5"), tmp_out_path, overwrite=True) h5vds = H5Reader( tmp_out_path, + global_objects_name=self.config.global_name, ) print( - f"Created combined virtual dataset with {h5vds.num_jets} jets at {tmp_out_path}", + f"Created combined virtual dataset with {h5vds.num_global_objects} " + f"objects at {tmp_out_path}", flush=True, ) yield tmp_out_path @@ -367,13 +379,18 @@ def create_meta_data(self): files[split][flavour].append(str(file[0])) - num_jets = { - split: {flavour: H5Reader(files[split][flavour]).num_jets for flavour in files[split]} + num_global_objects = { + split: { + flavour: H5Reader( + files[split][flavour], global_objects_name=self.config.global_name + ).num_global_objects + for flavour in files[split] + } for split in files } metadata = { "files": files, - "num_jets": num_jets, + "num_global_objects": num_global_objects, } output_file = output_dir / "organised-components.yaml" diff --git a/upp/utils/check_input_samples.py b/upp/utils/check_input_samples.py index ba6230a3..b60fa98e 100644 --- a/upp/utils/check_input_samples.py +++ b/upp/utils/check_input_samples.py @@ -69,7 +69,7 @@ def check_within_factor( Raises ------ ValueError - If n_jets is zero for a sample + If n_global_objects is zero for a sample If the spread factor is too high If the geometric mean deviation is too high """ @@ -80,7 +80,7 @@ def check_within_factor( # Check that each sample is not zero for inner_key, inner_value in inner.items(): if inner_value == 0: - raise ValueError(f"Found zero jets in group {name} / sample {inner_key}!") + raise ValueError(f"Found zero objects in group {name} / sample {inner_key}!") inner_values = {k: v for k, v in inner.items() if v > 0} @@ -112,7 +112,7 @@ def run_input_sample_check( deviation_factor: float, verbose: bool = True, ) -> None: - """Run the input sample checks on the number of the jets. + """Run the input sample checks on the number of the objects. Parameters ---------- @@ -158,7 +158,7 @@ def run_input_sample_check( sample_type_dict[config_blocks] = {} sample_type_dict[config_blocks]["pattern"] = config.config[config_blocks]["pattern"] - # Setup H5Reader for the different samples to read the total number of jets + # Setup H5Reader for the different samples to read the total number of objects for sample_type, sample_list in sample_type_dict.items(): # Log the status if verbose: @@ -206,13 +206,13 @@ def run_input_sample_check( else: entry_name = sample - # Create the H5 reader for each sample and read the number of jets from it + # Create the H5 reader for each sample and read the number of objects from it sample_list[entry_name] = H5Reader( fname=config.ntuple_dir / sample, batch_size=config.batch_size, - jets_name=config.jets_name, + global_objects_name=config.global_name, vds_dir=config.vds_dir, - ).num_jets + ).num_global_objects # Drop the pattern del sample_list["pattern"] @@ -220,15 +220,15 @@ def run_input_sample_check( # Check that all the samples per group are not too different check_within_factor(groups=sample_type_dict, factor=deviation_factor) - # Printing the dict with all the number of jets to the terminal, if wanted + # Printing the dict with all the number of objects to the terminal, if wanted if verbose: - log.info("Available jets in given groups:\n") + log.info("Available objects in given groups:\n") for sample_type, sample_dict in sample_type_dict.items(): log.info(f"Group: {sample_type}") - for entry_name, n_jets in sample_dict.items(): - log.info(f" - Sample: {entry_name}, N_Jets: {n_jets:,}") + for entry_name, n_global_objects in sample_dict.items(): + log.info(f" - Sample: {entry_name}, N_Objects: {n_global_objects:,}") def main(args: Any | None = None) -> None: