How to implement Beta Function effectively in Pymc3

That must’ve been it. I’ve made some adjustments such as to the priors so that they’re more easily digested for for HMC/NUTS.

import numpy as np
import pymc3 as pm
import arviz as az
import theano.tensor as tt
from math import gamma as gamma_f
from theano.compile.ops import as_op
from pymc3.distributions.dist_math import betaln

survived = np.array([869, 743, 653, 593, 551, 517, 491])
churned = np.array([ 131, 126, 90,  60,  42,  34,  26])
n_obs = len(survived)

with pm.Model() as sbg_model:
    alpha = pm.HalfNormal('alpha', sigma=1) # MLE est was 0.668
    beta = pm.HalfNormal('beta', sigma=1) # MLE est was 3.806
    churn = pm.Beta('churn', alpha, beta)
    renewal = pm.Deterministic('renewal', 1-churn)
    
    def logp(survived, churned, alpha = alpha, beta = beta):   
        # Calculate the final surivival probability (eq. 6)
        ln_surv_prob = betaln(alpha, beta + n_obs) - betaln(alpha, beta)
        
        # Find the probability of churn for all values prior
        ln_prob_vec = []
        for i in range(1, n_obs + 1):
            ln_prob_vec.append(betaln(alpha + 1, beta + i - 1) - betaln(alpha, beta))
        ln_prob_vec = tt.as_tensor_variable(ln_prob_vec, 'ln_prob_vec', ndim = 1)

        return pm.math.dot(churned, ln_prob_vec) + survived[-1] * ln_surv_prob   
    
    likelihood = pm.DensityDist('likelihood', logp, observed = {'survived': survived, 'churned': churned})
    sbg_trace = pm.sample(chains=4) 
2 Likes