From 8ffbfafb1d4e40c06647132991f81066f0a2bf4d Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Fri, 12 Jun 2026 11:45:16 +0200
Subject: [PATCH 01/10] add project configuration files and update invisible
 padding logic

---
 network/loss/assignment.py    | 2 +-
 network/metrics/assignment.py | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/network/loss/assignment.py b/network/loss/assignment.py
index e048595..b30e727 100644
--- a/network/loss/assignment.py
+++ b/network/loss/assignment.py
@@ -301,7 +301,7 @@ def loss_single_process(
     for symmetry_group in event_permutations:
         for symmetry_element in symmetry_group:
             symmetry_element = np.array(symmetry_element)
-            detection = detections[symmetry_element[0]]
+            detection = detections[symmetry_element[0]] # All the symmetry group detection/assignment outputs are just duplicated
             detection_target = torch.stack([detections_target[symmetry_index] for symmetry_index in symmetry_element])
             detection_target = detection_target.sum(0).long()
 
diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index 15b59c2..90406ed 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -159,6 +159,8 @@ def predict(assignments: List[Tensor],
             softmax = torch.nn.Softmax(dim=-1)
 
             detection_prob = softmax(detection_result)
+            # survival[:, r - 1] = P(N >= r)
+            survival = detection_prob[:, 1:].flip(-1).cumsum(-1).flip(-1)
 
             assignment_tmp = torch.stack([assignments_indices[element] for element in symmetry_element])
             assignment_probability_tmp = torch.stack(
@@ -170,13 +172,10 @@ def predict(assignments: List[Tensor],
             assignment_sorted = torch.gather(assignment_tmp, dim=0, index=expanded_sort_index)
             assignment_probability = torch.gather(assignment_probability_tmp, dim=0, index=sort_index)
 
-            init_probabilities = torch.ones_like(assignment_probability[0])
             for iorder in range(len(symmetry_element)):
                 final_assignments_indices.append(assignment_sorted[iorder])
                 final_assignments_probabilities.append(assignment_probability[iorder])
-                detections_probabilities = 1.0 - (detection_prob[:, iorder] / init_probabilities)
-                init_probabilities = detections_probabilities
-                final_detections_probabilities.append(detections_probabilities)
+                final_detections_probabilities.append(survival[:, iorder])
 
     return {
         "best_indices": final_assignments_indices,

From 51500a4edb6d9d45004158ef17b7ee05799f7dc5 Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Fri, 12 Jun 2026 12:05:36 +0200
Subject: [PATCH 02/10] refactor mask_jet function to handle valid indices and
 improve data assignment logic

---
 network/metrics/predict_assignment.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/network/metrics/predict_assignment.py b/network/metrics/predict_assignment.py
index 739bfba..1564e13 100644
--- a/network/metrics/predict_assignment.py
+++ b/network/metrics/predict_assignment.py
@@ -51,18 +51,21 @@ def maximal_prediction(predictions: List[torch.Tensor]) -> Tuple[int, int, float
 
 
 def mask_jet(data: torch.Tensor, num_partons: int, max_jets: int, index: torch.Tensor, value: float):
+    valid = index >= 0  # Ignore padded jet indices from lower-prong assignments.
     batch_size = data.shape[0]
+    batch_index = torch.arange(batch_size, device = data.device)[valid]
+    index = index[valid]
     if num_partons == 1:
-        data[torch.arange(batch_size), index] = value
+        data[batch_index, index] = value
     elif num_partons == 2:
         data = data.reshape((batch_size, max_jets, max_jets))
-        data[torch.arange(batch_size), index, :] = value
-        data[torch.arange(batch_size), :, index] = value
+        data[batch_index, index, :] = value
+        data[batch_index, :, index] = value
     elif num_partons == 3:
         data = data.reshape((batch_size, max_jets, max_jets, max_jets))
-        data[torch.arange(batch_size), index, :, :] = value
-        data[torch.arange(batch_size), :, index, :] = value
-        data[torch.arange(batch_size), :, :, index] = value
+        data[batch_index, index, :, :] = value
+        data[batch_index, :, index, :] = value
+        data[batch_index, :, :, index] = value
     else:
         raise NotImplementedError("num_partons > 3 not yet implemented in PyTorch version")
 
@@ -88,7 +91,7 @@ def extract_prediction(predictions: List[torch.Tensor], num_partons: torch.Tenso
 
         for i in range(num_targets):
             predictions[i] = predictions[i] + (torch.where(best_prediction == i, float_neg_inf, 0)).unsqueeze(1)
-            for i_parton in range(num_partons[i].item()):
+            for i_parton in range(max_partons):
                 jet = best_jets.reshape(batch_size, -1)[:, i_parton] # (batch_size,)
                 mask_jet(predictions[i], num_partons[i].item(), max_jets, jet, float_neg_inf)
 

From a2c8265abc863c19c305e4e80e9d92226af1680e Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Fri, 12 Jun 2026 12:46:36 +0200
Subject: [PATCH 03/10] refactor assignment logic to incorporate scaled
 assignments and improve selection process

---
 network/metrics/assignment.py | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index 90406ed..51c84c1 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -123,17 +123,45 @@ def predict(assignments: List[Tensor],
             product_symbolic_groups,
             event_permutations):
     device = assignments[0].device
+    scaled_assignments = list(assignments)
+    for symmetry_group in event_permutations:
+        for symmetry_element in symmetry_group:
+            symmetry_element = np.asarray(symmetry_element)
+            detection_prob = torch.softmax(
+                detections[symmetry_element[0]],
+                dim=-1,
+            )
+            survival = detection_prob[:, 1:].flip(-1).cumsum(-1).flip(-1)
+            expected_count = survival[:, :len(symmetry_element)].sum(-1)
+            expected_count = expected_count.clamp_min(1e-6)
+            for index in symmetry_element:
+                assignment = assignments[index]
+                broadcast_shape = (-1,) + (1,) * (assignment.ndim - 1)
+
+                scaled_assignments[index] = (
+                        assignment
+                        + expected_count.log().view(broadcast_shape)
+                )
+
+    selection_assignments = [
+        assignment + np.log(symmetries.order())
+        for assignment, symmetries in zip(
+            scaled_assignments,
+            product_symbolic_groups.values(),
+        )
+    ]
+
     assignments_indices = extract_predictions(
         [
             torch.nan_to_num(assignment, nan=-float('inf'))
-            for assignment in assignments
+            for assignment in selection_assignments
         ]
     )
-
     assignment_probabilities = []
     dummy_index = torch.arange(assignments_indices[0].shape[0])
+
     for assignment_probability, assignment, symmetries in zip(
-            assignments,
+            scaled_assignments,
             assignments_indices,
             product_symbolic_groups.values()
     ):
@@ -157,7 +185,6 @@ def predict(assignments: List[Tensor],
             symmetry_element = np.sort(np.array(symmetry_element))
             detection_result = detections[symmetry_element[0]]
             softmax = torch.nn.Softmax(dim=-1)
-
             detection_prob = softmax(detection_result)
             # survival[:, r - 1] = P(N >= r)
             survival = detection_prob[:, 1:].flip(-1).cumsum(-1).flip(-1)

From d7d24ddeb9932bae46830805475cf97d2328d153 Mon Sep 17 00:00:00 2001
From: Yulei Zhang <yulei.zhang@cern.ch>
Date: Mon, 15 Jun 2026 17:33:31 +0200
Subject: [PATCH 04/10] fix bug

---
 network/evenet_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/network/evenet_model.py b/network/evenet_model.py
index 2420d53..bf2417a 100644
--- a/network/evenet_model.py
+++ b/network/evenet_model.py
@@ -386,7 +386,7 @@ def forward(
         if 'x_invisible' in x and self.include_neutrino_generation:
             invisible_point_cloud = x['x_invisible']
         else:
-            invisible_point_cloud = torch.zeros(B, 1, self.invisible_input_dim, device=input_point_cloud.device)
+            invisible_point_cloud = torch.zeros(B, 1, self.sequential_input_dim, device=input_point_cloud.device)
 
         invisible_point_cloud_mask = x['x_invisible_mask'].unsqueeze(
             -1) if 'x_invisible_mask' in x else torch.zeros_like(input_point_cloud_mask[:, [0], :]).bool()

From af1a2a7b96f6d5dcfacaf9bfaebbe3840cb53eaa Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 16:33:07 +0200
Subject: [PATCH 05/10] add 2D score metrics and confusion detection plots for
 improved analysis visualization

---
 network/metrics/assignment.py | 250 ++++++++++++++++++++++++++++++++++
 1 file changed, 250 insertions(+)

diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index 51c84c1..bdb791b 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -21,6 +21,8 @@
 
 import wandb
 from matplotlib.lines import Line2D
+from matplotlib.patches import Patch
+from matplotlib.colors import LogNorm
 import logging
 
 logger = logging.getLogger(__name__)
@@ -281,6 +283,8 @@ def __init__(
 
         self.bins_score = np.linspace(0, 1, self.num_bins + 1)
         self.bin_centers_score = 0.5 * (self.bins_score[:-1] + self.bins_score[1:])
+        self.num_score2d_bins = 50
+        self.bins_score2d = np.linspace(0, 1, self.num_score2d_bins + 1)
 
         self.truth_metrics = dict(
             {f"{i + 1}{cluster_name}": {
@@ -312,6 +316,28 @@ def __init__(
                 for i in range(len(particle_name))
             })
 
+        self.score2d_metrics = dict(
+            {f"{i + 1}{cluster_name}": {
+                "correct": np.zeros((self.num_score2d_bins, self.num_score2d_bins)),
+                "wrong": np.zeros((self.num_score2d_bins, self.num_score2d_bins)),
+            }
+                for cluster_name, particle_name, orbit in self.clusters
+                for i in range(len(particle_name))
+            })
+
+        self.detection_count_confusion = {
+            cluster_name: np.zeros((len(particle_name) + 1, len(particle_name) + 1))
+            for cluster_name, particle_name, orbit in self.clusters
+        }
+
+        self.detection_survival_metrics = {
+            cluster_name: {
+                truth_count: np.zeros((len(particle_name), self.num_bins))
+                for truth_count in range(1, len(particle_name) + 1)
+            }
+            for cluster_name, particle_name, orbit in self.clusters
+        }
+
         self.train_metrics_correct = None
         self.train_metrics_wrong = None
 
@@ -410,9 +436,25 @@ def update(
             )
             correct_reco = torch.stack([correct_assigned[iorbit] for iorbit in list(sorted(orbit))], dim=0)
 
+            predict_count_wp = (predict_detection > detection_cut).sum(dim=0).long()
+            count_bins = np.arange(len(names) + 2) - 0.5
+            confusion, _, _ = np.histogram2d(
+                truth_count.detach().cpu().numpy(),
+                predict_count_wp.detach().cpu().numpy(),
+                bins=[count_bins, count_bins]
+            )
+            self.detection_count_confusion[cluster_name] += confusion
+
             for num_resonance in range(len(names)):
                 truth_mask = (truth_count == (num_resonance + 1))
                 hist_name = f"{num_resonance + 1}{cluster_name}"
+                for detection_order in range(len(names)):
+                    hist, _ = np.histogram(
+                        predict_detection[detection_order, truth_mask].detach().cpu().numpy(),
+                        bins=self.bins_score,
+                    )
+                    self.detection_survival_metrics[cluster_name][num_resonance + 1][detection_order] += hist
+
                 for local_resonance in range(len(names)):
                     truth_local = truth[local_resonance, :, :]
                     truth_mask_local = truth_mask
@@ -442,6 +484,13 @@ def update(
                     predict_correct = prediction_local[correct_local]
                     detection_correct = detection_local[correct_local]
                     assign_score_correct = assign_score_local[correct_local]
+                    if detection_correct.size()[0] > 0:
+                        hist2d, _, _ = np.histogram2d(
+                            assign_score_correct.detach().cpu().numpy(),
+                            detection_correct.detach().cpu().numpy(),
+                            bins=[self.bins_score2d, self.bins_score2d]
+                        )
+                        self.score2d_metrics[hist_name]["correct"] += hist2d
 
                     if prediction_local.size()[0] > 0:
                         reco_mass_correct = reconstruct_mass_peak(
@@ -462,6 +511,13 @@ def update(
                     prediction_false = prediction_local[~correct_local]
                     detection_false = detection_local[~correct_local]
                     assign_score_false = assign_score_local[~correct_local]
+                    if detection_false.size()[0] > 0:
+                        hist2d, _, _ = np.histogram2d(
+                            assign_score_false.detach().cpu().numpy(),
+                            detection_false.detach().cpu().numpy(),
+                            bins=[self.bins_score2d, self.bins_score2d]
+                        )
+                        self.score2d_metrics[hist_name]["wrong"] += hist2d
                     if (prediction_false.size()[0] > 0) and (prediction_false >= 0).all():
                         reco_mass_false = reconstruct_mass_peak(
                             jet[~correct_local], prediction_false, input_mask[~correct_local]
@@ -521,6 +577,17 @@ def reset(self):
             for key in hist.keys():
                 self.predict_metrics_wrong[name][key] = np.zeros(self.num_bins)
 
+        for name, hist in self.score2d_metrics.items():
+            for key in hist.keys():
+                self.score2d_metrics[name][key] = np.zeros((self.num_score2d_bins, self.num_score2d_bins))
+
+        for cluster_name, confusion in self.detection_count_confusion.items():
+            self.detection_count_confusion[cluster_name] = np.zeros_like(confusion)
+
+        for cluster_name, truth_count_metrics in self.detection_survival_metrics.items():
+            for truth_count, hist in truth_count_metrics.items():
+                self.detection_survival_metrics[cluster_name][truth_count] = np.zeros_like(hist)
+
         for name in self.full_log:
             for key in self.full_log[name].keys():
                 self.full_log[name][key] = 0
@@ -548,6 +615,23 @@ def reduce_across_gpus(self):
                     torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
                     self.predict_metrics_wrong[name][key] = tensor.cpu().numpy()
 
+            for name, hist in self.score2d_metrics.items():
+                for key in hist.keys():
+                    tensor = torch.tensor(hist[key], dtype=torch.long, device=self.device)
+                    torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
+                    self.score2d_metrics[name][key] = tensor.cpu().numpy()
+
+            for cluster_name, confusion in self.detection_count_confusion.items():
+                tensor = torch.tensor(confusion, dtype=torch.long, device=self.device)
+                torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
+                self.detection_count_confusion[cluster_name] = tensor.cpu().numpy()
+
+            for cluster_name, truth_count_metrics in self.detection_survival_metrics.items():
+                for truth_count, hist in truth_count_metrics.items():
+                    tensor = torch.tensor(hist, dtype=torch.long, device=self.device)
+                    torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
+                    self.detection_survival_metrics[cluster_name][truth_count] = tensor.cpu().numpy()
+
             for name, log in self.full_log.items():
                 for key in log.keys():
                     tensor = torch.tensor(log[key], dtype=torch.long, device=self.device)
@@ -783,6 +867,148 @@ def plot_score(self, target="detection_score"):
             )
         return return_plot
 
+    def plot_detection_count_confusion(self):
+        return_plot = dict()
+        for cluster_name, confusion in self.detection_count_confusion.items():
+            fig, ax = plt.subplots(figsize=(6, 5))
+            row_sum = confusion.sum(axis=1, keepdims=True)
+            normalized = np.divide(
+                confusion,
+                np.maximum(row_sum, 1),
+                out=np.zeros_like(confusion, dtype=float),
+                where=row_sum > 0,
+            )
+
+            im = ax.imshow(normalized, origin="lower", vmin=0, vmax=1, cmap="Blues")
+            fig.colorbar(im, ax=ax, label="Row-normalized density")
+
+            max_count = confusion.shape[0] - 1
+            ax.set_xticks(np.arange(max_count + 1))
+            ax.set_yticks(np.arange(max_count + 1))
+            ax.set_xlabel(f"Predicted N (WP: {self.detection_cut})")
+            ax.set_ylabel("Truth N")
+            ax.set_title(f"Detection Count Confusion: {cluster_name}")
+
+            for truth_count in range(max_count + 1):
+                for predicted_count in range(max_count + 1):
+                    count = int(confusion[truth_count, predicted_count])
+                    if count == 0:
+                        continue
+                    percent = normalized[truth_count, predicted_count]
+                    text_color = "white" if percent > 0.5 else "black"
+                    ax.text(
+                        predicted_count,
+                        truth_count,
+                        f"{count}\n{percent:.2f}",
+                        ha="center",
+                        va="center",
+                        fontsize=8,
+                        color=text_color,
+                    )
+
+            fig.tight_layout()
+            return_plot[cluster_name] = fig
+        return return_plot
+
+    def plot_score2d_density(self):
+        return_plot = dict()
+        for cluster_name, names, orbit in self.clusters:
+            num_panels = len(names)
+            fig, axes = plt.subplots(
+                1,
+                num_panels,
+                figsize=(4.5 * num_panels, 4),
+                sharex=True,
+                sharey=True,
+                squeeze=False,
+            )
+            axes = axes[0]
+
+            for index, ax in enumerate(axes):
+                hist_name = f"{index + 1}{cluster_name}"
+                correct = np.ma.masked_less_equal(self.score2d_metrics[hist_name]["correct"].T, 0)
+                wrong = np.ma.masked_less_equal(self.score2d_metrics[hist_name]["wrong"].T, 0)
+                max_density = max(
+                    self.score2d_metrics[hist_name]["correct"].max(),
+                    self.score2d_metrics[hist_name]["wrong"].max(),
+                )
+                norm = LogNorm(vmin=1, vmax=max(2, max_density))
+
+                ax.imshow(
+                    wrong,
+                    extent=[0, 1, 0, 1],
+                    origin="lower",
+                    aspect="auto",
+                    cmap="Oranges",
+                    norm=norm,
+                    alpha=0.65,
+                )
+                ax.imshow(
+                    correct,
+                    extent=[0, 1, 0, 1],
+                    origin="lower",
+                    aspect="auto",
+                    cmap="Blues",
+                    norm=norm,
+                    alpha=0.65,
+                )
+                ax.axhline(self.detection_cut, color="black", linestyle="--", linewidth=1)
+                ax.set_title(f"Truth N = {index + 1}")
+                ax.set_xlabel("Assignment score")
+                ax.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
+
+            axes[0].set_ylabel("Detection score")
+            legend_handles = [
+                Patch(facecolor=plt.cm.Blues(0.75), alpha=0.65, label="Correct assign"),
+                Patch(facecolor=plt.cm.Oranges(0.75), alpha=0.65, label="Wrong assign"),
+                Line2D([0], [0], color="black", linestyle="--", linewidth=1,
+                       label=f"Detection WP: {self.detection_cut}"),
+            ]
+            axes[-1].legend(handles=legend_handles, loc="lower right")
+            fig.suptitle(f"Assignment Score vs Detection Score: {cluster_name}")
+            fig.tight_layout()
+            return_plot[cluster_name] = fig
+        return return_plot
+
+    def plot_detection_survival_distribution(self):
+        return_plot = dict()
+        for cluster_name, names, orbit in self.clusters:
+            num_panels = len(names)
+            fig, axes = plt.subplots(
+                1,
+                num_panels,
+                figsize=(4.5 * num_panels, 4),
+                sharex=True,
+                sharey=True,
+                squeeze=False,
+            )
+            axes = axes[0]
+            bin_widths = np.diff(self.bins_score)
+
+            for truth_count, ax in enumerate(axes, start=1):
+                hist_by_order = self.detection_survival_metrics[cluster_name][truth_count]
+                for detection_order, hist in enumerate(hist_by_order, start=1):
+                    density = hist / np.maximum(1.0, hist.sum() * bin_widths)
+                    ax.step(
+                        self.bin_centers_score,
+                        density,
+                        where="mid",
+                        linewidth=1.8,
+                        label=f"P(N >= {detection_order})",
+                    )
+
+                ax.axvline(self.detection_cut, color="black", linestyle="--", linewidth=1)
+                ax.set_title(f"Truth N = {truth_count}")
+                ax.set_xlabel("Probability")
+                ax.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
+
+            axes[0].set_ylabel("Density")
+            axes[-1].legend(loc="best")
+            fig.suptitle(f"Detection Survival Probability: {cluster_name}")
+            fig.tight_layout()
+            return_plot[cluster_name] = fig
+        return return_plot
+
     def summary_log(self):
         return_log = dict()
 
@@ -1023,6 +1249,30 @@ def shared_epoch_end(
             for _, fig in figs.items():
                 plt.close(fig)
 
+            figs = metrics_valid[process].plot_detection_count_confusion()
+            wandb.log({
+                f"assignment_detection_count/{process}/{name}": wandb.Image(fig)
+                for name, fig in figs.items()
+            })
+            for _, fig in figs.items():
+                plt.close(fig)
+
+            figs = metrics_valid[process].plot_detection_survival_distribution()
+            wandb.log({
+                f"assignment_detection_survival/{process}/{name}": wandb.Image(fig)
+                for name, fig in figs.items()
+            })
+            for _, fig in figs.items():
+                plt.close(fig)
+
+            figs = metrics_valid[process].plot_score2d_density()
+            wandb.log({
+                f"assignment_score2d/{process}/{name}": wandb.Image(fig)
+                for name, fig in figs.items()
+            })
+            for _, fig in figs.items():
+                plt.close(fig)
+
     for _, metric in metrics_valid.items():
         metric.reset()
 

From b221e1a6f036b7bcdf974a1b99aa892d16dc5e14 Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 16:41:38 +0200
Subject: [PATCH 06/10] refactor score metrics and confusion detection to use
 int64 for improved accuracy

---
 network/metrics/assignment.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index bdb791b..cfec50c 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -318,21 +318,21 @@ def __init__(
 
         self.score2d_metrics = dict(
             {f"{i + 1}{cluster_name}": {
-                "correct": np.zeros((self.num_score2d_bins, self.num_score2d_bins)),
-                "wrong": np.zeros((self.num_score2d_bins, self.num_score2d_bins)),
+                "correct": np.zeros((self.num_score2d_bins, self.num_score2d_bins), dtype=np.int64),
+                "wrong": np.zeros((self.num_score2d_bins, self.num_score2d_bins), dtype=np.int64),
             }
                 for cluster_name, particle_name, orbit in self.clusters
                 for i in range(len(particle_name))
             })
 
         self.detection_count_confusion = {
-            cluster_name: np.zeros((len(particle_name) + 1, len(particle_name) + 1))
+            cluster_name: np.zeros((len(particle_name) + 1, len(particle_name) + 1), dtype=np.int64)
             for cluster_name, particle_name, orbit in self.clusters
         }
 
         self.detection_survival_metrics = {
             cluster_name: {
-                truth_count: np.zeros((len(particle_name), self.num_bins))
+                truth_count: np.zeros((len(particle_name), self.num_bins), dtype=np.int64)
                 for truth_count in range(1, len(particle_name) + 1)
             }
             for cluster_name, particle_name, orbit in self.clusters
@@ -443,7 +443,7 @@ def update(
                 predict_count_wp.detach().cpu().numpy(),
                 bins=[count_bins, count_bins]
             )
-            self.detection_count_confusion[cluster_name] += confusion
+            self.detection_count_confusion[cluster_name] += confusion.astype(np.int64, copy=False)
 
             for num_resonance in range(len(names)):
                 truth_mask = (truth_count == (num_resonance + 1))
@@ -490,7 +490,7 @@ def update(
                             detection_correct.detach().cpu().numpy(),
                             bins=[self.bins_score2d, self.bins_score2d]
                         )
-                        self.score2d_metrics[hist_name]["correct"] += hist2d
+                        self.score2d_metrics[hist_name]["correct"] += hist2d.astype(np.int64, copy=False)
 
                     if prediction_local.size()[0] > 0:
                         reco_mass_correct = reconstruct_mass_peak(
@@ -517,7 +517,7 @@ def update(
                             detection_false.detach().cpu().numpy(),
                             bins=[self.bins_score2d, self.bins_score2d]
                         )
-                        self.score2d_metrics[hist_name]["wrong"] += hist2d
+                        self.score2d_metrics[hist_name]["wrong"] += hist2d.astype(np.int64, copy=False)
                     if (prediction_false.size()[0] > 0) and (prediction_false >= 0).all():
                         reco_mass_false = reconstruct_mass_peak(
                             jet[~correct_local], prediction_false, input_mask[~correct_local]

From 9965db725b82469ea2bd6927470f68d2d5cf3bdc Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 16:54:16 +0200
Subject: [PATCH 07/10] refactor score metrics and confusion detection to use
 int64 for improved accuracy

---
 network/metrics/assignment.py | 214 +++++++++++-----------------------
 1 file changed, 71 insertions(+), 143 deletions(-)

diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index cfec50c..4cf4452 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -21,8 +21,6 @@
 
 import wandb
 from matplotlib.lines import Line2D
-from matplotlib.patches import Patch
-from matplotlib.colors import LogNorm
 import logging
 
 logger = logging.getLogger(__name__)
@@ -283,8 +281,6 @@ def __init__(
 
         self.bins_score = np.linspace(0, 1, self.num_bins + 1)
         self.bin_centers_score = 0.5 * (self.bins_score[:-1] + self.bins_score[1:])
-        self.num_score2d_bins = 50
-        self.bins_score2d = np.linspace(0, 1, self.num_score2d_bins + 1)
 
         self.truth_metrics = dict(
             {f"{i + 1}{cluster_name}": {
@@ -316,15 +312,6 @@ def __init__(
                 for i in range(len(particle_name))
             })
 
-        self.score2d_metrics = dict(
-            {f"{i + 1}{cluster_name}": {
-                "correct": np.zeros((self.num_score2d_bins, self.num_score2d_bins), dtype=np.int64),
-                "wrong": np.zeros((self.num_score2d_bins, self.num_score2d_bins), dtype=np.int64),
-            }
-                for cluster_name, particle_name, orbit in self.clusters
-                for i in range(len(particle_name))
-            })
-
         self.detection_count_confusion = {
             cluster_name: np.zeros((len(particle_name) + 1, len(particle_name) + 1), dtype=np.int64)
             for cluster_name, particle_name, orbit in self.clusters
@@ -484,13 +471,6 @@ def update(
                     predict_correct = prediction_local[correct_local]
                     detection_correct = detection_local[correct_local]
                     assign_score_correct = assign_score_local[correct_local]
-                    if detection_correct.size()[0] > 0:
-                        hist2d, _, _ = np.histogram2d(
-                            assign_score_correct.detach().cpu().numpy(),
-                            detection_correct.detach().cpu().numpy(),
-                            bins=[self.bins_score2d, self.bins_score2d]
-                        )
-                        self.score2d_metrics[hist_name]["correct"] += hist2d.astype(np.int64, copy=False)
 
                     if prediction_local.size()[0] > 0:
                         reco_mass_correct = reconstruct_mass_peak(
@@ -511,13 +491,6 @@ def update(
                     prediction_false = prediction_local[~correct_local]
                     detection_false = detection_local[~correct_local]
                     assign_score_false = assign_score_local[~correct_local]
-                    if detection_false.size()[0] > 0:
-                        hist2d, _, _ = np.histogram2d(
-                            assign_score_false.detach().cpu().numpy(),
-                            detection_false.detach().cpu().numpy(),
-                            bins=[self.bins_score2d, self.bins_score2d]
-                        )
-                        self.score2d_metrics[hist_name]["wrong"] += hist2d.astype(np.int64, copy=False)
                     if (prediction_false.size()[0] > 0) and (prediction_false >= 0).all():
                         reco_mass_false = reconstruct_mass_peak(
                             jet[~correct_local], prediction_false, input_mask[~correct_local]
@@ -577,10 +550,6 @@ def reset(self):
             for key in hist.keys():
                 self.predict_metrics_wrong[name][key] = np.zeros(self.num_bins)
 
-        for name, hist in self.score2d_metrics.items():
-            for key in hist.keys():
-                self.score2d_metrics[name][key] = np.zeros((self.num_score2d_bins, self.num_score2d_bins))
-
         for cluster_name, confusion in self.detection_count_confusion.items():
             self.detection_count_confusion[cluster_name] = np.zeros_like(confusion)
 
@@ -615,12 +584,6 @@ def reduce_across_gpus(self):
                     torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
                     self.predict_metrics_wrong[name][key] = tensor.cpu().numpy()
 
-            for name, hist in self.score2d_metrics.items():
-                for key in hist.keys():
-                    tensor = torch.tensor(hist[key], dtype=torch.long, device=self.device)
-                    torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
-                    self.score2d_metrics[name][key] = tensor.cpu().numpy()
-
             for cluster_name, confusion in self.detection_count_confusion.items():
                 tensor = torch.tensor(confusion, dtype=torch.long, device=self.device)
                 torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
@@ -910,101 +873,81 @@ def plot_detection_count_confusion(self):
             return_plot[cluster_name] = fig
         return return_plot
 
-    def plot_score2d_density(self):
-        return_plot = dict()
-        for cluster_name, names, orbit in self.clusters:
-            num_panels = len(names)
-            fig, axes = plt.subplots(
-                1,
-                num_panels,
-                figsize=(4.5 * num_panels, 4),
-                sharex=True,
-                sharey=True,
-                squeeze=False,
-            )
-            axes = axes[0]
-
-            for index, ax in enumerate(axes):
-                hist_name = f"{index + 1}{cluster_name}"
-                correct = np.ma.masked_less_equal(self.score2d_metrics[hist_name]["correct"].T, 0)
-                wrong = np.ma.masked_less_equal(self.score2d_metrics[hist_name]["wrong"].T, 0)
-                max_density = max(
-                    self.score2d_metrics[hist_name]["correct"].max(),
-                    self.score2d_metrics[hist_name]["wrong"].max(),
-                )
-                norm = LogNorm(vmin=1, vmax=max(2, max_density))
-
-                ax.imshow(
-                    wrong,
-                    extent=[0, 1, 0, 1],
-                    origin="lower",
-                    aspect="auto",
-                    cmap="Oranges",
-                    norm=norm,
-                    alpha=0.65,
-                )
-                ax.imshow(
-                    correct,
-                    extent=[0, 1, 0, 1],
-                    origin="lower",
-                    aspect="auto",
-                    cmap="Blues",
-                    norm=norm,
-                    alpha=0.65,
-                )
-                ax.axhline(self.detection_cut, color="black", linestyle="--", linewidth=1)
-                ax.set_title(f"Truth N = {index + 1}")
-                ax.set_xlabel("Assignment score")
-                ax.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
-
-            axes[0].set_ylabel("Detection score")
-            legend_handles = [
-                Patch(facecolor=plt.cm.Blues(0.75), alpha=0.65, label="Correct assign"),
-                Patch(facecolor=plt.cm.Oranges(0.75), alpha=0.65, label="Wrong assign"),
-                Line2D([0], [0], color="black", linestyle="--", linewidth=1,
-                       label=f"Detection WP: {self.detection_cut}"),
-            ]
-            axes[-1].legend(handles=legend_handles, loc="lower right")
-            fig.suptitle(f"Assignment Score vs Detection Score: {cluster_name}")
-            fig.tight_layout()
-            return_plot[cluster_name] = fig
-        return return_plot
+    def detection_count_summary_log(self):
+        return_log = dict()
+        for cluster_name, confusion in self.detection_count_confusion.items():
+            row_sum = confusion.sum(axis=1)
+            max_count = confusion.shape[0] - 1
+            for truth_count in range(max_count + 1):
+                denominator = row_sum[truth_count]
+                if denominator <= 0:
+                    continue
+
+                exact_rate = confusion[truth_count, truth_count] / denominator
+                under_rate = confusion[truth_count, :truth_count].sum() / denominator
+                over_rate = confusion[truth_count, truth_count + 1:].sum() / denominator
+                return_log[
+                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/exact_wp"
+                ] = exact_rate
+                return_log[
+                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/under_wp"
+                ] = under_rate
+                return_log[
+                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/over_wp"
+                ] = over_rate
 
-    def plot_detection_survival_distribution(self):
+                for predicted_count in range(max_count + 1):
+                    return_log[
+                        f"detection/{self.process}/{cluster_name}/truth_{truth_count}/pred_{predicted_count}_wp"
+                    ] = confusion[truth_count, predicted_count] / denominator
+        return return_log
+
+    def plot_detection_survival_summary(self):
         return_plot = dict()
         for cluster_name, names, orbit in self.clusters:
-            num_panels = len(names)
-            fig, axes = plt.subplots(
-                1,
-                num_panels,
-                figsize=(4.5 * num_panels, 4),
-                sharex=True,
-                sharey=True,
-                squeeze=False,
-            )
-            axes = axes[0]
-            bin_widths = np.diff(self.bins_score)
-
-            for truth_count, ax in enumerate(axes, start=1):
+            max_count = len(names)
+            means = np.zeros((max_count, max_count))
+            wp_rates = np.zeros((max_count, max_count))
+            for truth_count in range(1, max_count + 1):
                 hist_by_order = self.detection_survival_metrics[cluster_name][truth_count]
                 for detection_order, hist in enumerate(hist_by_order, start=1):
-                    density = hist / np.maximum(1.0, hist.sum() * bin_widths)
-                    ax.step(
-                        self.bin_centers_score,
-                        density,
-                        where="mid",
-                        linewidth=1.8,
-                        label=f"P(N >= {detection_order})",
+                    total = hist.sum()
+                    if total <= 0:
+                        continue
+                    means[truth_count - 1, detection_order - 1] = (
+                        hist * self.bin_centers_score
+                    ).sum() / total
+                    wp_rates[truth_count - 1, detection_order - 1] = (
+                        hist[self.bin_centers_score > self.detection_cut].sum() / total
                     )
 
-                ax.axvline(self.detection_cut, color="black", linestyle="--", linewidth=1)
-                ax.set_title(f"Truth N = {truth_count}")
-                ax.set_xlabel("Probability")
-                ax.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
+            fig, ax = plt.subplots(figsize=(1.6 * max_count + 3.2, 1.2 * max_count + 2.6))
+            im = ax.imshow(means, origin="lower", vmin=0, vmax=1, cmap="viridis")
+            fig.colorbar(im, ax=ax, label="Mean survival probability")
+
+            ax.set_xticks(np.arange(max_count))
+            ax.set_yticks(np.arange(max_count))
+            ax.set_xticklabels([f"P(N>={i})" for i in range(1, max_count + 1)])
+            ax.set_yticklabels([f"Truth {i}{cluster_name}" for i in range(1, max_count + 1)])
+            ax.set_xlabel("Detection survival output")
+            ax.set_ylabel("Truth count")
+            ax.set_title(f"Detection Survival Summary: {cluster_name} (WP: {self.detection_cut})")
+
+            for truth_index in range(max_count):
+                for detection_index in range(max_count):
+                    mean_value = means[truth_index, detection_index]
+                    wp_value = wp_rates[truth_index, detection_index]
+                    text_color = "white" if mean_value < 0.35 else "black"
+                    ax.text(
+                        detection_index,
+                        truth_index,
+                        f"{mean_value:.2f}\nWP {wp_value:.2f}",
+                        ha="center",
+                        va="center",
+                        fontsize=8,
+                        color=text_color,
+                    )
 
-            axes[0].set_ylabel("Density")
-            axes[-1].legend(loc="best")
-            fig.suptitle(f"Detection Survival Probability: {cluster_name}")
             fig.tight_layout()
             return_plot[cluster_name] = fig
         return return_plot
@@ -1206,6 +1149,7 @@ def shared_epoch_end(
 
             logs = metrics_valid[process].summary_log()
             logger.log(logs)
+            logger.log(metrics_valid[process].detection_count_summary_log())
 
             # if metrics_train[process] is not None:
             #     training_logs = metrics_train[process].summary_log()
@@ -1233,14 +1177,6 @@ def shared_epoch_end(
             for _, fig in figs.items():
                 plt.close(fig)
 
-            figs = metrics_valid[process].plot_score(target="detection_score")
-            wandb.log({
-                f"assignment_reco_detection/{process}/{name}": wandb.Image(fig)
-                for name, fig in figs.items()
-            })
-            for _, fig in figs.items():
-                plt.close(fig)
-
             figs = metrics_valid[process].plot_score(target="assignment_score")
             wandb.log({
                 f"assignment_score/{process}/{name}": wandb.Image(fig)
@@ -1251,23 +1187,15 @@ def shared_epoch_end(
 
             figs = metrics_valid[process].plot_detection_count_confusion()
             wandb.log({
-                f"assignment_detection_count/{process}/{name}": wandb.Image(fig)
-                for name, fig in figs.items()
-            })
-            for _, fig in figs.items():
-                plt.close(fig)
-
-            figs = metrics_valid[process].plot_detection_survival_distribution()
-            wandb.log({
-                f"assignment_detection_survival/{process}/{name}": wandb.Image(fig)
+                f"detection/{process}/{name}/count_confusion": wandb.Image(fig)
                 for name, fig in figs.items()
             })
             for _, fig in figs.items():
                 plt.close(fig)
 
-            figs = metrics_valid[process].plot_score2d_density()
+            figs = metrics_valid[process].plot_detection_survival_summary()
             wandb.log({
-                f"assignment_score2d/{process}/{name}": wandb.Image(fig)
+                f"detection/{process}/{name}/survival_mean": wandb.Image(fig)
                 for name, fig in figs.items()
             })
             for _, fig in figs.items():

From 8e068b4142fc95ee00219565b57a9c970ab8e5fd Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 17:04:08 +0200
Subject: [PATCH 08/10] refactor detection metrics and visualization to enhance
 summary and plotting functionality

---
 network/metrics/assignment.py | 105 ++++++++++++++++++++--------------
 1 file changed, 61 insertions(+), 44 deletions(-)

diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index 4cf4452..20df86b 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -423,11 +423,19 @@ def update(
             )
             correct_reco = torch.stack([correct_assigned[iorbit] for iorbit in list(sorted(orbit))], dim=0)
 
-            predict_count_wp = (predict_detection > detection_cut).sum(dim=0).long()
+            count_probabilities = torch.cat(
+                [
+                    1 - predict_detection[0:1],
+                    predict_detection[:-1] - predict_detection[1:],
+                    predict_detection[-1:],
+                ],
+                dim=0,
+            ).clamp_min(0)
+            predict_count = torch.argmax(count_probabilities, dim=0).long()
             count_bins = np.arange(len(names) + 2) - 0.5
             confusion, _, _ = np.histogram2d(
                 truth_count.detach().cpu().numpy(),
-                predict_count_wp.detach().cpu().numpy(),
+                predict_count.detach().cpu().numpy(),
                 bins=[count_bins, count_bins]
             )
             self.detection_count_confusion[cluster_name] += confusion.astype(np.int64, copy=False)
@@ -830,45 +838,71 @@ def plot_score(self, target="detection_score"):
             )
         return return_plot
 
-    def plot_detection_count_confusion(self):
+    def plot_detection_count_distribution(self):
         return_plot = dict()
         for cluster_name, confusion in self.detection_count_confusion.items():
-            fig, ax = plt.subplots(figsize=(6, 5))
+            max_count = confusion.shape[0] - 1
+            predicted_counts = np.arange(max_count + 1)
             row_sum = confusion.sum(axis=1, keepdims=True)
-            normalized = np.divide(
+            pmf = np.divide(
                 confusion,
                 np.maximum(row_sum, 1),
                 out=np.zeros_like(confusion, dtype=float),
                 where=row_sum > 0,
             )
 
-            im = ax.imshow(normalized, origin="lower", vmin=0, vmax=1, cmap="Blues")
-            fig.colorbar(im, ax=ax, label="Row-normalized density")
+            fig, (ax_matrix, ax_pdf) = plt.subplots(
+                2,
+                1,
+                figsize=(7.5, 8),
+                gridspec_kw={"height_ratios": [1.15, 1]},
+            )
 
-            max_count = confusion.shape[0] - 1
-            ax.set_xticks(np.arange(max_count + 1))
-            ax.set_yticks(np.arange(max_count + 1))
-            ax.set_xlabel(f"Predicted N (WP: {self.detection_cut})")
-            ax.set_ylabel("Truth N")
-            ax.set_title(f"Detection Count Confusion: {cluster_name}")
+            im = ax_matrix.imshow(pmf, origin="lower", vmin=0, vmax=1, cmap="Blues")
+            fig.colorbar(im, ax=ax_matrix, label="P(predicted N | truth N)")
+            ax_matrix.set_xticks(predicted_counts)
+            ax_matrix.set_yticks(predicted_counts)
+            ax_matrix.set_xlabel("Predicted N (argmax)")
+            ax_matrix.set_ylabel("Truth N")
+            ax_matrix.set_title(f"Detection Count Confusion: {cluster_name}")
 
             for truth_count in range(max_count + 1):
                 for predicted_count in range(max_count + 1):
-                    count = int(confusion[truth_count, predicted_count])
-                    if count == 0:
+                    probability = pmf[truth_count, predicted_count]
+                    if probability <= 0:
                         continue
-                    percent = normalized[truth_count, predicted_count]
-                    text_color = "white" if percent > 0.5 else "black"
-                    ax.text(
+                    text_color = "white" if probability > 0.5 else "black"
+                    ax_matrix.text(
                         predicted_count,
                         truth_count,
-                        f"{count}\n{percent:.2f}",
+                        f"{probability:.2f}",
                         ha="center",
                         va="center",
                         fontsize=8,
                         color=text_color,
                     )
 
+            for truth_count in range(max_count + 1):
+                total = row_sum[truth_count, 0]
+                if total <= 0:
+                    continue
+                expected_count = (pmf[truth_count] * predicted_counts).sum()
+                ax_pdf.step(
+                    predicted_counts,
+                    pmf[truth_count],
+                    where="mid",
+                    linewidth=2,
+                    marker="o",
+                    label=f"Truth {truth_count}{cluster_name}: E[N]={expected_count:.2f}",
+                )
+
+            ax_pdf.set_xticks(predicted_counts)
+            ax_pdf.set_ylim(0, 1)
+            ax_pdf.set_xlabel("Predicted N (argmax)")
+            ax_pdf.set_ylabel("P(predicted N | truth N)")
+            ax_pdf.set_title("Count Distribution by Truth")
+            ax_pdf.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
+            ax_pdf.legend(loc="best", fontsize=8)
             fig.tight_layout()
             return_plot[cluster_name] = fig
         return return_plot
@@ -878,28 +912,16 @@ def detection_count_summary_log(self):
         for cluster_name, confusion in self.detection_count_confusion.items():
             row_sum = confusion.sum(axis=1)
             max_count = confusion.shape[0] - 1
+            predicted_counts = np.arange(max_count + 1)
             for truth_count in range(max_count + 1):
                 denominator = row_sum[truth_count]
                 if denominator <= 0:
                     continue
 
-                exact_rate = confusion[truth_count, truth_count] / denominator
-                under_rate = confusion[truth_count, :truth_count].sum() / denominator
-                over_rate = confusion[truth_count, truth_count + 1:].sum() / denominator
-                return_log[
-                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/exact_wp"
-                ] = exact_rate
-                return_log[
-                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/under_wp"
-                ] = under_rate
+                predicted_mean = (confusion[truth_count] * predicted_counts).sum() / denominator
                 return_log[
-                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/over_wp"
-                ] = over_rate
-
-                for predicted_count in range(max_count + 1):
-                    return_log[
-                        f"detection/{self.process}/{cluster_name}/truth_{truth_count}/pred_{predicted_count}_wp"
-                    ] = confusion[truth_count, predicted_count] / denominator
+                    f"detection/{self.process}/{cluster_name}/truth_{truth_count}/expected_predicted_n"
+                ] = predicted_mean
         return return_log
 
     def plot_detection_survival_summary(self):
@@ -907,7 +929,6 @@ def plot_detection_survival_summary(self):
         for cluster_name, names, orbit in self.clusters:
             max_count = len(names)
             means = np.zeros((max_count, max_count))
-            wp_rates = np.zeros((max_count, max_count))
             for truth_count in range(1, max_count + 1):
                 hist_by_order = self.detection_survival_metrics[cluster_name][truth_count]
                 for detection_order, hist in enumerate(hist_by_order, start=1):
@@ -917,9 +938,6 @@ def plot_detection_survival_summary(self):
                     means[truth_count - 1, detection_order - 1] = (
                         hist * self.bin_centers_score
                     ).sum() / total
-                    wp_rates[truth_count - 1, detection_order - 1] = (
-                        hist[self.bin_centers_score > self.detection_cut].sum() / total
-                    )
 
             fig, ax = plt.subplots(figsize=(1.6 * max_count + 3.2, 1.2 * max_count + 2.6))
             im = ax.imshow(means, origin="lower", vmin=0, vmax=1, cmap="viridis")
@@ -931,17 +949,16 @@ def plot_detection_survival_summary(self):
             ax.set_yticklabels([f"Truth {i}{cluster_name}" for i in range(1, max_count + 1)])
             ax.set_xlabel("Detection survival output")
             ax.set_ylabel("Truth count")
-            ax.set_title(f"Detection Survival Summary: {cluster_name} (WP: {self.detection_cut})")
+            ax.set_title(f"Detection Survival Summary: {cluster_name}")
 
             for truth_index in range(max_count):
                 for detection_index in range(max_count):
                     mean_value = means[truth_index, detection_index]
-                    wp_value = wp_rates[truth_index, detection_index]
                     text_color = "white" if mean_value < 0.35 else "black"
                     ax.text(
                         detection_index,
                         truth_index,
-                        f"{mean_value:.2f}\nWP {wp_value:.2f}",
+                        f"{mean_value:.2f}",
                         ha="center",
                         va="center",
                         fontsize=8,
@@ -1185,9 +1202,9 @@ def shared_epoch_end(
             for _, fig in figs.items():
                 plt.close(fig)
 
-            figs = metrics_valid[process].plot_detection_count_confusion()
+            figs = metrics_valid[process].plot_detection_count_distribution()
             wandb.log({
-                f"detection/{process}/{name}/count_confusion": wandb.Image(fig)
+                f"detection/{process}/{name}/count_distribution": wandb.Image(fig)
                 for name, fig in figs.items()
             })
             for _, fig in figs.items():

From 01221e9914d76fa430befb544a11c3ec448044c8 Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 17:19:09 +0200
Subject: [PATCH 09/10] refactor detection metrics and visualization to enhance
 summary and plotting functionality

---
 network/loss/assignment.py    | 26 ++++++++++++---------
 network/metrics/assignment.py | 43 +++++++++++++++++++++++++++++------
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/network/loss/assignment.py b/network/loss/assignment.py
index b30e727..d6003be 100644
--- a/network/loss/assignment.py
+++ b/network/loss/assignment.py
@@ -293,6 +293,13 @@ def loss_single_process(
     ## Detection Loss ##
     ####################
 
+    particle_balance_weight = torch.ones_like(targets_mask[0], dtype=torch.float32)
+    masks_for_balance = torch.stack(targets_mask).int()
+
+    if particle_index_tensor is not None:
+        class_indices = (masks_for_balance * particle_index_tensor.to(masks_for_balance.device).unsqueeze(1)).sum(0).int()
+        particle_balance_weight *= particle_weights_tensor.to(masks_for_balance.device)[class_indices]
+
     detections = detections
     detections_target = targets_mask
     detection_losses = []
@@ -322,13 +329,15 @@ def loss_single_process(
 
     process_masking = torch.stack(process_masking).float()
     process_weighting = torch.stack(process_weighting).float()
+    particle_balance_weight = particle_balance_weight.unsqueeze(0)
 
     if event_weight is not None:
-        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * event_weight.view(-1, *([1] * (process_masking.dim() - 1)))
-        valid_process = torch.sum(process_masking * process_weighting * event_weight.view(-1, *([1] * (process_masking.dim() - 1))))
+        event_balance_weight = event_weight.view(-1, *([1] * (process_masking.dim() - 1)))
+        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * event_balance_weight * particle_balance_weight
+        valid_process = torch.sum(process_masking * process_weighting * event_balance_weight * particle_balance_weight)
     else:
-        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting
-        valid_process = torch.sum(process_masking * process_weighting)
+        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * particle_balance_weight
+        valid_process = torch.sum(process_masking * process_weighting * particle_balance_weight)
 
     if valid_process > 0:
         detection_loss = torch.sum(detection_losses) / valid_process
@@ -351,15 +360,10 @@ def loss_single_process(
         focal_gamma
     )
 
-    particle_balance_weight = torch.ones_like(symmetric_losses)
-    masks_for_balance = torch.stack(targets_mask).int()
-
-    if particle_balance_weight is not None and particle_index_tensor is not None:
-        class_indices = (masks_for_balance * particle_index_tensor.to(masks_for_balance.device).unsqueeze(1)).sum(0).int()
-        particle_balance_weight *= particle_weights_tensor.to(masks_for_balance.device)[class_indices]
+    particle_balance_weight = particle_balance_weight.squeeze(0).to(symmetric_losses.device)
 
     if process_weight[0] is not None:
-        particle_balance_weight *= process_weight[0].unsqueeze(0)
+        particle_balance_weight *= process_weight[0]
 
     targets_mask_finite = torch.stack(targets_mask).float()
     if not torch.isfinite(symmetric_losses).all():
diff --git a/network/metrics/assignment.py b/network/metrics/assignment.py
index 20df86b..abc53a5 100644
--- a/network/metrics/assignment.py
+++ b/network/metrics/assignment.py
@@ -317,6 +317,11 @@ def __init__(
             for cluster_name, particle_name, orbit in self.clusters
         }
 
+        self.detection_count_probability_sum = {
+            cluster_name: np.zeros((len(particle_name) + 1, len(particle_name) + 1), dtype=np.float64)
+            for cluster_name, particle_name, orbit in self.clusters
+        }
+
         self.detection_survival_metrics = {
             cluster_name: {
                 truth_count: np.zeros((len(particle_name), self.num_bins), dtype=np.int64)
@@ -439,6 +444,15 @@ def update(
                 bins=[count_bins, count_bins]
             )
             self.detection_count_confusion[cluster_name] += confusion.astype(np.int64, copy=False)
+            truth_count_np = truth_count.detach().cpu().numpy()
+            count_probabilities_np = count_probabilities.detach().cpu().numpy()
+            for count_value in range(len(names) + 1):
+                count_mask = truth_count_np == count_value
+                if not np.any(count_mask):
+                    continue
+                self.detection_count_probability_sum[cluster_name][count_value] += (
+                    count_probabilities_np[:, count_mask].sum(axis=1)
+                )
 
             for num_resonance in range(len(names)):
                 truth_mask = (truth_count == (num_resonance + 1))
@@ -561,6 +575,9 @@ def reset(self):
         for cluster_name, confusion in self.detection_count_confusion.items():
             self.detection_count_confusion[cluster_name] = np.zeros_like(confusion)
 
+        for cluster_name, probability_sum in self.detection_count_probability_sum.items():
+            self.detection_count_probability_sum[cluster_name] = np.zeros_like(probability_sum)
+
         for cluster_name, truth_count_metrics in self.detection_survival_metrics.items():
             for truth_count, hist in truth_count_metrics.items():
                 self.detection_survival_metrics[cluster_name][truth_count] = np.zeros_like(hist)
@@ -597,6 +614,11 @@ def reduce_across_gpus(self):
                 torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
                 self.detection_count_confusion[cluster_name] = tensor.cpu().numpy()
 
+            for cluster_name, probability_sum in self.detection_count_probability_sum.items():
+                tensor = torch.tensor(probability_sum, dtype=torch.float64, device=self.device)
+                torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
+                self.detection_count_probability_sum[cluster_name] = tensor.cpu().numpy()
+
             for cluster_name, truth_count_metrics in self.detection_survival_metrics.items():
                 for truth_count, hist in truth_count_metrics.items():
                     tensor = torch.tensor(hist, dtype=torch.long, device=self.device)
@@ -850,6 +872,12 @@ def plot_detection_count_distribution(self):
                 out=np.zeros_like(confusion, dtype=float),
                 where=row_sum > 0,
             )
+            soft_pmf = np.divide(
+                self.detection_count_probability_sum[cluster_name],
+                np.maximum(row_sum, 1),
+                out=np.zeros_like(self.detection_count_probability_sum[cluster_name], dtype=float),
+                where=row_sum > 0,
+            )
 
             fig, (ax_matrix, ax_pdf) = plt.subplots(
                 2,
@@ -886,21 +914,21 @@ def plot_detection_count_distribution(self):
                 total = row_sum[truth_count, 0]
                 if total <= 0:
                     continue
-                expected_count = (pmf[truth_count] * predicted_counts).sum()
+                expected_count = (soft_pmf[truth_count] * predicted_counts).sum()
                 ax_pdf.step(
                     predicted_counts,
-                    pmf[truth_count],
+                    soft_pmf[truth_count],
                     where="mid",
                     linewidth=2,
                     marker="o",
-                    label=f"Truth {truth_count}{cluster_name}: E[N]={expected_count:.2f}",
+                    label=f"Truth {truth_count}{cluster_name}: soft E[N]={expected_count:.2f}",
                 )
 
             ax_pdf.set_xticks(predicted_counts)
             ax_pdf.set_ylim(0, 1)
-            ax_pdf.set_xlabel("Predicted N (argmax)")
-            ax_pdf.set_ylabel("P(predicted N | truth N)")
-            ax_pdf.set_title("Count Distribution by Truth")
+            ax_pdf.set_xlabel("N")
+            ax_pdf.set_ylabel("Mean predicted P(N)")
+            ax_pdf.set_title("Soft Count Probability by Truth")
             ax_pdf.grid(True, linestyle=":", linewidth=0.5, alpha=0.5)
             ax_pdf.legend(loc="best", fontsize=8)
             fig.tight_layout()
@@ -913,12 +941,13 @@ def detection_count_summary_log(self):
             row_sum = confusion.sum(axis=1)
             max_count = confusion.shape[0] - 1
             predicted_counts = np.arange(max_count + 1)
+            probability_sum = self.detection_count_probability_sum[cluster_name]
             for truth_count in range(max_count + 1):
                 denominator = row_sum[truth_count]
                 if denominator <= 0:
                     continue
 
-                predicted_mean = (confusion[truth_count] * predicted_counts).sum() / denominator
+                predicted_mean = (probability_sum[truth_count] * predicted_counts).sum() / denominator
                 return_log[
                     f"detection/{self.process}/{cluster_name}/truth_{truth_count}/expected_predicted_n"
                 ] = predicted_mean

From 927ed5a2a18aff960cd9e5baac7ffecae2d1f364 Mon Sep 17 00:00:00 2001
From: tihsu99 <b07202037@ntu.edu.tw>
Date: Tue, 23 Jun 2026 17:23:21 +0200
Subject: [PATCH 10/10] refactor detection metrics and visualization to enhance
 summary and plotting functionality

---
 network/loss/assignment.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/network/loss/assignment.py b/network/loss/assignment.py
index d6003be..abe2567 100644
--- a/network/loss/assignment.py
+++ b/network/loss/assignment.py
@@ -293,12 +293,15 @@ def loss_single_process(
     ## Detection Loss ##
     ####################
 
-    particle_balance_weight = torch.ones_like(targets_mask[0], dtype=torch.float32)
+    base_particle_balance_weight = torch.ones_like(targets_mask[0], dtype=torch.float32)
     masks_for_balance = torch.stack(targets_mask).int()
 
     if particle_index_tensor is not None:
         class_indices = (masks_for_balance * particle_index_tensor.to(masks_for_balance.device).unsqueeze(1)).sum(0).int()
-        particle_balance_weight *= particle_weights_tensor.to(masks_for_balance.device)[class_indices]
+        base_particle_balance_weight = (
+                base_particle_balance_weight *
+                particle_weights_tensor.to(masks_for_balance.device)[class_indices]
+        )
 
     detections = detections
     detections_target = targets_mask
@@ -329,15 +332,15 @@ def loss_single_process(
 
     process_masking = torch.stack(process_masking).float()
     process_weighting = torch.stack(process_weighting).float()
-    particle_balance_weight = particle_balance_weight.unsqueeze(0)
+    detection_particle_balance_weight = base_particle_balance_weight.unsqueeze(0)
 
     if event_weight is not None:
         event_balance_weight = event_weight.view(-1, *([1] * (process_masking.dim() - 1)))
-        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * event_balance_weight * particle_balance_weight
-        valid_process = torch.sum(process_masking * process_weighting * event_balance_weight * particle_balance_weight)
+        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * event_balance_weight * detection_particle_balance_weight
+        valid_process = torch.sum(process_masking * process_weighting * event_balance_weight * detection_particle_balance_weight)
     else:
-        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * particle_balance_weight
-        valid_process = torch.sum(process_masking * process_weighting * particle_balance_weight)
+        detection_losses = torch.stack(detection_losses) * process_masking * process_weighting * detection_particle_balance_weight
+        valid_process = torch.sum(process_masking * process_weighting * detection_particle_balance_weight)
 
     if valid_process > 0:
         detection_loss = torch.sum(detection_losses) / valid_process
@@ -360,10 +363,10 @@ def loss_single_process(
         focal_gamma
     )
 
-    particle_balance_weight = particle_balance_weight.squeeze(0).to(symmetric_losses.device)
+    particle_balance_weight = base_particle_balance_weight.to(symmetric_losses.device)
 
     if process_weight[0] is not None:
-        particle_balance_weight *= process_weight[0]
+        particle_balance_weight = particle_balance_weight * process_weight[0]
 
     targets_mask_finite = torch.stack(targets_mask).float()
     if not torch.isfinite(symmetric_losses).all():