File size: 9,457 Bytes

d561d1e

"""
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_<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()