""" ============================================================= MMA-CCT HARD v2: HARDER Neutrino Detection - Overlapping signal/background distributions - Realistic noise and ambiguity ============================================================= """ import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import Dataset, DataLoader import math import numpy as np from typing import Dict from dataclasses import dataclass # ============================================================================= # CONFIG # ============================================================================= @dataclass class Config: d_model: int = 64 n_heads: int = 4 n_layers: int = 2 max_hits: int = 50 n_features: int = 12 n_classes: int = 2 train_samples: int = 5000 val_samples: int = 1000 test_samples: int = 1000 signal_ratio: float = 0.1 batch_size: int = 32 epochs: int = 5 lr: float = 1e-3 CONFIG = Config() # ============================================================================= # HARD REALISTIC SIMULATOR # ============================================================================= class HardNeutrinoSimulator: def __init__(self, seed=42): self.rng = np.random.RandomState(seed) def simulate(self, is_signal: bool) -> Dict: if is_signal: base_energy = self.rng.exponential(5.0) if self.rng.rand() < 0.2: base_energy = self.rng.uniform(0.5, 2.0) base_hits = int(self.rng.poisson(25)) if self.rng.rand() < 0.15: base_hits = self.rng.randint(5, 15) flavor = self.rng.randint(0, 3) theta = self.rng.uniform(0, np.pi) phi = self.rng.uniform(0, 2 * np.pi) direction = [ np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta) ] label = 1 else: base_energy = self.rng.exponential(3.0) if self.rng.rand() < 0.1: base_energy = self.rng.uniform(5.0, 15.0) base_hits = int(self.rng.poisson(15)) if self.rng.rand() < 0.15: base_hits = self.rng.randint(20, 35) flavor = -1 direction = [0, 0, 0] label = 0 # Position noise hit_noise_scale = 30 if is_signal else 50 hits_base = self.rng.randn(base_hits, 3) * hit_noise_scale offset = self.rng.randn(3) * 20 hits = hits_base + offset if self.rng.rand() < 0.1: hits += self.rng.randn(1, 3) * 10 # Features features = np.zeros(12) energy_noisy = base_energy * self.rng.uniform(0.8, 1.2) features[0] = np.log10(energy_noisy + 1e-6) / 3.0 if is_signal: features[1:4] = direction else: features[1:4] = self.rng.randn(3) * 0.5 features[4] = base_hits / 50.0 if base_hits > 0: features[5] = (hits[:, 0].mean() + self.rng.randn() * 10) / 500.0 features[6] = (hits[:, 1].mean() + self.rng.randn() * 10) / 500.0 features[7] = (hits[:, 2].mean() + self.rng.randn() * 10) / 500.0 features[8] = (hits[:, 0].std() + self.rng.randn() * 5) / 50.0 features[9] = (hits[:, 1].std() + self.rng.randn() * 5) / 50.0 features[10] = (hits[:, 2].std() + self.rng.randn() * 5) / 50.0 else: features[5:11] = 0.0 features[11] = flavor / 3.0 if flavor >= 0 else -1 # Prepare hits hits = hits / 500.0 if len(hits) < CONFIG.max_hits: padding = np.zeros((CONFIG.max_hits - len(hits), 3)) hits = np.vstack([hits, padding]) else: hits = hits[:CONFIG.max_hits] hits = hits + self.rng.randn(*hits.shape) * 0.01 return { 'hits': torch.FloatTensor(hits), 'features': torch.FloatTensor(features), 'label': torch.LongTensor([label]), # Will squeeze later 'energy': base_energy, 'num_hits': base_hits } class HardDataset(Dataset): def __init__(self, num_samples, signal_ratio, simulator): self.num_samples = num_samples num_signal = int(num_samples * signal_ratio) num_bg = num_samples - num_signal self.data = [] for _ in range(num_signal): self.data.append(simulator.simulate(is_signal=True)) for _ in range(num_bg): self.data.append(simulator.simulate(is_signal=False)) np.random.shuffle(self.data) def __len__(self): return self.num_samples def __getitem__(self, idx): return self.data[idx] # ============================================================================= # MODEL # ============================================================================= class TinyAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.proj = nn.Linear(d_model, d_model * 3) def forward(self, x, mask=None): B, L, D = x.shape qkv = self.proj(x).reshape(B, L, 3, self.n_heads, self.d_k) qkv = qkv.permute(2, 0, 3, 1, 4) Q, K, V = qkv[0], qkv[1], qkv[2] scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn = F.softmax(scores, dim=-1) out = torch.matmul(attn, V).transpose(1, 2).reshape(B, L, D) return out class TinyLayer(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.attn = TinyAttention(d_model, n_heads) self.ff = nn.Sequential( nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x, mask=None): x = x + self.attn(self.norm1(x), mask) x = x + self.ff(self.norm2(x)) return x class LightCCT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.hit_embed = nn.Sequential( nn.Linear(3, config.d_model), nn.LayerNorm(config.d_model), nn.GELU() ) self.hit_pos = nn.Parameter(torch.randn(1, config.max_hits, config.d_model) * 0.02) self.hit_layers = nn.ModuleList([ TinyLayer(config.d_model, config.n_heads) for _ in range(config.n_layers) ]) self.feature_net = nn.Sequential( nn.Linear(config.n_features, config.d_model), nn.LayerNorm(config.d_model), nn.GELU(), nn.Linear(config.d_model, config.d_model) ) self.fusion = nn.Linear(config.d_model * 2, config.d_model) self.classifier = nn.Sequential( nn.LayerNorm(config.d_model), nn.Linear(config.d_model, config.d_model // 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(config.d_model // 2, config.n_classes) ) def forward(self, hits, features): B = hits.shape[0] hit_x = self.hit_embed(hits) + self.hit_pos hit_mask = (hits.abs().sum(dim=-1) > 1e-5).float() for layer in self.hit_layers: hit_x = layer(hit_x, mask=hit_mask.unsqueeze(1).unsqueeze(2)) hit_mask = hit_mask.unsqueeze(-1) hit_pooled = (hit_x * hit_mask).sum(dim=1) / (hit_mask.sum(dim=1) + 1e-6) feat_rep = self.feature_net(features) joint = torch.cat([hit_pooled, feat_rep], dim=-1) fused = F.gelu(self.fusion(joint)) logits = self.classifier(fused) return {'logits': logits} # ============================================================================= # LOSS # ============================================================================= class ProperCCTLoss(nn.Module): def __init__(self, signal_weight=9.0): super().__init__() self.signal_weight = signal_weight def forward(self, outputs, labels): # labels shape: [batch, 1] -> squeeze to [batch] labels = labels.squeeze(-1) # FIX: make 1D ce = F.cross_entropy(outputs['logits'], labels, reduction='none') weight = torch.where(labels == 1, self.signal_weight, 1.0) ce_weighted = (ce * weight).mean() probs = F.softmax(outputs['logits'], dim=-1) entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1).mean() total = ce_weighted + 0.05 * entropy return total, { 'ce': ce_weighted.item(), 'entropy': entropy.item(), 'total': total.item() } # ============================================================================= # TRAINING # ============================================================================= def train_hard(config, device='cpu'): print("=" * 50) print("MMA-CCT HARD v2 - Realistic Difficulty") print("=" * 50) print("\n[1/4] Creating datasets...") train_sim = HardNeutrinoSimulator(seed=42) val_sim = HardNeutrinoSimulator(seed=43) test_sim = HardNeutrinoSimulator(seed=44) train_ds = HardDataset(config.train_samples, config.signal_ratio, train_sim) val_ds = HardDataset(config.val_samples, config.signal_ratio, val_sim) test_ds = HardDataset(config.test_samples, config.signal_ratio, test_sim) train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True) val_loader = DataLoader(val_ds, batch_size=config.batch_size) test_loader = DataLoader(test_ds, batch_size=config.batch_size) print(f" Train: {len(train_ds)} samples ({int(config.train_samples * config.signal_ratio)} signal)") print(f" Val: {len(val_ds)} samples") print(f" Test: {len(test_ds)} samples") print("\n[2/4] Creating model...") model = LightCCT(config).to(device) num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f" Parameters: {num_params:,}") print("\n[3/4] Training...") print("-" * 50) criterion = ProperCCTLoss() optimizer = optim.AdamW(model.parameters(), lr=config.lr, weight_decay=0.01) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.epochs) best_val_acc = 0 best_state = None for epoch in range(config.epochs): model.train() train_loss = 0 train_correct = 0 train_total = 0 for batch in train_loader: hits = batch['hits'].to(device) features = batch['features'].to(device) labels = batch['label'].squeeze(-1).to(device) # FIX: squeeze here too optimizer.zero_grad() outputs = model(hits, features) loss, _ = criterion(outputs, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() train_loss += loss.item() preds = outputs['logits'].argmax(1) train_correct += (preds == labels).sum().item() train_total += labels.size(0) scheduler.step() # Validation model.eval() val_loss = 0 val_correct = 0 val_total = 0 with torch.no_grad(): for batch in val_loader: hits = batch['hits'].to(device) features = batch['features'].to(device) labels = batch['label'].squeeze(-1).to(device) outputs = model(hits, features) loss, _ = criterion(outputs, labels) val_loss += loss.item() preds = outputs['logits'].argmax(1) val_correct += (preds == labels).sum().item() val_total += labels.size(0) train_acc = train_correct / train_total val_acc = val_correct / val_total print(f"Epoch {epoch+1:2d}/{config.epochs} | " f"Loss: {train_loss/len(train_loader):.3f} | " f"Train: {train_acc:.3f} | Val: {val_acc:.3f}") if val_acc > best_val_acc: best_val_acc = val_acc best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} print("-" * 50) model.load_state_dict(best_state) print("\n[4/4] Evaluating...") model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for batch in test_loader: hits = batch['hits'].to(device) features = batch['features'].to(device) labels = batch['label'].squeeze(-1).to(device) outputs = model(hits, features) preds = outputs['logits'].argmax(1) all_preds.extend(preds.cpu().tolist()) all_labels.extend(labels.cpu().tolist()) all_preds = np.array(all_preds) all_labels = np.array(all_labels) accuracy = (all_preds == all_labels).mean() tp = ((all_preds == 1) & (all_labels == 1)).sum() tn = ((all_preds == 0) & (all_labels == 0)).sum() fp = ((all_preds == 1) & (all_labels == 0)).sum() fn = ((all_preds == 0) & (all_labels == 1)).sum() precision = tp / (tp + fp + 1e-6) recall = tp / (tp + fn + 1e-6) f1 = 2 * precision * recall / (precision + recall + 1e-6) print("\n" + "=" * 50) print("RESULTS (Harder Distribution)") print("=" * 50) print(f"Accuracy: {accuracy:.4f}") print(f"Precision: {precision:.4f}") print(f"Recall: {recall:.4f}") print(f"F1 Score: {f1:.4f}") print(f"\nConfusion Matrix:") print(f" TP: {tp:4d} FN: {fn:4d}") print(f" FP: {fp:4d} TN: {tn:4d}") print("=" * 50) return model if __name__ == "__main__": device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") model = train_hard(CONFIG, device)