PyMC 4.0.1中带MVN先验的分层变效应模型维度错误解决

需求说明 #

已在PyMC 4.0.1中实现过多元先验，但无法在分层模型中正确使用。需构建含协变量x1、x2和结果变量y的回归模型，数据包含两个分类层级：dim1为高层分类，dim2为低层分类。模型包含截距、x1和x2的斜率；高层设多元正态超先验B_1，按dim1定义[截距、x1斜率、x2斜率]的变系数，为低层dim2的先验B_2提供信息。

问题描述 #

采样模型时触发错误：ValueError: Invalid dimension for value: 3。推测B_1和B_2的dims参数中的'3'有误，但手动评估模型变量显示B_2.eval().shape为(6,2,3)，mu中的线性模型计算正常。错误回溯指向multivariate.py的quaddist_parse函数，提示维度不能超过2，需修正MVN先验的定义方式。

核心问题分析 #

错误根源在于：

• 非法维度名称：不能直接使用数字'3'作为维度标识，PyMC要求所有维度必须是预先定义的合法坐标名。
• 多元分布维度匹配：MvNormal的联合分布维度（即参数向量维度）需要作为最后一维，且必须在模型coords中显式定义。

修正步骤 #

1. 在模型初始化时的coords参数中，新增params维度，对应3个回归参数（截距、x1斜率、x2斜率）。
2. 将B_1和B_2的dims参数中的'3'替换为params，确保维度引用合法。
3. 补充obs_id坐标定义（原代码遗漏，明确观测维度可避免潜在索引问题）。

修正后的模型代码 #

import numpy as np
import pymc as pm

with pm.Model(coords={
    'dim1': np.arange(2), 
    'dim2': np.arange(6),
    'params': ['intercept', 'x1_slope', 'x2_slope']  # 新增参数维度
}) as mdl:
    # 补充观测维度定义
    pm.add_coords('obs_id', np.arange(len(X)), mutable=True)
    
    d1 = pm.MutableData('d1', X['d1'].values, dims='obs_id')
    d2 = pm.MutableData('d2', X['d2'].values, dims='obs_id')
    
    # 高层先验
    sd_1 = pm.HalfNormal.dist(1, shape=3)
    chol_1, _, _ = pm.LKJCholeskyCov(
        'chol_1', n=3, eta=1,
        sd_dist=sd_1, compute_corr=True
    )
    B_1 = pm.MvNormal(
        'B_1', mu=[5, 0, 0], chol=chol_1, 
        dims=('dim1', 'params')  # 用合法维度名替代数字3
    )
    
    # 低层先验
    sd_2 = pm.HalfNormal.dist(1, shape=3)
    chol_2, _, _ = pm.LKJCholeskyCov(
        'chol_2', n=3, eta=1,
        sd_dist=sd_2, compute_corr=True
    )
    B_2 = pm.MvNormal(
        'B_2', mu=B_1, chol=chol_2, 
        dims=('dim2', 'dim1', 'params')  # 用合法维度名替代数字3
    )
    
    # 线性模型计算mu
    sigma = pm.HalfNormal('sigma', 2)
    mu = pm.Deterministic(
        'mu', 
        B_2[d2, d1, 0] +  # 截距对应params第0位
        B_2[d2, d1, 1] * X['x1'].values +  # x1斜率对应params第1位
        B_2[d2, d1, 2] * X['x2'].values,  # x2斜率对应params第2位
        dims='obs_id'
    )
    
    outcome = pm.StudentT(
        'outcome', nu=3, mu=mu, sigma=sigma, 
        observed=y, dims="obs_id"
    )
    trace = pm.sample(draws=2000, tune=1000, target_accept=0.95, random_seed=0)

数据生成代码 #

import numpy as np
import pandas as pd
from scipy import stats
import pymc as pm

rng = np.random.default_rng(0)
N = 1000
# 生成协变量和分类特征
X = pd.DataFrame(
    {
        'x1': stats.halfnorm(loc=0,scale=3).rvs(N),
        'x2': stats.norm(loc=0,scale=2).rvs(N),
        'd1': rng.choice([0,1],size=N, p=[0.6,0.4]),
        'd2': rng.choice(np.arange(6),size=N, p=[0.1,0.2,0.1,0.3,0.1,0.2]),
    }
)
# 参数分布均值
intercept = np.array([
    [5,5,6,7,6,5],
    [4,4,5,6,6,4]
])
slope1 = np.array([
    [0,0.7,0.3,-0.2,-1,0],
    [0.5,1,-1,0.6,-0.2,0.3]
])
slope2 = np.array([
    [0,0.7,0.3,-0.2,-1,0],
    [0.5,1,-1,0.6,-0.2,0.3]
])*1.5
# 生成随机协方差矩阵
corrs = []
for _ in np.arange(6):
    _,corr,_ = pm.LKJCholeskyCov.dist(eta=1,n=3,sd_dist=pm.HalfNormal.dist(1,shape=3), compute_corr=True)
    corrs.append(corr.eval())

# 生成结果变量
y = np.zeros(N)
for d1 in [0,1]:
    for d2 in np.arange(6):
        ind = (X['d1']==d1)&(X['d2']==d2)
        mv = stats.multivariate_normal(mean=[intercept[d1,d2], slope1[d1,d2],slope2[d1,d2]],cov=corrs[d2]).rvs(1)
        y[ind] = mv[0] + X.loc[ind,'x1']*mv[1] + X.loc[ind,'x2']*mv[2] + rng.normal(loc=0,scale=1,size=ind.sum())

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