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.
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 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 length_timeperiod = data.shape 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)