Truncation errors (falling back to rejection sampling despite logcdf present)

1

I’ve been working with a truncated BetaBinomial distribution, and while trying to simulate data (via pm.draw() or pm.sample_prior_predictive), I often encounter Truncation Errors (TruncationError: Truncation did not converge in 10000 steps). From what I see in the documentation and the code behind pm.Truncated, this error should only arise if rejection sampling is used, and it should only be used when the truncated distribution does not implement a logcdf function. But pm.BetaBinomial has logcdf method, and so I am not sure why this is happening. Am I doing something wrong and are there ways to avoid / minimize the risk of such errors?

Sample code:

from pymc.distributions.mixture import _hurdle_mixture
class HurdleBetaBinomial:

    def __new__(cls, name, psi, n, alpha, beta, **kwargs):
        return _hurdle_mixture(
            name=name, nonzero_p=psi, nonzero_dist=pm.BetaBinomial.dist(n=n, alpha=alpha, beta=beta), dtype="int", **kwargs
        )


    @classmethod
    def dist(cls, psi, n, alpha, beta, **kwargs):
        return _hurdle_mixture(
            name=None, nonzero_p=psi, nonzero_dist=pm.BetaBinomial.dist(n=n, alpha=alpha, beta=beta), dtype="int", **kwargs
        )

N = 1
with pm.Model() as m: 
    treatment = pm.Bernoulli("treatment", p=0.5, shape=N)
    baseline = pm.Normal("inv_p", 0.2, 1, shape=N)        
    β = pm.Gamma("β", alpha=1, beta=4, shape=N)
    α = pm.Gamma("α", alpha=1, beta=2, shape=N)   

    ψ_treatment = pm.Normal("inv_ψ_effect", 0, 0.2)
    α_treatment = pm.Normal("α_effect", 0, 0.02)
    β_treatment = pm.Normal("β_effect", 0, 0.02)

    ψ = pm.Deterministic("ψ", pm.invlogit(baseline + ψ_treatment * treatment))
    
    days = HurdleBetaBinomial(
        "days", 
        n=7, 
        psi=ψ, 
        alpha=α + α_treatment * treatment, 
        beta=β + β_treatment * treatment, 
        shape=N, dims="tracks"
    )    
sim_model = pm.do(m, {
    "inv_ψ_effect": 0,
    "α_effect": 0.1,
    "β_effect": 0
})

#try 10 times
for i in range(10):
    with sim_model:
        simulate = pm.sample_prior_predictive(samples=10_000)
2

Logcdf isn’t enough, icdf is also needed to avoid rejection sampling. According to this issue, that method is still missing for BetaBinomial: Add icdf functions for distributions · Issue #6612 · pymc-devs/pymc · GitHub

3

In the meantime you can try to increase the number of steps for convergence, which is called max_n_steps

1 Like
4

Got it, thank you!

5

How exactly should this “max_n_steps” be specified?
I got the same TruncationError using HurdleNegativeBinomial. It didn’t work if I just passed it to HurdleNegativeBinomial.
It seems only pm.Truncated explicitly has ‘max_n_steps’ in it, and so when I followed the code of _hurdle_mixture and tried the following which worked, I didn’t feel surprised. But I presume there must be an easier way to pass the ‘max_n_steps’ parameter?
Btw, on one hand, HurdleNB provides more accurate estimation in one specific application I am working with, but painfully it comes with the slowness issue as in " Optimization Warning: The Op betainc does not provide a C implementation. As well as being potentially slow, this also disables loop fusion". Is there any thing I do about it?

nonzero_p = .1
with pm.Model():
    nonzero_p = pt.as_tensor_variable(nonzero_p)
    weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1)
    comp_dists = [
        pm.DiracDelta.dist(0),
        pm.Truncated.dist(pm.NegativeBinomial.dist(p=1-6*1e-5, n=4000), lower=1, max_n_steps=10000),
    ]
    pm.Mixture('ads', weights, comp_dists)
    prior = pm.sample_prior_predictive(samples=1000)
6

The warning can be ignored, it’s not always true that will result in a slow down. Regarding max_n_steps we should allow passing it easily as an optional kwarg. Can you open an issue on our repository for this (and feel free to open a Pull request to fix it as well): GitHub - pymc-devs/pymc: Bayesian Modeling and Probabilistic Programming in Python

Non-code wise, will the extra n_steps actually help you? It’s very inefficient already with this many steps.

7

Sure, I will organize the example into an issue later :slight_smile: .
Just tried max_n_steps=100,000 instead of 10,000 and it did finish without raising the TruncationError, though much much slower as expected.
I am still trying to determine whether it is because my model is illy built or the low convergence is just what it is given how extreme the parameters are in my case (extreme but agreeing with real data): non-zero percentage (psi) is low(<5%) and probability for non-zeros is somewhat also low (NegativeBinomial(n=4000, p=0.9999x)).
Real data of size 200 looks something like [190 zeros + 51 + 42 + 1*3], so to model the non-zeros as accurate as possible, I chose Hurdle model (zero-inflated models can model the zero part ok but the non-zero parts not as accurate as hurdle models).
Here is how the model looks like:

    with pm.Model() as m:
        m.add_coord('sample', ...)
        m.add_coord('k4', ...)
        TD_obs = pm.MutableData('TD_obs', ..., 'k4'))
        AD_obs = pm.MutableData('AD_obs', ..., dims=('sample', 'k4'))
        zp_obs = pm.MutableData('zp_obs', ..., dims='k4')

        mu_k4 = pm.TruncatedNormal('mu_k4', mu=6, sigma=3, lower=2, upper=12, dims='k4')
        AD_predicted = pm.HurdleNegativeBinomial(
            'AD_predicted',
            psi=1 - zp_obs,
            n=TD_obs,
            p=1 - 1e-5 * mu_k4,
            observed=AD_obs)
        data = pm.sampling.jax.sample_numpyro_nuts(**inferenceParams)

I’m reading into the codes to learn how “convergence” is determined in such Truncated functions.