準備

import jax.numpy as jnp
import numpy as np
import jax.random as random
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(style='darkgrid',palette='bright')

!pip install numpyro
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
numpyro.set_host_device_count(4)

import arviz as az

# データインポート
sales_climate = pd.read_csv("3-7-1-beer-sales-4.csv")

データは以下のような分布。

g = sns.relplot(data=sales_climate,x="temperature",y="sales",hue="weather",)
g.fig.set_figheight(6)
g.fig.set_figwidth(10)

weatherごとにsalesの分布が異なり、temperature が高いほど sales が相関して高くなることが推察できます。

モデル

それでは、これまで実装してきた weather と temperature の２変数を利用したモデルに拡張します。

そのまま１変数ごとに説明変数&それに応じたパラメータを指定していっても良いんですが、変数が多くなるほど記述が冗長になって面倒です。
そこで、デザイン行列を利用したモデル記述方法に変更します。
これはRstan本でも紹介されています。

デザイン行列について詳しくはこちら。

ja.wikipedia.org

# モデル
def model(
    N,
    C,
    X,
    sales,
  
):
  beta = numpyro.sample("beta",dist.Normal(0,100),sample_shape=(C,))
  sigma = numpyro.sample("sigma",dist.HalfNormal(100))

  with numpyro.plate("N",N):
    mu = jnp.dot(X,beta)
    numpyro.sample("sales",dist.Normal(mu, sigma),obs = sales)

# データ
sales_climate_2 = pd.get_dummies(sales_climate).drop("weather_cloudy",axis=1)
sales_climate_2["Intercept"] = 1
X = sales_climate_2[["Intercept","weather_rainy","weather_sunny","temperature"]].values
data_dict = {
    "N":X.shape[0],
    "C":X.shape[1],
    "X":X,
    "sales":sales_climate_2["sales"].values
}

ポイントはパラメータを beta : sample_sape(C , )としていること。
C は X の変数（カラム数）に充当しており、今回で言えば4行1列の行列になります。
X は切片Intercept を含めた N行4列の行列です。
これで mu を X と beta の内積とし、加法モデルとして線形予測子を表現しています。
これなら説明変数がいくら増えようが、同じモデル内容で推論に回すことができ便利です。

推論・解釈

では、推論回していきます。

kernel = NUTS(model)
sample_kwargs = dict(
    sampler=kernel, 
    num_warmup=2000, 
    num_samples=2000, 
    num_chains=4, 
    chain_method="parallel"
)
mcmc = MCMC(**sample_kwargs)
mcmc.run(random.PRNGKey(0), **data_dict)

サンプリング結果を確認します。

az.summary(mcmc)

デザイン行列を利用したモデルによるパラメータは解釈に注意が必要です。

beta[0]　～　beta[3] は、それぞれ説明変数 X の列番号に該当する回帰係数です。
つまり、 beta[0] はIntercept、beta[1]はweather_rainy、beta[2]はweather_sunny、beta[3]はtemperatureです。
無事、書籍と同様の結果になりました。

モデルを利用した回帰直線のため、事後予測分布から描写します。

#事後予測分布のサンプリング
mcmc_samples=mcmc.get_samples()
predictive = numpyro.infer.Predictive(model, mcmc_samples)

pred_dict = {
    "N":X.shape[0],
    "C":X.shape[1],
    "X":X,
    "sales":None
}

ppc_samples = predictive(random.PRNGKey(0),**pred_dict)
idata_ppc = az.from_numpyro(mcmc, posterior_predictive=ppc_samples)
sales_pred =  idata_ppc.posterior_predictive['sales']

# 描写
pred_df = sales_climate.copy()
pred_df["sales_pred"] = sales_pred.mean(axis=(0,1))
g = sns.lmplot(x="temperature", y="sales_pred", hue="weather", data=pred_df,scatter=False)
g.map(sns.scatterplot,x="temperature",y="sales",hue="weather",data=pred_df,alpha=0.2)
g.fig.set_figheight(6)
g.fig.set_figwidth(10)