Add decode_and_test_sequences.py

src/decode_and_test_sequences.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+"""
+Decode all 80 generated sequences and test them with HMD-AMP.
+"""
+
+import torch
+import numpy as np
+import pandas as pd
+from Bio import SeqIO
+from Bio.SeqRecord import SeqRecord
+from Bio.Seq import Seq
+import os
+from datetime import datetime
+from tqdm import tqdm
+import sys
+
+# Import the decoder
+from final_sequence_decoder import EmbeddingToSequenceConverter
+
+# Import HMD-AMP components
+sys.path.append('/home/edwardsun/flow/HMD-AMP')
+from sklearn.utils import shuffle
+import esm
+from deepforest import CascadeForestClassifier
+from src.utils import *
+
+def load_generated_embeddings():
+    """Load all generated embeddings from today."""
+    base_path = '/data2/edwardsun/generated_samples'
+    today = '20250829'
+    
+    files = [
+        f'generated_amps_best_model_no_cfg_{today}.pt',
+        f'generated_amps_best_model_weak_cfg_{today}.pt', 
+        f'generated_amps_best_model_strong_cfg_{today}.pt',
+        f'generated_amps_best_model_very_strong_cfg_{today}.pt'
+    ]
+    
+    all_embeddings = []
+    all_labels = []
+    
+    for file in files:
+        file_path = os.path.join(base_path, file)
+        if os.path.exists(file_path):
+            print(f"Loading {file}...")
+            embeddings = torch.load(file_path, map_location='cpu')
+            
+            # Extract config type from filename
+            if 'no_cfg' in file:
+                cfg_type = 'no_cfg'
+            elif 'weak_cfg' in file:
+                cfg_type = 'weak_cfg'
+            elif 'strong_cfg' in file and 'very' not in file:
+                cfg_type = 'strong_cfg'
+            elif 'very_strong_cfg' in file:
+                cfg_type = 'very_strong_cfg'
+            
+            # Each file contains 20 sequences
+            for i in range(embeddings.shape[0]):
+                all_embeddings.append(embeddings[i])
+                all_labels.append(f"{cfg_type}_{i+1}")
+                
+    print(f"✓ Loaded {len(all_embeddings)} embeddings total")
+    return all_embeddings, all_labels
+
+def decode_embeddings_to_sequences(embeddings, labels):
+    """Decode embeddings to sequences."""
+    print("Initializing sequence decoder...")
+    decoder = EmbeddingToSequenceConverter(device='cuda')
+    
+    sequences = []
+    sequence_ids = []
+    
+    print("Decoding embeddings to sequences...")
+    for i, (embedding, label) in enumerate(tqdm(zip(embeddings, labels), total=len(embeddings))):
+        # Decode using diverse method for better results
+        sequence = decoder.embedding_to_sequence(
+            embedding, 
+            method='diverse', 
+            temperature=0.8
+        )
+        sequences.append(sequence)
+        sequence_ids.append(f"generated_seq_{i+1}_{label}")
+        
+    return sequences, sequence_ids
+
+def save_sequences_as_fasta(sequences, sequence_ids, filename):
+    """Save sequences as FASTA file."""
+    records = []
+    for seq_id, seq in zip(sequence_ids, sequences):
+        record = SeqRecord(Seq(seq), id=seq_id, description="")
+        records.append(record)
+    
+    SeqIO.write(records, filename, "fasta")
+    print(f"✓ Saved {len(sequences)} sequences to {filename}")
+
+def test_with_hmd_amp(sequences, sequence_ids):
+    """Test sequences with HMD-AMP classifier."""
+    print("\n🧬 Testing sequences with HMD-AMP classifier...")
+    
+    # Set device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Load models
+    ftmodel_save_path = '/home/edwardsun/flow/HMD-AMP/AMP/ft_parts.pth'
+    clsmodel_save_path = '/home/edwardsun/flow/HMD-AMP/AMP/clsmodel'
+    
+    # Create temporary FASTA file for HMD-AMP
+    temp_fasta = 'temp_sequences.fasta'
+    save_sequences_as_fasta(sequences, sequence_ids, temp_fasta)
+    
+    try:
+        # Generate sequence features using HMD-AMP's feature extraction
+        seq_embeddings, _, seq_ids = amp_feature_extraction(ftmodel_save_path, device, temp_fasta)
+        
+        # Load classifier
+        cls_model = CascadeForestClassifier()
+        cls_model.load(clsmodel_save_path)
+        
+        # Make predictions
+        binary_pred = cls_model.predict(seq_embeddings)
+        
+        print(f"📊 HMD-AMP Results:")
+        print(f"Total sequences: {len(sequences)}")
+        print(f"Predicted AMPs: {np.sum(binary_pred)} ({np.sum(binary_pred)/len(sequences)*100:.1f}%)")
+        print(f"Predicted non-AMPs: {len(sequences) - np.sum(binary_pred)} ({(len(sequences) - np.sum(binary_pred))/len(sequences)*100:.1f}%)")
+        
+        # Analyze results by CFG type
+        results_df = pd.DataFrame({
+            'ID': sequence_ids,
+            'Sequence': sequences,
+            'AMP_Prediction': binary_pred,
+            'CFG_Type': [seq_id.split('_')[-2] for seq_id in sequence_ids]
+        })
+        
+        # Group by CFG type
+        cfg_analysis = results_df.groupby('CFG_Type')['AMP_Prediction'].agg(['count', 'sum', 'mean']).round(3)
+        cfg_analysis.columns = ['Total', 'Predicted_AMPs', 'AMP_Rate']
+        
+        print(f"\n📋 Results by CFG Configuration:")
+        print(cfg_analysis)
+        
+        # Show predicted AMPs
+        amp_results = results_df[results_df['AMP_Prediction'] == 1]
+        if len(amp_results) > 0:
+            print(f"\n🏆 Sequences predicted as AMPs ({len(amp_results)}):")
+            for idx, row in amp_results.iterrows():
+                seq = row['Sequence']
+                cationic = seq.count('K') + seq.count('R')
+                net_charge = seq.count('K') + seq.count('R') + seq.count('H') - seq.count('D') - seq.count('E')
+                print(f"  {row['ID']}: {seq}")
+                print(f"    Length: {len(seq)}, Cationic (K+R): {cationic}, Net charge: {net_charge:+d}")
+        else:
+            print(f"\n❌ No sequences predicted as AMPs")
+        
+        # Save detailed results
+        results_df.to_csv('hmd_amp_detailed_results.csv', index=False)
+        cfg_analysis.to_csv('hmd_amp_cfg_analysis.csv')
+        
+        print(f"\n💾 Results saved:")
+        print(f"  - hmd_amp_detailed_results.csv (detailed per-sequence results)")
+        print(f"  - hmd_amp_cfg_analysis.csv (summary by CFG type)")
+        
+        return results_df, cfg_analysis
+        
+    finally:
+        # Clean up temporary file
+        if os.path.exists(temp_fasta):
+            os.remove(temp_fasta)
+
+def main():
+    print("🚀 Starting sequence decoding and HMD-AMP testing...")
+    
+    # Load embeddings
+    embeddings, labels = load_generated_embeddings()
+    
+    # Decode to sequences
+    sequences, sequence_ids = decode_embeddings_to_sequences(embeddings, labels)
+    
+    # Save sequences as FASTA
+    fasta_filename = f'generated_sequences_{datetime.now().strftime("%Y%m%d_%H%M%S")}.fasta'
+    save_sequences_as_fasta(sequences, sequence_ids, fasta_filename)
+    
+    # Test with HMD-AMP
+    results_df, cfg_analysis = test_with_hmd_amp(sequences, sequence_ids)
+    
+    print(f"\n✅ Complete! Generated and tested {len(sequences)} sequences")
+    print(f"📁 Sequences saved as: {fasta_filename}")
+    
+    # Final summary
+    total_amps = results_df['AMP_Prediction'].sum()
+    print(f"\n📊 FINAL SUMMARY:")
+    print(f"Generated sequences: {len(sequences)}")
+    print(f"HMD-AMP predicted AMPs: {total_amps}/{len(sequences)} ({total_amps/len(sequences)*100:.1f}%)")
+    
+    if total_amps > 0:
+        print(f"✨ Success! Your flow model generated {total_amps} sequences that HMD-AMP classifies as AMPs!")
+    else:
+        print(f"🔍 No sequences classified as AMPs - this may indicate the need for stronger AMP conditioning.")
+
+if __name__ == "__main__":
+    main()