speedup-distance-computation.py

"""Speed up distance computation."""

import numpy as np


######################################################################
#                 _   _                            _                 #
#                | \ | | ___  _ __ _ __ ___   __ _| |                #
#                |  \| |/ _ \| '__| '_ ` _ \ / _` | |                #
#                | |\  | (_) | |  | | | | | | (_| | |                #
#                |_| \_|\___/|_|  |_| |_| |_|\__,_|_|                #
#                                                                    #
######################################################################


def norm_euclidean(A, B):
    """Compute Euclidean distance row-wise on two matrices.

    Parameters
    ----------
    A : np.ndarray, shape=(n_samples, n_features)
    B : np.ndarray, shape=(n_samples, n_features)

    Returns
    -------
    np.ndarray
        Row-wise distance matrix with shape (n_samples_of_A, n_samples_of_B).
    """
    return np.sqrt(((A[:, :, None] - B[:, :, None].T)**2).sum(1))


######################################################################
#              ____                      _ _   _                     #
#             / ___| _ __   ___  ___  __| | | | |_ __                #
#             \___ \| '_ \ / _ \/ _ \/ _` | | | | '_ \               #
#              ___) | |_) |  __/  __/ (_| | |_| | |_) |              #
#             |____/| .__/ \___|\___|\__,_|\___/| .__/               #
#                   |_|                         |_|                  #
#                                                                    #
######################################################################


def _ext_A(A):
    """Extend A to [I, A, A^2]

    """
    nA, dim = A.shape
    A_ext = _get_ones(nA, dim * 3)
    A_ext[:, dim:2 * dim] = A
    A_ext[:, 2 * dim:] = A ** 2
    return A_ext


def _ext_B(B):
    """Extend B to [B^2^T, -2B^T, I]^T

    """
    nB, dim = B.shape
    B_ext = _get_ones(dim * 3, nB)
    B_ext[:dim] = (B ** 2).T
    B_ext[dim: 2 * dim] = -2.0 * B.T
    return B_ext


def _euclidean(A_ext, B_ext):
    sqdist = A_ext.dot(B_ext).clip(min=0)
    return np.sqrt(sqdist)


# the Cosine distance is nothing more but the “normed” Euclidean distance.
# It can be simply built on top of _euclidean function by feeding _norm(A)
# and _norm(B) to it:
def _norm(A):
    return A / np.linalg.norm(A, ord=2, axis=1, keepdims=True)


def _cosine(A_norm, B_norm):
    return _euclidean(A_norm, B_norm)