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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.1.16
+Version: 0.1.17
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT
@@ -87,11 +87,12 @@ The script is `working_points.py` and can be run after installing this package w
 ```
 wps \
     --ttbar "path/to/ttbar/*.h5" \
-    --tagger GN120220509 \
-    --fx 0.1
+    --tagger GN2v01 \
+    --fc 0.1
 ```
 
-Both the `--tagger` and `--fx` options accept a list if you want to get the WPs for multiple taggers.
+Both the `--tagger` and `--fc` options accept a list if you want to get the WPs for multiple taggers.
+If you are doing c-tagging or xbb-tagging, dedicated fx arguments are available ()you can find them all with `-h`.
 
 If you want to use the `ttbar` WPs get the efficiencies and rejections for the `zprime` sample, you can add `--zprime "path/to/zprime/*.h5"` to the command.
 Note that a default selection of $p_T > 250 ~GeV$ to jets in the `zprime` sample.
@@ -110,7 +111,7 @@ The script `working_points.py` can be run after intalling this package as follow
 ```
 wps \
     --ttbar "path/to/ttbar/*.h5" \
-    --tagger GN120220509 \
+    --tagger GN2v01 \
     --fx 0.1
     --disc_cuts 1.0 1.5
 ```

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

@@ -1,23 +1,24 @@
-ftag/__init__.py,sha256=PN9ZbZHVfxreetH3HnhEbZf1G-J4Up-5si-dQ4rrezg,629
+ftag/__init__.py,sha256=BZi1ffhTco63EC7njY1GtrGgoWXLbBQ-V5ptdfwvjRE,629
 ftag/cuts.py,sha256=RYAfK3MkEhYhlKQFWQTKu72ZrUwlExFeT8IWLSIgeTU,2798
-ftag/flavour.py,sha256=BLifDbJCoszPzgrU5X3Txff9fTMuVGEjUeU0OtOfllc,2701
+ftag/flavour.py,sha256=LbvTQzPKHaGFO5-6XHE-VOV1OgTBy93gYsVcRDMkM64,3570
 ftag/flavours.yaml,sha256=9ifKyz1_VoHlOaWuf3JEqMLSYyLFedYJf9x1D6dCTnM,5335
 ftag/git_check.py,sha256=TvF502eqDrYzhI-SgruVolx1BPJi-J0mswc4pmgaYY8,1621
-ftag/mock.py,sha256=T2YGeuCkgDEzKPVYwU7Vnq5TNIDrx9eIDVrAQt8s0Mw,4865
+ftag/mock.py,sha256=DZ72l8vi7h5SzgwjZklXdrZK1TI7Np7VApNQMNOVyxU,5321
 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=PG-NCJdpmK5X2l8i8YWBHx6CY8vosShfqp36facKaYM,3383
+ftag/vds.py,sha256=VuN62n-8JC2t-79vlcwKYJsLGRb0If0CVk_PXK3yLyA,3288
 ftag/hdf5/__init__.py,sha256=pZva2TI8nvpBwoawcm_ucVZbGsQJW_u8GGoJgt5mKEw,354
 ftag/hdf5/h5move.py,sha256=1XxiJZ96DYSp8JF0ry3lbRSQaFX72DjUV5vBA6hYw-0,873
 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=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/wps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ftag/wps/discriminant.py,sha256=J2Cst9slZCLoHZYTeltQSmN1Uoa8GUX529rwNlHmyfI,3519
+ftag/wps/working_points.py,sha256=H602ikwrLVoe2Vq13DPDYSxBpeJzN-2sIKZUeAYL3Pc,9694
+atlas_ftag_tools-0.1.17.dist-info/METADATA,sha256=EUSiZ7BKKcjPdD1M0qT1rZkGdv9UnF3iXySdYjOkpTs,5169
+atlas_ftag_tools-0.1.17.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+atlas_ftag_tools-0.1.17.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.1.17.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.1.17.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.16"
+__version__ = "v0.1.17"
 
 
 from ftag import hdf5

ftag/flavour.py

@@ -35,6 +35,10 @@ class Flavour:
     def rej_str(self) -> str:
         return self.label.replace("jets", "jet") + " rejection"
 
+    @property
+    def frac_str(self) -> str:
+        return "f" + self.name.replace("jets", "jet")
+
     def __str__(self) -> str:
         return self.name
 
@@ -62,6 +66,13 @@ class FlavourContainer:
             flavour = flavour.name
         return flavour in self.flavours
 
+    def __eq__(self, other) -> bool:
+        if isinstance(other, FlavourContainer):
+            return self.flavours == other.flavours
+        if isinstance(other, list) and all(isinstance(f, str) for f in other):
+            return {f.name for f in self} == set(other)
+        return False
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({', '.join([f.name for f in self])})"
 
@@ -95,5 +106,15 @@ class FlavourContainer:
 
         return cls(flavours_dict)
 
+    @classmethod
+    def from_list(cls, flavours: list[Flavour]) -> FlavourContainer:
+        return cls({f.name: f for f in flavours})
+
+    def backgrounds(self, flavour: Flavour, keep_possible_signals: bool = True) -> FlavourContainer:
+        bkg = [f for f in self if f.category == flavour.category and f != flavour]
+        if not keep_possible_signals:
+            bkg = [f for f in bkg if f.name not in ["ujets", "qcd"]]
+        return FlavourContainer.from_list(bkg)
+
 
 Flavours = FlavourContainer.from_yaml()

ftag/mock.py

@@ -57,18 +57,31 @@ def softmax(x, axis=None):
     return e_x / e_x.sum(axis=axis, keepdims=True)
 
 
-def get_mock_scores(labels: np.ndarray):
+def get_mock_scores(labels: np.ndarray, inc_tau: bool = False):
     rng = np.random.default_rng(42)
-    scores = np.zeros((len(labels), 3))
+    nclass = 3 + inc_tau
+    scores = np.zeros((len(labels), nclass))
+
     for label, count in zip(*np.unique(labels, return_counts=True)):
-        if label in (0, 15):
-            scores[labels == label] = rng.normal(loc=[2, 0, 0], scale=1, size=(count, 3))
+        if label == 0:
+            scores[labels == label] = rng.normal(
+                loc=[2, 0, 0] + [0] * inc_tau, scale=1, size=(count, nclass)
+            )
         elif label == 4:
-            scores[labels == label] = rng.normal(loc=[0, 1, 0], scale=2.5, size=(count, 3))
+            scores[labels == label] = rng.normal(
+                loc=[0, 1, 0] + [0] * inc_tau, scale=2.5, size=(count, nclass)
+            )
         elif label == 5:
-            scores[labels == label] = rng.normal(loc=[0, 0, 3.5], scale=5, size=(count, 3))
+            scores[labels == label] = rng.normal(
+                loc=[0, 0, 3.5] + [0] * inc_tau, scale=5, size=(count, nclass)
+            )
+        elif label == 15:
+            scores[labels == label] = rng.normal(
+                loc=[0, 0, 0] + [1] * inc_tau, scale=1, size=(count, nclass)
+            )
     scores = softmax(scores, axis=1)
-    cols = [f"MockTagger_p{x}" for x in ["u", "c", "b"]]
+    cols = [f"MockTagger_p{x}" for x in ["u", "c", "b"]] + (["MockTagger_ptau"] if inc_tau else [])
+
     return u2s(scores, dtype=np.dtype([(name, "f4") for name in cols]))
 
 
@@ -94,6 +107,7 @@ def get_mock_file(
     fname: str | None = None,
     tracks_name: str = "tracks",
     num_tracks: int = 40,
+    inc_tau: bool = False,
 ) -> tuple[str, h5py.File]:
     # setup jets
     rng = np.random.default_rng(42)
@@ -109,7 +123,7 @@ def get_mock_file(
     jets["n_truth_promptLepton"] = 0
 
     # add tagger scores
-    scores = get_mock_scores(jets["HadronConeExclTruthLabelID"])
+    scores = get_mock_scores(jets["HadronConeExclTruthLabelID"], inc_tau=inc_tau)
     xbb_scores = get_mock_xbb_scores(jets["R10TruthLabel_R22v1"])
     jets = join_structured_arrays([jets, scores, xbb_scores])
 

ftag/vds.py

@@ -72,10 +72,6 @@ def create_virtual_file(
     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:

ftag/wps/discriminant.py

@@ -1,42 +1,80 @@
 from __future__ import annotations
 
+from typing import Callable
+
 import numpy as np
 
 from ftag.flavour import Flavour, Flavours
 
 
-def btag_discriminant(jets, tagger, fc=0.1, epsilon=1e-10):
-    pb, pc, pu = (jets[f"{tagger}_pb"], jets[f"{tagger}_pc"], jets[f"{tagger}_pu"])
-    return np.log((pb + epsilon) / ((1.0 - fc) * pu + fc * pc + epsilon))
+def discriminant(
+    jets: np.ndarray,
+    tagger: str,
+    signal: Flavour,
+    fxs: dict[str, float],
+    epsilon: float = 1e-10,
+) -> np.ndarray:
+    """
+    Get the tagging discriminant.
 
+    Calculated as the logarithm of the ratio of a specified signal probability
+    to a weighted sum ofbackground probabilities.
 
-def ctag_discriminant(jets, tagger, fb=0.2, epsilon=1e-10):
-    pb, pc, pu = (jets[f"{tagger}_pb"], jets[f"{tagger}_pc"], jets[f"{tagger}_pu"])
-    return np.log((pc + epsilon) / ((1.0 - fb) * pu + fb * pb + epsilon))
+    Parameters
+    ----------
+    jets : np.ndarray
+        Structed jet array containing tagger scores.
+    tagger : str
+        Name of the tagger, used to construct field names.
+    signal : str
+        Type of signal.
+    fxs : dict[str, float]
+        Dict of background probability names and their fractions.
+        If a fraction is None, it is calculated as (1 - sum of provided fractions).
+    epsilon : float, optional
+        A small value added to probabilities to prevent division by zero, by default 1e-10.
+
+    Returns
+    -------
+    np.ndarray
+        The tagger discriminant values for the jets.
+
+    Raises
+    ------
+    ValueError
+        If a fraction is specified for a denominator that is not present in the input array.
+    """
+    denominator = 0.0
+    for d, fx in fxs.items():
+        name = f"{tagger}_{d}"
+        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))
+
+
+def btag_discriminant(jets, tagger, fc, ftau=0, epsilon=1e-10):
+    fxs = {"pc": fc, "ptau": ftau, "pu": 1 - fc - ftau}
+    return discriminant(jets, tagger, Flavours.bjets, 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)
 
 
 def hbb_discriminant(jets, tagger, ftop=0.25, fhcc=0.02, epsilon=1e-10):
-    phbb = jets[f"{tagger}_phbb"]
-    phcc = jets[f"{tagger}_phcc"]
-    ptop = jets[f"{tagger}_ptop"]
-    pqcd = jets[f"{tagger}_pqcd"]
-    return np.log(phbb / (ftop * ptop + fhcc * phcc + (1 - ftop - fhcc) * pqcd + epsilon))
+    fxs = {"phcc": fhcc, "ptop": ftop, "pqcd": 1 - ftop - fhcc}
+    return discriminant(jets, tagger, Flavours.hbb, fxs, epsilon=epsilon)
 
 
 def hcc_discriminant(jets, tagger, ftop=0.25, fhbb=0.3, epsilon=1e-10):
-    phbb = jets[f"{tagger}_phbb"]
-    phcc = jets[f"{tagger}_phcc"]
-    ptop = jets[f"{tagger}_ptop"]
-    pqcd = jets[f"{tagger}_pqcd"]
-    return np.log(phcc / (ftop * ptop + fhbb * phbb + (1 - ftop - fhbb) * pqcd + epsilon))
+    fxs = {"phbb": fhbb, "ptop": ftop, "pqcd": 1 - ftop - fhbb}
+    return discriminant(jets, tagger, Flavours.hcc, fxs, epsilon=epsilon)
 
 
 def get_discriminant(
-    jets: np.ndarray,
-    tagger: str,
-    signal: Flavour | str,
-    fx: float | tuple[float, ...],
-    epsilon: float = 1e-10,
+    jets: np.ndarray, tagger: str, signal: Flavour | str, epsilon: float = 1e-10, **fxs
 ):
     """Calculate the b-tag or c-tag discriminant for a given tagger.
 
@@ -48,27 +86,25 @@ def get_discriminant(
         Name of the tagger
     signal : Flavour
         Signal flavour (bjets/cjets or hbb/hcc)
-    fx : float, optional
-        Value fb or fc (fhbb or fhcc and ftop for Xbb)
     epsilon : float, optional
         Small number to avoid division by zero, by default 1e-10
+    **fxs : dict
+        Fractions for the different background flavours.
 
     Returns
     -------
     np.ndarray
         Array of discriminant values.
     """
-    if not isinstance(fx, tuple | list):
-        fx = (fx,)
-    tagger_funcs = {
+    tagger_funcs: dict[str, Callable] = {
         "bjets": btag_discriminant,
         "cjets": ctag_discriminant,
         "hbb": hbb_discriminant,
         "hcc": hcc_discriminant,
     }
 
-    func = tagger_funcs.get(str(Flavours[signal]), None)
-    if func is None:
-        raise ValueError(f"Signal flavour must be among {list(tagger_funcs.keys())}, not {signal}")
+    if str(signal) not in tagger_funcs:
+        raise ValueError(f"Signal flavour must be one of {list(tagger_funcs.keys())}, not {signal}")
 
-    return func(jets, tagger, *fx, epsilon)  # type: ignore
+    func: Callable = tagger_funcs[str(Flavours[signal])]
+    return func(jets, tagger, **fxs, epsilon=epsilon)

ftag/wps/working_points.py

@@ -12,12 +12,11 @@ from ftag.hdf5 import H5Reader
 from ftag.wps.discriminant import get_discriminant
 
 
-def parse_args(args):
+def parse_args(args):  # noqa: PLR0912
     parser = argparse.ArgumentParser(
         description="Calculate tagger working points",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-
     parser.add_argument(
         "--ttbar",
         required=True,
@@ -45,14 +44,6 @@ def parse_args(args):
         type=str,
         help="tagger name(s)",
     )
-    parser.add_argument(
-        "-f",
-        "--fx",
-        nargs="+",
-        required=True,
-        type=float,
-        help="fb or fc value(s) for each tagger",
-    )
     parser.add_argument(
         "-s",
         "--signal",
@@ -107,8 +98,90 @@ def parse_args(args):
         action="store_true",
         help="Enable Xbb tagging which expects two fx values ftop and fhcc/fhbb for each tagger",
     )
+    parser.add_argument(
+        "--fb",
+        nargs="+",
+        type=float,
+        help="fb value(s) for each tagger",
+    )
+    parser.add_argument(
+        "--fc",
+        nargs="+",
+        type=float,
+        help="fc value(s) for each tagger",
+    )
+    parser.add_argument(
+        "--ftau",
+        nargs="+",
+        type=float,
+        help="ftau value(s) for each tagger",
+    )
+    parser.add_argument(
+        "--ftop",
+        nargs="+",
+        type=float,
+        help="ftop value(s) for each tagger",
+    )
+    parser.add_argument(
+        "--fhbb",
+        nargs="+",
+        type=float,
+        help="fhbb value(s) for each tagger",
+    )
+    parser.add_argument(
+        "--fhcc",
+        nargs="+",
+        type=float,
+        help="fhcc value(s) for each tagger",
+    )
+    args = parser.parse_args(args)
 
-    return parser.parse_args(args)
+    args.signal = Flavours[args.signal]
+
+    if args.effs and args.disc_cuts:
+        raise ValueError("Cannot specify both --effs and --disc_cuts")
+    if not args.effs and not args.disc_cuts:
+        raise ValueError("Must specify either --effs or --disc_cuts")
+
+    if args.xbb:
+        if args.signal not in [Flavours.hbb, Flavours.hcc]:
+            raise ValueError("Xbb tagging only supports hbb or hcc signal flavours")
+        if args.fb or args.fc or args.ftau:
+            raise ValueError("For Xbb tagging, fb, fc and ftau should not be specified")
+        if not args.ftop:
+            raise ValueError("For Xbb tagging, ftop should be specified")
+        if args.signal == "hbb" and not args.fhcc:
+            raise ValueError("For hbb tagging, fhcc should be specified")
+        if args.signal == "hcc" and not args.fhbb:
+            raise ValueError("For hcc tagging, fhbb should be specified")
+    else:
+        if args.ftop or args.fhbb or args.fhcc:
+            raise ValueError("For single-b tagging, ftop, fhbb and fhcc should not be specified")
+        if args.signal == "bjets" and not args.fc:
+            raise ValueError("For bjets tagging, fc should be specified")
+        if args.signal == "cjets" and not args.fb:
+            raise ValueError("For cjets tagging, fb should be specified")
+        if args.ftau is None:
+            args.ftau = [0.0] * len(args.tagger)
+
+    for fx in ["fb", "fc", "ftau", "ftop", "fhbb", "fhcc"]:
+        if getattr(args, fx) and len(getattr(args, fx)) != len(args.tagger):
+            raise ValueError(f"Number of {fx} values must match number of taggers")
+
+    return args
+
+
+def get_fxs_from_args(args):
+    if args.signal == Flavours.bjets:
+        fxs = {"fc": args.fc, "ftau": args.ftau}
+    elif args.signal == Flavours.cjets:
+        fxs = {"fb": args.fb, "ftau": args.ftau}
+    elif args.signal == Flavours.hbb:
+        fxs = {"ftop": args.ftop, "fhcc": args.fhcc}
+    elif args.signal == Flavours.hcc:
+        fxs = {"ftop": args.ftop, "fhbb": args.fhbb}
+    assert fxs is not None
+    return [{k: v[i] for k, v in fxs.items()} for i in range(len(args.tagger))]
 
 
 def get_eff_rej(jets, disc, wp, flavs):
@@ -121,40 +194,53 @@ def get_eff_rej(jets, disc, wp, flavs):
     return out
 
 
-def get_working_points(args=None):
-    if args.xbb:
-        if len(args.fx) != 2 * len(args.tagger):
-            raise ValueError("For Xbb tagging, each tagger must have two fx values")
-        fx_values = list(zip(args.fx[::2], args.fx[1::2]))
-    else:
-        if len(args.fx) != len(args.tagger):
-            raise ValueError("Number of taggers must match number of fx values")
-        fx_values = [(fx,) for fx in args.fx]
+def get_rej_eff_at_disc(jets, tagger, signal, disc_cuts, **fxs):
+    disc = get_discriminant(jets, tagger, signal, **fxs)
+    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
 
-    # setup cuts and variables
-    flavs = Flavours.by_category("single-btag") if not args.xbb else Flavours.by_category("xbb")
 
+def setup_common_parts(args):
+    flavs = Flavours.by_category("single-btag") if not args.xbb else Flavours.by_category("xbb")
     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
+
+    # prepare to load jets
     all_vars = next(iter(flavs)).cuts.variables
+    reader = H5Reader(args.ttbar)
+    jet_vars = reader.dtypes()["jets"].names
     for tagger in args.tagger:
-        all_vars += [f"{tagger}_{f.px}" for f in flavs if "tau" not in f.px]
+        all_vars += [f"{tagger}_{f.px}" for f in flavs if (f"{tagger}_{f.px}" in jet_vars)]
 
     # load jets
-    reader = H5Reader(args.ttbar)
     jets = reader.load({"jets": all_vars}, args.num_jets, cuts=ttbar_cuts)["jets"]
+    zp_jets = None
     if args.zprime:
         zp_reader = H5Reader(args.zprime)
         zp_jets = zp_reader.load({"jets": all_vars}, args.num_jets, cuts=zprime_cuts)["jets"]
 
+    return jets, zp_jets, flavs
+
+
+def get_working_points(args=None):
+    jets, zp_jets, flavs = setup_common_parts(args)
+    fxs = get_fxs_from_args(args)
+
     # loop over taggers
     out = {}
-    for tagger, fx in zip(args.tagger, fx_values):
-        out[tagger] = {"signal": args.signal, "fx": fx}
-
+    for i, tagger in enumerate(args.tagger):
         # calculate discriminant
-        disc = get_discriminant(jets, tagger, args.signal, fx)
+        out[tagger] = {"signal": str(args.signal), **fxs[i]}
+        disc = get_discriminant(jets, tagger, args.signal, **fxs[i])
 
         # loop over efficiency working points
         for eff in args.effs:
@@ -173,7 +259,7 @@ def get_working_points(args=None):
 
             # calculate for zprime
             if args.zprime:
-                zp_disc = get_discriminant(zp_jets, tagger, Flavours[args.signal], fx)
+                zp_disc = get_discriminant(zp_jets, tagger, Flavours[args.signal], **fxs[i])
                 d["zprime"] = get_eff_rej(zp_jets, zp_disc, wp, flavs)
 
     if args.outfile:
@@ -184,50 +270,20 @@ def get_working_points(args=None):
         return out
 
 
-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]
+    jets, zp_jets, flavs = setup_common_parts(args)
+    fxs = get_fxs_from_args(args)
 
-    # 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}
+    for i, tagger in enumerate(args.tagger):
+        out[tagger] = {"signal": str(args.signal), **fxs[i]}
 
-        out[tagger]["ttbar"] = get_rej_eff_at_disc(jets, tagger, args.signal, fx, args.disc_cuts)
+        out[tagger]["ttbar"] = get_rej_eff_at_disc(
+            jets, tagger, args.signal, args.disc_cuts, **fxs[i]
+        )
         if args.zprime:
             out[tagger]["zprime"] = get_rej_eff_at_disc(
-                zp_jets, tagger, args.signal, fx, args.disc_cuts
+                zp_jets, tagger, args.signal, args.disc_cuts, **fxs[i]
             )
 
     if args.outfile:
@@ -241,18 +297,15 @@ def get_efficiencies(args=None):
 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