umami-preprocessing 0.2.2__tar.gz → 0.2.4__tar.gz

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/PKG-INFO

@@ -1,25 +1,25 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: umami-preprocessing
-Version: 0.2.2
+Version: 0.2.4
 Summary: Preprocessing for jet tagging
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing
 Requires-Python: <3.12,>=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: pyyaml-include==1.3
-Requires-Dist: PyYAML==6.0.1
+Requires-Dist: PyYAML>=6.0.1
 Requires-Dist: rich==12.6.0
-Requires-Dist: scipy==1.10.1
-Requires-Dist: puma-hep==0.4.2
-Requires-Dist: atlas-ftag-tools==0.2.8
+Requires-Dist: scipy>=1.15.2
+Requires-Dist: puma-hep==0.4.5
+Requires-Dist: atlas-ftag-tools==0.2.10
 Requires-Dist: dotmap==1.3.30
 Provides-Extra: dev
-Requires-Dist: ruff==0.1.6; extra == "dev"
-Requires-Dist: mypy==1.5.1; extra == "dev"
+Requires-Dist: ruff==0.6.2; extra == "dev"
+Requires-Dist: mypy==1.11.2; extra == "dev"
 Requires-Dist: pre-commit==3.5.0; extra == "dev"
-Requires-Dist: pytest>=7.0.1; extra == "dev"
+Requires-Dist: pytest>=7.2.2; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
 Requires-Dist: pytest-mock==3.11.1; extra == "dev"
-Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![codecov](https://codecov.io/gh/umami-hep/umami-preprocessing/graph/badge.svg?token=K8MJI20UZO)](https://codecov.io/gh/umami-hep/umami-preprocessing)

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/pyproject.toml

@@ -8,22 +8,22 @@ requires-python = "<3.12,>=3.8"
 
 dependencies = [
     "pyyaml-include==1.3",
-    "PyYAML==6.0.1",
+    "PyYAML>=6.0.1",
     "rich==12.6.0",
-    "scipy==1.10.1",
-    "puma-hep==0.4.2",
-    "atlas-ftag-tools==0.2.8",
+    "scipy>=1.15.2",
+    "puma-hep==0.4.5",
+    "atlas-ftag-tools==0.2.10",
     "dotmap==1.3.30"
 ]
 
 [project.optional-dependencies]
 dev = [
-  "ruff==0.1.6",
-  "mypy==1.5.1",
+  "ruff==0.6.2",
+  "mypy==1.11.2",
   "pre-commit==3.5.0",
-  "pytest>=7.0.1",
+  "pytest>=7.2.2",
+  "pytest-cov>=4.0.0",
   "pytest-mock==3.11.1",
-  "pytest-cov>=3.0.0",
 ]
 
 [project.urls]

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/umami_preprocessing.egg-info/PKG-INFO

@@ -1,25 +1,25 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: umami-preprocessing
-Version: 0.2.2
+Version: 0.2.4
 Summary: Preprocessing for jet tagging
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing
 Requires-Python: <3.12,>=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: pyyaml-include==1.3
-Requires-Dist: PyYAML==6.0.1
+Requires-Dist: PyYAML>=6.0.1
 Requires-Dist: rich==12.6.0
-Requires-Dist: scipy==1.10.1
-Requires-Dist: puma-hep==0.4.2
-Requires-Dist: atlas-ftag-tools==0.2.8
+Requires-Dist: scipy>=1.15.2
+Requires-Dist: puma-hep==0.4.5
+Requires-Dist: atlas-ftag-tools==0.2.10
 Requires-Dist: dotmap==1.3.30
 Provides-Extra: dev
-Requires-Dist: ruff==0.1.6; extra == "dev"
-Requires-Dist: mypy==1.5.1; extra == "dev"
+Requires-Dist: ruff==0.6.2; extra == "dev"
+Requires-Dist: mypy==1.11.2; extra == "dev"
 Requires-Dist: pre-commit==3.5.0; extra == "dev"
-Requires-Dist: pytest>=7.0.1; extra == "dev"
+Requires-Dist: pytest>=7.2.2; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
 Requires-Dist: pytest-mock==3.11.1; extra == "dev"
-Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![codecov](https://codecov.io/gh/umami-hep/umami-preprocessing/graph/badge.svg?token=K8MJI20UZO)](https://codecov.io/gh/umami-hep/umami-preprocessing)

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

@@ -0,0 +1,15 @@
+pyyaml-include==1.3
+PyYAML>=6.0.1
+rich==12.6.0
+scipy>=1.15.2
+puma-hep==0.4.5
+atlas-ftag-tools==0.2.10
+dotmap==1.3.30
+
+[dev]
+ruff==0.6.2
+mypy==1.11.2
+pre-commit==3.5.0
+pytest>=7.2.2
+pytest-cov>=4.0.0
+pytest-mock==3.11.1

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/upp/__init__.py

@@ -2,4 +2,4 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.2"
+__version__ = "v0.2.4"

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/upp/classes/preprocessing_config.py

@@ -53,42 +53,56 @@ class PreprocessingConfig:
 
     Parameters
     ----------
+    config_path : Path
+        Path to the config yaml file that is used. Does not need to be set in config.
+    split : Split
+        For which part the preprocessing is run. Either train, val or test. This needs
+        to be set as a command line argument when running the programm. Does not need
+        to be set in config.
+    config : dict
+        Dict with the loaded config. Does not need to be set in config.
     base_dir : Path
         Base directory for all other paths.
-    ntuple_dir : Path
+    ntuple_dir : Path, optional
         Directory containing the input h5 ntuples. If a relative path is given, it is
-        interpreted as relative to base_dir.
-    components_dir : Path
+        interpreted as relative to base_dir. By default Path("ntuples")
+    components_dir : Path, optional
         Directory for intermediate component files. If a relative path is given, it is
-        interpreted as relative to base_dir.
-    out_dir : Path
+        interpreted as relative to base_dir. By default Path("components")
+    out_dir : Path, optional
         Directory for output files. If a relative path is given, it is interpreted as
-        relative to base_dir.
-    out_fname : Path
-        Filename stem for the output files.
-    batch_size : int
+        relative to base_dir. By default Path("output")
+    out_fname : Path, optional
+        Filename stem for the output files. By default Path("pp_output.h5")
+    batch_size : int, optional
         Batch size for the preprocessing. For each batch select
         `sampling_fraction*batch_size_after_cuts`. It is recommended to choose high batch sizes
         especially to the `countup` method to achive best agreement of target and resampled
-        distributions.
-    num_jets_estimate : int
+        distributions. By default 100_000
+    num_jets_estimate : int, optional
         Any of the further three arguments that are not specified will default to this value
         Is equal to 1_000_000 by default.
-    num_jets_estimate_available : int | None
+    num_jets_estimate_available : int, optional
         A sabsample taken from the whole sample to estimate the number of jets after the cuts.
         Please keep this number high in order to not get poisson error of more then 5%.
         If time allows you can use -1 to get a precise number of jets and not just an estimate
         although it will be slow for large datasets. Is equal to num_jets_estimate by default.
-    num_jets_estimate_hist : int
+    num_jets_estimate_hist : int, optional
         Number of jets of each flavour that are used to construct histograms for probability
         density function estimation. Larger numbers give a better quality estmate of the pdfs.
         Is equal to num_jets_estimate by default.
-    num_jets_estimate_norm : int
+    num_jets_estimate_norm : int, optional
         Number of jets of each flavour that are used to estimate shifting and scaling during
         normalisation step. Larger numbers give a better quality estmates.
         Is equal to num_jets_estimate by default.
-    jets_name : str
-        Name of the jets dataset in the input file.
+    num_jets_estimate_plotting : int, optional
+        Number of jets of each flavour used for plotting the initial and the final resampling
+        variable distributions. Larger numbers give a better estimate of the full distributions.
+        Is equal to num_jets_estimate by default.
+    merge_test_samples : bool, optional
+        Merge the test samples of the different processes into one file. By default False.
+    jets_name : str, optional
+        Name of the jets dataset in the input file. By default "jets".
     """
 
     config_path: Path
@@ -104,9 +118,11 @@ class PreprocessingConfig:
     num_jets_estimate_available: int | None = None
     num_jets_estimate_hist: int | None = None
     num_jets_estimate_norm: int | None = None
+    num_jets_estimate_plotting: int | None = None
     merge_test_samples: bool = False
     jets_name: str = "jets"
     flavour_config: Path | None = None
+    num_jets_per_output_file: int | None = None
 
     def __post_init__(self):
         # postprocess paths
@@ -117,6 +133,8 @@ class PreprocessingConfig:
                 self.num_jets_estimate_hist = self.num_jets_estimate
             if self.num_jets_estimate_norm is None:
                 self.num_jets_estimate_norm = self.num_jets_estimate
+            if self.num_jets_estimate_plotting is None:
+                self.num_jets_estimate_plotting = self.num_jets_estimate
 
         for field in dataclasses.fields(self):
             if field.type == "Path" and field.name != "out_fname" and field.name != "base_dir":

umami_preprocessing-0.2.4/upp/logger.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import logging
+import sys
+from functools import partial
+
+from rich.console import Console
+from rich.logging import RichHandler
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+# Detect if the program is executed in an interactive terminal
+_IS_TTY = sys.stderr.isatty()
+
+# One console object is reused everywhere so that Rich keeps a consistent idea
+# of whether it may emit ANSI control codes / animations.
+_console = Console(
+    width=100,
+    force_terminal=_IS_TTY,
+    force_interactive=_IS_TTY,
+    no_color=not _IS_TTY,
+)
+
+# Template for the progress bar
+ProgressBar = partial(
+    Progress,
+    TextColumn("[task.description]{task.description}"),
+    BarColumn(),
+    TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+    TextColumn("•"),
+    TimeRemainingColumn(),
+    TextColumn("•"),
+    TimeElapsedColumn(),
+    refresh_per_second=1 if _IS_TTY else 0.05,
+    speed_estimate_period=30 if _IS_TTY else 120,
+    console=_console,
+    disable=not _IS_TTY,
+    transient=_IS_TTY,
+)
+
+
+# Helper for setup the logger
+def setup_logger(level: str = "INFO"):
+    """Set up the logger.
+
+    Configure Rich logging so that colourful / interactive output is used when
+    the program is attached to a terminal and plain text is written when it is
+    executed under a batch system such as Slurm (where stdout / stderr are files).
+    """
+    FORMAT = "%(message)s"
+
+    # In a batch job we create a console that never emits colour codes.
+    console = None
+    if not _IS_TTY:
+        console = Console(
+            width=120,
+            force_terminal=False,
+            force_interactive=False,
+            no_color=True,
+        )
+
+    handler = RichHandler(
+        show_time=False,
+        show_path=False,
+        markup=True,
+        rich_tracebacks=True,
+        console=console,
+    )
+
+    logging.basicConfig(level=level, format=FORMAT, handlers=[handler])
+    return logging

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/upp/main.py

@@ -21,7 +21,7 @@ from upp.logger import setup_logger
 from upp.stages.hist import create_histograms
 from upp.stages.merging import Merging
 from upp.stages.normalisation import Normalisation
-from upp.stages.plot import plot_initial_resampling_dists, plot_resampled_dists
+from upp.stages.plot import plot_resampling_dists
 from upp.stages.resampling import Resampling
 
 
@@ -41,10 +41,16 @@ def parse_args(args):
     parser.add_argument("--no-plot", dest="plot", action="store_false")
     splits = ["train", "val", "test", "all"]
     parser.add_argument("--split", default="train", choices=splits, help="Which file to produce")
+    parser.add_argument(
+        "--component", default=None, help="Component which is processed during --prep"
+    )
+    parser.add_argument(
+        "--region", default=None, help="Region which is processed during --resample"
+    )
 
     args = parser.parse_args(args)
     d = vars(args)
-    ignore = ["config", "split"]
+    ignore = ["config", "split", "component", "region"]
     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:
@@ -65,12 +71,15 @@ def run_pp(args) -> None:
 
     # create virtual datasets and pdf files
     if args.prep and args.split == "train":
-        create_histograms(config)
+        create_histograms(
+            config=config,
+            component_to_run=args.component,
+        )
 
     # run the resampling
     if args.resample:
         resampling = Resampling(config)
-        resampling.run()
+        resampling.run(region=args.region, component=args.component)
 
     # run the merging
     if args.merge:
@@ -86,8 +95,8 @@ def run_pp(args) -> None:
     if args.plot:
         title = " Plotting "
         log.info(f"[bold green]{title:-^100}")
-        plot_initial_resampling_dists(config=config)
-        plot_resampled_dists(config=config, stage=args.split)
+        plot_resampling_dists(config=config, stage="initial")
+        plot_resampling_dists(config=config, stage=args.split)
 
     # print end info
     end = datetime.now()

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.4}/upp/stages/hist.py

@@ -5,6 +5,7 @@ import logging as log
 import math
 from dataclasses import dataclass
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import h5py
 import numpy as np
@@ -13,6 +14,9 @@ from scipy.stats import binned_statistic_dd
 
 from upp.logger import setup_logger
 
+if TYPE_CHECKING:  # pragma: no cover
+    from upp.classes.preprocessing_config import PreprocessingConfig
+
 
 def bin_jets(array: dict, bins: list) -> np.ndarray:
     """Create the histogram and bins for the given resampling variables.
@@ -117,24 +121,39 @@ class Hist:
             return f["pbin"][:]
 
 
-def create_histograms(config) -> None:
+def create_histograms(
+    config: PreprocessingConfig,
+    component_to_run: str | None = None,
+) -> None:
     """Create the virtual datasets and pdf files.
 
     Parameters
     ----------
     config : PreprocessingConfig object
         PreprocessingConfig object of the current preprocessing.
+    component_to_run : str | None
+        Component which should be run. By default (None), all components
+        are processed sequentially.
     """
+    # Setup the logger and load the variables used for resampling
     setup_logger()
+    sampl_vars = config.sampl_cfg.vars
 
     title = " Writing PDFs "
     log.info(f"[bold green]{title:-^100}")
-
     log.info(f"[bold green]Estimating PDFs using {config.num_jets_estimate_hist:,} jets...")
-    sampl_vars = config.sampl_cfg.vars
+
+    # Create check variable to ensure at least one component was processed
+    component_processed = not component_to_run
+
+    # Process the different components
     for component in config.components:
+        # Check if only one component should be processed
+        if isinstance(component_to_run, str) and component_to_run != component.name:
+            continue
+
         log.info(f"Estimating {component} PDF using {config.num_jets_estimate_hist:,} samples...")
-        component.setup_reader(config.batch_size, config.jets_name)
+        component.setup_reader(batch_size=config.batch_size, jets_name=config.jets_name)
         cuts_no_split = component.cuts.ignore(["eventNumber"])
 
         ###
@@ -146,7 +165,29 @@ def create_histograms(config) -> None:
             silent=False,
             raise_error=False,
         )
-        jets = component.get_jets(sampl_vars, config.num_jets_estimate_hist, cuts_no_split)
-        component.hist.write_hist(jets, sampl_vars, config.sampl_cfg.flat_bins)
+
+        # Load the jets from file used for resampling
+        jets = component.get_jets(
+            variables=sampl_vars,
+            num_jets=config.num_jets_estimate_hist,
+            cuts=cuts_no_split,
+        )
+
+        # Write out the hist used for resampling
+        component.hist.write_hist(
+            jets=jets,
+            resampling_vars=sampl_vars,
+            bins=config.sampl_cfg.flat_bins,
+        )
+
+        # Set the check variable to true
+        component_processed = True
+
+    # Raise error of no region was processed
+    if component_processed is False:
+        raise ValueError(
+            "No component processed during resampling! Check that you correctly spelled "
+            "the component name when running with --component!"
+        )
 
     log.info(f"[bold green]Saved to {config.components[0].hist.path.parent}/")