如何在 numpy 中向量化矩阵和的一小部分（期望最大化）？

Question

我正在尝试使用 numpy 对二维高斯分布的以下期望最大化/聚类方程进行向量化。我有一个天真的方法，我将在问题末尾添加：

对于上下文，变量和维度定义如下：

• = 数据点索引（即 1-1000）
• = 簇索引（即 1-3）
• = 数据点 位于集群 中的条件概率（在 [0,1] 中）
• = 数据点的值 （形状 (2,)）
• = 集群当前估计的多变量均值（形状 (2,)）

最终产品是一个分子，它是 (2, 2) 形状矩阵的和，分母是一个标量。最终值是 (2, 2) 协变量矩阵估计。对于“k”的每个值（1、2、3）也必须执行此操作。

我通过定义以下 numpy 数组实现了其他值的矢量化方法：

• = 每个数据点、簇的估计概率值
• = 多元数据矩阵
• = 估计集群意味着

我的幼稚代码如下：

for kk in range(k):
    numsum = 0
    for ii in range(X.shape[0]):
        diff = (X[ii, :]-mu[kk, :]).reshape(-1, 1)
        numsum = numsum + Z[ii, kk]*np.matmul(diff, diff.T)
    sigma[kk] = numsum / np.sum(Z[:, kk])

长话短说 - 有没有更好的方法来做到这一点？

Answer 1

以下应该有效：

diff = X[np.newaxis, :, :] - mu[:, np.newaxis, :]  # kxnx2
numsum = np.matmul(Z.T[:, np.newaxis, :] * diff.transpose(0, 2, 1), diff)  # kx2x2
sigma_proposed = numsum / Z.sum(axis=0)[:, np.newaxis, np.newaxis]  # kx2x2

总之，我用以下代码检查了它：

import numpy as np

n, k = 1000, 3

# Create some data
rand = np.random.default_rng(seed=0xC0FFEE)  # For reproducibility
Z = rand.uniform(size=(n, k))
X = rand.normal(size=(n, 2))
mu = rand.normal(size=(k, 2))
sigma = np.zeros((k, 2, 2))

# Code from question
for kk in range(k):
    numsum = 0
    for ii in range(X.shape[0]):
        diff = (X[ii, :]-mu[kk, :]).reshape(-1, 1)
        numsum = numsum + Z[ii, kk]*np.matmul(diff, diff.T)
    sigma[kk] = numsum / np.sum(Z[:, kk])
    
# Proposed
diff = X[np.newaxis, :, :] - mu[:, np.newaxis, :]  # kxnx2
numsum = np.matmul(Z.T[:, np.newaxis, :] * diff.transpose(0, 2, 1), diff)  # kx2x2
sigma_proposed = numsum / Z.sum(axis=0)[:, np.newaxis, np.newaxis]  # kx2x2

assert np.allclose(sigma, sigma_proposed)