all_sequences_analysis.py

import pandas as pd
import numpy as np
from PepBD_surrogate import ScoreModel  # Ensure ScoreModel is accessible in your PYTHONPATH
from tqdm import tqdm
import os
import argparse

def load_pepbd_sequences():
    """Load all PepBD sequences from different plastic datasets."""
    pepbd_sequences = []
    plastics = ['PET', 'PE', 'PP', 'PVC', 'Nylon']
    for plastic in plastics:
        df = pd.read_csv(f'Data/{plastic}.csv')
        pepbd_sequences.extend(df['Sequence'].tolist())
    return list(set(pepbd_sequences))  # Remove duplicates

def load_generated_sequences(path):
    """Load generated sequences from the simulated annealing results."""
    generated_path = path
    #path = 'simulated_annealing/Results/PP-PET/initial_temperature-2.0_cooling_rate-0.99_temperature_min-0.1_max_iterations-21525_substitution_prob-0.75_swap_prob-0.25_cooling_schedule-exponential_embedding_type-one_hot_plastic_type-PP-PET/unique_sequences_analysis.csv'
    generated = pd.read_csv(generated_path)
    return generated['Generated_Sequence'].tolist()

def load_score_models(task_type):
    """Load score models for each plastic type based on task type."""
    base_params = {
        'model_type': 'lstm',
        'epochs': 500, 
        'learning_rate': 0.001,
        'batch_size': 32,
        'device': 'cuda:0'
    }
    
    if task_type == 'all_plastics':
        n_samples_dict = {
            'PET': 353581,
            'PE': 572406,
            'PP': 346789,
            'PVC': 166886,
            'Nylon': 114091
        }
    elif task_type == 'PP_PET':
        n_samples_dict = {
            'PET': 353581,
            'PP': 346789,
        }
    else:
        raise ValueError("task_type must be either 'all_plastics' or 'PP_PET'")
    
    score_models = []
    for plastic, sample_size in n_samples_dict.items():
        model = ScoreModel(
            session_name=f'exp_onehot_{plastic}_l_2_h_512',
            n_samples=sample_size,
            **base_params
        )
        model_dir = (
            f'ScoreModel_Results_onehot/exp_onehot_lstm_{plastic}_l_2_h_512/'
            f'model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-{sample_size}_'
            'early_stopping_patience-10/Trained_Models'
        )
        model.load_model(model_dir=model_dir)
        score_models.append((plastic, model))
    return score_models

def peptide_to_one_hot(peptide):
    """Convert peptide sequence to one-hot encoding."""
    amino_acids = 'ADEFGHIKLMNQRSTVWY'
    aa_to_index = {aa: idx for idx, aa in enumerate(amino_acids)}
    one_hot = np.zeros((12, 18))
    for i, aa in enumerate(peptide):
        idx = aa_to_index.get(aa)
        if idx is not None:
            one_hot[i, idx] = 1
        else:
            raise ValueError(f'Invalid amino acid {aa} in peptide {peptide}')
    return one_hot[np.newaxis, :, :]

def get_scores(sequences, score_models, task_type):
    """Calculate scores for a list of sequences using the provided score models."""
    scores_list = []
    for i in tqdm(range(len(sequences)), desc="Scoring sequences"):
        seq = sequences[i]
        seq_scores = []
        one_hot = peptide_to_one_hot(seq)
        for _, model in score_models:
            score = model.predict(one_hot)[0]
            seq_scores.append(score)
        
        # Create score dictionary based on task type
        score_dict = {'Sequence': seq}
        
        if task_type == 'all_plastics':
            score_dict.update({
                'PET_Score': seq_scores[0],
                'PE_Score': seq_scores[1],
                'PP_Score': seq_scores[2],
                'PVC_Score': seq_scores[3],
                'Nylon_Score': seq_scores[4],
                'Average_Score': np.mean(seq_scores)
            })
        elif task_type == 'PP_PET':
            score_dict.update({
                'PET_Score': seq_scores[0],
                'PP_Score': seq_scores[1],
            })
        
        scores_list.append(score_dict)
    return pd.DataFrame(scores_list)

def main():
    """Main function to execute the scoring process."""
    parser = argparse.ArgumentParser(description='Score peptide sequences using trained LSTM models')
    parser.add_argument('--task_type', type=str, choices=['all_plastics', 'PP_PET'], 
                       required=True, help='Type of scoring task to perform')
    parser.add_argument('--sequence_type', type=str, choices=['pepbd', 'generated'], 
                       required=True, help='Type of sequences to score')
    parser.add_argument('--path', type=str, required=True, help='Path to the generated sequences')
    
    args = parser.parse_args()
    
    # Load sequences based on sequence type
    if args.sequence_type == 'pepbd':
        sequences = load_pepbd_sequences()
        output_dir = f'All_sequences_analysis'
        output_file = f'pepbd_scores.csv'
    elif args.sequence_type == 'generated':
        if args.task_type == 'PP_PET':
            sequences = load_generated_sequences(args.path)
            output_dir = f'All_sequences_analysis'
            output_file = f'PP_PET_scores.csv'
        elif args.task_type == 'all_plastics':
            sequences = load_generated_sequences(args.path)
            output_dir = f'All_sequences_analysis'
            output_file = f'generated_scores.csv'
        else:
            raise ValueError("task_type must be either 'all_plastics' or 'PP_PET'")
    else:
        raise ValueError("sequence_type must be either 'pepbd' or 'generated'")
    
    # Load models
    score_models = load_score_models(args.task_type)
    
    # Get scores
    scores_df = get_scores(sequences, score_models, args.task_type)
    
    # Create output directory and save results
    os.makedirs(output_dir, exist_ok=True)
    scores_df.to_csv(f'{output_dir}/{output_file}', index=False)
    print(f"Scoring completed and results saved to {output_dir}/{output_file}")

if __name__ == "__main__":
    main()