HADES/main.py at main · GENTEL-lab/HADES · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import torch
import numpy as np
import logging
import random
import argparse
import time
import os
import pandas as pd
import directed_evolution as de

def setup_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
     torch.backends.cudnn.deterministic = True

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", default=2024, type=int)
    parser.add_argument("--task", default="GB1", type=str, choices=["GB1", "PhoQ"])
    parser.add_argument("--mut_pos", type=int, nargs=4, default=[38, 39, 40, 53], help="mutation positions (0-indexed)")
    parser.add_argument("--evo_steps", default=10, type=int, help="number of directed evolution rounds")
    parser.add_argument("--max_oracle_call_per_step", default=100, type=int)
    parser.add_argument("--max_sample_steps", default=100, type=int)
    parser.add_argument("--internal_steps", default=100, type=int)
    parser.add_argument("--device", default="cuda", type=str, choices=["cpu", "cuda"])
    parser.add_argument("--batch_size", default=64, type=int)
    parser.add_argument("--oracle_batch_size", default=100, type=int)
    parser.add_argument("--parallel_samples", default=128, type=int)
    parser.add_argument("--lr", default=0.001, type=float)
    parser.add_argument("--oracle", default="wetlab", type=str, choices=["wetlab"])
    parser.add_argument("--dropout", default=0.1, type=float)
    parser.add_argument("--embed_dim", default=256, type=int)
    parser.add_argument("--hidden_dim", default=256, type=int)
    parser.add_argument("--patience", default=3, type=int)
    parser.add_argument("--eta", default=0.1, type=float)
    parser.add_argument("--ensemble", default=3, type=int)
    parser.add_argument("--repeat_runs", default=10, type=int)
    parser.add_argument("--use_structure", action='store_true')
    parser.add_argument('--sampler', default='hmc', type=str, choices=['lmc', 'hmc', 'random'])
    parser.add_argument('--exp_name', default='default', type=str)
    parser.add_argument('--virtual_barrier', default=1, type=int, choices=[0,1])
    parser.add_argument('--pdb_cache_dir', default='pdb_cache', type=str)
    args = parser.parse_args()

    setup_seed(args.seed)

    tm = time.localtime()

    log_dir = f"log/{tm.tm_year}/{tm.tm_mon}/{tm.tm_mday}"
    os.makedirs(log_dir, exist_ok=True)
    log_filename = f"{log_dir}/log-{tm.tm_hour}-{tm.tm_min}-{tm.tm_sec}-{args.exp_name}.txt"
    logging.basicConfig(filename=log_filename, format='%(asctime)s-%(levelname)s-%(filename)s-%(lineno)d-%(message)s', level=logging.INFO)
    logger = logging.getLogger('main')

    logger.info("START")
    try:
        os.makedirs(f'results/{args.task}', exist_ok=True)
        data_filename = f'results/HADES-{args.exp_name}_{args.task}_results_wetlab.xlsx'
        with pd.ExcelWriter(data_filename, engine='xlsxwriter') as writer:
            for i in range(1, args.repeat_runs + 1):
                results = de.run(args)
                df = pd.DataFrame(results)
                df.to_excel(writer, sheet_name=f'{i}', index=False)
    except Exception as ex:
        logger.exception(ex)