<jupyter_start><jupyter_text>Embedding Antibody Sequences using SapiensTry out this notebook in your browser: [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/Merck/Sapiens/main?labpath=notebooks%2F02_sapiens_antibody_embedding.ipynb)<jupyter_code>import sapiens<jupyter_output>/Users/prihodad/miniconda3/envs/pip-sapiens/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm<jupyter_text>Get antibody embeddings Get a vector representation of each position in a sequence<jupyter_code>residue_embed = sapiens.predict_residue_embedding(
    "QVKLQESGAELARPGASVKLSCKASGYTFTNYWMQWVKQRPGQGLDWIGAIYPGDGNTRYTHKFKGKATLTADKSSSTAYMQLSSLASEDSGVYYCARGEGNYAWFAYWGQGTTVTVSS", 
    "H", 
    layer=None
)
residue_embed.shape
# (layer, position in sequence, features)<jupyter_output><empty_output><jupyter_text>Get a single vector for each sequence<jupyter_code>seq_embed = sapiens.predict_sequence_embedding(
    "QVKLQESGAELARPGASVKLSCKASGYTFTNYWMQWVKQRPGQGLDWIGAIYPGDGNTRYTHKFKGKATLTADKSSSTAYMQLSSLASEDSGVYYCARGEGNYAWFAYWGQGTTVTVSS", 
    "H", 
    layer=None
)
seq_embed.shape
# (layer, features)
# Sequence embedding from the last layer of our model
seq_embed[-1]<jupyter_output><empty_output><jupyter_text>Embedding sequences from CovAbDab<jupyter_code>import sys
import pandas as pd
import numpy as np
!curl http://opig.stats.ox.ac.uk/webapps/covabdab/static/downloads/CoV-AbDab_310122.csv > CovAbDab.csv
def get_species(row):
    for column in ['Heavy V Gene', 'Heavy J Gene', 'Light V Gene', 'Light J Gene']:
        if type(row[column]) == str and "(" in row[column]:
            species = row[column][row[column].find("(")+1:row[column].find(")")]
            return species
        return np.nan
cov = pd.read_csv("CovAbDab.csv")
cov["Species"] = cov.apply(get_species, axis=1)
cov["VH or VHH"] = cov["VH or VHH"].str.replace(' ', '') # remove whitespace
cov = cov[cov["VH or VHH"].str.len() < 140] # remove long sequences
print(f'{len(cov):,} sequences')
cov.head()
cov['Species'].value_counts()<jupyter_output><empty_output><jupyter_text>Embed all sequences<jupyter_code>%%time

embeds = cov["VH or VHH"].apply(lambda seq: pd.Series(sapiens.predict_sequence_embedding(seq, chain_type="H", layer=-1)))
embeds.head()<jupyter_output>CPU times: user 1min 37s, sys: 9.36 s, total: 1min 46s
Wall time: 1min 33s<jupyter_text>Visualize Embeddings<jupyter_code># Uncomment to install seaborn and scikit-learn
#
#!pip install seaborn scikit-learn
from sklearn.manifold import TSNE
import seaborn as sns
import matplotlib.pyplot as plt
%%time

tsne = TSNE(n_components=2, init='pca', learning_rate='auto', perplexity=10)
X_embedded = tsne.fit_transform(embeds)

df = cov.copy()
df.insert(0, 't-SNE 1', X_embedded[:,0])
df.insert(1, 't-SNE 2', X_embedded[:,1])
df.head()
plt.figure(figsize=(12,12))
sns.scatterplot(
    x="t-SNE 1", y="t-SNE 2",
    hue="Species",
    data=df,
    legend="full"
)
plt.show()
v_genes = df["Heavy V Gene"].apply(lambda s: s.strip().split(' ')[0])
common_v_genes = v_genes.value_counts()[:30].index
common_v_genes
genes_by_family = common_v_genes.groupby(common_v_genes.map(lambda g: g.split('-')[0].split('S')[0]))
genes_by_family
scales = {
    'IGHV1': 'Blues',
    'IGHV2': 'Purples',
    'IGHV3': 'Reds',
    'IGHV4': 'Greens',
    'IGHV5': 'Oranges',
}
palette = {'Other': '#999999', 'ND': '#444444'}
for family, genes in genes_by_family.items():
    if family == 'ND':
        continue
    for i, gene in enumerate(genes):
        palette[gene] = plt.get_cmap(scales[family])(i / len(genes) * 0.7 + 0.4)
fig = plt.figure(figsize=(12,12))
ax = sns.scatterplot(
    x="t-SNE 1", y="t-SNE 2",
    hue=v_genes.apply(lambda g: g if g in common_v_genes else 'Other'),
    hue_order=palette.keys(),
    style='Species',
    palette=palette,
    data=df,
    legend="full"
)
ax.legend(loc='upper left', bbox_to_anchor=(1, 1))
plt.show()
fig.savefig('Embedding_t-SNE.png', bbox_inches='tight');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1464061746261457, 'fraction_numerical': 0.02001929570670526, 'mean_word_length': 5.645833333333333, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '266889', 'license': 'MIT', 'n_tokens_mistral': 1747, 'n_tokens_neox': 1631, 'n_words': 300, 'path': 'notebooks/02_sapiens_antibody_embedding.ipynb', 'repo_name': 'Merck/Sapiens', 'size': '4146'}