<jupyter_start><jupyter_text>Helpers<jupyter_code>def umap_fit_transform(X: pd.DataFrame, **kwargs) -> pd.DataFrame:
    
    # Change defaults
    kwargs["verbose"] = kwargs.get("verbose", True)
    kwargs["low_memory"] = kwargs.get("low_memory", True)

    umap_transformer = umap_package.UMAP(**kwargs)

    X_umap = umap_transformer.fit_transform(X).astype("float32")

    # return as DataFrame
    return pd.DataFrame(
        data=X_umap,
        index=X.index,
        columns=[f"UMAP_{i}" for i in range(X_umap.shape[1])],
    )<jupyter_output><empty_output><jupyter_text>Load data<jupyter_code>swiss_df = pd.read_hdf('../data/bacterial_swissprot.h5')
swiss_df.head()
embed_df = pd.read_hdf('../data/reduced_deep_embeddings.h5')
# embed_df = pd.read_hdf('../data/reduced_3mers_embeddings.h5')
embedding_columns = embed_df.columns

embed_df.head()
swiss_df = swiss_df.join(embed_df)
annotations = [
#     'Phylum',
#     'Order',
#     'Family',
#     'Genus',
    'SUPFAM',
    'Gene3D',
    'InterPro',
    'KO',
    'GO',
    'eggNOG',
    'Pfam',
    'EC number'
]<jupyter_output><empty_output><jupyter_text>Get EC 2.7.2 proteins<jupyter_code>ec_272_pattern = "2.7.2.[\d]"
ec_272_df = supercluster_df = swiss_df.loc[
    swiss_df['EC number'].fillna('').str.match(ec_272_pattern)
]
print(supercluster_df.shape)
supercluster_df.head()<jupyter_output>(1311, 73)<jupyter_text>Get rid of rare pfam architectures<jupyter_code>vc = ec_272_df['Pfam'].value_counts()
top_pfams = vc[vc > 3].index
ec_272_df = ec_272_df[ec_272_df['Pfam'].isin(top_pfams)]
ec_272_df['Pfam'].value_counts()
ec_272_df = ec_272_df.query('`EC number` != "2.7.2.4||1.1.1.3"')
ec_272_df['EC number'].value_counts()<jupyter_output><empty_output><jupyter_text>UMAP on EC 2.7.2<jupyter_code>umap_df = umap_fit_transform(ec_272_df.filter(like='pca_'), n_neighbors=15, min_dist=0.5)
umap_df.head()
ec_272_df = ec_272_df.join(umap_df)<jupyter_output><empty_output><jupyter_text>Plot UMAP scatterplot<jupyter_code>ec_272_df['Domains'] = ec_272_df['Pfam'].map({
    'PF00871': 'PF00871',
    'PF00696': 'PF00696',
    'PF00162': 'PF00162',
    'PF00696||PF01472': 'PF00696 and PF01472'
})
ec_272_df['Domains'] = 'Including ' + ec_272_df['Domains']
palette = [
    "#8d5eac",
    "#3b66b6",
    "#d3759e",
    "#2fc7bd",
    "#22967b",
    "#a22f6f",
    "#d9b063",
    "#8c6d34",
]
plt.figure(figsize=(12,8))
sns.scatterplot(
    data=ec_272_df,
    x='UMAP_0',
    y='UMAP_1',
    hue='EC number',
    style='Domains',
    alpha=0.8,
    s=200,
    palette=reversed(palette)
)

plt.legend(bbox_to_anchor=(1, 1.03))<jupyter_output><empty_output><jupyter_text>Compare distances (ArdiMiPE vs. Sequence-based)  Random proteins as a background<jupyter_code>random_df = swiss_df.loc[swiss_df.index.difference(ec_272_df.index)].sample(500, random_state=0)
print(random_df.shape)<jupyter_output>(500, 73)<jupyter_text>Generate ArdiMiPE based distance<jupyter_code>from sklearn.metrics import euclidean_distances

def distance(df):
    return pd.DataFrame(
        data=euclidean_distances(df),
        index=df.index,
        columns=df.index
    )

def post_process(distance_matrix):    
    distance_matrix = distance_matrix.stack()
    distance_matrix.index.names = ['protein_a', 'protein_b']
    return distance_matrix
proteins_df = pd.concat([ec_272_df, random_df])
ardimipe_distances = post_process(distance(proteins_df[embedding_columns])).to_frame('ArdiMiPE')
ardimipe_distances.head()<jupyter_output><empty_output><jupyter_text>Generate sequence based distance using Clustal Omega<jupyter_code># !apt install timelimit
# %conda install -c bioconda clustalo

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
tmp_dir = Path(tempfile.mkdtemp())

sequence_file = str(tmp_dir / 'cluster.fasta')
distance_matrix = str(tmp_dir / 'distamat.txt')
clustering = str(tmp_dir / 'clustering.txt')
alignment = str(tmp_dir / 'alignment.txt')

clustalw_out = str(tmp_dir / 'clustalw.fasta')
SeqIO.write([
    SeqRecord(Seq(row['sequence']), id=row_id, description=row['description'])
    for row_id, row in proteins_df.iterrows()
], sequence_file, "fasta")
!timelimit -t 1500 clustalo --infile $sequence_file \
                          --seqtype=Protein \
                          --distmat-out $distance_matrix \
                          --clustering-out=$clustering \
                          --outfile=$alignment \
                          --threads=16 \
                          --percent-id \
                          --full -v<jupyter_output>Using 16 threads
Read 1795 sequences (type: Protein) from /tmp/tmpbw7hswr9/cluster.fasta
Calculating pairwise ktuple-distances...
Pairwise distance matrix written to /tmp/tmpbw7hswr9/distamat.txt0)
Ktuple-distance calculation progress done. CPU time: 262.46u 0.20s 00:04:22.65 Elapsed: 00:00:32
Guide-tree computation done.
^Cogressive alignment progress: 0 % (1 out of 1794)
timelimit: terminated by signal 2 during the run period<jupyter_text>Parse clustal results<jupyter_code>clustal_dist = 100 - pd.read_csv(distance_matrix, header=None, index_col=0, skiprows=1, sep='\s+')

clustal_dist.columns = clustal_dist.index
clustal_distances = post_process(clustal_dist).to_frame('Clustal')
clustal_distances.head()<jupyter_output><empty_output><jupyter_text>Merge distances<jupyter_code>distances = pd.concat([ardimipe_distances, clustal_distances], axis=1).reset_index()
distances
annotations = ['EC number', 'Pfam']
distances = (distances
 .join(swiss_df[annotations], on='protein_a')
 .join(swiss_df[annotations], on='protein_b', rsuffix='_b')
)<jupyter_output><empty_output><jupyter_text>Clean EC classes (Set EC number from random proteins to "Random")<jupyter_code>distances.loc[distances['protein_a'].isin(random_df.index), ['EC number', 'Pfam']] = 'Random'
distances.loc[distances['protein_b'].isin(random_df.index), ['EC number_b', 'Pfam_b']] = 'Random'
distances.head()<jupyter_output><empty_output><jupyter_text>Plot Heatmap<jupyter_code>order = ['2.7.2.1', '2.7.2.15', '2.7.2.7', '2.7.2.2','2.7.2.8','2.7.2.4','2.7.2.11','2.7.2.3', 'Random']
# order = ['2.7.2.1', '2.7.2.2', '2.7.2.3', '2.7.2.4','2.7.2.7','2.7.2.8','2.7.2.11','2.7.2.15', 'Random']
def format_annot(mean, std):
    return f'{mean:.2f}\n± {std:.2f}'

def plot_distance_heatmap(distance_df, distance_col, upper=False, cmap='Reds'):
    
    # Get distance per group
    grouped_distance = distance_df.groupby(['EC number', 'EC number_b'])[distance_col].mean()
    std = distance_df.groupby(['EC number', 'EC number_b'])[distance_col].std()
        
    annot_df = pd.concat([grouped_distance, std], axis=1).apply(lambda x: format_annot(*x), axis=1)

    plt.figure(figsize=(11,11))
    mask = np.zeros_like(grouped_distance.unstack())
    
    if upper:
        mask[np.tril_indices_from(mask, k=-1)] = True
    else:
        mask[np.triu_indices_from(mask, k=1)] = True

    sns.heatmap(
        grouped_distance.unstack().loc[order, order],
        annot=True,
    #     annot=annot_df.unstack(),
#         fmt='',
#         annot_kws=dict(fontsize=15),
        cbar=False,
        mask=mask,
        cmap=cmap,
        linewidths=1,
    )
    
#     plt.xlabel('EC number')
    plt.xlabel(None)
    plt.ylabel(None)
    
    ax = plt.gca()
    
    if upper:
        ax.yaxis.set_ticks_position('right')
        ax.yaxis.set_label_position('right')
        ax.xaxis.set_ticks_position('top')
        ax.xaxis.set_label_position('top')
        
        plt.xticks(rotation='vertical')
        plt.yticks(rotation='horizontal')
        
    return ax
sns.set_theme(
    context="poster", style="white", rc=rc
)<jupyter_output><empty_output><jupyter_text>ArdiMiPE<jupyter_code>plot_distance_heatmap(distances, 'ArdiMiPE', cmap='Greens', upper=False)
plt.ylabel('Embeddings distance', labelpad=10)<jupyter_output><empty_output><jupyter_text>Clustal<jupyter_code>plot_distance_heatmap(distances, 'Clustal', upper=True)
plt.xlabel('Sequence-based distance [%]', labelpad=10)
plt.title('Distances between 2.7.2 groups', pad=100)<jupyter_output><empty_output><jupyter_text>Plot jointplot<jupyter_code>distances['class'] = (distances['EC number'] == distances['EC number_b']).map({True:'Inside cluster', False: 'Cross cluster'})
distances.loc[(distances['EC number_b'] == 'Random'), 'class'] = 'Random'

# distances['class'] = (distances['Pfam'] == distances['Pfam_b']).map({True:'Inside cluster', False: 'Cross cluster'})
# distances.loc[(distances['Pfam_b'] == 'Random'), 'class'] = 'Random'
sns.set_theme(
    context="poster", style="ticks", rc=rc
)

g = sns.JointGrid(   
    data=distances.query('(`EC number` == "2.7.2.1") & (protein_a != protein_b)'),
    x='Clustal',
    y='ArdiMiPE',
    hue='class',
    hue_order=['Random', 'Cross cluster', 'Inside cluster'],
    size=14,
    palette=palette[4:7]
)
# Add the joint and marginal histogram plots
g.plot_joint(
    sns.histplot, discrete=(False, False), legend=False, pmax=.8, alpha=0.95,
)
g.plot_marginals(sns.histplot, element="step", stat="frequency", common_norm=False)

plt.legend(
    title='Distances',
    labels=[
        'between 2.7.2.1 proteins',
        'from 2.7.2.1 to 2.7.2 protein',
        'from 2.7.2.1 to random SwissProt proteins',
    ],
    loc='upper left',
    bbox_to_anchor=(-6, 1)
)
g.ax_joint.grid(alpha=0.25)
g.ax_joint.set_xlabel('Sequence-based distance [%]')
g.ax_joint.set_ylabel('Embeddings distance')<jupyter_output>/opt/conda/lib/python3.7/site-packages/seaborn/axisgrid.py:1556: UserWarning: The `size` parameter has been renamed to `height`; please update your code.
  warnings.warn(msg, UserWarning)
{'fraction_non_alphanumeric': 0.15957557474253614, 'fraction_numerical': 0.03963382918963903, 'mean_word_length': 4.440860215053763, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 63, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 22, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '16', 'dir': 'jupyter-structured-clean-dedup', 'id': '338519', 'license': 'MIT', 'n_tokens_mistral': 3928, 'n_tokens_neox': 3662, 'n_words': 552, 'path': '03-ec-2.7.2/ec-2.7.2-analysis.ipynb', 'repo_name': 'ardigen/microbiome-protein-embeddings', 'size': '9613'}