In-silico directed evolution using Metropolis–Hastings Markov chain Monte Carlo sampling¶

In the notebooks before this one, we tackled the challenge of understanding the relationship between protein sequences and their functions. We took two different paths: first notebook, training a regressor using fixed embeddings of our sequences from a protein language model in transformers, and second notebook, fine-tuning a pretrained protein transformer model with our DMS data.

Now, we're diving into the realm of in-silico directed evolution to navigate the vast protein sequence landscape. Think of it like using a surrogate fitness landscape created by our earlier sequence-to-function model. Our mission here is to implement the Metropolis–Hastings Markov chain Monte Carlo algorithm. This fancy term essentially means we're using a computational method for simulated evolution to propose potentially optimized protein designs.

Goals

• Dive into the Metropolis–Hastings Markov chain Monte Carlo (MH MCMC) algorithm for in-silico directed evolution on our surrogate fitness landscape.

Metropolis– Hastings Markov chain Monte Carlo algorithm¶

Algorithm description inspired from: https://github.com/SuperSecretBioTech/evolutionary_monte_carlo_search and https://www.nature.com/articles/s41592-021-01100-y

In the Metropolis-Hastings algorithm, especially when applied to scenarios like in silico directed evolution, the acceptance probability when proposing a move from a current state $s_i$ with fitness $f(s_i)$ to a new state $s_{i+1}$ with fitness $f(s_{i+1})$ is typically expressed as:

$$ \text{Accept with probability } \min\left(1, \exp\left(\frac{f(s_{i+1}) - f(s_i)}{T}\right)\right) $$

For our in silico directed evolution algorithm we need:

• a. An initial sequence.
• b. A sequence-to-function model that predicts the quantitative function or fitness of an amino acid sequence.
• c. Temperature, T
• d. Trust radius (the number of mutations relative to wild-type allowed in proposed designs).

The protocol for the algorithm is as follow:

1. Start with a state sequence $s_{i}$ with fitness $f(s_i)$. Set state sequence s equal to a provided initial sequence.
2. Propose a new sequence $s_{i+1}$ with fitness $f(s_{i+1})$. The new sequence is proposed by randomly adding m ~ Poisson(μ − 1) + 1 mutations to $s_{i}$, where μ is the sequence proposal mutation rate. The sequence proposal mutation rate μ for each trajectory was set to be a random draw from a uniform(1, 2.5) distribution.
3. If $f(s_{i+1})$ > $f(s_i)$, we accept the new sequence and the algorithm starts over. In this scenario $f(s_{i+1})$ - $f(s_i)$ is positive, making $exp\left(\frac{f_{s+1} - f_{s}}{T}\right)$ greater than 1. Here $, min(1,value) $ ensures the probability doesn't exceed 1.
4. If $f(s_{i+1}) < f(s_i)$, the expression $\exp\left(\frac{f(s_{i+1}) - f(s_i)}{T}\right)$ results in a value between 0 and 1, representing the probability of accepting a move to a less fit state.

Note that, if the sequence proposal has more mutations than the input trust radius, its predicted fitness is set, post hoc, to negative infinity, thereby forcing rejection of the proposal.

5. Iterate steps 3 and 4 for a predetermined number of iterations.

Metropolis–Hastings Markov chain Monte Carlo (MH_MCMC) has two main hyperparameters: temperature and the number of iterations. Typically, we determine the number of iterations by experimenting with different values and visually inspecting the fitness-versus-iterations plot to find a convergence point with acceptable fitness.

The temperature (called the temperature because that's what the term represents in the boltzmann distribution of energy states at a finite temperature) is the hyperparameter which decides how liberally non-optimal moves are accepted.

The temperature, named so because of its role in the Boltzmann distribution, dictates how generously we accept non-optimal moves. High temperatures increase acceptance probabilities, leading to overly explorative trajectories lingering in low predicted fitness areas. Excessively high temperatures, accepting more than 50-60% non-optimal moves, hinder convergence as the algorithm struggles to reject moves near an optimum, favoring broad searches.

Conversely, low temperatures reduce acceptance probabilities, resulting in overly exploitative trajectories with consistently increasing fitness traces. Lower temperatures make MH_MCMC more conservative, strictly accepting moves that marginally increase or decrease fitness. This facilitates movement towards an optimum but limits exploration to directions where fitness only increases, akin to an unguided steepest descent algorithm.

Selecting the right temperature is crucial, balancing search depth (exploitation) and breadth (exploration). A common guideline is choosing a temperature that leads to convergence while accepting about 15-20% non-optimal moves during the simulation, though this can vary depending on the problem.

MH_MCMC's strength lies in its simplicity, low computational cost, and ease of debugging/analysis. However, its effectiveness heavily relies on the initial sequence and the local landscape around it, determining whether convergence to a satisfactory fitness occurs.

Import and setup¶

Load the protein sequence-to-function model we built in my previous notebook Learning sequence-to-function relationship using language models.

Basic sampling¶

Here we will run one evolutionary trajectory on rep78 sequence using the MH MCMC algorithm described in the beginning. Each evolutionary trajectory consists of thousands of iterations depending on the length of protein. For basic sampling we need:

• starting sequence
• sequence-to-function-model
• Monte Carlo sampling alongside the Metropolis-Hastings criteria to score and rank-order new sequences

Starting sequence¶

Sequence to function model¶

Monte Carlo sampling alongside the Metropolis-Hastings criteria¶

For us a temperature of 0.01 yielded good trajectory behavior - fitness traces improved on average but were not monotonic.

Parallel sampling¶

Here we will run many evolutionary trajectories in parallel.

Find best protein sequence in each evolutionary trajectory

Check MCMC trajectories for sampled seqs¶

The fitness traces in evolutionary paths generally improved but weren't consistently increasing, indicating a balanced mix of exploration and exploitation, ensuring effective optimization while exploring diverse options.

Closing statement¶

In the quest for high-functioning protein sequences, scientists need to be mindful of the well-known saying, 'you get what you screen for.' This cautionary maxim underscores the risk of overoptimizing a protein's sequence based on functional assays that may not perfectly align with the ultimate design objective. It serves as a reminder to ensure that the screening process aligns closely with the intended goals to avoid unintended consequences in the final design.