使用 Arviz 绘制的旧 PyMC3 样式分组跟踪图

Question

我有一篇旧博文，其中我正在训练 PyMC3 模型。您可以找到博文 here，但模型的要点如下所示。

with pm.Model() as model:
    mu_intercept = pm.Normal('mu_intercept', mu=40, sd=5)
    mu_slope = pm.HalfNormal('mu_slope', 10, shape=(n_diets,))
    mu = mu_intercept + mu_slope[df.diet-1] * df.time
    sigma_intercept = pm.HalfNormal('sigma_intercept', sd=2)
    sigma_slope = pm.HalfNormal('sigma_slope', sd=2, shape=n_diets)
    sigma = sigma_intercept + sigma_slope[df.diet-1] * df.time
    weight = pm.Normal('weight', mu=mu, sd=sigma, observed=df.weight)
    approx = pm.fit(20000, random_seed=42, method="fullrank_advi")

在此数据集中，我正在估算 Diet 对鸡体重的影响。这就是跟踪图的样子。

你看多漂亮啊！每种饮食都有自己的路线！漂亮！

Arviz 变化

此跟踪图是使用较旧的 PyMC3 API 制作的。如今，此功能已转移到 arviz。所以尝试重做这项工作但是......情节看起来非常不同。

我在这里使用的代码略有不同。我现在正在使用 pm.Data，但我怀疑这会导致这种差异。

with pm.Model() as mod: 
    time_in = pm.Data("time_in", df['time'].astype(float))
    diet_in = pm.Data("diet_in", dummies)
    
    intercept = pm.Normal("intercept", 0, 2)
    time_effect = pm.Normal("time_weight_effect", 0, 2, shape=(4,))
    diet = pm.Categorical("diet", p=[0.25, 0.25, 0.25, 0.25], shape=(4,), observed=diet_in)
    sigma = pm.HalfNormal("sigma", 2)
    sigma_time_effect = pm.HalfNormal("time_sigma_effect", 2, shape=(4,))
    weight = pm.Normal("weight", 
                       mu=intercept + time_effect.dot(diet_in.T)*time_in, 
                       sd=sigma + sigma_time_effect.dot(diet_in.T)*time_in, 
                       observed=df.weight)
    trace = pm.sample(5000, return_inferencedata=True)

我需要做什么才能恢复每种 DIET 的不同颜色？

Answer 1

新的 plot_trace 函数中有一个参数。这样就可以了；

az.plot_trace(trace, compact=True)

使用 Arviz 绘制的旧 PyMC3 样式分组跟踪图

Old PyMC3 style grouping traceplot plotted with Arviz

python

pymc3

arviz

Arviz 变化