atlas-ftag-tools 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

{atlas_ftag_tools-0.1.14.dist-info → atlas_ftag_tools-0.1.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.1.14
+Version: 0.1.16
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT
@@ -102,6 +102,19 @@ By default the working points are printed to the terminal, but you can save the
 
 See `wps --help` for more options and information.
 
+## Calculate efficiency at discriminant cut 
+
+The same script can be used to calculate the efficiency and rejection values at a given discriminant cut value.
+The script `working_points.py` can be run after intalling this package as follows
+
+```
+wps \
+    --ttbar "path/to/ttbar/*.h5" \
+    --tagger GN120220509 \
+    --fx 0.1
+    --disc_cuts 1.0 1.5
+```
+The `--tagger`, `--fx`, and `--outfile` follow the same procedure as in the 'Calculate WPs' script as described above.
 
 ## H5 Utils
 

{atlas_ftag_tools-0.1.14.dist-info → atlas_ftag_tools-0.1.16.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ftag/__init__.py,sha256=cFcVSJvJIJEYUcFPAaSn0xuKE5lytiF-E3jcbnG6Hpo,629
+ftag/__init__.py,sha256=PN9ZbZHVfxreetH3HnhEbZf1G-J4Up-5si-dQ4rrezg,629
 ftag/cuts.py,sha256=RYAfK3MkEhYhlKQFWQTKu72ZrUwlExFeT8IWLSIgeTU,2798
 ftag/flavour.py,sha256=BLifDbJCoszPzgrU5X3Txff9fTMuVGEjUeU0OtOfllc,2701
 ftag/flavours.yaml,sha256=9ifKyz1_VoHlOaWuf3JEqMLSYyLFedYJf9x1D6dCTnM,5335
@@ -7,17 +7,17 @@ ftag/mock.py,sha256=T2YGeuCkgDEzKPVYwU7Vnq5TNIDrx9eIDVrAQt8s0Mw,4865
 ftag/region.py,sha256=-WxdC0Gy9zz3zEJ2pN779RcxXPG-QEROuMwMoP-Qs0g,353
 ftag/sample.py,sha256=cd-rNHsEY2aWSZdy3V4bOKi3aDMtHTCpjXS8Hl9zwUY,2597
 ftag/transform.py,sha256=uEGGJSnqoKOzLYQv650XdK_kDNw4Aw-5dc60z9Dp_y0,3963
-ftag/vds.py,sha256=l2JttxscusBWUSog8SjSkhPn-4lhOLzZmLaqe4hGiS0,2929
+ftag/vds.py,sha256=PG-NCJdpmK5X2l8i8YWBHx6CY8vosShfqp36facKaYM,3383
 ftag/hdf5/__init__.py,sha256=pZva2TI8nvpBwoawcm_ucVZbGsQJW_u8GGoJgt5mKEw,354
 ftag/hdf5/h5move.py,sha256=1XxiJZ96DYSp8JF0ry3lbRSQaFX72DjUV5vBA6hYw-0,873
-ftag/hdf5/h5reader.py,sha256=D3a9l4jVnvJPlp9fprvnlp2arHPRl_1PT5SEDFO0i_U,13390
+ftag/hdf5/h5reader.py,sha256=et-_LXt942xegqc14bPapUgIO7MUfC2m04uJslLkXxI,13579
 ftag/hdf5/h5split.py,sha256=BlhpsUlqBSDCjVRWuyEq1OImyzwp7VyVkDrCz7pvQKc,2508
 ftag/hdf5/h5utils.py,sha256=wjbAmFY5GoFkWW_AvEKTPbwYMFroHKKFuIcehd91dhM,3222
 ftag/hdf5/h5writer.py,sha256=5jm3vSk4m77lFSUyWm-i_y_USzQRVoKpLL8F_cii65Q,4826
-ftag/wps/discriminant.py,sha256=CnRr1GwfT12GYEI8i-iJxmBLavZN6ZAcv2NIeBP-Xqc,2258
-ftag/wps/working_points.py,sha256=BdVpAh9nPyb2B90Pr48RD_7A17_PzofIojS1iHs5th0,5601
-atlas_ftag_tools-0.1.14.dist-info/METADATA,sha256=3pMQRYJl1OiE2MGMzYM_PMnb8B6D14I7ALfJEnIBMdI,4591
-atlas_ftag_tools-0.1.14.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-atlas_ftag_tools-0.1.14.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
-atlas_ftag_tools-0.1.14.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.1.14.dist-info/RECORD,,
+ftag/wps/discriminant.py,sha256=0YmI3-ieSWReO_uY4-3Sc_85hLVpoCHQ7LfuU1SC_Sg,2318
+ftag/wps/working_points.py,sha256=pLoe8RaVmbtiGM9TxMtWocMohBWrY6JcMCLE_e3XtVY,8033
+atlas_ftag_tools-0.1.16.dist-info/METADATA,sha256=H547FpfWZ6plg5PXm7E9SJhppKGWM8iKM_EpFJaUNJs,5064
+atlas_ftag_tools-0.1.16.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+atlas_ftag_tools-0.1.16.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.1.16.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.1.16.dist-info/RECORD,,

ftag/__init__.py

@@ -1,7 +1,7 @@
 """atlas-ftag-tools - Common tools for ATLAS flavour tagging software."""
 from __future__ import annotations
 
-__version__ = "v0.1.14"
+__version__ = "v0.1.16"
 
 
 from ftag import hdf5

ftag/hdf5/h5reader.py

@@ -354,6 +354,11 @@ class H5Reader:
         int
             Estimated number of jets available after selection cuts, rounded down.
         """
+        # reset rngs to ensure same jets are used for each sample
+        self.rng = np.random.default_rng(42)
+        for r in self.readers:
+            r.rng = np.random.default_rng(42)
+
         # if equal jets is True, available jets is based on the smallest sample
         if self.equal_jets:
             num_jets = []

ftag/vds.py

@@ -62,10 +62,24 @@ def create_virtual_file(
     if not overwrite and out_fname.is_file():
         return out_fname
 
+    # identify common groups across all files
+    common_groups: set[str] = set()
+    for fname in fnames:
+        with h5py.File(fname) as f:
+            groups = set(f.keys())
+            common_groups = groups if not common_groups else common_groups.intersection(groups)
+
+    if not common_groups:
+        raise ValueError("No common groups found across files")
+
+    print("Common groups found:")
+    for group in common_groups:
+        print(f"  {group}")
+
     # create virtual file
     out_fname.parent.mkdir(exist_ok=True)
     with h5py.File(out_fname, "w") as f:
-        for group in h5py.File(fnames[0]):
+        for group in common_groups:
             layout = get_virtual_layout(fnames, group)
             f.create_virtual_dataset(group, layout)
             attrs_dict: dict = {}

ftag/wps/discriminant.py

@@ -58,6 +58,8 @@ def get_discriminant(
     np.ndarray
         Array of discriminant values.
     """
+    if not isinstance(fx, tuple | list):
+        fx = (fx,)
     tagger_funcs = {
         "bjets": btag_discriminant,
         "cjets": ctag_discriminant,

ftag/wps/working_points.py

@@ -36,7 +36,6 @@ def parse_args(args):
         nargs="+",
         type=float,
         help="efficiency working point(s). If -r is specified, values should be 1/efficiency",
-        default=[60, 70, 77, 85],
     )
     parser.add_argument(
         "-t",
@@ -69,6 +68,13 @@ def parse_args(args):
         choices=["ujets", "cjets", "bjets", "hbb", "hcc", "top", "qcd"],
         help="use rejection of specified background class to determine working points",
     )
+    parser.add_argument(
+        "-d",
+        "--disc_cuts",
+        nargs="+",
+        type=float,
+        help="D_x value(s) to calculate efficiency at",
+    )
     parser.add_argument(
         "-n",
         "--num_jets",
@@ -116,8 +122,6 @@ def get_eff_rej(jets, disc, wp, flavs):
 
 
 def get_working_points(args=None):
-    args = parse_args(args)
-
     if args.xbb:
         if len(args.fx) != 2 * len(args.tagger):
             raise ValueError("For Xbb tagging, each tagger must have two fx values")
@@ -180,10 +184,76 @@ def get_working_points(args=None):
         return out
 
 
-def main():
-    out = get_working_points()
+def get_rej_eff_at_disc(jets, tagger, signal, fx, disc_cuts):
+    disc = get_discriminant(jets, tagger, signal, fx)
+    d = {}
+    flavs = Flavours.by_category("single-btag")
+    for dcut in disc_cuts:
+        d[str(dcut)] = {"eff": {}, "rej": {}}
+        for f in flavs:
+            e_discs = disc[f.cuts(jets).idx]
+            eff = sum(e_discs > dcut) / len(e_discs)
+            d[str(dcut)]["eff"][str(f)] = float(f"{eff:.3g}")
+            d[str(dcut)]["rej"][str(f)] = 1 / float(f"{eff:.3g}")
+    return d
+
+
+def get_efficiencies(args=None):
+    if len(args.tagger) != len(args.fx):
+        raise ValueError("Must provide fb/fc for each tagger")
+
+    fx_values = [(fx,) for fx in args.fx]
+    # setup cuts and variables
+    flavs = Flavours.by_category("single-btag")
+    default_cuts = Cuts.from_list(["eta > -2.5", "eta < 2.5"])
+    ttbar_cuts = Cuts.from_list(args.ttbar_cuts) + default_cuts
+    zprime_cuts = Cuts.from_list(args.zprime_cuts) + default_cuts
+    all_vars = next(iter(flavs)).cuts.variables
+    for tagger in args.tagger:
+        all_vars += [f"{tagger}_{f.px}" for f in flavs if "tau" not in f.px]
+
+    # load jets
+    reader = H5Reader(args.ttbar)
+    jets = reader.load({"jets": all_vars}, args.num_jets, cuts=ttbar_cuts)["jets"]
+    if args.zprime:
+        zp_reader = H5Reader(args.zprime)
+        zp_jets = zp_reader.load({"jets": all_vars}, args.num_jets, cuts=zprime_cuts)["jets"]
+
+    # loop over taggers
+    out = {}
+    for tagger, fx in zip(args.tagger, fx_values):
+        out[tagger] = {"signal": args.signal, "fx": fx}
+
+        out[tagger]["ttbar"] = get_rej_eff_at_disc(jets, tagger, args.signal, fx, args.disc_cuts)
+        if args.zprime:
+            out[tagger]["zprime"] = get_rej_eff_at_disc(
+                zp_jets, tagger, args.signal, fx, args.disc_cuts
+            )
+
+    if args.outfile:
+        with open(args.outfile, "w") as f:
+            yaml.dump(out, f, sort_keys=False)
+            return None
+    else:
+        return out
+
+
+def main(args=None):
+    args = parse_args(args)
+
+    if args.effs and args.disc_cuts:
+        raise ValueError("Cannot specify both --effs and --disc_cuts")
+
+    if args.effs:
+        out = get_working_points(args)
+    elif args.disc_cuts:
+        out = get_efficiencies(args)
+    else:
+        raise ValueError("Must specify either --effs or --disc_cuts")
     if out:
         print(yaml.dump(out, sort_keys=False))
+        return out
+    return None
 
 
 if __name__ == "__main__":