Shape Parameter for Multiple Treatment Groups

Questions
6

Got it working thanks for your help lucianopaz! So it seems like with pm.switch it’s easier to better to pass data in a wide format, not a tall format (tall formats work fine on other distributions using numpy fancy indexing).

Curious to know the best practice if you have different numbers of observations for each treatment group (NANs give the warning: Data in likelihood contains missing values and will be automatically imputed from the sampling distribution.)?

Final code if it will help others:

fake_count_data = pd.DataFrame(np.random.poisson(lam = 15, size = (74, 3)))
fake_count_data.columns = ['p0', 'p1', 'p2']

with pm.Model() as model:

    n_persons = fake_count_data.shape[1]
    n_count_data = len(fake_count_data)
    
    alpha = 1.0/fake_count_data.mean().mean()
    
    lambda_1 = pm.Exponential("lambda_1", alpha, shape=n_persons)
    lambda_2 = pm.Exponential("lambda_2", alpha, shape=n_persons)
    tau = pm.DiscreteUniform("tau", lower=0, upper=n_count_data - 1, shape=n_persons)

    idx = np.arange(n_count_data)[:, None] # Thanks lucianopaz!
    lambda_ = pm.math.switch(tau > idx, lambda_1, lambda_2)

    observation = pm.Poisson("obs", mu=lambda_, observed=fake_count_data)