用Arviz绘制的旧PyMC3风格分组跟踪图

Question

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

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)

我需要做些什么才能让每个饮食中的不同颜色恢复原样？

Answer 1

在新的plot_trace函数中有一个参数。这起作用很大；

az.plot_trace(trace, compact=True)