PyMC零膨胀泊松模型形状不匹配错误解决求助

错误根源 #

你遇到的形状不匹配问题，是因为代码里的mu和p是标量（单个值），但给ZeroInflatedPoisson指定了dims='shop'，PyMC会默认要求p和mu的维度与shop（长度5）一致，但观测数据是5个样本（每个样本对应不同店铺），两者维度无法对齐。同时原代码里mu的范围设为0-1完全不符合实际销售数据（有100+的销售额），这也会导致后续推断失效。

两种可行修复方案 #

方案1：每家店铺独立参数

让mu和p都对应5家店铺的维度，再根据观测数据中的店铺索引，提取对应店铺的参数值：

import pymc as pm
import pandas as pd

# 构建示例数据（替换成你的实际数据）
data = pd.DataFrame({
    'sales': [117, 0, 0, 134, 0],
    'shop': [0, 0, 0, 1, 2]
})

with pm.Model(coords=dict(shop=range(5))) as model:
    # 为每家店铺定义独立的mu（调整范围适配实际销售额）和p
    mu = pm.Uniform('mu', 0, 1000, dims='shop')
    p = pm.Beta('p', alpha=1, beta=1, dims='shop')
    
    # 根据观测数据的shop列，匹配对应店铺的参数
    mu_obs = mu[data.shop.values]
    p_obs = p[data.shop.values]
    
    # 定义似然函数，参数与观测数据维度对齐
    likelihood = pm.ZeroInflatedPoisson(
        "likelihood", 
        p=p_obs, 
        mu=mu_obs, 
        observed=data.sales.values
    )
    
    # 采样先验预测
    prior = pm.sample_prior_predictive()

方案2：共享部分参数

如果假设所有店铺共享零膨胀概率p，仅让mu按店铺变化，代码可以简化为：

import pymc as pm
import pandas as pd

data = pd.DataFrame({
    'sales': [117, 0, 0, 134, 0],
    'shop': [0, 0, 0, 1, 2]
})

with pm.Model(coords=dict(shop=range(5))) as model:
    # 所有店铺共享的p（标量）
    p = pm.Beta('p', alpha=1, beta=1)
    # 每家店铺独立的mu
    mu = pm.Uniform('mu', 0, 1000, dims='shop')
    
    # 直接用mu的索引匹配观测数据，p会自动广播适配维度
    likelihood = pm.ZeroInflatedPoisson(
        "likelihood", 
        p=p, 
        mu=mu[data.shop.values], 
        observed=data.sales.values
    )
    
    prior = pm.sample_prior_predictive()

关于at.stack的说明 #

如果确实需要手动堆叠参数（比如批量生成与观测数据长度一致的参数数组），可以用pm.math.stack（at是aesara.tensor的别名，PyMC中pm.math等价于at）。比如共享p时，也可以手动堆叠成与观测数据同长度的数组：

p_obs = pm.math.stack([p]*len(data))

不过PyMC的广播机制会自动处理标量与数组的维度匹配，所以多数情况下不需要手动堆叠。

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