ゼロがなくなったポアソン分布のパラメーターをData augmentationで推定する

ゼロがなくなったポアソン分布のパラメーターをData augmentation（データ拡大）で推定するということをやってみたメモです。

【追記 2017-12-09 10:30】 すみません。Nの計算を間違えていました（松浦さん、ご指摘ありがとうございます）。ということで、修正しました。

元データを生成します。

library(rstan)

set.seed(123)
N <- 100
lambda <- 2

## Original data
x1 <- rpois(N, lambda)
ggplot(data.frame(x = x1)) +
  geom_bar(aes(x))

> mean(x1)
[1] 2.02

ゼロの部分を消去して切断データにします。

## Truncated data
x2 <- x1[x1 > 0]
ggplot(data.frame(x = x2)) +
  geom_bar(aes(x)) + xlim(0, NA)

> length(x2)
[1] 87
> mean(x2)
[1] 2.321839

ダミーデータを加えます（Data augmentation）。

## Data Augmentation
x3 <- c(x2, rep(0, 400))

元データ（切断前のデータ）が拡大後のデータにしめる割合です。

> N / length(x3)
[1] 0.2053388

Stanのプログラムです。Zero-Infrated modelと同型です。omegaは、元データ（切断前のデータ）が拡大後のデータにしめる割合です。【修正】前のコードではomegaの範囲を[C/M, 1]としていましたが、Nの計算方法を修正したので、ここも[0, 1]としました。

data {
  int<lower = 0> M;
  int<lower = 0> X[M];
}

transformed data {
  int C = 0; // Number of non-zero data

  for (m in 1:M)
    C = C + ((X[m] > 0) ? 1 : 0);
}

parameters {
  real<lower = 0> lambda;
  real<lower = 0, upper = 1> omega;
}

model {
  for (m in 1:M) {
    if (X[m] > 0) {
      1 ~ bernoulli(omega);
      X[m] ~ poisson(lambda);
    } else {
      target += log_sum_exp(bernoulli_lpmf(0 | omega),
                            bernoulli_lpmf(1 | omega)
                            + poisson_lpmf(0 | lambda));
    }
  }
}

generated quantities {
  int<lower = C, upper = M> N;
  {
    // Bern(1 | ω) Pois(0 | λ) /
    //     (Bern(1 | ω) Pois(0 | λ) + Bern(0 | ω))
    real p = (omega * exp(-lambda))
              / (omega * exp(-lambda) + (1 - omega));
    N = C + binomial_rng(M - C, p);
  }
}

【修正】前のコードではp = omega - real_C / M;としていましたが、上のように修正しました。

RStanから実行しました。

fit <- stan("DA.stan",
            data = list(M = length(x3), X = x3))

結果です。

Inference for Stan model: DA.
4 chains, each with iter=2000; warmup=1000; thin=1; 
post-warmup draws per chain=1000, total post-warmup draws=4000.

          mean se_mean   sd    2.5%     25%     50%     75%   97.5% n_eff Rhat
lambda    2.02    0.00 0.17    1.70    1.91    2.02    2.14    2.37  2503    1
omega     0.21    0.00 0.02    0.17    0.19    0.21    0.22    0.25  2432    1
N       100.86    0.09 4.69   93.00   98.00  100.00  104.00  111.00  2976    1
lp__   -252.00    0.02 0.98 -254.61 -252.37 -251.69 -251.29 -251.03  1639    1

タグ：STAn R