Also apply reconstructable filter for denonminator (#325)

gnn-tracking · May 9, 2023 · 4ffd7db · 4ffd7db
1 parent f79f37f
commit 4ffd7db
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/src/gnn_tracking/metrics/cluster_metrics.py b/src/gnn_tracking/metrics/cluster_metrics.py
@@ -135,6 +135,8 @@ def tracking_metrics(
     # Assume that negative cluster labels mean that the cluster was labeled as
     # invalid
     unique_predicted, predicted_counts = np.unique(predicted, return_counts=True)
+    # Cluster mask: all clusters that are not labeled as noise and have minimum number
+    # of hits.
     c_valid_cluster = (unique_predicted >= 0) & (
         predicted_counts >= predicted_count_thld
     )
@@ -177,15 +179,15 @@ def tracking_metrics(
 
         h_pt_mask = pts >= pt
         c_pt_mask = c_maj_pts >= pt
-        n_particles = len(np.unique(truth[h_pt_mask]))
-        n_clusters = len(unique_predicted[c_pt_mask & c_valid_cluster])
+        n_particles = len(np.unique(truth[h_pt_mask & reconstructable.astype(bool)]))
+        n_clusters = len(
+            unique_predicted[c_pt_mask & c_valid_cluster & c_maj_reconstructable]
+        )
 
         fake_pm = n_clusters - perfect_match
         fake_dm = n_clusters - double_majority
         fake_lhc = n_clusters - lhc_match
 
-        # breakpoint()
-
         r: TrackingMetrics = {
             "n_particles": n_particles,
             "n_cleaned_clusters": n_clusters,

diff --git a/src/gnn_tracking/metrics/test_cluster_metrics.py b/src/gnn_tracking/metrics/test_cluster_metrics.py
@@ -240,6 +240,7 @@ def key_reduction(dct, keys):
     # masking some particles
     ClusterMetricTestCase(
         # fmt: off
+        # particles: 0 (pt masked), 1 (reco masked), 2, 3, 4 (pt masked), 5 ==> Total 3
         truth=[
             0, 0, 0, 0, 0, 0,  # lhc, dm  (pt-masked)
             1, 1, 1, 1, 1, 5,  # lhc, dm (reco-masked)
@@ -262,25 +263,25 @@ def key_reduction(dct, keys):
             True, True, True, True, True, True,
             False, False, False, False, False, True,
             True, False, False, True,
-            True, False, True, True,
+            True, False, False, True,
             True, True,
             True,
         ],
         predicted=[
             0, 0, 0, 0, 0, 0,  # (pt-masked)
-            1, 1, 1, 1, 1, 1,
-            2, 2, 2, 2,
-            3, 3, 3, 3,
+            1, 1, 1, 1, 1, 1,  # (reco-masked)
+            2, 2, 2, 2,  # (reco-masked)
+            3, 3, 3, 3,  # (reco-masked)
             4, 4,  # (pt-masked)
             5,
         ],
         # fmt: on
         pt_thld=0.5,
-        n_particles=4,
-        n_cleaned_clusters=4,
-        perfect=0 / 3,
-        lhc=1 / 4,
-        double_majority=0 / 3,
+        n_particles=3,
+        n_cleaned_clusters=1,
+        perfect=0 / 2,
+        lhc=1 / 1,
+        double_majority=0 / 1,
     ),
 ]