Replacing nested for loops from JAGS code for cognitive model of temporal discounting

It seems to be caused by underflow in the Phi function which leads to zero when (V_B - V_A) / alpha are very negative. In fact anything smaller than -9 underflows to zero:

Phi(np.array([-13.0, -11.0, -9.0, -7.0])).eval()
array([0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.27980959e-12])

One solution is to work in the logit scale:

...
# Import the helper normal log cdf so that we can get the log of `Phi`
from pymc3.distributions.dist_math import normal_lcdf

    #P = pm.Deterministic('P', Phi((V_B - V_A)/alpha[:, None]))

    # logit(p) = log(p) - log(1-p)
    log_P = normal_lcdf(0, 1, (V_B - V_A) / alpha[:, None])  # log(p)
    log_1m_P = pm.math.log1mexp(-log_P)  # log(1-p)
    logit_P = pm.Deterministic('logit_P', log_P - log_1m_P)
    # Observed
    h = pm.Bernoulli('h', logit_p=logit_P, observed = R)


Discounting.check_test_point()

K_mu                         -3.22
K_sigma_interval__           -1.39
Alpha_mu_interval__          -1.39
Alpha_sigma_interval__       -1.39
k_lowerbound__              -38.23
alpha_lowerbound__          -30.23
h                        -68137.95  # This is no longer -inf!
Name: Log-probability of test_point, dtype: float64

If this works, please make sure the values make sense and that I did not introduce any mistake with my changes. In particular if you wrap the logit_P in new_P = pm.Deterministic('new_P', pm.math.invlogit(logit_p)) you should see that new_P and your original P (if you uncomment it) give the same results, except for the cases where your original P underflowed to zero.

3 Likes