求助：基于PyMC实现查表式MCMC参数优化的代码完善

问题描述 #

我在pandas DataFrame中存储了一组测量数据，前两列为参数Param1和Param2，每行对应4个特征值。现有一组新观测特征值，目标是在查表数据集中找到匹配的最优参数组合。

计划用MCMC智能搜索替代穷举网格搜索：遍历参数组合，获取对应特征后计算与观测特征的卡方（chi2）值，寻找最小卡方对应的参数。已用emcee和PyMC3实现，但PyMC3速度过慢，尝试改用PyMC时因参数为分布变量，不知如何实现查表匹配，恳请协助完善代码。

原代码 #

import pandas as pd
import pytensor.tensor as pt
import pymc as pm
import numpy as np

#this function should return features(data) for given parameters
def model_getFeaturesFromParams(Param1, Param2, lookup_table_arr):
    lookup_table_tensor = pt.as_tensor_variable(lookup_table_arr)
    # Compute the index based on the given Param1 and Param2
    print(lookup_table_tensor[:,0])
    a = pt.eq(lookup_table_tensor[:, 0], Param1)
    b = pt.eq(lookup_table_tensor[:, 1], Param1)
    print(a,b)
    index = (lookup_table_tensor[:, 0] == Param1) & (lookup_table_tensor[:, 1] == Param2)
    
    #  Check if there's a matching row in the lookup table
    if pt.all(index):
        # Extract the features from the matching row
        features = lookup_table_tensor[index][0][2:]
        return features
    else:
        # Handle the case when Param1 and Param2 are not found in the table
        return [0, 0, 0, 0]  # Modify this as needed

def chiSquaredLikelihood(observed_features, observed_errors, predicted_features):
        chi2 = pt.sum(pt.sqr((observed_features - predicted_features) / observed_errors))
        return -2 * chi2


def run_pymc():
    observed_features = np.array([2,2,4,4])
    observed_errors = np.array([0.03, 0.02, 0.02, 0.03])

    lookup_table = [
        [1,2, 1,1,2,2],
        [2,4, 2,2,4,4],
        [4,8, 4,4,8,8],
    ]
    columns = ["Param1", "Param2", "feature1", "feature2", "feature3", "feature4"]
    df = pd.DataFrame(lookup_table, columns=columns)
    lookup_table_arr = df.values

    with pm.Model() as model:
        Param1 = pm.Uniform('Param1', lower=1, upper=4)
        Param2 = pm.Uniform('Param2', lower=2, upper=8)
        Param1_tensor = pt.as_tensor_variable(Param1)
        Param2_tensor = pt.as_tensor_variable(Param2)

        # predicted_features = model_getFeaturesFromParams(Param1_tensor, Param2_tensor, df)
        predicted_features = model_getFeaturesFromParams(Param1_tensor, Param2_tensor, lookup_table_arr)

        likelihood = pm.DensityDist(
            "likelihood",  logp=lambda x: chiSquaredLikelihood(observed_features, observed_errors, x), 
                    observed=predicted_features.eval()
                    )
        trace = pm.sample(100,100, chains=4, return_inferencedata=True,
                                progressbar=True, random_seed=42, discard_tuned_samples=True, cores=4, step=pm.Metropolis())
    return trace

if __name__ == "__main__":
     run_pymc()

代码问题分析 #

1. 符号张量条件判断错误：原函数使用Python原生if语句判断张量条件，这在PyTensor符号计算中不允许，必须用PyTensor提供的pt.where或pt.switch做符号化条件分支。
2. 参数类型不匹配：查表中的Param1、Param2是离散值，但原代码定义为连续Uniform分布，导致几乎无法匹配到查表参数组合。
3. 代码笔误：b = pt.eq(lookup_table_tensor[:, 1], Param1)应为Param2。
4. DensityDist使用错误：observed参数不应调用.eval()提前计算张量，且logp函数输入逻辑不符合PyMC模型变量定义。
5. 采样参数格式错误：PyMC的pm.sample参数顺序已变更，需明确指定draws和tune参数。

修正后的代码 #

import pandas as pd
import pytensor.tensor as pt
import pymc as pm
import numpy as np

def chi_squared_logp(predicted_features, observed_features, observed_errors):
    # 计算卡方对数似然（负2倍卡方）
    chi2 = pt.sum(pt.sqr((observed_features - predicted_features) / observed_errors))
    return -2 * chi2

def run_pymc():
    observed_features = np.array([2, 2, 4, 4])
    observed_errors = np.array([0.03, 0.02, 0.02, 0.03])

    # 定义查表数据
    lookup_table = [
        [1, 2, 1, 1, 2, 2],
        [2, 4, 2, 2, 4, 4],
        [4, 8, 4, 4, 8, 8],
    ]
    columns = ["Param1", "Param2", "feature1", "feature2", "feature3", "feature4"]
    df = pd.DataFrame(lookup_table, columns=columns)
    
    # 提取参数和特征数组（转为PyTensor张量）
    params_tensor = pt.as_tensor_variable(df[["Param1", "Param2"]].values)
    features_tensor = pt.as_tensor_variable(df[["feature1", "feature2", "feature3", "feature4"]].values)

    with pm.Model() as model:
        # 定义离散参数：用Categorical选择查表的行索引
        row_idx = pm.Categorical('row_idx', p=np.ones(len(df))/len(df))
        
        # 根据索引获取对应的参数和特征
        selected_params = params_tensor[row_idx]
        predicted_features = features_tensor[row_idx]
        
        # 定义卡方似然
        pm.DensityDist(
            'likelihood',
            logp=lambda x: chi_squared_logp(x, observed_features, observed_errors),
            observed=predicted_features
        )
        
        # 采样：使用Metropolis步长（适合离散变量）
        trace = pm.sample(
            draws=1000,
            tune=500,
            chains=4,
            return_inferencedata=True,
            progressbar=True,
            random_seed=42,
            discard_tuned_samples=True,
            cores=4,
            step=pm.Metropolis()
        )
    
    # 输出结果：统计每个行索引的出现次数，对应最优参数
    idx_counts = np.bincount(trace.posterior['row_idx'].values.flatten())
    best_row_idx = np.argmax(idx_counts)
    best_params = df.iloc[best_row_idx][["Param1", "Param2"]]
    print(f"最优参数组合：Param1={best_params['Param1']}, Param2={best_params['Param2']}")
    
    return trace

if __name__ == "__main__":
    run_pymc()

关键修正说明 #

1. 离散参数定义：改用pm.Categorical直接选择查表的行索引，避免连续参数无法匹配离散查表值的问题，同时简化查表逻辑。
2. 符号化查表：提前将参数和特征转为张量，通过索引直接获取对应特征，完全符合PyTensor符号计算要求，避免原代码中的条件判断错误。
3. 似然函数优化：将卡方似然函数调整为符合PyMCDensityDist的输入格式，直接使用符号张量计算。
4. 采样参数修正：明确指定draws（采样数）和tune（调优数）参数，符合PyMC最新API规范。
5. 结果解析：通过统计采样中各索引的出现次数，找到最可能的参数组合（对应卡方值最小的行）。

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