Implement Single-Head Multi-Track Finder

QTracker_training/.gitignore

@@ -3,7 +3,12 @@
 checkpoints/
 Slurm_Files/
 local/
+docs/
 
 run.sh
 
-__pycache__/
+__pycache__/
+
+.cursor/settings.json
+.claude/settings.local.json
+CLAUDE.local.md

QTracker_training/CLAUDE.md

@@ -264,3 +264,5 @@ Pre-commit hook configured with ruff (v0.14.4):
 6. **File size considerations**: Training datasets can be very large (500K-19M events). Use `data/skim.py` or `data/skim_flat.py` to create manageable subsets for development.
 
 7. **Mixed precision**: Enabled by default in TrackFinder for performance. Final output layers cast to FP32 for numerical stability.
+
+8. **Commit and PR Rules**: Do NOT mention Claude Code. Preface commit msg with tag like `feat` or `fix`.

QTracker_training/data/multi_track/gen_training_random.py

@@ -23,9 +23,14 @@ def _get_value(x):
 
 
 def combine_files(
-    file1, file2, output_file, pairsmup, pairsmum, use_random, at_least_one_pair
+    file1,
+    file2,
+    output_file,
+    pairsmup,
+    pairsmum,
+    use_random,
+    at_least_one_pair,
 ):
-
     if os.path.exists(output_file):
         os.remove(output_file)
 
@@ -108,7 +113,6 @@ def combine_files(
     ev = 0
 
     while idx_mup < n1 and idx_mum < n2:
-
         if use_random:
             lower_bound = 1 if at_least_one_pair else 0
             max_possible_pairs = min(pairsmup, pairsmum)
@@ -156,7 +160,6 @@ def combine_files(
 
         # -------- Fill paired tracks --------
         for k in range(current_pairs):
-
             tree1.GetEntry(idx_mup)
             tree2.GetEntry(idx_mum)
 
@@ -203,6 +206,15 @@ def combine_files(
             idx_mup += 1
             idx_mum += 1
 
+        # --- Canonical ordering: sort filled pair slots by mu+ element ID sum ---
+        if current_pairs > 1:
+            pair_indices = list(range(current_pairs))
+            pair_indices.sort(
+                key=lambda k: int(np.sum(hitarray_mup[k][hitarray_mup[k] > 0]))
+            )
+            hitarray_mup[:current_pairs] = hitarray_mup[pair_indices]
+            hitarray_mum[:current_pairs] = hitarray_mum[pair_indices]
+
         # -------- Extra mu+ --------
         for _ in range(extra_mup):
             tree1.GetEntry(idx_mup)
@@ -364,6 +376,18 @@ def add_hit_array(input_file, output_file):
         action="store_true",
         help="Ensure at least one pair of mu+ and mu- in each event",
     )
+    parser.add_argument(
+        "--output_mom1",
+        type=str,
+        default="momentum_training-1.root",
+        help="Output path for mu+ momentum training file",
+    )
+    parser.add_argument(
+        "--output_mom2",
+        type=str,
+        default="momentum_training-2.root",
+        help="Output path for mu- momentum training file",
+    )
 
     args = parser.parse_args()
 
@@ -383,8 +407,8 @@ def add_hit_array(input_file, output_file):
     if pairsmup <= 0 or pairsmum <= 0:
         raise ValueError("pairsmup and pairsmum must be >= 1")
 
-    file1_array_output = "momentum_training-1.root"
-    file2_array_output = "momentum_training-2.root"
+    file1_array_output = args.output_mom1
+    file2_array_output = args.output_mom2
     for file_name in [args.output, file1_array_output, file2_array_output]:
         if os.path.exists(file_name):
             os.remove(file_name)

QTracker_training/eval_multi_track.py

@@ -7,6 +7,7 @@
 absl.logging.set_verbosity("error")
 
 import argparse
+import itertools
 import ROOT  # noqa: F401
 import numpy as np
 import tensorflow as tf
@@ -56,6 +57,62 @@ def chi_squared(y_true, y_pred):
     return chi2_mean
 
 
+def match_predictions(y_test, y_pred_argmax, mask):
+    """Reorder predicted pair slots per event to best match the active GT pairs.
+
+    For each event we restrict to the active GT pairs (slots with any nonzero
+    hit; GT is assumed already canonically sorted). We build a residual-based
+    cost between every predicted slot ``p`` and every active GT slot ``g``:
+
+        cost(p, g) = sum_{unmasked det} |pred_p - gt_g| for mu+
+                   + sum_{unmasked det} |pred_p - gt_g| for mu-
+
+    We enumerate permutations of predicted slots and pick the one minimizing
+    the total cost over the active GT slots. The chosen permutation is applied
+    so the predicted slot matched to GT rank ``k`` lands at index ``k``.
+    GT and model pair-slot counts may differ (e.g. 3-pair GT vs 5-pair model).
+    """
+    matched = y_pred_argmax.copy()
+    det_idx = np.where(mask)[0]
+
+    for ev in range(len(y_test)):
+        gt = y_test[ev].astype(np.int32)  # (n_gt, 2, 62)
+        pred = y_pred_argmax[ev]  # (n_pred, 2, 62)
+        n_gt = gt.shape[0]
+        n_pred = pred.shape[0]
+
+        n_active = max(
+            (k + 1 for k in range(n_gt) if np.any(gt[k] != 0)),
+            default=0,
+        )
+        if n_active <= 1:
+            continue
+
+        n_match = min(n_active, n_pred)
+
+        cost = np.zeros((n_pred, n_match))
+        for p in range(n_pred):
+            for g in range(n_match):
+                c_plus = np.sum(np.abs(pred[p, 0, det_idx] - gt[g, 0, det_idx]))
+                c_minus = np.sum(np.abs(pred[p, 1, det_idx] - gt[g, 1, det_idx]))
+                cost[p, g] = c_plus + c_minus
+
+        best_perm = None
+        best_cost = None
+        for perm in itertools.permutations(range(n_pred), n_match):
+            total = sum(cost[perm[k], k] for k in range(n_match))
+            if best_cost is None or total < best_cost:
+                best_cost = total
+                best_perm = perm
+
+        reordered = pred.copy()
+        for k in range(n_match):
+            reordered[k] = pred[best_perm[k]]
+        matched[ev] = reordered
+
+    return matched
+
+
 def evaluate_model(args):
     # Load data - existing loader handles both formats
     load_result = data_loader.load_data(
@@ -76,6 +133,25 @@ def evaluate_model(args):
         print(f"{'=' * 70}\n")
         y_test = np.stack([y_muPlus_test, y_muMinus_test], axis=2)
         # Shape: (num_events, max_pairs, 2, 62)
+
+        # Apply canonical ordering to GT pairs so evaluation matches training ordering
+        for ev_idx in range(len(y_test)):
+            n_active = max(
+                (
+                    k + 1
+                    for k in range(y_test.shape[1])
+                    if np.any(y_test[ev_idx, k] != 0)
+                ),
+                default=0,
+            )
+            if n_active > 1:
+                active_indices = np.arange(n_active)
+                sort_keys = [
+                    np.sum(y_test[ev_idx, k, 0, :][y_test[ev_idx, k, 0, :] > 0])
+                    for k in active_indices
+                ]
+                sorted_order = active_indices[np.argsort(sort_keys)]
+                y_test[ev_idx, :n_active] = y_test[ev_idx, sorted_order]
     else:
         print("\n" + "=" * 70)
         print("Single-track format detected")
@@ -118,60 +194,28 @@ def evaluate_model(args):
 
     # Extract argmax predictions
     y_pred_argmax = np.argmax(y_pred, axis=-1).astype(np.int32)
-    # Shape: (num_events, max_pairs, 2, 62)
+    # Shape: (num_events, pred_max_pairs, 2, 62)
+
+    gt_max_pairs = y_test.shape[1]
+    pred_max_pairs = y_pred_argmax.shape[1]
+    if gt_max_pairs != pred_max_pairs:
+        print(
+            f"WARNING: GT has {gt_max_pairs} pair slots but model outputs "
+            f"{pred_max_pairs}. Evaluating the first "
+            f"{min(gt_max_pairs, pred_max_pairs)} pair(s)."
+        )
+    eval_max_pairs = min(gt_max_pairs, pred_max_pairs)
+
+    # Reorder predicted pair slots per event to best match the active GT pairs
+    y_pred_argmax = match_predictions(y_test, y_pred_argmax, mask)
 
     # Evaluate each pair
-    for pair_idx in range(max_pairs):
+    for pair_idx in range(eval_max_pairs):
         print(f"\n{'=' * 70}")
         print(f"Evaluating Pair {pair_idx}")
         print(f"{'=' * 70}")
 
-        # ============================================================
-        # Pair Existence Evaluation (captures FP and FN)
-        # ============================================================
-
-        print("\n--- Pair Existence Metrics ---")
-
-        # Ground truth existence
-        gt_exists = np.any(y_test[:, pair_idx, :, :] != 0, axis=(1, 2))
-
-        # Prediction existence (after argmax)
-        pred_exists = np.any(y_pred_argmax[:, pair_idx, :, :] != 0, axis=(1, 2))
-
-        TP = np.sum(gt_exists & pred_exists)
-        TN = np.sum(~gt_exists & ~pred_exists)
-        FP = np.sum(~gt_exists & pred_exists)
-        FN = np.sum(gt_exists & ~pred_exists)
-
-        total = len(gt_exists)
-
-        accuracy_exist = (TP + TN) / total if total > 0 else 0.0
-        precision = TP / (TP + FP) if (TP + FP) > 0 else 0.0
-        recall = TP / (TP + FN) if (TP + FN) > 0 else 0.0
-        specificity = TN / (TN + FP) if (TN + FP) > 0 else 0.0
-        f1 = (
-            2 * precision * recall / (precision + recall)
-            if (precision + recall) > 0
-            else 0.0
-        )
-
-        print(f"Total events: {total}")
-        print(f"True Positives : {TP}")
-        print(f"True Negatives : {TN}")
-        print(f"False Positives: {FP}")
-        print(f"False Negatives: {FN}")
-
-        print(f"\nExistence Accuracy : {accuracy_exist:.4f}")
-        print(f"Precision          : {precision:.4f}")
-        print(f"Recall             : {recall:.4f}")
-        print(f"Specificity        : {specificity:.4f}")
-        print(f"F1 Score           : {f1:.4f}")
-
-        if np.sum(~gt_exists) > 0:
-            fp_rate_empty = FP / np.sum(~gt_exists)
-            print(f"\nFalse Positive Rate on Empty Pairs: {fp_rate_empty:.4f}")
-
-        # Check for non-zero ground truth to determine valid events
+        # Restrict to events where this GT rank is active (any nonzero hit)
         valid_mask = np.any(y_test[:, pair_idx, :, :] != 0, axis=(1, 2))
 
         num_valid = np.sum(valid_mask)
@@ -181,7 +225,7 @@ def evaluate_model(args):
 
         print(f"Valid events: {num_valid}/{len(y_test)}")
 
-        # Extract predictions and ground truth for this pair
+        # Extract matched predictions and ground truth for this pair
         y_p_raw = y_pred_argmax[valid_mask, pair_idx, 0, :]  # (valid_events, 62)
         y_m_raw = y_pred_argmax[valid_mask, pair_idx, 1, :]
 

QTracker_training/models/MultiTrackFinder.py

@@ -19,15 +19,12 @@
 
 from backbones import unetpp_backbone
 from data_loader import load_data_denoise
-from losses import multi_track_loss, weighted_bce
+from losses import min_perm_multi_track_loss, weighted_bce
 
 # Set seeds
 tf.random.set_seed(42)
 np.random.seed(42)
 
-# Ensure the checkpoints directory exists
-os.makedirs("checkpoints", exist_ok=True)
-
 # Set mixed precision policy for better performance
 mixed_precision.set_global_policy("mixed_float16")
 
@@ -185,9 +182,12 @@ def train_model(args: argparse.Namespace) -> None:
             optimizer=optimizer,
             loss={
                 "denoise": weighted_bce(pos_weight=args.pos_weight),
-                "segment": multi_track_loss(
+                "segment": min_perm_multi_track_loss(
+                    max_pairs=args.max_pairs,
                     lambda_presence=args.lambda_presence,
                     pos_weight_presence=args.pos_weight_presence,
+                    focal_gamma=args.focal_gamma,
+                    lambda_diversity=args.lambda_diversity,
                 ),
             },
             loss_weights={
@@ -336,6 +336,9 @@ def train_model(args: argparse.Namespace) -> None:
             verbose=2,
         )
 
+    output_dir = os.path.dirname(args.output_model)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
     model.save(args.output_model)
     print(f"Model saved to {args.output_model}")
 
@@ -500,7 +503,7 @@ def train_model(args: argparse.Namespace) -> None:
     parser.add_argument(
         "--lambda_presence",
         type=float,
-        default=0.2,
+        default=1.0,
         help="Weight for presence term in multi-track loss.",
     )
     parser.add_argument(
@@ -509,6 +512,18 @@ def train_model(args: argparse.Namespace) -> None:
         default=5.0,
         help="Positive class weight for presence term in multi-track loss.",
     )
+    parser.add_argument(
+        "--focal_gamma",
+        type=float,
+        default=2.0,
+        help="Gamma for focal loss in presence term (0 = standard BCE).",
+    )
+    parser.add_argument(
+        "--lambda_diversity",
+        type=float,
+        default=0.05,
+        help="Weight for inter-pair diversity penalty.",
+    )
     args = parser.parse_args()
 
     train_model(args)