如何高效运行5000样本的Beta分布MLE估计代码？

针对5000个含100个数据点的Beta分布样本MLE估计的性能瓶颈，以下是具体优化建议及修改后的代码：

核心优化方向 #

1. 替换np.append为预分配数组 #

np.append每次调用都会重新分配内存并复制现有数据，5000次循环会产生大量冗余开销。直接预分配(5000, 2)形状的数组存储MLE结果，避免内存反复申请。

2. 向量化生成所有样本 #

放弃循环内逐个生成样本的方式，一次性生成(5000, 100)的样本矩阵，利用scipy的向量化随机数生成能力，减少循环带来的开销。

3. 优化核心计算函数 #

• 用np.mean替代手动求和除法：numpy的mean是底层优化的向量化操作，比sum(...) / n效率更高。
• 移除Hessian函数冗余参数：当前Hessiana计算仅依赖theta，与输入样本x无关，删除无用的x参数，减少参数传递开销。
• Numba编译加速：对迭代过程中频繁调用的U_score、Hessiana、H_inv、max_likelihood使用numba.jit装饰，将Python代码编译为机器码，大幅提升计算速度。

4. MSE计算修正与优化 #

原MSE函数存在变量名笔误，修正后直接利用numpy向量化运算完成计算，无需额外循环，保证效率的同时避免错误。

优化后的完整代码 #

import numpy as np
import scipy.special as sp
from scipy import stats as st
from numba import jit

@jit(nopython=True)
def U_score(x, theta):
    a = theta[0]
    b = theta[1]
    epsilon = 1e-10  # 避免log(0)
    log_x = np.log(x + epsilon)
    log_1mx = np.log(1 - x + epsilon)
    d_a = -sp.digamma(a) + sp.digamma(a + b) + np.mean(log_x)
    d_b = -sp.digamma(b) + sp.digamma(a + b) + np.mean(log_1mx)
    return np.array((d_a, d_b))

@jit(nopython=True)
def Hessiana(theta):
    a = theta[0]
    b = theta[1]
    poly_ab = sp.polygamma(1, a + b)
    h11 = poly_ab - sp.polygamma(1, a)
    h12 = poly_ab
    h21 = poly_ab
    h22 = poly_ab - sp.polygamma(1, b)
    return np.array([[h11, h12], [h21, h22]])

@jit(nopython=True)
def H_inv(theta):
    H = Hessiana(theta)
    ridge = 1e-6  # 小常数防止矩阵奇异
    H_ridge = H + ridge * np.eye(H.shape[0])
    return np.linalg.inv(H_ridge)

@jit(nopython=True)
def max_likelihood(x, theta, tol):
    iter = 0
    while iter < 1000:
        theta_new = theta - H_inv(theta) @ U_score(x, theta)
        if np.linalg.norm(theta_new - theta) <= tol:
            return theta_new, iter + 1
        theta = theta_new
        iter += 1
    print('Não convergiu!')
    return theta_new, iter

def mse(arr, theta_real):
    a_true, b_true = theta_real
    mse_a = np.mean((arr[:, 0] - a_true) ** 2)
    mse_b = np.mean((arr[:, 1] - b_true) ** 2)
    return mse_a, mse_b

# 主流程优化
n_amostras = 5000
sample_size = 100
theta_init = np.array([1, 1])
tol = 1e-6

# 一次性生成所有样本
all_samples = st.beta.rvs(2, 5, size=(n_amostras, sample_size))

# 预分配结果数组
results = np.zeros((n_amostras, 2))

for i in range(n_amostras):
    x = all_samples[i]
    emv, iteracoes = max_likelihood(x, theta_init, tol)
    results[i] = emv

# 计算并输出MSE
theta_real = (2, 5)
mse_a, mse_b = mse(results, theta_real)
print(f"MSE para a: {mse_a:.6f}, MSE para b: {mse_b:.6f}")
print(results)

额外加速建议 #

若需进一步提升速度，可采用多进程并行处理：将5000个样本拆分到多个进程中同时计算MLE，利用多核CPU资源。例如使用multiprocessing.Pool或concurrent.futures.ProcessPoolExecutor实现并行化。

内容的提问来源于stack exchange，提问作者Gabriel Ligabô