atlas-ftag-tools 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

{atlas_ftag_tools-0.2.7.dist-info → atlas_ftag_tools-0.2.9.dist-info}/METADATA

@@ -1,23 +1,23 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: atlas-ftag-tools
-Version: 0.2.7
+Version: 0.2.9
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/atlas-ftag-tools/
 Requires-Python: <3.12,>=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: h5py >=3.0
+Requires-Dist: h5py>=3.0
 Requires-Dist: numpy
-Requires-Dist: PyYAML >=5.1
+Requires-Dist: PyYAML>=5.1
 Provides-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'
-Requires-Dist: pytest-notebook ==0.10.0 ; extra == 'dev'
-Requires-Dist: ipykernel ==6.21.3 ; 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"
+Requires-Dist: pytest_notebook==0.10.0; extra == "dev"
+Requires-Dist: ipykernel==6.21.3; extra == "dev"
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![PyPI version](https://badge.fury.io/py/atlas-ftag-tools.svg)](https://badge.fury.io/py/atlas-ftag-tools)

{atlas_ftag_tools-0.2.7.dist-info → atlas_ftag_tools-0.2.9.dist-info}/RECORD

@@ -1,12 +1,12 @@
-ftag/__init__.py,sha256=k5qBmtC7Ieh0trgm2Ba9Qj_6A2wQSpmAmXo2iIOAaI0,737
+ftag/__init__.py,sha256=YRug5UslRbNoQACbEhdenDS6wXmsmeLjlz4JaKP6eHs,737
 ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
 ftag/cuts.py,sha256=9_ooLZHaO3SnIQBNxwbaPZn-qptGdKnB27FdKQGTiTY,2933
 ftag/flavours.py,sha256=ShH4M2UjQZpZ_NlCctTm2q1tJbzYxjmGteioQ2GcqEU,114
-ftag/flavours.yaml,sha256=E_vpn38qJ3-Tygg2aHlH4wkn_rR1On_lMeaG8OemHCQ,8285
+ftag/flavours.yaml,sha256=87xBvLkMDkicuRMaXtxcao8gjEAgvlTbgjAzpvx4YFM,9021
 ftag/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
 ftag/labeller.py,sha256=IXUgU9UBir39PxVWRKs5r5fqI66Tv0x7nJD3-RYpbrg,2780
 ftag/labels.py,sha256=C7IylPTnc32dFXq8C2Ks2wuljYK3WaY2EsPLGrhtXy8,3932
-ftag/mock.py,sha256=Eyj3tkkaSSnqvS3G6NS7fq8sB__Nx8YE9-OM2_lpdoQ,4992
+ftag/mock.py,sha256=P2D7nNKAz2jRBbmfpHTDj9sBVU9r7HGd0rpWZOJYZ90,5980
 ftag/region.py,sha256=ANv0dGI2W6NJqD9fp7EfqAUReH4FOjc1gwl_Qn8llcM,360
 ftag/sample.py,sha256=3N0FrRcu9l1sX8ohuGOHuMYGD0See6gMO4--7NzR2tE,2538
 ftag/track_selector.py,sha256=fJNk_kIBQriBqV4CPT_3ReJbOUnavDDzO-u3EQlRuyk,2654
@@ -19,10 +19,10 @@ ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
 ftag/hdf5/h5utils.py,sha256=-4zKTMtNCrDZr_9Ww7uzfsB7M7muBKpmm_1IkKJnHOI,3222
 ftag/hdf5/h5writer.py,sha256=9FkClV__UbBqmFsq_h2jwiZnbWVm8QFRL_4mDZZBbTs,5316
 ftag/wps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ftag/wps/discriminant.py,sha256=VJdZlJJUwaTeyxmIDEk23rQSAuvWs6wDA3XRjDI6-_c,4277
-ftag/wps/working_points.py,sha256=cvStSpP8Cbb_FWM8v59tFsscUvdeqi831tLn5BiHUEg,9741
-atlas_ftag_tools-0.2.7.dist-info/METADATA,sha256=oo5m85dK467AWuK-L8xaIbLDVmRUO3r7vA_J1vgR5b8,5169
-atlas_ftag_tools-0.2.7.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-atlas_ftag_tools-0.2.7.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
-atlas_ftag_tools-0.2.7.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.2.7.dist-info/RECORD,,
+ftag/wps/discriminant.py,sha256=GKa0zZlLREdm0mCYSbcWXITYe3VEn3PXOBQiPg5WvgM,2521
+ftag/wps/working_points.py,sha256=jXyikB-bf73EaYFkngjE977-Ytvb9nDTqIdHxWW6WQQ,15960
+atlas_ftag_tools-0.2.9.dist-info/METADATA,sha256=lXC-e0iHMDtvJH8h3i7PcCEKh4_CFz5vlqdGXKSEoV4,5153
+atlas_ftag_tools-0.2.9.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+atlas_ftag_tools-0.2.9.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.2.9.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.9.dist-info/RECORD,,

{atlas_ftag_tools-0.2.7.dist-info → atlas_ftag_tools-0.2.9.dist-info}/WHEEL

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

ftag/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.7"
+__version__ = "v0.2.9"
 
 from ftag import hdf5
 from ftag.cuts import Cuts

ftag/flavours.yaml

@@ -60,12 +60,24 @@
   colour: tab:orange
   category: single-btag-ghost
   _px: pc
-- name: ghostujets
-  label: Light-jets
-  cuts: ["HadronGhostTruthLabelID == 0"]
+- name: ghostsjets
+  label: $s$-jets
+  cuts: ["HadronGhostTruthLabelID == 0", "PartonTruthLabelID == 3"]
+  colour: tab:red
+  category: single-btag-ghost
+  _px: ps
+- name: ghostudjets
+  label: Light-quark-jets
+  cuts: ["HadronGhostTruthLabelID == 0", "PartonTruthLabelID <= 2"]
   colour: tab:green
   category: single-btag-ghost
-  _px: pu
+  _px: pud
+- name: ghostgjets
+  label: Gluon-jets
+  cuts: ["HadronGhostTruthLabelID == 0", "PartonTruthLabelID == 21"]
+  colour: tab:gray
+  category: single-btag-ghost
+  _px: pg
 - name: ghosttaujets
   label: $\tau$-jets
   cuts: ["HadronGhostTruthLabelID == 15"]
@@ -119,6 +131,21 @@
   cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount == 0", "GhostCHadronsFinalCount == 0"]
   colour: "green"
   category: xbb
+- name: htauel
+  label: $H \rightarrow \tau e$
+  cuts: ["R10TruthLabel_R22v1 == 14"]
+  colour: "#b40612"
+  category: xbb
+- name: htaumu
+  label: $H \rightarrow \tau\mu$
+  cuts: ["R10TruthLabel_R22v1 == 15"]
+  colour: "#b40657"
+  category: xbb
+- name: htauhad
+  label: $H \rightarrow \tau\tau$
+  cuts: ["R10TruthLabel_R22v1 == 16"]
+  colour: "#b406a0"
+  category: xbb
 
 # extended Xbb tagging
 - name: tqqb
@@ -272,7 +299,7 @@
   category: isolation
 - name: npxall
   label: non-prompt lepton
-  cuts: ["iffClass notin (2,3,4,11)"]
+  cuts: ["iffClass notin (0,1,2,3,4,11)"]
   colour: "#264653"
   category: isolation
 - name: npxtau

ftag/mock.py

@@ -54,33 +54,74 @@ TRACK_VARS = [
 ]
 
 
-def softmax(x, axis=None):
+def softmax(x: np.ndarray, axis: int | None = None) -> np.ndarray:
+    """Softmax function for numpy arrays.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Input array for the softmax
+    axis : int | None, optional
+        Axis along which the softmax is calculated, by default None
+
+    Returns
+    -------
+    np.ndarray
+        Output array with the softmax output
+    """
     e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
     return e_x / e_x.sum(axis=axis, keepdims=True)
 
 
-def get_mock_scores(labels: np.ndarray, is_xbb: bool = False):
-    means = [
-        [2, 0, 0, 0],
-        [0, 1, 0, 0],
-        [0, 0, 3.5, 0],
-        [0, 0, 0, 1],
-    ]
+def get_mock_scores(labels: np.ndarray, is_xbb: bool = False) -> np.ndarray:
     if not is_xbb:
         label_dict = {"u": 0, "c": 4, "b": 5, "tau": 15}
-        label_mapping = dict(zip(label_dict.values(), means))
-    else:
-        label_dict = {"hbb": 11, "hcc": 12, "top": 1, "qcd": 10}
-        label_mapping = dict(zip(label_dict.values(), means))
 
+    else:
+        label_dict = {
+            "hbb": 11,
+            "hcc": 12,
+            "top": 1,
+            "qcd": 10,
+            "htauel": 14,
+            "htaumu": 15,
+            "htauhad": 16,
+        }
+
+    # Set random seed
     rng = np.random.default_rng(42)
-    nclass = len(label_dict)
-    scores = np.zeros((len(labels), nclass))
-    scales = [1, 2.5, 5, 1]
+
+    # Set a list of possible means/scales
+    mean_scale_list = [1, 2, 2.5, 3.5]
+
+    # Get the number of classes
+    n_classes = len(label_dict)
+
+    # Init a scores array
+    scores = np.zeros((len(labels), n_classes))
+
+    # Generate means/scales
+    means = []
+    scales = []
+    for i in range(n_classes):
+        tmp_means = []
+        tmp_means = [
+            0 if j != i else mean_scale_list[np.random.randint(0, len(mean_scale_list))]
+            for j in range(n_classes)
+        ]
+        means.append(tmp_means)
+        scales.append(mean_scale_list[np.random.randint(0, len(mean_scale_list))])
+
+    # Map the labels to the means
+    label_mapping = dict(zip(label_dict.values(), means))
+
+    # Generate random mock scores
     for i, (label, count) in enumerate(zip(*np.unique(labels, return_counts=True))):
         scores[labels == label] = rng.normal(
-            loc=label_mapping[label], scale=scales[i], size=(count, nclass)
+            loc=label_mapping[label], scale=scales[i], size=(count, n_classes)
         )
+
+    # Pipe scores through softmax
     scores = softmax(scores, axis=1)
     name = "MockXbbTagger" if is_xbb else "MockTagger"
     cols = [f"{name}_p{x}" for x in label_dict]
@@ -103,7 +144,7 @@ def mock_jets(num_jets=1000) -> np.ndarray:
     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)
+    jets["R10TruthLabel_R22v1"] = rng.choice([1, 10, 11, 12, 14, 15, 16], size=num_jets)
     scores = get_mock_scores(jets["HadronConeExclTruthLabelID"])
     xbb_scores = get_mock_scores(jets["R10TruthLabel_R22v1"], is_xbb=True)
     return join_structured_arrays([jets, scores, xbb_scores])

ftag/wps/discriminant.py

@@ -1,22 +1,22 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import TYPE_CHECKING
 
 import numpy as np
 
-from ftag import Flavours
-from ftag.labels import Label, remove_suffix
+if TYPE_CHECKING:  # pragma: no cover
+    from ftag.labels import Label, LabelContainer
 
 
-def discriminant(
+def get_discriminant(
     jets: np.ndarray,
     tagger: str,
     signal: Label,
-    fxs: dict[str, float],
+    flavours: LabelContainer,
+    fraction_values: dict[str, float],
     epsilon: float = 1e-10,
 ) -> np.ndarray:
-    """
-    Get the tagging discriminant.
+    """Calculate the tagging discriminant for a given tagger.
 
     Calculated as the logarithm of the ratio of a specified signal probability
     to a weighted sum ofbackground probabilities.
@@ -24,108 +24,61 @@ def discriminant(
     Parameters
     ----------
     jets : np.ndarray
-        Structed jet array containing tagger scores.
+        Structured array of jets containing tagger outputs
     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).
+        Name of the tagger
+    signal : Label
+        Signal flavour (bjets/cjets or hbb/hcc)
+    fraction_values : dict
+        Dict with the fraction values for the background classes for the given tagger
     epsilon : float, optional
-        A small value added to probabilities to prevent division by zero, by default 1e-10.
+        Small number to avoid division by zero, by default 1e-10
 
     Returns
     -------
     np.ndarray
-        The tagger discriminant values for the jets.
+        Array of discriminant values.
 
     Raises
     ------
     ValueError
-        If a fraction is specified for a denominator that is not present in the input array.
+        If the signal flavour is not recognised.
     """
+    # Init the denominator
     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
-    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):
-    fxs = {"pb": fb, "pc": fc, "pu": 1 - fb - fc}
-    return discriminant(jets, tagger, Flavours.taujets, fxs, epsilon=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 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)
+    # Loop over background flavours
+    for flav in flavours:
+        # Skip signal flavour for denominator
+        if flav == signal:
+            continue
 
+        # Get the probability name of the tagger/flavour combo + fraction value
+        prob_name = f"{tagger}_{flav.px}"
+        fraction_value = fraction_values[flav.frac_str]
 
-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)
+        # If fraction_value for the given flavour is zero, skip it
+        if fraction_value == 0:
+            continue
 
+        # Check that the probability value for the flavour is available
+        if fraction_value > 0 and prob_name not in jets.dtype.names:
+            raise ValueError(
+                f"Nonzero fraction value for {flav.name}, but '{prob_name}' "
+                "not found in input array."
+            )
 
-def hbb_discriminant(jets, tagger, ftop=0.25, fhcc=0.02, epsilon=1e-10):
-    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):
-    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: Label | str, epsilon: float = 1e-10, **fxs
-):
-    """Calculate the b-tag or c-tag discriminant for a given tagger.
+        # Update denominator
+        denominator += jets[prob_name] * fraction_value if prob_name in jets.dtype.names else 0
 
-    Parameters
-    ----------
-    jets : np.ndarray
-        Structured array of jets containing tagger outputs
-    tagger : str
-        Name of the tagger
-    signal : Label
-        Signal flavour (bjets/cjets or hbb/hcc)
-    epsilon : float, optional
-        Small number to avoid division by zero, by default 1e-10
-    **fxs : dict
-        Fractions for the different background flavours.
+    # Calculate numerator
+    signal_field = f"{tagger}_{signal.px}"
 
-    Returns
-    -------
-    np.ndarray
-        Array of discriminant values.
+    # Check that the probability of the signal is available
+    if signal_field not in jets.dtype.names:
+        raise ValueError(
+            f"No signal probability value(s) found for tagger {tagger}. "
+            f"Missing variable: {signal_field}"
+        )
 
-    Raises
-    ------
-    ValueError
-        If the signal flavour is not recognised.
-    """
-    tagger_funcs: dict[str, Callable] = {
-        "bjets": btag_discriminant,
-        "cjets": ctag_discriminant,
-        "taujets": tautag_dicriminant,
-        "hbb": hbb_discriminant,
-        "hcc": hcc_discriminant,
-        "ghostbjets": ghostbtag_discriminant,
-    }
-
-    if str(signal) not in tagger_funcs:
-        raise ValueError(f"Signal flavour must be one of {list(tagger_funcs.keys())}, not {signal}")
-
-    func: Callable = tagger_funcs[str(Flavours[signal])]
-    return func(jets, tagger, **fxs, epsilon=epsilon)
+    return np.log((jets[signal_field] + epsilon) / (denominator + epsilon))

ftag/wps/working_points.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import argparse
+import sys
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 import yaml
@@ -14,30 +16,109 @@ from ftag.cuts import Cuts
 from ftag.hdf5 import H5Reader
 from ftag.wps.discriminant import get_discriminant
 
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Sequence
+
+    from ftag.labels import Label, LabelContainer
+
+
+def parse_args(args: Sequence[str]) -> argparse.Namespace:
+    """Parse the input arguments into a Namespace.
+
+    Parameters
+    ----------
+    args : Sequence[str] | None
+        Sequence of string inputs to the script
+
+    Returns
+    -------
+    argparse.Namespace
+        Namespace with the parsed arguments
+
+    Raises
+    ------
+    ValueError
+        When both --effs and --disc_cuts are provided
+    ValueError
+        When neither --effs nor --disc_cuts are provided
+    ValueError
+        When the number of fraction values is not conistent
+    ValueError
+        When the sum of fraction values for a tagger is not equal to one
+    """
+    # Define the pre-parser which checks the --category
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument(
+        "-c",
+        "--category",
+        default="single-btag",
+        type=str,
+        help="Label category to use for the working point calculation",
+    )
+
+    pre_parser.add_argument(
+        "-s",
+        "--signal",
+        default="bjets",
+        type=str,
+        help="Signal flavour which is to be used",
+    )
+
+    # Parse only --category/--signal and ignore for now all other args
+    pre_args, remaining_argv = pre_parser.parse_known_args(args=args)
 
-def parse_args(args):
+    # Create the "real" parser
     parser = argparse.ArgumentParser(
         description=__doc__,
         formatter_class=HelpFormatter,
     )
+
+    # Add --category/--signal so the help is correctly shown
+    parser.add_argument(
+        "-c",
+        "--category",
+        default="single-btag",
+        type=str,
+        help="Label category to use for the working point calculation",
+    )
+    parser.add_argument(
+        "-s",
+        "--signal",
+        default="bjets",
+        type=str,
+        help="Signal flavour which is to be used",
+    )
+
+    # Check which label category was chosen and load the corresponding flavours
+    flavours = Flavours.by_category(pre_args.category)
+
+    # Build the fraction value arguments for all classes (besides signal)
+    for flav in flavours:
+        # Skip signal
+        if flav.name == pre_args.signal:
+            continue
+
+        # Built fraction values for all background classes
+        parser.add_argument(
+            f"--{flav.frac_str}",
+            nargs="+",
+            required=True,
+            type=float,
+            help=f"{flav.frac_str} value(s) for each tagger",
+        )
+
+    # # Adding the other arguments
     parser.add_argument(
         "--ttbar",
         required=True,
         type=Path,
-        help="path to ttbar sample (supports globbing)",
+        help="Path to ttbar sample (supports globbing)",
     )
     parser.add_argument(
         "--zprime",
         required=False,
         type=Path,
-        help="path to zprime (supports globbing). WPs from ttbar will be reused for zprime",
-    )
-    parser.add_argument(
-        "-e",
-        "--effs",
-        nargs="+",
-        type=float,
-        help="efficiency working point(s). If -r is specified, values should be 1/efficiency",
+        help="Path to zprime (supports globbing). WPs from ttbar will be reused for zprime",
     )
     parser.add_argument(
         "-t",
@@ -48,19 +129,17 @@ def parse_args(args):
         help="tagger name(s)",
     )
     parser.add_argument(
-        "-s",
-        "--signal",
-        default="bjets",
-        choices=["bjets", "cjets", "hbb", "hcc"],
-        type=str,
-        help='signal flavour ("bjets" or "cjets" for b-tagging, "hbb" or "hcc" for Xbb)',
+        "-e",
+        "--effs",
+        nargs="+",
+        type=float,
+        help="Efficiency working point(s). If -r is specified, values should be 1/efficiency",
     )
     parser.add_argument(
         "-r",
         "--rejection",
         default=None,
-        choices=["ujets", "cjets", "bjets", "hbb", "hcc", "top", "qcd"],
-        help="use rejection of specified background class to determine working points",
+        help="Use rejection of specified background class to determine working points",
     )
     parser.add_argument(
         "-d",
@@ -74,219 +153,357 @@ def parse_args(args):
         "--num_jets",
         default=1_000_000,
         type=int,
-        help="use this many jets (post selection)",
+        help="Use this many jets (post selection)",
     )
     parser.add_argument(
         "--ttbar_cuts",
         nargs="+",
         default=["pt > 20e3"],
         type=list,
-        help="selection to apply to ttbar (|eta| < 2.5 is always applied)",
+        help="Selection to apply to ttbar (|eta| < 2.5 is always applied)",
     )
     parser.add_argument(
         "--zprime_cuts",
         nargs="+",
         default=["pt > 250e3"],
         type=list,
-        help="selection to apply to zprime (|eta| < 2.5 is always applied)",
+        help="Selection to apply to zprime (|eta| < 2.5 is always applied)",
     )
     parser.add_argument(
         "-o",
         "--outfile",
         type=Path,
-        help="save results to yaml instead of printing",
-    )
-    parser.add_argument(
-        "--xbb",
-        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",
+        help="Save results to yaml instead of printing",
     )
-    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)
 
-    args.signal = Flavours[args.signal]
+    # Final parse of all arguments
+    parsed_args = parser.parse_args(remaining_argv)
+
+    # Define the signal as an instance of Flavours
+    parsed_args.signal = Flavours[parsed_args.signal]
 
-    if args.effs and args.disc_cuts:
+    # Check that only --effs or --disc_cuts is given
+    if parsed_args.effs and parsed_args.disc_cuts:
         raise ValueError("Cannot specify both --effs and --disc_cuts")
-    if not args.effs and not args.disc_cuts:
+    if not parsed_args.effs and not parsed_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):
-    out = {"eff": {}, "rej": {}}
-    for bkg in list(flavs):
-        bkg_disc = disc[bkg.cuts(jets).idx]
-        eff = sum(bkg_disc > wp) / len(bkg_disc)
-        out["eff"][str(bkg)] = float(f"{eff:.3g}")
-        out["rej"][str(bkg)] = float(f"{1 / eff:.3g}")
+    # Check that all fraction values have the same length
+    for flav in flavours:
+        if flav.name != parsed_args.signal.name and len(getattr(parsed_args, flav.frac_str)) != len(
+            parsed_args.tagger
+        ):
+            raise ValueError(f"Number of {flav.frac_str} values must match number of taggers")
+
+    # Check that all fraction value combinations add up to one
+    for tagger_idx in range(len(parsed_args.tagger)):
+        fraction_value_sum = 0
+        for flav in flavours:
+            if flav.name != parsed_args.signal.name:
+                fraction_value_sum += getattr(parsed_args, flav.frac_str)[tagger_idx]
+
+        # Round the value to take machine precision into account
+        fraction_value_sum = np.round(fraction_value_sum, 8)
+
+        # Check it's equal to one
+        if fraction_value_sum != 1:
+            raise ValueError(
+                "Sum of the fraction values must be one! You gave "
+                f"{fraction_value_sum} for tagger {parsed_args.tagger[tagger_idx]}"
+            )
+    return parsed_args
+
+
+def get_fxs_from_args(args: argparse.Namespace, flavours: LabelContainer) -> list:
+    """Get the fraction values for each tagger from the argparsed inputs.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Input arguments parsed by the argparser
+    flavours : LabelContainer
+        LabelContainer instance of the labels that are used
+
+    Returns
+    -------
+    list
+        List of dicts with the fraction values. Each dict is for one tagger.
+    """
+    # Init the fraction_dict dict
+    fraction_dict = {}
+
+    # Add the fraction values to the dict
+    for flav in flavours:
+        if flav.name != args.signal.name:
+            fraction_dict[flav.frac_str] = vars(args)[flav.frac_str]
+
+    return [{k: v[i] for k, v in fraction_dict.items()} for i in range(len(args.tagger))]
+
+
+def get_eff_rej(
+    jets: np.ndarray,
+    disc: np.ndarray,
+    wp: float,
+    flavours: LabelContainer,
+) -> dict:
+    """Calculate the efficiency/rejection for each flavour.
+
+    Parameters
+    ----------
+    jets : np.ndarray
+        Loaded jets
+    disc : np.ndarray
+        Discriminant values of the jets
+    wp : float
+        Working point that is used
+    flavours : LabelContainer
+        LabelContainer instance of the flavours used
+
+    Returns
+    -------
+    dict
+        Dict with the efficiency/rejection values for each flavour
+    """
+    # Init an out dict
+    out: dict[str, dict] = {"eff": {}, "rej": {}}
+
+    # Loop over the flavours
+    for flav in flavours:
+        # Calculate discriminant values and efficiencies/rejections
+        flav_disc = disc[flav.cuts(jets).idx]
+        eff = sum(flav_disc > wp) / len(flav_disc)
+        out["eff"][flav.name] = float(f"{eff:.3g}")
+        out["rej"][flav.name] = float(f"{1 / eff:.3g}")
+
     return out
 
 
-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")
+def get_rej_eff_at_disc(
+    jets: np.ndarray,
+    tagger: str,
+    signal: Label,
+    disc_cuts: list,
+    flavours: LabelContainer,
+    fraction_values: dict,
+) -> dict:
+    """Calculate the efficiency/rejection at a certain discriminant values.
+
+    Parameters
+    ----------
+    jets : np.ndarray
+        Loaded jets used
+    tagger : str
+        Name of the tagger
+    signal : Label
+        Label instance of the signal flavour
+    disc_cuts : list
+        List of discriminant cut values for which the efficiency/rejection is calculated
+    flavours : LabelContainer
+        LabelContainer instance of the flavours that are used
+
+    Returns
+    -------
+    dict
+        Dict with the discriminant cut values and their respective efficiencies/rejections
+    """
+    # Calculate discriminants
+    disc = get_discriminant(
+        jets=jets,
+        tagger=tagger,
+        signal=signal,
+        flavours=flavours,
+        fraction_values=fraction_values,
+    )
+
+    # Init out dict
+    ref_eff_dict: dict[str, dict] = {}
+
+    # Loop over the disc cut values
     for dcut in disc_cuts:
-        d[str(dcut)] = {"eff": {}, "rej": {}}
-        for f in flavs:
-            e_discs = disc[f.cuts(jets).idx]
+        ref_eff_dict[str(dcut)] = {"eff": {}, "rej": {}}
+
+        # Loop over the flavours
+        for flav in flavours:
+            e_discs = disc[flav.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
+            ref_eff_dict[str(dcut)]["eff"][str(flav)] = float(f"{eff:.3g}")
+            ref_eff_dict[str(dcut)]["rej"][str(flav)] = 1 / float(f"{eff:.3g}")
+
+    return ref_eff_dict
+
+
+def setup_common_parts(
+    args: argparse.Namespace,
+) -> tuple[np.ndarray, np.ndarray | None, LabelContainer]:
+    """Load the jets from the files and setup the taggers.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Input arguments from the argparser
 
+    Returns
+    -------
+    tuple[dict, dict | None, list]
+        Outputs the ttbar jets, the zprime jets (if wanted, else None), and the flavours used.
+    """
+    # Get the used flavours
+    flavours = Flavours.by_category(args.category)
 
-def setup_common_parts(args):
-    flavs = Flavours.by_category("single-btag") if not args.xbb else Flavours.by_category("xbb")
+    # Get the cuts for the samples
     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 = list(set(sum((flav.cuts.variables for flav in flavs), [])))
+    # Prepare the loading of the jets
+    all_vars = list(set(sum((flav.cuts.variables for flav in flavours), [])))
     reader = H5Reader(args.ttbar)
     jet_vars = reader.dtypes()["jets"].names
+
+    # Create for all taggers the fraction values
     for tagger in args.tagger:
-        all_vars += [f"{tagger}_{f.px}" for f in flavs if (f"{tagger}_{f.px}" in jet_vars)]
+        all_vars += [
+            f"{tagger}_{flav.px}" for flav in flavours if (f"{tagger}_{flav.px}" in jet_vars)
+        ]
+
+    # Load ttbar jets
+    ttbar_jets = reader.load({"jets": all_vars}, args.num_jets, cuts=ttbar_cuts)["jets"]
+    zprime_jets = None
 
-    # load jets
-    jets = reader.load({"jets": all_vars}, args.num_jets, cuts=ttbar_cuts)["jets"]
-    zp_jets = None
+    # Load zprime jets if needed
     if args.zprime:
-        zp_reader = H5Reader(args.zprime)
-        zp_jets = zp_reader.load({"jets": all_vars}, args.num_jets, cuts=zprime_cuts)["jets"]
+        zprime_reader = H5Reader(args.zprime)
+        zprime_jets = zprime_reader.load({"jets": all_vars}, args.num_jets, cuts=zprime_cuts)[
+            "jets"
+        ]
+
+    else:
+        zprime_jets = None
+
+    return ttbar_jets, zprime_jets, flavours
 
-    return jets, zp_jets, flavs
 
+def get_working_points(args: argparse.Namespace) -> dict | None:
+    """Calculate the working points.
 
-def get_working_points(args=None):
-    jets, zp_jets, flavs = setup_common_parts(args)
-    fxs = get_fxs_from_args(args)
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Input arguments from the argparser
 
-    # loop over taggers
+    Returns
+    -------
+    dict | None
+        Dict with the working points. If args.outfile is given, the function returns None and
+        stored the resulting dict in a yaml file in args.outfile.
+    """
+    # Load the jets and flavours and get the fraction values
+    ttbar_jets, zprime_jets, flavours = setup_common_parts(args=args)
+    fraction_values = get_fxs_from_args(args=args, flavours=flavours)
+
+    # Init an out dict
     out = {}
+
+    # Loop over taggers
     for i, tagger in enumerate(args.tagger):
-        # calculate discriminant
-        out[tagger] = {"signal": str(args.signal), **fxs[i]}
-        disc = get_discriminant(jets, tagger, args.signal, **fxs[i])
+        # Calculate discriminant
+        out[tagger] = {"signal": str(args.signal), **fraction_values[i]}
+        disc = get_discriminant(
+            jets=ttbar_jets,
+            tagger=tagger,
+            signal=args.signal,
+            flavours=flavours,
+            fraction_values=fraction_values[i],
+        )
 
-        # loop over efficiency working points
+        # Loop over efficiency working points
         for eff in args.effs:
             d = out[tagger][f"{eff:.0f}"] = {}
 
+            # Set the working point
             wp_flavour = args.signal
             if args.rejection:
                 eff = 100 / eff  # noqa: PLW2901
                 wp_flavour = args.rejection
 
-            wp_disc = disc[flavs[wp_flavour].cuts(jets).idx]
+            # Calculate the discriminant value of the working point
+            wp_disc = disc[flavours[wp_flavour].cuts(ttbar_jets).idx]
             wp = d["cut_value"] = round(float(np.percentile(wp_disc, 100 - eff)), 3)
 
-            # calculate eff and rej for each flavour
-            d["ttbar"] = get_eff_rej(jets, disc, wp, flavs)
+            # Calculate efficiency and rejection for each flavour
+            d["ttbar"] = get_eff_rej(
+                jets=ttbar_jets,
+                disc=disc,
+                wp=wp,
+                flavours=flavours,
+            )
 
             # calculate for zprime
             if args.zprime:
-                zp_disc = get_discriminant(zp_jets, tagger, Flavours[args.signal], **fxs[i])
-                d["zprime"] = get_eff_rej(zp_jets, zp_disc, wp, flavs)
+                zprime_disc = get_discriminant(
+                    jets=zprime_jets,
+                    tagger=tagger,
+                    signal=args.signal,
+                    flavours=flavours,
+                    fraction_values=fraction_values[i],
+                )
+                d["zprime"] = get_eff_rej(
+                    jets=zprime_jets,
+                    disc=zprime_disc,
+                    wp=wp,
+                    flavours=flavours,
+                )
 
     if args.outfile:
         with open(args.outfile, "w") as f:
             yaml.dump(out, f, sort_keys=False)
             return None
+
     else:
         return out
 
 
-def get_efficiencies(args=None):
-    jets, zp_jets, _ = setup_common_parts(args)
-    fxs = get_fxs_from_args(args)
+def get_efficiencies(args: argparse.Namespace) -> dict | None:
+    """Calculate the efficiencies for the given jets.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Input arguments from the argparser
 
+    Returns
+    -------
+    dict | None
+        Dict with the efficiencies. If args.outfile is given, the function returns None and
+        stored the resulting dict in a yaml file in args.outfile.
+    """
+    # Load the jets and flavours and get the fraction values
+    ttbar_jets, zprime_jets, flavours = setup_common_parts(args=args)
+    fraction_values = get_fxs_from_args(args=args, flavours=flavours)
+
+    # Init an out dict
     out = {}
+
+    # Loop over the taggers
     for i, tagger in enumerate(args.tagger):
-        out[tagger] = {"signal": str(args.signal), **fxs[i]}
+        out[tagger] = {"signal": str(args.signal), **fraction_values[i]}
 
         out[tagger]["ttbar"] = get_rej_eff_at_disc(
-            jets, tagger, args.signal, args.disc_cuts, **fxs[i]
+            jets=ttbar_jets,
+            tagger=tagger,
+            signal=args.signal,
+            disc_cuts=args.disc_cuts,
+            flavours=flavours,
+            fraction_values=fraction_values[i],
         )
         if args.zprime:
             out[tagger]["zprime"] = get_rej_eff_at_disc(
-                zp_jets, tagger, args.signal, args.disc_cuts, **fxs[i]
+                jets=zprime_jets,
+                tagger=tagger,
+                signal=args.signal,
+                disc_cuts=args.disc_cuts,
+                flavours=flavours,
+                fraction_values=fraction_values[i],
             )
 
     if args.outfile:
@@ -297,13 +514,27 @@ def get_efficiencies(args=None):
         return out
 
 
-def main(args=None):
-    args = parse_args(args)
+def main(args: Sequence[str]) -> dict | None:
+    """Main function to run working point calculation.
+
+    Parameters
+    ----------
+    args : Sequence[str] | None, optional
+        Input arguments, by default None
+
+    Returns
+    -------
+    dict | None
+        The output dict with the calculated values. When --outfile
+        was given, the return value is None
+    """
+    parsed_args = parse_args(args=args)
+
+    if parsed_args.effs:
+        out = get_working_points(args=parsed_args)
 
-    if args.effs:
-        out = get_working_points(args)
-    elif args.disc_cuts:
-        out = get_efficiencies(args)
+    elif parsed_args.disc_cuts:
+        out = get_efficiencies(args=parsed_args)
 
     if out:
         print(yaml.dump(out, sort_keys=False))
@@ -312,5 +543,5 @@ def main(args=None):
     return None
 
 
-if __name__ == "__main__":
-    main()
+if __name__ == "__main__":  # pragma: no cover
+    main(args=sys.argv[1:])