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

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

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

{atlas_ftag_tools-0.2.3.dist-info → atlas_ftag_tools-0.2.4.dist-info}/RECORD

@@ -1,28 +1,28 @@
-ftag/__init__.py,sha256=i5XMWwkyGmwIIIaO3veRyOnTsJ0qKyW1V6HsuYc9Dm4,629
+ftag/__init__.py,sha256=Mx2Emsw4TM1YL0wTQMHK36EaQE_ImeV6ukiz1X5BZAU,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/mock.py,sha256=_oy-r3eLllFy33NAoZaKfAx-Rp2vrCdrGj3UsTMks94,4740
 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/track_selector.py,sha256=piSYAN_IkOsrXxKXjXbJpMSseUig5P2BJW5mCwsMUDM,2535
 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/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/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.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/__init__.py

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

ftag/hdf5/h5reader.py

@@ -150,8 +150,13 @@ class H5Reader:
     transform : Transform | None, optional
         Transform to apply to data, by default None
     equal_jets : bool, optional
-        Take the same number of jets (weighted) from each sample, by default True
-        If False, use all jets in each sample.
+        Take the same number of jets (weighted) from each sample, by default True.
+        This is useful when you specify a list of DSIDs for the sample and they are
+        qualitatively different, and you want to ensure that you always return batches
+        with jets from all DSIDs. This is used for example in the QCD resampling for Xbb.
+        If False, use all jets in each sample, allowing for the full available statistics
+        to be used. Useful for example if you have multiple ttbar samples and you want to
+        use all available jets from each sample.
     """
 
     fname: Path | str | list[Path | str]
@@ -162,17 +167,10 @@ class H5Reader:
     weights: list[float] | None = None
     do_remove_inf: bool = False
     transform: Transform | None = None
-    equal_jets: bool = True
+    equal_jets: bool = False
 
     def __post_init__(self) -> None:
         self.rng = np.random.default_rng(42)
-        if not self.equal_jets:
-            log.warning(
-                "equal_jets is set to False, which will result in different number of jets taken"
-                " from each sample. Be aware that this can affect the resampling, so make sure you"
-                " know what you are doing."
-            )
-
         if isinstance(self.fname, (str, Path)):
             self.fname = [self.fname]
 
@@ -283,8 +281,8 @@ class H5Reader:
                     try:
                         samples.append(next(stream))
 
-                    # if equal_jets is True, we can stop when any stream is done
-                    # otherwise if sample is exhausted, mark it as done
+                    # if equal_jets is True, stop when any sample is done
+                    # otherwise if stream is exhausted, mark it as such and continue
                     except StopIteration:
                         if self.equal_jets:
                             return

ftag/mock.py

@@ -109,6 +109,11 @@ def mock_tracks(num_jets=1000, num_tracks=40) -> np.ndarray:
     tracks_dtype = np.dtype(TRACK_VARS)
     tracks = u2s(rng.random((num_jets, num_tracks, len(TRACK_VARS))), tracks_dtype)
     tracks["d0"] *= 5
+
+    # for the shared hits, add some reasonable integer values
+    tracks["numberOfPixelSharedHits"] = rng.integers(0, 3, size=(num_jets, num_tracks))
+    tracks["numberOfSCTSharedHits"] = rng.integers(0, 3, size=(num_jets, num_tracks))
+
     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])

ftag/track_selector.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from ftag import Cuts
+from ftag.cuts import Cut, Cuts
 
 
 @dataclass
@@ -34,7 +34,8 @@ class TrackSelector:
         # 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
+            keep_idx = self._nshared_cut(cut, tracks) if cut.variable == "NSHARED" else cut(tracks)
+            rm_idx[tracks[self.valid_str] & ~keep_idx] = True
 
         # set the values of the tracks that do not pass the cuts to
         for var in tracks.dtype.names:
@@ -51,3 +52,19 @@ class TrackSelector:
         tracks[rm_idx][self.valid_str] = False
 
         return tracks
+
+    def _nshared_cut(self, cut: Cut, tracks: np.ndarray) -> np.ndarray:
+        # hack to apply the FTAG shared hit cut, which requires an intermediate step
+        if cut.variable == "NSHARED" and "NSHARED" in tracks.dtype.names:
+            raise ValueError("NSHARED is a reserved variable name")
+
+        # compute
+        n_pix_shared = tracks["numberOfPixelSharedHits"]
+        n_sct_shared = tracks["numberOfSCTSharedHits"]
+        n_module_shared = n_pix_shared + n_sct_shared / 2
+
+        # convert n_module_shared to structured array
+        n_module_shared = n_module_shared.view(dtype=[(cut.variable, n_module_shared.dtype)])
+
+        # select
+        return cut(n_module_shared)