Load ptau in Results class (#242)

* taujets in hlapi * changelog * bump tools * update plots * fix test
umami-hep · Feb 26, 2024 · 6e0264f · 6e0264f
1 parent a0f8edf
commit 6e0264f
Show file tree

Hide file tree

Showing 7 changed files with 18 additions and 2 deletions.
diff --git a/changelog.md b/changelog.md
@@ -1,7 +1,8 @@
 # Changelog
 
 ### [Latest]
-- Update ruff [!241](https://github.com/umami-hep/puma/pull/241)
+- Load ptau in Results class [!242](https://github.com/umami-hep/puma/pull/242)
+- Update ruff, improve coverage [!241](https://github.com/umami-hep/puma/pull/241)
 - Extend ftau support and remove discriminant code [!240](https://github.com/umami-hep/puma/pull/240)
 - Fix ROC x-range config for yuma [!239](https://github.com/umami-hep/puma/pull/239)
 - Allow for base dir Yuma loading [!238](https://github.com/umami-hep/puma/pull/238)

diff --git a/puma/hlplots/results.py b/puma/hlplots/results.py
@@ -203,6 +203,8 @@ def check_nan(data: np.ndarray) -> np.ndarray:
         # set tagger output nodes
         for tagger in taggers:
             tagger.output_flavours = self.flavours
+            if "ftau" in tagger.fxs:
+                tagger.output_flavours += [Flavours.taujets]
 
         # get a list of all variables to be loaded from the file
         if not isinstance(cuts, Cuts):

diff --git a/puma/tests/expected_plots/test_bjets_fraction_scan.png b/puma/tests/expected_plots/test_bjets_fraction_scan.png
diff --git a/puma/tests/expected_plots/test_cjets_fraction_scan.png b/puma/tests/expected_plots/test_cjets_fraction_scan.png
diff --git a/puma/tests/hlplots/test_results.py b/puma/tests/hlplots/test_results.py
@@ -98,6 +98,15 @@ def test_add_taggers_with_cuts(self):
         results.add_taggers_from_file(taggers, fname, cuts=cuts)
         self.assertEqual(list(results.taggers.values()), taggers)
 
+    def test_add_taggers_taujets(self):
+        # get mock file and rename variables match taujets
+        fname = get_mock_file()[0]
+        results = Results(signal="bjets", sample="test")
+        taggers = [Tagger("MockTagger", fxs={"fc": 0.1, "fb": 0.1, "ftau": 0.1})]
+        results.add_taggers_from_file(taggers, fname)
+        assert "MockTagger_ptau" in taggers[0].scores.dtype.names
+        taggers[0].discriminant("bjets")
+
     def test_add_taggers_hbb(self):
         # get mock file and rename variables match hbb
         f = get_mock_file()[1]

diff --git a/pyproject.toml b/pyproject.toml
@@ -33,6 +33,10 @@ lint.ignore = [
     "PTH113", "S113", "TCH"
 ]
 
+[tool.ruff.lint.flake8-pytest-style]
+fixture-parentheses = false
+mark-parentheses = false
+
 [tool.ruff.lint.isort]
 required-imports = ["from __future__ import annotations"]
 

diff --git a/requirements.txt b/requirements.txt
@@ -14,4 +14,4 @@ scipy==1.10.1
 tables==3.7.0
 testfixtures==7.0.0
 palettable==3.3.0
-atlas-ftag-tools==0.1.18
+atlas-ftag-tools==0.1.19