umami-preprocessing 0.2.6__tar.gz → 0.2.7__tar.gz

{umami_preprocessing-0.2.6/umami_preprocessing.egg-info → umami_preprocessing-0.2.7}/PKG-INFO

@@ -1,26 +1,33 @@
 Metadata-Version: 2.4
 Name: umami-preprocessing
-Version: 0.2.6
-Summary: Preprocessing for jet tagging
+Version: 0.2.7
+Summary: ATLAS Flavour Tagging Preprocessing - Umami PreProcessing (UPP)
+Author: Alexander Froch
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing
-Requires-Python: <3.12,>=3.8
+Project-URL: Issue Tracker, https://github.com/umami-hep/umami-preprocessing/issues
+Requires-Python: <3.12,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: atlas-ftag-tools==0.2.15
-Requires-Dist: dotmap==1.3.30
-Requires-Dist: puma-hep==0.4.10
+Requires-Dist: atlas-ftag-tools==0.2.17
+Requires-Dist: dotmap>=1.3.30
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: puma-hep==0.4.11
 Requires-Dist: pyyaml-include==1.3
-Requires-Dist: PyYAML>=6.0.1
-Requires-Dist: rich==12.6.0
-Requires-Dist: scipy>=1.15.2
+Requires-Dist: PyYAML>=6.0.2
+Requires-Dist: rich>=14.1.0
+Requires-Dist: scipy>=1.15.3
 Provides-Extra: dev
-Requires-Dist: mypy==1.11.2; extra == "dev"
-Requires-Dist: pre-commit==3.5.0; extra == "dev"
-Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
-Requires-Dist: pytest-mock==3.11.1; extra == "dev"
-Requires-Dist: pytest>=7.2.2; extra == "dev"
-Requires-Dist: ruff==0.6.2; extra == "dev"
+Requires-Dist: coverage>=7.10.6; extra == "dev"
+Requires-Dist: ipykernel>=6.30.1; extra == "dev"
+Requires-Dist: mypy>=1.18.1; extra == "dev"
+Requires-Dist: pre-commit>=4.3.0; extra == "dev"
+Requires-Dist: pydoclint>=0.7.3; extra == "dev"
+Requires-Dist: pytest_notebook>=0.10.0; extra == "dev"
+Requires-Dist: pytest-cov>=7.0.0; extra == "dev"
+Requires-Dist: pytest-randomly>=4.0.1; extra == "dev"
+Requires-Dist: pytest>=8.4.2; extra == "dev"
+Requires-Dist: ruff>=0.13.0; extra == "dev"
 Dynamic: license-file
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/pyproject.toml

@@ -1,33 +1,40 @@
 [project]
 name = "umami-preprocessing"
-description = "Preprocessing for jet tagging"
+description = "ATLAS Flavour Tagging Preprocessing - Umami PreProcessing (UPP)"
+authors = [{name="Alexander Froch"}]
 dynamic = ["version"]
 license = {text = "MIT"}
 readme = "README.md"
-requires-python = "<3.12,>=3.8"
+requires-python = ">=3.10,<3.12"
 
 dependencies = [
-  "atlas-ftag-tools==0.2.15",
-  "dotmap==1.3.30",
-  "puma-hep==0.4.10",
+  "atlas-ftag-tools==0.2.17",
+  "dotmap>=1.3.30",
+  "numpy>=2.2.6",
+  "puma-hep==0.4.11",
   "pyyaml-include==1.3",
-  "PyYAML>=6.0.1",
-  "rich==12.6.0",
-  "scipy>=1.15.2",
+  "PyYAML>=6.0.2",
+  "rich>=14.1.0",
+  "scipy>=1.15.3",
 ]
 
 [project.optional-dependencies]
 dev = [
-  "mypy==1.11.2",
-  "pre-commit==3.5.0",
-  "pytest-cov>=4.0.0",
-  "pytest-mock==3.11.1",
-  "pytest>=7.2.2",
-  "ruff==0.6.2",
+  "coverage>=7.10.6",
+  "ipykernel>=6.30.1",
+  "mypy>=1.18.1",
+  "pre-commit>=4.3.0",
+  "pydoclint>=0.7.3",
+  "pytest_notebook>=0.10.0",
+  "pytest-cov>=7.0.0",
+  "pytest-randomly>=4.0.1",
+  "pytest>=8.4.2",
+  "ruff>=0.13.0",
 ]
 
 [project.urls]
 "Homepage" = "https://github.com/umami-hep/umami-preprocessing"
+"Issue Tracker" = "https://github.com/umami-hep/umami-preprocessing/issues"
 
 [project.scripts]
 preprocess = "upp.main:main"

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7/umami_preprocessing.egg-info}/PKG-INFO

@@ -1,26 +1,33 @@
 Metadata-Version: 2.4
 Name: umami-preprocessing
-Version: 0.2.6
-Summary: Preprocessing for jet tagging
+Version: 0.2.7
+Summary: ATLAS Flavour Tagging Preprocessing - Umami PreProcessing (UPP)
+Author: Alexander Froch
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing
-Requires-Python: <3.12,>=3.8
+Project-URL: Issue Tracker, https://github.com/umami-hep/umami-preprocessing/issues
+Requires-Python: <3.12,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: atlas-ftag-tools==0.2.15
-Requires-Dist: dotmap==1.3.30
-Requires-Dist: puma-hep==0.4.10
+Requires-Dist: atlas-ftag-tools==0.2.17
+Requires-Dist: dotmap>=1.3.30
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: puma-hep==0.4.11
 Requires-Dist: pyyaml-include==1.3
-Requires-Dist: PyYAML>=6.0.1
-Requires-Dist: rich==12.6.0
-Requires-Dist: scipy>=1.15.2
+Requires-Dist: PyYAML>=6.0.2
+Requires-Dist: rich>=14.1.0
+Requires-Dist: scipy>=1.15.3
 Provides-Extra: dev
-Requires-Dist: mypy==1.11.2; extra == "dev"
-Requires-Dist: pre-commit==3.5.0; extra == "dev"
-Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
-Requires-Dist: pytest-mock==3.11.1; extra == "dev"
-Requires-Dist: pytest>=7.2.2; extra == "dev"
-Requires-Dist: ruff==0.6.2; extra == "dev"
+Requires-Dist: coverage>=7.10.6; extra == "dev"
+Requires-Dist: ipykernel>=6.30.1; extra == "dev"
+Requires-Dist: mypy>=1.18.1; extra == "dev"
+Requires-Dist: pre-commit>=4.3.0; extra == "dev"
+Requires-Dist: pydoclint>=0.7.3; extra == "dev"
+Requires-Dist: pytest_notebook>=0.10.0; extra == "dev"
+Requires-Dist: pytest-cov>=7.0.0; extra == "dev"
+Requires-Dist: pytest-randomly>=4.0.1; extra == "dev"
+Requires-Dist: pytest>=8.4.2; extra == "dev"
+Requires-Dist: ruff>=0.13.0; extra == "dev"
 Dynamic: license-file
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/umami_preprocessing.egg-info/SOURCES.txt

@@ -15,6 +15,7 @@ upp/classes/components.py
 upp/classes/preprocessing_config.py
 upp/classes/region.py
 upp/classes/resampling_config.py
+upp/classes/reweight_config.py
 upp/classes/variable_config.py
 upp/stages/__init__.py
 upp/stages/hist.py
@@ -23,6 +24,9 @@ upp/stages/merging.py
 upp/stages/normalisation.py
 upp/stages/plot.py
 upp/stages/resampling.py
+upp/stages/reweight.py
+upp/stages/rw_merge.py
+upp/stages/split_containers.py
 upp/utils/__init__.py
 upp/utils/check_input_samples.py
 upp/utils/logger.py

umami_preprocessing-0.2.7/umami_preprocessing.egg-info/requires.txt

@@ -0,0 +1,20 @@
+atlas-ftag-tools==0.2.17
+dotmap>=1.3.30
+numpy>=2.2.6
+puma-hep==0.4.11
+pyyaml-include==1.3
+PyYAML>=6.0.2
+rich>=14.1.0
+scipy>=1.15.3
+
+[dev]
+coverage>=7.10.6
+ipykernel>=6.30.1
+mypy>=1.18.1
+pre-commit>=4.3.0
+pydoclint>=0.7.3
+pytest_notebook>=0.10.0
+pytest-cov>=7.0.0
+pytest-randomly>=4.0.1
+pytest>=8.4.2
+ruff>=0.13.0

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.6"
+__version__ = "v0.2.7"
 
 from . import classes, stages, utils
 from .main import run_pp

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/classes/components.py

@@ -336,6 +336,7 @@ class Components:
                 pattern=pattern,
                 ntuple_dir=config.ntuple_dir,
                 name=component["sample"]["name"],
+                skip_checks=config.skip_checks,
             )
 
             # Create the Component instances for the different flavours
@@ -360,7 +361,7 @@ class Components:
         components = cls(component_list)
 
         # Check the flavour ratios
-        if config.sampl_cfg.method is not None:
+        if config.sampl_cfg and config.sampl_cfg.method is not None:
             components.check_flavour_ratios()
 
         return components
@@ -510,7 +511,7 @@ class Components:
     def __getitem__(self, index):
         if isinstance(index, int):
             return self.components[index]
-        if isinstance(index, (str, Label)):
+        if isinstance(index, str | Label):
             return self.components[self.flavours.index(index)]
 
     def __len__(self):

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/classes/preprocessing_config.py

@@ -20,6 +20,7 @@ from yamlinclude import YamlIncludeConstructor
 from upp import __version__
 from upp.classes.components import Components
 from upp.classes.resampling_config import ResamplingConfig
+from upp.classes.reweight_config import ReweightConfig
 from upp.classes.variable_config import VariableConfig
 from upp.utils.tools import path_append
 
@@ -114,6 +115,8 @@ class PreprocessingConfig:
         than this number, the final h5 output files are splitted in multiple smaller
         files with this number of jets per file. By default None which produces one
         huge output file.
+    skip_checks : bool, optional
+        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
     """
@@ -137,6 +140,7 @@ class PreprocessingConfig:
     flavour_config: Path | None = None
     flavour_category: str = "standard"
     num_jets_per_output_file: int | None = None
+    skip_checks: bool = False
     skip_config_copy: bool = False
 
     def __post_init__(self):
@@ -154,11 +158,10 @@ class PreprocessingConfig:
         for field in dataclasses.fields(self):
             if field.type == "Path" and field.name != "out_fname" and field.name != "base_dir":
                 setattr(self, field.name, self.get_path(Path(getattr(self, field.name))))
-        if not self.ntuple_dir.exists():
+        if not self.ntuple_dir.exists() and not self.skip_checks:
             raise FileNotFoundError(f"Path {self.ntuple_dir} does not exist")
         self.components_dir = self.components_dir / self.split
         self.out_fname = self.out_dir / path_append(self.out_fname, self.split)
-
         # Define the content of the flavour label container
         if self.flavour_config:
             self.flavour_cont = LabelContainer.from_yaml(
@@ -177,12 +180,15 @@ class PreprocessingConfig:
                 "flavours! If you want to use your own flavour config yaml file, please "
                 "provide flavour_config!"
             )
-
         # configure classes
-        sampl_cfg = copy(self.config["resampling"])
-        if self.is_test:
-            sampl_cfg["method"] = None
-        self.sampl_cfg = ResamplingConfig(**sampl_cfg)
+        if sampl_cfg := self.config.get("resampling", None):
+            sampl_cfg = copy(sampl_cfg)
+            if self.is_test:
+                sampl_cfg["method"] = None
+            self.sampl_cfg = ResamplingConfig(**sampl_cfg)
+        else:
+            self.sampl_cfg = None
+
         self.components = Components.from_config(self)
 
         # get track selectors
@@ -196,13 +202,21 @@ class PreprocessingConfig:
         self.variables = VariableConfig(
             self.config["variables"], self.jets_name, self.is_test, selectors
         )
-        self.variables = self.variables.add_jet_vars(
-            list(self.config["resampling"]["variables"].keys()), "labels"
-        )
+        if self.sampl_cfg is not None:
+            self.variables = self.variables.add_jet_vars(
+                list(self.config["resampling"]["variables"].keys()), "labels"
+            )
         self.transform = (
             Transform(**self.config["transform"]) if "transform" in self.config else None
         )
 
+        self.rw_config = (
+            ReweightConfig(
+                **self.config["reweighting"],
+            )
+            if "reweighting" in self.config
+            else None
+        )
         # reproducibility
         self.git_hash = get_git_hash(Path(__file__).parent)
         if self.git_hash is None:
@@ -214,7 +228,13 @@ class PreprocessingConfig:
             self.copy_config()
 
     @classmethod
-    def from_file(cls, config_path: Path, split: Split, skip_config_copy: bool = False):
+    def from_file(
+        cls,
+        config_path: Path,
+        split: Split,
+        skip_checks: bool = False,
+        skip_config_copy: bool = False,
+    ):
         if not config_path.exists():
             raise FileNotFoundError(f"{config_path} does not exist - check your --config arg")
         with open(config_path) as file:
@@ -225,6 +245,7 @@ class PreprocessingConfig:
                 config=config,
                 skip_config_copy=skip_config_copy,
                 **config["global"],
+                skip_checks=skip_checks,
             )
 
     def get_path(self, path: Path):
@@ -238,7 +259,7 @@ class PreprocessingConfig:
     def global_cuts(self):
         cuts_list = self.config["global_cuts"].get("common", [])
         cuts_list += self.config["global_cuts"][self.split]
-        if not self.is_test:
+        if not self.is_test and self.config.get("resampling", None) is not None:
             for resampling_var, cfg in self.config["resampling"]["variables"].items():
                 cuts_list.append([resampling_var, ">", cfg["bins"][0][0]])
                 cuts_list.append([resampling_var, "<", cfg["bins"][-1][1]])
@@ -357,3 +378,34 @@ class PreprocessingConfig:
                 f"Option value {option} is not supported! "
                 "Only resampled and resampled_scaled_shuffled are."
             )
+
+    # Static because otherwise config paths end up getting messed up
+    @staticmethod
+    def get_input_files_with_split_components(config_path):
+        """Return a nested dictionary of the form.
+
+        {
+            "container_1" : {
+                "train_bjets" : list[str] : -> which are cuts
+                "test_bjets" : list[str] : -> which are cuts
+                ...
+                "test_cjets" : list[str] : -> which are cuts
+
+            },
+            "container_2" : ...
+        }
+        which represents all splits and components to be run for each input container
+        """
+        containers_with_splits = {}
+
+        for split in ["train", "val", "test"]:
+            split_config = PreprocessingConfig.from_file(config_path, split, skip_checks=True)
+
+            for component in split_config.components.components:
+                for container in component.sample.pattern:
+                    if container not in containers_with_splits:
+                        containers_with_splits[container] = {}
+                    component_cuts = component.cuts
+                    containers_with_splits[container][f"{split}_{component.name}"] = component_cuts
+
+        return containers_with_splits

umami_preprocessing-0.2.7/upp/classes/reweight_config.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+@dataclass
+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
+    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")
+
+        parsed_reweights = []
+        for rw in self.reweights:
+            parsed_reweights.append(SingleReweightConfig(**rw))
+        self.reweights = parsed_reweights
+
+
+@dataclass
+class SingleReweightConfig:
+    group: str  # The group our variables in the h5 file are in
+    reweight_vars: list[str]  # The variables we want to reweight
+    bins: list[np.ndarray]  # The bins we want to use for the reweighting
+    class_var: str  # The variable which contains the label we resample over, e.g. flavour
+    class_target: int | tuple | str | None = None
+    add_overflow: bool = True  # Whether to add overflow bins
+
+    target_hist_func: Callable | None = None
+    target_hist_func_name: str | None = None
+
+    # TODO - this is the same as in resampling, maybe can cleanup
+    def get_bins_x(self, bins_x, upscale=1):
+        flat_bins = []
+        for i, sub_bins_x in enumerate(bins_x):
+            start, stop, nbins = sub_bins_x
+            b = np.linspace(start, stop, nbins * upscale + 1)
+            if i > 0:
+                b = b[1:]
+            flat_bins.append(b)
+        if self.add_overflow:
+            flat_bins = [np.array([-np.inf])] + flat_bins + [np.array([np.inf])]
+        return np.concatenate(flat_bins)
+
+    @property
+    def flat_bins(self):
+        return [self.get_bins_x(self.bins[k]) for k in self.reweight_vars]
+
+    def __post_init__(self):
+        if isinstance(self.class_target, str) and self.class_target not in [
+            "mean",
+            "min",
+            "max",
+            "uniform",
+        ]:
+            raise ValueError("class_target must be either 'mean', 'min', 'max' or an integer")
+
+        if self.target_hist_func is not None and self.target_hist_func_name is None:
+            self.target_hist_func_name = self.target_hist_func.__name__
+
+    def __repr__(self):
+        target_str = "target_"
+        if self.target_hist_func_name is not None:
+            target_str += f"{self.target_hist_func_name}_"
+        if self.class_target is not None:
+            if isinstance(self.class_target, list | tuple):
+                target_str += "_".join(map(str, self.class_target))
+            else:
+                target_str += f"{self.class_target}_{self.class_var}"
+        else:
+            target_str += "none"
+        return f"weight_{self.group}_{'_'.join(self.reweight_vars)}_{target_str}"

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/main.py

@@ -23,6 +23,9 @@ from upp.stages.merging import Merging
 from upp.stages.normalisation import Normalisation
 from upp.stages.plot import plot_resampling_dists
 from upp.stages.resampling import Resampling
+from upp.stages.reweight import Reweight
+from upp.stages.rw_merge import RWMerge
+from upp.stages.split_containers import SplitContainers
 from upp.utils.check_input_samples import run_input_sample_check
 from upp.utils.logger import setup_logger
 
@@ -116,6 +119,28 @@ def parse_args(args: Any) -> argparse.Namespace:
         default=None,
         help="Component which is processed during --prep",
     )
+    parser.add_argument(
+        "--split-components",
+        action="store_true",
+        default=False,
+        help="Split containers into components",
+    )
+    parser.add_argument(
+        "--reweight", "--rw", action="store_true", default=False, help="Run the reweighting stage"
+    )
+    parser.add_argument(
+        "--rw-merge", "--rwm", action="store_true", default=False, help="Run the reweighting stage"
+    )
+    parser.add_argument(
+        "--rw-merge-idx",
+        "--rwm-idx",
+        type=str,
+        default=None,
+        help=(
+            "Commar seperated pair of indices representing the range of output "
+            "files to create, e.g '0,10' will create files 0 to 9"
+        ),
+    )
     parser.add_argument(
         "--region",
         default=None,
@@ -126,10 +151,37 @@ def parse_args(args: Any) -> argparse.Namespace:
         action="store_true",
         help="Skip the inital input sample check",
     )
+    parser.add_argument(
+        "--grid", action="store_true", help="Use when running the split stage on the grid. "
+    )
+    parser.add_argument(
+        "--container",
+        default=None,
+        type=str,
+        help="Container to use during the 'split-containers' stage. "
+        "If not specified, all containers in the config will be used.",
+    )
+    parser.add_argument(
+        "--files",
+        default=None,
+        help="comma-separated list of files to use during the 'split-containers' stage ",
+    )
 
     args = parser.parse_args(args)
     d = vars(args)
-    ignore = ["config", "split", "component", "region"]
+    ignore = [
+        "config",
+        "split",
+        "component",
+        "region",
+        "container",
+        "files",
+        "grid",
+        "split_components",
+        "reweight",
+        "rw_merge",
+        "rw_merge_idx",
+    ]
     if not any(v for a, v in d.items() if a not in ignore):
         for v in d:
             if v not in ignore and d[v] is None:
@@ -151,10 +203,24 @@ def run_pp(args: argparse.Namespace) -> None:
     log.info("[bold green]Starting preprocessing...")
     start = datetime.now()
     log.info(f"Start time: {start.strftime('%Y-%m-%d %H:%M:%S')}")
-
     # load config
-    config = PreprocessingConfig.from_file(args.config, args.split)
+    config = PreprocessingConfig.from_file(args.config, args.split, skip_checks=args.grid)
 
+    if args.split_components:
+        log.info("Splitting containers...")
+        split = SplitContainers(args.config)
+        split.run(
+            args.container,
+            args.files,
+            "." if args.grid else None,  # Use current directory if not on grid
+        )
+        # 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)
+        reweight.run()
     # create virtual datasets and pdf files
     if args.prep:
         # Check the input samples sizes
@@ -180,6 +246,16 @@ def run_pp(args: argparse.Namespace) -> None:
     if args.merge:
         merging = Merging(config)
         merging.run()
+    if args.rw_merge:
+        if args.rw_merge_idx:
+            rw_merge_idx = args.rw_merge_idx
+            assert "," in rw_merge_idx, "rw-merge-idx must be a comma-separated pair of indices"
+            rw_merge_idx = tuple(map(int, rw_merge_idx.split(",")))
+            assert len(rw_merge_idx) == 2, "rw-merge-idx must be a pair of indices"
+        else:
+            rw_merge_idx = None
+        rw_merge = RWMerge(config, rw_merge_idx)
+        rw_merge.run()
 
     # run the normalisation
     if args.norm and args.split == "train":
@@ -209,10 +285,10 @@ def main(args: Any | None = None) -> None:
         d = vars(args)
         for split in ["train", "val", "test"]:
             d["split"] = split
-            log.info(f"[bold blue]{'-'*100}")
+            log.info(f"[bold blue]{'-' * 100}")
             title = f" {args.split} "
             log.info(f"[bold blue]{title:-^100}")
-            log.info(f"[bold blue]{'-'*100}")
+            log.info(f"[bold blue]{'-' * 100}")
             run_pp(args)
     else:
         run_pp(args)

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/stages/normalisation.py

@@ -191,7 +191,7 @@ class Normalisation:
                             "Class dict A has arrays of different lengths for the same"
                             " variable. This should not happen."
                         )
-                    counts_A = dict(zip(*class_dict_A[name][v]))
+                    counts_A = dict(zip(*class_dict_A[name][v], strict=False))
                     counts[i] += counts_A.get(label, 0)
                 var[v] = (labels, counts)
         return class_dict_B
@@ -233,9 +233,11 @@ class Normalisation:
         """Run the normalisation calculation."""
         title = " Computing Normalisations "
         log.info(f"[bold green]{title:-^100}")
+        if self.config.rw_config is not None:
+            fname = str(self.config.out_fname).replace(".h5", "_vds.h5")
 
         # Get the correct output names if multiple output files were written
-        if self.config.num_jets_per_output_file:
+        elif self.config.num_jets_per_output_file:
             fname = self.config.out_fname.parent / f"{self.config.out_fname.stem}*.h5"
 
         else:

{umami_preprocessing-0.2.6 → umami_preprocessing-0.2.7}/upp/stages/resampling.py

@@ -15,7 +15,8 @@ from upp.stages.interpolation import subdivide_bins, upscale_array_regionally
 from upp.utils.logger import ProgressBar
 
 if TYPE_CHECKING:  # pragma: no cover
-    from typing import Any, Generator
+    from collections.abc import Generator
+    from typing import Any
 
     from rich.progress import Progress