<jupyter_start><jupyter_text>Adeno-Associated Virus Capsid Design>Using MRL to design AAV capsids  Adeno-Associated Virus Capsid DesignThis tutorial runs an end to end workflow for designing AAV capsidsAdeno-Associated Viruses (AAV) are small viruses that generally do not cause disease. The non-pathogenic nature of these viruses makes them an attractive target for gene therapy delivey vectors. One drawback to using AAV capsids for gene deivery is they can be neutralized by natural immunity. This sets up a protein engineering problem. We want to develop new variants of the virus that caan evade the human immune system.We could approach this with a scanning mutation approach, but this would result in a high number of invalid sequences. Instead, we will build a score function based off laboratory data and use a generative model to exploit this score function. We hope this approach will give more realistic variants.  Performance NotesThe workflow in this notebook is more CPU-constrained than GPU-constrained due to the need to evaluate samples on CPU. If you have a multi-core machine, it is recommended that you uncomment and run the `set_global_pool` cells in the notebook. This will trigger the use of multiprocessing, which will result in 2-4x speedups.This notebook may run slow on Collab due to CPU limitations.If running on Collab, remember to change the runtime to GPU<jupyter_code>import sys
sys.path.append('..')

from mrl.imports import *
from mrl.core import *
from mrl.chem import *
from mrl.templates.all import *

from mrl.torch_imports import *
from mrl.torch_core import *
from mrl.layers import *
from mrl.dataloaders import *
from mrl.g_models.all import *
from mrl.vocab import *
from mrl.policy_gradient import *
from mrl.train.all import *
from mrl.model_zoo import *
from sklearn.metrics import r2_score
os.makedirs('untracked_files', exist_ok=True)<jupyter_output><empty_output><jupyter_text>DataThe dataset comes from the paper [Generative AAV capsid diversification by latent interpolation ](https://www.biorxiv.org/content/10.1101/2021.04.16.440236v1.full.pdf). The authors looked at a 28 amino acid section of the AAV2 VP3 protein shown to have immunological significance.The authors created a mutation library based on 564 naturally occurring sequences. They then trained a VAE model on this dataset, sampled new 28-AA sequences from this model, and tested them in the lab.We will use their laboratory data to build a score function. Then we will run a generative screen against that score function<jupyter_code># download data
# ! wget https://raw.githubusercontent.com/churchlab/Generative_AAV_design/main/Data/vae2021_processed_data.csv
df = pd.read_csv('vae2021_processed_data.csv')
df.dropna(subset=['VAE_virus_S'], inplace=True)
df.head()<jupyter_output><empty_output><jupyter_text>Our metric of interest is `VAE_virus_S`. This is the log-2 ratio of the frequency of a variant in a virus pool relative to the frequency of the corresponding plasmid in the plasmid pool. Higher values indicate higher viability. The goal of the design is to produce variants predicted to beat the wildtype sequence. The dataset indicates which sequences are currently performing at this level<jupyter_code>df.VAE_virus_S.hist(alpha=0.5, density=True, label='full_dataset')
df[(df.viable==1) & (df.beats_wt==1)].VAE_virus_S.hist(alpha=0.5, density=True, label='Beats WT')
plt.legend();<jupyter_output><empty_output><jupyter_text>Score FunctionNow we want to develop a score function for predicting antimicrobial activity. We will use a CNN encoder with a MLP head to predict the score described above.Our input data will be token integers for amino acids. Note that fingerprint representations are a poor fit for peptide work because peptides contain many repeating substructures.We will train on 90% of the data and validate on the 10% held out.<jupyter_code>train_df = df.sample(frac=0.9, random_state=42).copy()
valid_df = df[~df.index.isin(train_df.index)].copy()
aa_vocab = CharacterVocab(AMINO_ACID_VOCAB)

train_ds = Text_Prediction_Dataset(train_df.aa.values, train_df.VAE_virus_S.values, aa_vocab)
test_ds = Text_Prediction_Dataset(valid_df.aa.values, valid_df.VAE_virus_S.values, aa_vocab)<jupyter_output><empty_output><jupyter_text>This is the model we will use:<jupyter_code>class Predictive_CNN(nn.Module):
    def __init__(self,
                 d_vocab,
                 d_embedding,
                 d_latent,
                 filters,
                 kernel_sizes,
                 strides,
                 dropouts,
                 mlp_dims,
                 mlp_drops,
                 d_out,
                 outrange
                ):
        super().__init__()
        
        self.conv_encoder = Conv_Encoder(
                                        d_vocab,
                                        d_embedding,
                                        d_latent,
                                        filters,
                                        kernel_sizes,
                                        strides,
                                        dropouts,
                                    )
        
        self.mlp_head = MLP(
                            d_latent,
                            mlp_dims,
                            d_out,
                            mlp_drops,
                            outrange=outrange
                            )
        
    def forward(self, x):
        encoded = self.conv_encoder(x)
        out = self.mlp_head(encoded)
        return out
d_vocab = len(aa_vocab.itos)
d_embedding = 256
d_latent = 512
filters = [256, 512, 1024]
kernel_sizes = [7, 7, 7]
strides = [2,2,2]
dropouts = [0.2, 0.2, 0.2]
mlp_dims = [512, 256, 128]
mlp_drops = [0.2, 0.2, 0.2]
d_out = 1
outrange = [-10, 10]


virus_model = Predictive_CNN(
                    d_vocab,
                    d_embedding,
                    d_latent,
                    filters,
                    kernel_sizes,
                    strides,
                    dropouts,
                    mlp_dims,
                    mlp_drops,
                    d_out,
                    outrange
                )
r_agent = PredictiveAgent(virus_model, MSELoss(), train_ds, opt_kwargs={'lr':1e-3})
r_agent.train_supervised(32, 20, 1e-3)<jupyter_output><empty_output><jupyter_text>Optional: save score function weights<jupyter_code># r_agent.save_weights('untracked_files/virus_predictor.pt')<jupyter_output><empty_output><jupyter_text>Optional: to load the exact weights used, run the following:<jupyter_code># r_agent.load_state_dict(model_from_url('virus_predictor.pt'))
# validate

valid_dl = test_ds.dataloader(256, num_workers=0, shuffle=False)
r_agent.model.eval();

preds = []
targs = []

with torch.no_grad():
    for i, batch in enumerate(valid_dl):
        batch = to_device(batch)
        x,y = batch
        pred = r_agent.model(x)
        preds.append(pred.detach().cpu())
        targs.append(y.detach().cpu())
        
preds = torch.cat(preds).numpy()
targs = torch.cat(targs).numpy()

preds = preds.squeeze(-1)<jupyter_output><empty_output><jupyter_text>Our score function has an r^2 value of about 0.883 on the validation dataset<jupyter_code>fig, ax = plt.subplots()

ax.scatter(targs, preds, alpha=0.5, s=1)
plt.xlabel('Target')
plt.ylabel('Prediction')

slope, intercept = np.polyfit(targs, preds, 1)
ax.plot(np.array(ax.get_xlim()), intercept + slope*np.array(ax.get_xlim()), c='r')

plt.text(5., 9., 'R-squared = %0.3f' % r2_score(targs, preds));<jupyter_output><empty_output><jupyter_text>We should also take a look at the prediction distribution for our known actives. The predicted values will differ somewhat from the actual values. This will give us a sense of what score we want to see from the model<jupyter_code>df['preds'] = r_agent.predict_data(df.aa.values).detach().cpu().numpy()
df.preds.max()
np.percentile(df.preds, 99)<jupyter_output><empty_output><jupyter_text>The maximum predicted value is `8.63`. A compound scoring `5.33` or higher would be in the top 1% of all known sequences  Chemical SpaceNext we need to develop our chemical space. This is where we decide what sequences will be allowed and what sequences will be removed.Getting the right filtering parameters makes a huge difference in sequence quality. In practice, finding the right constraints is an interative process. First run a generative screen. Then examine the highest scoring sequences and look for undesirable properties or structural features. Then update the template and iterate.For peptides, the presence of Arginine has shown to be toxic [[ref](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5625148/)]. We will apply a template filter for the number of Arginine residues per unit length.We will also limit the maximum residue frequency in a sample to 0.3. This prevents a common failure mode of seeing the same residue repeated multiple times (ie `MSSSSSSSRP`). This is a flaw in the simplistic score functions we are using<jupyter_code>aa_vocab = CharacterVocab(AMINO_ACID_VOCAB)

template = Template([ValidityFilter(),
                     CharacterCountFilter(['A'], min_val=0, max_val=0.1, per_length=True, mode='protein'),
                     CharacterCountFilter(aa_vocab.itos[4:], min_val=0, max_val=0.3, 
                                          per_length=True, mode='protein')], 
                    [], fail_score=-10., log=False, use_lookup=False, mode='protein')

template_cb = TemplateCallback(template, prefilter=True)<jupyter_output><empty_output><jupyter_text>Load ModelWe load the `LSTM_LM_Small_Swissprot` model. This is a basic LSTM-based language model trained on part of the Swissprot protein database<jupyter_code>agent = LSTM_LM_Small_Swissprot(drop_scale=0.3, opt_kwargs={'lr':5e-5})<jupyter_output><empty_output><jupyter_text>Fine-Tune ModelThe pretrained model we loaded is a very general model that can produce a high diversity of structures. However, what we actually want are structues with antimicrobial activity. To induce this, we fine-tune on high scoring sequences.The dataset denotes two categories for high scoring sequences. The `viable` category contains ~3000 samples that were able to successfully assemble and package genetic material. The `beats_wt` category contains ~300 samples with higher expression than the wildtype variant.The `beats_wt` dataset feels a bit small for fine-tuning, so we will first fine-tune on the `viable` dataset, then the `beats_wt` dataset<jupyter_code>agent.update_dataset_from_inputs(df[df.viable==1].aa.values)
agent.train_supervised(32, 8, 5e-5)

agent.update_dataset_from_inputs(df[df.beats_wt==1].aa.values)
agent.train_supervised(32, 6, 5e-5)

agent.base_to_model()<jupyter_output><empty_output><jupyter_text>Optional: save fine-tuned weights<jupyter_code># agent.save_weights('untracked_files/finetuned_model.pt')
# agent.load_weights('untracked_files/finetuned_model.pt')<jupyter_output><empty_output><jupyter_text>Reinforcement LearningNow we enter the reinforcement learning stage  LossWe use `PPO` as our policy gradient loss<jupyter_code>pg = PPO(0.99,
        0.5,
        lam=0.95,
        v_coef=0.5,
        cliprange=0.3,
        v_cliprange=0.3,
        ent_coef=0.01,
        kl_target=0.03,
        kl_horizon=3000,
        scale_rewards=True)

loss = PolicyLoss(pg, 'PPO', 
                   value_head=ValueHead(256), 
                   v_update_iter=2, 
                   vopt_kwargs={'lr':1e-3})<jupyter_output><empty_output><jupyter_text>RewardHere we pass the reward agent we trained earlier to a callback.<jupyter_code>aa_vocab = CharacterVocab(AMINO_ACID_VOCAB)

d_vocab = len(aa_vocab.itos)
d_embedding = 256
d_latent = 512
filters = [256, 512, 1024]
kernel_sizes = [7, 7, 7]
strides = [2,2,2]
dropouts = [0.2, 0.2, 0.2]
mlp_dims = [512, 256, 128]
mlp_drops = [0.2, 0.2, 0.2]
d_out = 1
outrange = [-10, 10]


reward_model = Predictive_CNN(
                    d_vocab,
                    d_embedding,
                    d_latent,
                    filters,
                    kernel_sizes,
                    strides,
                    dropouts,
                    mlp_dims,
                    mlp_drops,
                    d_out,
                    outrange
                )

r_ds = Text_Prediction_Dataset(['M'], [0.], aa_vocab)

r_agent = PredictiveAgent(reward_model, MSELoss(), r_ds, opt_kwargs={'lr':1e-3})

r_agent.load_weights('untracked_files/virus_predictor.pt')
# r_agent.load_state_dict(model_from_url('virus_predictor.pt')) # optional - load exact weights

reward_model.eval();

freeze(reward_model)

reward_function = Reward(r_agent.predict_data, weight=1)

virus_reward = RewardCallback(reward_function, 'virus')<jupyter_output><empty_output><jupyter_text>Optional Reward: Stability MetricThere has been a lot of great work recently looking at large scale transformer language models for unsupervised learning of protein structures. One interesting thing that has emerged is a rough relationship between the protein sequence log probability given by a generative model and the stability of the protein sequence.We can use the log probability values from a pretrained protein transformer model as a proxy for stability. Including this as a reward function can help keep the generated peptides realistic.To include this as a reward term, run the code below to install the [ESM](https://github.com/facebookresearch/esm) library to access a pretrained protein transformer model.In the interest of time, we will use the smallest ESM model with 43M parameters, rather than the large scale 630M parameter model. Note that even with the smaller model, this reward term adds significanly to the training runtime<jupyter_code># Optional: insall ESM
# ! pip install fair-esm
import esm
protein_model, alphabet = esm.pretrained.esm1_t6_43M_UR50S()
batch_converter = alphabet.get_batch_converter()
class PeptideStability():
    def __init__(self, model, alphabet, batch_converter):
        self.model = model
        to_device(self.model)
        self.alphabet = alphabet
        self.batch_converter = batch_converter
        
    def __call__(self, samples):
        
        data = [
            (f'protein{i}', samples[i]) for i in range(len(samples))
        ]
        
        batch_labels, batch_strs, batch_tokens = self.batch_converter(data)

        with torch.no_grad():
            results = self.model(to_device(batch_tokens))

        lps = F.log_softmax(results['logits'], -1)

        mean_lps = lps.gather(2, to_device(batch_tokens).unsqueeze(-1)).squeeze(-1).mean(-1)
        
        return mean_lps
ps = PeptideStability(protein_model, alphabet, batch_converter)
stability_reward = Reward(ps, weight=0.1, bs=300)
stability_cb = RewardCallback(stability_reward, name='stability')
stability_reward(df.aa.values[:10])<jupyter_output><empty_output><jupyter_text>SamplersWe create the following samplers:- `sampler1 ModelSampler`: this samples from the main model. This sample will add 1000 compounds to the buffer each buffer build, and sample 40% of each batch on the fly from the main model.- `sampler2 ModelSampler`: this samples from the baseline model and is not sampled live on each batch- `sampler3 LogSampler`: this samples high scoring samples from the lig- `sampler4 TokenSwapSampler`: this uses token swap comibichem to generate new samples from high scoring samples- `sampler5 DatasetSampler`: this sprinkles in a small amount of high scoring samples into each buffer build.<jupyter_code>gen_bs = 1500

sampler1 = ModelSampler(agent.vocab, agent.model, 'live', 1000, 0., gen_bs)
sampler2 = ModelSampler(agent.vocab, agent.base_model, 'base', 1000, 0., gen_bs)
sampler3 = LogSampler('samples', 'rewards', 10, 98, 200)
sampler4 = TokenSwapSampler('samples', 'rewards', 10, 98, 200, aa_vocab, 0.2)
sampler5 = DatasetSampler(df[(df.beats_wt==1)].aa.values, 
                          'data', buffer_size=6)

samplers = [sampler1, sampler2, sampler3, sampler4, sampler5]<jupyter_output><empty_output><jupyter_text>CallbacksAdditional callbacks- `SupervisedCB`: runs supervised training on the top 3% of samples every 400 batches- `MaxCallback`: prints the max reward for each batch- `PercentileCallback`: prints the 90th percentile score each batch<jupyter_code>supervised_cb = SupervisedCB(agent, 20, 0.5, 98, 1e-4, 64)
live_max = MaxCallback('rewards', 'live')
live_p90 = PercentileCallback('rewards', 'live', 90)

cbs = [supervised_cb, live_p90, live_max]<jupyter_output><empty_output><jupyter_text>Environment and TrainNow we can put together our Environment and run the training process<jupyter_code>env = Environment(agent, template_cb, samplers=samplers, rewards=[virus_reward, stability_cb], losses=[loss],
                 cbs=cbs)
# set_global_pool(min(12, os.cpu_count()))
env.fit(128, 28, 400, 25)
env.log.plot_metrics()<jupyter_output><empty_output><jupyter_text>EvaluationBased on our score function, we determined a sequence with a predicted score of 5.33 or higher would be in the top 1% of sequences relative to the training data. A sequence with a predicted score of 8.63 or higher would beat out all sequences in the datasetWe found 1670 sequences with a predicted score of 5.33 or higher and 13 sequences with a predicted score of 8.63 or higher<jupyter_code>env.log.df[(env.log.df.virus>5.33) & ~(env.log.df.sources=='data_buffer')].shape
env.log.df[(env.log.df.virus>8.63) & ~(env.log.df.sources=='data_buffer')].shape<jupyter_output><empty_output><jupyter_text>We can generate logo stack plots to visualize what residues are favored<jupyter_code>def plot_logo(seqs):
    freqs = []

    for i in range(len(seqs[0])):
        aas = [j[i] for j in seqs]
        counts = Counter(aas)
        total = sum(counts.values())
        norm_counts = defaultdict(lambda: 0)
        for k in counts.keys():
            norm_counts[k] = counts[k]/total
        freqs.append(norm_counts)
        
    aas = aa_vocab.itos[4:]
    dfs = []

    for i, f in enumerate(freqs):
        df_iter = pd.DataFrame([[f[aa] for aa in aas]], columns=aas)
        df_iter['Position'] = i
        dfs.append(df_iter)
        
    dfs = pd.concat(dfs)
    dfs = dfs.set_index('Position')

    return dfs.groupby('Position').mean().plot.bar(stacked=True, figsize=(12,8))<jupyter_output><empty_output><jupyter_text>Here's the logo plot for high scoring sequences in the dataset<jupyter_code>plot_logo(df[(df.beats_wt==1) & (df.VAE_virus_S>7)].aa.values)<jupyter_output><empty_output><jupyter_text>Here's the residue plot for high scoring generated sequences<jupyter_code>plot_logo(env.log.df[(env.log.df.virus>8.9) & ~(env.log.df.sources=='data_buffer')].samples.values)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08918456770332371, 'fraction_numerical': 0.021801859570375122, 'mean_word_length': 3.512900892211237, 'pattern_counts': {'":': 0, '<': 93, '<?xml version=': 0, '>': 98, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '23', 'dir': 'jupyter-structured-clean-dedup', 'id': '497821', 'license': 'MIT', 'n_tokens_mistral': 6125, 'n_tokens_neox': 5660, 'n_words': 1719, 'path': 'nbs/tutorials/tutorials.proteins.capsid.ipynb', 'repo_name': 'DarkMatterAI/mrl', 'size': '18714'}