import os import random import numpy as np import pandas as pd from tqdm import tqdm from array import array import ROOT import torch from sklearn.model_selection import train_test_split from sklearn.metrics import roc_auc_score from tabpfn import TabPFNClassifier # ensure TabPFN is deterministic torch.manual_seed(42) np.random.seed(42) random.seed(42) os.environ["PYTHONHASHSEED"] = str(42) torch.set_num_threads(1) # Note: torch.set_num_interop_threads(1) removed to avoid runtime error def tabpfn(signal, bkgd, batch_size=20_000, test_size=0.5, random_state=42): # Set random seeds for reproducibility torch.manual_seed(random_state) np.random.seed(random_state) random.seed(random_state) os.environ["PYTHONHASHSEED"] = str(random_state) torch.set_num_threads(1) # Try to set interop threads, but handle the case where it's already been set try: torch.set_num_interop_threads(1) except RuntimeError: pass # Interop threads already set, continue signal = np.nan_to_num(signal).astype(np.float32) bkgd = np.nan_to_num(bkgd).astype(np.float32) columns = ['ph1_pt', 'ph1_eta', 'ph1_phi', 'ph2_pt', 'ph2_eta', 'ph2_phi', \ 'lep1_pt', 'lep1_eta', 'lep1_phi', 'lep2_pt', 'lep2_eta', 'lep2_phi', \ 'jet1_pt', 'jet1_eta', 'jet1_phi', 'jet2_pt', 'jet2_eta', 'jet2_phi', \ 'jet3_pt', 'jet3_eta', 'jet3_phi', 'jet4_pt', 'jet4_eta', 'jet4_phi', \ 'jet5_pt', 'jet5_eta', 'jet5_phi', 'jet6_pt', 'jet6_eta', 'jet6_phi', \ 'met_pt', 'met_phi', 'weight', 'SumWeights', 'XSection', \ 'ph1_isTightID', 'ph2_isTightID', \ 'scaleFactor_PILEUP', 'scaleFactor_PHOTON', 'scaleFactor_PhotonTRIGGER', \ 'scaleFactor_ELE', 'scaleFactor_MUON', 'scaleFactor_LepTRIGGER', 'scaleFactor_BTAG', \ 'm_yy', 'pt_yy'] classifier_columns = ['ph1_pt', 'ph2_pt', 'ph1_eta', 'ph2_eta', 'delta_phi'] signal_scores = np.zeros(signal.shape[0]) bkgd_scores = np.zeros(bkgd.shape[0]) signal_df = pd.DataFrame(signal, columns=columns) signal_df['delta_phi'] = signal_df['ph2_phi'] - signal_df['ph1_phi'] signal_df['ph1_pt'] /= signal_df['m_yy'] signal_df['ph2_pt'] /= signal_df['m_yy'] signal_df = signal_df[classifier_columns] signal_df.replace([np.inf, -np.inf], 0.0, inplace=True) signal_df.fillna(0.0, inplace=True) bkgd_df = pd.DataFrame(bkgd, columns=columns) bkgd_df['delta_phi'] = bkgd_df['ph2_phi'] - bkgd_df['ph1_phi'] bkgd_df['ph1_pt'] /= bkgd_df['m_yy'] bkgd_df['ph2_pt'] /= bkgd_df['m_yy'] bkgd_df = bkgd_df[classifier_columns] bkgd_df.replace([np.inf, -np.inf], 0.0, inplace=True) bkgd_df.fillna(0.0, inplace=True) signal_df['target'] = 1 bkgd_df['target'] = 0 signal_df_temp = signal_df.iloc[0:batch_size] bkgd_df_temp = bkgd_df.iloc[0:batch_size] df = pd.concat([bkgd_df_temp, signal_df_temp]) df = df.sort_values(by='ph1_pt') df = df.sample(frac=1, random_state=random_state) x_train, x_test, y_train, y_test = train_test_split(df, df['target'], test_size=test_size, random_state=random_state) device = 'cuda' if torch.cuda.is_available() else 'cpu' print('Device: ', device) clf = TabPFNClassifier(ignore_pretraining_limits=True, device=device) clf.fit(x_train[df.columns[0:-1]], y_train) prediction_probabilities = clf.predict_proba(x_test[df.columns[0:-1]]) print('ROC AUC:', roc_auc_score(y_test, prediction_probabilities[:, 1])) start_idx = 0 bar_format = '{l_bar}{bar:20}{r_bar}{bar:-10b}' n_iterations = (signal.shape[0] + batch_size - 1) // batch_size # evaluating on full sample takes ~10 min for _ in tqdm(range(n_iterations), desc='Signal score inference', unit='batch', bar_format=bar_format, total=n_iterations): stop_idx = min(start_idx + batch_size, signal.shape[0]) signal_df_temp = signal_df.iloc[start_idx:stop_idx] signal_scores[start_idx:stop_idx] = clf.predict_proba(signal_df_temp.iloc[:, :-1])[:,1] start_idx += batch_size start_idx = 0 n_iterations = (bkgd.shape[0] + batch_size - 1) // batch_size for _ in tqdm(range(n_iterations), desc='Bkgd score inference', unit='batch', bar_format=bar_format, total=n_iterations): stop_idx = min(start_idx + batch_size, bkgd.shape[0]) bkgd_df_temp = bkgd_df.iloc[start_idx:stop_idx] bkgd_scores[start_idx:stop_idx] = clf.predict_proba(bkgd_df_temp.iloc[:, :-1])[:,1] start_idx += batch_size return signal_scores, bkgd_scores def load_datasets(signal, bkgd, signal_scores, bkgd_scores): signal_weights = signal[:,32] bkgd_weights = bkgd[:,32] # Determine job-specific output directory for ROOT files # Prefer environment variable OUTPUT_DIR (set by runner), else fallback to CWD output_dir = os.environ.get('OUTPUT_DIR', os.getcwd()) results_dir = os.path.join(output_dir, 'results') os.makedirs(results_dir, exist_ok=True) signal_root_path = os.path.join(results_dir, 'signal.root') bkgd_root_path = os.path.join(results_dir, 'bkgd.root') signal_tree = ROOT.TTree('output', 'output') s_score = array('d', [0.0]) s_weight = array('d', [0.0]) signal_tree.Branch('ml_score', s_score, 'ml_score/D') signal_tree.Branch('normalized_weight', s_weight, 'normalized_weight/D') for i in range(len(signal_scores)): s_score[0] = signal_scores[i] s_weight[0] = signal_weights[i] signal_tree.Fill() signal_file = ROOT.TFile(signal_root_path, 'RECREATE') signal_tree.Write() signal_file.Close() bkgd_tree = ROOT.TTree('output', 'output') b_score = array('d', [0.0]) b_weight = array('d', [0.0]) bkgd_tree.Branch('ml_score', b_score, 'ml_score/D') bkgd_tree.Branch('normalized_weight', b_weight, 'normalized_weight/D') for i in range(len(bkgd_scores)): b_score[0] = bkgd_scores[i] b_weight[0] = bkgd_weights[i] bkgd_tree.Fill() bkgd_file = ROOT.TFile(bkgd_root_path, 'RECREATE') bkgd_tree.Write() bkgd_file.Close() signal_df = ROOT.RDataFrame('output', signal_root_path) bkgd_df = ROOT.RDataFrame('output', bkgd_root_path) return signal_df, bkgd_df def place_boundary(signal_df, bkgd_df, boundaries, num_bins, min_events): boundaries = np.array(boundaries) b_candidates = [] Z_candidates = [] for idx in range(boundaries.shape[0]-1): start_score = boundaries[idx] stop_score = boundaries[idx+1] b, _ = get_optimal_cut_sb(signal_df, bkgd_df, start_score, stop_score, num_bins, min_events) b_candidates.append(b) boundaries_copy = np.copy(boundaries) if b<0: Z_candidates.append(0) continue i = np.searchsorted(boundaries, b) boundaries_copy = np.insert(boundaries, i, b) Z = get_significance(signal_df, bkgd_df, boundaries_copy) Z_candidates.append(Z) best_idx = np.argmax(Z_candidates) return float(b_candidates[best_idx]), float(Z_candidates[best_idx]) def get_optimal_cut_sb(signal_df, bkgd_df, start_score, stop_score, num_bins, min_events): bin_edges = np.linspace(0, 1, num_bins + 1) score = 'ml_score' title = 'Signal/Background;ML Score;Event Count' signal_hist = signal_df.Histo1D(('signal_histogram', title, num_bins, 0, 1), score, 'normalized_weight') bkgd_hist = bkgd_df.Histo1D(('bkgd_histogram', title, num_bins, 0, 1), score, 'normalized_weight') signal_hist_unweighted = signal_df.Histo1D(('signal_histogram_unweighted', title, num_bins, 0, 1), score) bkgd_hist_unweighted = bkgd_df.Histo1D(('bkgd_histogram_unweighted', title, num_bins, 0, 1), score) # ROOT histogram bins defined s.t. bin containing bin boundary *starts* at bin boundary # since we want to include start_score and exclude stop_score, we should define stop_bin to be bin *below* bin containing stop_score start_bin = signal_hist.FindBin(float(start_score)) stop_bin = signal_hist.FindBin(float(stop_score))-1 ZZ = [] candidate_boundaries = [] for b in range(start_bin + 1, stop_bin): signal_lower_yield = signal_hist.Integral(start_bin, b-1) signal_upper_yield = signal_hist.Integral(b, stop_bin) bkgd_lower_yield = bkgd_hist.Integral(start_bin, b-1) bkgd_upper_yield = bkgd_hist.Integral(b, stop_bin) signal_lower_counts = signal_hist_unweighted.Integral(start_bin, b-1) signal_upper_counts = signal_hist_unweighted.Integral(b, stop_bin) bkgd_lower_counts = bkgd_hist_unweighted.Integral(start_bin, b-1) bkgd_upper_counts = bkgd_hist_unweighted.Integral(b, stop_bin) if check_counts_sb(signal_lower_counts, signal_upper_counts, bkgd_lower_counts, bkgd_upper_counts, min_events): Z_lower = Z_sb(signal_lower_yield, bkgd_lower_yield) Z_upper = Z_sb(signal_upper_yield, bkgd_upper_yield) Z_lower = np.nan_to_num(Z_lower, nan=0.0) Z_upper = np.nan_to_num(Z_upper, nan=0.0) Z_tot = Z_comb(np.array([Z_lower, Z_upper])) ZZ.append(Z_tot) else: ZZ.append(0) candidate_boundaries.append(bin_edges[b]) ZZ = np.array(ZZ) if len(ZZ) > 0: optimal_cut = candidate_boundaries[np.argmax(ZZ)] else: optimal_cut = -1 return optimal_cut, ZZ def check_counts_sb(signal_lower_counts, signal_upper_counts, bkgd_lower_counts, bkgd_upper_counts, min_events): return min(signal_lower_counts, signal_upper_counts, bkgd_lower_counts, bkgd_upper_counts) > min_events def Z_sb(s, b): s = np.array(s, ndmin=1) b = np.array(b, ndmin=1) ZZ = np.zeros_like(b, dtype=np.float64) mask = b > 0 ZZ[mask] = np.sqrt(2 * ((s[mask] + b[mask]) * np.log(1 + s[mask] / b[mask]) - s[mask])) return ZZ def Z_comb(zz): return np.sqrt(np.sum(zz**2)) def get_significance(signal_df, bkgd_df, boundaries): boundaries = np.array(boundaries) ZZ = [] score = 'ml_score' for idx in range(boundaries.shape[0]-1): start_score = boundaries[idx] stop_score = boundaries[idx+1] selection = f'{score} >= {start_score} && {score} < {stop_score}' s = signal_df.Filter(selection).Sum('normalized_weight').GetValue() b = bkgd_df.Filter(selection).Sum('normalized_weight').GetValue() ZZ.append(Z_sb(s, b)) return float(Z_comb(np.array(ZZ)))