feat: add Docker/Apptainer support and FDR investigation tools

- Add Dockerfile for containerized deployment
- Add apptainer.def for HPC environments
- Add scripts to investigate FDR calibration data discrepancy
- Add precomputed probability verification test
- Remove ESM embedding (separate heavy dependency)

The FDR investigation checks for data leakage between calibration sets.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

apptainer.def

@@ -0,0 +1,68 @@
+Bootstrap: docker
+From: pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
+
+%labels
+    Author Ron Boger <ronboger@berkeley.edu>
+    Version 1.0
+    Description Conformal Protein Retrieval - Functional protein mining with statistical guarantees
+
+%post
+    # Update and install system dependencies
+    apt-get update && apt-get install -y \
+        git \
+        wget \
+        && rm -rf /var/lib/apt/lists/*
+
+    # Install Python dependencies
+    pip install --no-cache-dir \
+        numpy \
+        pandas \
+        scipy \
+        scikit-learn \
+        matplotlib \
+        seaborn \
+        tqdm \
+        faiss-gpu \
+        biopython \
+        pytorch-lightning \
+        h5py \
+        transformers \
+        sentencepiece \
+        gradio>=4.0.0
+
+    # Create workspace
+    mkdir -p /workspace/data /workspace/results /workspace/protein_vec_models
+
+%environment
+    export PYTHONPATH=/workspace
+    export GRADIO_SERVER_NAME=0.0.0.0
+    export GRADIO_SERVER_PORT=7860
+
+%runscript
+    echo "Conformal Protein Retrieval (CPR)"
+    echo "Usage:"
+    echo "  apptainer run cpr.sif cpr --help"
+    echo "  apptainer run cpr.sif python -m protein_conformal.gradio_app"
+    exec "$@"
+
+%help
+    Conformal Protein Retrieval (CPR)
+
+    This container provides tools for functional protein mining with
+    conformal guarantees, as described in:
+    "Functional protein mining with conformal guarantees"
+    Nature Communications (2025) 16:85
+
+    Usage:
+      # Run CLI
+      apptainer exec cpr.sif cpr embed --input seqs.fasta --output emb.npy
+      apptainer exec cpr.sif cpr search --query q.npy --database db.npy -o results.csv
+
+      # Run Gradio UI
+      apptainer exec cpr.sif python -m protein_conformal.gradio_app
+
+      # Interactive shell
+      apptainer shell cpr.sif
+
+    Build:
+      apptainer build cpr.sif apptainer.def

protein_conformal/cli.py

@@ -46,10 +46,9 @@ def cmd_embed(args):
         embeddings = _embed_protein_vec(sequences, device, args)
     elif args.model == 'clean':
         embeddings = _embed_clean(sequences, device, args)
-    elif args.model == 'esm':
-        embeddings = _embed_esm(sequences, device, args)
     else:
         print(f"Unknown model: {args.model}")
+        print("Available models: protein-vec, clean")
         sys.exit(1)
 
     print(f"Embeddings shape: {embeddings.shape}")
@@ -157,50 +156,6 @@ def _embed_clean(sequences, device, args):
     return clean_embeddings
 
 
-def _embed_esm(sequences, device, args):
-    """Embed using ESM-1b or ESM2 model."""
-    import numpy as np
-    import torch
-
-    try:
-        import esm
-    except ImportError:
-        print("Error: ESM package not installed.")
-        print("Install with: pip install fair-esm")
-        sys.exit(1)
-
-    esm_version = getattr(args, 'esm_version', '1b')
-    print(f"Loading ESM-{esm_version} model...")
-
-    if esm_version == '2':
-        model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
-    else:
-        model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
-
-    batch_converter = alphabet.get_batch_converter()
-    model = model.to(device).eval()
-
-    # Process in batches
-    batch_size = getattr(args, 'batch_size', 4)
-    embeddings = []
-
-    for i in range(0, len(sequences), batch_size):
-        batch_seqs = sequences[i:i+batch_size]
-        batch_data = [(f"seq_{j}", seq) for j, seq in enumerate(batch_seqs)]
-
-        _, _, batch_tokens = batch_converter(batch_data)
-        batch_tokens = batch_tokens.to(device)
-
-        with torch.no_grad():
-            results = model(batch_tokens, repr_layers=[33], return_contacts=False)
-            # Mean pooling over sequence length
-            batch_emb = results["representations"][33].mean(dim=1).cpu().numpy()
-            embeddings.append(batch_emb)
-
-        if (i + batch_size) % 20 == 0 or i + batch_size >= len(sequences):
-            print(f"  Processed {min(i + batch_size, len(sequences))}/{len(sequences)}")
-
-    return np.concatenate(embeddings)
 
 
 def cmd_search(args):
@@ -431,15 +386,11 @@ def main():
     p_embed.add_argument('--input', '-i', required=True, help='Input FASTA file')
     p_embed.add_argument('--output', '-o', required=True, help='Output .npy file for embeddings')
     p_embed.add_argument('--model', '-m', default='protein-vec',
-                         choices=['protein-vec', 'clean', 'esm'],
+                         choices=['protein-vec', 'clean'],
                          help='Embedding model (default: protein-vec)')
     p_embed.add_argument('--cpu', action='store_true', help='Force CPU even if GPU available')
     p_embed.add_argument('--clean-model', default='split100',
                          help='CLEAN model variant (default: split100)')
-    p_embed.add_argument('--esm-version', default='1b', choices=['1b', '2'],
-                         help='ESM version (default: 1b)')
-    p_embed.add_argument('--batch-size', type=int, default=4,
-                         help='Batch size for ESM (default: 4)')
     p_embed.set_defaults(func=cmd_embed)
 
     # search command

scripts/investigate_fdr.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+"""
+Investigate FDR calibration discrepancy between datasets.
+Checks for data leakage and compares calibration results.
+"""
+
+import numpy as np
+import sys
+sys.path.insert(0, '.')
+from protein_conformal.util import get_sims_labels, get_thresh_FDR
+
+print("=" * 60)
+print("FDR Calibration Dataset Investigation")
+print("=" * 60)
+print()
+
+# Load both calibration datasets
+print("Loading datasets...")
+pfam_new = np.load('data/pfam_new_proteins.npy', allow_pickle=True)
+backup_data = np.load('/groups/doudna/projects/ronb/conformal_backup/protein-conformal/data/conformal_pfam_with_lookup_dataset.npy', allow_pickle=True)
+
+print(f'pfam_new_proteins.npy: {len(pfam_new)} samples')
+print(f'backup dataset: {len(backup_data)} samples')
+print()
+
+# Check for overlap (potential leakage)
+print("Checking for overlap between datasets...")
+pfam_metas = set(d['meta'] for d in pfam_new)
+backup_metas = set(d['meta'] for d in backup_data)
+overlap = pfam_metas & backup_metas
+print(f"  Unique in pfam_new: {len(pfam_metas)}")
+print(f"  Unique in backup: {len(backup_metas)}")
+print(f"  Overlap: {len(overlap)} ({len(overlap)/len(pfam_metas)*100:.1f}% of pfam_new)")
+print()
+
+# Compare similarity distributions
+print("Similarity score distributions:")
+sims_new, labels_new = get_sims_labels(pfam_new[:500], partial=False)
+sims_backup, labels_backup = get_sims_labels(backup_data[:500], partial=False)
+
+print(f"  pfam_new (500 samples):")
+print(f"    Similarity: min={sims_new.min():.6f}, max={sims_new.max():.6f}, mean={sims_new.mean():.6f}")
+print(f"    Labels: {labels_new.sum()}/{labels_new.size} positive ({labels_new.mean()*100:.1f}%)")
+print()
+print(f"  backup (500 samples):")
+print(f"    Similarity: min={sims_backup.min():.6f}, max={sims_backup.max():.6f}, mean={sims_backup.mean():.6f}")
+print(f"    Labels: {labels_backup.sum()}/{labels_backup.size} positive ({labels_backup.mean()*100:.1f}%)")
+print()
+
+# Run FDR calibration on both with same parameters
+print("Running FDR calibration (alpha=0.1, n_calib=1000, 10 trials)...")
+print()
+
+def run_fdr_trials(data, name, n_trials=10, n_calib=1000):
+    lhats = []
+    risks = []
+    tprs = []
+
+    for trial in range(n_trials):
+        np.random.seed(42 + trial)
+        np.random.shuffle(data)
+        cal_data = data[:n_calib]
+        test_data = data[n_calib:n_calib+500]
+
+        X_cal, y_cal = get_sims_labels(cal_data, partial=False)
+        X_test, y_test = get_sims_labels(test_data, partial=False)
+
+        lhat, fdr_cal = get_thresh_FDR(y_cal, X_cal, alpha=0.1, delta=0.5, N=100)
+        lhats.append(lhat)
+
+        # Calculate test risk and TPR
+        preds = (X_test >= lhat).astype(int)
+        tp = np.sum((preds == 1) & (y_test == 1))
+        fp = np.sum((preds == 1) & (y_test == 0))
+        fn = np.sum((preds == 0) & (y_test == 1))
+
+        tpr = tp / (tp + fn) if (tp + fn) > 0 else 0
+        risk = fp / (fp + tp) if (fp + tp) > 0 else 0
+
+        tprs.append(tpr)
+        risks.append(risk)
+
+    print(f"{name}:")
+    print(f"  λ (threshold): {np.mean(lhats):.10f} ± {np.std(lhats):.10f}")
+    print(f"  Risk (FDR):    {np.mean(risks):.4f} ± {np.std(risks):.4f}")
+    print(f"  TPR:           {np.mean(tprs)*100:.1f}% ± {np.std(tprs)*100:.1f}%")
+    print()
+    return lhats, risks, tprs
+
+lhats_new, risks_new, tprs_new = run_fdr_trials(pfam_new.copy(), "pfam_new_proteins")
+lhats_backup, risks_backup, tprs_backup = run_fdr_trials(backup_data.copy(), "backup_dataset")
+
+print("=" * 60)
+print("CONCLUSION")
+print("=" * 60)
+if abs(np.mean(lhats_new) - np.mean(lhats_backup)) < 0.00001:
+    print("✓ Thresholds are similar - datasets likely compatible")
+else:
+    print("⚠ Thresholds differ significantly!")
+    print(f"  Difference: {abs(np.mean(lhats_new) - np.mean(lhats_backup)):.10f}")
+
+if len(overlap) > len(pfam_metas) * 0.5:
+    print("⚠ High overlap between datasets - potential data source")
+else:
+    print("✓ Low overlap - datasets appear independent")

scripts/slurm_investigate.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+#SBATCH --job-name=cpr-investigate
+#SBATCH --output=logs/cpr-investigate-%j.out
+#SBATCH --error=logs/cpr-investigate-%j.err
+#SBATCH --time=1:00:00
+#SBATCH --mem=32G
+#SBATCH --cpus-per-task=4
+
+# CPR Investigation - FDR calibration and precomputed probability verification
+set -e
+mkdir -p logs data
+
+source ~/.bashrc
+eval "$(conda shell.bash hook)"
+conda activate conformal-s
+
+cd /groups/doudna/projects/ronb/conformal-protein-retrieval
+
+echo "========================================"
+echo "CPR Investigation"
+echo "Date: $(date)"
+echo "Node: $(hostname)"
+echo "========================================"
+echo ""
+
+echo "=== Part 1: FDR Calibration Investigation ==="
+python scripts/investigate_fdr.py
+echo ""
+
+echo "=== Part 2: Precomputed Probability Verification ==="
+python scripts/test_precomputed_probs.py
+echo ""
+
+echo "========================================"
+echo "Completed: $(date)"
+echo "========================================"

scripts/test_precomputed_probs.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+"""
+Test that precomputed probability lookup gives same results as computing from scratch.
+"""
+
+import numpy as np
+import pandas as pd
+import sys
+sys.path.insert(0, '.')
+from protein_conformal.util import simplifed_venn_abers_prediction, get_sims_labels
+
+print("=" * 60)
+print("Precomputed Probability Verification")
+print("=" * 60)
+print()
+
+# Load calibration data
+print("Loading calibration data...")
+cal_data = np.load('data/pfam_new_proteins.npy', allow_pickle=True)
+np.random.seed(42)
+np.random.shuffle(cal_data)
+cal_subset = cal_data[:100]
+
+X_cal, y_cal = get_sims_labels(cal_subset, partial=False)
+X_cal = X_cal.flatten()
+y_cal = y_cal.flatten()
+print(f"  Calibration pairs: {len(X_cal)}")
+print(f"  Similarity range: [{X_cal.min():.6f}, {X_cal.max():.6f}]")
+print()
+
+# Create precomputed lookup table
+print("Creating precomputed lookup table (100 bins)...")
+min_sim, max_sim = X_cal.min(), X_cal.max()
+bins = np.linspace(min_sim, max_sim, 100)
+
+lookup = []
+for sim in bins:
+    p0, p1 = simplifed_venn_abers_prediction(X_cal, y_cal, sim)
+    lookup.append({'similarity': sim, 'p0': p0, 'p1': p1, 'prob': (p0+p1)/2})
+
+lookup_df = pd.DataFrame(lookup)
+print(f"  Lookup table: {len(lookup_df)} entries")
+print()
+
+# Test on random similarity values
+print("Testing lookup vs direct computation on 20 random values...")
+test_sims = np.random.uniform(min_sim, max_sim, 20)
+
+print(f"{'Similarity':>12} | {'Direct':>8} | {'Lookup':>8} | {'Diff':>8}")
+print("-" * 50)
+
+max_diff = 0
+for sim in test_sims:
+    # Direct computation
+    p0, p1 = simplifed_venn_abers_prediction(X_cal, y_cal, sim)
+    prob_direct = (p0 + p1) / 2
+
+    # Lookup with interpolation
+    lower = lookup_df[lookup_df['similarity'] <= sim].iloc[-1] if len(lookup_df[lookup_df['similarity'] <= sim]) > 0 else lookup_df.iloc[0]
+    upper = lookup_df[lookup_df['similarity'] >= sim].iloc[0] if len(lookup_df[lookup_df['similarity'] >= sim]) > 0 else lookup_df.iloc[-1]
+    prob_lookup = (lower['prob'] + upper['prob']) / 2
+
+    diff = abs(prob_direct - prob_lookup)
+    max_diff = max(max_diff, diff)
+    print(f"{sim:12.8f} | {prob_direct:8.4f} | {prob_lookup:8.4f} | {diff:8.4f}")
+
+print()
+print("=" * 60)
+if max_diff < 0.01:
+    print(f"✓ VERIFICATION PASSED (max diff: {max_diff:.4f})")
+    print("  Precomputed lookup matches direct computation")
+else:
+    print(f"⚠ VERIFICATION WARNING (max diff: {max_diff:.4f})")
+    print("  Consider using more bins for better accuracy")
+print("=" * 60)
+
+# Save the lookup table
+output_path = 'data/sim2prob_lookup.csv'
+lookup_df.to_csv(output_path, index=False)
+print(f"\nSaved lookup table to: {output_path}")