From b8d1b09204f7675ba4a1a478d9ae82c54949a8b8 Mon Sep 17 00:00:00 2001 From: Donghwa Shin Date: Sun, 7 Jun 2026 14:00:16 -0400 Subject: [PATCH 1/2] feat: implement ensemble fixed-N multi-track finder pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five-component ensemble framework for finding multiple dimuon pairs per event, replacing the dynamic-occupancy MultiTrackFinder: - data/multi_track/gen_fixed_pairs.py: fork of gen_training_random.py that pins current_pairs=N (fixed) with random extra single-muon noise within the remaining budget; HitArray sized to exactly (N,62) - models/losses.py: add min_perm_loss(n_pairs) — permutation-invariant loss that brute-forces N! assignments (N<=3) to find minimum-cost pair slot matching; reproduces custom_loss exactly at N=1 - models/data_loader.py: add load_data_counter() returning (X, counts) for training the track counter classifier - models/MultiTrackFinder.py: unified N code path (no N=1 special case); output always (B,N,2,62,201); uses min_perm_loss; CLI --n_pairs replaces --max_pairs/--lambda_presence/--pos_weight_presence - models/TrackCounter.py: new classification model predicting number of dimuon pairs (0..max_pairs) using denoiser U-Net++ backbone + GlobalAveragePooling2D + Dense head with curriculum learning - eval_multi_track.py: rewrite with brute-force permutation matching before metrics, per-pair reporting, no chi-squared; --n_pairs arg - scripts/preprocess_fixed.slurm: generates fixed-N and random data - scripts/train_counter.slurm: trains TrackCounter (single A100) - scripts/train_finder.slurm: trains MultiTrackFinder for given N (4xA100), parameterized via N_PAIRS env var, runs eval after training --- .../data/multi_track/gen_fixed_pairs.py | 396 +++++++++++++++++ QTracker_training/eval_multi_track.py | 268 ++++-------- QTracker_training/models/MultiTrackFinder.py | 58 +-- QTracker_training/models/TrackCounter.py | 412 ++++++++++++++++++ QTracker_training/models/data_loader.py | 60 +++ QTracker_training/models/losses.py | 94 ++++ .../scripts/preprocess_fixed.slurm | 219 ++++++++++ QTracker_training/scripts/train_counter.slurm | 53 +++ QTracker_training/scripts/train_finder.slurm | 63 +++ 9 files changed, 1412 insertions(+), 211 deletions(-) create mode 100644 QTracker_training/data/multi_track/gen_fixed_pairs.py create mode 100644 QTracker_training/models/TrackCounter.py create mode 100644 QTracker_training/scripts/preprocess_fixed.slurm create mode 100644 QTracker_training/scripts/train_counter.slurm create mode 100644 QTracker_training/scripts/train_finder.slurm diff --git a/QTracker_training/data/multi_track/gen_fixed_pairs.py b/QTracker_training/data/multi_track/gen_fixed_pairs.py new file mode 100644 index 0000000..041612c --- /dev/null +++ b/QTracker_training/data/multi_track/gen_fixed_pairs.py @@ -0,0 +1,396 @@ +import ROOT +import argparse +import numpy as np +from array import array +import time +import os +import random + +""" +This version of gen_training creates events with EXACTLY n_pairs ground-truth +dimuon pairs (fixed, not random), plus random extra unpaired single-muon noise +tracks within the remaining budget. + +--pairsmup A : Maximum number of mu+ tracks per event (must be >= n_pairs) +--pairsmum B : Maximum number of mu- tracks per event (must be >= n_pairs) +--n_pairs N : Fixed number of ground-truth dimuon pairs per event +""" + +NUM_LAYERS = 62 + + +def _get_value(x): + if hasattr(x, "__getitem__"): + return x[0] + return x + + +def combine_files(file1, file2, output_file, pairsmup, pairsmum, n_pairs): + if os.path.exists(output_file): + os.remove(output_file) + + f1 = ROOT.TFile.Open(file1, "READ") + f2 = ROOT.TFile.Open(file2, "READ") + + tree1 = f1.Get("tree") + tree2 = f2.Get("tree") + + fout = ROOT.TFile.Open(output_file, "RECREATE", "", ROOT.kLZMA) + fout.SetCompressionLevel(5) + + output_tree = ROOT.TTree("tree", "Tree with aligned dimuon pairs") + output_tree.SetAutoSave(0) + + # ---------------- Event-level branches ---------------- + event_id = array("i", [0]) + nPairs = array("i", [0]) + nExtraMup = array("i", [0]) + nExtraMum = array("i", [0]) + + output_tree.Branch("eventID", event_id, "eventID/I") + output_tree.Branch("nPairs", nPairs, "nPairs/I") + output_tree.Branch("nExtraMup", nExtraMup, "nExtraMup/I") + output_tree.Branch("nExtraMum", nExtraMum, "nExtraMum/I") + + # ---------------- Flat hit storage ---------------- + mu_id = ROOT.std.vector("int")() + element_id = ROOT.std.vector("int")() + detector_id = ROOT.std.vector("int")() + drift_distance = ROOT.std.vector("double")() + tdc_time = ROOT.std.vector("double")() + hit_id = ROOT.std.vector("int")() + hit_track_id = ROOT.std.vector("int")() + gprocess_id = ROOT.std.vector("int")() + + gcharge = ROOT.std.vector("int")() + gtrack_id = ROOT.std.vector("int")() + gpx = ROOT.std.vector("double")() + gpy = ROOT.std.vector("double")() + gpz = ROOT.std.vector("double")() + gvx = ROOT.std.vector("double")() + gvy = ROOT.std.vector("double")() + gvz = ROOT.std.vector("double")() + + output_tree.Branch("muID", mu_id) + output_tree.Branch("elementID", element_id) + output_tree.Branch("detectorID", detector_id) + output_tree.Branch("driftDistance", drift_distance) + output_tree.Branch("tdcTime", tdc_time) + output_tree.Branch("hitID", hit_id) + output_tree.Branch("hitTrackID", hit_track_id) + output_tree.Branch("gProcessID", gprocess_id) + + output_tree.Branch("gCharge", gcharge) + output_tree.Branch("gTrackID", gtrack_id) + output_tree.Branch("gpx", gpx) + output_tree.Branch("gpy", gpy) + output_tree.Branch("gpz", gpz) + output_tree.Branch("gvx", gvx) + output_tree.Branch("gvy", gvy) + output_tree.Branch("gvz", gvz) + + # ---------------- Structured paired arrays ---------------- + hitarray_mup = np.zeros((n_pairs, NUM_LAYERS), dtype=np.int32) + hitarray_mum = np.zeros((n_pairs, NUM_LAYERS), dtype=np.int32) + + output_tree.Branch( + "HitArray_mup", hitarray_mup, f"HitArray_mup[{n_pairs}][{NUM_LAYERS}]/I" + ) + output_tree.Branch( + "HitArray_mum", hitarray_mum, f"HitArray_mum[{n_pairs}][{NUM_LAYERS}]/I" + ) + + n1 = tree1.GetEntries() + n2 = tree2.GetEntries() + + idx_mup = 0 + idx_mum = 0 + ev = 0 + + while idx_mup < n1 and idx_mum < n2: + # Fixed n_pairs pairs, random extra singles within budget + current_pairs = n_pairs + extra_mup = random.randint(0, pairsmup - n_pairs) + extra_mum = random.randint(0, pairsmum - n_pairs) + + if idx_mup + current_pairs + extra_mup > n1: + break + if idx_mum + current_pairs + extra_mum > n2: + break + + # Reset event + event_id[0] = ev + nPairs[0] = current_pairs + nExtraMup[0] = extra_mup + nExtraMum[0] = extra_mum + + mu_id.clear() + element_id.clear() + detector_id.clear() + drift_distance.clear() + tdc_time.clear() + hit_id.clear() + hit_track_id.clear() + gprocess_id.clear() + gcharge.clear() + gtrack_id.clear() + gpx.clear() + gpy.clear() + gpz.clear() + gvx.clear() + gvy.clear() + gvz.clear() + + hitarray_mup.fill(0) + hitarray_mum.fill(0) + + track_counter = 0 + + # -------- Fill paired tracks -------- + for k in range(current_pairs): + tree1.GetEntry(idx_mup) + tree2.GetEntry(idx_mum) + + # Fill structured arrays + for elem, det in zip(tree1.elementID, tree1.detectorID): + if 1 <= det <= NUM_LAYERS: + hitarray_mup[k, det - 1] = elem + + for elem, det in zip(tree2.elementID, tree2.detectorID): + if 1 <= det <= NUM_LAYERS: + hitarray_mum[k, det - 1] = elem + + # Append flat hits from both + for source in [tree1, tree2]: + for elem, det, drift, tdc, hid, trk, proc in zip( + source.elementID, + source.detectorID, + source.driftDistance, + source.tdcTime, + source.hitID, + source.hitTrackID, + source.gProcessID, + ): + element_id.push_back(elem) + detector_id.push_back(det) + drift_distance.push_back(drift) + tdc_time.push_back(tdc) + hit_id.push_back(hid) + hit_track_id.push_back(trk) + gprocess_id.push_back(proc) + + mu_id.push_back(track_counter + 1) + gcharge.push_back(_get_value(source.gCharge)) + gtrack_id.push_back(_get_value(source.gTrackID)) + gpx.push_back(_get_value(source.gpx)) + gpy.push_back(_get_value(source.gpy)) + gpz.push_back(_get_value(source.gpz)) + gvx.push_back(_get_value(source.gvx)) + gvy.push_back(_get_value(source.gvy)) + gvz.push_back(_get_value(source.gvz)) + + track_counter += 1 + + idx_mup += 1 + idx_mum += 1 + + # -------- Extra mu+ -------- + for _ in range(extra_mup): + tree1.GetEntry(idx_mup) + for elem, det, drift, tdc, hid, trk, proc in zip( + tree1.elementID, + tree1.detectorID, + tree1.driftDistance, + tree1.tdcTime, + tree1.hitID, + tree1.hitTrackID, + tree1.gProcessID, + ): + element_id.push_back(elem) + detector_id.push_back(det) + drift_distance.push_back(drift) + tdc_time.push_back(tdc) + hit_id.push_back(hid) + hit_track_id.push_back(trk) + gprocess_id.push_back(proc) + + idx_mup += 1 + + # -------- Extra mu− -------- + for _ in range(extra_mum): + tree2.GetEntry(idx_mum) + for elem, det, drift, tdc, hid, trk, proc in zip( + tree2.elementID, + tree2.detectorID, + tree2.driftDistance, + tree2.tdcTime, + tree2.hitID, + tree2.hitTrackID, + tree2.gProcessID, + ): + element_id.push_back(elem) + detector_id.push_back(det) + drift_distance.push_back(drift) + tdc_time.push_back(tdc) + hit_id.push_back(hid) + hit_track_id.push_back(trk) + gprocess_id.push_back(proc) + + idx_mum += 1 + + output_tree.Fill() + ev += 1 + + output_tree.Write("", ROOT.TObject.kOverwrite) + fout.Close() + f1.Close() + f2.Close() + + fout = ROOT.TFile.Open(output_file, "READ") + out_tree = fout.Get("tree") + print(f"Events in output tree after writing: {out_tree.GetEntries()}") + print( + "Trees in output file:", + [ + key.GetName() + for key in fout.GetListOfKeys() + if key.GetClassName() == "TTree" + ], + ) + cycles = [ + key.GetCycle() for key in fout.GetListOfKeys() if key.GetClassName() == "TTree" + ] + print("Tree cycles in output file:", cycles) + fout.Close() + + +def add_hit_array(input_file, output_file): + if os.path.exists(output_file): + os.remove(output_file) + + f_in = ROOT.TFile.Open(input_file, "READ") + tree = f_in.Get("tree") + + print(f"Entries in input file {input_file}: {tree.GetEntries()}") + print( + "Trees in input file:", + [ + key.GetName() + for key in f_in.GetListOfKeys() + if key.GetClassName() == "TTree" + ], + ) + + fout = ROOT.TFile.Open(output_file, "RECREATE", "", ROOT.kLZMA) + fout.SetCompressionLevel(5) + + output_tree = tree.CloneTree(0) + output_tree.SetAutoSave(0) + + hit_array = np.zeros((NUM_LAYERS, 2), dtype=np.float64) + output_tree.Branch("HitArray", hit_array, f"HitArray[{NUM_LAYERS}][2]/D") + + fill_count = 0 + for i in range(tree.GetEntries()): + tree.GetEntry(i) + hit_array.fill(0) + + for elem, det, drift in zip( + tree.elementID, tree.detectorID, tree.driftDistance + ): + if 1 <= det <= NUM_LAYERS: + hit_array[det - 1, 0] = elem + hit_array[det - 1, 1] = drift if elem != 0 else 0.0 + + output_tree.Fill() + fill_count += 1 + + print(f"Fill() called {fill_count} times for {output_file}") + print(f"Events in output tree before writing: {output_tree.GetEntries()}") + + output_tree.Write("", ROOT.TObject.kOverwrite) + fout.Close() + f_in.Close() + + fout = ROOT.TFile.Open(output_file, "READ") + out_tree = fout.Get("tree") + print(f"Events in output tree after writing: {out_tree.GetEntries()}") + print( + "Trees in output file:", + [ + key.GetName() + for key in fout.GetListOfKeys() + if key.GetClassName() == "TTree" + ], + ) + cycles = [ + key.GetCycle() for key in fout.GetListOfKeys() if key.GetClassName() == "TTree" + ] + print("Tree cycles in output file:", cycles) + fout.Close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Combine ROOT trees with a fixed number of dimuon pairs per event." + ) + parser.add_argument("file1", type=str, help="First ROOT file (mu+)") + parser.add_argument("file2", type=str, help="Second ROOT file (mu-)") + parser.add_argument( + "--output", type=str, default="finder_training.root", help="Output file" + ) + parser.add_argument( + "--pairsmup", type=int, default=None, help="Max number of mu+ tracks per event" + ) + parser.add_argument( + "--pairsmum", type=int, default=None, help="Max number of mu- tracks per event" + ) + parser.add_argument( + "--n_pairs", + type=int, + required=True, + help="Fixed number of ground-truth dimuon pairs per event", + ) + + args = parser.parse_args() + + if args.n_pairs < 1: + raise ValueError("--n_pairs must be >= 1") + + if args.pairsmup is None and args.pairsmum is None: + args.pairsmup = args.n_pairs + args.pairsmum = args.n_pairs + elif args.pairsmup is None: + args.pairsmup = args.pairsmum + elif args.pairsmum is None: + args.pairsmum = args.pairsmup + + if args.pairsmup < args.n_pairs: + raise ValueError("--pairsmup must be >= --n_pairs") + if args.pairsmum < args.n_pairs: + raise ValueError("--pairsmum must be >= --n_pairs") + + file1_array_output = "momentum_training-1.root" + file2_array_output = "momentum_training-2.root" + for file_name in [args.output, file1_array_output, file2_array_output]: + if os.path.exists(file_name): + os.remove(file_name) + + print(f"ROOT version: {ROOT.gROOT.GetVersion()}") + start_time = time.time() + combine_files( + args.file1, + args.file2, + args.output, + args.pairsmup, + args.pairsmum, + args.n_pairs, + ) + end_time = time.time() + + start_time = time.time() + add_hit_array(args.file1, file1_array_output) + add_hit_array(args.file2, file2_array_output) + end_time = time.time() + + print(f"Generated: {args.output}, {file1_array_output}, {file2_array_output}") diff --git a/QTracker_training/eval_multi_track.py b/QTracker_training/eval_multi_track.py index 381a736..c2caa86 100644 --- a/QTracker_training/eval_multi_track.py +++ b/QTracker_training/eval_multi_track.py @@ -7,19 +7,19 @@ absl.logging.set_verbosity("error") import argparse +import itertools import ROOT # noqa: F401 import numpy as np import tensorflow as tf import matplotlib.pyplot as plt -# core TrackFinder loaders / custom loss from models import data_loader from models.layers import AxialAttention import QTracker from refine import refine_hit_arrays -def plot_residuals(det_ids, res_plus, res_minus, model_path, stage_label): +def plot_residuals(det_ids, res_plus, res_minus, model_path, stage_label, pair_idx): mean_p = np.nanmean(np.abs(res_plus), axis=0) std_p = np.nanstd(np.abs(res_plus), axis=0) mean_m = np.nanmean(np.abs(res_minus), axis=0) @@ -31,167 +31,114 @@ def plot_residuals(det_ids, res_plus, res_minus, model_path, stage_label): plt.axhline(0, linestyle="--", linewidth=1) plt.xlabel("Detector Layer (skipping masked slots)") plt.ylabel("Absolute Residual (predicted − true)") - plt.title(f"Per-layer Absolute Residual ({stage_label.capitalize()})") + plt.title( + f"Pair {pair_idx + 1} Per-layer Absolute Residual ({stage_label.capitalize()})" + ) plt.legend() plt.tight_layout() base = os.path.splitext(os.path.basename(model_path))[0] - fname = f"{base}_{stage_label}_residuals.png" + fname = f"{base}_pair{pair_idx + 1}_{stage_label}_residuals.png" plot_dir = os.path.join(os.path.dirname(__file__), "plots", "multi_track") os.makedirs(plot_dir, exist_ok=True) plt.savefig(os.path.join(plot_dir, fname)) plt.show() -def chi_squared(y_true, y_pred): - # y_true.shape = (num_events, 62) - # y_pred.shape = (num_events, 62) - residuals = y_true - y_pred - sigma = np.std(y_true, axis=0) + 1e-6 # Prevent division by zero - - res_norm = residuals / sigma - - chi2 = np.sum((res_norm**2), axis=1) # Chi-squared per event - chi2_mean = np.mean(chi2) # Mean chi-squared over all events - return chi2_mean +def min_perm_match(y_pred_argmax, y_true, n_pairs): + """ + For each event, reorder predicted pair slots to minimize total L1 residual vs GT. + + Args: + y_pred_argmax: (B, N, 2, 62) int32 predicted element IDs + y_true: (B, N, 2, 62) int32 GT element IDs + n_pairs: int, N + + Returns: + matched_pred: (B, N, 2, 62) int32 -- predictions reordered to best match GT slot order + """ + perms = list(itertools.permutations(range(n_pairs))) + B = y_pred_argmax.shape[0] + matched_pred = np.zeros_like(y_pred_argmax) + + for b in range(B): + best_cost = np.inf + best_perm = list(range(n_pairs)) + for perm in perms: + # cost: sum over i of L1(pred[i], true[perm[i]]) + cost = sum( + np.sum(np.abs(y_pred_argmax[b, i] - y_true[b, perm[i]])) + for i in range(n_pairs) + ) + if cost < best_cost: + best_cost = cost + best_perm = perm + # matched_pred[b, best_perm[i], :, :] = y_pred_argmax[b, i, :, :] + # (pred slot i is matched to GT slot best_perm[i]) + for i, j in enumerate(best_perm): + matched_pred[b, j] = y_pred_argmax[b, i] + + return matched_pred def evaluate_model(args): - # Load data - existing loader handles both formats - load_result = data_loader.load_data( - args.root_file, multi_track=True, max_pairs=args.max_pairs - ) - X_test, y_muPlus_test, y_muMinus_test = load_result[:3] + n_pairs = args.n_pairs + # 1. Load data (always multi_track=True) + X_test, y_muPlus_test, y_muMinus_test = data_loader.load_data( + args.root_file, multi_track=True, max_pairs=n_pairs + ) if X_test is None: return + # y_muPlus_test: (B, N, 62), y_muMinus_test: (B, N, 62) + # Stack to (B, N, 2, 62): + y_true = np.stack([y_muPlus_test, y_muMinus_test], axis=2) - # Detect format from shape - is_multi_track = len(y_muPlus_test.shape) == 3 # (num_events, max_pairs, 62) - - if is_multi_track: - num_events, max_pairs, num_detectors = y_muPlus_test.shape - print(f"\n{'=' * 70}") - print(f"Multi-track format detected: max_pairs={max_pairs}") - print(f"{'=' * 70}\n") - y_test = np.stack([y_muPlus_test, y_muMinus_test], axis=2) - # Shape: (num_events, max_pairs, 2, 62) - else: - print("\n" + "=" * 70) - print("Single-track format detected") - print("=" * 70 + "\n") - y_test = np.stack([y_muPlus_test, y_muMinus_test], axis=1) - # Reshape to (num_events, 1, 2, 62) for uniform processing - y_test = y_test[:, np.newaxis, :, :] - max_pairs = 1 - + # 2. Load detector data for refinement det_test, elem_test, _, _, _ = QTracker.load_detector_element_data(args.root_file) + # 3. Detector mask (same as evaluate.py) mask = np.ones(62, dtype=bool) mask[6:12] = False mask[54:62] = False - custom_objects = {"AxialAttention": AxialAttention} + # 4. Load model model = tf.keras.models.load_model( args.model_path, compile=False, - custom_objects=custom_objects, + custom_objects={"AxialAttention": AxialAttention}, ) - # Run predictions + # 5. Predict in chunks (segment output = index 1, always 5D (B,N,2,62,201)) preds = [] chunk_size = 128 - for i in range(0, len(X_test), chunk_size): X_chunk = tf.cast(X_test[i : i + chunk_size], tf.float32) y_chunk = model.predict(X_chunk, verbose=0) - preds.append(y_chunk[1]) # Segment output - - y_pred = np.concatenate(preds, axis=0) - - # Check prediction shape to determine format - if len(y_pred.shape) == 5: # Multi-track: (num_events, max_pairs, 2, 62, 201) - print("Multi-track model predictions") - elif len(y_pred.shape) == 4: # Single-track: (num_events, 2, 62, 201) - print("Single-track model predictions - reshaping for uniform processing\n") - y_pred = y_pred[:, np.newaxis, :, :, :] # Add pair dimension - - # Extract argmax predictions - y_pred_argmax = np.argmax(y_pred, axis=-1).astype(np.int32) - # Shape: (num_events, max_pairs, 2, 62) - - # Evaluate each pair - for pair_idx in range(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}") + preds.append(y_chunk[1]) # segment output + y_pred = np.concatenate(preds, axis=0) # (B, N, 2, 62, 201) - # Check for non-zero ground truth to determine valid events - valid_mask = np.any(y_test[:, pair_idx, :, :] != 0, axis=(1, 2)) + # 6. argmax over element IDs + y_pred_argmax = np.argmax(y_pred, axis=-1).astype(np.int32) # (B, N, 2, 62) - num_valid = np.sum(valid_mask) - if num_valid == 0: - print(f"No valid events for pair {pair_idx}, skipping...") - continue + # 7. Permutation matching + matched_pred = min_perm_match(y_pred_argmax, y_true, n_pairs) # (B, N, 2, 62) - print(f"Valid events: {num_valid}/{len(y_test)}") + # 8. Per-pair metrics + for pair_idx in range(n_pairs): + print(f"\n{'=' * 60}") + print(f" Pair {pair_idx + 1} of {n_pairs}") + print(f"{'=' * 60}") - # Extract 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, :] + y_p_true = y_true[:, pair_idx, 0, :].astype(np.int32) # (B, 62) + y_m_true = y_true[:, pair_idx, 1, :].astype(np.int32) # (B, 62) + y_p_raw = matched_pred[:, pair_idx, 0, :] # (B, 62) + y_m_raw = matched_pred[:, pair_idx, 1, :] # (B, 62) - y_p_true = y_test[valid_mask, pair_idx, 0, :].astype(np.int32) - y_m_true = y_test[valid_mask, pair_idx, 1, :].astype(np.int32) - - # Compute raw residuals (SAME as single-track) raw_p_res = y_p_true - y_p_raw raw_m_res = y_m_true - y_m_raw + # Raw residuals per detector (same as evaluate.py) print("\n--- Raw Residuals (Before Refinement) ---") print("Det | μ+ mean | μ+ std | μ- mean | μ- std") for det in np.where(mask)[0]: @@ -201,24 +148,22 @@ def evaluate_model(args): s_m = np.std(np.abs(raw_m_res[:, det])) print(f"{det + 1:3d} | {m_p:8.3f} | {s_p:8.3f} | {m_m:8.3f} | {s_m:8.3f}") - # Plot residuals with pair index in filename dets_used = np.where(mask)[0] + 1 plot_residuals( dets_used, raw_p_res[:, mask], raw_m_res[:, mask], args.model_path, - f"pair{pair_idx}_raw", + "raw", + pair_idx, ) - # Refinement (filter detector data for valid events) - det_valid = [det_test[i] for i in np.where(valid_mask)[0]] - elem_valid = [elem_test[i] for i in np.where(valid_mask)[0]] - - ref_p, ref_m = refine_hit_arrays(y_p_raw, y_m_raw, det_valid, elem_valid) + # Refine + ref_p, ref_m = refine_hit_arrays(y_p_raw, y_m_raw, det_test, elem_test) ref_p_res = y_p_true - ref_p ref_m_res = y_m_true - ref_m + # Refined residuals per detector print("\n--- Refined Residuals (After Refinement) ---") print("Det | μ+ mean | μ+ std | μ- mean | μ- std") for det in np.where(mask)[0]: @@ -228,88 +173,63 @@ def evaluate_model(args): s_m = np.std(np.abs(ref_m_res[:, det])) print(f"{det + 1:3d} | {m_p:8.3f} | {s_p:8.3f} | {m_m:8.3f} | {s_m:8.3f}") - # Plot refined residuals plot_residuals( dets_used, ref_p_res[:, mask], ref_m_res[:, mask], args.model_path, - f"pair{pair_idx}_refined", + "refined", + pair_idx, ) - # Calculate accuracy and chi-squared (SAME metrics as single-track) + # Accuracy metrics (NO chi-squared) acc_p = np.mean(np.abs(raw_p_res) == 0) acc_m = np.mean(np.abs(raw_m_res) == 0) - print(f"\nRaw μ+ accuracy: {acc_p:.4f}") - print(f"Raw μ- accuracy: {acc_m:.4f}") + print(f"\nRaw μ+ exact-match accuracy: {acc_p:.4f}") + print(f"Raw μ- exact-match accuracy: {acc_m:.4f}") acc_p = np.mean(np.abs(raw_p_res) <= 2) acc_m = np.mean(np.abs(raw_m_res) <= 2) - print(f"\nRaw μ+ within-2 accuracy: {acc_p:.4f}") + print(f"Raw μ+ within-2 accuracy: {acc_p:.4f}") print(f"Raw μ- within-2 accuracy: {acc_m:.4f}") - chi2_p = chi_squared(y_p_true, y_p_raw) - chi2_m = chi_squared(y_m_true, y_m_raw) - print(f"\nRaw μ+ Chi-squared: {chi2_p:.3f}") - print(f"Raw μ- Chi-squared: {chi2_m:.3f}") - - # Calculate accuracy and chi-squared after refinement acc_p = np.mean(np.abs(ref_p_res) == 0) acc_m = np.mean(np.abs(ref_m_res) == 0) - print(f"\nRefined μ+ accuracy: {acc_p:.4f}") - print(f"Refined μ- accuracy: {acc_m:.4f}") + print(f"\nRefined μ+ exact-match accuracy: {acc_p:.4f}") + print(f"Refined μ- exact-match accuracy: {acc_m:.4f}") acc_p = np.mean(np.abs(ref_p_res) <= 2) acc_m = np.mean(np.abs(ref_m_res) <= 2) - print(f"\nRefined μ+ within-2 accuracy: {acc_p:.4f}") + print(f"Refined μ+ within-2 accuracy: {acc_p:.4f}") print(f"Refined μ- within-2 accuracy: {acc_m:.4f}") - chi2_p = chi_squared(y_p_true, ref_p) - chi2_m = chi_squared(y_m_true, ref_m) - print(f"\nRefined μ+ Chi-squared: {chi2_p:.3f}") - print(f"Refined μ- Chi-squared: {chi2_m:.3f}") - - print("\n--- Raw Absolute Residuals (Before Refinement) ---") - print("μ+ mean | μ+ std | μ- mean | μ- std") + # Global absolute residuals + print("\n--- Raw Global Absolute Residuals ---") m_p, s_p = np.mean(np.abs(raw_p_res)), np.std(np.abs(raw_p_res)) m_m, s_m = np.mean(np.abs(raw_m_res)), np.std(np.abs(raw_m_res)) - print(f"{m_p:8.3f} | {s_p:8.3f} | {m_m:8.3f} | {s_m:8.3f}") + print(f"μ+ mean={m_p:.3f} std={s_p:.3f} | μ- mean={m_m:.3f} std={s_m:.3f}") - print("\n--- Refined Absolute Residuals (After Refinement) ---") - print("μ+ mean | μ+ std | μ- mean | μ- std") + print("\n--- Refined Global Absolute Residuals ---") m_p, s_p = np.mean(np.abs(ref_p_res)), np.std(np.abs(ref_p_res)) m_m, s_m = np.mean(np.abs(ref_m_res)), np.std(np.abs(ref_m_res)) - print(f"{m_p:8.3f} | {s_p:8.3f} | {m_m:8.3f} | {s_m:8.3f}") + print(f"μ+ mean={m_p:.3f} std={s_p:.3f} | μ- mean={m_m:.3f} std={s_m:.3f}") if __name__ == "__main__": parser = argparse.ArgumentParser( - description="Evaluate pre-trained TrackFinder models (single or multi-track)." + description="Evaluate pre-trained fixed-N multi-track finder models." ) parser.add_argument("root_file", type=str, help="Path to the val/test ROOT file.") parser.add_argument( - "model_path", type=str, help="Path to the saved model file (.h5 or .keras)." - ) - parser.add_argument( - "--batch_norm", - type=int, - default=0, - help="Flag to set batch normalization: [0 = False, 1 = True].", - ) - parser.add_argument( - "--base", - type=int, - default=64, - help="Flag to set batch normalization: [0 = False, 1 = True].", + "model_path", type=str, help="Path to the saved model (.keras or .h5)." ) - parser.add_argument("--model", type=str, default=None, help="Model name.") parser.add_argument( - "--max_pairs", + "--n_pairs", type=int, - default=5, - help="Maximum number of dimuon pairs (for multi-track files).", + default=1, + help="Fixed number of dimuon pairs the model was trained with (1-3).", ) args = parser.parse_args() - print(f"\nResults for {args.model_path}...") + print(f"\nEvaluating {args.model_path} (n_pairs={args.n_pairs})...") evaluate_model(args) diff --git a/QTracker_training/models/MultiTrackFinder.py b/QTracker_training/models/MultiTrackFinder.py index 693b93b..37cebb4 100644 --- a/QTracker_training/models/MultiTrackFinder.py +++ b/QTracker_training/models/MultiTrackFinder.py @@ -19,7 +19,7 @@ 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_loss, weighted_bce # Set seeds tf.random.set_seed(42) @@ -46,7 +46,7 @@ def build_model( use_attn: bool = False, use_attn_ffn: bool = True, dropout_attn: float = 0.0, - max_pairs: int = 5, + n_pairs: int = 1, ) -> tf.keras.Model: """ This function builds the joint denoising and segmentation model using two U-Net++ backbones. @@ -63,6 +63,7 @@ def build_model( use_attn (bool): Whether to use axial attention mechanism in segmentation U-Net++ (default: False). use_attn_ffn (bool): Whether to use feed-forward layers in attention (default: True). dropout_attn (float): Dropout rate for attention block (default: 0.0). + n_pairs (int): Fixed number of dimuon pairs per event (default: 1). Returns: tf.keras.Model: The constructed joint denoising and segmentation model. @@ -100,13 +101,11 @@ def build_model( ) # Segmentation Head - x = layers.Conv2D(max_pairs * 2, kernel_size=1)( + x = layers.Conv2D(n_pairs * 2, kernel_size=1)(x) # (batch, det, elem, n_pairs*2) + x = layers.Permute((3, 1, 2))(x) # (batch, n_pairs*2, det, elem) + x = layers.Reshape((n_pairs, 2, num_detectors, num_elementIDs))( x - ) # (batch, det, elem, max_pairs*2) - x = layers.Permute((3, 1, 2))(x) # (batch, max_pairs*2, det, elem) - x = layers.Reshape((max_pairs, 2, num_detectors, num_elementIDs))( - x - ) # (batch, max_pairs, 2, det, elem) + ) # (batch, n_pairs, 2, det, elem) seg_output = layers.Softmax(axis=-1, name="segment", dtype=tf.float32)( x ) # softmax over elementID @@ -137,22 +136,22 @@ def train_model(args: argparse.Namespace) -> None: y_muPlus_train_low, y_muMinus_train_low, ) = load_data_denoise( - args.train_root_file_low, multi_track=True, max_pairs=args.max_pairs + args.train_root_file_low, multi_track=True, max_pairs=args.n_pairs ) if X_train_low is None or X_clean_train_low is None: return y_train_low = np.stack( [y_muPlus_train_low, y_muMinus_train_low], axis=2 - ) # Shape: (num_events, max_pairs, 2, 62) + ) # Shape: (num_events, n_pairs, 2, 62) X_val, X_clean_val, y_muPlus_val, y_muMinus_val = load_data_denoise( - args.val_root_file, multi_track=True, max_pairs=args.max_pairs + args.val_root_file, multi_track=True, max_pairs=args.n_pairs ) if X_val is None or X_clean_val is None: return y_val = np.stack( [y_muPlus_val, y_muMinus_val], axis=2 - ) # Shape: (num_events, max_pairs, 2, 62) + ) # Shape: (num_events, n_pairs, 2, 62) with strategy.scope(): model = build_model( @@ -166,7 +165,7 @@ def train_model(args: argparse.Namespace) -> None: use_attn=args.use_attn, use_attn_ffn=args.use_attn_ffn, dropout_attn=args.dropout_attn, - max_pairs=args.max_pairs, + n_pairs=args.n_pairs, ) model.summary() @@ -185,10 +184,7 @@ def train_model(args: argparse.Namespace) -> None: optimizer=optimizer, loss={ "denoise": weighted_bce(pos_weight=args.pos_weight), - "segment": multi_track_loss( - lambda_presence=args.lambda_presence, - pos_weight_presence=args.pos_weight_presence, - ), + "segment": min_perm_loss(args.n_pairs), }, loss_weights={ "denoise": 10.0, @@ -244,13 +240,13 @@ def train_model(args: argparse.Namespace) -> None: y_muPlus_train_med, y_muMinus_train_med, ) = load_data_denoise( - args.train_root_file_med, multi_track=True, max_pairs=args.max_pairs + args.train_root_file_med, multi_track=True, max_pairs=args.n_pairs ) if X_train_med is None or X_clean_train_med is None: return y_train_med = np.stack( [y_muPlus_train_med, y_muMinus_train_med], axis=2 - ) # Shape: (num_events, max_pairs, 2, 62) + ) # Shape: (num_events, n_pairs, 2, 62) K.set_value(model.optimizer.learning_rate, args.lr_med) lr_scheduler = ReduceLROnPlateau( @@ -281,13 +277,13 @@ def train_model(args: argparse.Namespace) -> None: y_muPlus_train_high, y_muMinus_train_high, ) = load_data_denoise( - args.train_root_file_high, multi_track=True, max_pairs=args.max_pairs + args.train_root_file_high, multi_track=True, max_pairs=args.n_pairs ) if X_train_high is None or X_clean_train_high is None: return y_train_high = np.stack( [y_muPlus_train_high, y_muMinus_train_high], axis=2 - ) # Shape: (num_events, max_pairs, 2, 62) + ) # Shape: (num_events, n_pairs, 2, 62) K.set_value(model.optimizer.learning_rate, args.lr_high) lr_scheduler = ReduceLROnPlateau( @@ -365,7 +361,7 @@ def train_model(args: argparse.Namespace) -> None: parser.add_argument( "--output_model", type=str, - default="checkpoints/track_finder_joint.keras", + default="checkpoints/multi_track_finder.keras", help="Path to save the trained model.", ) parser.add_argument( @@ -492,22 +488,10 @@ def train_model(args: argparse.Namespace) -> None: help="Fraction of epochs for medium complexity data.", ) parser.add_argument( - "--max_pairs", + "--n_pairs", type=int, - default=5, - help="Maximum number of possible dimuon pairs in an event.", - ) - parser.add_argument( - "--lambda_presence", - type=float, - default=0.2, - help="Weight for presence term in multi-track loss.", - ) - parser.add_argument( - "--pos_weight_presence", - type=float, - default=5.0, - help="Positive class weight for presence term in multi-track loss.", + default=1, + help="Fixed number of dimuon pairs per event (1-3).", ) args = parser.parse_args() diff --git a/QTracker_training/models/TrackCounter.py b/QTracker_training/models/TrackCounter.py new file mode 100644 index 0000000..6256fc6 --- /dev/null +++ b/QTracker_training/models/TrackCounter.py @@ -0,0 +1,412 @@ +"""Denoiser U-Net++ based track-counter: classifies number of dimuon pairs (0..max_pairs)""" + +# ruff: noqa: E402 +import argparse +import gc +import os + +os.environ["TF_GPU_ALLOCATOR"] = ( + "cuda_malloc_async" # Enable asynchronous GPU memory allocation for better performance +) + +import numpy as np +import ROOT # noqa: F401 +import tensorflow as tf +from tensorflow.keras import layers, mixed_precision +import tensorflow.keras.backend as K +from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau +from tensorflow.keras.optimizers import AdamW +from sklearn.metrics import classification_report, confusion_matrix + +from backbones import unetpp_backbone +from data_loader import load_data_counter + +# 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") + +NUM_DETECTORS = 62 +NUM_ELEMENT_IDS = 201 + + +def build_model( + num_detectors: int = 62, + num_elementIDs: int = 201, + use_bn: bool = False, + dropout_bn: float = 0.0, + dropout_enc: float = 0.0, + denoise_base: int = 64, + max_pairs: int = 3, +) -> tf.keras.Model: + """ + Build the track-counter model using a single denoiser U-Net++ backbone followed + by a global pooling head that predicts how many dimuon pairs are in an event. + + Args: + num_detectors (int): Number of detectors (default: 62). + num_elementIDs (int): Number of element IDs (default: 201). + use_bn (bool): Whether to use batch normalization (default: False). + dropout_bn (float): Dropout rate for bottleneck layer (default: 0.0). + dropout_enc (float): Dropout rate for encoder blocks (default: 0.0). + denoise_base (int): Number of base channels in U-Net++ backbone (default: 64). + max_pairs (int): Maximum number of dimuon pairs; output has max_pairs+1 classes + corresponding to counts {0, 1, ..., max_pairs} (default: 3). + + Returns: + tf.keras.Model: The constructed track-counter model. + """ + + input_layer = layers.Input(shape=(num_detectors, num_elementIDs, 1)) + + # Denoiser backbone (attention disabled — counter does not need axial attention) + x = unetpp_backbone( + input_layer, + num_detectors, + num_elementIDs, + use_bn, + dropout_bn, + dropout_enc, + denoise_base, + use_attn=False, + ) + + # Classification head + x = layers.GlobalAveragePooling2D()(x) + x = layers.Dense(256, activation="relu")(x) + x = layers.Dropout(0.5)(x) + x = layers.Dense(128, activation="relu")(x) + count_output = layers.Dense( + max_pairs + 1, activation="softmax", name="count", dtype=tf.float32 + )(x) + + model = tf.keras.Model(inputs=input_layer, outputs=[count_output]) + return model + + +def train_model(args: argparse.Namespace) -> None: + """ + Train the track-counter model using the provided arguments. + Supports curriculum learning with low, medium, and high complexity datasets. + Utilizes MirroredStrategy for multi-GPU distributed training. + + Args: + args (argparse.Namespace): Command-line arguments for training configuration. + """ + + # Distributed training + strategy = tf.distribute.MirroredStrategy() + print(f"Number of devices: {strategy.num_replicas_in_sync}") + + # Load low complexity training data and validation data + X_train_low, counts_train_low = load_data_counter( + args.train_root_file_low, args.max_pairs + ) + if X_train_low is None: + return + + X_val, counts_val = load_data_counter(args.val_root_file, args.max_pairs) + if X_val is None: + return + + with strategy.scope(): + model = build_model( + num_detectors=NUM_DETECTORS, + num_elementIDs=NUM_ELEMENT_IDS, + use_bn=args.batch_norm, + dropout_bn=args.dropout_bn, + dropout_enc=args.dropout_enc, + denoise_base=args.denoise_base, + max_pairs=args.max_pairs, + ) + model.summary() + + optimizer = AdamW( + learning_rate=args.lr_low, + weight_decay=args.weight_decay, + clipnorm=args.clipnorm, + ) + + model.compile( + optimizer=optimizer, + loss=tf.keras.losses.SparseCategoricalCrossentropy(), + metrics=["accuracy"], + ) + + if args.train_root_file_med and args.train_root_file_high: + # Curriculum learning: low → med → high complexity + print("Curriculum learning enabled.") + + epochs_low = int(args.epochs * args.low_ratio) + epochs_med = int(args.epochs * args.med_ratio) + epochs_high = args.epochs + + # --- Stage 1: low complexity --- + lr_scheduler = ReduceLROnPlateau( + monitor="val_loss", + factor=args.factor, + patience=args.lr_patience, + min_lr=1e-6, + ) + early_stopping = EarlyStopping( + monitor="val_loss", patience=args.patience, restore_best_weights=False + ) + history = model.fit( + X_train_low, + counts_train_low, + initial_epoch=0, + epochs=epochs_low, + batch_size=args.batch_size, + validation_data=(X_val, counts_val), + callbacks=[lr_scheduler, early_stopping], + verbose=2, + ) + print("Stage 1 (low) history:", history.history) + del X_train_low, counts_train_low + gc.collect() + + # --- Stage 2: medium complexity --- + X_train_med, counts_train_med = load_data_counter( + args.train_root_file_med, args.max_pairs + ) + if X_train_med is None: + return + + K.set_value(model.optimizer.learning_rate, args.lr_med) + lr_scheduler = ReduceLROnPlateau( + monitor="val_loss", + factor=args.factor, + patience=args.lr_patience, + min_lr=1e-6, + ) + early_stopping = EarlyStopping( + monitor="val_loss", patience=args.patience, restore_best_weights=False + ) + history = model.fit( + X_train_med, + counts_train_med, + initial_epoch=epochs_low, + epochs=epochs_med, + batch_size=args.batch_size, + validation_data=(X_val, counts_val), + callbacks=[lr_scheduler, early_stopping], + verbose=2, + ) + print("Stage 2 (med) history:", history.history) + del X_train_med, counts_train_med + gc.collect() + + # --- Stage 3: high complexity --- + X_train_high, counts_train_high = load_data_counter( + args.train_root_file_high, args.max_pairs + ) + if X_train_high is None: + return + + K.set_value(model.optimizer.learning_rate, args.lr_high) + lr_scheduler = ReduceLROnPlateau( + monitor="val_loss", + factor=args.factor, + patience=args.lr_patience, + min_lr=1e-6, + ) + early_stopping = EarlyStopping( + monitor="val_loss", patience=args.patience, restore_best_weights=True + ) + history = model.fit( + X_train_high, + counts_train_high, + initial_epoch=epochs_med, + epochs=epochs_high, + batch_size=args.batch_size, + validation_data=(X_val, counts_val), + callbacks=[lr_scheduler, early_stopping], + verbose=2, + ) + print("Stage 3 (high) history:", history.history) + del X_train_high, counts_train_high + gc.collect() + + else: + # Standard single-stage training + print("Standard training without curriculum learning.") + + lr_scheduler = ReduceLROnPlateau( + monitor="val_loss", + factor=args.factor, + patience=args.lr_patience, + min_lr=1e-6, + ) + early_stopping = EarlyStopping( + monitor="val_loss", patience=args.patience, restore_best_weights=False + ) + history = model.fit( + X_train_low, + counts_train_low, + initial_epoch=0, + epochs=args.epochs, + batch_size=args.batch_size, + validation_data=(X_val, counts_val), + callbacks=[lr_scheduler, early_stopping], + verbose=2, + ) + print("Training history:", history.history) + + # Evaluation on validation set + val_preds = model.predict(X_val, batch_size=args.batch_size) + val_pred_counts = np.argmax(val_preds, axis=-1) + print("Confusion Matrix:") + print(confusion_matrix(counts_val, val_pred_counts)) + print("\nClassification Report:") + print(classification_report(counts_val, val_pred_counts, zero_division=0)) + + model.save(args.output_model) + print(f"Model saved to {args.output_model}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Train a TensorFlow model to count the number of dimuon pairs in an event." + ) + parser.add_argument( + "train_root_file_low", + type=str, + help="Path to the low-complexity train ROOT file.", + ) + parser.add_argument( + "val_root_file", type=str, help="Path to the validation ROOT file." + ) + parser.add_argument( + "--train_root_file_med", + type=str, + default=None, + help="Train ROOT file with medium complexity for curriculum learning.", + ) + parser.add_argument( + "--train_root_file_high", + type=str, + default=None, + help="Train ROOT file with high complexity for curriculum learning.", + ) + parser.add_argument( + "--output_model", + type=str, + default="checkpoints/track_counter.keras", + help="Path to save the trained model.", + ) + parser.add_argument( + "--max_pairs", + type=int, + default=3, + help="Maximum number of pairs (counter predicts 0..max_pairs).", + ) + parser.add_argument( + "--lr_low", + type=float, + default=0.0003, + help="Learning rate for low complexity data.", + ) + parser.add_argument( + "--lr_med", + type=float, + default=0.0001, + help="Learning rate for medium complexity data.", + ) + parser.add_argument( + "--lr_high", + type=float, + default=0.00003, + help="Learning rate for high complexity data.", + ) + parser.add_argument( + "--factor", + type=float, + default=0.3, + help="Factor for ReduceLROnPlateau.", + ) + parser.add_argument( + "--patience", type=int, default=12, help="Patience for EarlyStopping." + ) + parser.add_argument( + "--lr_patience", + type=int, + default=4, + help="Patience for learning rate scheduler.", + ) + parser.add_argument( + "--batch_norm", + type=int, + default=0, + help="Flag to set batch normalization: [0 = False, 1 = True].", + ) + parser.add_argument( + "--use_attn", + type=int, + default=0, + help="Flag kept for CLI parity with TrackFinder (ignored; counter always uses use_attn=False).", + ) + parser.add_argument( + "--dropout_bn", + type=float, + default=0.0, + help="Dropout rate for bottleneck layer.", + ) + parser.add_argument( + "--dropout_enc", + type=float, + default=0.0, + help="Dropout rate for encoder blocks.", + ) + parser.add_argument( + "--denoise_base", + type=int, + default=64, + help="Number of base channels in U-Net++ backbone.", + ) + parser.add_argument( + "--epochs", + type=int, + default=40, + help="Total number of training epochs.", + ) + parser.add_argument( + "--batch_size", + type=int, + default=32, + help="Batch size for mini-batch gradient descent.", + ) + parser.add_argument( + "--weight_decay", + type=float, + default=1e-4, + help="Weight decay for AdamW optimizer.", + ) + parser.add_argument( + "--clipnorm", + type=float, + default=1.0, + help="Gradient clipping norm for AdamW optimizer.", + ) + parser.add_argument( + "--low_ratio", + type=float, + default=0.5, + help="Fraction of total epochs to use for low complexity stage.", + ) + parser.add_argument( + "--med_ratio", + type=float, + default=0.8, + help="Fraction of total epochs to use for medium complexity stage.", + ) + args = parser.parse_args() + + # batch_norm and use_attn are stored as ints from argparse; convert to bool + args.batch_norm = bool(args.batch_norm) + + train_model(args) diff --git a/QTracker_training/models/data_loader.py b/QTracker_training/models/data_loader.py index 8a3c018..2b114d5 100644 --- a/QTracker_training/models/data_loader.py +++ b/QTracker_training/models/data_loader.py @@ -143,3 +143,63 @@ def load_data_denoise( y_muMinus = np.array(y_muMinus) return X, X_clean, y_muPlus, y_muMinus + + +def load_data_counter(root_file: str, max_pairs: int) -> Tuple[np.ndarray, np.ndarray]: + """ + Load data from a ROOT file for track-counter model training. + + Returns the noisy hit matrix alongside an integer label indicating how many + dimuon pairs are present in each event (capped at max_pairs). + + Args: + root_file (str): Path to the ROOT file. + max_pairs (int): Maximum number of pairs to predict; counts are clamped + to this value via ``min(n_pairs, max_pairs)``. + + Returns: + Tuple[np.ndarray, np.ndarray]: + - X: Noisy hit matrix of shape (num_events, 62, 201, 1), float32. + Built by setting ``matrix[det_id, elem_id] = 1`` for every hit in + ``event.detectorID`` / ``event.elementID``. + - counts: Integer pair-count labels of shape (num_events,), int32. + Read from the ``nPairs`` branch when available; otherwise falls + back to counting nonzero rows in ``HitArray_mup``. + """ + + f = ROOT.TFile.Open(root_file, "READ") + tree = f.Get("tree") + + if not tree: + print("Error: Tree not found in file.") + return None, None + + num_detectors = 62 + num_elementIDs = 201 + + X = [] + counts = [] + + for event in tree: + event_hits_matrix = np.zeros((num_detectors, num_elementIDs), dtype=np.float32) + + for det_id, elem_id in zip(event.detectorID, event.elementID): + if 0 <= det_id < num_detectors and 0 <= elem_id < num_elementIDs: + event_hits_matrix[det_id, elem_id] = 1 + + try: + n_pairs = int(event.nPairs) + except AttributeError: + # Fallback: count nonzero rows in HitArray_mup. + # HitArray_mup is stored flat as (max_pairs * num_detectors,); reshape + # so each row corresponds to one pair, then count rows with any hit. + hit_arr = np.array(list(event.HitArray_mup), dtype=np.int32) + n_pairs = int(np.any(hit_arr.reshape(-1, num_detectors) != 0, axis=1).sum()) + + counts.append(min(n_pairs, max_pairs)) + X.append(event_hits_matrix) + + X = np.array(X)[..., np.newaxis] # Shape: (num_events, 62, 201, 1) + counts = np.array(counts, dtype=np.int32) # Shape: (num_events,) + + return X, counts diff --git a/QTracker_training/models/losses.py b/QTracker_training/models/losses.py index eab9b05..793a427 100644 --- a/QTracker_training/models/losses.py +++ b/QTracker_training/models/losses.py @@ -1,3 +1,5 @@ +import itertools + import tensorflow as tf from typing import Callable @@ -111,6 +113,98 @@ def loss(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: return loss +def min_perm_loss(n_pairs: int) -> Callable: + """ + Permutation-invariant loss for a fixed-N multi-track finder. + + Brute-forces all N! permutations of predicted track slots vs. ground-truth + track slots and picks the minimum-cost assignment, so track slot ordering + does not affect training. + + For N=1, the computed scalar is mathematically identical to ``custom_loss`` + when called with equivalently shaped inputs (``y_true: (B,1,2,62)``, + ``y_pred: (B,1,2,62,201)`` vs the single-track ``(B,2,62)``/``(B,2,62,201)``). + Note that ``custom_loss`` and ``min_perm_loss(1)`` are not interchangeable + at the call site — inputs must carry the N=1 pair axis. + + Args: + n_pairs (int): Number of dimuon pairs N. Must satisfy 1 <= n_pairs <= 3 + (max 6 permutations at N=3). + + Returns: + A loss function with signature ``loss(y_true, y_pred) -> tf.Tensor``. + + Where: + y_true: Ground truth tensor with shape ``(B, N, 2, 62)``. Integer + element IDs (0 = no hit); axis 2 holds mu+/mu-. + y_pred: Predicted softmax probabilities with shape + ``(B, N, 2, 62, 201)``; axis 2 holds mu+/mu-. + """ + if n_pairs < 1: + raise ValueError(f"n_pairs must be >= 1, got {n_pairs}") + if n_pairs > 3: + raise ValueError( + f"n_pairs > 3 not supported (would generate {__import__('math').factorial(n_pairs)} " + f"permutations); got {n_pairs}" + ) + + # Precompute permutation tensors once at factory time (not per batch). + perms = list(itertools.permutations(range(n_pairs))) + perm_tensors = [tf.constant(list(p), dtype=tf.int32) for p in perms] + + def loss(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ + Args: + y_true (tf.Tensor): Shape ``(B, N, 2, 62)``. + y_pred (tf.Tensor): Shape ``(B, N, 2, 62, 201)``. + + Returns: + tf.Tensor: Scalar loss value. + """ + y_pred = tf.cast(y_pred, tf.float32) + y_true = tf.cast(y_true, tf.int32) + + # --- Cost matrix: cost[b, i, j] = CE cost of assigning pred slot i to GT slot j --- + # Expand pred: (B, N, 1, 2, 62, 201) + pred_expand = tf.expand_dims(y_pred, axis=2) + # Expand true: (B, 1, N, 2, 62) + true_expand = tf.expand_dims(y_true, axis=1) + + # CE per (pred slot, GT slot, muon, detector): (B, N, N, 2, 62) + # axis 3 = muon dim (size 2), axis 4 = detector dim (size 62) + ce = tf.keras.losses.sparse_categorical_crossentropy(true_expand, pred_expand) + + # Sum over muon dim (axis=3): (B, N, N, 62) + # Mean over detector dim (now axis=3): (B, N, N) + cost = tf.reduce_mean(tf.reduce_sum(ce, axis=3), axis=3) + + # --- Find minimum-cost permutation --- + perm_costs = [] + for perm_tensor in perm_tensors: + # Reorder columns of cost by permutation, then trace = sum_i cost[b, i, perm[i]] + cost_perm = tf.gather(cost, perm_tensor, axis=2) # (B, N, N) cols reordered + perm_costs.append(tf.linalg.trace(cost_perm)) # (B,) + + # Stack along axis=1 then take minimum: (B,) + min_cost = tf.reduce_min(tf.stack(perm_costs, axis=1), axis=1) + + # --- Overlap penalty: discourage mu+/mu- collapsing to same position --- + p_plus = y_pred[:, :, 0, :, :] # (B, N, 62, 201) + p_minus = y_pred[:, :, 1, :, :] # (B, N, 62, 201) + # Sum over 201 element IDs per (pair, detector): (B, N, 62) + overlap = tf.reduce_sum(tf.square(p_plus - p_minus), axis=-1) + + # Mean over detectors (axis=2), sum over N pairs (axis=1): (B,) + # For N=1 this reduces to mean_d(sum_c (p+−p−)²), matching custom_loss exactly. + overlap_penalty = OVERLAP_LAMBDA * tf.reduce_sum( + tf.reduce_mean(overlap, axis=2), axis=1 + ) + + return tf.reduce_mean(min_cost + overlap_penalty) + + return loss + + def weighted_bce(pos_weight: float = 1.0) -> Callable: """ Returns a weighted binary cross-entropy loss function. False negatives are penalized more heavily diff --git a/QTracker_training/scripts/preprocess_fixed.slurm b/QTracker_training/scripts/preprocess_fixed.slurm new file mode 100644 index 0000000..3b33468 --- /dev/null +++ b/QTracker_training/scripts/preprocess_fixed.slurm @@ -0,0 +1,219 @@ +#!/bin/bash +#SBATCH -A spinquest_standard +#SBATCH -p standard +#SBATCH -c 4 +#SBATCH -t 72:00:00 +#SBATCH -J preprocess_fixed +#SBATCH -o Slurm_Files/preprocess_fixed.out +#SBATCH -e Slurm_Files/preprocess_fixed.err +#SBATCH --mem=128000 + +module purge + +# For apptainer exec +APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif +CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/ + +# # --- 0. Create train/val/test sets --- +# echo "Split dimuon file into train, val, and test sets" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/split.py \ +# /project/ptgroup/spinquest/QTracker-Data/Dimuons-Target_19M.rus.root \ +# --train_output data/raw_files/Dimuons_Train.root \ +# --val_output data/raw_files/Dimuons_Val.root \ +# --test_output data/raw_files/Dimuons_Test.root + +# echo "Split muon positive file into train, val, and test sets" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/split.py \ +# /project/ptgroup/spinquest/QTracker-Data/MUP-Dump_500M.rus.root \ +# --train_output data/raw_files/MUP_Dump_Train.root \ +# --val_output data/raw_files/MUP_Dump_Val.root \ +# --test_output data/raw_files/MUP_Dump_Test.root + +# echo "Split muon negative file into train, val, and test sets" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/split.py \ +# /project/ptgroup/spinquest/QTracker-Data/MUM-Dump_500M.rus.root \ +# --train_output data/raw_files/MUM_Dump_Train.root \ +# --val_output data/raw_files/MUM_Dump_Val.root \ +# --test_output data/raw_files/MUM_Dump_Test.root + +# # --- 1. Skim train and val files to reduce sample size (for downstream ML) --- +# echo "Sample 500,000 events from raw train file (stratified)" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/skim_flat.py \ +# data/raw_files/Dimuons_Train.root \ +# data/raw_files/Dimuons_Train_500K.root + +# echo "Sample 200,000 events from raw val file (simple random)" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/skim.py \ +# data/raw_files/Dimuons_Val.root \ +# --output_file data/raw_files/Dimuons_Val_200K.root \ +# --max_events 200000 \ +# --random 1 + +# # --- 2. Split signal ROOT file into μ⁺ and μ⁻ tracks --- +# echo "Separate dimuons train set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/separate.py data/raw_files/Dimuons_Train_500K.root + +# echo "Separate dimuons val set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/separate.py data/raw_files/Dimuons_Val_200K.root + +# echo "Separate dimuons test set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/separate.py data/raw_files/Dimuons_Test.root + +# # --- 3. Merge two single-muon ROOT files --- +# echo "Combine MUP and MUM train set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/combine.py data/raw_files/MUP_Dump_Train.root data/raw_files/MUM_Dump_Train.root \ +# --output data/processed_files/single_muons_train.root \ +# --max_output_events 220000000 + +# echo "Combine MUP and MUM val set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/combine.py data/raw_files/MUP_Dump_Val.root data/raw_files/MUM_Dump_Val.root \ +# --output data/processed_files/single_muons_val.root \ +# --max_output_events 27500000 + +# echo "Combine MUP and MUM test set" +# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ +# python3 data/combine.py data/raw_files/MUP_Dump_Test.root data/raw_files/MUM_Dump_Test.root \ +# --output data/processed_files/single_muons_test.root \ +# --max_output_events 27500000 + +# --- 4a. Generate random multi-track data for the counter --- +echo "Generate random training data for train set (counter)" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/gen_training_random.py \ + data/raw_files/Dimuons_Train_500K_track1.root \ + data/raw_files/Dimuons_Train_500K_track2.root \ + --output data/multi_track/processed_files/finder_training_train_random.root \ + --pairsmup 5 \ + --pairsmum 5 \ + --random + +mv momentum_training-1.root data/multi_track/processed_files/momentum_training-1_train.root +mv momentum_training-2.root data/multi_track/processed_files/momentum_training-2_train.root + +echo "Generate random training data for val set (counter)" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/gen_training_random.py \ + data/raw_files/Dimuons_Val_200K_track1.root \ + data/raw_files/Dimuons_Val_200K_track2.root \ + --output data/multi_track/processed_files/finder_training_val_random.root \ + --pairsmup 5 \ + --pairsmum 5 \ + --random + +mv momentum_training-1.root data/multi_track/processed_files/momentum_training-1_val.root +mv momentum_training-2.root data/multi_track/processed_files/momentum_training-2_val.root + +# --- 4b. Generate fixed-pair training data for N=1, 2, 3 --- +for N in 1 2 3; do + echo "Generate fixed-N training data (N=${N}) for train set" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 /mnt/code/data/multi_track/gen_fixed_pairs.py \ + data/raw_files/Dimuons_Train_500K_track1.root \ + data/raw_files/Dimuons_Train_500K_track2.root \ + --output data/multi_track/processed_files/finder_training_train_N${N}.root \ + --n_pairs ${N} \ + --pairsmup 5 \ + --pairsmum 5 + # discard momentum files for finder training (not needed here) + rm -f momentum_training-1.root momentum_training-2.root + + echo "Generate fixed-N training data (N=${N}) for val set" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 /mnt/code/data/multi_track/gen_fixed_pairs.py \ + data/raw_files/Dimuons_Val_200K_track1.root \ + data/raw_files/Dimuons_Val_200K_track2.root \ + --output data/multi_track/processed_files/finder_training_val_N${N}.root \ + --n_pairs ${N} \ + --pairsmup 5 \ + --pairsmum 5 + rm -f momentum_training-1.root momentum_training-2.root +done + +# --- 5a. Inject background into random (counter) training data --- +echo "Inject low-level background tracks into random train set" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_random.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_low_random.root \ + --uniform_tracks 0 \ + --lower_bound 0 \ + --num_tracks 16 + +echo "Inject mid-level background tracks into random train set" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_random.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_med_random.root \ + --uniform_tracks 0 \ + --lower_bound 17 \ + --num_tracks 33 + +echo "Inject high-level background tracks into random train set" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_random.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_high_random.root \ + --uniform_tracks 0 \ + --lower_bound 34 \ + --num_tracks 50 + +echo "Inject background tracks into random val set" +${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_val_random.root \ + data/multi_track/processed_files/single_muons_val.root \ + --output data/multi_track/processed_files/mc_events_val_random.root + +# --- 5b. Inject background into fixed-pair training data for N=1, 2, 3 --- +for N in 1 2 3; do + echo "Inject low-level background tracks into fixed-N train set (N=${N})" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_N${N}.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_low_N${N}.root \ + --uniform_tracks 0 \ + --lower_bound 0 \ + --num_tracks 16 + + echo "Inject mid-level background tracks into fixed-N train set (N=${N})" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_N${N}.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_med_N${N}.root \ + --uniform_tracks 0 \ + --lower_bound 17 \ + --num_tracks 33 + + echo "Inject high-level background tracks into fixed-N train set (N=${N})" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_train_N${N}.root \ + data/multi_track/processed_files/single_muons_train.root \ + --output data/multi_track/processed_files/mc_events_train_high_N${N}.root \ + --uniform_tracks 0 \ + --lower_bound 34 \ + --num_tracks 50 + + echo "Inject background tracks into fixed-N val set (N=${N})" + ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 data/multi_track/messy_gen.py \ + data/multi_track/processed_files/finder_training_val_N${N}.root \ + data/multi_track/processed_files/single_muons_val.root \ + --output data/multi_track/processed_files/mc_events_val_N${N}.root +done diff --git a/QTracker_training/scripts/train_counter.slurm b/QTracker_training/scripts/train_counter.slurm new file mode 100644 index 0000000..e6d2eec --- /dev/null +++ b/QTracker_training/scripts/train_counter.slurm @@ -0,0 +1,53 @@ +#!/bin/bash +#SBATCH -A spinquest_standard +#SBATCH -p gpu +#SBATCH --gres=gpu:a100:1 +#SBATCH --ntasks-per-node=1 +#SBATCH -C gpupod +#SBATCH -c 4 +#SBATCH -t 72:00:00 +#SBATCH -J track-counter +#SBATCH -o Slurm_Files/track-counter.out +#SBATCH -e Slurm_Files/track-counter.err +#SBATCH --mem=256000 + +module purge + +# For apptainer exec +APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif +CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/ + + +### Main Body ### +# --- Training Track Counter --- +BATCH_NORM=1 +USE_ATTENTION=1 # unused by counter but kept for CLI parity +USE_ATTENTION_FFN=0 # unused +MAX_PAIRS=3 + +# GPU validation +${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + echo $CUDA_VISIBLE_DEVICES + +${APPTAINER} exec --nv --env CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 /mnt/code/models/TrackCounter.py \ + /mnt/code/data/multi_track/processed_files/mc_events_train_low_random.root \ + /mnt/code/data/multi_track/processed_files/mc_events_val_random.root \ + --train_root_file_med /mnt/code/data/multi_track/processed_files/mc_events_train_med_random.root \ + --train_root_file_high /mnt/code/data/multi_track/processed_files/mc_events_train_high_random.root \ + --output_model /mnt/code/checkpoints/track_counter.keras \ + --lr_low 0.0003 \ + --lr_med 0.0001 \ + --lr_high 0.00003 \ + --patience 10 \ + --lr_patience 3 \ + --batch_norm $BATCH_NORM \ + --use_attn $USE_ATTENTION \ + --denoise_base 32 \ + --epochs 60 \ + --batch_size 64 \ + --dropout_bn 0.5 \ + --dropout_enc 0.4 \ + --factor 0.3 \ + --max_pairs $MAX_PAIRS diff --git a/QTracker_training/scripts/train_finder.slurm b/QTracker_training/scripts/train_finder.slurm new file mode 100644 index 0000000..791f5b5 --- /dev/null +++ b/QTracker_training/scripts/train_finder.slurm @@ -0,0 +1,63 @@ +#!/bin/bash +#SBATCH -A spinquest_standard +#SBATCH -p gpu +#SBATCH --gres=gpu:a100:4 +#SBATCH --ntasks-per-node=1 +#SBATCH -C gpupod +#SBATCH -c 4 +#SBATCH -t 72:00:00 +#SBATCH -J fixed-track-finder +#SBATCH -o Slurm_Files/fixed-track-finder.out +#SBATCH -e Slurm_Files/fixed-track-finder.err +#SBATCH --mem=256000 + +module purge + +# For apptainer exec +APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif +CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/ + + +### Main Body ### +# --- Training and Evaluating Fixed-N Track Finder --- +BATCH_NORM=1 +USE_ATTENTION=1 +USE_ATTENTION_FFN=0 +N=${N_PAIRS:-1} # default to N=1; override with: sbatch --export=N_PAIRS=2 train_finder.slurm + +# GPU validation +${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + echo $CUDA_VISIBLE_DEVICES + +${APPTAINER} exec --nv --env CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 /mnt/code/models/MultiTrackFinder.py \ + /mnt/code/data/multi_track/processed_files/mc_events_train_low_N${N}.root \ + /mnt/code/data/multi_track/processed_files/mc_events_val_N${N}.root \ + --train_root_file_med /mnt/code/data/multi_track/processed_files/mc_events_train_med_N${N}.root \ + --train_root_file_high /mnt/code/data/multi_track/processed_files/mc_events_train_high_N${N}.root \ + --output_model /mnt/code/checkpoints/finder_N${N}.keras \ + --lr_low 0.0003 \ + --lr_med 0.0001 \ + --lr_high 0.00003 \ + --patience 10 \ + --lr_patience 3 \ + --batch_norm $BATCH_NORM \ + --use_attn $USE_ATTENTION \ + --use_attn_ffn $USE_ATTENTION_FFN \ + --denoise_base 32 \ + --base 64 \ + --epochs 60 \ + --batch_size 64 \ + --dropout_bn 0.5 \ + --dropout_enc 0.4 \ + --dropout_attn 0.1 \ + --factor 0.3 \ + --pos_weight 20.0 \ + --n_pairs $N + +${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \ + python3 /mnt/code/eval_multi_track.py \ + /mnt/code/data/multi_track/processed_files/mc_events_val_N${N}.root \ + /mnt/code/checkpoints/finder_N${N}.keras \ + --n_pairs $N From 9d9396412a975e542b4fbf282299a738e49f8214 Mon Sep 17 00:00:00 2001 From: Donghwa Shin Date: Sun, 7 Jun 2026 16:59:25 -0400 Subject: [PATCH 2/2] fix(slurm): update apptainer version and GPU allocation count --- QTracker_training/scripts/preprocess_fixed.slurm | 2 +- QTracker_training/scripts/train_counter.slurm | 4 ++-- QTracker_training/scripts/train_finder.slurm | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/QTracker_training/scripts/preprocess_fixed.slurm b/QTracker_training/scripts/preprocess_fixed.slurm index 3b33468..95577a0 100644 --- a/QTracker_training/scripts/preprocess_fixed.slurm +++ b/QTracker_training/scripts/preprocess_fixed.slurm @@ -11,7 +11,7 @@ module purge # For apptainer exec -APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/ diff --git a/QTracker_training/scripts/train_counter.slurm b/QTracker_training/scripts/train_counter.slurm index e6d2eec..d9919e1 100644 --- a/QTracker_training/scripts/train_counter.slurm +++ b/QTracker_training/scripts/train_counter.slurm @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH -A spinquest_standard #SBATCH -p gpu -#SBATCH --gres=gpu:a100:1 +#SBATCH --gres=gpu:a100:4 #SBATCH --ntasks-per-node=1 #SBATCH -C gpupod #SBATCH -c 4 @@ -14,7 +14,7 @@ module purge # For apptainer exec -APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/ diff --git a/QTracker_training/scripts/train_finder.slurm b/QTracker_training/scripts/train_finder.slurm index 791f5b5..4cf7a30 100644 --- a/QTracker_training/scripts/train_finder.slurm +++ b/QTracker_training/scripts/train_finder.slurm @@ -14,7 +14,7 @@ module purge # For apptainer exec -APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer +APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/