python - `numpy.random.multivariate_normal` 的向量化实现

Question

我正在尝试使用 numpy.random.multivariate_normal 生成多个样本，其中每个样本都来自具有不同 mean 和 cov 的多元正态分布。比如我想画2个样本，我试过

from numpy import random as rand

means = np.array([[-1., 0.], [1., 0.]])
covs = np.array([np.identity(2) for k in xrange(2)]) 
rand.multivariate_normal(means, covs)

但这会导致 ValueError: mean must be 1 dimensional。我必须为此做一个for循环吗？我认为对于像 rand.binomial 这样的函数，这是可能的。

Answer 1

正如@hpaulj 所建议的那样，您可以从标准多元正态分布生成样本，然后使用例如 einsum 和/或广播来转换样本。缩放是通过将标准样本点乘以协方差矩阵的平方根来完成的。在下文中，我使用 scipy.linalg.sqrtm 计算矩阵平方根，并使用 numpy.einsum 进行矩阵乘法。

import numpy as np
from scipy.linalg import sqrtm
import matplotlib.pyplot as plt


# Sequence of means
means = np.array([[-15., 0.], [15., 0.], [0., 0.]])
# Sequence of covariance matrices.  Must be the same length as means.
covs = np.array([[[ 3, -1],
                  [-1,  2]],
                 [[ 1,  2],
                  [ 2,  5]],
                 [[ 1,  0],
                  [ 0,  1]]])
# Number of samples to generate for each (mean, cov) pair.
nsamples = 4000

# Compute the matrix square root of each covariance matrix.
sqrtcovs = np.array([sqrtm(c) for c in covs])

# Generate samples from the standard multivariate normal distribution.
dim = len(means[0])
u = np.random.multivariate_normal(np.zeros(dim), np.eye(dim),
                                  size=(len(means), nsamples,))
# u has shape (len(means), nsamples, dim)

# Transform u.
v = np.einsum('ijk,ikl->ijl', u, sqrtcovs)
m = np.expand_dims(means, 1)
t = v + m

# t also has shape (len(means), nsamples, dim).
# t[i] holds the nsamples sampled from the distribution with mean means[i]
# and covariance cov[i].

plt.subplot(2, 1, 1)
plt.plot(t[...,0].ravel(), t[...,1].ravel(), '.', alpha=0.02)
plt.axis('equal')
plt.xlim(-25, 25)
plt.ylim(-8, 8)
plt.grid()

# Make another plot, where we generate the samples by passing the given
# means and covs to np.random.multivariate_normal.  This plot should look
# the same as the first plot.
plt.subplot(2, 1, 2)
p0 = np.random.multivariate_normal(means[0], covs[0], size=nsamples)
p1 = np.random.multivariate_normal(means[1], covs[1], size=nsamples)
p2 = np.random.multivariate_normal(means[2], covs[2], size=nsamples)

plt.plot(p0[:,0], p0[:,1], 'b.', alpha=0.02)
plt.plot(p1[:,0], p1[:,1], 'g.', alpha=0.02)
plt.plot(p2[:,0], p2[:,1], 'r.', alpha=0.02)
plt.axis('equal')
plt.xlim(-25, 25)
plt.ylim(-8, 8)
plt.grid()

这种方法可能不会比遍历 means 和 covs 数组并为每对调用一次 multivariate_normal 更快(mean，cov ).当您有许多 不同的均值和协方差并且每对生成少量样本时，此方法会发挥最大作用。即便如此，它也可能不会更快，因为该脚本在 covs 数组上使用 Python 循环来为每个协方差矩阵调用 sqrtm。如果性能至关重要，请使用您的实际数据进行测试。