principal-coordinates-analysis.py

import numpy as np
import matplotlib.pyplot as plt
from skbio.stats.ordination import pcoa
from skbio import DistanceMatrix

from scipy.spatial.distance import pdist, squareform

# embedding = < your language model embedding > # dim L x D where L is the length, D is the dimension
rr_dist = squareform(pdist(embedding))
dm = DistanceMatrix(rr_dists_d.mean(2))
ord_res = pcoa(dm)
cm = plt.cm.get_cmap('viridis')

sc = plt.scatter(ord_res.samples['PC1'], ord_res.samples['PC2'], 
                 c=np.arange(ord_res.samples.shape[0]), cmap=cm)
plt.colorbar(sc)

plt.xlabel('PC1')
plt.ylabel('PC2')
plt.title('Residue PCA')