Speed up posterior predictive sampling

[brandonwillard]

Posterior predictive sampling takes considerably longer than model estimation itself, and that's completely unreasonable!

Last time I tried, the new fast_sample_posterior_predictive function wouldn't work with our custom distribution classes. If I recall, it looked like fast_sample_posterior_predictive was attempting to sample from a non-Distribution object, which would raise an exception, of course.

We should either attempt to get fast_sample_posterior_predictive working, or simply write our own posterior sampling functionality.

The latter case might be a lot simpler than it seems, since we could convert the PyMC3 models into native Theano graphs using symbolic-pymc, then compile those graphs into functions that sample directly from the given models. (In general, we could start doing most of our work in Theano using symbolic-pymc's random variables and interface with PyMC3 when we want to run samplers.)

[brandonwillard]

FYI: I've added a lot of relevant functionality in pymc-devs/symbolic-pymc#113 and pymc-devs/symbolic-pymc#114.

Part of the functionality introduced there allows us to automatically transform Theano graphs with random variable outputs from a Scan into inputs to the same Scan. This transform is currently being used to produce log-likelihood graphs, since log-likelihoods take the values of random variables as inputs, but it's also what we need in order to construct a fast Theano function that produces posterior predictive samples. Writing a posterior predictive sampler that applies this transform and produces the desired function is the next step.

However, for us to make use of this new posterior predictive sampler, we would need to create a separate Theano model graph for each PyMC3 model. This approach involves tedious and error-prone duplication and would make the modeling process much more tedious. Alternatively, we can make exclusive use of Theano model graphs to specify our models and automatically convert them to PyMC3 Model objects when needed (e.g. when model terms need to be estimated via NUTS using pm.sample).

The existing symbolic_pymc.theano.pymc3.graph_model function does this for simple Theano model graphs, but it does not work with the Theano Scan operator used to define non-trivial time-series/HMM model graphs.

In order to represent Scans in a PyMC3 model object, we need to automatically construct a PyMC3 Distribution class for the Scans. The core requirement for that involves the construction of a valid Distribution.logp function, which can now be done using the features introduced in the aforementioned PRs!

[brandonwillard]

Since the recent improvements and fixes to the shape handling in our custom Distributions, we've been able to dramatically increase the speed of posterior predictive sampling by constructing the posterior predictive distributions by hand.

For example, the following is a handwritten sampler for a two-state negative-binomial HMM with time-varying transition matrices:

def sample_posterior_predictive(trace, X=None, X_xi=None, chain=0, var_names=None):
    ppc_trace = {}

    X_xi = X if X_xi is None else X_xi

    if hasattr(trace, "xi") and X_xi is not None:
        xi_samples = trace.xi[chain].values
        z_samples = np.tensordot(X_xi, xi_samples, axes=((1,), (1,)))
        z_samples = np.swapaxes(z_samples, 0, 1)
        Gamma_samples = multilogit_inv(z_samples)
    else:
        Gamma_samples = trace.Gamma[chain].values

    if var_names is None or 'Gamma' in var_names:
        ppc_trace['Gamma'] = Gamma_samples

    if X is not None:
        V_shape = trace.V_t[chain].shape[-1]

        v_dist = HMMStateSeq.dist(Gamma_samples,
                                  trace.gamma_0[chain].values,
                                  shape=V_shape)

        V_samples = v_dist.random()
        p_t_samples = np.exp(1.0 + trace.beta[chain].values.dot(X.T))

        ppc_trace['V_t'] = V_samples
    else:
        V_samples = trace.V_t[chain]
        p_t_samples = trace.p_t[chain]

    nb_comp = pm.NegativeBinomial.dist(p_t_samples,
                                       1.0 / trace.alpha.values[chain][..., None],
                                       shape=p_t_samples.shape)
    y_dist = SwitchingProcess.dist(
        [
            pm.Constant.dist(0), nb_comp
        ],
        V_samples.astype(np.int32))

    samples = y_dist.random()

    ppc_trace['Y_t'] = samples

    return ppc_trace

While this dramatically improves the speed of sampling, it requires a new sampler function for every new model (or relevant update to an existing model). This slows down the modeling process and is very error prone, so we (still) need to do something about it.

[rlouf]

Will posterior sampling speed still be an issue with PyMC3 version 4?

[brandonwillard]

Will posterior [predictive] sampling speed still be an issue with PyMC3 version 4?

No, not at all.

[rlouf]

That's what I thought. Shall we close this issue then?