使用 numpy 向量化方程答案

【问题标题】：Vectorising an equation using numpy使用 numpy 向量化方程
【发布时间】：2014-06-28 19:54:15
【问题描述】：

我正在尝试将上述公式实现为矢量化形式。 K=3 这里，X 是 150x4 numpy 数组。 mu 是 3x4 numpy 数组。 Gamma 是一个 150x3 numpy 数组。 Sigma 是一个 kx4x4 numpy 数组。因此Sigma[k] 是一个4x4 numpy 数组。 N=150

N_k = np.sum(Gamma, axis=0)
for k in range(K): # Correct
         x_new = X - mu[k] #Correct
         a = np.dot(x_new.T, x_new) #Incorrect from here I feel
         for i in range(len(data)):
             sigma[k] = Gamma[i][k] * a
         sigma[k]=sigma[k]/N_k #totally incorrect

如何解决这个问题？

【问题讨论】：

标签： python python-2.7 python-3.x numpy scipy

【解决方案1】：

产品的总和？听起来像是np.einsum 的工作：

import numpy as np
N = 150
K = 3
M = 4
x = np.random.random((N,M))
mu = np.random.random((K,M))
gamma = np.random.random((N,K))

xbar = x-mu[:,None,:] # shape (3, 150, 4)
sigma = np.einsum('nk,knm,kno->kmo', gamma, xbar, xbar)
sigma /= gamma.sum(axis=0)[:,None,None]

解码'nk,knm,kno->kmo'：

此下标规范在数组的左侧 (->) 具有三个组件，其后跟在右侧的一个组件。

左边的三个组件对应gamma、xbar和xbar的下标，操作数被传递给np.einsum。

gamma 具有下标nk，就像您发布的公式中一样。 xbar 的形状为 (3, 150, 4)。您可以将其视为具有下标knm，其中k 和n 与您发布的公式中的含义相同，而m 是表示长度为4 的轴的下标，在你的公式，但显然有你对数组形状的描述。

现在第三个下标组件是kno。使用o 下标是因为o 与m 下标扮演相同的角色，但我们不希望对m 求和。事实上，我们希望 m 和 o 下标能够独立迭代，而不是同步迭代。因此我们给第三个下标不同的字母。

请注意，n 出现在左侧的下标中 (nk, knm, kno)，但没有出现在右侧的下标中 (kmo)。这告诉np.einsum 求和n。

k 出现在左侧和右侧的下标中。这告诉np.einsum，我们希望将k 下标同步推进，但是（因为它出现在右侧）我们不想对k 求和。

由于kmo 出现在右侧，这些下标保留在结果中。这导致 sigma 的形状为 (K,M,M)（即 (3,4,4)）。

【讨论】：

这真是太棒了。谢谢你教我一些新东西！

【解决方案2】：

在 unubtu 的出色答案之上仅提供几个性能指标。

np.einsum 无法优化具有两个以上参数的调用。只要有可能，手动将计算分成两个参数组通常会更快，例如：

def unubtu():
    xbar = x-mu[:,None,:] # shape (3, 150, 4)
    sigma = np.einsum('nk,knm,kno->kmo', gamma, xbar, xbar)
    sigma /= gamma.sum(axis=0)[:,None,None]
    return sigma

def faster():
    xbar = x-mu[:,None,:] # shape (3, 150, 4)
    sigma = np.einsum('knm,kno->kmo', gamma.T[..., None] * xbar, xbar)
    sigma /= gamma.sum(axis=0)[:,None,None]
    return sigma

In [50]: %timeit unubtu()
10000 loops, best of 3: 147 µs per loop

In [51]: %timeit faster()
10000 loops, best of 3: 129 µs per loop

12% 的改进并不多，但随着阵列的增大，差异会变得（很多）更大。

此外，即使np.einsum 是一个很棒的工具，它使非常困难的事情变得简单，但如果你的 numpy 是用一个好的线性代数库构建的，它绝不像np.dot 那样优化。在您的情况下，鉴于 K 很小，使用 np.dot 和 Python 循环甚至更快：

def even_faster():
    sigma = np.empty((K, M, M))
    for k in xrange(K):
        x_ = x - mu[k]
        sigma[k] = np.dot((x_ * gamma[:, k, None]).T, x_)
    sigma /= gamma.sum(axis=0)[:,None,None]
    return sigma

In [52]: %timeit even_faster()
10000 loops, best of 3: 101 µs per loop

【讨论】：