BUG: Truncated gamma distribution has incorrect logp

[tvwenger]

Describe the issue:

A Gamma distribution that is wrapped in Truncated has the wrong logp when beta is small. For beta = 1.0 the results are correct, but for beta = 0.01 the results are incorrect. The code below produces the following distributions.

plot_gamma_prob(1.0, 5.0) # beta = 1.0, truncated at 5.0

plot_gamma_prob(0.01, 1000.0) # beta = 0.01, truncated at 1000.0

Reproduceable code example:

import numpy as np
import matplotlib.pyplot as plt
import pymc as pm
print(pm.__version__)

def plot_gamma_prob(beta, x_max):
    x = np.linspace(0.0, 2.0 * x_max, 1000)

    gamma = pm.Gamma.dist(alpha=3.0, beta=beta)
    p_gamma = np.exp(pm.logp(gamma, x).eval())

    trunc_gamma = pm.Truncated.dist(pm.Gamma.dist(alpha=3.0, beta=beta), upper=x_max)
    p_trunc_gamma = np.exp(pm.logp(trunc_gamma, x).eval())

    fig, ax = plt.subplots()
    ax.plot(x, p_gamma, "k-", label="Gamma")
    ax.plot(x, p_trunc_gamma, "r-", label="Truncated Gamma")
    ax.set_xlabel("x")
    ax.set_ylabel("P(x)")
    ax.legend(loc="best")
    fig.tight_layout()
    fig.savefig(f"gamma_beta{beta}.png")
    plt.close(fig)

plot_gamma_prob(1.0, 5.0)
plot_gamma_prob(0.01, 1000.0)

Error message:

No response

PyMC version information:

pymc v. 5.8.2

Context for the issue:

This seems like an important bug!

[larryshamalama]

This does seem important, I can have a cursory look. Thanks for bringing this up!

CC @ricardoV94

[ricardoV94]

I would guess float precision issues on the cdf

[tvwenger]

@ricardoV94 Seems that the CDF is doing ok:

from scipy.stats import gamma
import pymc as pm

alpha = 3.0
beta = 0.01

x = np.linspace(0.0, 2000.0, 1000)
gamma_scipy = gamma(a=alpha, scale=1.0/beta)
logpdf_scipy = gamma_scipy.logpdf(x)
logcdf_scipy = gamma_scipy.logcdf(x)

gamma_pymc = pm.Gamma.dist(alpha=alpha, beta=beta)
logpdf_pymc = pm.logp(gamma_pymc, x).eval()
logcdf_pymc = pm.logcdf(gamma_pymc, x).eval()

assert np.allclose(logpdf_scipy, logpdf_pymc)
assert np.allclose(logcdf_scipy, logcdf_pymc)

[tvwenger]

Here's logp and logcdf:

gamma_trunc_pymc = pm.Truncated.dist(
    pm.Gamma.dist(alpha=alpha, beta=beta),
    lower=0.0,
    upper=500.0)
logpdf_pymc_trunc = pm.logp(gamma_trunc_pymc, x).eval()
logcdf_pymc_trunc = pm.logcdf(gamma_trunc_pymc, x).eval()

fig, ax = plt.subplots()
ax.plot(x, logpdf_pymc, 'k-', label='Gamma')
ax.plot(x, logpdf_pymc_trunc, 'r-', label='Truncated Gamma')
ax.set_xlabel('x')
ax.set_ylabel('p(x)')
fig.tight_layout()
fig.savefig('gamma_logpdf.png')
plt.close(fig)

fig, ax = plt.subplots()
ax.plot(x, logcdf_pymc, 'k-', label='Gamma')
ax.plot(x, logcdf_pymc_trunc, 'r-', label='Truncated Gamma')
ax.set_xlabel('x')
ax.set_ylabel('F(x)')
fig.tight_layout()
fig.savefig('gamma_logcdf.png')
plt.close(fig)

[tvwenger]

The plot thickens! The truncated gamma distribution logp is wrong for all beta != 1.0. The logic in Truncated is right and reproduced below.

from scipy.stats import gamma
import pymc as pm

alpha = 3.0
beta = 3.0
upper = 2.5

x = np.linspace(0.0, 5.0, 1000)
gamma_scipy = gamma(a=alpha, scale=1.0/beta)
logpdf_scipy = gamma_scipy.logpdf(x)
logcdf_scipy = gamma_scipy.logcdf(x)

gamma_pymc = pm.Gamma.dist(alpha=alpha, beta=beta)
logpdf_pymc = pm.logp(gamma_pymc, x).eval()
logcdf_pymc = pm.logcdf(gamma_pymc, x).eval()

assert np.allclose(logpdf_scipy, logpdf_pymc)
assert np.allclose(logcdf_scipy, logcdf_pymc)

gamma_trunc_pymc = pm.Truncated.dist(
    gamma_pymc,
    upper=upper)
logpdf_pymc_trunc = pm.logp(gamma_trunc_pymc, x).eval()
logcdf_pymc_trunc = pm.logcdf(gamma_trunc_pymc, x).eval()

test_logp = pm.logp(gamma_pymc, x).eval()
upper_logcdf = pm.logcdf(gamma_pymc, upper).eval()
test_logp = test_logp - upper_logcdf
test_logp[x > upper] = -np.inf

fig, ax = plt.subplots()
ax.plot(x, np.exp(logpdf_pymc), 'k-', label='Gamma')
ax.plot(x, np.exp(logpdf_pymc_trunc), 'r-', label='Truncated Gamma')
ax.plot(x, np.exp(test_logp), 'b-', label='Correct')
ax.set_xlabel('x')
ax.set_ylabel('p(x)')
ax.legend(loc='best')
fig.tight_layout()
fig.savefig('gamma_pdf.png')
plt.close(fig)

I suspect that the problem has to do with the fact that the pytensor gamma RV stores the inverse of beta, whereas pymc parameterizes gamma in terms of beta.

[tvwenger]

My intuition was correct. After wrapping Gamma in Truncated, beta gets mapped to 1/beta. That is, Truncated(Gamma(alpha, 1/beta)) == Actual Truncated Gamma with (alpha, beta). See:

from scipy.stats import gamma
import pymc as pm

alpha = 3.0
beta = 3.0
upper = 2.5

x = np.linspace(0.0, 5.0, 1000)

gamma_pymc = pm.Gamma.dist(alpha=alpha, beta=beta)
logpdf_pymc = pm.logp(gamma_pymc, x).eval()

gamma_trunc_pymc = pm.Truncated.dist(
    pm.Gamma.dist(alpha=alpha, beta=beta),
    upper=upper)
logpdf_pymc_trunc = pm.logp(gamma_trunc_pymc, x).eval()

corrected_gamma_trunc_pymc = pm.Truncated.dist(
    pm.Gamma.dist(alpha=alpha, beta=1/beta),
    upper=upper)
corrected_logpdf_pymc_trunc = pm.logp(corrected_gamma_trunc_pymc, x).eval()

fig, ax = plt.subplots()
ax.plot(x, np.exp(logpdf_pymc), 'k-', label='Gamma')
ax.plot(x, np.exp(logpdf_pymc_trunc), 'r-', label='Truncated Gamma')
ax.plot(x, np.exp(corrected_logpdf_pymc_trunc), 'b-', label='Correct')
ax.set_xlabel('x')
ax.set_ylabel('p(x)')
ax.legend(loc='best')
fig.tight_layout()
fig.savefig('gamma_pdf.png')
plt.close(fig)

The easiest solution (to me) would be to update pytensor to use the same parameterization as pymc, or to create a new GammaRV in pymc that uses the pymc parameterization.

[ricardoV94]

That sounds like a bug in Truncated, the fact we flip the parameter of beta shouldn't matter because we only do that once. an unfortunate behavior in PyTensor

[ricardoV94]

Great investigation @tvwenger

The problem happens in this line:

pymc/pymc/distributions/shape_utils.py

Line 377 in f77372c

new_rv = rv_node.op(*dist_params, size=new_size, dtype=dtype)

Gamma.__call__ flips the meaning of beta every-time it is called: https://github.com/pymc-devs/pytensor/blob/a708732dfe507ae892fdd6c783d7a1bc15ebc9e2/pytensor/tensor/random/basic.py#L467

A safer solution would be to do rv_node.op.make_node, but I think we should just get rid of that behavior of Gamma in PyTensor. It's a source of confusion for no good reason.

To not break things we can stop flipping the sign in the __call__ and instead flip it in the random method. We can also switch from scipy to numpy gamma while we are at it. Once we update the PyTensor dependency in PyMC we have to be aware that the parameter is no longer flipped "automatically". Tests will probably fail so I don't think we will miss it.

[ricardoV94]

@tvwenger are you interested in doing the suggested fix in PyTensor instead?

[tvwenger]

@ricardoV94 Here they are:

pymc-devs/pytensor#460
#6934