Following #16 we should implement a Metropolis-within-Gibbs sampler similar to PyMC's (algorithms listed in this file) that uses NUTS to update the RVs with continuous variables and Gibbs samplers / others to update the RVs with discrete support.

This way we can start attracting users who will interact with a generic build_sampler (name tbd) function, and we will have a first performance benchmark. It will be a big step towards #3.

The order in which the sampling steps are chained together can influence the sampling results, we should thus allow users to change the scan order of the samplers we return.

After #76 it will be possible to inspect the sampler graph and identify the sampling steps according to their type. We can then rewire the graphs e.g. as follows:

More generally, we will need to write a series of operators that act on samplers: we may want to change the scan order, randomize the scan order, merge two samplers (e.g. with different scan orders for symmetric scan), vectorize the sampler, draw several samples using scan, etc. Related to #80.

The changes introduced in aesara-devs/aeppl#196 will be merged in AePPL. It should actually simplify the code in nuts.py: rvs_to_sample is now enough to do most things, and variables not sampled by NUTS are introduced in the logprob graph using the realized keyword.

I was wondering if it would make sens to move the implementation of PolyaGammaRV from this repo to aesara?

With the recently added aeppl mixture features and an HMM state sequence like the one implemented in aeppl-hmm, we're now able to make a function (e.g. create_ffbs_sampler) that constructs FFBS sample step graphs given a graph of an HMM model.

To do this, we'll need the intermediate forms provided by aeppl, and, in particular, the Mixture Op. Just as aeppl does to produce a log-probability graph, we can use the same steps to generate the intermediate form, search for eligible Mixture and HMM state sequence nodes, extract the necessary FFBS parameters (i.e. transition matrices, initial state probabilities, state sequence), and use those to construct an FFBS step graph.

Aside from constructing complete sampler loops per #80, we need to set some standard HMC/NUTS defaults (e.g. mass matrix adaptation) and clarify the settings/options interface to the sampler construction process.

Description of your problem or feature request

I want to perform a classification task which consists in predicting whether a document $i$ published at a time $t_i$ belongs to a certain class ( $y_i=1$ ) or not ( $y_i=0$ ). The contents of $i$ is given by a bag of words $b_{ij}$. The predictive features are encoded into a sparse tensor $x_{ijt} = b_{ij} \delta_{t,t_i}$. For this task I am considering a logistic regression with time-varying coefficients, such that the coefficients evolve through a random walk:

$$\text{logit} (p_{i}) = \sum_{jt} x_{ijt} \beta_{jt}$$

$$y_{i} \sim \text{Bernouilli}(p_i)$$

$$\beta_{j,t+1} \sim \text{Normal}(\beta_{j,t},\sigma_j^{-1})$$

$$\beta_{j,0} \sim \text{Normal}(0,\sigma_{j,0}^{-1})$$

$$\sigma_j \sim \text{Exponential}(1)$$

$$\sigma_{j,0} \sim \text{Exponential}(1)$$

For the sake of simplicity below I assume $\sigma_{j,0}=\sigma_j$, although I do not see any compelling reason to do so.

However:

• I cannot construct a sampler for this model (see details below)
• I wonder if there could be a more efficient implementation anyway: $x_{ijt}$ is very sparse, and at.tensordot(X, beta_t_rv, 2) involves too many needless operations. I would not have done that with Stan but in this case I found it easier to write the model this way.

Please provide a minimal, self-contained, and reproducible example.

As a clueless user, I tried to build a sampler based on my generative model by looking at other examples; which means I could be doing something completely idiotic!
Here is what I came up with:

import numpy as np

import aesara
import aesara.tensor as at
from aemcmc.basic import construct_sampler
from aesara.tensor.random.utils import RandomStream

def logistic_fit(X_val, y_val):
    N, M, T = X_val.shape

    srng = RandomStream(0)
    X = at.tensor3("X")

    sigma_rv = srng.exponential(1, size=X.shape[1])
    beta_t_rv = at.cumsum(srng.normal(0, 1/sigma_rv, size=(X.shape[1],X.shape[2])), axis=1)

    eta = at.tensordot(X, beta_t_rv, 2)
    p = at.sigmoid(-eta)
    Y_rv = srng.bernoulli(p, name="Y")

    y_vv = Y_rv.clone()
    y_vv.name = "y"

    sample_vars = [sigma_rv, beta_t_rv]

    sample_steps, updates, initial_values = construct_sampler({Y_rv: y_vv}, srng)

# X_val = np.load("X_val.npy")
# y_val = np.load("y_val.npy")

X_val = np.zeros((1000, 50, 10))
y_val = np.zeros(1000)

beta = logistic_fit(X_val, y_val)

Although thus far the script does not actually use the data, if at some point you want to try with actual data, download and unzip these files (X_val.npy and y_val.npy) to the folder the script is executed from: data.zip

However, this crashes:

Please provide the full traceback of any errors.

ERROR (aesara.graph.rewriting.basic): Rewrite failure due to: local_elemwise_dimshuffle_subsume
ERROR (aesara.graph.rewriting.basic): node: Elemwise{true_div,no_inplace}(InplaceDimShuffle{x}.0, exponential_rv{0, (0,), floatX, False}.out)
ERROR (aesara.graph.rewriting.basic): TRACEBACK:
ERROR (aesara.graph.rewriting.basic): Traceback (most recent call last):
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/graph/rewriting/basic.py", line 1933, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/graph/rewriting/basic.py", line 1092, in transform
    return self.fn(fgraph, node)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/rewriting.py", line 345, in local_elemwise_dimshuffle_subsume
    new_op = SubsumingElemwise(new_inputs, new_outputs, inline=True)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/rewriting.py", line 189, in __init__
    OpFromGraph.__init__(self, inputs, outputs, *args, **kwargs)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/compile/builders.py", line 343, in __init__
    raise TypeError(f"Constants not allowed as inputs; {i}")
TypeError: Constants not allowed as inputs; TensorConstant{1}

Traceback (most recent call last):
  File "examples/gibbs_sample.py", line 61, in <module>
    beta = logistic_fit(X_val, y_val)
  File "examples/gibbs_sample.py", line 26, in logistic_fit
    sample_steps, updates, initial_values = construct_sampler({Y_rv: y_vv}, srng)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/basic.py", line 108, in construct_sampler
    raise NotImplementedError(
NotImplementedError: Could not find a posterior samplers for {exponential_rv{0, (0,), floatX, False}.out, normal_rv{0, (0, 0), floatX, False}.out}

Versions and main components

• Aesara version: 2.8.9
• Aemcmc version: 0.06
• Python version: 3.8.10
• Operating system: macOS 10.15.4
• How did you install Aesara: pip

Consider the following observation model:

import aesara
import aesara.tensor as at

srng = at.random.RandomStream(4238)

alpha = at.scalar()
beta = at.scalar()
observed = at.scalar()

Z_rv = srng.gamma(alpha, beta, name="Z_1")
Y_rv = srng.poisson(Z_rv, name="Y")

Z_rv is a conjugate prior of Y_rv and if the observed value of the latter is observed the model is equivalent to the following:

Z_rv = srng.gamma(alpha + observed, beta + 1)

Tests are currently ran with Python 3.6 and 3.7. However Python 3.6 has reached End Of Life and versions until 3.10 should be supported.

We already have a function that builds FFBS sampling steps since #6, we need to add it to the construct_sampler logic.

Just like #4, we need to incorporate the recentering and rescaling from symbolic-pymc (see the example here).

aemcmc currently does not implement calibration tests for its samplers (i.e. it does not check that the samplers generate samples from the correct distribution), while aehmc's tests are very flaky and can be made to pass or fail with different RNG seeds. While this is still a very open area of research, making sure that the samplers that aemcmc build generate correct samples is critical.

This blog post summarizes the situation fairly well and links to the relevant literature.

My current understanding from the literature is that existing solutions are hypothesis tests which test hypotheses that we don't expect to be true. There seems to be (unsurprisingly) little hope that we can devise tests that require absolutely no supervision. However, there are a few things we could do to automate some of Gelman's suggestions in the post above, and the general idea would be to have some rough tests in CI to avoid obvious miscalibration (check that the chain has moved on a simple problem, check the monte carlo error, etc. )and give the users some tools to evaluate the calibration of the sampler on their use case.

We should set up scheduled nightly builds for the library.

See section 2.1. in [1].

[1] Makalic, Enes, and Daniel F. Schmidt. “High-Dimensional Bayesian Regularised Regression with the BayesReg Package.” ArXiv Preprint ArXiv:1611.06649, 2016.

We should add this sooner than later, so that there's less type-based catching up to do.

A couple of examples to illustrate the capabilities of the library. We can begin with :

• Closed-form posterior
• Sparse regression

Ratio distributions are distributions of random variables $Z$ that can be formed as the ratio $Z = X/Y$ of two other random variables $X$ and $Y$. Common examples are Cauchy distributions as the ratio of two normal distributions, the T distribution and the F distribution.

construct_sample is the high-level interface to AeMCMC. Instead we currently have o import it as from aemcmc.basic import construct_sampler.

It looks like we need to convert model graphs into AePPL IR and operate on those, because construct_sampler/construct_ir_fgraph only supports observed RandomVariables and not general MeasurableVariables.

An example:

import aesara.tensor as at
import aemcmc

srng = at.random.RandomStream(23920)

X_rv = srng.normal(0, 1, name="X")
Y_rv = 1 + X_rv

obs_rvs_to_values = {Y_rv: at.scalar("y")}

sample_steps, updates, initial_values, nuts_parameters = aemcmc.construct_sampler(
    obs_rvs_to_values, srng
)

aemcmc/basic.py:40: in construct_sampler
    fgraph, obs_rvs_to_values, memo, new_to_old_rvs = construct_ir_fgraph(
aemcmc/rewriting.py:91: in construct_ir_fgraph
    obs_rvs_to_values = {memo[k]: v for k, v in obs_rvs_to_values.items()}
aemcmc/rewriting.py:91: in <dictcomp>
    obs_rvs_to_values = {memo[k]: v for k, v in obs_rvs_to_values.items()}
E   KeyError: Elemwise{add,no_inplace}.0

We currently only check that the shape of the output is the expected shape, but we should at least make sure that the value of the mean & covariance of the samples are correct.

We need a framework for mapping samplers to Aesara model graphs.

In general, one of the main interfaces to the sampler implementations in this project would accept a model graph (i.e. an Aesara graph that represents the statistical model) and return a graph that represents the sampler steps that generate posterior samples for the model.

For example, our current Horseshoe Gibbs sampler implementation applies to normal regression models with Horseshoe prior regression parameters. We want to be able to parse a given model graph and determine whether or not those elements are present so that the Gibbs sampler can be proposed/applied.

The act of parsing such models is easily answered by Aesara's basic rewriting and unification functionality, but we need a good way to organize and deploy the process. For instance, how do we want to store the mappings between a unification "pattern" and a set of samplers?

Also, this framework needs to be flexible enough to identify and apply rewrites that could make a model graph amenable to samplers. The primary example is our current normal scale-mixture expansion for the negative-binomial (via the Polya-gamma); after identifying the Horseshoe regression portion of the model (e.g. X.dot(beta_rv), where beta_rv = at.random.normal(0, eta_rv * lambda2_rv) and eta_rv = at.random.halfcauchy(...), etc.), the negative-binomial shouldn't prevent the Horseshoe Gibbs sampler from being applied, because it can be converted into a normal scale-mixture.

This framework will clearly involve a rewrite process similar to AePPL, but there are some unique organizational aspects we need to address as well.

We need to add more scale-mixtures.

The rule is currently described as follows in the docstring:

$$ \begin{align*} p &\sim \operatorname{Beta}\left(\alpha, \beta\right)\\ Y &\sim \operatorname{Binom}\left(N, p\right) \end{align*} $$

If we observe $Y=y$, then $p$ follows a beta distribution:

$$ p \sim \operatorname{Beta}\left(\alpha + y, \beta + N - y\right) $$

Proposed change

$$ \frac{ Y \sim \operatorname{Binom}\left(N, p\right), \quad p \sim \operatorname{Beta}\left(\alpha, \beta\right) }{ \left(p|Y=y\right) \sim \operatorname{Beta}\left(\alpha+y, \beta+N-y\right) } $$

The Leemis chart represents relationships between univariate probability distributions that might be useful to simplify model graphs / rewrite model graphs to build better samplers.

See section 2.2 in [1]. We should be able to reuse most of the negative binomial regularized regression's implementation.

[1] Makalic, Enes, and Daniel F. Schmidt. “High-Dimensional Bayesian Regularised Regression with the BayesReg Package.” ArXiv Preprint ArXiv:1611.06649, 2016.

We should be able use both sparse and dense matrices for X in horseshoe_nbinom_gibbs. While Aesara is designed so that at.dot can be used with both, @ currently raises an error when used with a sparse matrices. I raised an issue on the Aesara repo aesara-devs/aesara#881, in the meantime we can use at.dot here.

Currently, this library provides a few functions that construct samplers for very specific models given the inputs/parameters of said models. We should standardize these constructor functions under a function interface that takes model graphs as inputs (i.e. graphs that represent the models supported by the constructor functions).

This is a first step toward the more general case of #3.

Again, since our current constructor functions apply to very specific types of models, we can keep things simple and employ some very basic/restricted unification/pattern matching and generalize from there.

For example:

import numpy as np
import aesara
import aesara.tensor as at


srng = at.random.RandomStream(seed=2309)

X = at.as_tensor(np.random.normal(size=(100, 10)))

# Horseshoe sub-graph
tau_rv = srng.halfcauchy(size=1)
lambda_rv = srng.halfcauchy(size=10)
beta_rv = srng.normal(0, lambda_rv * tau_rv, size=10)

# The "observation" model
eta = X @ beta_rv
p = at.sigmoid(-eta)
Y_rv = srng.nbinom(10, p)


gibbs_sampler = horseshoe_nbinom_gibbs(Y_rv, obs, n_samples)

where horseshoe_nbinom_gibbs would extract the X, beta_rv, and any other necessary components and subsequently perform the same actions as aemcmc.negative_binomial_horseshoe_gibbs.

1. If $Z \sim \operatorname{Norm}(0,1)$ and $\chi^2 \sim \operatorname{ChiSquared}(n)$ then $Y = \frac{Z}{\sqrt{\chi^2 /n}} \sim \operatorname{T}(n)$. Falls under #65.
2. If $X \sim \operatorname{Norm}(\mu, \sigma)$, $\tau = 1 / \sigma^2 \sim \operatorname{Gamma}(\alpha, \beta)$ then $X$ is a random variable with a Student's t distribution. This one falls under #49.

The closed-form expression for the sum of random variables is known in many cases. This can lead to significant simplifications.

We need to implement more conjugate relationships and posteriors.

Reversible Jump MCMC algorithms are a generalization of the Rosenbluth-Metropolis-Hastings algorithm for when the target density does not have a fixed number of dimensions. Since aesara can handle variables of varying dimensions we should be able to implement such algorithms in aehmc.

As a motivating example we could use the coal-mining disaster datasets and let the number of change points be a Poisson-distributed as in the reference below.

References

https://www2.stat.duke.edu/~scs/Courses/Stat376/Papers/TransdimMCMC/GreenRevJump.1995.pdf

sample_steps, updates, initial_values, parameters = aemcmc.construct_sampler({Y_rv: y}, srng)

should be

sampler, initial_values = aemcmc.construct_sampler({Y_rv: y}, srng)

We could implement the Gibbs sampler described in [1] for the parameters of the Horseshoe+ prior.

[1] Makalic, Enes, and Daniel F. Schmidt. “High-Dimensional Bayesian Regularised Regression with the BayesReg Package.” ArXiv Preprint ArXiv:1611.06649, 2016.

uniform is already available in Aesara, and t will be added in #1211. The change is minimal and consists in adding two facts in transforms.py.

After #45, we need to add a scale-mixture-expansion step to the sampler-generating process.

We can start by implementing the Polya-gamma expansion for negative-binomials and Bernoulli random variables. This will work together to produce nbinom_normal_posterior and bern_normal_posterior in #45, after expansion and subsequent detection of the resulting normal-normal posterior.

On a related note, we'll need to start walking through "graph rewrite space" in a more reasonable way. For instance, a scale-mixture expansion might aid in the construction of a sampler where one wasn't possible before (e.g. we don't have a sampler for the unexpanded term in our "database"). On the other hand, when there is a sampler for a term that can be expanded, we will need to consider the samplers resulting from both expanding and not expanding.

The current docstring contains the following relation:

$$ \begin{equation*} \frac{ X \sim \operatorname{N}\left(\mu_x, \sigma_x^2\right), \quad Y \sim \operatorname{N}\left(\mu_y, \sigma_y^2\right), \quad X + Y = Z }{ Z \sim \operatorname{N}\left(\mu_x + \mu_y, \sigma_x^2 + \sigma_y^2\right) } \end{equation*} $$

when we want:

$$ \begin{equation*} \frac{ Y \sim \operatorname{P}(0, 1), \quad X = \mu + \sigma,Y }{ X \sim \operatorname{P}\left(\mu, \sigma\right) } \end{equation*} $$

Our tests should not allow un-caught/checked warnings.

The addition to setup.cfg mentioned here should be sufficient. We'll need to update our tests to prevent them from failing, though.

The Gibbs sampler for the horseshoe prior does not work on the conditional distributions found in the model, but on the conditional distributions of a transformed version of the model.

In particular, it exploits the fact that a half-Cauchy distribution can be intepreted as a mixture of inverse-gamma distributions, and that if $Z \sim \operatorname{InvGamma}(1, a)$ then $Z \sim 1 / \operatorname{Exp}(a)$.

We should instead implement these relations and update construct_samplers so it finds this structure and builds the sampling steps for each variable in the original model expression.

Related to #46, #65

Currently, horseshoe_nbinom_gibbs fixes a prior for the beta intercept term, but we need all choices like that to be fundamentally configurable, and, eventually, specified exclusively by a given model graph (e.g. once #16 is in place).

Instead, let's make horseshoe_nbinom_gibbs work without an explicit intercept assumption (i.e. only X.dot(beta)). This will broaden the applicability of the sampler and should simplify the situation for #16.

Also, if there's an accuracy/performance cost to this change, that would be a good thing to investigate relative to our model graph-based goals (e.g. how could we reproduce these benefits without compromising the desired interface/functionality).

Copying @brandonwillard's message from #4. If we lift random/measurable variables through mixtures, we can enable some important closed-form posterior opportunities.

For example:

import aesara
import aesara.tensor as at


srng = at.random.RandomStream(4238)

I_rv = srng.bernoulli(0.5, name="I")

Z_1_rv = srng.gamma(10, 100, name="Z_1")
Z_2_rv = srng.gamma(1, 1, name="Z_2")

Z_rv = at.stack([Z_1_rv, Z_2_rv])

# Observation model
Y_rv = srng.poisson(Z_rv[I_rv], name="Y")

Conjugate updates are available between Y_rv and the two Z_*_rv, conditional on the values of I_rv.

The model after lifting should be equivalent to the following:

Z_1_new_rv = srng.poisson(Z_1_rv, name="Z_1_new")
Z_2_new_rv = srng.poisson(Z_2_rv, name="Z_2_new")

# New observation model
Y_new_rv = at.stack([Z_1_new_rv, Z_2_new_rv])
Y_new_rv.name = "Y_new"

The Z_*_new_rv terms are now amenable to the Poisson-gamma conjugate rewrites.

A PR has been created here: conda-forge/staged-recipes#17389.

Examples that come to mind are the Binomial distribution and the Wald distribution (but there are so many others, see the Leemis chart for more examples). We will need to add new RV types for these distributions.

This could be very useful, for instance the Wald distribution in its usual form can be reparametrized to a form that belongs to the scale family, thus enabling new rewrites.

The "X is a special case of Y when ..." fall under this category, the difference being we don't need to introduce new types.

Indeed these are not needed for many models, for instance the following toy example:

import aesara.tensor as at
from aesara.tensor.random import RandomStream

from aemcmc.utils import ModelInfo

srng = RandomStream(0)
a_rv = srng.normal(0, 1, name="a")
b_rv = srng.normal(a_rv, 1, name="b")

a_vv = a_rv.clone()
a_vv.name = "a_vv"

b_at = at.scalar(name="b_at")

model = ModelInfo((b_rv,), {a_rv: a_vv, b_rv: b_at}, (), ())

We need to copy over the kanren relations in symbolic_pymc so that we can sample directly from closed-form posteriors—when possible. See the basic example here.

Within the general process of constructing a sampler for a model graph (see #3), we could apply these rewrites as well, and—in some cases—simply return a graph of the joint posterior, instead of a sampler graph.

We expect the function to return a step function that generates posterior samples for $\tau$ and $\lambda$. However, it calls horseshoe_posterior:

which updates the values of $1/ \lambda$ and $1 / \tau$.

The current version of horseshoe_nbinom_gibbs includes a scan loop to get samples using the kernel defined inside the function. However, to be able to compose this kernel with others we need to be able to build its graph independently from any sampling loop.

Using the sampling steps built by AeMCMC in a scan loop is not straightforward:

import aesara
import aemcmc


sample_steps, sample_updates, initial_values = aemcmc.construct_sampler(
    {Y_rv: y_tt}, srng
)

to_sample_rvs: List[TensorVariable]
inputs = [initial_values[rv] for rv in to_sample_rvs]
outputs = [sample_steps[rv] for rv in to_sample_rvs]

def step_fn(*values):
    from aesara.compile.function.pfunc import rebuild_collect_shared

    vv_to_values = {inputs[i]: val for i, val in enumerate(values)}

    _, new_values, [_, new_updates, _, _] = rebuild_collect_shared(
        outputs, inputs=inputs, replace=vv_to_values, updates=sample_updates
    )

    return new_values, new_updates

n_samples = at.iscalar("n_samples")
outputs, updates = aesara.scan(step_fn, outputs_info=inputs, n_steps=n_samples)

sample_fn = aesara.function(inputs + [n_samples], outputs, updates=updates)

but easily generalizable. We should implement a utility function, e.g. aemcmc.sampling_loop which, given the outputs of construct_sampler and a number of iterations n_samples returns a graph that generate n_samples.

We need to assign the NUTS sampler to the remaining (continuous) variables here in aemcmc.basic.construct_sampler.

aesara-devs/aesara#1054 in Aesara changes the use of "optimize" to "rewrite". We will need to conform to this new API. Since we use TransformValuesOp we will need to wait until aesara-devs/aeppl#160 is closed.

The relation implemented in #29 could not be expanded both ways due to a limitation in aesara. We nevertheless implemented the corresponding test, currently marked as xfail. We should be able to remove the mark when aesara-devs/aesara#1002 is merged.