""" EMR (Extended Mass Range) binary classifier -- cropped channels 700:1018 (318 ch). Task: distinguish the 34 known EMR spectra (class 1) from unlabelled background spectra (class 0) in the cropped high-mass region, then rank the full unknown pool by EMR-likeness. Imbalance handling (34 pos vs ~1000 sampled neg, ~30:1): - Focal Loss (alpha=0.85, gamma=2.0) - WeightedRandomSampler (inverse class frequency) - Shallow 1D CNN + dropout to limit overfitting on the 34 positives - 5-seed ensemble (one model per stratified 5-fold) for robustness Outputs (EMR Classifier/): models/emr_model_seed_.pth -- 5 ensemble members models/metrics.json -- per-fold AUPRC / precision / recall / F1 """ import json import os import numpy as np import pandas as pd import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler from sklearn.model_selection import StratifiedKFold from sklearn.metrics import precision_recall_curve, auc, f1_score, precision_score, recall_score # --- Paths --- ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # EMR Classifier/ MODELS_DIR = os.path.join(HERE, 'models') os.makedirs(MODELS_DIR, exist_ok=True) POS_PATH = os.path.join(ROOT, 'Data Parquet', 'training data for extended mass range.parquet') NEG_PATH = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter_cropped.parquet') # --- Hyperparameters --- TARGET_LEN = 318 # channels 700..1017 N_NEGATIVES = 1000 # sampled background (class 0) BATCH_SIZE = 64 EPOCHS = 100 PATIENCE = 15 LEARNING_RATE = 1e-3 WEIGHT_DECAY = 1e-4 SEEDS = [42, 100, 2026, 777, 888] DATA_SEED = 42 # fixed negative subsample # --- Focal Loss --- class FocalLoss(nn.Module): def __init__(self, alpha=0.85, gamma=2.0, reduction='mean'): super().__init__() self.alpha = alpha self.gamma = gamma self.reduction = reduction def forward(self, inputs, targets): bce = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction='none') targets = targets.type(torch.float32) pt = torch.exp(-bce) loss = self.alpha * (1 - pt) ** self.gamma * bce return torch.mean(loss) if self.reduction == 'mean' else loss # --- Model (same shallow 1D CNN as the HMOC classifier) --- class Spectra1DCNN(nn.Module): def __init__(self, seq_len): super().__init__() self.conv = nn.Sequential( nn.Conv1d(1, 32, kernel_size=7, padding=3), nn.BatchNorm1d(32), nn.ReLU(), nn.Conv1d(32, 64, kernel_size=5, padding=2), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.AdaptiveAvgPool1d(1), ) self.fc = nn.Sequential( nn.Linear(128, 64), nn.ReLU(), nn.Dropout(0.4), nn.Linear(64, 1), ) def forward(self, x): x = self.conv(x).squeeze(-1) return self.fc(x) class SpectraDataset(Dataset): def __init__(self, X, y): self.X = torch.tensor(X, dtype=torch.float32).unsqueeze(1) self.y = torch.tensor(y, dtype=torch.float32).unsqueeze(1) def __len__(self): return len(self.X) def __getitem__(self, idx): return self.X[idx], self.y[idx] def minmax(X): lo = X.min(axis=1, keepdims=True) hi = X.max(axis=1, keepdims=True) return (X - lo) / (hi - lo + 1e-8) def load_data(): print('Loading EMR positives...') pos = pd.read_parquet(POS_PATH) X_pos = np.stack(pos['spectrum_cropped'].to_numpy()).astype(np.float32) print(f' positives: {X_pos.shape}') print('Loading unknown (cropped) for negative sampling...') neg = pd.read_parquet(NEG_PATH, columns=['spectrum_cropped']) X_neg_all = np.stack(neg['spectrum_cropped'].to_numpy()).astype(np.float32) print(f' unknown pool: {X_neg_all.shape}') rng = np.random.default_rng(DATA_SEED) neg_idx = rng.choice(len(X_neg_all), size=N_NEGATIVES, replace=False) X_neg = X_neg_all[neg_idx] print(f' sampled negatives: {X_neg.shape}') X = np.concatenate([X_pos, X_neg], axis=0) y = np.concatenate([np.ones(len(X_pos)), np.zeros(len(X_neg))]).astype(np.int64) return X, y def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'Using device: {device}') X, y = load_data() print(f'Combined: X={X.shape} class counts: {np.bincount(y)}') X = minmax(X) skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=DATA_SEED) folds = list(skf.split(X, y)) all_metrics = {} for seed_idx, seed in enumerate(SEEDS): fold = folds[seed_idx % 5] train_idx, val_idx = fold[0], fold[1] X_train, y_train = X[train_idx], y[train_idx] X_val, y_val = X[val_idx], y[val_idx] print(f'\n--- Model {seed_idx + 1}/{len(SEEDS)} (seed {seed}) ---') print(f' train: {len(train_idx)} (pos={int(y_train.sum())}) ' f'val: {len(val_idx)} (pos={int(y_val.sum())})') torch.manual_seed(seed) np.random.seed(seed) train_ds = SpectraDataset(X_train, y_train) val_ds = SpectraDataset(X_val, y_val) class_count = np.bincount(y_train, minlength=2) weight = 1.0 / class_count samples_weight = torch.tensor(weight[y_train], dtype=torch.float64) sampler = WeightedRandomSampler(samples_weight, num_samples=len(samples_weight), replacement=True) train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, sampler=sampler, num_workers=0) val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) model = Spectra1DCNN(TARGET_LEN).to(device) criterion = FocalLoss(alpha=0.85, gamma=2.0) optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) best_auprc, patience_counter = 0.0, 0 best_state = None for epoch in range(EPOCHS): model.train() train_loss = 0.0 for bx, by in train_loader: bx, by = bx.to(device), by.to(device) optimizer.zero_grad() loss = criterion(model(bx), by) loss.backward() optimizer.step() train_loss += loss.item() model.eval() preds, targs = [], [] with torch.no_grad(): for bx, by in val_loader: bx = bx.to(device) preds.extend(torch.sigmoid(model(bx)).cpu().numpy().flatten()) targs.extend(by.numpy().flatten()) preds, targs = np.array(preds), np.array(targs) precision, recall, _ = precision_recall_curve(targs, preds) auprc = auc(recall, precision) if auprc > best_auprc: best_auprc = auprc best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} patience_counter = 0 else: patience_counter += 1 if (epoch + 1) % 5 == 0 or np.isnan(auprc): print(f' Epoch {epoch + 1:03d} | loss {train_loss / len(train_loader):.4f} ' f'| val AUPRC {auprc:.4f}') if patience_counter >= PATIENCE: print(f' Early stopping at epoch {epoch + 1}') break # Restore best weights and compute fold metrics at best-AUPRC checkpoint model.load_state_dict(best_state) model.eval() preds, targs = [], [] with torch.no_grad(): for bx, by in val_loader: bx = bx.to(device) preds.extend(torch.sigmoid(model(bx)).cpu().numpy().flatten()) targs.extend(by.numpy().flatten()) preds, targs = np.array(preds), np.array(targs) # Best-F1 threshold scan best_f1, best_thr = 0.0, 0.5 for thr in np.arange(0.05, 0.96, 0.05): f1 = f1_score(targs, (preds >= thr).astype(int), zero_division=0) if f1 > best_f1: best_f1, best_thr = f1, thr pred_bin = (preds >= best_thr).astype(int) all_metrics[f'seed_{seed}'] = { 'val_auprc': float(best_auprc), 'val_precision': float(precision_score(targs, pred_bin, zero_division=0)), 'val_recall': float(recall_score(targs, pred_bin, zero_division=0)), 'val_f1': float(best_f1), 'best_threshold': float(best_thr), } print(f' Best val AUPRC {best_auprc:.4f} | P {all_metrics[f"seed_{seed}"]["val_precision"]:.3f} ' f'R {all_metrics[f"seed_{seed}"]["val_recall"]:.3f} F1 {best_f1:.3f} @ thr {best_thr:.2f}') model_path = os.path.join(MODELS_DIR, f'emr_model_seed_{seed}.pth') torch.save(model.state_dict(), model_path) print(f' Saved {model_path}') with open(os.path.join(MODELS_DIR, 'metrics.json'), 'w') as f: json.dump(all_metrics, f, indent=2) print(f'\nMetrics saved to {os.path.join(MODELS_DIR, "metrics.json")}') if __name__ == '__main__': main()