二维数组的Python Earth Mover距离

Question

我想计算两个 2D 阵列（这些不是图像）之间的 Earth Mover 距离。

现在我浏览了两个库：scipy (https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.wasserstein_distance.html) 和 pyemd (https://pypi.org/project/pyemd/)。

#define a sampeling method
def sampeling2D(n, mu1, std1, mu2, std2):
   #sample from N(0, 1) in the 2D hyperspace
   x = np.random.randn(n, 2)

   #scale N(0, 1) -> N(mu, std)
   x[:,0] = (x[:,0]*std1) + mu1
   x[:,1] = (x[:,1]*std2) + mu2

   return x

#generate two sets
Y1 = sampeling2D(1000, 0, 1, 0, 1)
Y2 = sampeling2D(1000, -1, 1, -1, 1)

#compute the distance
distance = pyemd.emd_samples(Y1, Y2)

虽然 scipy 版本不接受 2D 数组并返回错误，但 pyemd 方法返回一个值。如果您从文档中看到，它说它只接受一维数组，所以我认为输出是错误的。在这种情况下如何计算这个距离？

Answer 1

因此，如果我理解正确，您正在尝试传输采样分布，即计算所有集群的权重为 1 的设置的距离。通常，您可以将 EMD 的计算视为最小值的实例成本流，在您的情况下，这归结为linear assignment problem：您的两个数组是二分图中的分区，两个顶点之间的权重是您选择的距离。假设您想使用欧几里得范数作为度量，则可以使用scipy.spatial.distance.cdist 获得边缘的权重，即地面距离，实际上 SciPy 在@987654329 中也提供了线性求和分配问题的求解器@（其中recently 看到了 SciPy 1.4 中提供的巨大性能改进。如果您遇到性能问题，这可能会让您感兴趣；对于 1000x1000 输入，1.3 实现有点慢）。

换句话说，你想做的事情归结为

from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment

d = cdist(Y1, Y2)
assignment = linear_sum_assignment(d)
print(d[assignment].sum() / n)

也可以使用scipy.sparse.csgraph.min_weight_bipartite_full_matching 作为linear_sum_assignment 的替代品；虽然针对稀疏输入（您的输入肯定不是），但它可能会在某些情况下提供性能改进。

验证此计算的结果是否与您从最小成本流求解器中获得的结果相匹配可能是有益的； NetworkX 中提供了一个这样的求解器，我们可以在其中手动构建图形：

import networkx as nx

G = nx.DiGraph()

# Represent elements in Y1 by 0, ..., 999, and elements in
# Y2 by 1000, ..., 1999.
for i in range(n):
    G.add_node(i, demand=-1)
    G.add_node(n + i, demand=1)

for i in range(n):
    for j in range(n):
        G.add_edge(i, n + j, capacity=1, weight=d[i, j])

此时，我们可以验证上述方法是否符合最小成本流：

In [16]: d[assignment].sum() == nx.algorithms.min_cost_flow_cost(G)
Out[16]: True

同样，对于一维输入，看到结果与scipy.stats.wasserstein_distance 一致也很有启发意义：

from scipy.stats import wasserstein_distance

np.random.seed(0)
n = 100

Y1 = np.random.randn(n)
Y2 = np.random.randn(n) - 2
d =  np.abs(Y1 - Y2.reshape((n, 1)))

assignment = linear_sum_assignment(d)
print(d[assignment].sum() / n)       # 1.9777950447866477
print(wasserstein_distance(Y1, Y2))  # 1.977795044786648