PoET: Aliphatic amidase sequence analysis (substitutions and deletions)

This tutorial demonstrates how to use PoET to score substitution and deletion variants of a sequence using aliphatic amidase (AMIE_PSEAE) as an example.

We’ll also compare substitution variant scores from PoET, ESM1b, and a site-independent model (PSSM) with activities assayed in a deep mutational scanning study by {AUTHORS}.

[1]:
%matplotlib inline
[2]:
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import scipy
from scipy.stats import pearsonr,spearmanr,kendalltau
import scipy.special
import json
import time
from mpl_toolkits.axes_grid1 import make_axes_locatable

import seaborn as sns
sns.set()
[3]:
import openprotein
import openprotein.fasta as fasta
[4]:
name = 'AMIE_PSEAE'
wt = b'MRHGDISSSNDTVGVAVVNYKMPRLHTAAEVLDNARKIAEMIVGMKQGLPGMDLVVFPEYSLQGIMYDPAEMMETAVAIPGEETEIFSRACRKANVWGVFSLTGERHEEHPRKAPYNTLVLIDNNGEIVQKYRKIIPWCPIEGWYPGGQTYVSEGPKGMKISLIICDDGNYPEIWRDCAMKGAELIVRCQGYMYPAKDQQVMMAKAMAWANNCYVAVANAAGFDGVYSYFGHSAIIGFDGRTLGECGEEEMGIQYAQLSLSQIRDARANDQSQNHLFKILHRGYSGLQASGDGDRGLAECPFEFYRTWVTDAEKARENVERLTRSTTGVAQCPVGRLPYEGLEKEA'

print('>' + name)
for i in range(0, len(wt), 80):
    print(wt[i:i+80].decode())
>AMIE_PSEAE
MRHGDISSSNDTVGVAVVNYKMPRLHTAAEVLDNARKIAEMIVGMKQGLPGMDLVVFPEYSLQGIMYDPAEMMETAVAIP
GEETEIFSRACRKANVWGVFSLTGERHEEHPRKAPYNTLVLIDNNGEIVQKYRKIIPWCPIEGWYPGGQTYVSEGPKGMK
ISLIICDDGNYPEIWRDCAMKGAELIVRCQGYMYPAKDQQVMMAKAMAWANNCYVAVANAAGFDGVYSYFGHSAIIGFDG
RTLGECGEEEMGIQYAQLSLSQIRDARANDQSQNHLFKILHRGYSGLQASGDGDRGLAECPFEFYRTWVTDAEKARENVE
RLTRSTTGVAQCPVGRLPYEGLEKEA

Connect to the OpenProtein.AI API

[5]:
with open('secrets.config', 'r') as f:
    config = json.load(f)

session = openprotein.connect(config['username'], config['password'])

Create the MSA and define an ensemble prompt

At the core of PoET is the prompt, a set of sequences that encode information about the fitness landscape of our protein of interest. A typical way to create a prompt is by constructing an MSA from homologues of the protein of interest. We support this natively via our API (see below).

OpenProtein uses an asynchronous API, where potentially long running functions return a job ID that can be used to query for completed results. The wait_until_done function can be used to poll for completion.

[6]:
# search for homologs to automatically create an MSA for the seed sequence
msa = session.poet.create_msa(wt)
print(msa)
msa.wait_until_done(verbose=True)

# create the prompt, set the seed for reproducibility
prompt = msa.sample_prompt(num_ensemble_prompts=10, random_seed=1)
print(prompt)
prompt.wait_until_done(verbose=True)
status=<JobStatus.SUCCESS: 'SUCCESS'> job_id='05da7510-3f63-44ce-a512-2475a418a873' job_type='/align/align' created_date=datetime.datetime(2023, 11, 3, 3, 46, 42, 839364) start_date=None end_date=datetime.datetime(2023, 11, 3, 3, 46, 42, 967065) prerequisite_job_id=None progress_message=None progress_counter=None num_records=None msa_id='05da7510-3f63-44ce-a512-2475a418a873'
Waiting: 100%|██████████| 100/100 [00:00<00:00, 110.99it/s, status=SUCCESS]
status=<JobStatus.SUCCESS: 'SUCCESS'> job_id='76e8fca5-7f41-44ef-9724-456f49acfb82' job_type='/align/prompt' created_date=datetime.datetime(2023, 11, 2, 7, 36, 9, 887755) start_date=datetime.datetime(2023, 11, 2, 7, 36, 18, 793486) end_date=datetime.datetime(2023, 11, 2, 7, 36, 33, 822849) prerequisite_job_id=None progress_message=None progress_counter=None num_records=None
Waiting: 100%|██████████| 100/100 [00:00<00:00, 374.70it/s, status=SUCCESS]
[6]:
True

Above, you can see we first create an MSA from the wildtype (wt) protein. We then create 10 different prompts based on this, using the default prompt arguments. Each prompt will meet the prompt definition parameters but contain different sequences from the MSA to create a more robust fitness score.

Run all single substitution predictions for the parent sequence

All single substitutions can be scored using the session.poet.single_site function.

[7]:
future_ssp = session.poet.single_site(prompt, wt)
print(future_ssp.job)
future_ssp.wait_until_done(verbose=True)
status=<JobStatus.SUCCESS: 'SUCCESS'> job_id='780766ef-d263-42c3-9456-8f221c490c47' job_type='/poet/single_site' created_date=datetime.datetime(2023, 9, 16, 17, 28, 3, 449996) start_date=datetime.datetime(2023, 9, 16, 17, 35, 22, 965839) end_date=datetime.datetime(2023, 9, 16, 17, 37, 45, 496463) prerequisite_job_id=None progress_message=None progress_counter=None num_records=None
Waiting: 100%|██████████| 100/100 [00:00<00:00, 236.24it/s, status=SUCCESS]
[7]:
True
[8]:
# retrieve the results for the SSP job
results = future_ssp.get()
[9]:
# parse results into matrix of variants at each site
result_dict = {}
for r in results:
    result_dict[r.sequence] = np.array(r.score)

baseline = result_dict[b'input']
amino_acids = b'ARNDCQEGHILKMFPSTWYV'

scores_ensemble = np.zeros((len(wt), 20, len(baseline))) + baseline
for i in range(scores_ensemble.shape[0]):
    for j in range(len(amino_acids)):
        pos = i + 1
        fr = wt[i:i+1]
        to = amino_acids[j:j+1]
        if fr != to:
            code = fr + str(pos).encode() + to
            val = result_dict[code]
            scores_ensemble[i, j] = val
scores_ensemble.shape
[9]:
(346, 20, 10)

Plot the single site predictions

[10]:
scores = np.stack(list(result_dict.values()), axis=0) - baseline
scores = scores.mean(axis=-1)
plt.title("Sequence fitness vs wildtype")
_ = plt.hist(scores.ravel(), bins=30)
_ = plt.xlabel('Log likelihood ratio (p(variant)/p(wt))');
plt.axvline(0, color='blue', ls="--");
../_images/demos_AMIE_substitution_deletion_analysis_poet_15_0.png

We can see above that a small subset of single site mutants have higher predicted fitness than wildtype aliphatic amidase.

[11]:
scores = np.mean(scores_ensemble - baseline, axis=-1)
fig, ax = plt.subplots(figsize=(55, 10))
plt.title("Single site mutation fitness heatmap", size=32)
# Plot the data
im = ax.imshow(scores.T, cmap='viridis')

# Set y-axis labels
y_labels = amino_acids.decode()
y_ticks = np.arange(len(y_labels))
ax.set_yticks(y_ticks)
ax.set_yticklabels(y_labels, fontsize=12)

# Label axes
plt.ylabel('Substitution', size=32)
plt.xlabel('Position', size=32)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="1%", pad=0.05)

plt.colorbar(im, cax=cax)
plt.colorbar(im, cax=cax)
plt.show()

../_images/demos_AMIE_substitution_deletion_analysis_poet_17_0.png

We can look at a subset of the positions more closely:

[12]:

fig, ax = plt.subplots(figsize=(15, 5)) plt.title("Positions 70-100 only", size=18) # Plot the data im = ax.imshow(scores.T, cmap='viridis') # Set y-axis labels y_labels = amino_acids.decode() y_ticks = np.arange(len(y_labels)) ax.set_yticks(y_ticks) ax.set_yticklabels(y_labels, fontsize=12) # Adjust fontsize as needed # Label axes plt.ylabel('Substitution', size=32) plt.xlabel('Position', size=32) plt.xlim(70,100) plt.colorbar(im, ax=ax) plt.show()
../_images/demos_AMIE_substitution_deletion_analysis_poet_19_0.png

Plot scores from an individual prompts in the ensemble

Because we used a prompt ensemble, we get scores for each prompt in the ensemble. They can be combined into a single score via averaging, but can also be examined individually.

[13]:

heatmap = scores_ensemble[..., 0] - baseline[..., 0] fig, ax = plt.subplots(figsize=(55, 10)) plt.title("Single site mutation fitness heatmap (prompt 1)", size=32) # Plot the data im = ax.imshow(heatmap.T, cmap='viridis') # Set y-axis labels y_labels = amino_acids.decode() y_ticks = np.arange(len(y_labels)) ax.set_yticks(y_ticks) ax.set_yticklabels(y_labels, fontsize=12) # Label axes plt.ylabel('Substitution', size=32) plt.xlabel('Position', size=32) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="1%", pad=0.05) plt.colorbar(im, cax=cax) plt.colorbar(im, cax=cax) plt.show()
../_images/demos_AMIE_substitution_deletion_analysis_poet_21_0.png

Rank PoET variant predictions

[14]:
variants, scores = zip(*result_dict.items())
variants = np.array(variants)
scores = np.array(scores)
variants.shape, scores.shape
[14]:
((6575,), (6575, 10))
[15]:
# rank the variants
order = np.argsort(-scores.mean(axis=1))
for i in order[:10]:
    print(variants[i].decode(), scores[i].mean(axis=-1), (scores[i]-baseline).mean(axis=-1))
L62T -140.06779766082764 5.570239067077635
A28K -140.98878440856933 4.649252319335934
I279L -141.4669720649719 4.1710646629333485
A28R -141.7436710357666 3.8943656921386776
Q288H -142.2893359184265 3.3487008094787543
V77A -142.6376898765564 3.00034685134888
K93E -143.45875930786133 2.1792774200439426
T327V -143.4867042541504 2.1513324737548816
S285T -143.63027153015136 2.0077651977539035
N125K -143.64069890975952 1.9973378181457506

Compare the PoET substitution scores with measurements from deep mutational scanning

Deep mutational scanning of AMIE_PSEAE has been performed by AUTHORS. In this study, they measured activites of most single substitution variants of the wildtype protein against three substrates. Let’s load this data and see how well PoET predicted the effects of these variants compared with a few baselines.

As baselines, we’ll consider * a position-specific scoring matrix (PSSM, and additive model) fit on the MSA found above * variant effect predictions from the per-position amino acid probabilities given by ESM1b

Load the DMS data

[16]:
path = "data/AMIE_PSEAE_Whitehead.csv"
table = pd.read_csv(path, index_col=0)
table.head()
[16]:
mutant acetamide_normalized_fitness isobutyramide_normalized_fitness propionamide_normalized_fitness mutation_effect_prediction_vae_ensemble mutation_effect_prediction_vae_1 mutation_effect_prediction_vae_2 mutation_effect_prediction_vae_3 mutation_effect_prediction_vae_4 mutation_effect_prediction_vae_5 mutation_effect_prediction_pairwise mutation_effect_prediction_independent
0 M1W NaN -0.5174 NaN NaN NaN NaN NaN NaN NaN NaN NaN
1 M1Y NaN -0.5253 NaN NaN NaN NaN NaN NaN NaN NaN NaN
2 M1P -2.1514 -0.5154 -1.1457 NaN NaN NaN NaN NaN NaN NaN NaN
3 M1M 0.0000 0.0000 0.0000 NaN NaN NaN NaN NaN NaN NaN NaN
4 M1I -0.1227 -0.3640 -0.1212 NaN NaN NaN NaN NaN NaN NaN NaN
[17]:
cols = ['acetamide_normalized_fitness', 'isobutyramide_normalized_fitness', 'propionamide_normalized_fitness']
variants = table['mutant']

Match against the DMS variants and compare PoET log probabilities with activity measurements

PoET log probabilities accurately predict the activity of AMIE variants for all three substrates.

[18]:
baseline = result_dict[b'input']
variants = table['mutant'].values
logp = np.array([result_dict.get(code.encode(), baseline) for code in variants])
logp.shape
[18]:
(6819, 10)
[19]:
scores = np.mean(logp, axis=-1) - baseline.mean()
sigma = np.std(logp - baseline, axis=-1)

# Create subplots with labels and formatting
fig, axs = plt.subplots(1, len(cols), figsize=(6 * len(cols), 4))
for i in range(len(cols)):
    c = cols[i]
    y = table[c].values
    axs[i].scatter(scores, y, s=1.5, alpha=0.7)
    axs[i].set_title(c)
    axs[i].set_xlabel('PoET likelihood score')
    axs[i].set_ylabel('Assay measurement')

    mask = np.isnan(y)
    r = pearsonr(scores[~mask], y[~mask])[0]
    rho = spearmanr(scores[~mask], y[~mask])[0]

    # Add correlation values on the graph
    axs[i].text(0.6, 0.15, f'Pearson r: {r:.2f}', transform=axs[i].transAxes, fontsize=12)
    axs[i].text(0.6, 0.05, f'Spearman rho: {rho:.2f}', transform=axs[i].transAxes, fontsize=12)

plt.show()

../_images/demos_AMIE_substitution_deletion_analysis_poet_31_0.png

Compare predictiveness of single prompts with the ensemble

We can also check each prompts results vs the assay data to examine the variance from prompting:

[20]:
scores = np.mean(logp, axis=-1)
sigma = np.std(logp - baseline, axis=-1)

rows = []
for i in range(logp.shape[-1]):
    pred = logp[..., i]
    row = [f'Prompt {i+1}']
    for j in range(len(cols)):
        c = cols[j]
        y = table[c].values
        mask = np.isnan(y)
        rho = spearmanr(pred[~mask], y[~mask])[0]
        row.append(rho)
    rows.append(row)

arr = pd.DataFrame(rows).iloc[:, 1:].values
means = arr.mean(axis=0)
rows.append(['Average'] + means.tolist())

stdevs = arr.std(axis=0)
rows.append(['stdevs'] + stdevs.tolist())

pred = scores
row = [f'Ensemble']
for j in range(len(cols)):
    c = cols[j]
    y = table[c].values
    mask = np.isnan(y)
    rho = spearmanr(pred[~mask], y[~mask])[0]
    row.append(rho)
rows.append(row)

ensemble_comparison_table = pd.DataFrame(rows, columns=['Prompt'] + cols)
ensemble_comparison_table.round(2)
[20]:
Prompt acetamide_normalized_fitness isobutyramide_normalized_fitness propionamide_normalized_fitness
0 Prompt 1 0.66 0.61 0.65
1 Prompt 2 0.67 0.62 0.66
2 Prompt 3 0.61 0.57 0.59
3 Prompt 4 0.65 0.62 0.64
4 Prompt 5 0.66 0.62 0.65
5 Prompt 6 0.68 0.53 0.64
6 Prompt 7 0.65 0.63 0.64
7 Prompt 8 0.63 0.62 0.62
8 Prompt 9 0.62 0.61 0.62
9 Prompt 10 0.64 0.62 0.63
10 Average 0.65 0.60 0.63
11 stdevs 0.02 0.03 0.02
12 Ensemble 0.67 0.63 0.66

Having a diversity of prompts is useful in finding a more accurate mean fitness score that is less reliant on specific sequences (that might be outliers). However we can see here that the prompts built on this MSA are all reasonable similar (stdev <= 0.03).

Compare the PoET scores with a PSSM and ESM1b

We can compare OpenProtein’s proprietary model with open source models such as PSSM and ESM1b. From this comparison we find PoET is significantly more predictive, especially for acetamide and propionamide.

[21]:
# retrieve the MSA and fit the PSSM
msa_sequences = []
for name, sequence in msa.get():
    msa_sequences.append(sequence)
msa_sequences = np.array([[c for c in msa_sequences[i]] for i in range(len(msa_sequences))])
msa_sequences.shape
[21]:
(1244, 346)

Calculate the PSSM scores

[22]:
# count the amino acids at each site and calculate site frequencies with a small pseudocount
from collections import Counter
pseudocount = 1

pssm = []
for i in range(msa_sequences.shape[1]):
    counts = Counter(msa_sequences[:, i])
    freqs = {}
    total = sum(counts.get(a, 0) + pseudocount for a in amino_acids.decode())
    for a in amino_acids.decode():
        c = counts.get(a, 0) + pseudocount
        freqs[a] = np.log(c) - np.log(total)
    pssm.append(freqs)
[23]:
# score the variants in the DMS dataset using the PSSM
scores_pssm = np.zeros(len(table))

# first, score the WT
logp_pssm_wt = 0
for j, aa in enumerate(wt.decode()):
    logp_pssm_wt += pssm[j][aa]

for i,mut in enumerate(table['mutant'].values):
    # parse the mutant sequence
    wt_aa, site, var_aa = mut[:1], int(mut[1:-1])-1, mut[-1:]
    assert wt_aa == wt[site:site+1].decode(), f'{wt_aa}, {wt[site]}, {site}'
    s = wt[:site].decode() + var_aa + wt[site+1:].decode()
    logp = 0
    for j, aa in enumerate(s):
        logp += pssm[j][aa]
    scores_pssm[i] = logp
scores_pssm = scores_pssm - logp_pssm_wt
scores_pssm.shape
[23]:
(6819,)
[24]:
_, axs = plt.subplots(1, len(cols), figsize=(6 * len(cols), 4))
for i in range(len(cols)):
    c = cols[i]
    y = table[c].values
    axs[i].scatter(scores_pssm, y, s=1.5, alpha=0.7)
    axs[i].set_title(c)
    axs[i].set_xlabel('PSSM likelihood score')
    axs[i].set_ylabel('Assay measurement')

    mask = np.isnan(y)
    r = pearsonr(scores_pssm[~mask], y[~mask])[0]
    rho = spearmanr(scores_pssm[~mask], y[~mask])[0]

    # Add correlation values on the graph
    axs[i].text(0.6, 0.15, f'Pearson r: {r:.2f}', transform=axs[i].transAxes, fontsize=12)
    axs[i].text(0.6, 0.05, f'Spearman rho: {rho:.2f}', transform=axs[i].transAxes, fontsize=12)

plt.show()

../_images/demos_AMIE_substitution_deletion_analysis_poet_40_0.png

Calculate the ESM1b scores using the API

[25]:
esm1b = session.embedding.get_model('esm1b_t33_650M_UR50S')
[26]:
# to make the ESM1b scores, we get the probability of the amino acids at each site if the site is masked
# ESM1b will return logits flanked by start and stop token positions
queries = []
for i in range(len(wt)):
    q = wt[:i] + b'X' + wt[i+1:]
    queries.append(q)

future_logits_esm1b = esm1b.logits(queries)
print(future_logits_esm1b.job)
status=<JobStatus.PENDING: 'PENDING'> job_id='eca5283d-20b8-4b71-9380-b9c28bc8a939' job_type='/embeddings/logits' created_date=datetime.datetime(2023, 11, 3, 3, 47, 11, 505991, tzinfo=TzInfo(UTC)) start_date=None end_date=None prerequisite_job_id=None progress_message=None progress_counter=0 num_records=346
[27]:

future_logits_esm1b.wait_until_done(verbose=True)
Waiting: 100%|██████████| 100/100 [00:00<00:00, 349.10it/s, status=SUCCESS]
[27]:
True
[28]:
results_esm1b = future_logits_esm1b.get()
# parse the site potentials out from the results for each masked site
logits_esm1b = [None]*len(wt)
for sequence, array in results_esm1b:
    # discard the leading <cls> token and the trailing <eos> token
    array = array[1:-1]
    # find the masked position to get the logits there
    i = sequence.index(b'X')
    logits_esm1b[i] = array[i]
logits_esm1b = np.stack(logits_esm1b, axis=0)

# parse the tokens
tokens = esm1b.metadata.output_tokens[:logits_esm1b.shape[1]]
esm1b_scoring_matrix = []
for i in range(len(wt)):
    s = scipy.special.log_softmax(logits_esm1b[i])
    site_scores = {}
    for j,t in enumerate(tokens):
        site_scores[t] = s[j]
    esm1b_scoring_matrix.append(site_scores)

logits_esm1b.shape
[28]:
(346, 33)
[29]:
# score the variants in the DMS dataset using the PSSM
scores_esm1b = np.zeros(len(table))

# first, score the WT
logp_esm1b_wt = 0
for j, aa in enumerate(wt.decode()):
    logp_esm1b_wt += esm1b_scoring_matrix[j][aa]

for i,mut in enumerate(table['mutant'].values):
    # parse the mutant sequence
    wt_aa, site, var_aa = mut[:1], int(mut[1:-1])-1, mut[-1:]
    assert wt_aa == wt[site:site+1].decode(), f'{wt_aa}, {wt[site]}, {site}'
    s = wt[:site].decode() + var_aa + wt[site+1:].decode()
    logp = 0
    for j, aa in enumerate(s):
        logp += esm1b_scoring_matrix[j][aa]
    scores_esm1b[i] = logp
scores_esm1b = scores_esm1b - logp_esm1b_wt
scores_esm1b.shape
[29]:
(6819,)
[30]:
score_map_esm1b = np.zeros((len(wt), 20))
for i in range(len(wt)):
    for j in range(20):
        aa = amino_acids[j:j+1].decode()
        score_map_esm1b[i, j] = esm1b_scoring_matrix[i][aa] - esm1b_scoring_matrix[i][wt[i:i+1].decode()]

[31]:

heatmap = scores_ensemble[..., 0] - baseline[..., 0] fig, ax = plt.subplots(figsize=(55, 10)) plt.title("Single site mutation fitness heatmap (ESM1b)", size=32) # Plot the data im = ax.imshow(score_map_esm1b.T, cmap='viridis') # Set y-axis labels y_labels = amino_acids.decode() y_ticks = np.arange(len(y_labels)) ax.set_yticks(y_ticks) ax.set_yticklabels(y_labels, fontsize=12) # Label axes plt.ylabel('Substitution', size=32) plt.xlabel('Position', size=32) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="1%", pad=0.05) plt.colorbar(im, cax=cax) plt.colorbar(im, cax=cax) plt.show()
../_images/demos_AMIE_substitution_deletion_analysis_poet_48_0.png
[32]:

fig, ax = plt.subplots(figsize=(15, 5)) plt.title("Positions 70-100 only (ESM1b)", size=18) # Plot the data im = ax.imshow(score_map_esm1b.T, cmap='viridis') # Set y-axis labels y_labels = amino_acids.decode() y_ticks = np.arange(len(y_labels)) ax.set_yticks(y_ticks) ax.set_yticklabels(y_labels, fontsize=12) # Adjust fontsize as needed # Label axes plt.ylabel('Substitution', size=32) plt.xlabel('Position', size=32) plt.xlim(70,100) plt.colorbar(im, ax=ax) plt.show()
../_images/demos_AMIE_substitution_deletion_analysis_poet_49_0.png
[33]:
_, axs = plt.subplots(1, len(cols), figsize=(6 * len(cols), 4))
for i in range(len(cols)):
    c = cols[i]
    y = table[c].values
    axs[i].scatter(scores_esm1b, y, s=1.5, alpha=0.7)
    axs[i].set_title(c)
    axs[i].set_xlabel('ESM1b likelihood score')
    axs[i].set_ylabel('Assay measurement')

    mask = np.isnan(y)
    r = pearsonr(scores_esm1b[~mask], y[~mask])[0]
    rho = spearmanr(scores_esm1b[~mask], y[~mask])[0]

    # Add correlation values on the graph
    axs[i].text(0.6, 0.15, f'Pearson r: {r:.2f}', transform=axs[i].transAxes, fontsize=12)
    axs[i].text(0.6, 0.05, f'Spearman rho: {rho:.2f}', transform=axs[i].transAxes, fontsize=12)

plt.show()

../_images/demos_AMIE_substitution_deletion_analysis_poet_50_0.png

Fitness comparison summary

We can see that the correlations between assay data and predicted fitness for acetamine, isobutyramide and propionamide with Poet are significantly better than the PSSM and ESM1b models (see the table below):

[34]:
from IPython.display import display

scores_vae = table['mutation_effect_prediction_vae_ensemble'].values
vae_mask = np.isnan(scores_vae)

score_names = ['PSSM', 'ESM1b', 'PoET']
score_arrays = [scores_pssm, scores_esm1b, scores]

rows = []
for i in range(len(cols)):
    c = cols[i]
    y = table[c].values
    mask = np.isnan(y)
    rhos = []
    for j in range(len(score_names)):
        rho = spearmanr(score_arrays[j][~mask], y[~mask])[0]
        rhos.append(rho)
    axs[i].set_title(c)

    row = [c] + rhos
    rows.append(row)

result_table = pd.DataFrame(rows, columns=['Property'] + score_names)
print("Correlation (assay vs model) results per property and model:")
display(result_table.round(2))
Correlation (assay vs model) results per property and model:
Property PSSM ESM1b PoET
0 acetamide_normalized_fitness 0.50 0.44 0.67
1 isobutyramide_normalized_fitness 0.51 0.59 0.63
2 propionamide_normalized_fitness 0.49 0.46 0.66

Analyze deletion variants of AMIE_PSEAE

Single deletion screen

Predict the deletion tolerance for each site in AMIE_PSEAE using the score function of PoET to evaluate the log-likelihoods of specified sequences.

[35]:
queries = [wt]
for i in range(len(wt)):
    q = wt[:i] + wt[i+1:]
    queries.append(q)
len(queries)
[35]:
347
[36]:
future_score = session.poet.score(prompt, queries)
print(future_score.job)

status=<JobStatus.SUCCESS: 'SUCCESS'> job_id='8d9ae707-ac82-4dd0-afab-75a1f973255e' job_type='/poet' created_date=datetime.datetime(2023, 9, 16, 17, 38, 39, 225981) start_date=datetime.datetime(2023, 9, 16, 17, 39, 57, 752930) end_date=datetime.datetime(2023, 9, 16, 17, 40, 13, 406974) prerequisite_job_id=None progress_message=None progress_counter=None num_records=None
[37]:
future_score.wait_until_done(verbose=True)
Waiting: 100%|██████████| 100/100 [00:00<00:00, 101.56it/s, status=SUCCESS]
[37]:
True
[38]:
results_del_scan = future_score.get()
[39]:
scores_del_ensemble = np.stack([r.score for r in results_del_scan], axis=0)
scores_del_ensemble.shape
[39]:
(347, 10)
[40]:
import numpy as np
import matplotlib.pyplot as plt

scores_del = np.mean(scores_del_ensemble - baseline, axis=-1)

# Create a colormap for coloring points based on height
colormap = plt.get_cmap('viridis')  # You can choose a different colormap

_, ax = plt.subplots(figsize=(20, 4))
scatter = ax.scatter(np.arange(0, len(wt) + 1), scores_del, c=scores_del, cmap=colormap)
ax.set_xlim(0, len(scores_del) + 2)
ax.set_xticks(np.arange(0, len(scores_del), 50), ['WT'] + [str(i) for i in np.arange(50, len(scores_del), 50)])
plt.xlabel('Position')
plt.ylabel('PoET Likelihood Score')
plt.title('Deletion Tolerance Scan')

# Add a colorbar to indicate the values associated with the colors
#cbar = plt.colorbar(scatter, ax=ax)
#cbar.set_label('PoET Likelihood Score')

plt.show()

../_images/demos_AMIE_substitution_deletion_analysis_poet_60_0.png

Deletion fragment screen

Now, let’s predict the deletion tolerance for fragments (subsequence) of up to 16 amino acids by enumerating and scoring the log-likelihoods for all of them. For AMIE_PSEAE, there are about 5,400 possible such deletions.

[41]:
# enumerate WT variants with all deletions of up to size 16
# NOTE - empty string sequences are not scores by the PoET API

queries = [wt]
for i in range(len(wt)):
    for j in range(i, min(i+16, len(wt))):
        q = wt[:i] + wt[j+1:]
        queries.append(q)
len(queries)
[41]:
5417
[42]:
# submit this in multiple chunks, because the number of query sequences is large
futures_score_del_span = []
batch_size = 2000
for i in range(0, len(queries), batch_size):
    batch = queries[i:i+batch_size]
    f = session.poet.score(prompt, batch)
    print(i, 'to', min(len(queries), i+batch_size), f.job)
    futures_score_del_span.append(f)
0 to 2000 status=<JobStatus.PENDING: 'PENDING'> job_id='4badf9df-3a42-4bc8-9035-daa5fe740e30' job_type='/poet' created_date=datetime.datetime(2023, 11, 3, 3, 48, 55, 361793) start_date=None end_date=None prerequisite_job_id=None progress_message=None progress_counter=None num_records=None parent_id=None s3prefix=None page_size=None page_offset=None num_rows=None result=None n_completed=None
2000 to 4000 status=<JobStatus.PENDING: 'PENDING'> job_id='66a80f62-3d39-49ba-bf91-3abb79b67c1f' job_type='/poet' created_date=datetime.datetime(2023, 11, 3, 3, 49, 0, 732854) start_date=None end_date=None prerequisite_job_id=None progress_message=None progress_counter=None num_records=None parent_id=None s3prefix=None page_size=None page_offset=None num_rows=None result=None n_completed=None
4000 to 5417 status=<JobStatus.PENDING: 'PENDING'> job_id='85e52351-c1d0-4960-8e43-b02c5e232bf8' job_type='/poet' created_date=datetime.datetime(2023, 11, 3, 3, 49, 5, 963940) start_date=None end_date=None prerequisite_job_id=None progress_message=None progress_counter=None num_records=None parent_id=None s3prefix=None page_size=None page_offset=None num_rows=None result=None n_completed=None
[43]:
# wait for all of the score jobs to finish
for f in futures_score_del_span:
    f.wait_until_done(verbose=True)
Waiting: 100%|██████████| 100/100 [07:39<00:00,  4.60s/it, status=SUCCESS]
Waiting: 100%|██████████| 100/100 [00:00<00:00, 155.09it/s, status=SUCCESS]
Waiting: 100%|██████████| 100/100 [00:00<00:00, 116.61it/s, status=SUCCESS]

Fetch the results in batches too:

[44]:
# get all of the results for the deletion region scan
# NOTE - empty string sequences are not scores by the PoET API
results_del_scan_dict = {}
for f in futures_score_del_span:
    res = f.get()
    for r in res:
        results_del_scan_dict[r.sequence] = r.score
    print('GET', len(results_del_scan_dict))
len(results_del_scan_dict)
GET 1900
GET 3758
GET 5092
[44]:
5092
[45]:
# match against the queries
scores_del_scan_ensemble = []
for q in queries:
    s = results_del_scan_dict.get(q, np.array([np.nan]*10)) # fill in NaN for "" fragments
    scores_del_scan_ensemble.append(s)
scores_del_scan_ensemble = np.stack(scores_del_scan_ensemble)
scores_del_scan_ensemble.shape
[45]:
(5417, 10)
[46]:
# visualize scores for the fragment deletions
max_length = 16
scores_mat = np.zeros((len(wt), max_length)) + np.nan
for i in range(len(wt)):
    for j in range(i, min(i+max_length, len(wt))):
        q = wt[:i] + wt[j+1:]
        s = results_del_scan_dict.get(q, [np.nan]*10) # fill in NaN for "" fragments
        s = np.array(s)
        scores_mat[i, j-i] = s.mean() - baseline.mean()
scores_mat.shape
[46]:
(346, 16)
[47]:
_, ax = plt.subplots(figsize=(20, 8))
im = ax.imshow(scores_mat.T, cmap='viridis')
ax.set_xlabel('Position')
ax.set_ylabel('Deletion Size')
ax.set_title('Deletion Size Tolerance');
../_images/demos_AMIE_substitution_deletion_analysis_poet_69_0.png
[48]:
# sort out the most favorable deletion starting from each position
order = np.argsort(-scores_mat.max(axis=1).ravel())
print('Score   ', 'Start,End', 'Length', 'Sequence', sep='\t')
for k in order[:10]:
    i = k
    j = np.argmax(scores_mat[i])
    #i, j = np.unravel_index(k, scores_mat.shape)
    span = (i+1, i+1+j)
    deleted_fragment = wt[i:i+j+1].decode()
    print(f'{scores_mat[i, j]:8.5f}', span, j+1, deleted_fragment, sep='\t')
Score           Start,End       Length  Sequence
10.30301        (331, 346)      16      QCPVGRLPYEGLEKEA
10.30301        (330, 345)      16      AQCPVGRLPYEGLEKE
 4.94593        (329, 344)      16      VAQCPVGRLPYEGLEK
 2.35968        (327, 342)      16      TGVAQCPVGRLPYEGL
 1.12998        (328, 343)      16      GVAQCPVGRLPYEGLE
 0.70272        (325, 340)      16      STTGVAQCPVGRLPYE
 0.54333        (326, 341)      16      TTGVAQCPVGRLPYEG
 0.13023        (62, 77)        16      LQGIMYDPAEMMETAV
-1.06118        (323, 337)      15      TRSTTGVAQCPVGRL
-1.06118        (322, 336)      15      LTRSTTGVAQCPVGR
[49]:
# sort out the most favorable deletion starting from each position
order = np.argsort(-scores_mat[:320].max(axis=1).ravel())
print('Score   ', 'Start,End', 'Length', 'Sequence', sep='\t')
for k in order[:10]:
    i = k
    j = np.argmax(scores_mat[i])
    #i, j = np.unravel_index(k, scores_mat.shape)
    span = (i+1, i+1+j)
    deleted_fragment = wt[i:i+j+1].decode()
    print(f'{scores_mat[i, j]:8.5f}', span, j+1, deleted_fragment, sep='\t')
Score           Start,End       Length  Sequence
 0.13023        (62, 77)        16      LQGIMYDPAEMMETAV
-1.42150        (77, 89)        13      VAIPGEETEIFSR
-1.72312        (63, 78)        16      QGIMYDPAEMMETAVA
-1.73653        (285, 297)      13      SGLQASGDGDRGL
-1.83856        (284, 296)      13      YSGLQASGDGDRG
-1.83856        (283, 295)      13      GYSGLQASGDGDR
-1.83856        (282, 294)      13      RGYSGLQASGDGD
-2.00928        (76, 89)        14      AVAIPGEETEIFSR
-2.22272        (286, 298)      13      GLQASGDGDRGLA
-2.24300        (75, 89)        15      TAVAIPGEETEIFSR

Deletion analysis summary

We can see that PoET identifies that C-terminal deletions are most favorable and are predicted to be highly tolerated. Examining the protein on PDB, 2UXY, we can see that several N-terminal residues are missing from the crystal structure and that the remaining observed N-terminal residues form a relatively unstructured loop away from the active site, validating that this deletion would be plausibly tolerated.

Note that PoET sometimes assigns higher likelihoods to protein fragments, so large deletion scans should be performed with some caution. This is because shorter sequences contain fewer amino acids that need to be explained by the model and sequence fragments often occur in natural protein databases leading to some PoET tending to prefer early stop tokens.

Summary

We have here used a combination of PoET (an OpenProtein proprietary model) and ESM, PSSM models (open-source) to look at sequence properties of aliphatic amidase. We have used PoET to examine every possible single site mutation and its predicted effect on fitness. Such an approach can be used to design novel mutants or find positions of interest. The flexibility of PoET also allows us to score variable length deletions, which seem to explain structural features of aliphatic amidase, further mutational analysis could verify these results.

Lastly, we have demonstrated that wetlab DMS data correlates well with our PoET scores (spearman r> 0.6 across 3 distinct properties), and that the DMS data is better explained by PoET (ranged of Spearman r = 0.63-0.67) than either PSSM (r = 0.49-0.5) and ESM1b (r= 0.44-0.59).

Overall, PoET is a powerful and flexible protein model to understand and exploit sequence constraints.