Very long sampling time using simple model

Questions v5
1

I want to sample from this model

image
image470×724 18.2 KB

To do that I have created this pymc model :

with pm.Model() as inference_model_categorical:
    w = pm.Dirichlet("w", a=np.ones(max_state)) 
    latent_z = pm.Categorical('z', p=w,shape=N)

    alpha = pm.Uniform("alpha", lower=0, upper=0.1) 
    mu = pt.shared(expected_mean)[latent_z]

    obs_distrib = pm.NegativeBinomial('obs',alpha=1/alpha, mu=mu, observed=obs_draw.T)

    obs_sample = pm.sample() 
pm.model_to_graphviz(model=inference_model_categorical)

However, even when my observation dataset is relatively small (50x100) it takes 3 minutes. Obviously, it does not scale when the dataset increases in size, e.g for a 100x500 dataset it takes more than 1h30.

I wonder what is the slow bit here? Is it sampling from Negative Binomial (I have also previously designed a mixture of 10 Negative Binomial distributions to solve the same problem and it was equally slow).

Or am I doing something absurd/completely inefficient ?

Thank you for your help

2

Maybe marginalizing the discrete latent_z will help the sampler. You can try to use MarginalModel to do that automatically.

It also allows you to recover latent_z later with recover_marginals so you don’t lose anything.

https://www.pymc.io/projects/experimental/en/latest/generated/pymc_experimental.MarginalModel.html