bioemu_example.tsv added

run_bioemu.py added
run_bioemu.sh added
bioemu run added in run_PMGen and run_utils.py

Author: AmirAsgary

README.md

@@ -1,5 +1,5 @@
 # PMGen
-## Still under development, mino bugs and issues might happen
+## Still under development, minor bugs and issues might happen
 **PMGen** (Peptide-MHC Predictive, Modeling and Generative) is a powerful and flexible framework 
 for Peptide-MHC (pMHC) complex modeling, binding prediction, and neoantigen design. It integrates 
 cutting-edge tools such as **PANDORA** for template generation, **AlphaFold (via AFfine)** for 
@@ -279,7 +279,25 @@ python run_PMGen.py \
   --peptide_random_fix_fraction 0 \
   --iterative_peptide_gen 50
 ```
-
+10. BioEmu sampling on the final iteration in  Iterative Peptide Generation.
+```bash
+python run_PMGen.py \
+  --mode wrapper \
+  --run single \
+  --df data/example/bioemu_exmple.tsv \
+  --no_pandora \
+  --output_dir outputs/bioemu_mhcreps_forpaper \
+  --bioemu_batch_size_100 30 \
+  --iterative_peptide_gen 5 \
+  --return_all_outputs \
+  --fix_anchors \
+  --peptide_random_fix_fraction 0.6 \
+  --batch_size 2 \
+  --peptide_design \
+  --binder_pred \
+  --bioemu_run_on_iter 5 \
+  --run_bioemu
+```
 ### Output
 
 Results are saved in --output_dir with the following structure:

data/example/bioemu_exmple.tsv

@@ -0,0 +1,3 @@
+peptide	mhc_seq	mhc_type	anchors	id
+QHIRCNIPKRIGPSKVATLVPR	EEHVIIQAEFYLNPDQSGEFMFDFDGDEIFHVDMAKKETVWRLEEFGRFASFEAQGALANIAVDKANLEIMTKRSNY/TRPRFLELLKSECHFFNGTERVRFLERYFHNQEEFVRFDSDVGEYRAVTELGRPVAESWNSQKDLLEQKRGQVDTYCRHNYGVVESFTVQRRVH	2		4H25
+FLNKDLEVDGHFVTM	GSHSMRYFFTSVSRPGRGEPRFIAVGYVDDTQFVRFDSDAASQRMEPRAPWIEQEGPEYWDGETRKVKAHSQTHRVDLGTLRGYYNQSEAGSHTVQRMYGCDVGSDWRFLRGYHQYAYDGKDYIALKEDLRSWTAADMAAQTTKHKWEAAHVAEQLRAYLEGTCVEWLRRYLENGKETLQRTDFLNKDLEVDGHFVTM	1		4U6Y

run_PMGen.py

@@ -1,9 +1,10 @@
 import argparse
 import pandas as pd
-from run_utils import run_PMGen_wrapper, run_PMGen_modeling, protein_mpnn_wrapper, MultipleAnchors, get_best_structres, retrieve_anchors_and_fixed_positions, assert_iterative_mode, collect_generated_binders, create_new_input_and_fixed_positions
+from run_utils import run_PMGen_wrapper, run_PMGen_modeling, protein_mpnn_wrapper, bioemu_assertions, MultipleAnchors, get_best_structres, retrieve_anchors_and_fixed_positions, assert_iterative_mode, collect_generated_binders, create_new_input_and_fixed_positions
 import shutil
 from Bio import SeqIO
 import warnings
+import subprocess
 import os
 from Bio import BiopythonDeprecationWarning
 warnings.filterwarnings("ignore", category=BiopythonDeprecationWarning)
@@ -74,13 +75,25 @@ def main():
     parser.add_argument('--mhc_design', action='store_true', help='Enables whole mhc design. we recommend to use only_pseudo_sequence_design mode.')
     parser.add_argument('--num_sequences_peptide', type=int, default=10, help='Number of peptide sequences to be generated. Works only with --peptide_design')
     parser.add_argument('--num_sequences_mhc', type=int, default=10, help='Number of mhc sequences to be generated. Works only with --only_pseudo_sequence_design or --mhc_design')
-    parser.add_argument('--sampling_temp', type=float, default=0.1, help='ProteinMPNN sampling temperature.')
+    parser.add_argument('--sampling_temp', type=float, default=1.5, help='ProteinMPNN sampling temperature.')
     parser.add_argument('--batch_size', type=int, default=1, help='ProteinMPNN batch size.')
     parser.add_argument('--hot_spot_thr', type=float, default=6.0, help='Distance threshold to peptide, to define hot-spots on mhc.')
     parser.add_argument('--binder_pred', action='store_true', help='Enables binder prediction from ProteinMPNN generated peptides.')
     parser.add_argument("--fix_anchors", action='store_true', help='If set, does not design anchor positions in peptide generation')
     parser.add_argument("--peptide_random_fix_fraction", type=float, default=0., help="Disables design for a random fraction of amino acids in peptide")
 
+    # BioEmu Argumetns
+    parser.add_argument('--run_bioemu', action='store_true', help='Enables bioemu pMHC sampling.')
+    parser.add_argument('--bioemu_num_samples', type=int, default=100, help='Sampling rounds in bioemu. You might get lower number of structures'
+                                                                           'if --filter_samples is active')
+    parser.add_argument('--bioemu_batch_size_100', type=int, default=10, help='Batch size you use for a sequence of length 100. The batch size '
+                                                                              'will be calculated from this, assuming that the memory requirement to compute '
+                                                                              'each sample scales quadratically with the sequence length.')
+    parser.add_argument('--bioemu_filter_samples', action='store_true', help='Filter out unphysical samples with e.g. long bond distances or steric clashes.')
+    parser.add_argument('--bioemu_run_on_iter', type=int, default=None, help='Optional, only works when iterative_peptide_gen > 0. Runs bioemu on the structure taken'
+                                                                             'from the iteration number given by user. If not set, runs on the 0 iteration by default.')
+
+
     # Setting to Run only a part:
     parser.add_argument('--no_alphafold', action='store_false', help='does not run alphafold.')
     parser.add_argument('--only_protein_mpnn', action='store_true', help='Skips PANDORA and AF modeling, and runs ProteinMPNN for already available predictions.')
@@ -92,6 +105,7 @@ def main():
     parser.add_argument('--iterative_peptide_gen', type=int, default=0, help='If used, the iterative peptide generation is performed, defines the number of iterations.')
 
     args = parser.parse_args()
+    bioemu_assertions(args)
     for iteration in range(args.iterative_peptide_gen + 1):
         fixed_positions_path = None #only for iter > 0 in iteration mode
         # if we have entered the iterative generation mode
@@ -226,5 +240,44 @@ def main():
                 print("Collecting the best binders")
                 collected_generated_binders_path = collect_generated_binders(args, df, iteration)
 
+    if args.run_bioemu:
+        print('**BioEmu runs initiating**')
+
+        output_dir = args.output_dir
+        bioemu_input_df_path = args.df
+
+        if args.iterative_peptide_gen > 0: #iterative mode --> which iteration to run bioemu on? only one can be used.
+            bioemu_run_on_iter = 0
+            if args.bioemu_run_on_iter:
+                assert args.bioemu_run_on_iter <= args.iterative_peptide_gen, f'Please make sure --iterative_peptide_gen is less or equal to  --iterative_peptide_gen'
+                bioemu_run_on_iter = args.bioemu_run_on_iter
+            output_dir = os.path.join("/".join(args.output_dir.split("/")[:-1]), f"iter_{bioemu_run_on_iter}") # outputdir/
+            bioemu_input_df_path = os.path.join(output_dir, f'input_df_{bioemu_run_on_iter}.tsv')
+
+        bioemu_output_dir = os.path.join(output_dir, 'bioemu')
+        cache_embeds_dir = os.path.join(output_dir, 'alphafold')
+        os.makedirs(bioemu_output_dir, exist_ok=True)
+        assert os.path.isdir(cache_embeds_dir), f'{cache_embeds_dir} not found'
+        assert os.path.isdir(bioemu_output_dir), f'{bioemu_output_dir} not found'
+        assert os.path.isfile(bioemu_input_df_path), f'{bioemu_input_df_path} not found'
+        print(f'Running on cache_embeds_dir: {cache_embeds_dir}')
+
+        bioemu_df = pd.read_csv(bioemu_input_df_path, sep='\t')
+        for i, row in bioemu_df.iterrows():
+            sequence = str(row['mhc_seq'].replace('/','') + row['peptide'])
+            id = str(row['id'])
+            bioemu_output_dir_id = os.path.join(bioemu_output_dir, id)
+            cmd = ['./run_bioemu.sh', '--sequence', sequence, '--id', id, '--output_dir', bioemu_output_dir_id,
+                   '--cache_embeds_dir', cache_embeds_dir, '--bioemu_num_samples', args.bioemu_num_samples,
+                   '--bioemu_batch_size_100', args.bioemu_batch_size_100, '--bioemu_filter_samples', args.bioemu_filter_samples]
+            print(f"🌀 Starting iteration {i}: {' '.join(cmd)}")
+            try:
+                result = subprocess.run(cmd, check=True)
+                print(f"✅ Bioemurun-> Iteration {i}, id {id}, outpath {bioemu_output_dir_id} completed successfully\n")
+            except subprocess.CalledProcessError as e:
+                print(f"❌ Bioemurun-> Iteration {i}, id {id}, outpath {bioemu_output_dir_id} failed with error code {e.returncode}")
+                # Optional: stop if a command fails
+
+
 if __name__ == "__main__":
     main()

run_bioemu.py

@@ -0,0 +1,26 @@
+# Run on bioemu env
+from bioemu.sample import main as sample
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="BioEmu sampling script")
+
+    parser.add_argument('--sequence', type='str', required=True, help='MHC+peptide sequence with no separator. For MHC-II, Alpha+Beta+peptide')
+    parser.add_argument('--id', type='str', required=True, help='PMGen id given to sequence input.')
+    parser.add_argument('--output_dir', type='str', required=True, help='Output directory to save the results')
+    parser.add_argument('--cache_embeds_dir', required=True, tupe='str', help='Path to alphafold folder containing ids/*_representations.pkl')
+    parser.add_argument( '--bioemu_num_samples', type=int, default=10, help='Sampling rounds in BioEmu. You might get a lower number of structures if --bioemu_filter_samples is active.')
+    parser.add_argument('--bioemu_batch_size_100', type=int, default=10, help='Batch size for a sequence of length 100. Actual batch size '                                                                       'scales with the square of sequence length.')
+    parser.add_argument('--bioemu_filter_samples', action='store_true', help='Filter out unphysical samples (e.g., long bond distances or steric clashes).')
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    sample(sequence=args.sequence,
+            id=args.id,
+            num_samples=args.bioemu_num_samples,
+            output_dir=args.output_dir,
+            batch_size_100=args.bioemu_batch_size_100,
+            cache_embeds_dir=args.cache_embeds_dir,
+            filter_samples=args.bioemu_filter_samples)

run_bioemu.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+# Name of the environment to activate
+ENV_NAME="bioemu"
+
+# Detect available environment manager
+if command -v conda &> /dev/null; then
+    TOOL="conda"
+    eval "$($TOOL shell.bash hook)"
+elif command -v mamba &> /dev/null; then
+    TOOL="mamba"
+    eval "$($TOOL shell.bash hook)"
+elif command -v micromamba &> /dev/null; then
+    TOOL="micromamba"
+    eval "$($TOOL shell.bash hook)"
+else
+    echo "❌ Error: conda, mamba, or micromamba not found in PATH."
+    exit 1
+fi
+
+echo "✅ Using $TOOL to activate environment '$ENV_NAME'..."
+$TOOL activate "$ENV_NAME"
+
+# Now forward all arguments to the Python script
+python run_bioemu.py "$@"

run_utils.py

@@ -585,7 +585,7 @@ def __init__(self, PMGen_pdb, output_dir,
                  num_sequences_peptide=10, num_sequences_mhc=3,
                 peptide_chain='P', mhc_design=True, peptide_design=True,
                  only_pseudo_sequence_design=True, anchor_pred=True,
-                 sampling_temp=0.05, batch_size=1, hot_spot_thr=6.0,
+                 sampling_temp=1.5, batch_size=1, hot_spot_thr=6.0,
                  save_hotspots=True, binder_pred=False, fix_anchors=False,
                  anchor_and_peptide=None):
         '''
@@ -718,7 +718,6 @@ def __peptide_design(self):
             "--save_probs", "1",
             "--save_score", "1",
             "--omit_AAs", "X",
-            "--sampling_temp", "1.5"
         ]
         if self.fix_anchors:# to fix anchors, fixed_pdbs file and design_only_positions should be generated
             # we have anchors, we need to define designable positions which are non-anchor positions
@@ -1267,6 +1266,13 @@ def create_new_input_and_fixed_positions(args, best_generated_peptides_path, ite
 
 
 
+def bioemu_assertions(args):
+    if args.run_bioemu:
+        args.return_all_outputs = True  # make sure AF2 outputs are returned
+        assert args.mode == 'wrapper', f'Bioemu is availabe only in wrapper mode: use --mode wrapper'
+        if args.iterative_peptide_gen > 0: #iterative mode --> which iteration to run bioemu on? only one can be used.
+            if args.bioemu_run_on_iter:
+                assert args.bioemu_run_on_iter <= args.iterative_peptide_gen, f'Please make sure --iterative_peptide_gen is less or equal to  --iterative_peptide_gen'