Bayesian bootstrap for beta geometric model

Hey,

I am trying to model product retention based on Fader’s shifted beta geometric paper. That model is working well. I have a working version both in numpyro and pryo. I want to add a bootstrap, described here, as a way to model the variation in the retention curve. In both the pryo and numpryo version the posterior collapse to a single point. Any help would be great appreciated.

Thanks,
Ryan

import jax.numpy as np
import numpy as onp

class sBG(numpyro.distributions.Distribution):
    def __init__(self, alpha, beta):
        self.alpha = alpha
        self.beta = beta
        
    def sample(self):
        return None
    
    def log_prob(self, data):
        active = data[0,:]
        lost = data[1,:]
        
        n = active.shape[0]
        p = [0., self.alpha / (self.alpha + self.beta)]
        s = [0., 1 - p[-1]]
        for t in range(2, n):
            pt = ((self.beta + t - 2) / (self.alpha + self.beta + t - 1)) * p[-1]
            p.append(pt)
            s.append(s[-1] - p[-1])

        p = np.stack(p)
        s = np.stack(s)
        
        died = np.multiply(np.log(p[1:]), lost[1:])
        still_active = np.log(s[-1]) * active[-1]
        logp = np.sum(died) + still_active
        return logp 

def sBG_bs(data):

    n = data.shape[0]
    length_timeperiod = data.shape[1]
    ki = onp.random.randint(0,n, size = n)

    raw_data_bs = data[ki, :]
    active = raw_data_bs.sum(axis=0)
  
    lost = onp.zeros_like(active)
    lost[1:] = onp.flip(onp.diff(onp.flip(active, axis=0)), axis=0)
    
    data_bs = np.array([active, lost])

    alpha = numpyro.sample('alpha', dist.Uniform(0,5))
    beta  = numpyro.sample('beta' , dist.Gamma(2,2) )
    numpyro.sample('p', sBG(alpha, beta), obs=data_bs)

nuts_kernel = numpyro.infer.NUTS(sBG_bs)
mcmc= numpyro.infer.MCMC(nuts_kernel, num_warmup=1000, num_samples=1000)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, data)
1 Like