Repository: umami-hep/puma
Pull request #305: Small fixes and improvements for vertexing plots
Author: login: samvanstroud
Created: 2025-03-04T14:57:12Z | Merged: 2025-03-25T22:16:25Z

## Summary

This pull request introduces the following changes

* add a comment about producing track assignment recall and purity plots (should disable vertex quality requirements)
* improve atlas tag format
* add option to disable only the heavy flavour track requirement for GN2


## Conformity
- [x] [Changelog entry](https://github.com/umami-hep/puma/blob/main/changelog.md)
- [x] [Documentation](https://umami-hep.github.io/puma/)



Diff:
diff --git a/changelog.md b/changelog.md
index 3a672ad6..aaa03bfb 100644
--- a/changelog.md
+++ b/changelog.md
@@ -2,6 +2,7 @@
 
 ### [Latest]
 
+- Small improvements for vertexing plots [!305](https://github.com/umami-hep/puma/pull/305)
 - Minor change to how n_classes are calculated in confusion matrix [!306](https://github.com/umami-hep/puma/pull/306)
 - Fixing Coverage [!308](https://github.com/umami-hep/puma/pull/308)
 - Bump atlas-ftag-tools to v0.2.9 [!307](https://github.com/umami-hep/puma/pull/307)
diff --git a/puma/hlplots/aux_results.py b/puma/hlplots/aux_results.py
index 1d0bacfa..c1e50eae 100644
--- a/puma/hlplots/aux_results.py
+++ b/puma/hlplots/aux_results.py
@@ -261,6 +261,9 @@ def plot_var_vtx_perf(
 
         if vertex_match_requirement is None:
             vertex_match_requirement = {"eff_req": 0.65, "purity_req": 0.5}
+            # if you want to make the track level recall and purity plots,
+            # don't require any quality requirements
+            # vertex_match_requirement = {"eff_req": 0, "purity_req": 0}
         eff_req = round(vertex_match_requirement["eff_req"] * 100, 1)
         purity_req = round(vertex_match_requirement["purity_req"] * 100, 1)
         vtx_match_str = rf"Recall $\geq {eff_req}\%$, Purity $\geq {purity_req}\%$"
@@ -272,13 +275,16 @@ def plot_var_vtx_perf(
         vtx_metrics = {}
         for tagger in self.taggers.values():
             if "vertexing" not in tagger.aux_tasks:
-                logger.warning(f"{tagger.name} does not have vertexing aux task defined. Skipping.")
-            assert perf_var in tagger.perf_vars, f"{perf_var} not in tagger {tagger.name} data!"
+                logger.warning(
+                    f"{tagger.label} does not have vertexing aux task defined. Skipping."
+                )
+            assert perf_var in tagger.perf_vars, f"{perf_var} not in tagger {tagger.label} data!"
 
             # get cleaned vertex indices and calculate vertexing metrics
             truth_indices, reco_indices = tagger.vertex_indices(incl_vertexing=incl_vertexing)
-            vtx_metrics[tagger.name] = calculate_vertex_metrics(
-                reco_indices, truth_indices, **vertex_match_requirement
+            max_vertices = 20 if not incl_vertexing else 1
+            vtx_metrics[tagger.label] = calculate_vertex_metrics(
+                reco_indices, truth_indices, max_vertices=max_vertices, **vertex_match_requirement
             )
 
         if not vtx_metrics:
@@ -296,7 +302,7 @@ def plot_var_vtx_perf(
                 xlabel=xlabel,
                 logy=False,
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f", {flav.label}\n{vtx_match_str}",
+                atlas_second_tag=atlas_second_tag + f", {flav.label.lower()}\n{vtx_match_str}",
                 y_scale=1.4,
             )
             # $n_{vtx}^{match}/n_{vtx}^{reco}$
@@ -306,41 +312,41 @@ def plot_var_vtx_perf(
                 xlabel=xlabel,
                 logy=False,
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f", {flav.label}\n{vtx_match_str}",
+                atlas_second_tag=atlas_second_tag + f", {flav.label.lower()}\n{vtx_match_str}",
                 y_scale=1.4,
             )
             # $n_{trk}^{match}/n_{trk}^{true}$
             plot_vtx_trk_eff = VarVsVtxPlot(
                 mode="efficiency",
-                ylabel="Track Assignment Efficiency",
+                ylabel="Track assignment efficiency",
                 xlabel=xlabel,
                 logy=False,
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f", {flav.label}\n{vtx_match_str}",
+                atlas_second_tag=atlas_second_tag + f", {flav.label.lower()}",
                 y_scale=1.4,
             )
             # $n_{trk}^{match}/n_{trk}^{reco}$
             plot_vtx_trk_purity = VarVsVtxPlot(
                 mode="purity",
-                ylabel="Track Assignment Purity",
+                ylabel="Track assignment purity",
                 xlabel=xlabel,
                 logy=False,
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f", {flav.label}\n{vtx_match_str}",
+                atlas_second_tag=atlas_second_tag + f", {flav.label.lower()}",
                 y_scale=1.4,
             )
 
             for tagger in self.taggers.values():
-                if tagger.name not in vtx_metrics:
+                if tagger.label not in vtx_metrics:
                     continue
                 is_flavour = tagger.is_flav(flav)
-                include_sum = vtx_metrics[tagger.name]["track_overlap"][is_flavour] >= 0
+                include_sum = vtx_metrics[tagger.label]["track_overlap"][is_flavour] >= 0
 
                 vtx_perf = VarVsVtx(
                     x_var=tagger.perf_vars[perf_var][is_flavour],
-                    n_match=vtx_metrics[tagger.name]["n_match"][is_flavour],
-                    n_true=vtx_metrics[tagger.name]["n_ref"][is_flavour],
-                    n_reco=vtx_metrics[tagger.name]["n_test"][is_flavour],
+                    n_match=vtx_metrics[tagger.label]["n_match"][is_flavour],
+                    n_true=vtx_metrics[tagger.label]["n_ref"][is_flavour],
+                    n_reco=vtx_metrics[tagger.label]["n_test"][is_flavour],
                     label=tagger.label,
                     colour=tagger.colour,
                     **kwargs,
@@ -348,17 +354,17 @@ def plot_var_vtx_perf(
                 vtx_trk_perf = VarVsVtx(
                     x_var=tagger.perf_vars[perf_var][is_flavour],
                     n_match=np.sum(
-                        vtx_metrics[tagger.name]["track_overlap"][is_flavour],
+                        vtx_metrics[tagger.label]["track_overlap"][is_flavour],
                         axis=1,
                         where=include_sum,
                     ),
                     n_true=np.sum(
-                        vtx_metrics[tagger.name]["ref_vertex_size"][is_flavour],
+                        vtx_metrics[tagger.label]["ref_vertex_size"][is_flavour],
                         axis=1,
                         where=include_sum,
                     ),
                     n_reco=np.sum(
-                        vtx_metrics[tagger.name]["test_vertex_size"][is_flavour],
+                        vtx_metrics[tagger.label]["test_vertex_size"][is_flavour],
                         axis=1,
                         where=include_sum,
                     ),
@@ -393,24 +399,24 @@ def plot_var_vtx_perf(
 
             plot_vtx_fakes = VarVsVtxPlot(
                 mode="fakes",
-                ylabel="Fake Rate",
+                ylabel="Vertex rate",
                 xlabel=xlabel,
                 logy=False,
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f", {flav.label}",
+                atlas_second_tag=atlas_second_tag + f", {flav.label.lower()}",
                 y_scale=1.4,
             )
 
             for tagger in self.taggers.values():
-                if tagger.name not in vtx_metrics:
+                if tagger.label not in vtx_metrics:
                     continue
                 is_flavour = tagger.is_flav(flav)
 
                 vtx_perf = VarVsVtx(
                     x_var=tagger.perf_vars[perf_var][is_flavour],
-                    n_match=vtx_metrics[tagger.name]["n_match"][is_flavour],
-                    n_true=vtx_metrics[tagger.name]["n_ref"][is_flavour],
-                    n_reco=vtx_metrics[tagger.name]["n_test"][is_flavour],
+                    n_match=vtx_metrics[tagger.label]["n_match"][is_flavour],
+                    n_true=vtx_metrics[tagger.label]["n_ref"][is_flavour],
+                    n_reco=vtx_metrics[tagger.label]["n_test"][is_flavour],
                     label=tagger.label,
                     colour=tagger.colour,
                     **kwargs,
@@ -459,9 +465,12 @@ def plot_vertex_mass(
                 xlabel="$m_{SV}$ [GeV]",
                 ylabel="Normalised number of vertices",
                 atlas_first_tag=self.atlas_first_tag,
-                atlas_second_tag=atlas_second_tag + f"\n{vertexing_text} vertexing, {flav.label}",
+                atlas_second_tag=atlas_second_tag
+                + f"\n{vertexing_text} vertexing, {flav.label.lower()}",
                 y_scale=1.7,
                 n_ratio_panels=1,
+                ymin_ratio=[0.8],
+                ymax_ratio=[1.4],
                 **kwargs,
             )
 
@@ -472,7 +481,7 @@ def plot_vertex_mass(
                     ylabel="Normalised number of vertices",
                     atlas_first_tag=self.atlas_first_tag,
                     atlas_second_tag=atlas_second_tag
-                    + f"\n{vertexing_text} vertexing, {flav.label}",
+                    + f"\n{vertexing_text} vertexing, {flav.label.lower()}",
                     y_scale=1.4,
                     **kwargs,
                 )
@@ -480,7 +489,7 @@ def plot_vertex_mass(
             for i, tagger in enumerate(self.taggers.values()):
                 if "vertexing" not in tagger.aux_tasks:
                     logger.warning(
-                        f"{tagger.name} does not have vertexing aux task defined. Skipping."
+                        f"{tagger.label} does not have vertexing aux task defined. Skipping."
                     )
 
                 assert {"pt", "eta", "dphi"}.issubset(
@@ -623,3 +632,4 @@ def plot_track_origin_confmat(
             plot_cm.draw(cm)
             base = tagger.name + "_trackOrigin_cm"
             plot_cm.savefig(self.get_filename(base))
+            print("saved file with name", self.get_filename(base))
diff --git a/puma/hlplots/tagger.py b/puma/hlplots/tagger.py
index 5dbe2b19..52d916ae 100644
--- a/puma/hlplots/tagger.py
+++ b/puma/hlplots/tagger.py
@@ -42,6 +42,7 @@ class Tagger:
     output_flavours: list = None
     disc_cut: float = None
     working_point: float = None
+    vertexing_require_hf_track: bool = True
 
     # Used only by YUMA
     yaml_name: str = None
@@ -343,7 +344,10 @@ def vertex_indices(self, incl_vertexing=False):
                 else self.aux_scores["track_origin"][i]
             )
             reco_indices[i] = clean_reco_vertices(
-                reco_indices[i], reco_track_origin, incl_vertexing=incl_vertexing
+                reco_indices[i],
+                reco_track_origin,
+                incl_vertexing=incl_vertexing,
+                require_hf_track=self.vertexing_require_hf_track,
             )
 
         return truth_indices, reco_indices
diff --git a/puma/utils/vertexing.py b/puma/utils/vertexing.py
index 45f77375..5d31543d 100644
--- a/puma/utils/vertexing.py
+++ b/puma/utils/vertexing.py
@@ -277,7 +277,9 @@ def clean_truth_vertices(truth_vertices, truth_track_origin, incl_vertexing=Fals
     return truth_vertices
 
 
-def clean_reco_vertices(reco_vertices, reco_track_origin=None, incl_vertexing=False):
+def clean_reco_vertices(
+    reco_vertices, reco_track_origin=None, incl_vertexing=False, require_hf_track=True
+):
     """
     Clean reconstructed vertices for each track in a single jet. This function
     removes the vertex with the most reco PV tracks if track origin classification
@@ -293,6 +295,8 @@ def clean_reco_vertices(reco_vertices, reco_track_origin=None, incl_vertexing=Fa
         Array containing reco track origin labels for each track in a jet.
     incl_vertexing: bool, optional
         Whether to merge all vertex indices, by default False.
+    require_hf_track: bool, optional
+        Whether to require at least one track from HF to keep a vertex, by default True.
 
     Returns
     -------
@@ -314,12 +318,13 @@ def clean_reco_vertices(reco_vertices, reco_track_origin=None, incl_vertexing=Fa
             )
 
         # remove vertices with no tracks from HF
-        hf_vertex_indices = np.unique(reco_vertices[np.isin(reco_track_origin, [3, 4, 5])])
-        reco_vertices = clean_indices(
-            reco_vertices,
-            np.isin(reco_vertices, hf_vertex_indices, invert=True),
-            mode="remove",
-        )
+        if require_hf_track:
+            hf_vertex_indices = np.unique(reco_vertices[np.isin(reco_track_origin, [3, 4, 5])])
+            reco_vertices = clean_indices(
+                reco_vertices,
+                np.isin(reco_vertices, hf_vertex_indices, invert=True),
+                mode="remove",
+            )
 
         # merge remaining vertices for inclusive performance
         if incl_vertexing:


Reviews:
user: login: samvanstroud
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T14:59:26Z
user: login: jmw464
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T16:17:37Z
user: login: samvanstroud
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T16:33:57Z
user: login: jmw464
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T16:34:47Z
user: login: jmw464
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T16:39:43Z
user: login: samvanstroud
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T08:40:56Z
user: login: jmw464
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T22:08:44Z
user: login: afroch
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T22:16:09Z

Review comments:
diff_hunk: @@ -227,13 +227,10 @@ def calculate_vertex_metrics(
         max_index = min(metrics["n_match"][i], max_vertices)
 
         # write out vertexing purity metrics
-        metrics["track_overlap"][i, :max_index] = common_tracks[associations][:max_index]
-        metrics["test_vertex_size"][i, :max_index] = test_vertices[
-            associations.sum(axis=1).astype(bool)
-        ].sum(axis=1)[:max_index]
-        metrics["ref_vertex_size"][i, :max_index] = ref_vertices[
-            associations.sum(axis=0).astype(bool)
-        ].sum(axis=1)[:max_index]
+        # assert (common_tracks[associations] == common_tracks).all()
+        metrics["track_overlap"][i, :max_index] = common_tracks[:, :max_index]
+        metrics["test_vertex_size"][i, :max_index] = test_vertices.sum(axis=1)[:max_index]
+        metrics["ref_vertex_size"][i, :max_index] = ref_vertices.sum(axis=1)[:max_index]
path: puma/utils/vertexing.py
user: login: samvanstroud
body: @jmw464 I made these changes trying to disable the vertex quality requirements, but it didn't solve the problem so I should probably revert. I assume this will mess up the exclusive vertexing case?
created_at: 2025-03-04T14:58:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -261,6 +261,9 @@ def plot_var_vtx_perf(
 
         if vertex_match_requirement is None:
             vertex_match_requirement = {"eff_req": 0.65, "purity_req": 0.5}
+            # if you want to make the track level recall and purity plots,
+            # don't require any quality requirements
+            # vertex_match_requirement = {"eff_req": 0, "purity_req": 0}
path: puma/hlplots/aux_results.py
user: login: samvanstroud
body: @jmw464 do you see a better (and still easy) way to produce track assignment recall and purity plots _without_ any vertex quality selections applied?
created_at: 2025-03-04T14:59:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -261,6 +261,9 @@ def plot_var_vtx_perf(
 
         if vertex_match_requirement is None:
             vertex_match_requirement = {"eff_req": 0.65, "purity_req": 0.5}
+            # if you want to make the track level recall and purity plots,
+            # don't require any quality requirements
+            # vertex_match_requirement = {"eff_req": 0, "purity_req": 0}
path: puma/hlplots/aux_results.py
user: login: jmw464
body: I don't quite see why this shouldn't work out of the box if you do it the way you have it commented out. I'm not sure how useful it is though because this basically means that you'll be considering vertices as matched even if they have no tracks in common. In that case recall and purity should just be 0 for such a vertex pair.
created_at: 2025-03-05T16:17:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -261,6 +261,9 @@ def plot_var_vtx_perf(
 
         if vertex_match_requirement is None:
             vertex_match_requirement = {"eff_req": 0.65, "purity_req": 0.5}
+            # if you want to make the track level recall and purity plots,
+            # don't require any quality requirements
+            # vertex_match_requirement = {"eff_req": 0, "purity_req": 0}
path: puma/hlplots/aux_results.py
user: login: samvanstroud
body: That would desirable I think, the idea is to deconvolve the vertex reco efficiency metric from the track assignment metrics. So if a vertex is unmatched to a truth it would correspond to a vertex reconstructed with recall and purity in reality right?
created_at: 2025-03-05T16:33:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -227,13 +227,10 @@ def calculate_vertex_metrics(
         max_index = min(metrics["n_match"][i], max_vertices)
 
         # write out vertexing purity metrics
-        metrics["track_overlap"][i, :max_index] = common_tracks[associations][:max_index]
-        metrics["test_vertex_size"][i, :max_index] = test_vertices[
-            associations.sum(axis=1).astype(bool)
-        ].sum(axis=1)[:max_index]
-        metrics["ref_vertex_size"][i, :max_index] = ref_vertices[
-            associations.sum(axis=0).astype(bool)
-        ].sum(axis=1)[:max_index]
+        # assert (common_tracks[associations] == common_tracks).all()
+        metrics["track_overlap"][i, :max_index] = common_tracks[:, :max_index]
+        metrics["test_vertex_size"][i, :max_index] = test_vertices.sum(axis=1)[:max_index]
+        metrics["ref_vertex_size"][i, :max_index] = ref_vertices.sum(axis=1)[:max_index]
path: puma/utils/vertexing.py
user: login: jmw464
body: Actually to completely remove quality requirements you need to get rid of this: https://github.com/umami-hep/puma/blob/main/puma/utils/vertexing.py#L130
created_at: 2025-03-05T16:34:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -261,6 +261,9 @@ def plot_var_vtx_perf(
 
         if vertex_match_requirement is None:
             vertex_match_requirement = {"eff_req": 0.65, "purity_req": 0.5}
+            # if you want to make the track level recall and purity plots,
+            # don't require any quality requirements
+            # vertex_match_requirement = {"eff_req": 0, "purity_req": 0}
path: puma/hlplots/aux_results.py
user: login: jmw464
body: With 0 recall and purity yeah (if that's what you mean).
created_at: 2025-03-05T16:39:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -227,13 +227,10 @@ def calculate_vertex_metrics(
         max_index = min(metrics["n_match"][i], max_vertices)
 
         # write out vertexing purity metrics
-        metrics["track_overlap"][i, :max_index] = common_tracks[associations][:max_index]
-        metrics["test_vertex_size"][i, :max_index] = test_vertices[
-            associations.sum(axis=1).astype(bool)
-        ].sum(axis=1)[:max_index]
-        metrics["ref_vertex_size"][i, :max_index] = ref_vertices[
-            associations.sum(axis=0).astype(bool)
-        ].sum(axis=1)[:max_index]
+        # assert (common_tracks[associations] == common_tracks).all()
+        metrics["track_overlap"][i, :max_index] = common_tracks[:, :max_index]
+        metrics["test_vertex_size"][i, :max_index] = test_vertices.sum(axis=1)[:max_index]
+        metrics["ref_vertex_size"][i, :max_index] = ref_vertices.sum(axis=1)[:max_index]
path: puma/utils/vertexing.py
user: login: samvanstroud
body: I checked and this didn't have a visible impact on the results, so closing this
created_at: 2025-03-11T08:40:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-04T15:12:47Z
body: ## [Codecov](https://app.codecov.io/gh/umami-hep/puma/pull/305?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=umami-hep) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 98.04%. Comparing base [(`0bce607`)](https://app.codecov.io/gh/umami-hep/puma/commit/0bce6078213d7653554190e07e499251270fb082?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=umami-hep) to head [(`a1c33d5`)](https://app.codecov.io/gh/umami-hep/puma/commit/a1c33d55f888abde0bbfa5fb9972ed09d35324a5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=umami-hep).
> Report is 1 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #305   +/-   ##
=======================================
  Coverage   98.04%   98.04%           
=======================================
  Files          55       55           
  Lines        5669     5673    +4     
=======================================
+ Hits         5558     5562    +4     
  Misses        111      111           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/umami-hep/puma/pull/305?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=umami-hep).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=umami-hep).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: samvanstroud
created_at: 2025-03-11T08:42:09Z
body: @jmw464 could you have another quick look at this and merge if you are happy with it?
author_association: CONTRIBUTOR
user: login: jmw464
created_at: 2025-03-11T18:57:10Z
body: @samvanstroud looks good to me, thanks for doing this!
author_association: CONTRIBUTOR
user: login: samvanstroud
created_at: 2025-03-25T10:21:27Z
body: @jmw464 could you approve/merge this?
author_association: CONTRIBUTOR
user: login: jmw464
created_at: 2025-03-25T22:09:39Z
body: @samvanstroud Apologies, I think I have read-only permissions
author_association: CONTRIBUTOR