directed_evolution.py

import json
import logging
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import random
from tqdm import tqdm
from sampler import Sampler
from model.esmfold import ESMFold
from model.ensemble import Ensemble
from oracle.wetlab import WetLab_Landscape
from utils.seq_utils import hamming_distance, mutation_alphabet, random_mutate, sequence_to_onehot
from model.surrogate import SurrogateModel

logger = logging.getLogger('directed_evolution')

def run(args):
    wt_seq = json.load(open(f'./oracle_data/{args.task.lower()}_wt.json'))
    wt_raw = json.load(open(f'./oracle_data/{args.task.lower()}_wt_raw.json'))
    if args.use_structure:
        esmfold_model = ESMFold(args.task, wt_raw, args.mut_pos, args.pdb_cache_dir, 'cuda', 128)

    logger.info(args)

    surrogate_model_ensemble = Ensemble([SurrogateModel(args, len(mutation_alphabet), args.embed_dim, args.hidden_dim, len(wt_seq)).to(args.device) for i in range(args.ensemble)], 'ucb')
    
    data_store = dict({'step':[], 'sequence': [], 'fitness': [], 'distance': []})

    oracle = WetLab_Landscape(args, wt_seq)

    wt_fitness = oracle.get_fitness([wt_seq])[0][0]
    data_store['step'].append(0)
    data_store['sequence'].append(wt_seq)
    data_store['fitness'].append(wt_fitness)
    data_store['distance'].append(0)
    logger.info(f'wt fitness: {wt_fitness}')

    sample_model = Sampler(wt_fitness, args, surrogate_model_ensemble)

    seq_scores = dict({wt_seq: wt_fitness})
    max_fitness = wt_fitness
    max_seq = wt_seq
    seed_seqs = [wt_seq]
    all_seqs = set([wt_seq])


    for i in range(1, args.evo_steps + 1):
        logger.info("===EVO STEP {}===".format(i))
        
        samples = set(seed_seqs)
        # random samples for first round
        while i == 1 and len(samples) < args.max_oracle_call_per_step:
            for seq in seed_seqs:
                mutant = random_mutate(seq, mutation_alphabet, -1)
                if mutant not in all_seqs and oracle.is_valid_seq(mutant):
                    samples.add(mutant)
                    if len(samples) >= args.max_oracle_call_per_step:
                        break
        
        samples = list(samples)[:args.max_oracle_call_per_step]

        samples_fitness = []
        samples_seq_scores = dict()
        prev_max_fitness = max_fitness
        seed_seqs = [max_seq]
        for bi in tqdm(range(0, len(samples), args.oracle_batch_size)):
            sample_batch = samples[bi:bi+args.oracle_batch_size]
            fitness = oracle.get_fitness(sample_batch)[0]
            samples_fitness += fitness

            batch_max = np.max(fitness)
            if  batch_max > max_fitness:
                max_fitness = batch_max
                max_seq = sample_batch[np.argmax(fitness)]

            for si in range(len(sample_batch)):
                seq_scores[sample_batch[si]] = fitness[si]
                samples_seq_scores[sample_batch[si]] = fitness[si]
                if fitness[si] > prev_max_fitness:
                    seed_seqs.append(sample_batch[si])

        all_seqs.update(samples)
        logger.info('{} sampled sequences fitness: max: {:.4f}, min: {:.4f}, avg: {:.4f}'.format(len(samples_fitness), np.max(samples_fitness), np.min(samples_fitness), np.mean(samples_fitness)))

        logger.info('global measured sequences count: {}'.format(len(all_seqs)))
        logger.info('global max fitness: {:.4f}, distance to wt: {}'.format(max_fitness, hamming_distance(wt_seq, max_seq)))

        data_store['step'] += [i for j in range(len(samples))]
        data_store['sequence'] += samples
        data_store['fitness'] += samples_fitness
        data_store['distance'] += [hamming_distance(wt_seq, seq) for seq in samples]

        if i == args.evo_steps:
            # save model parameters
            for mi, surrogate_model in enumerate(surrogate_model_ensemble.models):
                torch.save(surrogate_model.state_dict(), f'results/{args.task}/{args.exp_name}_model_{mi}.ckpt')
            return data_store

        all_seq_scores = list(seq_scores.items())
        sample_seqs_scores = list(samples_seq_scores.items())
        if args.use_structure:
            sample_seq_rmsd = esmfold_model.calculate_rmsd([ss[0] for ss in sample_seqs_scores])
            sample_seq_rmsd = dict([[ss[0], sample_seq_rmsd[i]] for i, ss in enumerate(sample_seqs_scores)])
            for mi, surrogate_model in enumerate(surrogate_model_ensemble.models):
                optimizer = optim.Adam(params=surrogate_model.parameters(), lr=args.lr)
                loss_fn = nn.MSELoss()
                surrogate_model.train()
                logger.info(f"training structure model {mi}:")
                epoch_losses = []
                patience = 0
                min_loss = torch.inf
                epoch = 1
                while patience < args.patience:
                    losses = []
                    random.shuffle(sample_seqs_scores)
                    for bi in range(0, len(sample_seqs_scores), args.batch_size):
                        optimizer.zero_grad()
                        batch_seq = [d[0] for d in sample_seqs_scores[bi:bi+args.batch_size]]
                        structure_labels = torch.tensor([sample_seq_rmsd[seq] for seq in batch_seq], dtype=torch.float32).to(args.device)
                        input_onehot = torch.stack([sequence_to_onehot(data, mutation_alphabet) for data in batch_seq], dim=0).float().to(args.device)
                        with torch.autograd.set_detect_anomaly(True):
                            struct_output = surrogate_model.struct_forward(input_onehot) # bsz x seq_len
                            loss = loss_fn(struct_output, structure_labels) 
                            losses.append(loss.item())
                            loss.backward()
                            optimizer.step()
                    
                    epoch_loss = np.mean(losses)
                    epoch_loss_std = np.std(losses)
                    epoch_losses.append(epoch_loss)
                    if epoch == 1:
                        logger.info(f'epoch {epoch}\tloss: {epoch_losses[-1]:.4f}, std: {epoch_loss_std:.4f}')
                    epoch += 1
                    if epoch_loss < min_loss:
                        min_loss = epoch_loss
                        patience = 0
                    else:
                        patience += 1
                    
                    if patience >= args.patience:
                        logger.info(f'epoch {epoch}\tloss: {epoch_losses[-1]:.4f}, std: {epoch_loss_std:.4f}')
                
                surrogate_model.eval()

        # training fitness surrogate
        for mi, surrogate_model in enumerate(surrogate_model_ensemble.models):
            optimizer = optim.Adam(params=surrogate_model.parameters(), lr=args.lr)
            loss_fn = nn.MSELoss()
            surrogate_model.train()
            surrogate_model.struct_encoder.eval()
            
            logger.info(f"fine-tuning fitness model {mi}:")
            epoch_losses = []
            patience = 0
            min_loss = torch.inf
            epoch = 1
            while patience < args.patience:
                losses = []
                random.shuffle(all_seq_scores)
                for bi in range(0, len(all_seq_scores), args.batch_size):
                    optimizer.zero_grad()
                    batch_data = all_seq_scores[bi:bi+args.batch_size]
                    labels = torch.tensor([data[1] for data in batch_data], dtype=torch.float32).to(args.device)
                    input_onehot = torch.stack([sequence_to_onehot(data[0], mutation_alphabet) for data in batch_data], dim=0).float().to(args.device)

                    with torch.autograd.set_detect_anomaly(True):
                        output = surrogate_model.forward(input_onehot, False)
                        loss = loss_fn(output, labels)
                        losses.append(loss.item())
                        loss.backward()
                        optimizer.step()
                
                epoch_loss = np.mean(losses)
                epoch_loss_std = np.std(losses)
                epoch_losses.append(epoch_loss)
                if epoch == 1:
                    logger.info(f'epoch {epoch}\tloss: {epoch_losses[-1]:.4f}, std: {epoch_loss_std:.4f}')
                epoch += 1
                if epoch_loss < min_loss:
                    min_loss = epoch_loss
                    patience = 0
                else:
                    patience += 1
                
                if patience >= args.patience:
                    logger.info(f'epoch {epoch}\tloss: {epoch_losses[-1]:.4f}, std: {epoch_loss_std:.4f}')
            
            surrogate_model.eval()
        
        # test surrogate model performance on collected samples
        pred_scores = [[], [], []]
        for mi, surrogate_model in enumerate(surrogate_model_ensemble.models):
            surrogate_model.eval()
            for bi in range(0, len(all_seq_scores), args.batch_size):
                batch_data = all_seq_scores[bi:bi+args.batch_size]
                labels = torch.tensor([data[1] for data in batch_data], dtype=torch.float32).to(args.device)
                input_onehot = torch.stack([sequence_to_onehot(data[0], mutation_alphabet) for data in batch_data], dim=0).float().to(args.device)
                output = surrogate_model.forward(input_onehot)
                pred_scores[mi] += output.tolist()
        
        aes = []
        for i in range(len(all_seq_scores)):
            aes.append(sum([abs(pred_scores[m][i] - all_seq_scores[i][1]) for m in range(args.ensemble)])/args.ensemble)
        logger.info(f'uncertainty: {np.mean(aes)}')

        seed_seqs = [max_seq]
        if args.sampler == 'lmc':
            logger.info(f'lmc sampling:')
            esm_samples = sample_model.lmc_surrogate_sample(seed_seqs, all_seqs, set(oracle.wetlab_data.keys()))
        elif args.sampler == 'hmc':
            logger.info(f'hmc sampling:')
            esm_samples = sample_model.hmc_surrogate_sample(seed_seqs, all_seqs, set(oracle.wetlab_data.keys()))
        elif args.sampler == 'random':
            logger.info(f'random sampling:')
            esm_samples = sample_model.random_sample(seed_seqs, all_seqs, set(oracle.wetlab_data.keys()))
        else:
            assert args.sampler in ['lmc', 'hmc', 'random']
        
        seed_seqs = list(esm_samples)

    return data_store