Sampling Steps
Contents
Sampling Steps¶
NUTS¶
-
aemcmc.nuts.step(srng, to_sample_rvs, realized_rvs_to_values)[source]¶ Build a NUTS sampling step and its initial state.
This sampling step works with variables in their original space, to create the sampling step we thus need to:
Create the initial value for the variables to sample;
Create a log-density graph that works in the transformed space, and build a NUTS kernel that uses this graph;
Apply the default transformations to the initial values;
Apply the NUTS kernel to the transformed initial values;
Apply the backward transformation to the updated values.
- Parameters:
rvs_to_samples – A dictionary that maps the random variables whose posterior distribution we wish to sample from to their initial values.
realized_rvs_to_values – A dictionary that maps the random variables not sampled by NUTS to their realized value. These variables can either correpond to observations, or to variables whose value is set by a different sampler.
- Returns:
A NUTS sampling step for each random variable, their initial values, the
shared variable updates and the NUTS parameters.
FFBS¶
-
aemcmc.ffbs.ffbs_step(gamma_0, Gammas, log_lik, srng, lower_prec_bound=1e-20)[source]¶ Draw a discrete state sequence sample using forward-filtering backward-sampling (FFBS) [fs].
FFBS draws samples according to
\[\begin{split}S_T &\sim p(S_T | y_{1:T}) \\ S_t \mid S_{t+1} &\sim p(S_{t+1} | S_t) p(S_{t+1} \mid y_{1:T})\end{split}\]for discrete states in the sequence \(S_t\), \(t \in \{0, \dots, T\}\) and observations \(y_t\).
The argument
gamma_0corresponds to \(S_0\),Gammasto \(p(S_{t+1} | S_t)\), andlog_likto \(\log p(y_t \mid S_t)\). The latter is used to derive the forward quantities \(p(S_{t+1} \mid y_{1:T})\).This implementations is most similar to the one in [nr], which forgoes expensive log-space calculations by using fixed precision bounds and re-normalization/scaling.
- Parameters:
gamma_0 – The initial state probabilities as an array of base shape
(M,).Gamma – The transition probability matrices. This array should take the base shape
(N, M, M), whereNis the state sequence length andMis the number of distinct states. IfNis1, the single transition matrix will broadcast across all elements of the state sequence.log_lik – An array of base shape
(M, N)consisting of the log-likelihood values for each state value at each point in the sequence.lower_prec_bound – A number indicating when forward probabilities are too small and need to be re-normalized/scaled.
- Returns:
A tuple containing the tensor representing the FFBS sampled states and the
Scan-generated updates.
References
[fs]Sylvia Fruehwirth-Schnatter, “Markov chain Monte Carlo estimation of classical and dynamic switching and mixture models”, Journal of the American Statistical Association 96 (2001): 194–209
[nr]Press, William H., Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery. 2007. “Numerical Recipes 3rd Edition: The Art of Scientific Computing.” Cambridge university press.
Gibbs¶
-
aemcmc.gibbs.normal_regression_posterior(srng, omega, lmbdatau_inv, X, z)[source]¶ Sample from the posterior of a normal prior and normal observation model.
\[\begin{split}\begin{align*} Z &\sim \operatorname{N}\left( X \beta, \Omega^{-1} \right) \\ \beta &\sim \operatorname{N}\left( 0, \tau^2 \Lambda \right) \end{align*}\end{split}\]where \(X \in \mathbb{R}^{n \times p}\), \(\Lambda = \operatorname{diag}\left(\lambda^2_1, \dots, \lambda^2_p\right)\), and \(\Omega^{-1} = \operatorname{diag}\left(\omega_1, \dots, \omega_n\right)\).
The posterior distribution of \(\beta\) is given by
\[\begin{split}\begin{align*} \left( \beta \mid Z = z \right) &\sim \operatorname{N}\left( A^{-1} X^{\top} \Omega z, A^{-1} \right) \\ A &= X^{\top} X + \Lambda^{-1}_{*} \\ \Lambda_{*} &= \tau^2 \Lambda \end{align*}\end{split}\]This function chooses the best sampler for \(\beta \mid z\) based on the dimensions of \(X\).
- Parameters:
srng – The random number generator used to draw samples.
omega – The observation model diagonal std. dev. values \(\omega_i\). In other words, \(\operatorname{diag}\left(\Omega\right)\).
lmbdatau_inv – The inverse \(beta\) std. dev. values \(\tau^{-1} \lambda^{-1}_i\). In other words, \(\operatorname{diag}\left(\Lambda^{-1/2}_{*}\right)\).
X – Regression matrix \(X\).
z – Observed values \(z \sim Z\).
- Returns:
- Return type:
A sample from \(\beta \mid z\).
-
aemcmc.gibbs.horseshoe_posterior(srng, beta, sigma2, lambda2, tau2)[source]¶ Gibbs kernel to sample from the posterior distributions of the horseshoe prior shrinkage parameters.
This kernel generates samples from the posterior distribution of the local and global shrinkage parameters of a horseshoe prior, respectively \(\lambda\) and \(\tau\) in the following model:
\[\begin{split}\begin{align*} \beta_j &\sim \operatorname{N}\left(0, \lambda_j^2 \tau^2 \sigma^2\right) \\ \lambda_j &\sim \operatorname{C}^{+}(0, 1) \\ \tau &\sim \operatorname{C}^{+}(0, 1) \end{align*}\end{split}\]The graphs constructed by this function are \(\lambda \mid \beta, \tau\) and \(\tau \mid \lambda\), respectively.
We use the following observations [1]_ to sample from the posterior conditional probability of \(\tau\) and \(\lambda\):
1. The half-Cauchy distribution can be intepreted as a mixture of inverse-gamma distributions; 2. If \(Z \sim \operatorname{InvGamma}(1, a)\), \(Z \sim 1 / \operatorname{Exp}(a)\).
- Parameters:
srng – The random number generating object to be used during sampling.
beta – Regression coefficients.
sigma2 – Variance of the regression coefficients.
lambda2 – Square of the local shrinkage parameters.
tau2 – Square of the global shrinkage parameters.
- Returns:
- Return type:
Posteriors for \(lambda\) and \(tau\), respectively.
References
-
aemcmc.gibbs.nbinom_dispersion_posterior(srng, h, p, a, b, y)[source]¶ Sample the conditional posterior for the dispersion parameter under a negative-binomial and gamma prior.
In other words, this draws a sample from \(h \mid Y = y\) per
\[\begin{split}\begin{align*} Y_i &\sim \operatorname{NB}(h, p_i) \\ h &\sim \operatorname{Gamma}(a, b) \end{align*}\end{split}\]where \(\operatorname{NB}\) is a negative-binomial distribution.
The conditional posterior sample step is derived from the following decomposition:
\[\begin{gather*} Y_i = \sum_{j=1}^{l_i} u_{i j}, \quad u_{i j} \sim \operatorname{Log}(p_i), \quad l_i \sim \operatorname{Pois}\left(-h \log(1 - p_i)\right) \end{gather*}\]where \(\operatorname{Log}\) is the logarithmic distribution. Under a gamma prior, \(h\) is conjugate to \(l\). We draw samples from \(l\) according to \(l \sim \operatorname{CRT(y, h)}\), where \(y\) is a sample from \(y \sim Y\).
The resulting posterior is
\[\begin{gather*} \left(h \mid Y = y\right) \sim \operatorname{Gamma}\left(a + \sum_{i=1}^N l_i, \frac{1}{1/b + \sum_{i=1}^N \log(1 - p_i)} \right) \end{gather*}\]- Parameters:
srng – The random number generator from which samples are drawn.
h – The value of \(h\).
p – The success probability parameter in the negative-binomial distribution of \(Y\).
a – The shape parameter in the \(\operatorname{Gamma}\) prior on \(h\).
b – The rate parameter in the \(\operatorname{Gamma}\) prior on \(h\).
y – A sample from \(Y\).
- Returns:
- Return type:
A sample from the posterior \(h \mid y\).
References
-
aemcmc.gibbs.nbinom_normal_posterior(srng, beta, beta_std, X, h, y)[source]¶ Produce a Gibbs sample step for a negative binomial logistic-link regression with a normal prior.
The implementation follows the sampler described in [1]_. It is designed to sample efficiently from the following negative binomial regression model:
\[\begin{split}\begin{align*} Y &\sim \operatorname{NB}\left(h, p\right) \\ p &= \frac{\exp(\psi)}{1 + \exp(\psi)} \\ \psi &= X^\top \beta \\ \beta &\sim \operatorname{N}(0, \lambda^2) \\ \end{align*}\end{split}\]where \(\operatorname{NB}\) is a negative-binomial distribution.
- Parameters:
srng – The random number generator from which samples are drawn.
beta – The current/previous value of the regression parameter \(beta\).
beta_std – The std. dev. of the regression parameter \(beta\).
X – The regression matrix.
h – The \(h\) parameter in the negative-binomial distribution of \(Y\).
y – A sample from the observation distribution \(y \sim Y\).
- Returns:
- Return type:
A sample from the posterior \(\beta \mid y\).
Notes
The \(z\) expression in Section 2.2 of [1]_ seems to omit division by the Polya-Gamma auxilliary variables whereas [2] and [3] explicitly include it. We found that including the division results in accurate posterior samples for the regression coefficients. It is also worth noting that the \(\sigma^2\) parameter is not sampled directly in the negative binomial regression problem and thus set to 1 [2].
References
-
aemcmc.gibbs.bern_normal_posterior(srng, beta, beta_std, X, y)[source]¶ Produce a Gibbs sample step for a bernoulli logistic-link regression with a normal prior.
The implementation follows the sampler described in [1]_. It is designed to sample efficiently from the following binary logistic regression model:
\[\begin{split}\begin{align*} Y &\sim \operatorname{Bern}\left( p \right) \\ p &= \frac{1}{1 + \exp\left( -X^\top \beta\right)} \\ \beta_j &\sim \operatorname{N}\left( 0, \lambda_j^2 \right) \end{align*}\end{split}\]- Parameters:
beta – The current/previous value of the regression parameter \(beta\).
beta_std – The std. dev. of the regression parameter \(beta\).
X – The regression matrix.
y – A sample from the observation distribution \(y \sim Y\).
- Returns:
- Return type:
A sample from \(\beta \mid y\).
References