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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: umami-preprocessing
-Version: 0.2.2
+Version: 0.2.3
 Summary: Preprocessing for jet tagging
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: umami-preprocessing
-Version: 0.2.2
+Version: 0.2.3
 Summary: Preprocessing for jet tagging
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/umami-preprocessing

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

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

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

@@ -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)
 
     # run the merging
     if args.merge:

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.3}/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}/")

{umami_preprocessing-0.2.2 → umami_preprocessing-0.2.3}/upp/stages/resampling.py

@@ -356,9 +356,14 @@ class Resampling:
                 f" Jets are upsampled at most {np.max(component._ups_max):.0f} times"
             )
 
-    def set_component_sampling_fractions(self):
+    def set_component_sampling_fractions(self, component: Component):
         """Automatically set the sampling fraction for each of the components.
 
+        Parameters
+        ----------
+        component : Component
+            Component for which the sampling fraction is set.
+
         Raises
         ------
         ValueError
@@ -366,71 +371,86 @@ class Resampling:
         """
         # Check that the sampling fraction must be found automatically
         if self.config.sampling_fraction == "auto" or self.config.sampling_fraction is None:
-            log.info("[bold green]Sampling fraction chosen for each component automatically...")
+            log.info(
+                "[bold green]Sampling fraction will be chosen "
+                f"automatically for {component.name}..."
+            )
 
-            # Loop over each component
-            for component in self.components:
-                # Target component always gets one as sampling fraction
-                if component.is_target(self.config.target):
-                    component.sampling_fraction = 1
-
-                else:
-                    sam_frac = component.get_auto_sampling_fraction(
-                        num_jets=component.num_jets,
-                        cuts=component.cuts,
-                    )
+            # Target component always gets one as sampling fraction
+            if component.is_target(self.config.target):
+                component.sampling_fraction = 1
+
+            else:
+                sam_frac = component.get_auto_sampling_fraction(
+                    num_jets=component.num_jets,
+                    cuts=component.cuts,
+                )
 
-                    # Raise an error/warning if the sampling fraction is above one
-                    # for the countup/pdf method
-                    if sam_frac > 1:
-                        if self.config.method == "countup":
-                            raise ValueError(
-                                f"[bold red]Sampling fraction of {sam_frac:.3f}>1 is"
-                                f" needed for component {component} This is not supported for"
-                                " countup method."
-                            )
-                        else:
-                            log.warning(
-                                f"[bold yellow]sampling fraction of {sam_frac:.3f}>1 is"
-                                f" needed for component {component}"
-                            )
-
-                    # Ensure the sampling fraction is at least above 0.1
-                    component.sampling_fraction = max(sam_frac, 0.1)
+                # Raise an error/warning if the sampling fraction is above one
+                # for the countup/pdf method
+                if sam_frac > 1:
+                    if self.config.method == "countup":
+                        raise ValueError(
+                            f"[bold red]Sampling fraction of {sam_frac:.3f}>1 is"
+                            f" needed for component {component} This is not supported for"
+                            " countup method."
+                        )
+                    else:
+                        log.warning(
+                            f"[bold yellow]sampling fraction of {sam_frac:.3f}>1 is"
+                            f" needed for component {component}"
+                        )
+
+                # Ensure the sampling fraction is at least above 0.1
+                component.sampling_fraction = max(sam_frac, 0.1)
 
         else:
-            # Set the sampling fraction for each component to the value defined
+            # Set the sampling fraction for the component to the value defined
             # in the config
-            for component in self.components:
-                if component.is_target(self.config.target):
-                    component.sampling_fraction = 1
+            if component.is_target(self.config.target):
+                component.sampling_fraction = 1
+
+            else:
+                component.sampling_fraction = self.config.sampling_fraction
 
-                else:
-                    component.sampling_fraction = self.config.sampling_fraction
+    def run(self, region: str | None = None):
+        """Execute the resampling.
 
-    def run(self):
-        """Execute the resampling."""
+        Parameters
+        ----------
+        region : str | None, optional
+            Define which region is to be resampled, by default None
+            which works through the regions sequentially
+
+        Raises
+        ------
+        ValueError
+            If no region was processed during resampling
+        """
         title = " Running resampling "
         log.info(f"[bold green]{title:-^100}")
         log.info(f"Resampling method: {self.config.method}")
 
-        # Setup the different components and readers/writers
+        # Setup the different components and readers/writers and their sampling fraction
         for component in self.components:
+            # Check if the component is needed for the region
+            if region and region not in component.name:
+                continue
+
             # just used for the writer configuration
             component.setup_reader(
                 self.batch_size, jets_name=self.jets_name, transform=self.transform
             )
             component.setup_writer(self.variables, jets_name=self.jets_name)
 
-        # Set samplig fraction if needed
-        self.set_component_sampling_fractions()
+            # Set sampling fraction
+            self.set_component_sampling_fractions(component=component)
 
-        # Check samples
-        log.info(
-            "[bold green]Checking requested num_jets based on a sampling fraction of"
-            f" {self.config.sampling_fraction}..."
-        )
-        for component in self.components:
+            # Check that enough jets are available
+            log.info(
+                "[bold green]Checking requested num_jets based on a sampling fraction of"
+                f" {self.config.sampling_fraction}..."
+            )
             frac = component.sampling_fraction if self.select_func else 1
             component.check_num_jets(
                 component.num_jets,
@@ -438,16 +458,46 @@ class Resampling:
                 cuts=component.cuts,
             )
 
+        # Create check variable to ensure at least one region was processed
+        region_processed = not region
+
         # Run resampling
-        for region, components in self.components.groupby_region():
+        for iter_region, iter_components in self.components.groupby_region():
+            # Check if a specific region for resampling was chosen
+            if region and region != iter_region.name:
+                continue
+
             log.info(f"[bold green]Running over region {region}...")
-            self.run_on_region(components, region)
+            self.run_on_region(components=iter_components, region=iter_region)
+            region_processed = True
+
+        # Raise error of no region was processed
+        if region_processed is False:
+            raise ValueError(
+                "No region processed during resampling! Check that you correctly spelled "
+                "the region name when running with --region!"
+            )
 
         # Finalise the resampling
-        unique = sum(component.writer.get_attr("unique_jets") for component in self.components)
-        log.info(f"[bold green]Finished resampling a total of {self.components.num_jets:,} jets!")
-        log.info(f"[bold green]Estimated unqiue jets: {unique:,.0f}")
-        log.info(f"[bold green]Saved to {self.components.out_dir}/")
+        if region:
+            unique = 0
+            for component in self.components:
+                if region in component.name:
+                    unique += component.writer.get_attr("unique_jets")
+            log.info(
+                f"[bold green]Finished resampling of region {region}. "
+                f"A total of {self.components.num_jets:,} jets!"
+            )
+            log.info(f"[bold green]Estimated unqiue jets: {unique:,.0f}")
+            log.info(f"[bold green]Saved to {self.components.out_dir}/")
+
+        else:
+            unique = sum(component.writer.get_attr("unique_jets") for component in self.components)
+            log.info(
+                f"[bold green]Finished resampling a total of {self.components.num_jets:,} jets!"
+            )
+            log.info(f"[bold green]Estimated unqiue jets: {unique:,.0f}")
+            log.info(f"[bold green]Saved to {self.components.out_dir}/")
 
     def get_num_bins_from_config(self) -> list[list[int]]:
         """Get the lengths of the binning regions in each variable from the config.