EBPM problem with spike-and-slab prior

\[ \begin{align} & x_i \sim Pois(s_i \lambda_i)\\ & \lambda_i \sim g(.)\\ & g \in \mathcal{G} \end{align} \]

where \(\mathcal{G} = \{\pi_0 \delta(.) + (1-\pi_0) gamma(a,b): \pi_0 \in [0,1] \}\)
Now the goal is to compute \(\hat{\pi}_0,\hat{a}, \hat{b}\) with MLE, then compute posterior mean of \(\lambda_i\).

pg_nlm_fn <- function(par, x, s){
  pi = 1/(1+ exp(-par[1]))
  a = exp(par[2])
  b  =  exp(par[3])
  d <- dnbinom(x, a, b/(b+s), log = F) 
  c = as.integer(x ==  0) - d
  return(-sum(log(pi*c + d)))
}

transform_param <- function(par0){
  par = rep(0,length(par0))
  par[1] = log(par0[1]/(1-par0[1]))
  par[2] = log(par0[2])
  par[3] = log(par0[3])
  return(par)
}

transform_param_back <- function(par){
  par0 = rep(0,length(par))
  #par0[1] = log(par[1]) - log(1-par[1])
  par0[1] = 1/(1+ exp(-par[1]))
  par0[2] = exp(par[2])
  par0[3] = exp(par[3])
  return(par0)
}

sim_spike_one <- function(pi, a, b){
  if(rbinom(1,1, pi)){return(0)}
  else{return(rgamma(1,shape = a, rate = b))}
}

simulate_pm <- function(s, param){
  pi = param[1]
  a = param[2]
  b  = param[3]
  lam = replicate(length(s), sim_spike_one(pi, a, b))
  x = rpois(length(s), s*lam)
  ll = -pg_nlm_fn(transform_param(param), x, s)
  return(list(x = x, s= s, lam = lam, param = param, ll = ll))
}

n = 4000
s = replicate(n, 1)
pi  = 0.8
a = 100
b  = 1
param =  c(pi, a, b)
sim = simulate_pm(s, param)

init_par = c(0.5,1,1)
opt = nlm(pg_nlm_fn, transform_param(init_par), sim$x, sim$s)
opt_par = transform_param_back(opt$estimate)

[1] "oracle ll: -5166.363892"

[1] "opt    ll: -5165.270752"

[1] "oracle:pi, a, b"

[1]   0.8 100.0   1.0

[1] "estimate: pi, a, b"

[1]   0.80 107.79   1.08

Wrap up into ebpm algorithm

It is easy to deduce posterior mean:

\[ \begin{align} \text{posterior mean} = \frac{(1-\pi_0)NB(x; a, \frac{b}{b + s}) \frac{a+x}{b+s}}{\pi_0 \delta(x) + (1-\pi_0)NB(x; a , \frac{b}{b + s})} \end{align} \]

ebpm_point_gamma_demo <- function(x, s, init_par = c(0.5,1,1), seed = 123){
  set.seed(seed) ## though seems determined
  ## MLE
  opt = nlm(pg_nlm_fn, transform_param(init_par), x, s)
  opt_par = transform_param_back(opt$estimate)
  ll =  -pg_nlm_fn(transform_param(opt_par), x, s)

  ## posterior mean
  pi = opt_par[1]
  a =  opt_par[2]
  b =  opt_par[3]
  nb = dnbinom(x, size = a, prob = b/(b+s))
  pm = ((1-pi)*nb*(a+x)/(b+s))/(pi*as.integer(x ==  0) + (1-pi)*nb)
  return(list(param = opt_par, lam_pm = pm, ll = ll))
}

I have packaged the functions above into the ebpm package functin ebpm_point_gamma. Try it out!

library(ebpm)
start = proc.time()
fit <- ebpm_point_gamma(sim$x, sim$s)

Called from: f(x, ...)
debug at /Users/ontheroad/Desktop/git/ebpm/R/ebpm_point_gamma.R#61: return(-sum(log(pi * c + d)))
[1] 1.40673212 4.68017717 0.08098415

runtime = proc.time() - start
print(sprintf("fit %d data with runtime %f  seconds", n, runtime[[3]]))

[1] "fit 4000 data with runtime 0.128000  seconds"

Compare RMSE with \(\lambda_{oracle}\)

[1] "RMSE with lam_oracle:"

[1] "mle    : 4.204930"

[1] "fitted : 3.041565"

df <- data.frame(n = 1:length(sim$x), x = sim$x, s = sim$s, lam = sim$lam, lam_pm = fit$posterior$mean)
ggplot(df)  + geom_point(aes(x = x/s, y = lam_pm), color = "blue", cex = 0.5) +
    labs(x = "x/s", y = "lam_pm", title = "ebpm_point_gamma: x/s vs lam_posterior_mean") +
    guides(fill = "color")

Past versions of unnamed-chunk-10-1.png

[1] "max posterior mean when x = 0"

[1] 3.251273e-30

Let’s take a look at the nonzero (for x) parts.
Note that for \(x \neq 0\), we have posterior mean \(\frac{a+x}{b+s}\). Therefore we expect to see a line, with slope \(1/(1 + \frac{b}{s})\)

df_nz = df[df$x != 0, ]
ggplot(df_nz)  + geom_point(aes(x = x/s, y = lam_pm), color = "blue", cex = 0.5) +
    labs(x = "x/s", y = "lam_pm", title = "ebpm_point_gamma: x/s vs lam_posterior_mean") +
    geom_abline(slope = 1, intercept = 0)+
    guides(fill = "color")

Past versions of unnamed-chunk-12-1.png

now let’s compare \(\lambda_{true}, \lambda_{\text{posterior mean}}\)

ggplot(df_nz)  + geom_point(aes(x = lam, y = lam_pm), color = "blue", cex = 0.5) +
    labs(x = "lam_true", y = "lam_pm", title = "ebpm_point_gamma: lam _true lam_posterior_mean") +
    geom_abline(slope = 1, intercept = 0)+
    guides(fill = "color")

Past versions of unnamed-chunk-13-1.png

Session information

sessionInfo()

R version 3.5.1 (2018-07-02)
Platform: x86_64-apple-darwin15.6.0 (64-bit)
Running under: macOS  10.14

Matrix products: default
BLAS: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRlapack.dylib

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] ebpm_0.0.0.9000 ggplot2_3.2.1  

loaded via a namespace (and not attached):
 [1] Rcpp_1.0.2       knitr_1.25       whisker_0.3-2    magrittr_1.5    
 [5] workflowr_1.4.0  tidyselect_0.2.5 munsell_0.5.0    colorspace_1.4-1
 [9] R6_2.4.0         rlang_0.4.0      dplyr_0.8.1      stringr_1.4.0   
[13] tools_3.5.1      grid_3.5.1       gtable_0.3.0     xfun_0.8        
[17] withr_2.1.2      git2r_0.25.2     htmltools_0.3.6  assertthat_0.2.1
[21] yaml_2.2.0       lazyeval_0.2.2   rprojroot_1.3-2  digest_0.6.21   
[25] tibble_2.1.3     crayon_1.3.4     mixsqp_0.1-120   purrr_0.3.2     
[29] fs_1.3.1         glue_1.3.1       evaluate_0.14    rmarkdown_1.13  
[33] labeling_0.3     stringi_1.4.3    pillar_1.4.2     compiler_3.5.1  
[37] scales_1.0.0     backports_1.1.5  pkgconfig_2.0.3