goldingn
commented
3 months ago Here's a script with my current hacky functions, using the current greta_2 branch (which has time-varying parameters etc and bug fixes), applied to a problem where the population goes extinct.
# do 1D version of stochastic dynamics in greta.dynamics
# simulate a discrete stochastic growth rate population model
set.seed(3)
n_times <- 50
# daily population growth rate
r_true <- 1.01
# initial population (just before timeseries)
pop_init_true <- 25
time <- seq_len(n_times)
pop_true <- rep(NA, n_times)
pop_previous <- pop_init_true
for (i in 1:n_times) {
pop_true[i] <- rpois(1, pop_previous * r_true)
pop_previous <- pop_true[i]
}
# add an observation process (binomial with fixed detection probability)
obs_prob <- 0.7
pop_obs <- rbinom(n_times, size = pop_true, prob = obs_prob)
# define functions for continuous relaxation and reparameterisation of the
# Poisson distributions
# given poisson rate parameter lambda and random uniform deviate u, a continuous
# relaxation of poisson random variable generation is computed using the inverse
# of the incomplete gamma function. ie. igammainv(lambda, 1 - u) is
# approximately equal to qpois(u, lambda) (and exactly equal to qgamma(1 - u,
# lambda) in the R implementation)
gamma_continuous_poisson <- function(lambda, u) {
igammainv(lambda, 1 - u)
}
# # check:
# lambda <- as_data(pi)
# u <- uniform(0, 1)
# y <- gamma_continuous_poisson(lambda, u)
# sims <- calculate(y, u, nsim = 1e5)
# max(abs(sims$y - qgamma(1 - sims$u, pi)))
# quantile(round(sims$y))
# quantile(rpois(1e5, pi))
# the inverse incomplete gamma function (the major part of the quantile function
# of a gamma distribution)
igammainv <- function(a, p) {
op <- greta::.internals$nodes$constructors$op
op("igammainv", a, p,
tf_operation = "tf_igammainv"
)
}
tf_igammainv <- function(a, p) {
tfp <- greta:::tfp
tfp$math$igammainv(a, p)
}
# given random variables z (with standard normal distribution a priori), and
# Poisson rate parameter lambda, return a strictly positive continuous random
# variable with the same mean and variance as a poisson random variable with
# rate lambda, by approximating the poisson as a lognormal distribution.
lognormal_continuous_poisson <- function(lambda, z) {
sigma <- sqrt(log1p(lambda / exp(2 * log(lambda))))
# sigma2 <- log1p(1 / lambda)
mu <- log(lambda) - sigma^2 / 2
exp(mu + z * sigma)
}
# Working: The lognormal mean and variance should both equal lambda. The
# lognormal mean and variance can both be expressed in terms of the parameters
# mu and sigma.
# mean = lambda = exp(mu + sigma^2 / 2)
# variance = lambda = (exp(sigma ^ 2) - 1) * exp(2 * mu + sigma ^ 2)
# solve to get sigma and mu as a function of lambda:
# mu = log(lambda) - sigma^2 / 2
# lambda = (exp(sigma ^ 2) - 1) * exp(2 * mu + sigma ^ 2)
# lambda = (exp(sigma ^ 2) - 1) * exp(2 * (log(lambda) - sigma^2 / 2) + sigma ^ 2)
# lambda = (exp(sigma ^ 2) - 1) * exp(sigma ^ 2) * exp(2 * (log(lambda) - sigma^2 / 2))
# lambda = (exp(sigma ^ 2) - 1) * exp(sigma ^ 2) * exp(2 * log(lambda) - sigma^2)
# lambda = (exp(sigma ^ 2) - 1) * exp(sigma ^ 2) * exp(2 * log(lambda)) / exp(sigma^2)
# lambda / exp(2 * log(lambda)) = (exp(sigma ^ 2) - 1) * exp(sigma ^ 2) * 1 / exp(sigma^2)
# lambda / exp(2 * log(lambda)) = (exp(sigma ^ 2) - 1)
# log(lambda / exp(2 * log(lambda)) + 1) = sigma ^ 2
# sigma = sqrt(log(lambda / exp(2 * log(lambda)) + 1)) = sigma
# mu = log(lambda) - sigma^2 / 2
# # check these numerically
# library(tidyverse)
# compare <- tibble(
# lambda = seq(0.01, 1000, length.out = 100)
# ) %>%
# mutate(
# sigma = sqrt(log(lambda / exp(2 * log(lambda)) + 1)),
# mu = log(lambda) - sigma^2 / 2
# ) %>%
# mutate(
# mean = exp(mu + sigma^2 / 2),
# variance = (exp(sigma ^ 2) - 1) * exp(2 * mu + sigma ^ 2)
# ) %>%
# mutate(
# diff_mean_variance = abs(mean - variance),
# diff_mean_lambda = abs(mean - lambda),
# diff_variance_lambda = abs(variance - lambda)
# ) %>%
# summarise(
# across(
# starts_with("diff"),
# ~max(.x)
# )
# )
library(greta.dynamics)
#> Loading required package: greta
#>
#> Attaching package: 'greta'
#> The following objects are masked from 'package:stats':
#>
#> binomial, cov2cor, poisson
#> The following objects are masked from 'package:base':
#>
#> %*%, apply, backsolve, beta, chol2inv, colMeans, colSums, diag,
#> eigen, forwardsolve, gamma, identity, rowMeans, rowSums, sweep,
#> tapply
# growth rate
r <- normal(1, 0.1, truncation = c(0, Inf))
#> ℹ Initialising python and checking dependencies, this may take a moment.
#> ✔ Initialising python and checking dependencies ... done!
#>
# latent variable for stochastic transitions
latent_z_vec <- normal(0, 1, dim = n_times)
# latent_u_vec <- uniform(0, 1, dim = n_times)
# initial population size (with mean equal to truth)
init <- exponential(1 / pop_init_true)
# transition function for the population process as a difference equation in
# integer timestep
fun <- function(state, iter, r, latent_z) {
lambda <- state * r
state_new <- lognormal_continuous_poisson(lambda, latent_z)
# state_new <- gamma_continuous_poisson(lambda, latent_u)
state_new
}
# solve it (integer times)
pop_sim <- iterate_dynamic_function(
transition_function = fun,
initial_state = init,
niter = n_times,
tol = 0,
r = r,
latent_z = latent_z_vec,
parameter_is_time_varying = "latent_z",
# clamp the populations to reasonably values to avoid numerical under/overflow
state_limits = c(1e-3, 1e3)
)
# get the modelled true population
pop_modelled <- t(pop_sim$all_states)
# and the expected value of the observation distribution
pop_obs_expected <- pop_modelled * obs_prob
pop_obs_ga <- as_data(pop_obs)
distribution(pop_obs_ga) <- poisson(pop_obs_expected)
# fit this, reducing the stochasticity when finding the initial values
n_chains <- 4
inits <- replicate(n_chains,
initials(
init = pop_init_true,
latent_z_vec = rnorm(n_times, 0, 0.1)),
simplify = FALSE)
m <- model(r, latent_z_vec, init)
draws <- mcmc(m,
initial_values = inits)
#> running 4 chains simultaneously on up to 8 CPU cores
#>
#> warmup ====================================== 1000/1000 | eta: 0s
#> sampling ====================================== 1000/1000 | eta: 0s
# get posterior samples and plot summaries
sims <- calculate(pop_modelled, nsim = 1000, values = draws)
posterior_sims <- sims$pop_modelled[, , 1]
posterior_sims_discrete <- round(posterior_sims)
posterior_est <- colMeans(posterior_sims_discrete)
posterior_ci <- apply(posterior_sims_discrete, 2, quantile, c(0.025, 0.975))
# plot posterior draws, summary stats, and the truth
plot(posterior_est,
ylim = range(c(posterior_ci, pop_true)),
type = "n",
xlab = "day",
ylab = "population")
for (i in 1:50) {
lines(posterior_sims_discrete[i, ],
lwd = 0.1,
col = grey(0.4))
}
lines(posterior_ci[1, ], lty = 2)
lines(posterior_ci[2, ], lty = 2)
lines(posterior_est,
lwd = 1.5)
lines(pop_true,
col = "blue")
# noisy observations, naively adjusted for detection probability
points(pop_obs / obs_prob,
cex = 0.5)
Created on 2024-03-13 with reprex v2.0.2
Background
Gradient-based inference (like the HMC greta uses) can only operate on continuous parameter spaces. That means it cannot learn the values of parameters with discrete support (e.g. no unobserved Poisson random variables).
But demographic stochasticity due to discrete stochastic variation in population sizes between timesteps in models of populations and infections (e.g. the number of new infectees is Poisson, the number of individuals surviving is binomial) is often important, especially when populations reach low numbers and near extinction. We cannot directly model the values of these discrete random variables, but we can apply continuous relaxations to approximate them; keeping the state values continuous and replacing the discrete random variables with continuous random variables that matches the mean and variance (and ideally the full shape of the distribution) of the random variable we would like to model as stochastic.
Continuous approximations to distributions
E.g. a poisson random variable can be approximated with an appropriately-shaped gamma distribution that exactly matches the PMF at discrete values, or by some other distribution that is a a close-enough approximation:
We might write a discrete stochastic growth-rate model like this:
where $x$ takes integer values, $Poisson(\lambda)$ is the Poisson distribution, and $r$ is a positive-valued growth rate parameter. To estimate the posterior over the values represented by $x$ in this model, but using HMC, we could instead fit the model:
where $y_t$ is a (strictly positive) real-valued parameter, and $\pi(\lambda)$ is some probability distribution with support on positive real values that has similar moments (men, variance, skewness, etc.) to $Poisson(\lambda)$.
Reparameterisation
If we structure these probability distributions such that they can be reparameterised in terms of the parameter and some latent 'innovation' or noise, with known distribution, we can significantly improve the ability to sample these models, since we can decorrelate the posterior distribution in a similar way to the reparameterisation trick for hierarchical models. This can also provide some computational advantages in greta/tensorflow by working with arrays rather than scalars.
If we know the quantile function of the continuous distribution (e.g. $q_{\pi}(p, \lambda)$ as the quantile function of $\pi(\lambda)$ , with $p$ the probability argument), we can reparameterise the innovations as $u \sim U(0, 1)$, and then plug them into the quantile function to sample poisson values. Ie. equation 2 above is equivalent to:
3.
The vector of $u$ values can then be computed in advance, and passed into the solvers to be chopped up appropriately. The dependency structure in the model means $yt$ depends on $y{t-1}$, and so they are a priori (and therefore also a posteriori) correlated. But in this reparameterised formulation, HMC operates instead on $u$, and $ut$ doesn't depend on $u{t-1}$ so they are a priori uncorrelated, which removes a lot of correlation in the posterior and makes sampling much easier.
Note that if the quantile function is expensive to compute, other approximations and reparameterisations may be more appealing. E.g. for the Poisson, the inverse of the incomplete gamma function gives the quantile of the gamma distribution whose PDF matches the Poisson PMF at discrete values, but the function has no analytic form and is expensive to compute. A lognormal approximation (with either uniform or standard normal innovations) is imperfect but much more computationally efficient.
Currently this reparameterisation trick is only applicable in the
greta_2branch withiterate_dynamic_function(), since it requires functionality for indexing time-varying parameters.Implementation
We just need to provide functions for the relaxations, documentation, and examples of applying this approach. I have some, I just need to add them to the
greta_2branch and work out the neatest user interface.