python - Arviz でプロットされた古い PyMC3 スタイルのグループ化トレースプロット

Question

PyMC3 モデルをトレーニングしている古いブログ投稿があります。ブログ投稿はこちらにありますが、モデルの要点を以下に示します。

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、ニワトリの体重に対する の影響を推定しています。これがトレースプロットの外観です。

それがどれほどきれいか見てください！各ダイエットには独自のラインがあります！美しい！

アルビズの変更点

このトレースプロットは、古い 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)

ダイエットごとに異なる色を元に戻すにはどうすればよいですか?