Integrate tools package

.github/workflows/pages.yml

@@ -2,10 +2,10 @@ name: Pages
 on:
   pull_request:
     branches:
-      - '*'
+      - "*"
   push:
     branches:
-      - 'main'
+      - "main"
 
 jobs:
   build-sphinx-docs:

.github/workflows/testing.yml

@@ -2,12 +2,12 @@ name: Test and publish
 on:
   pull_request:
     branches:
-      - '*'
+      - "*"
   push:
     branches:
-      - 'main'
+      - "main"
     tags:
-      - '*'
+      - "*"
 
 jobs:
   unit_tests:
@@ -21,7 +21,7 @@ jobs:
       - run: echo "The name of your branch is ${{ github.ref }} and your repository is ${{ github.repository }}."
       - run: echo "The ${{ github.repository }} repository has been cloned to the runner."
       - run: pip install -r requirements.txt
-      - run: pip install -e . # install in editable mode
+      - run: pip install .
       - name: Run tests
         run: pytest -v --junitxml=coverage.xml --cov-report=term-missing:skip-covered --cov=puma puma/tests/
       - name: Install git
@@ -38,6 +38,7 @@ jobs:
       - run: echo "The job was automatically triggered by a ${{ github.event_name }} event."
       - run: echo "The name of your branch is ${{ github.ref }} and your repository is ${{ github.repository }}."
       - run: echo "The ${{ github.repository }} repository has been cloned to the runner."
+      - run: pip install -r requirements.txt
       - run: pip install .
       - run: mkdir plots
       - run: cd plots && for f in ../examples/*.py; do python "$f"; done

.github/workflows/update_todo_issue.yaml

@@ -2,7 +2,7 @@ name: Update TODO Issue
 on:
   push:
     branches:
-      - 'main'
+      - "main"
 
 jobs:
   update_issue:

changelog.md

@@ -2,6 +2,7 @@
 
 ### [Latest]
 
+- Integrate [atlas-ftag-tools](https://github.com/umami-hep/atlas-ftag-tools/) package [!168](https://github.com/umami-hep/puma/pull/168)
 - HLAPI and CI Updates [!165](https://github.com/umami-hep/puma/pull/165)
 - Extend format saving options [!160](https://github.com/umami-hep/puma/pull/160)
 

docker/Dockerfile

@@ -4,4 +4,4 @@ FROM python:3.8.15-slim
 COPY . /puma_repo
 # Install and remove the folder afterwards
 RUN pip install /puma_repo && rm -rf /puma_repo
-RUN pip install h5py==3.7.0
+RUN pip install h5py==3.8.0

docker/Dockerfile_dev

@@ -1,7 +1,6 @@
 FROM python:3.8.15
 
-RUN apt update && apt install -y \
-    vim
+RUN apt update && apt install -y vim
 
 # install requirements that are needed for development
 COPY requirements.txt .

docs/source/examples/high_level_api.md

@@ -17,14 +17,14 @@ The `Results` object is initialised with the signal class, by default this is `b
 to produce the c-tagging plots, or `Hbb`/`Hcc` for Xbb tagging.
 
 ```py
-§§§examples/high_level_plots.py:1:51§§§
+§§§examples/high_level_plots.py:1:48§§§
 ```
 
 ## Discriminant plots
 To plot the discriminant, you can now simply call one function and everything else is handled automatically,
 here for the _b_-jet discriminant
 ```py
-§§§examples/high_level_plots.py:53:55§§§
+§§§examples/high_level_plots.py:50:52§§§
 ```
 
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_disc_b.png width=500>
@@ -34,7 +34,7 @@ here for the _b_-jet discriminant
 
 In the same manner you can plot ROC curves, here for the _b_-tagging performance
 ```py
-§§§examples/high_level_plots.py:57:59§§§
+§§§examples/high_level_plots.py:54:56§§§
 ```
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_roc_b.png width=500>
 
@@ -43,15 +43,15 @@ In the same manner you can plot ROC curves, here for the _b_-tagging performance
 ## Performance vs a variable
 In this case we plot the performance as a function of the jet pT with the same syntax as above for an inclusive working point of 70%
 ```py
-§§§examples/high_level_plots.py:62:75§§§
+§§§examples/high_level_plots.py:59:72§§§
 ```
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_dummy_tagger_pt_bjets_eff.png width=500>
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_dummy_tagger_pt_cjets_rej.png width=500>
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_dummy_tagger_pt_ujets_rej.png width=500>
 
 and similar for a fixed b-efficiency per bin.
 ```py
-§§§examples/high_level_plots.py:77:87§§§
+§§§examples/high_level_plots.py:74:84§§§
 ```
 
 <img src=https://github.com/umami-hep/puma/raw/examples-material/hlplots_dummy_tagger_fixed_per_bin_pt_bjets_eff.png width=500>

examples/high_level_plots.py

@@ -11,9 +11,6 @@
 # define jet selections
 cuts = [("n_truth_promptLepton", "==", 0)]
 
-# can also use `signal="cjets"` to study c-tagging performance
-results = Results(signal="bjets")
-
 # define the taggers
 dips = Tagger(
     name="dips",
@@ -33,7 +30,7 @@
 
 # create the Results object
 # for c-tagging use signal="cjets"
-# for Hbb-tagging use signal="Hbb"
+# for Xbb/cc-tagging use signal="hbb"/"hcc"
 results = Results(signal="bjets")
 
 # load taggers from the file object

examples/plot_discriminant_scores.py

@@ -1,9 +1,10 @@
 """Produce histogram of discriminant from tagger output and labels."""
 
 import numpy as np
+from ftag import Flavours
 
 from puma import Histogram, HistogramPlot
-from puma.utils import get_dummy_2_taggers, get_good_linestyles, global_config
+from puma.utils import get_dummy_2_taggers, get_good_linestyles
 
 # The line below generates dummy data which is similar to a NN output
 df = get_dummy_2_taggers()
@@ -22,7 +23,6 @@
 is_c = df["HadronConeExclTruthLabelID"] == 4
 is_b = df["HadronConeExclTruthLabelID"] == 5
 
-flav_cat = global_config["flavour_categories"]
 taggers = ["dips", "rnnip"]
 linestyles = get_good_linestyles()[:2]
 
@@ -52,7 +52,7 @@
             # linecolour. Instead, we specify a "linestyle legend" further down in the
             # script
             label="Light-flavour jets" if tagger == "dips" else None,
-            colour=flav_cat["ujets"]["colour"],
+            colour=Flavours["ujets"].colour,
             ratio_group="ujets",
             linestyle=linestyle,
         ),
@@ -62,7 +62,7 @@
         Histogram(
             df[is_c][f"disc_{tagger}"],
             label="$c$-jets" if tagger == "dips" else None,
-            colour=flav_cat["cjets"]["colour"],
+            colour=Flavours["cjets"].colour,
             ratio_group="cjets",
             linestyle=linestyle,
         ),
@@ -72,7 +72,7 @@
         Histogram(
             df[is_b][f"disc_{tagger}"],
             label="$b$-jets" if tagger == "dips" else None,
-            colour=flav_cat["bjets"]["colour"],
+            colour=Flavours["bjets"].colour,
             ratio_group="bjets",
             linestyle=linestyle,
         ),

examples/plot_pie.py

@@ -1,6 +1,7 @@
 """Example plotting script for the puma.PiePlot class"""
+from ftag import Flavours
+
 from puma.pie import PiePlot
-from puma.utils import global_config
 
 HadrTruthLabel_fracs = [200_000, 34_000, 150_000, 5_000]
 HadrTruthLabel_labels = ["Light-flavour jets", "$c$-jets", "$b$-jets", "$\\tau$-jets"]
@@ -18,12 +19,11 @@
 # Another example with some styling
 
 # Get the flavour colours from the global config
-flav_cat = global_config["flavour_categories"]
 colours = [
-    flav_cat["ujets"]["colour"],
-    flav_cat["cjets"]["colour"],
-    flav_cat["bjets"]["colour"],
-    flav_cat["taujets"]["colour"],
+    Flavours["ujets"].colour,
+    Flavours["cjets"].colour,
+    Flavours["bjets"].colour,
+    Flavours["taujets"].colour,
 ]
 
 example_plot_2 = PiePlot(

puma/histogram.py

@@ -1,7 +1,11 @@
 """Histogram plot functions."""
+
+from __future__ import annotations
+
 import matplotlib as mpl
 import numpy as np
 import pandas as pd
+from ftag import Flavour, Flavours
 
 from puma.plot_base import PlotBase, PlotLineObject
 from puma.utils import get_good_colours, global_config, logger
@@ -21,7 +25,7 @@ def __init__(  # pylint: disable=too-many-arguments
         values: np.ndarray,
         weights: np.ndarray = None,
         ratio_group: str = None,
-        flavour: str = None,
+        flavour: str | Flavour = None,
         add_flavour_label: bool = True,
         histtype: str = "step",
         **kwargs,
@@ -41,7 +45,7 @@ def __init__(  # pylint: disable=too-many-arguments
             Name of the ratio group this histogram is compared with. The ratio group
             allows you to compare different groups of histograms within one plot.
             By default None
-        flavour: str, optional
+        flavour: str | Flavour, optional
             If set, the correct colour and a label prefix will be extracted from
             `puma.utils.global_config` set for this histogram.
             Allowed values are e.g. "bjets", "cjets", "ujets", "bbjets", ...
@@ -82,7 +86,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self.values = values
         self.weights = weights
         self.ratio_group = ratio_group
-        self.flavour = flavour
+        self.flavour = Flavours[flavour] if isinstance(flavour, str) else flavour
         self.add_flavour_label = add_flavour_label
         self.histtype = histtype
 
@@ -99,18 +103,14 @@ def __init__(  # pylint: disable=too-many-arguments
         )
         # If flavour was specified, extract configuration from global config
         if self.flavour is not None:
-            if self.flavour in global_config["flavour_categories"]:
+            if self.flavour in Flavours:
                 # Use globally defined flavour colour if not specified
                 if self.colour is None:
-                    self.colour = global_config["flavour_categories"][self.flavour][
-                        "colour"
-                    ]
+                    self.colour = self.flavour.colour
                     logger.debug("Histogram colour was set to %s", self.colour)
                 # Add globally defined flavour label if not suppressed
                 if self.add_flavour_label:
-                    global_flavour_label = global_config["flavour_categories"][
-                        self.flavour
-                    ]["legend_label"]
+                    global_flavour_label = self.flavour.label
                     self.label = f"{global_flavour_label} {label}"
                 else:
                     self.label = label