Thinking about statistical models, estimation

model.stan

data {
  int<lower=0> n;
  int y[n];
  vector[n] x;
}
parameters {
  real<lower=0> mu;
  real<lower=0> delta;
}
transformed parameters {
  vector[n] lambda = mu + delta * x;
}
model {
  y ~ poisson(lambda);
}

stat_model.R

# Jan 18 2023
# Michael Love

# premise: suppose we have some counts (of something)
# before and after treatment, and we want to estimate
# delta, the effect of treatment

n <- 20 # sample size
delta <- 3 # true effect
mu <- 5 # baseline counts

makeData <- function(n, mu, delta) {
  y0 <- rpois(n, mu)
  y1 <- rpois(n, mu + delta)
  data.frame(y=c(y0,y1), group=rep(c("t0","t1"),each=n))
}

set.seed(1)
dat <- makeData(n, mu, delta)

model <- function(par, dat) {
  mu <- par[1]
  delta <- par[2]
  # lambda = parameter of a probability distribution
  lambda <- rep(c(mu, mu+delta), each=n) 
  sum(dpois(dat$y, lambda, log=TRUE)) # the log likelihood
}

model(c(4.9, 3.1), dat)
model(c(5.1, 2.9), dat)

# general purpose optimization (maximize the likelihood)
res <- optim(c(10,10), model, control=list(fnscale=-1), dat=dat)
res

mle <- data.frame(mu=res$par[1], delta=res$par[2])
mle$ll <- model(res$par, dat)

library(dplyr)
library(tidyr)
library(purrr)
library(ggplot2)

grid <- expand_grid(mu=seq(4,6,length.out=100),
                    delta=seq(2,4,length.out=100))
grid %>%
  mutate(ll = map2_dbl(mu, delta, \(x, y) model(c(x, y), dat))) %>%
  ggplot(aes(mu, delta, z=ll)) +
  geom_contour() +
  geom_point(data=mle, shape=4, size=2)

# bayes fitting

library(rstan)
x <- ifelse(dat$group == "t1", 1, 0)
fit <- stan("model.stan", data=list(n=2*n, y=dat$y, x=x), pars=c("mu","delta"))
stan_plot(fit)
stan_scat(fit, pars=c("mu","delta"))