NumPy Broadcasting:计算两个数组之间平方差的总和

时间:2016-03-26 22:27:07

标签: python numpy numpy-broadcasting

我有以下代码。它永远在Python中。必须有办法将此计算转换为广播...

def euclidean_square(a,b):
    squares = np.zeros((a.shape[0],b.shape[0]))
    for i in range(squares.shape[0]):
        for j in range(squares.shape[1]):
            diff = a[i,:] - b[j,:]
            sqr = diff**2.0
            squares[i,j] = np.sum(sqr)
    return squares

3 个答案:

答案 0 :(得分:8)

您可以在计算np.einsum中的差异后使用broadcasted way,就像这样 -

ab = a[:,None,:] - b
out = np.einsum('ijk,ijk->ij',ab,ab)

或者使用scipy's cdist并将其可选的指标参数设置为'sqeuclidean',以便为我们提供问题所需的欧氏距离平方,如此 -

from scipy.spatial.distance import cdist
out = cdist(a,b,'sqeuclidean')

答案 1 :(得分:1)

除了使用 cdist 之外的另一个解决方案是以下

difference_squared = np.zeros((a.shape[0], b.shape[0]))
for dimension_iterator in range(a.shape[1]):
    difference_squared = difference_squared + np.subtract.outer(a[:, dimension_iterator], b[:, dimension_iterator])**2.

答案 2 :(得分:1)

我收集了这里提出的不同方法,并分为两个other questions,并测量了不同方法的速度:

import numpy as np
import scipy.spatial
import sklearn.metrics

def dist_direct(x, y):
    d = np.expand_dims(x, -2) - y
    return np.sum(np.square(d), axis=-1)

def dist_einsum(x, y):
    d = np.expand_dims(x, -2) - y
    return np.einsum('ijk,ijk->ij', d, d)

def dist_scipy(x, y):
    return scipy.spatial.distance.cdist(x, y, "sqeuclidean")

def dist_sklearn(x, y):
    return sklearn.metrics.pairwise.pairwise_distances(x, y, "sqeuclidean")

def dist_layers(x, y):
    res = np.zeros((x.shape[0], y.shape[0]))
    for i in range(x.shape[1]):
        res += np.subtract.outer(x[:, i], y[:, i])**2
    return res

# inspired by the excellent https://github.com/droyed/eucl_dist
def dist_ext1(x, y):
    nx, p = x.shape
    x_ext = np.empty((nx, 3*p))
    x_ext[:, :p] = 1
    x_ext[:, p:2*p] = x
    x_ext[:, 2*p:] = np.square(x)

    ny = y.shape[0]
    y_ext = np.empty((3*p, ny))
    y_ext[:p] = np.square(y).T
    y_ext[p:2*p] = -2*y.T
    y_ext[2*p:] = 1

    return x_ext.dot(y_ext)

# https://stackoverflow.com/a/47877630/648741
def dist_ext2(x, y):
    return np.einsum('ij,ij->i', x, x)[:,None] + np.einsum('ij,ij->i', y, y) - 2 * x.dot(y.T)

我使用timeit比较不同方法的速度。为了进行比较,我使用长度为10的向量,第一组为100个向量,第二组为1000个向量。

import timeit

p = 10
x = np.random.standard_normal((100, p))
y = np.random.standard_normal((1000, p))

for method in dir():
    if not method.startswith("dist_"):
        continue
    t = timeit.timeit(f"{method}(x, y)", number=1000, globals=globals())
    print(f"{method:12} {t:5.2f}ms")

在我的笔记本电脑上,结果如下:

dist_direct   5.07ms
dist_einsum   3.43ms
dist_ext1     0.20ms  <-- fastest
dist_ext2     0.35ms
dist_layers   2.82ms
dist_scipy    0.60ms
dist_sklearn  0.67ms

虽然两种方法dist_ext1dist_ext2都基于将(x-y)**2写为x**2 - 2*x*y + y**2的思想,但是它们却非常快,但有一个缺点: xy之间的值非常小,由于cancellation error,数值结果有时可能(非常轻微)为负。