如何在Python的emcee中处理离散变量进行模型拟合？

针对你的需求，emcee本身是为连续参数设计的MCMC采样器，但可以通过两种方式处理离散变量的采样，核心思路是让采样器只接受离散变量的合法取值，或者将离散变量转换为连续索引后再映射回目标值：

方法1：直接在验中过滤非法离散值 #

这种方法最简单，直接在lnprior函数中严格限制离散变量m和b只能取你指定的列表值，非法取值直接返回负无穷（被采样器拒绝）。

修改后的代码： #

import numpy as np
import emcee

# 定义离散变量的允许取值
m_allowed = [1, 2, 3, 4]
b_allowed = [16, 17, 18, 20]
# 连续变量的范围（自行替换为你的实际值）
xs_min, xs_max = -5, 5
ys_min, ys_max = -5, 5

# 数据（自行替换为你的实际数据）
x = np.array([...])
y = np.array([...])
yerr = np.array([...])

def lnlike(theta, x, y, yerr):
    m, b, xs, ys = theta
    # 读取对应离散参数的预生成模型
    model = np.genfromtxt(f'Generic_name_{m}{b}.txt')
    model[:,0] += xs
    model[:,1] += ys
    # 计算对数似然（注意原代码里y和model的维度要匹配，这里假设model[:,1]是拟合值）
    inv_sigma2 = 1.0 / (yerr**2)
    return -0.5 * np.sum((y - model[:,1])**2 * inv_sigma2)

def lnprior(theta):
    m, b, xs, ys = theta
    # 仅允许离散变量取指定值，连续变量在范围内
    if (m in m_allowed) and (b in b_allowed) and (xs_min < xs < xs_max) and (ys_min < ys < ys_max):
        return 0.0
    return -np.inf

def lnprob(theta, x, y, yerr):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, x, y, yerr)

# 采样器参数设置
nwalkers = 200
nchains = 1500
ndim = 4

# 初始化walker位置：离散变量从允许列表中选，连续变量给初始值
# 比如选m=1, b=16, xs=3, ys=4作为初始猜测
fguess = [1, 16, 3, 4]
# 离散变量的扰动要小，避免跳出合法值范围，连续变量正常扰动
sigma_gb = 0.1
pos = [
    fguess + sigma_gb * np.random.randn(ndim)
    for i in range(nwalkers)
]

# 创建采样器并运行
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(x, y, abs(yerr)))
sampler.run_mcmc(pos, nchains, progress=True)

注意事项： #

• 这种方法的缺点是会有大量提案被拒绝（比如采样器提出m=1.5时会被直接丢弃），采样效率较低。
• 初始化walker时，离散变量的初始值必须是m_allowed和b_allowed中的值，且扰动幅度不能太大，否则初始位置就会被拒绝。

方法2：将离散变量转换为连续索引采样（推荐） #

把离散变量的取值映射到连续的索引参数上，采样器对索引（连续值）进行采样，再在似然函数中把索引映射回对应的离散值。这种方法能大幅提升采样效率，避免大量无效提案。

修改后的代码： #

import numpy as np
import emcee

# 定义离散变量的允许取值
m_allowed = [1, 2, 3, 4]
b_allowed = [16, 17, 18, 20]
# 连续变量的范围（自行替换为你的实际值）
xs_min, xs_max = -5, 5
ys_min, ys_max = -5, 5

# 数据（自行替换为你的实际数据）
x = np.array([...])
y = np.array([...])
yerr = np.array([...])

def lnlike(theta, x, y, yerr):
    # theta现在是[m_idx, b_idx, xs, ys]，其中m_idx和b_idx是连续索引
    m_idx, b_idx, xs, ys = theta
    # 将连续索引四舍五入为整数，映射到对应的离散值
    m = m_allowed[round(m_idx)]
    b = b_allowed[round(b_idx)]
    # 读取预生成模型
    model = np.genfromtxt(f'Generic_name_{m}{b}.txt')
    model[:,0] += xs
    model[:,1] += ys
    inv_sigma2 = 1.0 / (yerr**2)
    return -0.5 * np.sum((y - model[:,1])**2 * inv_sigma2)

def lnprior(theta):
    m_idx, b_idx, xs, ys = theta
    # 索引的范围是0到len(allowed)-1，确保四舍五入后不会越界
    if (0 <= m_idx <= len(m_allowed)-1) and (0 <= b_idx <= len(b_allowed)-1) and (xs_min < xs < xs_max) and (ys_min < ys < ys_max):
        return 0.0
    return -np.inf

def lnprob(theta, x, y, yerr):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, x, y, yerr)

# 采样器参数设置
nwalkers = 200
nchains = 1500
ndim = 4

# 初始化walker位置：索引初始值取0（对应第一个离散值），连续变量给初始值
fguess = [0, 0, 3, 4]
sigma_gb = 0.1  # 索引的扰动幅度小，避免跨太多索引
pos = [
    fguess + sigma_gb * np.random.randn(ndim)
    for i in range(nwalkers)
]

# 创建采样器并运行
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(x, y, abs(yerr)))
sampler.run_mcmc(pos, nchains, progress=True)

# 采样完成后，将索引映射回离散值
samples = sampler.get_chain(flat=True)
# 处理m和b的索引，得到实际离散值
m_samples = [m_allowed[round(idx)] for idx in samples[:,0]]
b_samples = [b_allowed[round(idx)] for idx in samples[:,1]]
xs_samples = samples[:,2]
ys_samples = samples[:,3]

优势： #

• 采样器始终在连续参数空间内提案，不会有大量无效提案被拒绝，效率更高。
• 后续处理采样结果时，只需把索引映射回离散值即可。

额外说明 #

如果你的离散变量数量不多，也可以考虑用嵌套采样工具（比如dynesty），这类工具对混合离散/连续变量的支持更友好，但emcee通过上述两种方法也能满足需求。

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