wazy 0.9.0

Pretrained Bayesian Optimization of Sequences

wazy

Pretrained Bayesian Optimization of Amino Acid Sequences

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)

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