使用Arviz绘制的旧PyMC3样式分组轨迹图

Question

我有一篇老博客，在那里我正在训练一个PyMC3模型。您可以找到blogposthere，但模型的要点如下所示

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对鸡体重的影响。这就是追踪图的样子

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

阿维兹的变化

此traceplot是使用较旧的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)

我需要做什么才能让每种饮食中的不同颜色恢复

Answer 1

在新的plot_trace函数中有一个参数。这就是诀窍

az.plot_trace(trace, compact=True)