wazy

Pretrained Bayesian Optimization of Amino Acid Sequences. This is actively developed research code - things can break and the API may change. Please use caution and open an issue if things are unexpected!

Colab

Click the button below to use the algorithm in Google Colab

Colab Peptide Binder

To use AlphaFold with Wazy to design peptides that bind to specific proteins, click the button below.

installing

pip install wazy

Quickstart

You can use an ask/tell style interface to design a peptide.

We can tell a few examples of sequences we know and their scalar labels. Let's try a simple example where the label is the number of alanines. You'll also want your labels to vary from about -5 to 5. We'll start by importing and building a BOAlgorithm class. In this example, I re-use the same key for simplicity.

import wazy
import jax
key = jax.random.PRNGKey(0)
boa = wazy.BOAlgorithm()

Now we can tell it a few examples.

boa.tell(key, "GGGG", 0)
boa.tell(key, "GAHK", 1)
boa.tell(key, "DAAE", 2)
boa.tell(key, "DAAA", 3)

We can predict on new values. This will return both a predicted label and its uncertainty and its epistemic uncertainty.

boa.predict(key, "LPAH")
# Output:
(5.823452, 69.99278, 24.500998)

The accuracy is poor - $5.8\pm 70$. Let's now use Bayesian optimization to choose which sequence to try next:

boa.ask(key)
# Output
('DAAV', 6.901945)

The first value is the sequence to try next. The second is an indicator in how valuable (value of acquisition function) it finds that sequence. Now we can tell it the value:

boa.tell(key, "DAAV", 2)

We can also choose the sequence length:

boa.ask(key, length=6)
# Output
('DAAATA', 5.676821)

We can try our new prediction to see if it improved.

boa.tell(key, "DAAATA", 4)
boa.predict(key, "LPAH")
# Output
(2.0458677, 13.694655, 1.0933837)

Which is indeed closer to the true answer of 1. Finally, we can ask for the best sequence:

boa.ask(key, "max", length=5)
# Output
('DAAAA', 3.8262398)

Key

If you are going to use this process in a loop, be sure to split the key:

s = "START"
for i in range(10):
  key, _ = jax.random.split(key)
  boa.tell(key, s, 4)
  s, _ = boa.ask(key, "max", length=5)

Batching

You can increase the number of returned sequences by using the batch_ask, which uses an ad-hoc regret minimization strategy to spread out the proposed sequences:

boa.batch_ask(key, N=3)
# returns 3 seqs

and you can add a multiplier to batch sequences (no overhead), but they may be similar

boa.batch_ask(key, N=3, return_seqs = 10)
# returns 30 seqs

Citation

Please cite Yang et. al.

@article{yang2022now,
  title={Now What Sequence? Pre-trained Ensembles for Bayesian Optimization of Protein Sequences},
  author={Yang, Ziyue and Milas, Katarina A and White, Andrew D},
  journal={bioRxiv},
  year={2022},
  publisher={Cold Spring Harbor Laboratory}
}