code/probes/probe_7_contrastive_ablation.py

#!/usr/bin/env python3
"""
Probe 7: Contrastive Learning Ablation — V5-A (MLM-only) vs V5.1 (MLM+SimCLR)

Runs layer-wise CLS probing (Probe 2) and embedding vs structural similarity
(Probe 4) on V5-A, then generates comparative overlay figures against V5.1.

Usage:
  python probe_7_contrastive_ablation.py --device cuda
  python probe_7_contrastive_ablation.py --compare_only
"""

import os, sys, json, csv, argparse
import numpy as np
from pathlib import Path
from collections import Counter

PROJECT_ROOT = Path(__file__).resolve().parents[2]
VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
V5A_CKPT = PROJECT_ROOT / 'checkpoints_v5_bpe_topo' / 'best_v5_bpe_topo_model.pt'
V51_CKPT = PROJECT_ROOT / 'bert_v5.1_contrastive' / 'checkpoints' / 'best_v51_contrastive_model.pt'
BENCH_DIR = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data'
PROBE_DIR = PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'probe_results_v6'

sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'bert_training_v4'))
sys.path.insert(0, str(PROJECT_ROOT / 'bert_v6_contrastive' / 'scripts'))

from probe_2_layerwise_cls import (
    load_model, extract_layerwise_cls,
    load_domain_data, load_glycosylation_data, load_immunogenicity_data,
    train_linear_probe, run_layerwise_probing
)
from downstream_tasks.utils.tokenizer import WURCSTokenizer


def extract_cls_final(model, samples, device='cuda', max_len=256):
    """Extract final-layer CLS embeddings."""
    import torch, torch.nn.functional as F
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    embeddings = []
    for si, s in enumerate(samples):
        try:
            result = tokenizer.tokenize(s['wurcs'], max_length=max_len)
            token_ids = torch.tensor(result['token_ids'], dtype=torch.long)
            bd = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long)
            lt = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long)
            ml = min(len(token_ids), len(bd), len(lt))
            token_ids, bd, lt = token_ids[:ml], bd[:ml], lt[:ml]
            if ml > max_len:
                token_ids, bd, lt = token_ids[:max_len], bd[:max_len], lt[:max_len]
            elif ml < max_len:
                pad = max_len - ml
                token_ids = F.pad(token_ids, (0, pad), value=0)
                bd = F.pad(bd, (0, pad), value=0)
                lt = F.pad(lt, (0, pad), value=0)
            with torch.no_grad():
                hidden = model.seq_embeddings(token_ids.unsqueeze(0).to(device),
                    branch_depths=bd.unsqueeze(0).to(device), linkage_types=lt.unsqueeze(0).to(device))
                for layer in model.seq_layers:
                    hidden = layer(hidden)
                embeddings.append(hidden[0, 0, :].cpu().numpy())
        except:
            embeddings.append(np.zeros(768))
        if si > 0 and si % 500 == 0:
            print(f"    Processed {si}/{len(samples)}")
    return np.array(embeddings)


def load_glyco_with_iupac():
    csv_path = BENCH_DIR / 'glycan_link_wurcs_subset.csv'
    samples, labels, iupacs = [], [], []
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            w, link, iupac = row.get('wurcs',''), row.get('link',''), row.get('glycan','')
            if w.startswith('WURCS') and link in ('N','O') and iupac:
                samples.append({'wurcs': w}); labels.append(link); iupacs.append(iupac)
    print(f"  Glyco with IUPAC: {len(samples)} samples")
    return samples, labels, iupacs


def run_probe4(model, samples, iupacs, model_name, out_dir):
    from sklearn.metrics.pairwise import cosine_similarity
    print(f"\nExtracting CLS ({len(samples)} samples)...")
    embs = extract_cls_final(model, samples)
    embed_sim = cosine_similarity(embs)
    # Structural similarity
    cache = out_dir / 'struct_sim_cache.npy'
    if cache.exists():
        print("  Loading cached struct sim...")
        struct_sim = np.load(str(cache))
    else:
        try:
            from glycowork.motif.annotate import annotate_dataset
            print(f"  Computing motif fingerprints...")
            mdf = annotate_dataset(iupacs, feature_set=['known'])
            struct_sim = cosine_similarity(mdf.values.astype(float))
            np.save(str(cache), struct_sim)
        except Exception as e:
            print(f"  ERROR: {e}"); return None
    n = len(samples)
    iu = np.triu_indices(n, k=1)
    from scipy.stats import pearsonr, spearmanr
    pr, pp = pearsonr(embed_sim[iu], struct_sim[iu])
    sr, sp = spearmanr(embed_sim[iu], struct_sim[iu])
    res = {'model': model_name, 'pearson_r': float(pr), 'spearman_rho': float(sr),
           'n_pairs': int(len(embed_sim[iu])), 'n_glycans': n}
    print(f"  {model_name}: Pearson r={pr:.4f}, Spearman rho={sr:.4f}")
    return res


def plot_ablation(v5a_res, v51_res, out_dir):
    import matplotlib; matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    plt.rcParams.update({'font.size': 11, 'figure.dpi': 300, 'axes.spines.top': False, 'axes.spines.right': False})
    tasks = ['Domain', 'Immunogenicity', 'Glycosylation']
    colors = {'Domain': '#E63946', 'Immunogenicity': '#457B9D', 'Glycosylation': '#2A9D8F'}
    fig, axes = plt.subplots(1, 3, figsize=(18, 5.5))
    for i, task in enumerate(tasks):
        ax = axes[i]
        v5a_t = sorted([r for r in v5a_res if r['task']==task], key=lambda r: r['layer'])
        v51_t = sorted([r for r in v51_res if r['task']==task], key=lambda r: r['layer'])
        if not v5a_t or not v51_t: continue
        layers = [r['layer'] for r in v5a_t]
        v5a_a = [r['accuracy_mean'] for r in v5a_t]
        v51_a = [r['accuracy_mean'] for r in v51_t]
        v5a_s = [r['accuracy_std'] for r in v5a_t]
        v51_s = [r['accuracy_std'] for r in v51_t]
        ax.errorbar(layers, v5a_a, yerr=v5a_s, color='#888', marker='o', ms=6, lw=2,
                   ls='--', capsize=3, label='V5-A (MLM only)', zorder=2)
        ax.fill_between(layers, [a-s for a,s in zip(v5a_a,v5a_s)],
                       [a+s for a,s in zip(v5a_a,v5a_s)], color='#888', alpha=0.08)
        ax.errorbar(layers, v51_a, yerr=v51_s, color=colors[task], marker='s', ms=7,
                   lw=2.5, capsize=3, label='V5.1 (MLM+SimCLR)', zorder=3)
        ax.fill_between(layers, [a-s for a,s in zip(v51_a,v51_s)],
                       [a+s for a,s in zip(v51_a,v51_s)], color=colors[task], alpha=0.1)
        d = max(v51_a) - max(v5a_a)
        bl = v51_t[np.argmax(v51_a)]['layer']
        ax.annotate(f'Delta = {"+" if d>=0 else ""}{d:.3f}', xy=(bl, max(v51_a)),
                   xytext=(bl+1.5, max(v51_a)+0.01), fontsize=10, fontweight='bold',
                   color=colors[task], arrowprops=dict(arrowstyle='->', color=colors[task]))
        ax.set_xlabel('Layer'); ax.set_ylabel('Accuracy')
        ax.set_title(task, fontweight='bold')
        ax.set_xticks(range(13)); ax.set_xticklabels(['Emb']+[str(j) for j in range(1,13)])
        ax.legend(frameon=True, edgecolor='#ccc', loc='lower right')
        ax.grid(axis='y', alpha=0.3, ls='--')
        ax.text(-0.08, 1.05, f'({chr(97+i)})', transform=ax.transAxes, fontsize=14, fontweight='bold')
    fig.suptitle('Probe 7: Contrastive Pre-training Ablation', fontsize=15, fontweight='bold', y=1.02)
    plt.tight_layout()
    out = Path(out_dir)
    for fmt in ['png','pdf']:
        plt.savefig(out/f'probe7_ablation_layerwise.{fmt}', dpi=300, bbox_inches='tight', facecolor='white')
        print(f"  Saved: {out/f'probe7_ablation_layerwise.{fmt}'}")
    plt.close()


def plot_p4_comparison(v5a_p4, v51_p4, out_dir):
    import matplotlib; matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    if not v5a_p4 or not v51_p4: print("  Skip P4 plot (missing data)"); return
    fig, ax = plt.subplots(1, 1, figsize=(6, 5))
    models = ['V5-A\n(MLM only)', 'V5.1\n(MLM+SimCLR)']
    rhos = [v5a_p4['spearman_rho'], v51_p4['spearman_rho']]
    rs = [v5a_p4['pearson_r'], v51_p4['pearson_r']]
    x = np.arange(2); w = 0.3
    b1 = ax.bar(x-w/2, rhos, w, label='Spearman rho', color='#457B9D', zorder=3)
    b2 = ax.bar(x+w/2, rs, w, label='Pearson r', color='#A8DADC', zorder=3)
    for b in list(b1)+list(b2):
        ax.text(b.get_x()+b.get_width()/2, b.get_height()+0.005, f'{b.get_height():.3f}',
               ha='center', va='bottom', fontsize=11, fontweight='bold')
    dr = rhos[1]-rhos[0]
    ax.annotate(f'Delta_rho = {"+" if dr>=0 else ""}{dr:.3f}', xy=(1.5, max(rhos)+0.03),
               fontsize=12, fontweight='bold', color='#E63946', ha='center')
    ax.set_ylabel('Correlation'); ax.set_title('Embed-Structure Correlation', fontweight='bold')
    ax.set_xticks(x); ax.set_xticklabels(models, fontsize=12)
    ax.legend(frameon=True, edgecolor='#ccc'); ax.grid(axis='y', alpha=0.3, ls='--')
    ax.set_ylim(0, max(max(rhos), max(rs))+0.08)
    ax.spines['top'].set_visible(False); ax.spines['right'].set_visible(False)
    plt.tight_layout()
    out = Path(out_dir)
    for fmt in ['png','pdf']:
        plt.savefig(out/f'probe7_ablation_embed_vs_struct.{fmt}', dpi=300, bbox_inches='tight', facecolor='white')
        print(f"  Saved: {out/f'probe7_ablation_embed_vs_struct.{fmt}'}")
    plt.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--compare_only', action='store_true')
    parser.add_argument('--skip_probe4', action='store_true')
    args = parser.parse_args()

    out_dir = PROBE_DIR / 'probe_7_contrastive_ablation'
    out_dir.mkdir(parents=True, exist_ok=True)

    # Load existing V5.1 results
    v51_p2_path = PROBE_DIR / 'probe_2_layerwise_cls' / 'layerwise_results_v6.json'
    v51_p4_path = PROBE_DIR / 'probe_4_embed_vs_structure' / 'embed_vs_struct_v6.json'
    v51_results = json.load(open(v51_p2_path)) if v51_p2_path.exists() else []
    v51_p4 = json.load(open(v51_p4_path)) if v51_p4_path.exists() else None
    print(f"V5.1 Probe 2: {len(v51_results)} entries")

    v5a_p2_path = out_dir / 'v5a_layerwise_results.json'
    v5a_p4_path = out_dir / 'v5a_probe4_results.json'

    if args.compare_only:
        v5a_results = json.load(open(v5a_p2_path)) if v5a_p2_path.exists() else []
        v5a_p4 = json.load(open(v5a_p4_path)) if v5a_p4_path.exists() else None
    else:
        print(f"\n{'='*60}")
        print(f"  PROBE 7a: Layer-wise CLS on V5-A (MLM-only)")
        print(f"{'='*60}")
        model = load_model(str(V5A_CKPT), device=args.device)

        print("\nLoading datasets...")
        dom_s, dom_l = load_domain_data()
        gly_s, gly_l = load_glycosylation_data()
        imm_s, imm_l = load_immunogenicity_data()

        if len(dom_s) > 3000:
            np.random.seed(42)
            idx = np.random.choice(len(dom_s), 3000, replace=False)
            dom_s = [dom_s[i] for i in idx]; dom_l = [dom_l[i] for i in idx]
            print(f"  Subsampled domain to {len(dom_s)}")

        print("\nExtracting layer-wise CLS (V5-A)...")
        dom_emb = extract_layerwise_cls(model, dom_s, device=args.device)
        gly_emb = extract_layerwise_cls(model, gly_s, device=args.device)
        imm_emb = extract_layerwise_cls(model, imm_s, device=args.device)

        print("\nRunning linear probes...")
        v5a_results = (run_layerwise_probing(dom_emb, dom_l, 'Domain', 13) +
                       run_layerwise_probing(imm_emb, imm_l, 'Immunogenicity', 13) +
                       run_layerwise_probing(gly_emb, gly_l, 'Glycosylation', 13))

        with open(v5a_p2_path, 'w') as f: json.dump(v5a_results, f, indent=2, default=str)
        csv_out = out_dir / 'v5a_layerwise_results.csv'
        with open(csv_out, 'w', newline='') as f:
            w = csv.DictWriter(f, fieldnames=['task','layer','accuracy_mean','accuracy_std','f1_mean','f1_std','n_samples','n_classes'])
            w.writeheader()
            for r in v5a_results: w.writerow({k: r[k] for k in w.fieldnames})
        print(f"  Saved: {v5a_p2_path}")

        # Probe 4
        v5a_p4 = None
        if not args.skip_probe4:
            print(f"\n{'='*60}")
            print(f"  PROBE 7b: Embed vs Structure on V5-A")
            print(f"{'='*60}")
            gs, _, iupacs = load_glyco_with_iupac()
            v5a_p4 = run_probe4(model, gs, iupacs, 'V5-A', out_dir)
            if v5a_p4:
                with open(v5a_p4_path, 'w') as f: json.dump(v5a_p4, f, indent=2)

        import gc, torch; del model; torch.cuda.empty_cache(); gc.collect()

    # Generate comparison figures
    print(f"\n{'='*60}")
    print(f"  Generating Comparison Figures")
    print(f"{'='*60}")
    if v51_results and v5a_results:
        plot_ablation(v5a_results, v51_results, out_dir)
    if not v5a_p4 and v5a_p4_path.exists():
        v5a_p4 = json.load(open(v5a_p4_path))
    plot_p4_comparison(v5a_p4, v51_p4, out_dir)

    # Summary
    print(f"\n{'='*60}")
    print(f"  ABLATION SUMMARY")
    print(f"{'='*60}")
    for task in ['Domain', 'Immunogenicity', 'Glycosylation']:
        v5a_t = [r for r in v5a_results if r['task']==task]
        v51_t = [r for r in v51_results if r['task']==task]
        if v5a_t and v51_t:
            v5a_b = max(v5a_t, key=lambda r: r['accuracy_mean'])
            v51_b = max(v51_t, key=lambda r: r['accuracy_mean'])
            d = v51_b['accuracy_mean'] - v5a_b['accuracy_mean']
            print(f"\n{task}:")
            print(f"    V5-A best: L{v5a_b['layer']} = {v5a_b['accuracy_mean']:.4f}")
            print(f"    V5.1 best: L{v51_b['layer']} = {v51_b['accuracy_mean']:.4f}")
            print(f"    Delta = {'+'if d>=0 else ''}{d:.4f} ({'+'if d>=0 else ''}{d*100:.2f}pp)")
    if v5a_p4 and v51_p4:
        dr = v51_p4['spearman_rho'] - v5a_p4['spearman_rho']
        print(f"\nEmbed-Structure: V5-A rho={v5a_p4['spearman_rho']:.4f}, V5.1 rho={v51_p4['spearman_rho']:.4f}, Delta={'+'if dr>=0 else ''}{dr:.4f}")

    summary = {'v5a_probe2': v5a_results, 'v51_probe2': v51_results, 'v5a_probe4': v5a_p4, 'v51_probe4': v51_p4}
    with open(out_dir / 'ablation_summary.json', 'w') as f: json.dump(summary, f, indent=2, default=str)
    print(f"\nCOMPLETE")


if __name__ == '__main__':
    main()