atlas-ftag-tools 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

{atlas_ftag_tools-0.2.1.dist-info → atlas_ftag_tools-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.2.1
+Version: 0.2.3
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT
@@ -11,8 +11,8 @@ Requires-Dist: h5py >=3.0
 Requires-Dist: numpy
 Requires-Dist: PyYAML >=5.1
 Provides-Extra: dev
-Requires-Dist: ruff ==0.2.2 ; 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.1.1 ; extra == 'dev'
 Requires-Dist: pytest ==7.2.2 ; extra == 'dev'
 Requires-Dist: pytest-cov ==4.0.0 ; extra == 'dev'

atlas_ftag_tools-0.2.3.dist-info/RECORD

@@ -0,0 +1,28 @@
+ftag/__init__.py,sha256=i5XMWwkyGmwIIIaO3veRyOnTsJ0qKyW1V6HsuYc9Dm4,629
+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/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
+ftag/labeller.py,sha256=uDygOhVGSNn96DWw8aErHpTtFsFX0RnxYYpy4g1FRog,2457
+ftag/mock.py,sha256=FboI1Kq6aKZv43SpubiLkJvn4BSMDv2Fl2UniuhxspU,4502
+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=WQ6lzY8n6pumIVfLMLgjlTkWUaIQu-8Dq3FKsmCsG_0,1736
+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
+ftag/hdf5/h5move.py,sha256=oYpRu0IDCIJIQ2ML52HBAdoyDxmKkHTeM9JdbPEgKfI,947
+ftag/hdf5/h5reader.py,sha256=H_5Aw0lOyEzK_phMRhD-jR_OSCsXnCA3qJZnRvPqaRU,13569
+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/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.3.dist-info/METADATA,sha256=XP7QCxOKz-QUKKx5ds8ldz13n_qfTVcP-SZciExdIr0,5169
+atlas_ftag_tools-0.2.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+atlas_ftag_tools-0.2.3.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.2.3.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.3.dist-info/RECORD,,

{atlas_ftag_tools-0.2.1.dist-info → atlas_ftag_tools-0.2.3.dist-info}/WHEEL

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

ftag/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.1"
+__version__ = "v0.2.3"
 
 
 from ftag import hdf5

ftag/cuts.py

@@ -79,7 +79,9 @@ class Cuts:
     def ignore(self, variables: list[str]):
         return Cuts(tuple(c for c in self if c.variable not in variables))
 
-    def __call__(self, array) -> CutsResult:
+    def __call__(self, array: np.ndarray) -> CutsResult:
+        if array.ndim == 2:
+            raise ValueError("This interface only supports jet selections")
         keep = np.arange(len(array))
         for cut in self.cuts:
             idx = cut(array)

ftag/flavour.py

@@ -42,6 +42,9 @@ class Flavour:
     def __str__(self) -> str:
         return self.name
 
+    def __lt__(self, other) -> bool:
+        return self.name < other.name
+
 
 @dataclass
 class FlavourContainer:
@@ -81,7 +84,10 @@ class FlavourContainer:
         return list(dict.fromkeys(f.category for f in self))
 
     def by_category(self, category: str) -> FlavourContainer:
-        return FlavourContainer({k: v for k, v in self.flavours.items() if v.category == category})
+        f = FlavourContainer({k: v for k, v in self.flavours.items() if v.category == category})
+        if not f.flavours:
+            raise KeyError(f"No flavours with category '{category}' found")
+        return f
 
     def from_cuts(self, cuts: list | Cuts) -> Flavour:
         if isinstance(cuts, list):

ftag/flavours.yaml

@@ -47,6 +47,28 @@
   colour: gold
   category: single-btag-extended
 
+# single b-tagging (ghost association)
+- name: ghostbjets
+  label: $b$-jets
+  cuts: ["HadronGhostTruthLabelID == 5"]
+  colour: tab:blue
+  category: single-btag-ghost
+- name: ghostcjets
+  label: $c$-jets
+  cuts: ["HadronGhostTruthLabelID == 4"]
+  colour: tab:orange
+  category: single-btag-ghost
+- name: ghostujets
+  label: Light-jets
+  cuts: ["HadronGhostTruthLabelID == 0"]
+  colour: tab:green
+  category: single-btag-ghost
+- name: ghosttaujets
+  label: $\tau$-jets
+  cuts: ["HadronGhostTruthLabelID == 15"]
+  colour: tab:purple
+  category: single-btag-ghost
+
 # Xbb tagging
 - name: hbb
   label: $H \rightarrow b\bar{b}$

ftag/hdf5/h5reader.py

@@ -60,7 +60,7 @@ class H5SingleReader:
             for var in array.dtype.names:
                 isinf = np.isinf(array[var])
                 isinf = isinf if name == self.jets_name else isinf.any(axis=-1)
-                keep_idx = keep_idx & ~isinf
+                keep_idx &= ~isinf
                 if num_inf := isinf.sum():
                     log.warning(
                         f"{num_inf} inf values detected for variable {var} in"

ftag/hdf5/h5writer.py

@@ -41,7 +41,7 @@ class H5Writer:
     jets_name: str = "jets"
     add_flavour_label: bool = False
     compression: str = "lzf"
-    precision: str | None = None
+    precision: str = "full"
     shuffle: bool = True
 
     def __post_init__(self):
@@ -51,6 +51,13 @@ class H5Writer:
         assert len(set(self.num_jets)) == 1, "Must have same number of jets per group"
         self.num_jets = self.num_jets[0]
 
+        if self.precision == "full":
+            self.fp_dtype = np.float32
+        elif self.precision == "half":
+            self.fp_dtype = np.float16
+        else:
+            raise ValueError(f"Invalid precision: {self.precision}")
+
         self.dst = Path(self.dst)
         self.dst.parent.mkdir(parents=True, exist_ok=True)
         self.file = h5py.File(self.dst, "w")
@@ -77,11 +84,17 @@ class H5Writer:
         if name == self.jets_name and self.add_flavour_label:
             dtype = np.dtype([*dtype.descr, ("flavour_label", "i4")])
 
+        # adjust dtype based on specified precision
+        dtype = np.dtype([
+            (field, self.fp_dtype if np.issubdtype(dt, np.floating) else dt)
+            for field, dt in dtype.descr
+        ])
+
         # optimal chunking is around 100 jets, only aply for track groups
         shape = self.shapes[name]
         chunks = (100,) + shape[1:] if shape[1:] else None
 
-        # note: enabling the hd5 shuffle filter doesn't improve anything
+        # note: enabling the hd5 shuffle filter doesn't improve write performance
         self.file.create_dataset(
             name, dtype=dtype, shape=shape, compression=self.compression, chunks=chunks
         )

ftag/labeller.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ftag import Flavours
+from ftag.flavour import Flavour, FlavourContainer
+from ftag.hdf5 import join_structured_arrays, structured_from_dict
+
+
+@dataclass
+class Labeller:
+    """
+    Defines a labelling scheme.
+
+    Labels are [0, ..., n] and are assigned using pre-defined selections.
+
+    Parameters
+    ----------
+    labels : FlavourContainer | list[str | Flavour]
+        The labels to be use.
+    require_labels : bool
+        Whether to require that all objects are labelled.
+    """
+
+    labels: FlavourContainer | list[str | Flavour]
+    require_labels: bool = True
+
+    def __post_init__(self) -> None:
+        if isinstance(self.labels, FlavourContainer):
+            self.labels = list(self.labels)
+        self.labels = sorted([Flavours[label] for label in self.labels])
+
+    def get_labels(self, array: np.ndarray) -> np.ndarray:
+        """
+        Returns the labels for the given array.
+
+        Parameters
+        ----------
+        array : np.ndarray
+            The array to label.
+
+        Returns
+        -------
+        np.ndarray
+            The labels for the given array.
+
+        Raises
+        ------
+        ValueError
+            If the `require_labels` attribute is set to `True` and some objects were not labelled.
+        """
+        labels = -1 * np.ones_like(array, dtype=int)
+        for i, label in enumerate(self.labels):
+            labels[label.cuts(array).idx] = i
+
+        if self.require_labels and -1 in labels:
+            raise ValueError("Some objects were not labelled")
+
+        return labels[labels != -1]
+
+    def add_labels(self, array: np.ndarray, label_name: str = "labels") -> np.ndarray:
+        """
+        Adds the labels to the given array.
+
+        Parameters
+        ----------
+        array : np.ndarray
+            The array to label.
+        label_name : str
+            The name of the label column.
+
+        Returns
+        -------
+        np.ndarray
+            The array with the labels added.
+
+        Raises
+        ------
+        ValueError
+            If the `require_labels` attribute is set to `False`.
+        """
+        if not self.require_labels:
+            raise ValueError("Cannot add labels if require_labels is set to False")
+        labels = self.get_labels(array)
+        labels = structured_from_dict({label_name: labels})
+        return join_structured_arrays([array, labels])

ftag/mock.py

@@ -52,7 +52,6 @@ TRACK_VARS = [
 
 
 def softmax(x, axis=None):
-    """Compute softmax values for each sets of scores in x."""
     e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
     return e_x / e_x.sum(axis=axis, keepdims=True)
 
@@ -85,12 +84,7 @@ def get_mock_scores(labels: np.ndarray, is_xbb: bool = False):
     return u2s(scores, dtype=np.dtype([(name, "f4") for name in cols]))
 
 
-def get_mock_file(
-    num_jets=1000,
-    fname: str | None = None,
-    tracks_name: str = "tracks",
-    num_tracks: int = 40,
-) -> tuple[str, h5py.File]:
+def mock_jets(num_jets=1000) -> np.ndarray:
     # setup jets
     rng = np.random.default_rng(42)
     jets_dtype = np.dtype(JET_VARS)
@@ -107,7 +101,26 @@ def get_mock_file(
     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)
-    jets = join_structured_arrays([jets, scores, xbb_scores])
+    return join_structured_arrays([jets, scores, xbb_scores])
+
+
+def mock_tracks(num_jets=1000, num_tracks=40) -> np.ndarray:
+    rng = np.random.default_rng(42)
+    tracks_dtype = np.dtype(TRACK_VARS)
+    tracks = u2s(rng.random((num_jets, num_tracks, len(TRACK_VARS))), tracks_dtype)
+    tracks["d0"] *= 5
+    valid = rng.choice([True, False], size=(num_jets, num_tracks))
+    valid = valid.astype(bool).view(dtype=np.dtype([("valid", bool)]))
+    return join_structured_arrays([tracks, valid])
+
+
+def get_mock_file(
+    num_jets=1000,
+    fname: str | None = None,
+    tracks_name: str = "tracks",
+    num_tracks: int = 40,
+) -> tuple[str, h5py.File]:
+    jets = mock_jets(num_jets)
 
     # create a tempfile in a new folder
     if fname is None:
@@ -121,11 +134,7 @@ def get_mock_file(
 
     # setup tracks
     if tracks_name:
-        tracks_dtype = np.dtype(TRACK_VARS)
-        tracks = u2s(rng.random((num_jets, num_tracks, len(TRACK_VARS))), tracks_dtype)
-        valid = rng.choice([True, False], size=(num_jets, num_tracks))
-        valid = valid.astype(bool).view(dtype=np.dtype([("valid", bool)]))
-        tracks = join_structured_arrays([tracks, valid])
+        tracks = mock_tracks(num_jets, num_tracks)
         f.create_dataset(tracks_name, data=tracks)
 
     return fname, f

ftag/track_selector.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ftag import Cuts
+
+
+@dataclass
+class TrackSelector:
+    """
+    Apply track selections to a set of tracks stored in a structured numpy array.
+
+    The array is assumed to have shape (n_jets, n_tracks, n_features).
+    Applying cuts will NaN out the tracks that do not pass the cuts,
+    but leave the shape of the array unchanged.
+
+    Parameters
+    ----------
+    cuts : Cuts
+        The cuts to apply to the tracks
+    valid_str : str
+        The name of the field in the tracks that indicates whether the track is
+    """
+
+    cuts: Cuts
+    valid_str: str = "valid"
+
+    def __call__(self, tracks: np.ndarray) -> np.ndarray:
+        # get a bool array for all tracks passing before any cuts
+        rm_idx = np.zeros_like(tracks[self.valid_str], dtype=bool)
+
+        # apply the cuts
+        for cut in self.cuts.cuts:
+            # remove valid track indices that do not pass the selection
+            rm_idx[tracks[self.valid_str] & ~cut(tracks)] = True
+
+        # set the values of the tracks that do not pass the cuts to
+        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):
+                tracks[var][rm_idx] = -1
+            elif issubclass(tracks[var].dtype.type, np.bool_):
+                tracks[var][rm_idx] = False
+            else:
+                raise TypeError(f"Unknown dtype {tracks[var].dtype}")
+
+        # specifically set the valid flag to false (even though it's already false by now)
+        tracks[rm_idx][self.valid_str] = False
+
+        return tracks

ftag/wps/discriminant.py

@@ -100,6 +100,11 @@ def get_discriminant(
     -------
     np.ndarray
         Array of discriminant values.
+
+    Raises
+    ------
+    ValueError
+        If the signal flavour is not recognised.
     """
     tagger_funcs: dict[str, Callable] = {
         "bjets": btag_discriminant,

atlas_ftag_tools-0.2.1.dist-info/RECORD

@@ -1,26 +0,0 @@
-ftag/__init__.py,sha256=aCC6idmHdlETdgAFN4PeESN67WpeVqo85Nledw9VkH4,629
-ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
-ftag/cuts.py,sha256=a0BJj4cVRunc-hFLPloGvNoSFvRmZg2kVLv7sA0iAaI,2817
-ftag/flavour.py,sha256=qvgp4DarOdcQgjae_NWnd81k_YqdmFY74lOKky2lpb8,3568
-ftag/flavours.yaml,sha256=h0A2cw-je1oCe-rh5qhuL1BhDNhu2aLkf7W9y-Cpy3g,6959
-ftag/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
-ftag/mock.py,sha256=9V6sAT4_t-rhR67q9KHaj1NKAeqU7lQjWxiOxEzk8Sw,4338
-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/transform.py,sha256=uEGGJSnqoKOzLYQv650XdK_kDNw4Aw-5dc60z9Dp_y0,3963
-ftag/vds.py,sha256=nRViQZQIORB95nC7NZsW3KsSoGkLzEdOsuCViH5h8-U,3296
-ftag/hdf5/__init__.py,sha256=LFDNxVOCp58SvLHwQhdT68Q-KBMS_i6jBrbXoRpHzbM,354
-ftag/hdf5/h5move.py,sha256=oYpRu0IDCIJIQ2ML52HBAdoyDxmKkHTeM9JdbPEgKfI,947
-ftag/hdf5/h5reader.py,sha256=et-_LXt942xegqc14bPapUgIO7MUfC2m04uJslLkXxI,13579
-ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
-ftag/hdf5/h5utils.py,sha256=-4zKTMtNCrDZr_9Ww7uzfsB7M7muBKpmm_1IkKJnHOI,3222
-ftag/hdf5/h5writer.py,sha256=wVyurIgfSBtvZTX-v0v3R5-8JOwWK_yF1rUX-RewXzY,4826
-ftag/wps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ftag/wps/discriminant.py,sha256=c2-bal124yY1rZjK87iQHb9RRo4pus6SVue6SfMKrRY,3749
-ftag/wps/working_points.py,sha256=VTU6OD40ULAJQD0MlD1EZd33q8ociUvFX1YrhgJFvXc,9722
-atlas_ftag_tools-0.2.1.dist-info/METADATA,sha256=KpKapru-RWOoq01J5sm04jE2brwpoPxDGVPTZoppl8c,5168
-atlas_ftag_tools-0.2.1.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-atlas_ftag_tools-0.2.1.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
-atlas_ftag_tools-0.2.1.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.2.1.dist-info/RECORD,,