ディリクレ過程回帰で空間データを分析してみた

# はじめに

こんにちは、事業会社で働いているデータサイエンティストです：

https://dippeople.dip-net.jp/18130/

最近周りからディリクレ過程を説明してほしいとの要望が増えていますが、いきなり時系列モデルへの応用

https://qiita.com/Gotoubun_taiwan/items/e791161863a8f92a8a03

や国会議員のイデオロギーのベクトル化モデル

https://qiita.com/Gotoubun_taiwan/items/d450c05ffe7f7f818e99

に入っても大変なので、この記事では有名なカリフォルニア州の不動産価格データ

https://www.kaggle.com/datasets/camnugent/california-housing-prices

の予測でディリクレ過程の基礎を説明しまーす！

基礎内容の説明記事なので、わかりやすさを重視し、学術的な専門性は割愛ます。参考文献も特にない予定です。

# モデル説明

ディリクレ過程とはクラスター数の指定が不要なクラスタリングモデルです。

ディリクレ過程回帰とは

- 説明変数のクラスタリングをして、これをもとに
- クラスターごとのモデルを作成する

という考えから生まれたモデルです。

ここで強調しておきたいですが、クラスタリング用の説明変数と従属変数予測用の説明変数は同一である必要がありません。

「カリフォルニア全域の不動産価格を一本の線形回帰でうまく予測できるはずがない！」、確かにそうかもしれないです。でも、カリフォルニア全域を、線形回帰で説明できるような同質なエリアにセグメンテーションすれば、うまくいきそうですね。

なので、ディリクレ過程回帰は、一部の分析者でよくみられる恣意的に（？）都道府県でデータを切ってモデルを作るような行為を自動化し、かつ統計学的な正当性を付与することで、過学習を回避するだけでなく、効果の異質性にも目を向けてくれる有難い手法です。

また、ベイズなので、説明変数にカテゴリデータやカウントデータが入っても、柔軟に対応できます。

では早速定式化を確認しましょう！

# モデル定式化

$$切片 \sim Normal(0, 1)$$

$$sigma \sim Gamma(0.1, 0.1)$$

- 全体のディリクレ過程

$$d\\_alpha\sim Gamma(0.01, 0.01)$$

$$z\sim Stick\space Breaking(d\\_alpha)$$

- クラスターp

$$数値説明変数管轄中心点_{p}\sim Normal(0,1)$$

$$数値説明変数管轄範囲_{p} \sim Gamma(0.1, 0.1)$$

$$カテゴリ説明変数管轄中心点_{p} \sim Dirichlet(1/カテゴリの種類)$$

$$sigma_{beta_{p}} \sim Gamma(0.1, 0.1)$$

$$beta_{p} \sim Normal(0, sigma_{beta_{p}}^2)$$

- 観測値i

$$\eta_{i} \sim Categorical(z)$$

$$数値説明変数_{i} \sim Normal(数値説明変数管轄中心点_{\eta_{i}}, 数値説明変数管轄範囲_{\eta_{i}}^2)$$

$$カテゴリ説明変数_{i} \sim Categorical(カテゴリ説明変数管轄中心点_{\eta_{i}})$$

$$従属変数 \sim Normal(数値説明変数_{i} * beta_{数値,\eta_{i}} + カテゴリ説明変数_{i} * beta_{カテゴリ, \eta_{i}}, sigma^2)$$

はい、まずモデルを書きましたが、見てわかるように、一般的なベイズモデルは生成モデルと言われるものの、基本的には従属変数しかモデリングせず、説明変数を所与の条件として扱うことが多いです。

これに対して、ディリクレ過程回帰では、説明変数もクラスタリングの形でモデリングすることで、データをセグメントに分けてモデルを構築することを実現します。

# モデル実装

上記のモデルのStanコードです。

```cpp:dirichlet_housing.stan
functions {
  real partial_sum_lpdf(
    array[] real y,
    int start, int end,
    
    matrix x,
    matrix x_effect,
    array[] int ocean,
    int P,
    vector z,
    array[] vector middle_point_x,
    array[] vector spread_x,
    array[] vector beta_x,
    array[] vector middle_point_ocean,
    array[] vector beta_ocean,
    
    real intercept, real sigma
  ){
    vector[end - start + 1] lambda;
    int count = 1;
    for (i in start:end){
      vector[P] case_vector;
      for (p in 1:P){
        case_vector[p] = log(z[p]) +
                         normal_lpdf(x[i,] | middle_point_x[p], spread_x[p] .^ 2) + 
                         categorical_lpmf(ocean[i] | middle_point_ocean[p]) + 
                         normal_lpdf(y[count] | intercept + x_effect[i,] * beta_x[p] + beta_ocean[p, ocean[i]], sigma ^ 2);
      }
      lambda[count] = log_sum_exp(case_vector);
      count += 1;
    }
    return sum(lambda);
  }
}
data {
  int K;
  int K_effect;
  int P;
  int ocean_type;
  
  int N;
  matrix[N, K] x;
  matrix[N, K_effect] x_effect;
  array[N] int ocean;
  array[N] real y;
  
  int val_N;
  matrix[val_N, K] val_x;
  matrix[val_N, K_effect] val_x_effect;
  array[val_N] int val_ocean;
  array[val_N] real val_y;
}
parameters {
  real<lower=0> d_alpha;                                       // ディリクレ過程の全体のパラメータ
  vector<lower=0, upper=1>[P - 1] breaks;  // ディリクレ過程のstick-breaking representationのためのパラメータ
  array[P] vector[K] middle_point_x;
  array[P] vector<lower=0>[K] spread_x;
  array[P] vector<lower=0>[K_effect] sigma_beta_x;
  array[P] vector[K_effect] beta_x;
  array[P] simplex[ocean_type] middle_point_ocean;
  array[P] vector<lower=0>[ocean_type] sigma_beta_ocean;
  array[P] vector[ocean_type] beta_ocean;
  
  real intercept;
  real<lower=0> sigma;
}
transformed parameters {
  simplex[P] z; 
  {
    // stick-breaking representationの変換開始
    // https://discourse.mc-stan.org/t/better-way-of-modelling-stick-breaking-process/2530/2 を参考
    z[1] = breaks[1];
    real sum = z[1];
    for (p in 2:(P - 1)) {
      z[p] = (1 - sum) * breaks[p];
      sum += z[p];
    }
    z[P] = 1 - sum;
  }
}
model {
  d_alpha ~ gamma(0.01, 0.01);
  breaks ~ beta(1, d_alpha);
  
  for (p in 1:P){
    middle_point_x[p] ~ normal(0, 1);
    spread_x[p] ~ gamma(0.1, 0.1);
    sigma_beta_x[p] ~ gamma(0.1, 0.1);
    beta_x[p] ~ normal(0, sigma_beta_x[p] .^ 2);
    middle_point_ocean[p] ~ dirichlet(rep_vector(1.0/ocean_type, ocean_type));
    sigma_beta_ocean[p] ~ gamma(0.1, 0.1);
    beta_ocean[p] ~ normal(0, sigma_beta_ocean[p] .^ 2);
  }
  
  intercept ~ normal(0, 1);
  sigma ~ gamma(0.1, 0.1);
  
  int grainsize = 1;
  
  target += reduce_sum(
    partial_sum_lupdf, y, 
    
    grainsize,
    
    x,
    x_effect,
    ocean,
    P,
    z,
    middle_point_x,
    spread_x,
    beta_x,
    middle_point_ocean,
    beta_ocean,
    
    intercept, sigma
  );
  
}
generated quantities {
  array[val_N] vector[P] estimated_eta;
  array[val_N] real predicted;
  vector[val_N] log_likelihood;
  real sum_log_likelihood;
  for (i in 1:val_N){
    vector[P] case_vector;
    for (p in 1:P){
      case_vector[p] = log(z[p]) +
                       normal_lpdf(val_x[i,] | middle_point_x[p], spread_x[p] .^ 2) + 
                       categorical_lpmf(val_ocean[i] | middle_point_ocean[p]);
    }
    estimated_eta[i] = softmax(case_vector);
    int eta = categorical_rng(estimated_eta[i]);
    predicted[i] = normal_rng(intercept + val_x_effect[i,] * beta_x[eta] + beta_ocean[eta, val_ocean[i]], sigma ^ 2);
    log_likelihood[i] = normal_lpdf(val_y[i] | intercept + val_x_effect[i,] * beta_x[eta] + beta_ocean[eta, val_ocean[i]], sigma ^ 2);
  }
  sum_log_likelihood = sum(log_likelihood);
}

```

# 前処理

前処理として、まずカテゴリデータマスター表（海の近さ）を作成してID化して、他の数値データは標準化します。

もちろん、数値データの中にはカウントデータのようなものもあり、ポワソン分布でモデリングした方がいいものもありますが、コードのわかりやすさのために割愛します。

カウント説明変数を正規分布以外の分布でモデリングすることを読者への宿題とします笑。

また、緯度経度は説明変数のクラスタリングで使えるかもしれないですが、さすがに従属変数予測用の線形モデルに入れるのは少々不自然なので、クラスタリング用の行列と予測用の行列は別々で作ります。

```R
housing_df <- readr::read_csv("housing.csv") |>
  tidyr::drop_na() |>
  dplyr::mutate(dplyr::across(!ocean_proximity,
                              ~ (. - mean(.))/sd(.), .names = "{.col}_normalized"))

ocean_master <- housing_df |>
  dplyr::select(ocean_proximity) |>
  dplyr::distinct() |>
  dplyr::arrange(ocean_proximity) |>
  dplyr::mutate(
    ocean_id = dplyr::row_number()
  )

housing_df_with_id <- housing_df |>
  dplyr::left_join(ocean_master, by = "ocean_proximity")

# Rが勝手に切片を入れてくるので、-1でそうさせないようにする

X <- model.matrix(
  ~ - 1 + longitude_normalized + latitude_normalized + housing_median_age_normalized + 
    total_rooms_normalized + total_bedrooms_normalized + population_normalized + 
    households_normalized + median_income_normalized,
  data = housing_df_with_id
)

# 緯度経度が従属変数説明用の線形モデルに入るのは若干？？？なので別々で作ります
X_effect <- model.matrix(
  ~ -1 + housing_median_age_normalized + 
    total_rooms_normalized + total_bedrooms_normalized + population_normalized + 
    households_normalized + median_income_normalized,
  data = housing_df_with_id
)

set.seed(12345)
val_id <- sample(
  1:nrow(housing_df_with_id),
  1000,
  replace = FALSE
)

dh_data_list <- list(
  K = ncol(X),
  K_effect = ncol(X_effect),
  P = 10,
  ocean_type = nrow(ocean_master),
  
  N = nrow(X[-val_id,]),
  x = X[-val_id,],
  x_effect = X_effect[-val_id,],
  ocean = housing_df_with_id$ocean_id[-val_id],
  y = housing_df_with_id$median_house_value_normalized[-val_id],
  
  val_N = nrow(X[val_id,]),
  val_x = X[val_id,],
  val_x_effect = X_effect[val_id,],
  val_ocean = housing_df_with_id$ocean_id[val_id]
)

```

# モデル推定

では、早速モデル推定に入りましょう！

```R
> m_dh_init <- cmdstanr::cmdstan_model("dirichlet_housing.stan",
                                      cpp_options = list(
                                          stan_threads = TRUE
                                      )
 ) 
> m_dh_estimate <- m_dh_init$variational(
     seed = 12345,
     threads = 6,
     data = dh_data_list,
     iter = 10000
 )
------------------------------------------------------------ 
EXPERIMENTAL ALGORITHM: 
  This procedure has not been thoroughly tested and may be unstable 
  or buggy. The interface is subject to change. 
------------------------------------------------------------ 
Gradient evaluation took 0.032669 seconds 
1000 transitions using 10 leapfrog steps per transition would take 326.69 seconds. 
Adjust your expectations accordingly! 
Begin eta adaptation. 
Iteration:   1 / 250 [  0%]  (Adaptation) 
Iteration:  50 / 250 [ 20%]  (Adaptation) 
Iteration: 100 / 250 [ 40%]  (Adaptation) 
Iteration: 150 / 250 [ 60%]  (Adaptation) 
Iteration: 200 / 250 [ 80%]  (Adaptation) 
Success! Found best value [eta = 1] earlier than expected. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100      -228624.026             1.000            1.000 
   200      -221634.127             0.516            1.000 
   300      -205743.127             0.370            0.077 
   400      -201520.533             0.282            0.077 
   500      -192545.844             0.235            0.047 
   600      -190256.665             0.198            0.047 
   700      -189791.822             0.170            0.032 
   800      -189168.177             0.149            0.032 
   900      -178159.874             0.140            0.032 
  1000      -178635.085             0.126            0.032 
  1100      -174498.689             0.028            0.024 
  1200      -174950.818             0.025            0.021 
  1300      -174320.502             0.018            0.012 
  1400      -174522.857             0.016            0.004   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  101.4 seconds.
> m_dh_summary <- m_dh_estimate$summary()
```
5分で推定が終わりました！次は比較のため単純な線形モデルも推定します。ベイズ統計学らしく、検証データの対数尤度和で二つのモデルを比較します。

# 単純線形モデルとの比較

線形モデルのStanコードはこちらです：

```cpp:lm_housing.stan
functions {
  real partial_sum_lpdf(
    array[] real y,
    int start, int end,
    
    matrix x_effect,
    array[] int ocean,
    
    vector beta_x,
    vector beta_ocean,
    
    real intercept, real sigma
  ){
    vector[end - start + 1] lambda;
    int count = 1;
    for (i in start:end){
      lambda[count] = normal_lpdf(y[count] | intercept + x_effect[i,] * beta_x + beta_ocean[ocean[i]], sigma ^ 2);
      count += 1;
    }
    return sum(lambda);
  }
}
data {
  int K_effect;
  int ocean_type;
  
  int N;
  matrix[N, K_effect] x_effect;
  array[N] int ocean;
  array[N] real y;
  
  int val_N;
  matrix[val_N, K_effect] val_x_effect;
  array[val_N] int val_ocean;
  array[val_N] real val_y;
}
parameters {
  vector<lower=0>[K_effect] sigma_beta_x;
  vector[K_effect] beta_x;
  vector<lower=0>[ocean_type] sigma_beta_ocean;
  vector[ocean_type] beta_ocean;
  
  real intercept;
  real<lower=0> sigma;
}
model {
  sigma_beta_x ~ gamma(0.1, 0.1);
  beta_x ~ normal(0, sigma_beta_x .^ 2);
  sigma_beta_ocean ~ gamma(0.1, 0.1);
  beta_ocean ~ normal(0, sigma_beta_ocean ^ 2);
  
  intercept ~ normal(0, 1);
  sigma ~ gamma(0.1, 0.1);
  
  int grainsize = 1;
  
  target += reduce_sum(
    partial_sum_lupdf, y, 
    
    grainsize,
    
    x_effect,
    ocean,
    beta_x,
    beta_ocean,
    
    intercept, sigma
  );
  
}
generated quantities {
  array[val_N] real predicted;
  vector[val_N] log_likelihood;
  real sum_log_likelihood;
  for (i in 1:val_N){
    predicted[i] = normal_rng(intercept + val_x_effect[i,] * beta_x + beta_ocean[val_ocean[i]], sigma ^ 2);
    log_likelihood[i] = normal_lpdf(val_y[i] | intercept + val_x_effect[i,] * beta_x + beta_ocean[val_ocean[i]], sigma ^ 2);
  }
  sum_log_likelihood = sum(log_likelihood);
}

```

続いて単純線形モデルを推定します：

```R
lm_data_list <- list(
  K_effect = ncol(X_effect),
  ocean_type = nrow(ocean_master),
  
  N = nrow(X[-val_id,]),
  x_effect = X_effect[-val_id,],
  ocean = housing_df_with_id$ocean_id[-val_id],
  y = housing_df_with_id$median_house_value_normalized[-val_id],
  
  val_N = nrow(X[val_id,]),
  val_x_effect = X_effect[val_id,],
  val_ocean = housing_df_with_id$ocean_id[val_id],
  val_y = housing_df_with_id$median_house_value_normalized[val_id]
)
```

```R
> m_lm_init <- cmdstanr::cmdstan_model("lm_housing.stan",
                                      cpp_options = list(
                                          stan_threads = TRUE
                                      )
 ) 
> m_lm_estimate <- m_lm_init$variational(
     seed = 12345,
     threads = 6,
     data = lm_data_list,
     iter = 10000
 )
------------------------------------------------------------ 
EXPERIMENTAL ALGORITHM: 
  This procedure has not been thoroughly tested and may be unstable 
  or buggy. The interface is subject to change. 
------------------------------------------------------------ 
Gradient evaluation took 0.000584 seconds 
1000 transitions using 10 leapfrog steps per transition would take 5.84 seconds. 
Adjust your expectations accordingly! 
Begin eta adaptation. 
Iteration:   1 / 250 [  0%]  (Adaptation) 
Iteration:  50 / 250 [ 20%]  (Adaptation) 
Iteration: 100 / 250 [ 40%]  (Adaptation) 
Iteration: 150 / 250 [ 60%]  (Adaptation) 
Iteration: 200 / 250 [ 80%]  (Adaptation) 
Success! Found best value [eta = 1] earlier than expected. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100      -102845.218             1.000            1.000 
   200       -42848.414             1.200            1.400 
   300       -19248.705             1.209            1.226 
   400       -18091.131             0.923            1.226 
   500       -18824.729             0.746            1.000 
   600       -18377.223             0.626            1.000 
   700       -18115.129             0.538            0.064 
   800       -18248.659             0.472            0.064 
   900       -18118.901             0.420            0.039 
  1000       -18514.626             0.380            0.039 
  1100       -18221.356             0.282            0.024 
  1200       -18009.634             0.143            0.021 
  1300       -17973.216             0.021            0.016 
  1400       -18080.296             0.015            0.014 
  1500       -18190.332             0.012            0.012 
  1600       -17936.513             0.011            0.012 
  1700       -18047.411             0.010            0.007   MEAN ELBO CONVERGED   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  2.1 seconds.
> m_lm_summary <- m_lm_estimate$summary()
```
さて、ここで検証データに対する対数尤度でモデルのパフォーマンスを比較しましょう！


```R
m_dh_estimate$draws("sum_log_likelihood") |> 
  as.numeric() |>
  tibble::tibble(
    log_lik = _
  ) |>
  dplyr::mutate(
    model = "dirichlet"
  ) |>
  dplyr::bind_rows(
    m_lm_estimate$draws("sum_log_likelihood") |> 
      as.numeric() |>
      tibble::tibble(
        log_lik = _
      ) |>
      dplyr::mutate(
        model = "linear model"
      )
  ) |>
  ggplot2::ggplot() + 
  ggplot2::geom_density(ggplot2::aes(x = log_lik, fill = model),
                        alpha = 0.6) + 
  ggplot2::labs(title = "model comparison")
```

![model_comparison.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/d1c370b1-6aa1-35f6-dfe4-86437f2584fa.png)


対数尤度は高ければ高いほど、データの分布をよく表していることを示しますので、圧倒的にディリクレ過程回帰が勝っています！

# 地理特性の可視化

続いては、ディリクレ過程回帰の意思決定を可視化するため、クラスタリングの方でデータをどのように分類しているかを可視化します。

まず、検証データの地理分布を確認しましょう：

```R
usmap::plot_usmap("counties", include = "California", labels = TRUE) + 
  ggplot2::geom_sf(data = housing_df_with_id[val_id,] |>
                     dplyr::select(longitude, latitude) |>
                     dplyr::rename(lon = longitude, lat = latitude) |>
                     usmap::usmap_transform(),
                   color = "blue",
                   alpha = 0.1) 
```

![original.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/e7593cd4-6fbd-6550-7df7-263778099d19.png)

続いては、実際にモデルに利用されたクラスターの緯度経度管轄点の事後分布を可視化して、それぞれのクラスターがどの辺りの観測地の予測を管轄しているかを確認しよう：

```R
# 検証データのetaの事後分布の平均を出して、
estimated_eta <- m_dh_summary |>
  dplyr::filter(stringr::str_detect(variable, "estimated_eta")) |>
  dplyr::pull(mean) |>
  matrix(ncol = 10)

# その中から最大の数字が出ているクラスターをその観測値のクラスターとする
estimated_group <- purrr::map_int(
  estimated_eta |>
    nrow() |>
    seq_len(),
  \(x){
    which.max(estimated_eta[x,])
  }
)

lonlat_sample_df <- purrr::map(
  10 |>
    seq_len(),
  \(x){
    # 10のグループについてのforループのような処理で、
    # m_dh_estimateからmiddle_point_x[今のクラスター,1]と
    # middle_point_x[今のクラスター,2]を取り出す
    regular_expression <- stringr::str_c("\\[", x, ",1\\]$|\\[", x, ",2\\]$")
    
    df <- m_dh_estimate$draws("middle_point_x") |> 
      tibble::as_tibble() |> 
      dplyr::select(dplyr::matches(regular_expression)) |>
      `colnames<-`(c("lon", "lat")) |>
      dplyr::mutate(
        lon = as.numeric(lon) * sd(housing_df$longitude) + mean(housing_df$longitude),
        lat = as.numeric(lat) * sd(housing_df$latitude) + mean(housing_df$latitude),
        group = x
      )
    return(df)
  }
) |>
  dplyr::bind_rows() |>
  dplyr::filter(
    group %in% unique(estimated_group)
  ) |>
  usmap::usmap_transform()


usmap::plot_usmap("counties", include = "California", labels = TRUE) + 
  ggplot2::geom_sf(data = lonlat_sample_df,
                   ggplot2::aes(color = as.character(group)),
                   alpha = 0.3) + 
  ggplot2::theme(legend.position = "right")

```

![middle_points.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/1e95dee6-a0d0-ff8a-4fbe-118bcc64e045.png)

このクラスタリングの妥当性はカリフォルニア州に土地勘のある方にお任せしたいですが、ロサンゼルスとKern郡の方にすごく緯度経度分散の低いクラスターもありますが、残りのクラスターはそもそもあまり観測値がない気がします。

この現象の理由としては、モデルは緯度経度以外にも、さまざまな説明変数の分布や説明変数と従属変数の関係を参考にクラスタリングしているからです。

最後に、検証データがどのクラスターに分類されているかを可視化します：

```R
lonlat_origin <- housing_df_with_id[val_id,] |>
  dplyr::bind_cols(
    estimated_group = estimated_group
  ) |>
  dplyr::rename(lon = longitude, lat = latitude) |>
  usmap::usmap_transform()

usmap::plot_usmap("counties", include = "California", labels = TRUE) + 
  ggplot2::geom_sf(data = lonlat_origin,
                   ggplot2::aes(color = as.character(estimated_group)),
                   alpha = 0.2) +
  ggplot2::theme(legend.position = "right")
```

![clustered.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/cff27cbc-eb05-721c-dfb6-ebecba237623.png)

ロサンゼルス以外の地域はあまり明確なクラスタリングされているとは言えないですね。これは地理情報と従属変数モデルの関連はロサンゼルス以外ではあまりないことを意味します。

# 結論

いかがでしたか？このように、ディリクレ過程を利用して、「説明変数の分布」と「説明変数と従属変数の関係」でクラスタリングすれば、普通の線形回帰よりいい質のモデルが作れます。

皆さんもぜひディリクレ過程を活用してください！