#!/usr/bin/env python3 """ Script to generate zeroshot mutations for ESM and ESM-IF. Example usage: conda activate multievolve plm_zeroshot_ensemble.py \ --wt-file apex.fasta \ --pdb-files apex.cif, apex_2.cif \ --chain-id A \ --variants 24 \ --normalizing-method aa_substitution_type \ --excluded-positions 1,14,41,112 """ import argparse from Bio import SeqIO import pandas as pd import os from model import zero_shot_esm_dms, zero_shot_esm_if_dms def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser(description='Nominate mutations using a PLM zeroshot ensemble') parser.add_argument( '--wt-file', required=True, help='Path to the wildtype FASTA file' ) parser.add_argument( '--pdb-files', required=True, help='Comma-separated list of PDB/CIF structure files' ) parser.add_argument( '--chain-id', required=True, help='Chain ID to include in the zeroshot predictions' ) parser.add_argument( '--variants', type=int, required=True, help='Number of variants to nominate per method' ) parser.add_argument( '--normalizing-method', required=True, help='Method for normalizing fold-change scores to generate z-scores' ) parser.add_argument( '--excluded-positions', required=False, help='Comma-separated list of positions to exclude from mutation' ) args = parser.parse_args() # Process arguments args.pdb_files = [f.strip() for f in args.pdb_files.split(',')] args.excluded_positions = [int(p) for p in args.excluded_positions.split(',')] if args.excluded_positions else [] return args def main(): args = parse_args() wt_file = args.wt_file pdb_files = args.pdb_files variants = args.variants excluded_positions = args.excluded_positions normalizing_method = args.normalizing_method chain_id = args.chain_id wt_seq = str(SeqIO.read(wt_file, "fasta").seq) print('Running ESM zeroshot...') esm_zeroshot = zero_shot_esm_dms(wt_seq) print('Running ESM-IF zeroshot...') esm_if_zeroshot_ls = [] for pdb_file in pdb_files: esm_if_zeroshot_ls.append(zero_shot_esm_if_dms(wt_seq, pdb_file, chain_id = chain_id, scoring_strategy='wt-marginals')) def sample_mutations(df, total_muts, excluded_positions): muts = [] pos = excluded_positions.copy() # iterate over each row for index, row in df.iterrows(): if row['pos'] not in pos: muts.append(row.to_frame().T) pos.append(row['pos']) if len(muts) == total_muts: break result = pd.concat(muts, ignore_index=True) return result def calculate_z_scores(df, col_name, activity_col): """ Calculate z-scores for activity values grouped by a column and filter groups with sufficient samples. Args: df (pd.DataFrame): Input dataframe col_name (str): Column name to group by activity_col (str): Column name containing activity values to calculate z-scores for Returns: pd.DataFrame: Dataframe with z-scores calculated and sorted, filtered to groups with >= 5 samples """ dfs = [] col_values = df[col_name].unique() for value in col_values: subset = df[df[col_name] == value].copy() subset['z_logratio'] = (subset[activity_col] - subset[activity_col].mean()) / subset[activity_col].std() subset['n'] = len(subset) if len(subset) >= 5: dfs.append(subset) df = pd.concat(dfs, ignore_index=True) df.sort_values(by='z_logratio', ascending=False, inplace=True) return df # Function to merge dataframes def merge_mutation_dfs(df_dict): # Start with the first dataframe first_key = list(df_dict.keys())[0] result = df_dict[first_key][0][['mutations', df_dict[first_key][1]]].copy() # Merge all remaining dataframes for key in list(df_dict.keys())[1:]: df = df_dict[key][0] col = df_dict[key][1] result = pd.merge(result, df[['mutations', col]], on='mutations', how='outer') return result.fillna(0) # average results for esm if zeroshot across multiple structure models subset_ls = [] for j in range(len(esm_if_zeroshot_ls)): subset_ls.append(esm_if_zeroshot_ls[j][['mutations','logratio']].copy()) subset_ls[j].rename(columns={'logratio': f'logratio_model{j}'}, inplace=True) # Start with first dataframe esm_if_zeroshot = subset_ls[0].copy() # Merge remaining dataframes iteratively for j in range(1, len(subset_ls)): esm_if_zeroshot = pd.merge(esm_if_zeroshot, subset_ls[j], on='mutations', how='outer') # Calculate average across all model logratios logratio_cols = [f'logratio_model{j}' for j in range(len(subset_ls))] esm_if_zeroshot['average_model_logratio'] = esm_if_zeroshot[logratio_cols].mean(axis=1) # sort esm_zeroshot by number of total models with FC > 1 and then by FC value esm_zeroshot_ls = [] total_model_pass_list = list(set(esm_zeroshot['total_model_pass'].values)) total_model_pass_list = total_model_pass_list[::-1] for model_pass_value in total_model_pass_list: subset = esm_zeroshot[esm_zeroshot['total_model_pass'] == model_pass_value].copy() subset.sort_values(by='average_model_logratio', ascending=False, inplace=True) esm_zeroshot_ls.append(subset) esm_zeroshot_sorted = pd.concat(esm_zeroshot_ls) # modify dataframes with columns for amino acid substitution type and remove wt esm_zeroshot_sorted['aa_mutation'] = esm_zeroshot_sorted['mutations'].apply(lambda x: x[-1]) esm_zeroshot_sorted['aa_substitution_type'] = esm_zeroshot_sorted['mutations'].apply(lambda x: f'{x[0]}-{x[-1]}') esm_zeroshot_sorted['pos'] = esm_zeroshot_sorted['mutations'].apply(lambda x: int(x[1:-1])) esm_if_zeroshot['aa_mutation'] = esm_if_zeroshot['mutations'].apply(lambda x: x[-1]) esm_if_zeroshot['aa_substitution_type'] = esm_if_zeroshot['mutations'].apply(lambda x: f'{x[0]}-{x[-1]}') esm_if_zeroshot['pos'] = esm_if_zeroshot['mutations'].apply(lambda x: int(x[1:-1])) # ESM FC muts_esm = sample_mutations(esm_zeroshot_sorted, variants, excluded_positions) muts_esm['esm_sampled'] = 1 # # ESM-IF FC esm_if_zeroshot.sort_values(by='average_model_logratio', ascending=False, inplace=True) muts_esm_if = sample_mutations(esm_if_zeroshot, variants, excluded_positions) muts_esm_if['esm_if_sampled'] = 1 # ESM Z df = esm_zeroshot_sorted.copy() activity_col = 'average_model_logratio' df = calculate_z_scores(df, normalizing_method, activity_col) muts_esm_z = sample_mutations(df, variants, excluded_positions) muts_esm_z['esm_z_sampled'] = 1 # ESM-IF Z df = esm_if_zeroshot.copy() activity_col = 'average_model_logratio' df = calculate_z_scores(df, normalizing_method, activity_col) muts_esm_if_z = sample_mutations(df, variants, excluded_positions) muts_esm_if_z['esm_if_z_sampled'] = 1 # Define dataframes to combine dfs = { 'esm': [muts_esm, 'esm_sampled'], 'esm_if': [muts_esm_if, 'esm_if_sampled'], 'esm_z': [muts_esm_z, 'esm_z_sampled'], 'esm_if_z': [muts_esm_if_z, 'esm_if_z_sampled'] } muts_combined = merge_mutation_dfs(dfs) muts_combined.to_csv(os.path.join(os.path.dirname(wt_file), 'plm_zeroshot_ensemble_nominated_mutations.csv')) if __name__ == '__main__': main()