atlas-ftag-tools 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

{atlas_ftag_tools-0.2.4.dist-info → atlas_ftag_tools-0.2.6.dist-info}/METADATA

@@ -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.dist-info → atlas_ftag_tools-0.2.6.dist-info}/RECORD

@@ -1,15 +1,15 @@
-ftag/__init__.py,sha256=Mx2Emsw4TM1YL0wTQMHK36EaQE_ImeV6ukiz1X5BZAU,629
+ftag/__init__.py,sha256=edvBW3ySbO10OzSeqTQd9HPuXqiosdWibAhi3ZJFT4g,680
 ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
 ftag/cuts.py,sha256=9_ooLZHaO3SnIQBNxwbaPZn-qptGdKnB27FdKQGTiTY,2933
-ftag/flavour.py,sha256=EMZZLyl6lSdvkfrYxHhMcSn3aqP_FU7OpCFkvZpTksU,3761
-ftag/flavours.yaml,sha256=lFnVwjh_DwLhOc3mr5n6bSIWyHgxQvAXas4lEmEDncU,7520
+ftag/flavour.py,sha256=MUij7lknepnkidErDsyNsJDbAPotAEwS-uMaTSYCet4,3800
+ftag/flavours.yaml,sha256=E_vpn38qJ3-Tygg2aHlH4wkn_rR1On_lMeaG8OemHCQ,8285
 ftag/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
-ftag/labeller.py,sha256=uDygOhVGSNn96DWw8aErHpTtFsFX0RnxYYpy4g1FRog,2457
-ftag/mock.py,sha256=_oy-r3eLllFy33NAoZaKfAx-Rp2vrCdrGj3UsTMks94,4740
+ftag/labeller.py,sha256=iyDjRkOaly48jTaaFimA5eZ5YySp1lrLVEt3GyDZW-8,2553
+ftag/mock.py,sha256=Eyj3tkkaSSnqvS3G6NS7fq8sB__Nx8YE9-OM2_lpdoQ,4992
 ftag/region.py,sha256=ANv0dGI2W6NJqD9fp7EfqAUReH4FOjc1gwl_Qn8llcM,360
 ftag/sample.py,sha256=TFXMhDkbPmjkms9-b-bINJ32T3bO86JcU70C0nY7wa8,2500
 ftag/test_cli_utils.py,sha256=xa08vf6SEOow58SSFagYdAselb-dkNOVvWsWheMnW-g,1001
-ftag/track_selector.py,sha256=piSYAN_IkOsrXxKXjXbJpMSseUig5P2BJW5mCwsMUDM,2535
+ftag/track_selector.py,sha256=fJNk_kIBQriBqV4CPT_3ReJbOUnavDDzO-u3EQlRuyk,2654
 ftag/transform.py,sha256=uEGGJSnqoKOzLYQv650XdK_kDNw4Aw-5dc60z9Dp_y0,3963
 ftag/vds.py,sha256=nRViQZQIORB95nC7NZsW3KsSoGkLzEdOsuCViH5h8-U,3296
 ftag/hdf5/__init__.py,sha256=LFDNxVOCp58SvLHwQhdT68Q-KBMS_i6jBrbXoRpHzbM,354
@@ -17,12 +17,12 @@ ftag/hdf5/h5move.py,sha256=oYpRu0IDCIJIQ2ML52HBAdoyDxmKkHTeM9JdbPEgKfI,947
 ftag/hdf5/h5reader.py,sha256=i31pDAqmOSaxdeRhc4iSBlld8xJ0pmp4rNd7CugNzw0,13706
 ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
 ftag/hdf5/h5utils.py,sha256=-4zKTMtNCrDZr_9Ww7uzfsB7M7muBKpmm_1IkKJnHOI,3222
-ftag/hdf5/h5writer.py,sha256=j3Fy8snkiVVfimiUz3rrZOhSV8OF27978Y9pk0QcTGM,5277
+ftag/hdf5/h5writer.py,sha256=9FkClV__UbBqmFsq_h2jwiZnbWVm8QFRL_4mDZZBbTs,5316
 ftag/wps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ftag/wps/discriminant.py,sha256=kJFekUTPNIvCabJCon6OqOAQEzz5hj3XrWFFRLOgGOs,3836
-ftag/wps/working_points.py,sha256=VTU6OD40ULAJQD0MlD1EZd33q8ociUvFX1YrhgJFvXc,9722
-atlas_ftag_tools-0.2.4.dist-info/METADATA,sha256=f4aCu6JmItUBp5EmTzbrqhC5-Wsy7uiOiBO9yufyacQ,5169
-atlas_ftag_tools-0.2.4.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-atlas_ftag_tools-0.2.4.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
-atlas_ftag_tools-0.2.4.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.2.4.dist-info/RECORD,,
+ftag/wps/discriminant.py,sha256=rTASnt9s6xEHK5hLKHgvc9c7TiyyX6Upbp4LeTh6TB0,4270
+ftag/wps/working_points.py,sha256=fPfSZoTp-2qxL3YB7HdPxIB-7aSR_HozEtWWIA0BUtg,9749
+atlas_ftag_tools-0.2.6.dist-info/METADATA,sha256=gmEDL518a_xrRBsPsRYCFt-C06BD3K9bH6HCStV6mAU,5169
+atlas_ftag_tools-0.2.6.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+atlas_ftag_tools-0.2.6.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.2.6.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.6.dist-info/RECORD,,

{atlas_ftag_tools-0.2.4.dist-info → atlas_ftag_tools-0.2.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (75.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

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__",

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:

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

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

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)

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)

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:

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:

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: