Fast Euclidean Distance Matrix Computaiton

EDMExample.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def getEDM(X):
    """
    Compute an all pairs distance matrix using broadcasting
    and matrix multiplications to speed up computations
    Parameters
    ----------
    X: ndarray(N, d)
        N points in d dimensions
    
    Returns: ndarray(N, N)
        All pairs distances
    """
    dotX = np.reshape(np.sum(X*X, 1), (X.shape[0], 1))
    D = (dotX + dotX.T) - 2*(np.dot(X, X.T))
    D[D < 0] = 0
    D = np.sqrt(D)
    return D

if __name__ == '__main__':
    # Create a p-q torus knot in 3D with N samples
    N = 500
    p = 3
    q = 5
    t = np.linspace(0, 2*np.pi, N+1)[0:N]
    R = 3
    r = 1
    X = np.zeros((N, 3))
    X[:, 0] = (R + r*np.cos(p*t))*np.cos(q*t)
    X[:, 1] = (R + r*np.cos(p*t))*np.sin(q*t)
    X[:, 2] = r*np.sin(p*t)
    
    # Compute and plot the SSM
    D = getEDM(X)

    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(121, projection='3d')
    ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=np.arange(N))
    plt.title("%i Samples of The %i-%i Torus Knot"%(N, p, q))
    plt.subplot(122)
    plt.imshow(D, cmap='magma_r')
    plt.colorbar()
    plt.scatter(np.zeros(N), np.arange(N), c=np.arange(N))
    plt.scatter(np.arange(N), np.zeros(N), c=np.arange(N))
    plt.xlabel("Sample Number")
    plt.ylabel("Sample Number")
    plt.title("Euclidean Distance Matrix")
    plt.show()