replace-dirichlet-gamma

gistfile1.txt

import numpy as np
import pymc3 as pm
from pymc3.distributions.transforms import StickBreaking

genus_counts = np.random.multinomial(10, np.ones(5)/5, 20)

with pm.Model() as model:
    k = 3
    n, p = genus_counts.shape
    profile_gamma = pm.Gamma('profile_gamma', alpha=np.ones((k, p)), beta=np.ones((k, p)), shape=(k,p))
    weights_gamma = pm.Gamma('weights_gamma', alpha=np.ones((n, k)), beta=np.ones((n, k)), shape=(n, k))
    profiles = pm.Deterministic("profiles", profile_gamma/pm.math.sum(profile_gamma, axis=-1, keepdims=True))
    weights  = pm.Deterministic("weights", weights_gamma/pm.math.sum(weights_gamma,axis=-1, keepdims=True))
    apparent_abundance =  pm.Deterministic("apparent_abundance", pm.math.dot(weights,profiles))
    overdispersion = pm.Exponential("overdispersion", 1)
    read_counts = pm.NegativeBinomial("read_counts", genus_counts.sum(axis=1)[:,None]*apparent_abundance, 1/overdispersion, shape=(n,p), 
                                      observed=genus_counts)

    trace = pm.sample_prior_predictive()
    
    _ = [print(var, trace[var].shape) for var in trace.keys()]