ytdistmat-fast.py

def get_distmat(X, eps=1e-7):
    """
    A fast method for computing the distance matrix of a Euclidean
    point cloud
    
    Parameters
    ----------
    X: ndarray(N, d)
      Point cloud
    eps: float
      Value below which to clamp squared distance before taking the square root
    
    Returns
    -------
    ndarray(N, N)
      Euclidean distance matrix
    """
    XSqr = np.sum(X**2, axis=1, keepdims=True)
    DSqr = XSqr + XSqr.T - 2*X.dot(X.T) # First term, second term, third term
    DSqr[DSqr < eps] = eps # Make sure we don't underflow before taking the square root!
    return np.sqrt(DSqr)