atlas-ftag-tools 0.2.7__tar.gz → 0.2.9__tar.gz

{atlas_ftag_tools-0.2.7 → atlas_ftag_tools-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
-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

{atlas_ftag_tools-0.2.7 → atlas_ftag_tools-0.2.9}/atlas_ftag_tools.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-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

{atlas_ftag_tools-0.2.7 → atlas_ftag_tools-0.2.9}/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

{atlas_ftag_tools-0.2.7 → atlas_ftag_tools-0.2.9}/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

{atlas_ftag_tools-0.2.7 → atlas_ftag_tools-0.2.9}/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])

atlas_ftag_tools-0.2.9/ftag/wps/discriminant.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:  # pragma: no cover
+    from ftag.labels import Label, LabelContainer
+
+
+def get_discriminant(
+    jets: np.ndarray,
+    tagger: str,
+    signal: Label,
+    flavours: LabelContainer,
+    fraction_values: dict[str, float],
+    epsilon: float = 1e-10,
+) -> np.ndarray:
+    """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.
+
+    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)
+    fraction_values : dict
+        Dict with the fraction values for the background classes for the given tagger
+    epsilon : float, optional
+        Small number to avoid division by zero, by default 1e-10
+
+    Returns
+    -------
+    np.ndarray
+        Array of discriminant values.
+
+    Raises
+    ------
+    ValueError
+        If the signal flavour is not recognised.
+    """
+    # Init the denominator
+    denominator = 0.0
+
+    # 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]
+
+        # 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."
+            )
+
+        # Update denominator
+        denominator += jets[prob_name] * fraction_value if prob_name in jets.dtype.names else 0
+
+    # Calculate numerator
+    signal_field = f"{tagger}_{signal.px}"
+
+    # 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}"
+        )
+
+    return np.log((jets[signal_field] + epsilon) / (denominator + epsilon))