使用multiprocess并行运行EMCEE时log_prior未定义问题求助

问题原因 #

Jupyter Notebook的运行环境和普通Python脚本不同：主进程的__name__并非__main__，当multiprocess.Pool创建子进程时，子进程会尝试重新加载当前命名空间，但Notebook的代码执行逻辑无法让子进程正确获取到你在Notebook中定义的log_prior等函数，导致出现未定义错误。

解决办法 #

最可靠的方案是将所有MCMC相关的概率函数封装到单独的Python模块中，让子进程可以通过导入模块正确读取函数定义：

步骤1：创建单独的函数模块

新建一个名为mcmc_functions.py的文件，写入以下内容：

import numpy as np

def log_likelihood(theta, x, y, yerr):
    mu, sig, log_f = theta
    model = np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
    sigma2 = yerr**2 + model**2 * np.exp(2 * log_f)
    return -0.5 * np.sum((y - model) ** 2 / sigma2 + np.log(sigma2))

def log_prior(theta):
    mu, sig, log_f  = theta
    if 2 < mu < 10 and 0.0 < sig < 5 and -10.0 < log_f < 1.0:
        return 0.0
    return -np.inf

def log_probability(theta, x, y, yerr):
    lp = log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + log_likelihood(theta, x, y, yerr)

步骤2：修改Notebook中的代码

将原Notebook代码修改为以下形式，导入模块中的函数并调整进程启动方式：

import numpy as np
import emcee
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from multiprocess import Pool
# 导入封装好的函数
from mcmc_functions import log_likelihood, log_probability

# 生成合成数据部分不变
np.random.seed(123)
mu_true = 5
sig_true = 0.5
f_true = 0.534
N = 500
x = np.sort(10 * np.random.rand(N))
yerr = 0.03 + 0.05 * np.random.rand(N)
y = np.exp(-np.power(x - mu_true, 2.) / (2 * np.power(sig_true, 2.)))
y += 0.5 * np.abs(f_true * y) * np.random.randn(N)
y += yerr * np.random.randn(N)

# 最大似然拟合部分不变
np.random.seed(42)
nll = lambda *args: -log_likelihood(*args)
initial = np.array([mu_true, sig_true, np.log(f_true)]) + 0.1 * np.random.randn(3)
soln = minimize(nll, initial, args=(x, y, yerr))
mu_ml, sig_ml, log_f_ml = soln.x

# MCMC并行部分
pos = soln.x + 0.2 * np.random.randn(32, 3)
nwalkers, ndim = pos.shape
nsteps = 10000

if __name__ == "__main__":
    # 显式设置进程启动方式为spawn，适配Jupyter环境
    import multiprocess
    multiprocess.set_start_method('spawn', force=True)
    with Pool() as pool:
        sampler = emcee.EnsembleSampler(
            nwalkers, ndim, log_probability, args=(x, y, yerr), pool=pool
        )
        sampler.run_mcmc(pos, nsteps, progress=True)

备选方案（不推荐，稳定性差） #

如果不想单独创建模块，可以尝试将所有函数定义放在if __name__ == "__main__":代码块之前，并强制设置spawn启动方式，但这种方法受Jupyter命名空间限制，偶尔仍会出现问题：

# 所有导入、函数定义、数据生成、拟合代码放在这里

if __name__ == "__main__":
    import multiprocess
    multiprocess.set_start_method('spawn', force=True)
    with Pool() as pool:
        sampler = emcee.EnsembleSampler(...)
        sampler.run_mcmc(...)

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