atlas-ftag-tools 0.2.4__tar.gz → 0.2.6__tar.gz

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.2.4
+Version: 0.2.6
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/atlas_ftag_tools.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.2.4
+Version: 0.2.6
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/__init__.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.4"
+__version__ = "v0.2.6"
 
 
 from ftag import hdf5
 from ftag.cuts import Cuts
 from ftag.flavour import Flavour, Flavours
+from ftag.labeller import Labeller
 from ftag.mock import get_mock_file
 from ftag.sample import Sample
 from ftag.transform import Transform
@@ -18,6 +19,7 @@ __all__ = [
     "Cuts",
     "Flavour",
     "Flavours",
+    "Labeller",
     "Sample",
     "Transform",
     "__version__",

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/flavour.py

@@ -22,10 +22,11 @@ class Flavour:
     cuts: Cuts
     colour: str
     category: str
+    _px: str | None = None
 
     @property
     def px(self) -> str:
-        return f"p{remove_suffix(self.name, 'jets')}"
+        return self._px or f"p{remove_suffix(self.name, 'jets')}"
 
     @property
     def eff_str(self) -> str:

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/flavours.yaml

@@ -53,21 +53,25 @@
   cuts: ["HadronGhostTruthLabelID == 5"]
   colour: tab:blue
   category: single-btag-ghost
+  _px: pb
 - name: ghostcjets
   label: $c$-jets
   cuts: ["HadronGhostTruthLabelID == 4"]
   colour: tab:orange
   category: single-btag-ghost
+  _px: pc
 - name: ghostujets
   label: Light-jets
   cuts: ["HadronGhostTruthLabelID == 0"]
   colour: tab:green
   category: single-btag-ghost
+  _px: pu
 - name: ghosttaujets
   label: $\tau$-jets
   cuts: ["HadronGhostTruthLabelID == 15"]
   colour: tab:purple
   category: single-btag-ghost
+  _px: ptau
 
 # Xbb tagging
 - name: hbb
@@ -90,6 +94,31 @@
   cuts: ["R10TruthLabel_R22v1 == 10"]
   colour: "#38761D"
   category: xbb
+- name: qcdbb
+  label: QCD->bb
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount == 2"]
+  colour: "red"
+  category: xbb
+- name: qcdnonbb
+  label: QCD
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount != 2"]
+  colour: "silver"
+  category: xbb
+- name: qcdbx
+  label: QCD
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount == 1"]
+  colour: "gold"
+  category: xbb
+- name: qcdcx
+  label: QCD
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostCHadronsFinalCount >= 1", "GhostBHadronsFinalCount == 0"]
+  colour: "pink"
+  category: xbb
+- name: qcdll
+  label: QCD
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount == 0", "GhostCHadronsFinalCount == 0"]
+  colour: "green"
+  category: xbb
 
 # extended Xbb tagging
 - name: tqqb

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/hdf5/h5writer.py

@@ -81,7 +81,7 @@ class H5Writer:
         return cls(dtypes=dtypes, shapes=shapes, **kwargs)
 
     def create_ds(self, name: str, dtype: np.dtype) -> None:
-        if name == self.jets_name and self.add_flavour_label:
+        if name == self.jets_name and self.add_flavour_label and "flavour_label" not in dtype.names:
             dtype = np.dtype([*dtype.descr, ("flavour_label", "i4")])
 
         # adjust dtype based on specified precision

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/labeller.py

@@ -4,8 +4,7 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from ftag import Flavours
-from ftag.flavour import Flavour, FlavourContainer
+from ftag.flavour import Flavour, FlavourContainer, Flavours
 from ftag.hdf5 import join_structured_arrays, structured_from_dict
 
 
@@ -27,6 +26,10 @@ class Labeller:
     labels: FlavourContainer | list[str | Flavour]
     require_labels: bool = True
 
+    @property
+    def variables(self):
+        return sum((label.cuts.variables for label in self.labels), [])
+
     def __post_init__(self) -> None:
         if isinstance(self.labels, FlavourContainer):
             self.labels = list(self.labels)

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/mock.py

@@ -22,6 +22,8 @@ JET_VARS = [
     ("HadronConeExclTruthLabelID", "i4"),
     ("HadronConeExclTruthLabelPt", "f4"),
     ("R10TruthLabel_R22v1", "i4"),
+    ("GhostBHadronsFinalCount", "i4"),
+    ("GhostCHadronsFinalCount", "i4"),
     ("n_truth_promptLepton", "i4"),
     ("flavour_label", "i4"),
 ]
@@ -37,17 +39,18 @@ TRACK_VARS = [
     ("phiUncertainty", "f4"),
     ("thetaUncertainty", "f4"),
     ("qOverPUncertainty", "f4"),
-    ("numberOfPixelHits", "i4"),
-    ("numberOfSCTHits", "i4"),
-    ("numberOfInnermostPixelLayerHits", "i4"),
-    ("numberOfNextToInnermostPixelLayerHits", "i4"),
-    ("numberOfInnermostPixelLayerSharedHits", "i4"),
-    ("numberOfInnermostPixelLayerSplitHits", "i4"),
-    ("numberOfPixelSharedHits", "i4"),
-    ("numberOfPixelSplitHits", "i4"),
-    ("numberOfSCTSharedHits", "i4"),
-    ("numberOfPixelHoles", "i4"),
-    ("numberOfSCTHoles", "i4"),
+    ("numberOfPixelHits", "u1"),
+    ("numberOfSCTHits", "u1"),
+    ("numberOfInnermostPixelLayerHits", "u1"),
+    ("numberOfNextToInnermostPixelLayerHits", "u1"),
+    ("numberOfInnermostPixelLayerSharedHits", "u1"),
+    ("numberOfInnermostPixelLayerSplitHits", "u1"),
+    ("numberOfPixelSharedHits", "u1"),
+    ("numberOfPixelSplitHits", "u1"),
+    ("numberOfSCTSharedHits", "u1"),
+    ("numberOfPixelHoles", "u1"),
+    ("numberOfSCTHoles", "u1"),
+    ("leptonID", "i1"),
 ]
 
 
@@ -98,6 +101,8 @@ def mock_jets(num_jets=1000) -> np.ndarray:
 
     # add tagger scores
     jets["HadronConeExclTruthLabelID"] = rng.choice([0, 4, 5, 15], size=num_jets)
+    jets["GhostBHadronsFinalCount"] = rng.choice([0, 1, 2], size=num_jets)
+    jets["GhostCHadronsFinalCount"] = rng.choice([0, 1, 2], size=num_jets)
     jets["R10TruthLabel_R22v1"] = rng.choice([1, 10, 11, 12], size=num_jets)
     scores = get_mock_scores(jets["HadronConeExclTruthLabelID"])
     xbb_scores = get_mock_scores(jets["R10TruthLabel_R22v1"], is_xbb=True)

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/track_selector.py

@@ -41,8 +41,10 @@ class TrackSelector:
         for var in tracks.dtype.names:
             if issubclass(tracks[var].dtype.type, np.floating):
                 tracks[var][rm_idx] = np.nan
-            elif issubclass(tracks[var].dtype.type, np.integer):
+            elif issubclass(tracks[var].dtype.type, np.signedinteger):
                 tracks[var][rm_idx] = -1
+            elif issubclass(tracks[var].dtype.type, np.unsignedinteger):
+                tracks[var][rm_idx] = 0
             elif issubclass(tracks[var].dtype.type, np.bool_):
                 tracks[var][rm_idx] = False
             else:

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/wps/discriminant.py

@@ -4,7 +4,7 @@ from typing import Callable
 
 import numpy as np
 
-from ftag.flavour import Flavour, Flavours
+from ftag.flavour import Flavour, Flavours, remove_suffix
 
 
 def discriminant(
@@ -50,7 +50,10 @@ def discriminant(
         if fx > 0 and name not in jets.dtype.names:
             raise ValueError(f"Nonzero fx for {d}, but '{name}' not found in input array.")
         denominator += jets[name] * fx if name in jets.dtype.names else 0
-    return np.log((jets[f"{tagger}_{signal.px}"] + epsilon) / (denominator + epsilon))
+    signal_field = f"{tagger}_{signal.px}"
+    if signal_field not in jets.dtype.names:
+        signal_field = f"{tagger}_p{remove_suffix(signal.name, 'jets')}"
+    return np.log((jets[signal_field] + epsilon) / (denominator + epsilon))
 
 
 def tautag_dicriminant(jets, tagger, fb, fc, epsilon=1e-10):
@@ -63,6 +66,11 @@ def btag_discriminant(jets, tagger, fc, ftau=0, epsilon=1e-10):
     return discriminant(jets, tagger, Flavours.bjets, fxs, epsilon=epsilon)
 
 
+def ghostbtag_discriminant(jets, tagger, fc, ftau=0, epsilon=1e-10):
+    fxs = {"pghostc": fc, "pghosttau": ftau, "pghostu": 1 - fc - ftau}
+    return discriminant(jets, tagger, Flavours.ghostbjets, fxs, epsilon=epsilon)
+
+
 def ctag_discriminant(jets, tagger, fb, ftau=0, epsilon=1e-10):
     fxs = {"pb": fb, "ptau": ftau, "pu": 1 - fb - ftau}
     return discriminant(jets, tagger, Flavours.cjets, fxs, epsilon=epsilon)
@@ -112,6 +120,7 @@ def get_discriminant(
         "taujets": tautag_dicriminant,
         "hbb": hbb_discriminant,
         "hcc": hcc_discriminant,
+        "ghostbjets": ghostbtag_discriminant,
     }
 
     if str(signal) not in tagger_funcs:

{atlas_ftag_tools-0.2.4 → atlas_ftag_tools-0.2.6}/ftag/wps/working_points.py

@@ -218,7 +218,7 @@ def setup_common_parts(args):
     zprime_cuts = Cuts.from_list(args.zprime_cuts) + default_cuts
 
     # prepare to load jets
-    all_vars = next(iter(flavs)).cuts.variables
+    all_vars = list(set(sum((flav.cuts.variables for flav in flavs), [])))
     reader = H5Reader(args.ttbar)
     jet_vars = reader.dtypes()["jets"].names
     for tagger in args.tagger: