Python混合正态分布生成代码提速优化求助

问题描述 #

我实现了一个用于生成混合正态分布的Python函数，代码及调用方式如下：

# M=[float]*k = 分布中心值,
# S=[float]*k = 标准差,
# P=[float]*k = 每个分布的概率（总和为1.0）
# rng = 随机数生成器
# n = 返回数组的长度 [float]*n
# return [float]*n

def mixed_normal(rng, n, M, S, P):
    # 参见混合模型相关定义
    idx = np.random.choice(len(M), p=P, replace=True, size=n)
    return np.fromiter((rng.normal(M[i], S[i]) for i in idx),dtype=np.float64)

调用示例：

rng = np.random.default_rng()
def mixed_normal_3(rng, n):
    data = [(-5, 0, 5), (1, 1, 1), (1/3, 1/3, 1/3)]
    return mixed_normal(rng, n, *data)

当前实现的性能瓶颈：生成10^6个数据点时耗时约350秒，需要将耗时降低至约30秒。尝试将循环调用rng.normal改为单次Numpy调用但未找到可行方案，寻求优化方法。

最小可运行示例：

import numpy as np

rng = np.random.default_rng()

def mixed_normal(rng, n, M, S, P):
    # 参见混合模型相关定义
    idx = np.random.choice(len(M), p=P, replace=True, size=n) 
    # 需要优化的部分
    return np.fromiter((rng.normal(M[i], S[i]) for i in idx),dtype=np.float64)      

def mixed_normal_3(rng, n):
    data = [(-5, 0, 5), (1, 1, 1), (1/3, 1/3, 1/3)]
    return mixed_normal(rng, n, *data)
    
# 预期返回长度为10^6的浮点数组  
print( mixed_normal_3( rng , 10**6 ) );

优化方案 #

核心思路是利用Numpy的向量化操作替代Python循环，彻底消除逐个调用rng.normal的性能开销。优化后的代码如下：

import numpy as np

def mixed_normal(rng, n, M, S, P):
    # 将输入转换为Numpy数组，便于批量操作
    M = np.asarray(M)
    S = np.asarray(S)
    P = np.asarray(P)
    
    # 一次性计算每个分布需要生成的样本数量
    counts = np.random.multinomial(n, P)
    
    # 批量生成每个分布的正态样本
    samples = []
    for mu, sigma, cnt in zip(M, S, counts):
        if cnt > 0:
            # 单次调用生成对应数量的样本，利用Numpy向量化性能
            samples.append(rng.normal(mu, sigma, size=cnt))
    
    # 拼接所有样本并打乱顺序，保证混合随机性
    all_samples = np.concatenate(samples)
    rng.shuffle(all_samples)
    
    return all_samples

def mixed_normal_3(rng, n):
    data = [(-5, 0, 5), (1, 1, 1), (1/3, 1/3, 1/3)]
    return mixed_normal(rng, n, *data)

# 测试生成10^6个样本
rng = np.random.default_rng()
result = mixed_normal_3(rng, 10**6)
print(result)

优化要点说明 #

1. 批量分配样本数量：使用np.random.multinomial一次性确定每个分布的样本生成量，避免逐个选择索引的冗余操作
2. 向量化生成样本：对每个分布仅调用一次rng.normal批量生成样本，完全利用Numpy的底层C实现加速，避免Python循环的巨大开销
3. 保持随机性：最后通过rng.shuffle打乱样本顺序，确保生成的混合分布与原逻辑的随机性一致

性能测试显示，优化后的代码生成10^6个样本耗时可控制在几秒内，远低于30秒的目标。

内容的提问来源于stack exchange，提问作者ElectronicsStudent