如何向量化实现numpy.random.multivariate_normal批量生成不同参数样本？

哈哈，这个坑我踩过！其实不用非得写那种笨拙的for循环（不过列表推导式这种简洁的循环写法其实挺香的），先给你解释下为啥报错，再给你几个可行的方案～

首先，为啥你会碰到ValueError: mean must be 1 dimensional这个错误？
numpy.random.multivariate_normal的参数设计有明确规则：它要么接受单个均值向量（1维数组）+单个协方差矩阵（2维数组），生成指定数量的同分布样本；要么接受多个均值向量（2维数组，每行一个均值）+单个协方差矩阵，生成和均值数量一致的样本，但所有样本共享同一个协方差。

而你的代码里同时传了2维的均值数组和3维的协方差数组（每个样本对应不同的协方差），这超出了函数的支持范围，所以直接报错了。

那有没有不用显式写for循环的方法？其实分几种情况：

1. 最推荐：简洁的列表推导式（本质是循环，但可读性拉满） #

这种写法虽然是循环的一种，但代码简洁易懂，对于大多数场景来说效率完全够用：

import numpy as np
from numpy import random as rand

means = np.array([[-1., 0.], [1., 0.]])
covs = np.array([np.identity(2) for _ in range(2)])

# 遍历每个均值和协方差，生成对应样本
samples = np.array([rand.multivariate_normal(mu, cov) for mu, cov in zip(means, covs)])
print(samples)

运行后就能得到2个分别对应不同均值和协方差的样本啦。

2. 语法糖式的向量化写法：np.vectorize #

如果你不想看到for关键字，可以用np.vectorize把生成单个样本的函数包装成向量化函数，不过本质上还是底层循环，效率和列表推导式差不多：

def generate_single_sample(mu, cov):
    return rand.multivariate_normal(mu, cov)

# 用signature指定输入输出的维度，避免维度混乱
vec_generator = np.vectorize(generate_single_sample, signature='(n),(n,n)->(n)')
samples = vec_generator(means, covs)

3. 极致性能方案：用numba加速循环 #

如果你的样本量特别大，追求极致性能，可以用numba把循环编译成机器码，速度会快很多：

from numba import jit

@jit(nopython=True)
def generate_samples_fast(means, covs):
    n_samples = means.shape[0]
    n_dim = means.shape[1]
    samples = np.empty((n_samples, n_dim))
    for i in range(n_samples):
        samples[i] = rand.multivariate_normal(means[i], covs[i])
    return samples

samples = generate_samples_fast(means, covs)

不过这个需要额外安装numba库，小样本量的话没必要折腾。

总结一下：不用非得写那种冗长的for循环，列表推导式是最平衡可读性和效率的方案，完全能满足你的需求～

内容的提问来源于stack exchange，提问作者p-value