From bdbddab92d2cbd2f2aae56265a270f971387a1d7 Mon Sep 17 00:00:00 2001 From: yuanda Date: Wed, 27 May 2026 15:46:35 +0200 Subject: [PATCH 01/11] metadata injector reweighting workflow --- upp/classes/preprocessing_config.py | 11 +++- upp/configs/flavours-r10.yaml | 44 ++++++++++++++ upp/configs/metadata.yaml | 88 +++++++++++++++++++++++++++ upp/configs/variables.yaml | 6 +- upp/main.py | 17 ++++++ upp/stages/hist.py | 8 ++- upp/stages/metadata_injector.py | 92 +++++++++++++++++++++++++++++ upp/stages/normalisation.py | 12 +++- upp/stages/reweight.py | 42 ++++++++----- upp/stages/rw_merge.py | 24 ++++++-- upp/stages/split_containers.py | 27 +++++++-- 11 files changed, 340 insertions(+), 31 deletions(-) create mode 100644 upp/configs/flavours-r10.yaml create mode 100644 upp/configs/metadata.yaml create mode 100644 upp/stages/metadata_injector.py diff --git a/upp/classes/preprocessing_config.py b/upp/classes/preprocessing_config.py index 2c279f08..668f3e5f 100644 --- a/upp/classes/preprocessing_config.py +++ b/upp/classes/preprocessing_config.py @@ -119,6 +119,8 @@ class PreprocessingConfig: Skip checks for the input files. This is used for grid submission skip_config_copy : bool, optional Decide, if the config copying is skipped or not. By default False + keep_all_variables : bool, optional + If true, split and rw-merge keep all top-level HDF5 datasets / full jets fields (PW + RW on jets). """ config_path: Path @@ -142,6 +144,8 @@ class PreprocessingConfig: num_jets_per_output_file: int | None = None skip_checks: bool = False skip_config_copy: bool = False + # When true: split-containers and rw-merge keep all input datasets/fields (tracks, towers, …), not just variables.yaml subset. + keep_all_variables: bool = False def __post_init__(self): # postprocess paths @@ -198,9 +202,12 @@ def __post_init__(self): if selection := groups.get("selection", None): selectors[name] = TrackSelector(Cuts.from_list(selection)) - # configure variables + # configure variables (keep_all_variables is independent of train/val/test split) self.variables = VariableConfig( - self.config["variables"], self.jets_name, self.is_test, selectors + self.config["variables"], + self.jets_name, + self.keep_all_variables, + selectors, ) if self.sampl_cfg is not None: self.variables = self.variables.add_jet_vars( diff --git a/upp/configs/flavours-r10.yaml b/upp/configs/flavours-r10.yaml new file mode 100644 index 00000000..35e6d048 --- /dev/null +++ b/upp/configs/flavours-r10.yaml @@ -0,0 +1,44 @@ +# flavour-r10.yaml +# 已同步为 0-6 连续整数标签 + +- name: djets + category: jets + label: "d-jets" + colour: blue + cuts: + - [PartonTruthLabelID, "==", 1] + +- name: ujets + category: jets + label: "u-jets" + colour: cyan + cuts: + - [PartonTruthLabelID, "==", 2] + +- name: sjets + category: jets + label: "s-jets" + colour: purple + cuts: + - [PartonTruthLabelID, "==", 3] + +- name: cjets + category: jets + label: "c-jets" + colour: red + cuts: + - [PartonTruthLabelID, "==", 4] + +- name: bjets + category: jets + label: "b-jets" + colour: orange + cuts: + - [PartonTruthLabelID, "==", 5] + +- name: gluon + category: jets + label: "gluon" + colour: pink + cuts: + - [PartonTruthLabelID, "==", 21] \ No newline at end of file diff --git a/upp/configs/metadata.yaml b/upp/configs/metadata.yaml new file mode 100644 index 00000000..469923c9 --- /dev/null +++ b/upp/configs/metadata.yaml @@ -0,0 +1,88 @@ +global: + base_dir: /eos/home-y/yuanda/dev + # Metadata injection (--metadata) modifies files in-place; use copies under central_working, not origin. + ntuple_dir: /eos/home-y/yuanda/dev/sample/central_working + # When true: split + rw-merge preserve all top-level HDF5 datasets / full jet fields (PW + RW columns still added on jets). + keep_all_variables: true + out_dir: output_metadata + jets_name: jets + batch_size: 40000 + num_jets_estimate: 700000 + num_jets_per_output_file: 1_000_000 + # test split: plot.py uses merged VDS path (glob) instead of per-sample pp_output_test_*.h5 + merge_test_samples: true + flavour_config: /eos/home-y/yuanda/dev/umami-preprocessing/upp/configs/flavours-r10.yaml + +inputs: + train: + input_files: + - /eos/home-y/yuanda/dev/sample/central_working/*.h5 + +global_cuts: + common: [] + train: [[eventNumber, "%10<=", 7]] + val: [[eventNumber, "%10==", 8]] + test: [[eventNumber, "%10==", 9]] + +sample_all: &sample_all + name: sample_all + pattern: ["*.h5"] + +components: + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 1]] } + sample: { <<: *sample_all } + flavours: [djets] + num_jets: 60000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 2]] } + sample: { <<: *sample_all } + flavours: [ujets] + num_jets: 120000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 3]] } + sample: { <<: *sample_all } + flavours: [sjets] + num_jets: 15000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 4]] } + sample: { <<: *sample_all } + flavours: [cjets] + num_jets: 10000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 5]] } + sample: { <<: *sample_all } + flavours: [bjets] + num_jets: 6000 + - region: { name: inclusive, cuts: [[PartonTruthLabelID, "==", 21]] } + sample: { <<: *sample_all } + flavours: [gluon] + num_jets: 160000 + +variables: + jets: + inputs: [pt, eta, eventNumber] + labels: [PartonTruthLabelID, physicalWeight] + +resampling: + method: none + target: ujets + variables: + pt: + bins: [[20000, 2480000, 100]] + +reweighting: + num_jets_estimate: 700000 + reweights: + - group: jets + reweight_vars: [pt] + bins: + pt: [[20000, 2480000, 100]] + class_var: flavour_label + class_target: uniform + +flavour_categories: + jets: [djets, ujets, sjets, cjets, bjets, gluon] + +custom_flavours: + djets: { label: 0, cuts: [[PartonTruthLabelID, "==", 1]] } + ujets: { label: 1, cuts: [[PartonTruthLabelID, "==", 2]] } + sjets: { label: 2, cuts: [[PartonTruthLabelID, "==", 3]] } + cjets: { label: 3, cuts: [[PartonTruthLabelID, "==", 4]] } + bjets: { label: 4, cuts: [[PartonTruthLabelID, "==", 5]] } + gluon: { label: 5, cuts: [[PartonTruthLabelID, "==", 21]] } \ No newline at end of file diff --git a/upp/configs/variables.yaml b/upp/configs/variables.yaml index 1870c74e..25613a29 100644 --- a/upp/configs/variables.yaml +++ b/upp/configs/variables.yaml @@ -2,6 +2,9 @@ jets: inputs: - pt_btagJes - eta_btagJes + - R10TruthLabel_R22v1 + - PartonTruthLabelID + - physicalWeight labels: - HadronConeExclTruthLabelID - HadronConeExclExtendedTruthLabelID @@ -20,7 +23,8 @@ jets: - n_truth_promptLepton - eventNumber - jetFoldHash - + - physicalWeight + tracks: inputs: - d0 diff --git a/upp/main.py b/upp/main.py index cf14a790..b910b02f 100644 --- a/upp/main.py +++ b/upp/main.py @@ -29,6 +29,10 @@ from upp.utils.check_input_samples import run_input_sample_check from upp.utils.logger import setup_logger +from ftag.find_metadata import MetadataFinder +from upp.stages.metadata_injector import MetadataInjector + + def parse_args(args: Any) -> argparse.Namespace: """Parse the command line arguments. @@ -53,6 +57,12 @@ def parse_args(args: Any) -> argparse.Namespace: type=valid_path, help="Path to config file", ) + parser.add_argument( + "--metadata", + action="store_true", + default=False, + help="Run metadata injection stage before reweighting", + ) parser.add_argument( "--prep", action="store_true", @@ -206,6 +216,12 @@ def run_pp(args: argparse.Namespace) -> None: # load config config = PreprocessingConfig.from_file(args.config, args.split, skip_checks=args.grid) + # run metadata injection + if getattr(args, "metadata", False): + log.info("Running metadata injection...") + injector = MetadataInjector(config) + injector.run() + if args.split_components: log.info("Splitting containers...") split = SplitContainers(args.config) @@ -217,6 +233,7 @@ def run_pp(args: argparse.Namespace) -> None: # If we aren't running on the grid, we create the metadata after splitting if not args.grid: split.create_meta_data() + if args.reweight: log.info("Running reweighting...") reweight = Reweight(config) diff --git a/upp/stages/hist.py b/upp/stages/hist.py index 8f49238f..940424dd 100644 --- a/upp/stages/hist.py +++ b/upp/stages/hist.py @@ -18,7 +18,7 @@ from upp.classes.preprocessing_config import PreprocessingConfig -def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: +def bin_jets(array: dict, bins: list, weights: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]: """Create the histogram and bins for the given resampling variables. Parameters @@ -41,10 +41,12 @@ def bin_jets(array: dict, bins: list) -> tuple[np.ndarray, np.ndarray]: bin in which this observation falls. The representation depends on the `expand_binnumbers` argument. See `Notes` for details. """ + statistic_mode = "count" if weights is None else "sum" + hist, _, out_bins = binned_statistic_dd( sample=s2u(array), - values=None, - statistic="count", + values=weights, + statistic=statistic_mode, bins=bins, expand_binnumbers=True, ) diff --git a/upp/stages/metadata_injector.py b/upp/stages/metadata_injector.py new file mode 100644 index 00000000..a7e7bc9d --- /dev/null +++ b/upp/stages/metadata_injector.py @@ -0,0 +1,92 @@ +from ftag.find_metadata import MetadataFinder +from upp.utils.logger import setup_logger +import glob +import shutil +from pathlib import Path +import h5py +import numpy as np +from numpy.lib import recfunctions as rfn + +class MetadataInjector: + def __init__(self, config): + self.config = config + self.log = setup_logger() + + def run(self): + # Initialize input file list + input_files = [] + if hasattr(self.config, "config") and "inputs" in self.config.config: + raw_inputs = self.config.config["inputs"] + if "train" in raw_inputs and "input_files" in raw_inputs["train"]: + input_files.extend(raw_inputs["train"]["input_files"]) + + # Expand wildcards/glob patterns + expanded_files = [] + for f in input_files: + matched = glob.glob(f) + expanded_files.extend(matched) + + for fpath_str in expanded_files: + fpath = Path(fpath_str) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + try: + # 0. Create a physical backup + shutil.copy(fpath, backup_path) + + # 1. Inject metadata (this adds a new 'metadata' group) + finder = MetadataFinder(fpath_str) + finder.inject_metadata() + + with h5py.File(fpath_str, "a") as f: + # 2. Calculate parameters required for physical weights + dsid = list(f["metadata"].keys())[0] + xs = float(f[f"metadata/{dsid}/cross_section_pb"][()]) + eff = float(f[f"metadata/{dsid}/genFiltEff"][()]) + kf = float(f[f"metadata/{dsid}/kfactor"][()]) + + # Read Sum of Weights (SOW) with support for different formats + sow_ds = f["cutBookkeeper/nominal/counts"] + sow = float(sow_ds["sumOfWeights"][0]) if sow_ds.dtype.names else float(sow_ds[0]) + + # 3. Read original jets and their attributes + old_jets_ds = f["jets"] + original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata + jets_data = old_jets_ds[:] + + # 4. Calculate Physical Weight + if "mcEventWeight" not in jets_data.dtype.names: + raise KeyError(f"mcEventWeight missing in {fpath_str}") + + mcw = jets_data["mcEventWeight"].astype('f8') + # Formula: Weight = (XS * Efficiency * k-factor / SOW) * MC_Weight + physical_w = (xs * eff * kf / sow) * mcw + + # 5. Construct the updated structured array + if "physicalWeight" in jets_data.dtype.names: + jets_data = rfn.drop_fields(jets_data, "physicalWeight") + + updated_jets = rfn.append_fields( + jets_data, + "physicalWeight", + physical_w.astype('f4'), + usemask=False + ) + + # 6. Delete and recreate the dataset while restoring attributes + del f["jets"] + new_ds = f.create_dataset("jets", data=updated_jets, compression="gzip") + + # Restore all original attributes (e.g., descriptions for eventNumber, etc.) + for k, v in original_attrs.items(): + new_ds.attrs[k] = v + + # 7. Processing successful: remove the backup + backup_path.unlink() + self.log.info(f"Successfully updated {fpath.name} (Attributes and original fields preserved)") + + except Exception as e: + self.log.error(f"Failed for {fpath_str}: {e}") + if backup_path.exists(): + self.log.warning(f"Restoring {fpath.name} from backup...") + shutil.move(backup_path, fpath) \ No newline at end of file diff --git a/upp/stages/normalisation.py b/upp/stages/normalisation.py index 95bb5c45..75194473 100644 --- a/upp/stages/normalisation.py +++ b/upp/stages/normalisation.py @@ -255,10 +255,18 @@ def run(self): norm_dict = None class_dict = None total = None - vars = self.variables.combined() + combined_partial = self.variables.combined() + vars: dict[str, list[str] | None] = {} + for name in self.variables: + cols = combined_partial[name] + if cols is None: + # keep_all_variables: streaming needs explicit columns; still norm only declared inputs/labels + cols = self.variables[name]["inputs"] + self.variables[name].get("labels", []) + vars[name] = list(cols) 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") + if "flavour_label" not in vars[self.jets_name]: + vars[self.jets_name].append("flavour_label") stream = reader.stream(vars, self.num_jets) with ProgressBar() as progress: diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index 0be0e72d..9c91571c 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -116,6 +116,8 @@ def calculate_weights( all_vars[rw_group].extend(rw.reweight_vars) if "valid" in existing_vars[rw_group]: all_vars[rw_group] += ["valid"] + if "physicalWeight" in existing_vars.get(rw_group, []): + all_vars[rw_group] += ["physicalWeight"] if "jets" not in all_vars: all_vars["jets"] = ["pt"] all_vars = {k: list(set(v)) for k, v in all_vars.items()} @@ -171,10 +173,25 @@ def calculate_weights( assert "valid" in data.dtype.names data = data[data["valid"]] classes = np.unique(data[rw.class_var]) if rw.class_var is not None else [None] - for cls in classes: mask = data[rw.class_var] == cls - hist, outbins = bin_jets(data[mask][rw.reweight_vars], rw.flat_bins) + data_masked = data[mask] + + # Use physicalWeight column from the structured jets array if present; + # otherwise fall back to uniform weights. + if "physicalWeight" in data_masked.dtype.names: + PW_CAP = 10000.0 + w = np.asarray(data_masked["physicalWeight"], dtype=np.float64) + w = np.clip(w, 0, PW_CAP) + else: + w = np.ones(mask.sum(), dtype=float) + + hist, outbins = bin_jets( + data_masked[rw.reweight_vars], + rw.flat_bins, + weights=w, + ) + if rw.class_var is not None: cls = str(cls) if rw_group not in all_histograms: @@ -270,21 +287,16 @@ def calculate_weights( "rw_vars": rw.reweight_vars, "class_var": rw.class_var, } - idx_below_min = None + RW_CAP = 1e4 + HIST_FLOOR = 1e-6 + target = all_targets[rw_group][rw_rep] for cls, hist in all_histograms[rw_group][rw_rep]["histograms"].items(): - this_idx_below_min = hist == 0 # | (all_targets[rw_group][rw_rep] == 0) - output_weights[rw_group][rw_rep]["weights"][cls] = np.where( - hist > 0, all_targets[rw_group][rw_rep] / hist, 0 + rw = np.where( + hist > HIST_FLOOR, + target / np.maximum(hist, HIST_FLOOR), + 0.0, ) - if idx_below_min is None: - idx_below_min = this_idx_below_min - else: - idx_below_min |= this_idx_below_min - # If we have any bins where we have 0 of a given flavour, we set all the - # weights to 0 - if np.any(idx_below_min): - for cls in all_histograms[rw_group][rw_rep]["histograms"]: - output_weights[rw_group][rw_rep]["weights"][cls][idx_below_min] = 0 + output_weights[rw_group][rw_rep]["weights"][cls] = np.clip(rw, 0, RW_CAP) return output_weights diff --git a/upp/stages/rw_merge.py b/upp/stages/rw_merge.py index 31a7efbe..276f8b7c 100644 --- a/upp/stages/rw_merge.py +++ b/upp/stages/rw_merge.py @@ -5,6 +5,7 @@ from multiprocessing import Pool from pathlib import Path +import h5py import numpy as np import yaml from ftag.hdf5 import H5Reader, H5Writer, join_structured_arrays @@ -76,12 +77,27 @@ def run(self): batches_per_file = num_jets_per_file // batch_size or 1 num_batches = ( - total_jets // batch_size + (1 if total_jets % num_jets_per_file != 0 else 0) + total_jets // batch_size + (1 if total_jets % batch_size != 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"] + # H5Reader.stream(None) only loads jets; for full-ntuple merge pass {ds: None} per top-level dataset. + if self.config.variables.keep_all: + p0 = Path(all_files[0]) + with h5py.File(p0, "r") as hf: + variables = {k: None for k in hf if isinstance(hf[k], h5py.Dataset)} + else: + variables = ( + self.config.variables.combined() + if self.config.split != "test" + else None + ) + if ( + variables is not None + and not self.config.variables.keep_all + and isinstance(variables.get(self.config.jets_name), list) + and "flavour_label" not in variables[self.config.jets_name] + ): + variables[self.config.jets_name] += ["flavour_label"] args_list = [] for i, bi in enumerate(range(0, num_batches, batches_per_file)): args_list.append( diff --git a/upp/stages/split_containers.py b/upp/stages/split_containers.py index c9b08ab4..d791ce03 100644 --- a/upp/stages/split_containers.py +++ b/upp/stages/split_containers.py @@ -86,10 +86,12 @@ def get_all_fp_vars(file: Path | str) -> list[str]: all_vars = get_all_vars(file) # combine the values in this dict into a single list + target_keywords = ["pt", "energy", "mass", "weight", "physicalWeight"] + fp_vars = [ - v for v in all_vars if ("pt" in v.lower() or "energy" in v.lower() or "mass" in v.lower()) + v for v in all_vars + if any(key in v.lower() for key in target_keywords) ] - return fp_vars @@ -127,6 +129,7 @@ def split_file( output_name=None, variables: dict[str, dict[str, list[str]]] | None = None, flavour_label_list: list[str] | None = None, + keep_all_variables: bool = False, ): if isinstance(input_file, str): input_file = Path(input_file) @@ -138,6 +141,16 @@ def split_file( parsed_variables: dict[str, list[str]] | dict[str, None] = ( parse_variables(variables) if variables is not None else all_variables ) + # --- Modified patch logic as follows --- + if isinstance(parsed_variables, dict): + # Only process if the 'jets' group exists + if "jets" in parsed_variables: + # Force-add 'physicalWeight' only; do not add 'valid' + if "physicalWeight" not in parsed_variables["jets"]: + parsed_variables["jets"].append("physicalWeight") + print("FORCE RE-INTEGRATED: Only physicalWeight into jets", flush=True) + # ----------------------- + print("parsed variables: ", parsed_variables, flush=True) start = time.time() reader = H5Reader(input_file, batch_size=batch_size, shuffle=False) @@ -163,6 +176,7 @@ def split_file( print(f"At least 1 output file exists for {input_file}. Skipping it", flush=True) return + use_all_cols = keep_all_variables or ("test" in split) writers_by_sample_components[split] = H5Writer.from_file( input_file, num_jets=None, @@ -170,7 +184,7 @@ def split_file( precision="half", full_precision_vars=fp_vars, shuffle=False, - variables=all_variables if "test" in split else parsed_variables, + variables=all_variables if use_all_cols else parsed_variables, compression="gzip", add_flavour_label=add_flavour_label, ) @@ -308,8 +322,12 @@ def run( assert container is not None, "Can only specify files if a container is specified" for container, cuts_by_component in containers_with_split_cuts.items(): + # Sanitize container name when used as a directory: avoid literal '*' or '/' + # in the split-components subdir, otherwise downstream H5Reader will treat the + # path as a glob pattern and create a shared VDS that collides across flavours. + container_dir_name = container.replace("*", "all").replace("/", "_") or "default" this_out_dir = ( - Path(self.config.base_dir) / "split-components" / container + Path(self.config.base_dir) / "split-components" / container_dir_name if output_dir is None else Path(".") ) @@ -340,6 +358,7 @@ def run( variables=self.config.config["variables"], flavour_label_list=all_flavours, output_name="output", + keep_all_variables=self.config.keep_all_variables, ) def create_meta_data(self): From 94b4b94d5283001ec630e7dd2834ff23b8d4a9d7 Mon Sep 17 00:00:00 2001 From: yuanda Date: Wed, 27 May 2026 16:30:52 +0200 Subject: [PATCH 02/11] precommit fix --- upp/classes/preprocessing_config.py | 6 ++-- upp/configs/flavours-r10.yaml | 2 +- upp/configs/metadata.yaml | 12 ++++---- upp/configs/single-b-upgrade.yaml | 4 --- upp/configs/variables.yaml | 2 +- upp/main.py | 13 ++++----- upp/stages/hist.py | 7 ++++- upp/stages/metadata_injector.py | 45 ++++++++++++++++------------- upp/stages/normalisation.py | 10 ++++--- upp/stages/rw_merge.py | 13 +++------ upp/stages/split_containers.py | 19 +++++------- 11 files changed, 65 insertions(+), 68 deletions(-) diff --git a/upp/classes/preprocessing_config.py b/upp/classes/preprocessing_config.py index 668f3e5f..7beae5bd 100644 --- a/upp/classes/preprocessing_config.py +++ b/upp/classes/preprocessing_config.py @@ -120,7 +120,8 @@ class PreprocessingConfig: skip_config_copy : bool, optional Decide, if the config copying is skipped or not. By default False keep_all_variables : bool, optional - If true, split and rw-merge keep all top-level HDF5 datasets / full jets fields (PW + RW on jets). + When true: split-containers and rw-merge keep full jets fields and all top-level + HDF5 datasets (PW and RW columns still added on jets). """ config_path: Path @@ -144,7 +145,8 @@ class PreprocessingConfig: num_jets_per_output_file: int | None = None skip_checks: bool = False skip_config_copy: bool = False - # When true: split-containers and rw-merge keep all input datasets/fields (tracks, towers, …), not just variables.yaml subset. + # When true: keep all input top-level datasets (tracks, towers, …), not only the + # variables.yaml subset, through split-containers and rw-merge. keep_all_variables: bool = False def __post_init__(self): diff --git a/upp/configs/flavours-r10.yaml b/upp/configs/flavours-r10.yaml index 35e6d048..13d142ad 100644 --- a/upp/configs/flavours-r10.yaml +++ b/upp/configs/flavours-r10.yaml @@ -41,4 +41,4 @@ label: "gluon" colour: pink cuts: - - [PartonTruthLabelID, "==", 21] \ No newline at end of file + - [PartonTruthLabelID, "==", 21] diff --git a/upp/configs/metadata.yaml b/upp/configs/metadata.yaml index 469923c9..b51314e0 100644 --- a/upp/configs/metadata.yaml +++ b/upp/configs/metadata.yaml @@ -80,9 +80,9 @@ flavour_categories: jets: [djets, ujets, sjets, cjets, bjets, gluon] custom_flavours: - djets: { label: 0, cuts: [[PartonTruthLabelID, "==", 1]] } - ujets: { label: 1, cuts: [[PartonTruthLabelID, "==", 2]] } - sjets: { label: 2, cuts: [[PartonTruthLabelID, "==", 3]] } - cjets: { label: 3, cuts: [[PartonTruthLabelID, "==", 4]] } - bjets: { label: 4, cuts: [[PartonTruthLabelID, "==", 5]] } - gluon: { label: 5, cuts: [[PartonTruthLabelID, "==", 21]] } \ No newline at end of file + djets: { label: 0, cuts: [[PartonTruthLabelID, "==", 1]] } + ujets: { label: 1, cuts: [[PartonTruthLabelID, "==", 2]] } + sjets: { label: 2, cuts: [[PartonTruthLabelID, "==", 3]] } + cjets: { label: 3, cuts: [[PartonTruthLabelID, "==", 4]] } + bjets: { label: 4, cuts: [[PartonTruthLabelID, "==", 5]] } + gluon: { label: 5, cuts: [[PartonTruthLabelID, "==", 21]] } diff --git a/upp/configs/single-b-upgrade.yaml b/upp/configs/single-b-upgrade.yaml index 48648905..3a6a659c 100644 --- a/upp/configs/single-b-upgrade.yaml +++ b/upp/configs/single-b-upgrade.yaml @@ -18,8 +18,6 @@ zprime: &zprime 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" - - lowpt: &lowpt name: lowpt cuts: @@ -74,8 +72,6 @@ components: flavours: [taujets] num_jets: 805_555 - - resampling: target: cjets method: pdf #countup diff --git a/upp/configs/variables.yaml b/upp/configs/variables.yaml index 25613a29..be6d21e9 100644 --- a/upp/configs/variables.yaml +++ b/upp/configs/variables.yaml @@ -24,7 +24,7 @@ jets: - eventNumber - jetFoldHash - physicalWeight - + tracks: inputs: - d0 diff --git a/upp/main.py b/upp/main.py index 5623d2bd..028e17e9 100644 --- a/upp/main.py +++ b/upp/main.py @@ -20,6 +20,7 @@ from upp.classes.preprocessing_config import PreprocessingConfig from upp.stages.hist import create_histograms from upp.stages.merging import Merging +from upp.stages.metadata_injector import MetadataInjector from upp.stages.normalisation import Normalisation from upp.stages.plot import plot_resampling_dists from upp.stages.resampling import Resampling @@ -29,10 +30,6 @@ from upp.utils.check_input_samples import run_input_sample_check from upp.utils.logger import setup_logger -from ftag.find_metadata import MetadataFinder -from upp.stages.metadata_injector import MetadataInjector - - def parse_args(args: Any) -> argparse.Namespace: """Parse the command line arguments. @@ -58,10 +55,10 @@ def parse_args(args: Any) -> argparse.Namespace: help="Path to config file", ) parser.add_argument( - "--metadata", - action="store_true", - default=False, - help="Run metadata injection stage before reweighting", + "--metadata", + action="store_true", + default=False, + help="Run metadata injection stage before reweighting", ) parser.add_argument( "--prep", diff --git a/upp/stages/hist.py b/upp/stages/hist.py index 5f9fb189..48bc4369 100644 --- a/upp/stages/hist.py +++ b/upp/stages/hist.py @@ -18,7 +18,9 @@ from upp.classes.preprocessing_config import PreprocessingConfig -def bin_jets(array: dict, bins: list, weights: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]: +def bin_jets( + array: dict, bins: list, weights: np.ndarray | None = None +) -> tuple[np.ndarray, np.ndarray]: """Create the histogram and bins for the given resampling variables. Parameters @@ -28,6 +30,9 @@ def bin_jets(array: dict, bins: list, weights: np.ndarray | None = None) -> tupl variables. bins : list Flat list with the bins which are to be used. + weights : np.ndarray | None, optional + Per-jet weights for the histogram (e.g. physicalWeight). When None, bins are filled + by count; otherwise by sum of weights. Returns ------- diff --git a/upp/stages/metadata_injector.py b/upp/stages/metadata_injector.py index a7e7bc9d..fb7a338f 100644 --- a/upp/stages/metadata_injector.py +++ b/upp/stages/metadata_injector.py @@ -1,12 +1,16 @@ -from ftag.find_metadata import MetadataFinder -from upp.utils.logger import setup_logger +from __future__ import annotations + import glob import shutil from pathlib import Path + import h5py -import numpy as np +from ftag.find_metadata import MetadataFinder from numpy.lib import recfunctions as rfn +from upp.utils.logger import setup_logger + + class MetadataInjector: def __init__(self, config): self.config = config @@ -29,7 +33,7 @@ def run(self): for fpath_str in expanded_files: fpath = Path(fpath_str) backup_path = fpath.with_suffix(fpath.suffix + ".bak") - + try: # 0. Create a physical backup shutil.copy(fpath, backup_path) @@ -40,53 +44,54 @@ def run(self): with h5py.File(fpath_str, "a") as f: # 2. Calculate parameters required for physical weights - dsid = list(f["metadata"].keys())[0] + dsid = next(iter(f["metadata"].keys())) xs = float(f[f"metadata/{dsid}/cross_section_pb"][()]) eff = float(f[f"metadata/{dsid}/genFiltEff"][()]) kf = float(f[f"metadata/{dsid}/kfactor"][()]) - + # Read Sum of Weights (SOW) with support for different formats sow_ds = f["cutBookkeeper/nominal/counts"] - sow = float(sow_ds["sumOfWeights"][0]) if sow_ds.dtype.names else float(sow_ds[0]) - + sow = ( + float(sow_ds["sumOfWeights"][0]) if sow_ds.dtype.names else float(sow_ds[0]) + ) + # 3. Read original jets and their attributes old_jets_ds = f["jets"] - original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata + original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata jets_data = old_jets_ds[:] - + # 4. Calculate Physical Weight if "mcEventWeight" not in jets_data.dtype.names: raise KeyError(f"mcEventWeight missing in {fpath_str}") - - mcw = jets_data["mcEventWeight"].astype('f8') + + mcw = jets_data["mcEventWeight"].astype("f8") # Formula: Weight = (XS * Efficiency * k-factor / SOW) * MC_Weight physical_w = (xs * eff * kf / sow) * mcw # 5. Construct the updated structured array if "physicalWeight" in jets_data.dtype.names: jets_data = rfn.drop_fields(jets_data, "physicalWeight") - + updated_jets = rfn.append_fields( - jets_data, - "physicalWeight", - physical_w.astype('f4'), - usemask=False + jets_data, "physicalWeight", physical_w.astype("f4"), usemask=False ) # 6. Delete and recreate the dataset while restoring attributes del f["jets"] new_ds = f.create_dataset("jets", data=updated_jets, compression="gzip") - + # Restore all original attributes (e.g., descriptions for eventNumber, etc.) for k, v in original_attrs.items(): new_ds.attrs[k] = v # 7. Processing successful: remove the backup backup_path.unlink() - self.log.info(f"Successfully updated {fpath.name} (Attributes and original fields preserved)") + self.log.info( + f"Successfully updated {fpath.name} (Attributes and original fields preserved)" + ) except Exception as e: self.log.error(f"Failed for {fpath_str}: {e}") if backup_path.exists(): self.log.warning(f"Restoring {fpath.name} from backup...") - shutil.move(backup_path, fpath) \ No newline at end of file + shutil.move(backup_path, fpath) diff --git a/upp/stages/normalisation.py b/upp/stages/normalisation.py index 75194473..f551f24c 100644 --- a/upp/stages/normalisation.py +++ b/upp/stages/normalisation.py @@ -260,13 +260,15 @@ def run(self): for name in self.variables: cols = combined_partial[name] if cols is None: - # keep_all_variables: streaming needs explicit columns; still norm only declared inputs/labels + # keep_all_variables: explicit stream columns; norm still on yaml inputs/labels cols = self.variables[name]["inputs"] + self.variables[name].get("labels", []) vars[name] = list(cols) with h5py.File(reader.files[0]) as f: - if "flavour_label" in f[self.jets_name].dtype.names: - if "flavour_label" not in vars[self.jets_name]: - vars[self.jets_name].append("flavour_label") + if ( + "flavour_label" in f[self.jets_name].dtype.names + and "flavour_label" not in vars[self.jets_name] + ): + vars[self.jets_name].append("flavour_label") stream = reader.stream(vars, self.num_jets) with ProgressBar() as progress: diff --git a/upp/stages/rw_merge.py b/upp/stages/rw_merge.py index 276f8b7c..1a897c14 100644 --- a/upp/stages/rw_merge.py +++ b/upp/stages/rw_merge.py @@ -76,21 +76,16 @@ def run(self): num_jets_per_file = self.config.num_jets_per_output_file or total_jets batches_per_file = num_jets_per_file // batch_size or 1 - num_batches = ( - total_jets // batch_size + (1 if total_jets % batch_size != 0 else 0) - ) or 1 + num_batches = (total_jets // batch_size + (1 if total_jets % batch_size != 0 else 0)) or 1 - # H5Reader.stream(None) only loads jets; for full-ntuple merge pass {ds: None} per top-level dataset. + # stream(None) only loads jets; for full ntuple merge use {dataset: None} per + # top-level HDF5 dataset. if self.config.variables.keep_all: p0 = Path(all_files[0]) with h5py.File(p0, "r") as hf: variables = {k: None for k in hf if isinstance(hf[k], h5py.Dataset)} else: - variables = ( - self.config.variables.combined() - if self.config.split != "test" - else None - ) + variables = self.config.variables.combined() if self.config.split != "test" else None if ( variables is not None and not self.config.variables.keep_all diff --git a/upp/stages/split_containers.py b/upp/stages/split_containers.py index d791ce03..ede25e1d 100644 --- a/upp/stages/split_containers.py +++ b/upp/stages/split_containers.py @@ -88,10 +88,7 @@ def get_all_fp_vars(file: Path | str) -> list[str]: target_keywords = ["pt", "energy", "mass", "weight", "physicalWeight"] - fp_vars = [ - v for v in all_vars - if any(key in v.lower() for key in target_keywords) - ] + fp_vars = [v for v in all_vars if any(key in v.lower() for key in target_keywords)] return fp_vars @@ -141,14 +138,12 @@ def split_file( parsed_variables: dict[str, list[str]] | dict[str, None] = ( parse_variables(variables) if variables is not None else all_variables ) - # --- Modified patch logic as follows --- + # Ensure physicalWeight is requested for jets when using a column list (not None). if isinstance(parsed_variables, dict): - # Only process if the 'jets' group exists - if "jets" in parsed_variables: - # Force-add 'physicalWeight' only; do not add 'valid' - if "physicalWeight" not in parsed_variables["jets"]: - parsed_variables["jets"].append("physicalWeight") - print("FORCE RE-INTEGRATED: Only physicalWeight into jets", flush=True) + jets_cols = parsed_variables.get("jets") + if isinstance(jets_cols, list) and "physicalWeight" not in jets_cols: + jets_cols.append("physicalWeight") + print("FORCE RE-INTEGRATED: Only physicalWeight into jets", flush=True) # ----------------------- print("parsed variables: ", parsed_variables, flush=True) @@ -203,7 +198,7 @@ def split_file( assert all( len(_flavour_label_by_component[component]) == 1 for component in sample_components ), f"Each component must have exactly 1 flavour label not {_flavour_label_by_component}" - flavour_label_by_component: dict[str, int] = { # noqa: no-redef + flavour_label_by_component: dict[str, int] = { component: _flavour_label_by_component[component][0] for component in _flavour_label_by_component } From 76bb44ca5835916fd48dada30a0dfb51d2860724 Mon Sep 17 00:00:00 2001 From: yuanda Date: Thu, 28 May 2026 13:49:18 +0200 Subject: [PATCH 03/11] test logic added --- tests/integration/test_run.py | 11 +++++ tests/integration/test_run_rw.py | 21 +++++++- tests/unit/stages/test_metadata_injector.py | 55 +++++++++++++++++++++ tests/unit/test_main.py | 12 +++++ 4 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 tests/unit/stages/test_metadata_injector.py diff --git a/tests/integration/test_run.py b/tests/integration/test_run.py index b485588d..2f9cb8f2 100644 --- a/tests/integration/test_run.py +++ b/tests/integration/test_run.py @@ -8,6 +8,7 @@ import numpy as np import pytest from ftag import get_mock_file +from numpy.lib import recfunctions as rfn from upp.main import main @@ -17,6 +18,16 @@ class TestClass: def generate_mock(self, out_file, N=100_000): _, f = get_mock_file(num_jets=N, fname=out_file) + jets = f["jets"][:] + if "physicalWeight" not in jets.dtype.names: + jets2 = rfn.append_fields( + jets, + "physicalWeight", + np.ones(jets.shape[0], dtype="f4"), + usemask=False, + ) + del f["jets"] + f.create_dataset("jets", data=jets2) f.close() def setup_method(self, method): diff --git a/tests/integration/test_run_rw.py b/tests/integration/test_run_rw.py index b0ea8b8e..d7265141 100644 --- a/tests/integration/test_run_rw.py +++ b/tests/integration/test_run_rw.py @@ -7,6 +7,7 @@ import h5py import numpy as np from ftag.mock import JET_VARS, get_mock_file +from numpy.lib import recfunctions as rfn from upp.main import main @@ -18,7 +19,25 @@ class TestRunRW: def generate_mock(self, out_file, N=1_000): _, f = get_mock_file(num_jets=N, fname=out_file) - f["jets"]["eventNumber"] = np.arange(N, dtype="i4") + jets = f["jets"][:] + if "eventNumber" in jets.dtype.names: + jets["eventNumber"] = np.arange(N, dtype="i4") + else: + jets = rfn.append_fields( + jets, + "eventNumber", + np.arange(N, dtype="i4"), + usemask=False, + ) + if "physicalWeight" not in jets.dtype.names: + jets = rfn.append_fields( + jets, + "physicalWeight", + np.ones(jets.shape[0], dtype="f4"), + usemask=False, + ) + del f["jets"] + f.create_dataset("jets", data=jets) f.close() def setup_method(self, method): diff --git a/tests/unit/stages/test_metadata_injector.py b/tests/unit/stages/test_metadata_injector.py new file mode 100644 index 00000000..9a341baf --- /dev/null +++ b/tests/unit/stages/test_metadata_injector.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from pathlib import Path + +import h5py +import numpy as np +from numpy.lib import recfunctions as rfn + +import upp.stages.metadata_injector as mi + + +def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): + """Smoke-test MetadataInjector when dependencies are available. + + In CI environments where `MetadataFinder` is not available via the installed `ftag`, + importing this module should fail; in that case we skip. + """ + # Stub MetadataFinder to avoid depending on external metadata DB. + class _Finder: + def __init__(self, *_a, **_k): + pass + + def inject_metadata(self): + return None + + monkeypatch.setattr(mi, "MetadataFinder", _Finder, raising=False) + + # Create minimal input file. + fpath = tmp_path / "in.h5" + n = 10 + jets = np.zeros(n, dtype=[("mcEventWeight", "f4")]) + jets["mcEventWeight"] = np.arange(n, dtype="f4") + 1.0 + with h5py.File(fpath, "w") as f: + f.create_dataset("jets", data=jets) + cb = f.create_group("cutBookkeeper").create_group("nominal") + # MetadataInjector supports both a structured and a plain dataset here. Use a + # plain float dataset to keep the test robust across h5py versions. + cb.create_dataset("counts", data=np.array([float(np.sum(jets["mcEventWeight"]))], dtype="f8")) + md = f.create_group("metadata").create_group("dummy_dsid") + md.create_dataset("cross_section_pb", data=np.array(2.0)) + md.create_dataset("genFiltEff", data=np.array(0.5)) + md.create_dataset("kfactor", data=np.array(1.0)) + + class _Cfg: + def __init__(self, path: Path): + self.config = {"inputs": {"train": {"input_files": [str(path)]}}} + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + out = f["jets"][:] + assert "physicalWeight" in out.dtype.names + # Also ensure old fields remain. + assert "mcEventWeight" in out.dtype.names diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 896652fa..be4941d2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -25,6 +25,7 @@ def test_parse_args_with_config(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=None, resample=True, merge=None, @@ -51,6 +52,7 @@ def test_parse_args_flags_not_given(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -88,6 +90,7 @@ def test_parse_args_flags_negative(config_file): # Check if the parsed arguments match the expected values expected_args = Namespace( config=config_file, + metadata=False, prep=False, resample=False, merge=False, @@ -123,6 +126,7 @@ def test_parse_args_flags_positive(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -160,6 +164,7 @@ def test_parse_args_component(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -197,6 +202,7 @@ def test_parse_args_region(config_file): parsed_args = parse_args(args) expected_args = Namespace( config=config_file, + metadata=False, prep=True, resample=True, merge=True, @@ -216,3 +222,9 @@ def test_parse_args_region(config_file): ) assert parsed_args == expected_args + + +def test_parse_args_metadata_flag(config_file): + args = ["--config", str(config_file), "--metadata"] + parsed_args = parse_args(args) + assert parsed_args.metadata is True From 29ef221b0c7a4fb0afbec42f0a6db512de575262 Mon Sep 17 00:00:00 2001 From: yuanda Date: Thu, 28 May 2026 14:07:54 +0200 Subject: [PATCH 04/11] precommit and pytest fix --- tests/unit/stages/test_metadata_injector.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/stages/test_metadata_injector.py b/tests/unit/stages/test_metadata_injector.py index 9a341baf..6e6a1960 100644 --- a/tests/unit/stages/test_metadata_injector.py +++ b/tests/unit/stages/test_metadata_injector.py @@ -4,7 +4,6 @@ import h5py import numpy as np -from numpy.lib import recfunctions as rfn import upp.stages.metadata_injector as mi @@ -15,6 +14,7 @@ def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): In CI environments where `MetadataFinder` is not available via the installed `ftag`, importing this module should fail; in that case we skip. """ + # Stub MetadataFinder to avoid depending on external metadata DB. class _Finder: def __init__(self, *_a, **_k): @@ -35,7 +35,9 @@ def inject_metadata(self): cb = f.create_group("cutBookkeeper").create_group("nominal") # MetadataInjector supports both a structured and a plain dataset here. Use a # plain float dataset to keep the test robust across h5py versions. - cb.create_dataset("counts", data=np.array([float(np.sum(jets["mcEventWeight"]))], dtype="f8")) + cb.create_dataset( + "counts", data=np.array([float(np.sum(jets["mcEventWeight"]))], dtype="f8") + ) md = f.create_group("metadata").create_group("dummy_dsid") md.create_dataset("cross_section_pb", data=np.array(2.0)) md.create_dataset("genFiltEff", data=np.array(0.5)) From 3cd4a5780aa5443391c89497aa4b800c06d3960b Mon Sep 17 00:00:00 2001 From: yuanda Date: Mon, 8 Jun 2026 13:47:35 +0200 Subject: [PATCH 05/11] precommit fix and improve coverage --- tests/unit/stages/test_metadata_injector.py | 110 ++++++++++++++------ tests/unit/stages/test_normalisation.py | 25 +++++ tests/unit/stages/test_rw_merge.py | 77 ++++++++++++++ tests/unit/test_main.py | 76 +++++++++++++- upp/stages/reweight.py | 2 +- 5 files changed, 258 insertions(+), 32 deletions(-) diff --git a/tests/unit/stages/test_metadata_injector.py b/tests/unit/stages/test_metadata_injector.py index 6e6a1960..a3f23783 100644 --- a/tests/unit/stages/test_metadata_injector.py +++ b/tests/unit/stages/test_metadata_injector.py @@ -8,50 +8,100 @@ import upp.stages.metadata_injector as mi -def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): - """Smoke-test MetadataInjector when dependencies are available. - - In CI environments where `MetadataFinder` is not available via the installed `ftag`, - importing this module should fail; in that case we skip. - """ - - # Stub MetadataFinder to avoid depending on external metadata DB. - class _Finder: - def __init__(self, *_a, **_k): - pass - - def inject_metadata(self): - return None - - monkeypatch.setattr(mi, "MetadataFinder", _Finder, raising=False) - - # Create minimal input file. +def _make_injector_file(tmp_path, jets_dtype, set_attrs=False): + """Create a minimal HDF5 file for MetadataInjector tests.""" fpath = tmp_path / "in.h5" n = 10 - jets = np.zeros(n, dtype=[("mcEventWeight", "f4")]) - jets["mcEventWeight"] = np.arange(n, dtype="f4") + 1.0 + jets = np.zeros(n, dtype=jets_dtype) + if "mcEventWeight" in jets.dtype.names: + jets["mcEventWeight"] = np.arange(n, dtype="f4") + 1.0 + sow = float(np.sum(jets["mcEventWeight"])) if "mcEventWeight" in jets.dtype.names else 55.0 with h5py.File(fpath, "w") as f: - f.create_dataset("jets", data=jets) + ds = f.create_dataset("jets", data=jets) + if set_attrs: + ds.attrs["description"] = "test_attr" + ds.attrs["version"] = 42 cb = f.create_group("cutBookkeeper").create_group("nominal") - # MetadataInjector supports both a structured and a plain dataset here. Use a - # plain float dataset to keep the test robust across h5py versions. - cb.create_dataset( - "counts", data=np.array([float(np.sum(jets["mcEventWeight"]))], dtype="f8") - ) + cb.create_dataset("counts", data=np.array([sow], dtype="f8")) md = f.create_group("metadata").create_group("dummy_dsid") md.create_dataset("cross_section_pb", data=np.array(2.0)) md.create_dataset("genFiltEff", data=np.array(0.5)) md.create_dataset("kfactor", data=np.array(1.0)) + return fpath + - class _Cfg: - def __init__(self, path: Path): - self.config = {"inputs": {"train": {"input_files": [str(path)]}}} +class _Cfg: + def __init__(self, path: Path) -> None: + self.config = {"inputs": {"train": {"input_files": [str(path)]}}} + +def _stub_finder(monkeypatch): + class _Finder: + def __init__(self, *_a, **_k) -> None: + pass + + def inject_metadata(self) -> None: + return None + + monkeypatch.setattr(mi, "MetadataFinder", _Finder, raising=False) + + +def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): + """Smoke-test MetadataInjector: physicalWeight is appended and fields are preserved.""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")]) injector = mi.MetadataInjector(_Cfg(fpath)) injector.run() with h5py.File(fpath) as f: out = f["jets"][:] assert "physicalWeight" in out.dtype.names - # Also ensure old fields remain. assert "mcEventWeight" in out.dtype.names + + +def test_metadata_injector_missing_mcEventWeight_recovers(tmp_path, monkeypatch): + """Exception path: missing mcEventWeight raises KeyError, backup is restored.""" + _stub_finder(monkeypatch) + # jets without mcEventWeight triggers line 65 then lines 93-97 + fpath = _make_injector_file(tmp_path, [("someField", "f4")]) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() # exception caught internally, must not propagate + + assert fpath.exists() + assert not backup_path.exists() + + +def test_metadata_injector_drops_existing_physicalWeight(tmp_path, monkeypatch): + """Existing physicalWeight field is dropped before recomputing (line 73).""" + _stub_finder(monkeypatch) + dtype = [("mcEventWeight", "f4"), ("physicalWeight", "f4")] + fpath = _make_injector_file(tmp_path, dtype) + # Overwrite physicalWeight with sentinel value to verify it gets replaced + with h5py.File(fpath, "a") as f: + data = f["jets"][:] + data["physicalWeight"] = 999.0 + del f["jets"] + f.create_dataset("jets", data=data) + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + out = f["jets"][:] + assert "physicalWeight" in out.dtype.names + assert not np.all(out["physicalWeight"] == 999.0) + + +def test_metadata_injector_preserves_dataset_attrs(tmp_path, monkeypatch): + """Original dataset attrs are restored after jets dataset recreation (lines 84-85).""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")], set_attrs=True) + + injector = mi.MetadataInjector(_Cfg(fpath)) + injector.run() + + with h5py.File(fpath) as f: + assert f["jets"].attrs["description"] == "test_attr" + assert f["jets"].attrs["version"] == 42 diff --git a/tests/unit/stages/test_normalisation.py b/tests/unit/stages/test_normalisation.py index f2bbaeef..450cff20 100644 --- a/tests/unit/stages/test_normalisation.py +++ b/tests/unit/stages/test_normalisation.py @@ -4,6 +4,7 @@ import subprocess from pathlib import Path +import numpy as np import pytest from ftag import get_mock_file @@ -105,6 +106,30 @@ def test_combine_mean_std(): assert combined_mean_ref == combined_mean assert combined_std_ref == combined_std + def test_get_class_dict_integer_label(self): + """Integer-typed label vars are counted in get_class_dict (lines 160-161).""" + norm = Normalisation( + config=PreprocessingConfig.from_file( + Path(CFG_DIR / "test_config_pdf_auto_umami.yaml"), "train" + ) + ) + # config labels for jets: [pt, eta, mass] + resampling vars [pt, abs_eta] + # → effective labels: [pt, eta, mass, abs_eta]; make "pt" integer → lines 160-161 + jets_dtype = np.dtype( + [ + ("pt_btagJes", "f4"), + ("eta_btagJes", "f4"), + ("flavour_label", "i4"), + ("pt", "i4"), + ("eta", "f4"), + ("mass", "f4"), + ("abs_eta", "f4"), + ] + ) + batch = {"jets": np.zeros(10, dtype=jets_dtype)} + result = norm.get_class_dict(batch) + assert "pt" in result["jets"] + @staticmethod def test_combine_norm_dict(): # Test combination of mean and std diff --git a/tests/unit/stages/test_rw_merge.py b/tests/unit/stages/test_rw_merge.py index eb60cc45..8156f2db 100644 --- a/tests/unit/stages/test_rw_merge.py +++ b/tests/unit/stages/test_rw_merge.py @@ -2,6 +2,8 @@ from __future__ import annotations +from unittest.mock import MagicMock + import numpy as np import pytest @@ -54,3 +56,78 @@ def test_empty_input(self): result = RWMerge._assign_weights(rw, bins, classes) assert result.shape == (0,) + + +class TestRWMergeInit: + """Cover RWMerge.__init__ assertion for outfile_idx_range (line 26).""" + + def test_non_tuple_idx_range_raises(self): + config = MagicMock() + with pytest.raises(AssertionError): + RWMerge(config, outfile_idx_range=[0, 1]) + + def test_wrong_length_tuple_raises(self): + config = MagicMock() + with pytest.raises(AssertionError): + RWMerge(config, outfile_idx_range=(0, 1, 2)) + + +class TestStartMp: + """Cover start_mp multiprocess branch (lines 314-316).""" + + def test_start_mp_multiprocess(self, monkeypatch): + called_with: list = [] + + class FakePool: + def __init__(self, n: int) -> None: + pass + + def __enter__(self) -> FakePool: + return self + + def __exit__(self, *_: object) -> None: + pass + + def starmap(self, _fn: object, args_list: list) -> None: + for args in args_list: + called_with.append(args) + + monkeypatch.setattr("upp.stages.rw_merge.Pool", FakePool) + + def fn(x: int) -> int: + return x + + RWMerge.start_mp(fn, [(1,), (2,), (3,)], n_threads=2) + assert called_with == [(1,), (2,), (3,)] + + +class TestGetSampleWeightsException: + """Cover the except-and-reraise path in get_sample_weights (lines 185-187).""" + + def test_exception_propagates(self, monkeypatch): + n = 5 + jets = np.zeros(n, dtype=[("x", "f4"), ("class_var", "i4")]) + batch = {"jets": jets} + weights = { + "jets": { + "rw1": { + "rw_vars": ["x"], + "class_var": "class_var", + "bins": [np.linspace(0.0, 1.0, 6)], + "weights": {"0": np.ones(5)}, + } + } + } + + monkeypatch.setattr( + "upp.stages.rw_merge.bin_jets", + lambda *_a, **_k: (np.zeros(5), np.zeros((1, n), dtype=int)), + ) + + def bad_assign(*_a: object, **_k: object) -> None: + raise ValueError("injected error") + + monkeypatch.setattr(RWMerge, "_assign_weights", staticmethod(bad_assign)) + + with pytest.raises(ValueError, match="injected error"): + RWMerge.get_sample_weights(batch, weights) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index be4941d2..bcd1e98e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1,10 +1,12 @@ from __future__ import annotations +import argparse from argparse import Namespace +from unittest.mock import MagicMock, patch from pytest import fixture -from upp.main import parse_args +from upp.main import main, parse_args, run_pp @fixture @@ -228,3 +230,75 @@ def test_parse_args_metadata_flag(config_file): args = ["--config", str(config_file), "--metadata"] parsed_args = parse_args(args) assert parsed_args.metadata is True + + +def _base_args(config_file: object, **overrides: object) -> argparse.Namespace: + defaults: dict = dict( + config=config_file, + metadata=False, + prep=False, + resample=False, + merge=False, + norm=False, + plot=False, + split="train", + component=None, + region=None, + container=None, + grid=False, + split_components=False, + reweight=False, + rw_merge=False, + rw_merge_idx=None, + files=None, + skip_sample_check=False, + ) + defaults.update(overrides) + return argparse.Namespace(**defaults) + + +def test_run_pp_metadata_injection(tmp_path): + """run_pp with metadata=True constructs and runs MetadataInjector (lines 222-224).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + args = _base_args(config_file, metadata=True) + + mock_injector = MagicMock() + with ( + patch("upp.main.PreprocessingConfig.from_file", return_value=MagicMock()), + patch("upp.main.MetadataInjector", return_value=mock_injector), + ): + run_pp(args) + + mock_injector.run.assert_called_once() + + +def test_run_pp_rw_merge_with_idx(tmp_path): + """run_pp with rw_merge_idx parses the comma-separated pair (lines 269-272).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + args = _base_args(config_file, rw_merge=True, rw_merge_idx="0,1") + + mock_rw = MagicMock() + with ( + patch("upp.main.PreprocessingConfig.from_file", return_value=MagicMock()), + patch("upp.main.RWMerge", return_value=mock_rw), + ): + run_pp(args) + + mock_rw.run.assert_called_once() + + +def test_main_split_all(tmp_path): + """main() with split='all' calls run_pp three times (lines 303-310).""" + config_file = tmp_path / "config.yaml" + config_file.write_text("") + mock_args = _base_args(config_file, split="all") + + with ( + patch("upp.main.parse_args", return_value=mock_args), + patch("upp.main.run_pp") as mock_run_pp, + ): + main() + + assert mock_run_pp.call_count == 3 diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index d49a08de..43f2ee70 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -186,7 +186,7 @@ def calculate_weights( else: w = np.ones(mask.sum(), dtype=float) - hist, outbins = bin_jets( + hist, _outbins = bin_jets( data_masked[rw.reweight_vars], rw.flat_bins, weights=w, From 917cf3cc7b157440ecfa1000d1c10a13a30e543d Mon Sep 17 00:00:00 2001 From: yuanda Date: Mon, 8 Jun 2026 13:54:25 +0200 Subject: [PATCH 06/11] 3.13 and 3.14 fix --- tests/unit/stages/test_rw_merge.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/tests/unit/stages/test_rw_merge.py b/tests/unit/stages/test_rw_merge.py index 8156f2db..6cd8df4e 100644 --- a/tests/unit/stages/test_rw_merge.py +++ b/tests/unit/stages/test_rw_merge.py @@ -2,11 +2,12 @@ from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pytest +import upp.stages.rw_merge as rw_merge_module from upp.stages.rw_merge import RWMerge @@ -75,7 +76,7 @@ def test_wrong_length_tuple_raises(self): class TestStartMp: """Cover start_mp multiprocess branch (lines 314-316).""" - def test_start_mp_multiprocess(self, monkeypatch): + def test_start_mp_multiprocess(self): called_with: list = [] class FakePool: @@ -92,19 +93,19 @@ def starmap(self, _fn: object, args_list: list) -> None: for args in args_list: called_with.append(args) - monkeypatch.setattr("upp.stages.rw_merge.Pool", FakePool) - def fn(x: int) -> int: return x - RWMerge.start_mp(fn, [(1,), (2,), (3,)], n_threads=2) + with patch.object(rw_merge_module, "Pool", FakePool): + RWMerge.start_mp(fn, [(1,), (2,), (3,)], n_threads=2) + assert called_with == [(1,), (2,), (3,)] class TestGetSampleWeightsException: """Cover the except-and-reraise path in get_sample_weights (lines 185-187).""" - def test_exception_propagates(self, monkeypatch): + def test_exception_propagates(self): n = 5 jets = np.zeros(n, dtype=[("x", "f4"), ("class_var", "i4")]) batch = {"jets": jets} @@ -119,15 +120,15 @@ def test_exception_propagates(self, monkeypatch): } } - monkeypatch.setattr( - "upp.stages.rw_merge.bin_jets", - lambda *_a, **_k: (np.zeros(5), np.zeros((1, n), dtype=int)), - ) + def fake_bin_jets(*_a: object, **_k: object) -> tuple: + return (np.zeros(5), np.zeros((1, n), dtype=int)) def bad_assign(*_a: object, **_k: object) -> None: raise ValueError("injected error") - monkeypatch.setattr(RWMerge, "_assign_weights", staticmethod(bad_assign)) - - with pytest.raises(ValueError, match="injected error"): + with ( + patch.object(rw_merge_module, "bin_jets", fake_bin_jets), + patch.object(RWMerge, "_assign_weights", staticmethod(bad_assign)), + pytest.raises(ValueError, match="injected error"), + ): RWMerge.get_sample_weights(batch, weights) From ff06f3472d805ac7f7ce4814e51b58c69e4f9aa9 Mon Sep 17 00:00:00 2001 From: yuanda Date: Mon, 22 Jun 2026 12:10:32 +0200 Subject: [PATCH 07/11] Address PR #142 review feedback - Relocate metadata/flavours/variables-r10 configs under upp/configs/MetadataRW/; revert default variables/xbb/single-b configs to main - split_containers: gate physicalWeight append on column presence / keep_all - reweight: restore idx_below_min zeroing, unify WEIGHT_CAP applied only on the physicalWeight path, drop dead safe_hist and debug prints - metadata_injector: guard sow==0 / non-finite metadata, collect per-file failures and re-raise a summary instead of log-and-continue - preprocessing_config: keep_all = keep_all_variables or is_test (restore test split) - Condense inline block comments; add changelog entry and document --metadata stage and keep_all_variables flag Co-authored-by: Cursor --- changelog.md | 1 + docs/run.md | 34 ++++++++++++ tests/unit/stages/test_metadata_injector.py | 27 ++++++++-- upp/classes/preprocessing_config.py | 7 ++- .../{ => MetadataRW}/flavours-r10.yaml | 4 +- upp/configs/{ => MetadataRW}/metadata.yaml | 2 +- upp/configs/MetadataRW/variables-r10.yaml | 53 +++++++++++++++++++ upp/configs/single-b-upgrade.yaml | 4 ++ upp/configs/variables.yaml | 4 -- upp/configs/xbb-gn3x.yaml | 2 - upp/configs/xbb.yaml | 2 - upp/stages/metadata_injector.py | 19 +++++++ upp/stages/reweight.py | 41 ++++++++++---- upp/stages/split_containers.py | 16 +++--- 14 files changed, 177 insertions(+), 39 deletions(-) rename upp/configs/{ => MetadataRW}/flavours-r10.yaml (83%) rename upp/configs/{ => MetadataRW}/metadata.yaml (98%) create mode 100644 upp/configs/MetadataRW/variables-r10.yaml diff --git a/changelog.md b/changelog.md index ebb038db..0e4610b9 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,7 @@ ### [Latest] +- Add `--metadata` stage (metadata injection + physicalWeight) and `global.keep_all_variables` full-ntuple passthrough for centrally produced inputs [#142](https://github.com/umami-hep/umami-preprocessing/pull/142) - Update atlas-ftag-tools and puma [#143](https://github.com/umami-hep/umami-preprocessing/pull/143) - Sync pre-commit hooks and fix Ruff lint failures [#138](https://github.com/umami-hep/umami-preprocessing/pull/138) - Update docs and CI validation for supported Python versions [#137](https://github.com/umami-hep/umami-preprocessing/pull/137) diff --git a/docs/run.md b/docs/run.md index b78d9877..3ebad7c3 100644 --- a/docs/run.md +++ b/docs/run.md @@ -134,6 +134,40 @@ The normalize stage (`--norm`) calculates scaling and shifting values for all va The plotting stage (`--plot`) produces histograms of resampled variables to verify the resampling quality. You can find these plots in `/plots/`. +### Centrally produced inputs: metadata + reweighting workflow + +Centrally produced `.h5` ntuples from TDD often lack pre-computed cross-section metadata and a `physicalWeight` column. For these inputs UPP provides an alternative chain that replaces resample/merge with a metadata injection and a histogram-based reweighting: + +`--metadata` → `--split-components` → `--reweight` (`--rw`) → `--rw-merge` (`--rwm`) → `--norm` + +An example config is provided in [`upp/configs/MetadataRW/metadata.yaml`](https://github.com/umami-hep/umami-preprocessing/blob/main/upp/configs/MetadataRW/metadata.yaml). + +#### Metadata injection (`--metadata`) + +The metadata stage injects cross section, k-factor and generator filter efficiency (via `ftag.find_metadata.MetadataFinder`) and appends a `physicalWeight` column to the `jets` dataset: + +$$\text{physicalWeight} = \frac{\text{XS} \times \text{eff} \times \text{kfactor}}{\text{SOW}} \times \text{mcEventWeight}$$ + +```bash +preprocess --config upp/configs/MetadataRW/metadata.yaml --metadata +``` + +!!!warning "In-place modification" + The metadata stage edits the input files in place (a `.bak` backup is created per file and removed on success). Run it on **workspace copies**, not on central originals. If the sum-of-weights is zero or any metadata value is non-finite the stage raises, and any per-file failures are reported as a summary at the end of the run. + +The downstream reweighting stage automatically uses the `physicalWeight` column when present, and otherwise falls back to uniform weights. + +#### Full-ntuple passthrough (`global.keep_all_variables`) + +By default the split and rw-merge outputs only keep the variables declared in your `variables` config. Set `keep_all_variables: true` under `global` to instead preserve **all** top-level HDF5 datasets (e.g. `tracks`, `towers`, `flow`) and full `jets` fields in the train/val/test outputs, while still appending the reweight columns on `jets`: + +```yaml +global: + keep_all_variables: true +``` + +The `test` split always keeps all variables regardless of this flag. + ### Additional Scripts: Initial Sample Check The check for the initial samples from the prepare stage can also be run stand-alone. This is important if you plan to run in parallel mode. To do so, you can simply use the following command: diff --git a/tests/unit/stages/test_metadata_injector.py b/tests/unit/stages/test_metadata_injector.py index a3f23783..b42cbfd1 100644 --- a/tests/unit/stages/test_metadata_injector.py +++ b/tests/unit/stages/test_metadata_injector.py @@ -4,6 +4,7 @@ import h5py import numpy as np +import pytest import upp.stages.metadata_injector as mi @@ -60,14 +61,34 @@ def test_metadata_injector_appends_physical_weight(tmp_path, monkeypatch): def test_metadata_injector_missing_mcEventWeight_recovers(tmp_path, monkeypatch): - """Exception path: missing mcEventWeight raises KeyError, backup is restored.""" + """Exception path: missing mcEventWeight restores the backup and a summary is raised.""" _stub_finder(monkeypatch) - # jets without mcEventWeight triggers line 65 then lines 93-97 + # jets without mcEventWeight triggers the per-file failure + restore path fpath = _make_injector_file(tmp_path, [("someField", "f4")]) backup_path = fpath.with_suffix(fpath.suffix + ".bak") injector = mi.MetadataInjector(_Cfg(fpath)) - injector.run() # exception caught internally, must not propagate + # The file is restored from backup, but failures are surfaced as a summary error. + with pytest.raises(RuntimeError, match="Metadata injection failed"): + injector.run() + + assert fpath.exists() + assert not backup_path.exists() + + +def test_metadata_injector_zero_sow_raises(tmp_path, monkeypatch): + """A zero sum-of-weights is rejected instead of producing inf/nan weights.""" + _stub_finder(monkeypatch) + fpath = _make_injector_file(tmp_path, [("mcEventWeight", "f4")]) + # Force the stored sum-of-weights to zero. + with h5py.File(fpath, "a") as f: + del f["cutBookkeeper/nominal/counts"] + f["cutBookkeeper/nominal"].create_dataset("counts", data=np.array([0.0], dtype="f8")) + backup_path = fpath.with_suffix(fpath.suffix + ".bak") + + injector = mi.MetadataInjector(_Cfg(fpath)) + with pytest.raises(RuntimeError, match="Metadata injection failed"): + injector.run() assert fpath.exists() assert not backup_path.exists() diff --git a/upp/classes/preprocessing_config.py b/upp/classes/preprocessing_config.py index 7beae5bd..20ab770e 100644 --- a/upp/classes/preprocessing_config.py +++ b/upp/classes/preprocessing_config.py @@ -145,8 +145,7 @@ class PreprocessingConfig: num_jets_per_output_file: int | None = None skip_checks: bool = False skip_config_copy: bool = False - # When true: keep all input top-level datasets (tracks, towers, …), not only the - # variables.yaml subset, through split-containers and rw-merge. + # Keep all top-level datasets (not just the variables.yaml subset) through split/rw-merge. keep_all_variables: bool = False def __post_init__(self): @@ -204,11 +203,11 @@ def __post_init__(self): if selection := groups.get("selection", None): selectors[name] = TrackSelector(Cuts.from_list(selection)) - # configure variables (keep_all_variables is independent of train/val/test split) + # test split always keeps all variables; keep_all_variables also forces it for train/val. self.variables = VariableConfig( self.config["variables"], self.jets_name, - self.keep_all_variables, + self.keep_all_variables or self.is_test, selectors, ) if self.sampl_cfg is not None: diff --git a/upp/configs/flavours-r10.yaml b/upp/configs/MetadataRW/flavours-r10.yaml similarity index 83% rename from upp/configs/flavours-r10.yaml rename to upp/configs/MetadataRW/flavours-r10.yaml index 13d142ad..0003bea6 100644 --- a/upp/configs/flavours-r10.yaml +++ b/upp/configs/MetadataRW/flavours-r10.yaml @@ -1,5 +1,5 @@ -# flavour-r10.yaml -# 已同步为 0-6 连续整数标签 +# Flavour definitions for the metadata/reweight workflow on centrally produced +# large-R ntuples. Labels are contiguous integers 0-5. - name: djets category: jets diff --git a/upp/configs/metadata.yaml b/upp/configs/MetadataRW/metadata.yaml similarity index 98% rename from upp/configs/metadata.yaml rename to upp/configs/MetadataRW/metadata.yaml index b51314e0..0aef6040 100644 --- a/upp/configs/metadata.yaml +++ b/upp/configs/MetadataRW/metadata.yaml @@ -11,7 +11,7 @@ global: num_jets_per_output_file: 1_000_000 # test split: plot.py uses merged VDS path (glob) instead of per-sample pp_output_test_*.h5 merge_test_samples: true - flavour_config: /eos/home-y/yuanda/dev/umami-preprocessing/upp/configs/flavours-r10.yaml + flavour_config: /eos/home-y/yuanda/dev/umami-preprocessing/upp/configs/MetadataRW/flavours-r10.yaml inputs: train: diff --git a/upp/configs/MetadataRW/variables-r10.yaml b/upp/configs/MetadataRW/variables-r10.yaml new file mode 100644 index 00000000..534e559a --- /dev/null +++ b/upp/configs/MetadataRW/variables-r10.yaml @@ -0,0 +1,53 @@ +jets: + inputs: + - pt_btagJes + - eta_btagJes + - R10TruthLabel_R22v1 + - PartonTruthLabelID + - physicalWeight + labels: + - HadronConeExclTruthLabelID + - HadronConeExclExtendedTruthLabelID + - HadronConeExclTruthLabelPt + - HadronConeExclTruthLabelLxy + - HadronConeExclTruthLabelDR + - HadronGhostTruthLabelID + - HadronGhostExtendedTruthLabelID + - HadronGhostTruthLabelPt + - HadronGhostTruthLabelLxy + - HadronGhostTruthLabelDR + - pt + - eta + - mass + - n_tracks + - n_truth_promptLepton + - eventNumber + - jetFoldHash + - physicalWeight + +tracks: + inputs: + - d0 + - z0SinTheta + - dphi + - deta + - qOverP + - lifetimeSignedD0Significance + - lifetimeSignedZ0SinThetaSignificance + - phiUncertainty + - thetaUncertainty + - qOverPUncertainty + - numberOfPixelHits + - numberOfSCTHits + - numberOfInnermostPixelLayerHits + - numberOfNextToInnermostPixelLayerHits + - numberOfInnermostPixelLayerSharedHits + - numberOfInnermostPixelLayerSplitHits + - numberOfPixelSharedHits + - numberOfPixelSplitHits + - numberOfSCTSharedHits + - leptonID + labels: + - ftagTruthOriginLabel + - ftagTruthTypeLabel + - ftagTruthVertexIndex diff --git a/upp/configs/single-b-upgrade.yaml b/upp/configs/single-b-upgrade.yaml index 3a6a659c..48648905 100644 --- a/upp/configs/single-b-upgrade.yaml +++ b/upp/configs/single-b-upgrade.yaml @@ -18,6 +18,8 @@ zprime: &zprime 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" + + lowpt: &lowpt name: lowpt cuts: @@ -72,6 +74,8 @@ components: flavours: [taujets] num_jets: 805_555 + + resampling: target: cjets method: pdf #countup diff --git a/upp/configs/variables.yaml b/upp/configs/variables.yaml index be6d21e9..1870c74e 100644 --- a/upp/configs/variables.yaml +++ b/upp/configs/variables.yaml @@ -2,9 +2,6 @@ jets: inputs: - pt_btagJes - eta_btagJes - - R10TruthLabel_R22v1 - - PartonTruthLabelID - - physicalWeight labels: - HadronConeExclTruthLabelID - HadronConeExclExtendedTruthLabelID @@ -23,7 +20,6 @@ jets: - n_truth_promptLepton - eventNumber - jetFoldHash - - physicalWeight tracks: inputs: diff --git a/upp/configs/xbb-gn3x.yaml b/upp/configs/xbb-gn3x.yaml index 50fd9ff3..482957ca 100644 --- a/upp/configs/xbb-gn3x.yaml +++ b/upp/configs/xbb-gn3x.yaml @@ -111,8 +111,6 @@ resampling: bins: [[0, 2, 20]] mass: bins: [[50_000, 300_000, 50]] - mcCampaignYear: - bins: [[2015, 2026, 11]] # note: sensible defaults are defined in the PreprocessingConfig constructor global: diff --git a/upp/configs/xbb.yaml b/upp/configs/xbb.yaml index 866e41d7..dd5b379e 100644 --- a/upp/configs/xbb.yaml +++ b/upp/configs/xbb.yaml @@ -65,8 +65,6 @@ resampling: bins: [[0, 2, 20]] mass: bins: [[50_000, 300_000, 50]] - mcCampaignYear: - bins: [[2015, 2026, 11]] # note: sensible defaults are defined in the PreprocessingConfig constructor global: diff --git a/upp/stages/metadata_injector.py b/upp/stages/metadata_injector.py index fb7a338f..b4b9b6c9 100644 --- a/upp/stages/metadata_injector.py +++ b/upp/stages/metadata_injector.py @@ -1,6 +1,7 @@ from __future__ import annotations import glob +import math import shutil from pathlib import Path @@ -30,6 +31,7 @@ def run(self): matched = glob.glob(f) expanded_files.extend(matched) + failures: list[tuple[str, str]] = [] for fpath_str in expanded_files: fpath = Path(fpath_str) backup_path = fpath.with_suffix(fpath.suffix + ".bak") @@ -55,6 +57,14 @@ def run(self): float(sow_ds["sumOfWeights"][0]) if sow_ds.dtype.names else float(sow_ds[0]) ) + # Guard against div-by-zero / non-finite metadata (silent inf/nan weights). + if sow == 0 or not math.isfinite(sow): + raise ValueError(f"Invalid sum-of-weights (sow={sow}) in {fpath_str}") + if not all(math.isfinite(v) for v in (xs, eff, kf)): + raise ValueError( + f"Non-finite metadata in {fpath_str}: xs={xs}, eff={eff}, kfactor={kf}" + ) + # 3. Read original jets and their attributes old_jets_ds = f["jets"] original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata @@ -95,3 +105,12 @@ def run(self): if backup_path.exists(): self.log.warning(f"Restoring {fpath.name} from backup...") shutil.move(backup_path, fpath) + failures.append((fpath_str, str(e))) + + # Re-raise a summary so partial failures are not silently missed in a long run. + if failures: + summary = "\n".join(f" - {path}: {err}" for path, err in failures) + raise RuntimeError( + f"Metadata injection failed for {len(failures)}/{len(expanded_files)} file(s):\n" + f"{summary}" + ) diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index 43f2ee70..d872d823 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -14,6 +14,9 @@ from upp.classes.preprocessing_config import PreprocessingConfig from upp.stages.hist import bin_jets +# Upper bound for physical weights and reweight factors, which can span many orders of magnitude. +WEIGHT_CAP = 1e4 + class Reweight: def __init__(self, config: PreprocessingConfig): @@ -107,6 +110,7 @@ def calculate_weights( print("N per file : ", self.num_jets_estimate) # Get the variables we need to reweight + use_physical_weight = False for rw in reweights: rw_group = rw.group if rw_group not in all_vars: @@ -118,6 +122,7 @@ def calculate_weights( all_vars[rw_group] += ["valid"] if "physicalWeight" in existing_vars.get(rw_group, []): all_vars[rw_group] += ["physicalWeight"] + use_physical_weight = True if "jets" not in all_vars: all_vars["jets"] = ["pt"] all_vars = {k: list(set(v)) for k, v in all_vars.items()} @@ -177,12 +182,10 @@ def calculate_weights( mask = data[rw.class_var] == cls data_masked = data[mask] - # Use physicalWeight column from the structured jets array if present; - # otherwise fall back to uniform weights. + # Use physicalWeight if present, else uniform weights. if "physicalWeight" in data_masked.dtype.names: - PW_CAP = 10000.0 w = np.asarray(data_masked["physicalWeight"], dtype=np.float64) - w = np.clip(w, 0, PW_CAP) + w = np.clip(w, 0, WEIGHT_CAP) else: w = np.ones(mask.sum(), dtype=float) @@ -287,14 +290,30 @@ def calculate_weights( "rw_vars": rw.reweight_vars, "class_var": rw.class_var, } - RW_CAP = 1e4 - HIST_FLOOR = 1e-6 - target = all_targets[rw_group][rw_rep] + idx_below_min = None for cls, hist in all_histograms[rw_group][rw_rep]["histograms"].items(): - safe_hist = np.maximum(hist, HIST_FLOOR) - weights = np.zeros_like(hist, dtype=float) - np.divide(target, safe_hist, out=weights, where=hist > HIST_FLOOR) - output_weights[rw_group][rw_rep]["weights"][cls] = np.clip(weights, 0, RW_CAP) + this_idx_below_min = hist == 0 # | (all_targets[rw_group][rw_rep] == 0) + output_weights[rw_group][rw_rep]["weights"][cls] = np.where( + hist > 0, all_targets[rw_group][rw_rep] / hist, 0 + ) + if idx_below_min is None: + idx_below_min = this_idx_below_min + else: + idx_below_min |= this_idx_below_min + # If we have any bins where we have 0 of a given flavour, we set all the + # weights to 0 + if np.any(idx_below_min): + for cls in all_histograms[rw_group][rw_rep]["histograms"]: + output_weights[rw_group][rw_rep]["weights"][cls][idx_below_min] = 0 + # Cap final factors only on the physicalWeight path; default path unchanged. + if use_physical_weight: + for cls in output_weights[rw_group][rw_rep]["weights"]: + np.clip( + output_weights[rw_group][rw_rep]["weights"][cls], + 0, + WEIGHT_CAP, + out=output_weights[rw_group][rw_rep]["weights"][cls], + ) return output_weights diff --git a/upp/stages/split_containers.py b/upp/stages/split_containers.py index ede25e1d..318d9256 100644 --- a/upp/stages/split_containers.py +++ b/upp/stages/split_containers.py @@ -133,20 +133,18 @@ def split_file( add_flavour_label = flavour_label_list is not None # All variables for test file all_variables = get_all_datasets(input_file) - print("All variables: ", all_variables, flush=True) # Subset of variables for train/val files parsed_variables: dict[str, list[str]] | dict[str, None] = ( parse_variables(variables) if variables is not None else all_variables ) - # Ensure physicalWeight is requested for jets when using a column list (not None). + # Request physicalWeight only if present (added by --metadata), else reader raises. if isinstance(parsed_variables, dict): jets_cols = parsed_variables.get("jets") if isinstance(jets_cols, list) and "physicalWeight" not in jets_cols: - jets_cols.append("physicalWeight") - print("FORCE RE-INTEGRATED: Only physicalWeight into jets", flush=True) - # ----------------------- - - print("parsed variables: ", parsed_variables, flush=True) + with h5py.File(input_file, "r") as f: + has_pw = "jets" in f and "physicalWeight" in (f["jets"].dtype.names or ()) + if has_pw: + jets_cols.append("physicalWeight") start = time.time() reader = H5Reader(input_file, batch_size=batch_size, shuffle=False) if output_name is None: @@ -317,9 +315,7 @@ def run( assert container is not None, "Can only specify files if a container is specified" for container, cuts_by_component in containers_with_split_cuts.items(): - # Sanitize container name when used as a directory: avoid literal '*' or '/' - # in the split-components subdir, otherwise downstream H5Reader will treat the - # path as a glob pattern and create a shared VDS that collides across flavours. + # Sanitize '*' and '/' so the subdir isn't treated as a glob by H5Reader. container_dir_name = container.replace("*", "all").replace("/", "_") or "default" this_out_dir = ( Path(self.config.base_dir) / "split-components" / container_dir_name From ce5403f3c2185a765fbcdc43b1f414368febd345 Mon Sep 17 00:00:00 2001 From: yuanda Date: Mon, 22 Jun 2026 13:09:35 +0200 Subject: [PATCH 08/11] Add physical-weight aware plotting for the reweight route Weight --plot histograms by physicalWeight (capped) and reweight columns when present, and label plots "Pre/Post Reweighting" instead of "Resampling" on the reweight route. Default resampling route stays unweighted. Adds unit tests and updates changelog/docs. Co-authored-by: Cursor --- changelog.md | 2 +- docs/run.md | 2 + tests/unit/stages/test_plotting.py | 79 ++++++++++++++++++++++++ upp/stages/plot.py | 96 ++++++++++++++++++++++++++++-- 4 files changed, 174 insertions(+), 5 deletions(-) diff --git a/changelog.md b/changelog.md index be80dac0..e40fbfdb 100644 --- a/changelog.md +++ b/changelog.md @@ -2,7 +2,7 @@ ### [Latest] -- Add `--metadata` stage (metadata injection + physicalWeight) and `global.keep_all_variables` full-ntuple passthrough for centrally produced inputs [#142](https://github.com/umami-hep/umami-preprocessing/pull/142) +- Add `--metadata` stage (metadata injection + physicalWeight) and `global.keep_all_variables` full-ntuple passthrough for centrally produced inputs; `--plot` weights distributions by `physicalWeight` (capped) × reweight columns when present so plots reflect the reweight route [#142](https://github.com/umami-hep/umami-preprocessing/pull/142) ### [v0.3.1](https://github.com/umami-hep/umami-preprocessing/releases/tag/v0.3.1) (19.06.2026) - Make skip-resampling work end-to-end; support `num_jets: -1` to write all jets passing cuts, and record the resampling method in the output metadata [#153](https://github.com/umami-hep/umami-preprocessing/pull/153) diff --git a/docs/run.md b/docs/run.md index 3ebad7c3..204d0bab 100644 --- a/docs/run.md +++ b/docs/run.md @@ -157,6 +157,8 @@ preprocess --config upp/configs/MetadataRW/metadata.yaml --metadata The downstream reweighting stage automatically uses the `physicalWeight` column when present, and otherwise falls back to uniform weights. +On this route the plotting stage (`--plot`) also detects `physicalWeight` and the rw-merge weight columns and weights its histograms by their product (`physicalWeight` is capped at the reweight `WEIGHT_CAP` for consistency with the reweight histograms). When neither column is present (the default resampling route) plots remain unweighted. + #### Full-ntuple passthrough (`global.keep_all_variables`) By default the split and rw-merge outputs only keep the variables declared in your `variables` config. Set `keep_all_variables: true` under `global` to instead preserve **all** top-level HDF5 datasets (e.g. `tracks`, `towers`, `flow`) and full `jets` fields in the train/val/test outputs, while still appending the reweight columns on `jets`: diff --git a/tests/unit/stages/test_plotting.py b/tests/unit/stages/test_plotting.py index 54420d95..68bf1601 100644 --- a/tests/unit/stages/test_plotting.py +++ b/tests/unit/stages/test_plotting.py @@ -210,6 +210,85 @@ def test_post_resampling_paths_split_mode(tmp_path): assert paths == [tmp_path / "test" / "pp_output_test*.h5"] +def test_reweight_weight_fields_detects_present_columns(): + """Check reweight-route weight columns are picked up only when present.""" + from upp.classes.reweight_config import SingleReweightConfig + + rw = SingleReweightConfig( + group="jets", + reweight_vars=["pt"], + bins={"pt": [[20_000, 250_000, 5]]}, + class_var="flavour_label", + class_target="uniform", + ) + rw_name = repr(rw) + assert rw_name == "weight_jets_pt_target_uniform_flavour_label" + config = SimpleNamespace(rw_config=SimpleNamespace(reweights=[rw])) + + # both physicalWeight and the rw-merge column present + assert plot_mod._reweight_weight_fields(config, {"physicalWeight", rw_name, "pt"}) == [ + "physicalWeight", + rw_name, + ] + # rw column absent (e.g. raw ntuple after --metadata) -> only physicalWeight + assert plot_mod._reweight_weight_fields(config, {"physicalWeight", "pt"}) == ["physicalWeight"] + # default resampling route: no rw_config and no physicalWeight -> unweighted + assert plot_mod._reweight_weight_fields(SimpleNamespace(), {"pt"}) == [] + + +def test_make_hist_applies_capped_weights(monkeypatch, tmp_path): + """Check make_hist multiplies weight columns and caps physicalWeight.""" + import numpy as np + + from upp.stages.reweight import WEIGHT_CAP + + dtype = [("pt", "f4"), ("HadronConeExclTruthLabelID", "i4"), ("physicalWeight", "f4")] + arr = np.zeros(5, dtype=dtype) + arr["pt"] = [10_000, 20_000, 30_000, 40_000, 50_000] + arr["HadronConeExclTruthLabelID"] = 5 # bjets + arr["physicalWeight"] = [1.0, 2.0, 3.0 * WEIGHT_CAP, 4.0, 5.0] + + captured = {} + + class FakeHist: + def __init__(self, **kwargs): + captured["weights"] = kwargs.get("weights") + self.bin_edges = np.array([0.0, 100.0]) + + class FakePlot: + def __init__(self, **kwargs): + pass + + def add(self, **kwargs): + pass + + def draw(self): + pass + + def make_linestyle_legend(self, **kwargs): + pass + + def savefig(self, _path): + pass + + monkeypatch.setattr(plot_mod, "Histogram", FakeHist) + monkeypatch.setattr(plot_mod, "HistogramPlot", FakePlot) + + make_hist( + stage="initial", + values_dict={"": arr}, + flavours=[Flavours["bjets"]], + variable="pt", + out_dir=tmp_path, + out_format_list=["png"], + weight_fields=["physicalWeight"], + ) + + expected = np.clip(arr["physicalWeight"].astype(float), 0.0, WEIGHT_CAP) + assert captured["weights"] is not None + assert np.allclose(captured["weights"], expected) + + def test_plot_initial_uses_split_suffix_and_plotting_jet_count(monkeypatch, tmp_path): """Check initial plot calls include split-specific suffixes and plotting counts.""" diff --git a/upp/stages/plot.py b/upp/stages/plot.py index ee0c1a1b..42501854 100644 --- a/upp/stages/plot.py +++ b/upp/stages/plot.py @@ -1,5 +1,6 @@ from __future__ import annotations +import glob import logging as log import re from dataclasses import dataclass @@ -7,6 +8,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +import h5py +import numpy as np from ftag import Cuts from ftag.hdf5 import H5Reader from puma import Histogram, HistogramPlot @@ -418,6 +421,62 @@ def _stitching_regions(regions: list[PlotRegion], pt_variable: str | None) -> li return stitching_regions +def _available_jet_fields(config: PreprocessingConfig, in_paths: Any) -> set[str]: + """Return the jet field names present in the first matching input file. + + Parameters + ---------- + config : PreprocessingConfig + Active preprocessing configuration. + in_paths : Any + Input HDF5 file path, glob, or list of paths. + + Returns + ------- + set[str] + Jet dtype field names of the first existing file, or an empty set. + """ + paths = in_paths if isinstance(in_paths, list) else [in_paths] + for pattern in paths: + for fpath in sorted(glob.glob(str(pattern))): + with h5py.File(fpath, "r") as f: + if config.jets_name in f: + return set(f[config.jets_name].dtype.names or ()) + return set() + + +def _reweight_weight_fields(config: PreprocessingConfig, available: set[str]) -> list[str]: + """Return per-jet weight columns to apply on the reweight route. + + Combines ``physicalWeight`` (physics weight from --metadata) with the + reweight columns written by --rw-merge. Only columns present in ``available`` + are returned, so the default resampling route (which has neither) stays + unweighted. + + Parameters + ---------- + config : PreprocessingConfig + Active preprocessing configuration. + available : set[str] + Jet field names present in the input being plotted. + + Returns + ------- + list[str] + Weight column names to multiply together when histogramming. + """ + fields = [] + if "physicalWeight" in available: + fields.append("physicalWeight") + rw_config = getattr(config, "rw_config", None) + if rw_config is not None: + for rw in rw_config.reweights: + name = repr(rw) + if name in available: + fields.append(name) + return fields + + def _load_jets(config: PreprocessingConfig, in_paths: Any, vars_to_load: list[str]) -> Any: """Load jet variables for plotting. @@ -461,6 +520,7 @@ def make_hist( selection_cuts: Cuts | None = None, atlas_second_tag: str | None = None, plotting: PlottingConfig | None = None, + weight_fields: list[str] | None = None, ) -> None: """Make a flavour-split histogram for one variable. @@ -502,8 +562,14 @@ def make_hist( centre-of-mass energy label is shown. plotting : PlottingConfig | None, optional Plot labels and style settings. If ``None``, use the defaults. + weight_fields : list[str] | None, optional + Jet weight columns multiplied per jet on the reweight route (e.g. + ``physicalWeight`` and the rw-merge weight column). ``physicalWeight`` is + capped at ``WEIGHT_CAP`` to match the reweight histograms. If ``None`` or + empty, histograms are unweighted (default resampling behaviour). """ from upp.classes.plotting_config import PlottingConfig + from upp.stages.reweight import WEIGHT_CAP selection_cuts = selection_cuts or Cuts.empty() plotting = plotting or PlottingConfig() @@ -541,9 +607,20 @@ def make_hist( selected_values = cuts(values).values histo_values = _display_values(variable, selected_values[variable]) + # Reweight route: weight per jet by capped physicalWeight x rw-merge columns. + histo_weights = None + if weight_fields: + histo_weights = np.ones(len(selected_values), dtype=np.float64) + for field in weight_fields: + column = selected_values[field].astype(np.float64) + if field == "physicalWeight": + column = np.clip(column, 0.0, WEIGHT_CAP) + histo_weights *= column + # Get the histogram object histo = Histogram( values=histo_values, + weights=histo_weights, bins=plotting.bins, bins_range=bins_range, norm=plotting.norm, @@ -600,8 +677,12 @@ def _plot_initial(config: PreprocessingConfig) -> None: for flavour in config.components.flavours: vars_to_load += flavour.cuts.variables + in_paths = list(sample.path) + weight_fields = _reweight_weight_fields(config, _available_jet_fields(config, in_paths)) + vars_to_load += weight_fields + stage_status = "Pre Reweighting" if weight_fields else "Pre Resampling" values_dict = { - sample.name: _load_jets(config, list(sample.path), vars_to_load), + sample.name: _load_jets(config, in_paths, vars_to_load), } pt_range = _pt_bounds_from_cuts(selection_cuts, pt_var) if pt_var else None @@ -627,10 +708,11 @@ def _plot_initial(config: PreprocessingConfig) -> None: num_jets=_plotting_num_jets(config, region_components.num_jets) if config.plotting.show_num_jets else None, - resampling_status="Pre Resampling", + resampling_status=stage_status, ), plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, ) @@ -680,6 +762,9 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: if full_region := _full_region(base_regions, pt_var): plot_regions.append(full_region) + in_paths = _post_resampling_paths(config, stage) + weight_fields = _reweight_weight_fields(config, _available_jet_fields(config, in_paths)) + sample_names = [sample.name for sample in config.components.samples] atlas_second_tag = _atlas_second_tag( *sample_names, @@ -687,15 +772,16 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: num_jets=_plotting_num_jets(config, config.components.num_jets) if config.plotting.show_num_jets else None, - resampling_status="Post Resampling", + resampling_status="Post Reweighting" if weight_fields else "Post Resampling", ) vars_to_load = list(config.sampl_cfg.vars) + ["flavour_label"] for region in plot_regions: vars_to_load += region.cuts.variables + vars_to_load += weight_fields values_dict = { - "": _load_jets(config, _post_resampling_paths(config, stage), vars_to_load), + "": _load_jets(config, in_paths, vars_to_load), } for variable in config.sampl_cfg.vars: @@ -718,6 +804,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: atlas_second_tag=atlas_second_tag, plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, ) if _is_pt_variable(variable): @@ -735,6 +822,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: atlas_second_tag=atlas_second_tag, plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, + weight_fields=weight_fields, ) From 20548f27f961a2f972b17a871ce969394801b2bb Mon Sep 17 00:00:00 2001 From: yuanda Date: Thu, 25 Jun 2026 15:25:39 +0200 Subject: [PATCH 09/11] Address PR #142 latest review feedback - reweight: use masked np.divide so empty bins no longer re-emit the divide-by-zero RuntimeWarning that #140 removed - configs: restore mcCampaignYear resampling bins in xbb / xbb-gn3x - test_run: only inject physicalWeight on demand so the default (unweighted) resampling route stays exercised end-to-end - metadata_injector: preserve source compression/chunks/shuffle instead of hard-coding gzip on the recreated jets dataset Co-authored-by: Cursor --- tests/integration/test_run.py | 24 +++++++++++++----------- upp/configs/xbb-gn3x.yaml | 2 ++ upp/configs/xbb.yaml | 2 ++ upp/stages/metadata_injector.py | 9 ++++++++- upp/stages/reweight.py | 6 +++--- 5 files changed, 28 insertions(+), 15 deletions(-) diff --git a/tests/integration/test_run.py b/tests/integration/test_run.py index 055a4ebc..5737b14e 100644 --- a/tests/integration/test_run.py +++ b/tests/integration/test_run.py @@ -17,18 +17,20 @@ class TestClass: - def generate_mock(self, out_file, N=100_000): + def generate_mock(self, out_file, N=100_000, with_physical_weight=False): + # Inject physicalWeight only on demand; the weighted route is covered in test_run_rw.py. _, f = get_mock_file(num_jets=N, fname=out_file) - jets = f["jets"][:] - if "physicalWeight" not in jets.dtype.names: - jets2 = rfn.append_fields( - jets, - "physicalWeight", - np.ones(jets.shape[0], dtype="f4"), - usemask=False, - ) - del f["jets"] - f.create_dataset("jets", data=jets2) + if with_physical_weight: + jets = f["jets"][:] + if "physicalWeight" not in jets.dtype.names: + jets2 = rfn.append_fields( + jets, + "physicalWeight", + np.ones(jets.shape[0], dtype="f4"), + usemask=False, + ) + del f["jets"] + f.create_dataset("jets", data=jets2) f.close() def setup_method(self, method): diff --git a/upp/configs/xbb-gn3x.yaml b/upp/configs/xbb-gn3x.yaml index 482957ca..50fd9ff3 100644 --- a/upp/configs/xbb-gn3x.yaml +++ b/upp/configs/xbb-gn3x.yaml @@ -111,6 +111,8 @@ resampling: bins: [[0, 2, 20]] mass: bins: [[50_000, 300_000, 50]] + mcCampaignYear: + bins: [[2015, 2026, 11]] # note: sensible defaults are defined in the PreprocessingConfig constructor global: diff --git a/upp/configs/xbb.yaml b/upp/configs/xbb.yaml index dd5b379e..866e41d7 100644 --- a/upp/configs/xbb.yaml +++ b/upp/configs/xbb.yaml @@ -65,6 +65,8 @@ resampling: bins: [[0, 2, 20]] mass: bins: [[50_000, 300_000, 50]] + mcCampaignYear: + bins: [[2015, 2026, 11]] # note: sensible defaults are defined in the PreprocessingConfig constructor global: diff --git a/upp/stages/metadata_injector.py b/upp/stages/metadata_injector.py index b4b9b6c9..4df2d5de 100644 --- a/upp/stages/metadata_injector.py +++ b/upp/stages/metadata_injector.py @@ -68,6 +68,13 @@ def run(self): # 3. Read original jets and their attributes old_jets_ds = f["jets"] original_attrs = dict(old_jets_ds.attrs) # Backup all attribute metadata + # Preserve source storage layout so rewrite keeps compression/chunking. + create_kwargs = { + "compression": old_jets_ds.compression, + "compression_opts": old_jets_ds.compression_opts, + "chunks": old_jets_ds.chunks, + "shuffle": old_jets_ds.shuffle, + } jets_data = old_jets_ds[:] # 4. Calculate Physical Weight @@ -88,7 +95,7 @@ def run(self): # 6. Delete and recreate the dataset while restoring attributes del f["jets"] - new_ds = f.create_dataset("jets", data=updated_jets, compression="gzip") + new_ds = f.create_dataset("jets", data=updated_jets, **create_kwargs) # Restore all original attributes (e.g., descriptions for eventNumber, etc.) for k, v in original_attrs.items(): diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index d872d823..afb45fde 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -293,9 +293,9 @@ def calculate_weights( idx_below_min = None for cls, hist in all_histograms[rw_group][rw_rep]["histograms"].items(): this_idx_below_min = hist == 0 # | (all_targets[rw_group][rw_rep] == 0) - output_weights[rw_group][rw_rep]["weights"][cls] = np.where( - hist > 0, all_targets[rw_group][rw_rep] / hist, 0 - ) + weights = np.zeros_like(hist, dtype=float) + np.divide(all_targets[rw_group][rw_rep], hist, out=weights, where=hist > 0) + output_weights[rw_group][rw_rep]["weights"][cls] = weights if idx_below_min is None: idx_below_min = this_idx_below_min else: From 5d559b53307730e19d1a6ed4dabc2b83a1d84510 Mon Sep 17 00:00:00 2001 From: yuanda Date: Wed, 29 Jul 2026 07:26:32 +0200 Subject: [PATCH 10/11] Make the weight cap configurable via reweighting.weight_cap Co-authored-by: Cursor --- upp/classes/reweight_config.py | 6 ++++++ upp/stages/plot.py | 12 ++++++++++-- upp/stages/reweight.py | 8 +++++--- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/upp/classes/reweight_config.py b/upp/classes/reweight_config.py index a488621c..42f41b17 100644 --- a/upp/classes/reweight_config.py +++ b/upp/classes/reweight_config.py @@ -11,12 +11,18 @@ class ReweightConfig: # Number of jets to estimate, if None, use the global num jets estimate num_jets_estimate: None | int = None merge_num_proc: int = 1 # Number of processes to use for merging + # Upper bound for physical weights and reweight factors, + # which can span many orders of magnitude + weight_cap: float = 1e4 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.weight_cap <= 0: + raise ValueError("weight_cap must be a positive number") + parsed_reweights = [] for rw in self.reweights: parsed_reweights.append(SingleReweightConfig(**rw)) diff --git a/upp/stages/plot.py b/upp/stages/plot.py index 42501854..c8aad120 100644 --- a/upp/stages/plot.py +++ b/upp/stages/plot.py @@ -521,6 +521,7 @@ def make_hist( atlas_second_tag: str | None = None, plotting: PlottingConfig | None = None, weight_fields: list[str] | None = None, + weight_cap: float | None = None, ) -> None: """Make a flavour-split histogram for one variable. @@ -565,14 +566,19 @@ def make_hist( weight_fields : list[str] | None, optional Jet weight columns multiplied per jet on the reweight route (e.g. ``physicalWeight`` and the rw-merge weight column). ``physicalWeight`` is - capped at ``WEIGHT_CAP`` to match the reweight histograms. If ``None`` or + capped at ``weight_cap`` to match the reweight histograms. If ``None`` or empty, histograms are unweighted (default resampling behaviour). + weight_cap : float | None, optional + Upper bound applied to ``physicalWeight``. If ``None``, the default + ``WEIGHT_CAP`` from the reweight stage is used. """ from upp.classes.plotting_config import PlottingConfig from upp.stages.reweight import WEIGHT_CAP selection_cuts = selection_cuts or Cuts.empty() plotting = plotting or PlottingConfig() + if weight_cap is None: + weight_cap = WEIGHT_CAP # Setup the histogram plot = HistogramPlot( @@ -614,7 +620,7 @@ def make_hist( for field in weight_fields: column = selected_values[field].astype(np.float64) if field == "physicalWeight": - column = np.clip(column, 0.0, WEIGHT_CAP) + column = np.clip(column, 0.0, weight_cap) histo_weights *= column # Get the histogram object @@ -805,6 +811,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, weight_fields=weight_fields, + weight_cap=config.rw_config.weight_cap if config.rw_config else None, ) if _is_pt_variable(variable): @@ -823,6 +830,7 @@ def _plot_post_resampling(config: PreprocessingConfig, stage: str) -> None: plotting=config.plotting, out_dir=config.out_dir / config.plotting.output_directory, weight_fields=weight_fields, + weight_cap=config.rw_config.weight_cap if config.rw_config else None, ) diff --git a/upp/stages/reweight.py b/upp/stages/reweight.py index afb45fde..4b49e3b6 100644 --- a/upp/stages/reweight.py +++ b/upp/stages/reweight.py @@ -14,7 +14,8 @@ from upp.classes.preprocessing_config import PreprocessingConfig from upp.stages.hist import bin_jets -# Upper bound for physical weights and reweight factors, which can span many orders of magnitude. +# Default upper bound for physical weights and reweight factors, which can span +# many orders of magnitude. Configurable via `reweighting.weight_cap`. WEIGHT_CAP = 1e4 @@ -26,6 +27,7 @@ def __init__(self, config: PreprocessingConfig): assert self.rw_config is not None, ( "Reweighting configuration is not set in the preprocessing config" ) + self.weight_cap = self.rw_config.weight_cap self.organised_components_config = ( Path(config.base_dir) / "split-components/organised-components.yaml" ) @@ -185,7 +187,7 @@ def calculate_weights( # Use physicalWeight if present, else uniform weights. if "physicalWeight" in data_masked.dtype.names: w = np.asarray(data_masked["physicalWeight"], dtype=np.float64) - w = np.clip(w, 0, WEIGHT_CAP) + w = np.clip(w, 0, self.weight_cap) else: w = np.ones(mask.sum(), dtype=float) @@ -311,7 +313,7 @@ def calculate_weights( np.clip( output_weights[rw_group][rw_rep]["weights"][cls], 0, - WEIGHT_CAP, + self.weight_cap, out=output_weights[rw_group][rw_rep]["weights"][cls], ) From 7e95f87c400fa5c1c47a21b466e83f69c3c4b621 Mon Sep 17 00:00:00 2001 From: yuanda Date: Wed, 29 Jul 2026 07:31:00 +0200 Subject: [PATCH 11/11] Add requests dependency needed by ftag.find_metadata Co-authored-by: Cursor --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 7c65c075..4b1c7726 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "puma-hep==0.5.3", "pyyaml-include==1.3", "PyYAML>=6.0.2", + # Needed by ftag.find_metadata (missing from atlas-ftag-tools' own deps) + "requests>=2.32", "rich>=14.1.0", "scipy>=1.15.3", ]