Update demo.py

AntBO/demo.py

@@ -1,3 +1,5 @@
+import argparse
+
 import os
 import sys
 from pathlib import Path
@@ -14,59 +16,68 @@
 
 warnings.simplefilter('ignore', NumericalWarning)
 
-# TOFILL
-save_path = './results/'
-absolut_path = "" #TOFILL '/path/to/executable/Absolut/'
-antigen = "2DD8_S"
-n_init = 20
-max_iters = 40
-batch_size = 5
-device = "cuda:0"
-pre_evals_path="" #TOFILL with absolut path to csv of evaluated points "/path/to/evaluated_points.csv"
+if __name__ == "__main__":
 
-# -------- Create config
+    parser = argparse.ArgumentParser(add_help=True, description='Run AntBO')
+    parser.add_argument('--antigen', type=str, required=True, help='Name of the target antigen (e.g. 2DD8_S)')
+    parser.add_argument('--seq_len', type=int, required=True, help='Longer of the optimized antibody sequence')
+    parser.add_argument('--batch_size', type=int, default=1,
+                        help='Number of antibodies suggested at each step (default: 1)')
+    parser.add_argument('--pre_evals_csv', type=str,
+                        help='Path to csv file containing the binding energy of already evaluated antibody sequences.')
+    parser.add_argument('--cuda_id', type=int, default=0, help='ID of the cuda device to use.')
+    parser.add_argument('--seed', type=int, nargs="+", default=[42], help='Seed for reproducibility.')
+    parser.add_argument('--absolut_path', type=str, help='Path to Absolut! (if Absolut is needed.)')
+    args = parser.parse_args()
 
-config = {
-    'acq': 'ei',
-    'ard': True,
-    'n_init': n_init,
-    'max_iters': max_iters,
-    'min_cuda': 10,
-    'device': device,
-    'seq_len': 11,
-    'normalise': True,
-    'batch_size': batch_size,
-    'save_path': save_path,
-    'kernel_type': 'transformed_overlap',
-    'noise_variance': '1e-6',
-    'search_strategy': 'local',
-    'resume': False,
-    # 'bbox': {
-    #     'tool': 'Absolut',
-    #     'path': absolut_path,
-    #     'process': 4,
-    #     'startTask': 0,
-    #     'antigen': antigen
-    # },
-    'bbox': {
-        'tool': 'manual',
-        'antigen': antigen
-    },
-    'pre_evals': pre_evals_path,
-}
+    # TOFILL
+    save_path = './results/'
+    n_init = 20
+    max_iters = 400
+    device = f"cuda:{args.cuda_id}"
 
-if __name__ == "__main__":
-    for seed in [42]:
+    # -------- Create config
+
+    config = {
+        'pre_evals': args.pre_evals_csv,
+        'acq': 'ei',
+        'ard': True,
+        'n_init': n_init,
+        'max_iters': max_iters,
+        'min_cuda': 10,
+        'device': device,
+        'seq_len': args.seq_len,
+        'normalise': True,
+        'batch_size': args.batch_size,
+        'save_path': save_path,
+        'kernel_type': 'transformed_overlap',
+        'noise_variance': '1e-6',
+        'search_strategy': 'local',
+        'resume': False,
+        # 'bbox': {
+        #     'tool': 'Absolut',
+        #     'path': args.absolut_path,
+        #     'process': 4,
+        #     'startTask': 0,
+        #     'antigen': args.antigen
+        # },
+        'bbox': {
+            'tool': 'manual',
+            'antigen': args.antigen
+        },
+    }
+
+    for seed in args.seed:
         start_antigen = time.time()
         boexp = BOExperiments(config=config, cdr_constraints=True, seed=seed)
         boexp.run()
         end_antigen = time.time()
-        print(f"Time taken for antigen {antigen} = {end_antigen - start_antigen:.1f}s")
+        print(f"Time taken for antigen {args.antigen} = {end_antigen - start_antigen:.1f}s")
 
         result_dir = boexp.path + f"/results.csv"
         results = pd.read_csv(result_dir, index_col=0)
         results.head()
 
         print(
-            f"Best binder for target antigen {antigen}: {results.iloc[-1].BestProtein} "
+            f"Best binder for target antigen {args.antigen}: {results.iloc[-1].BestProtein} "
             f"with binding energy {results.iloc[-1].BestValue:.1f}")