# Inference issue using a model with Multinomial latent variables on multidimensional observations

**URL:** https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022
**Category:** v5
**Tags:** sampling
**Created:** [October 24, 2024, 5:19pm UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022 "2024-10-24T17:19:31Z")
**Posts on this page:** 5
**Page:** 1

<div class="post-metadata">

### Author: ![Sopchi](https://avatars.discourse-cdn.com/v4/letter/s/e19adc/32.png) [@Sopchi](https://discourse.pymc.io/u/Sopchi)
#### Post date: [October 24, 2024, 5:19pm UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022/1 "2024-10-24T17:19:31Z")

</div>

Let’s say that we have access to N observations of a phenomenon where for each observation i we observe a defined number m of independant draws of a Negative Binomial whose paremeters are \alpha\_i (Gamma distribution shape parameter) and \mu\_i (Gamma distribution mean).

The \alpha parameter is shared across every observations thus \forall i, \alpha\_i = \alpha

However, \mu\_i depends on a latent variable z\_i that assigns every observation i to a hidden state among k possible states.

The latent variables z can be seen as indicator variables and can be modeled using a Multinomial \mathcal{Multinomial}(n=1,p\_1,...,p\_k) where p\_1,...,p\_k the probabilities of each one of the different outcomes. In this case z\_i is a binary vector pf size k with one 1 on the state where observation i belongs.

Alternatively, the latent variables z can also be modeled using a Categorical distribution with probabilities p\_1,...,p\_k.

Additionnally, negative binomial mean of each state 1,...,k is known

For example, let’s imagine that observation i belongs to the second state, then z\_i =2, and this observation is a vector of m independant draws of a Negative Binomial with parameters \alpha and \mu\_j = mean of state 2

Given a bunch of observation following this model, we want to infer \alpha, the latent variable of each observation z and the probabilities of each state p\_1,...,p\_k.

## Observations simulation

```auto
N, M, max_state = 10, 5, 8
prob = np.random.rand(max_state)
state_prob = prob/sum(prob) #probabilities of each state
expected_mean = np.concatenate((np.array([1]),np.arange(1,max_state)*10)) #mean of each state
alpha_gt = 0.0001 # shape parameter

```

```auto
with pm.Model() as ind_simulation:
    w = state_prob 
    latent_z = pm.Multinomial('z',n=1, p=w, shape=(N,max_state))
    state_assignment = latent_z.eval() #For debuging purposes we do an evaluation of the latent variabl

    alpha = alpha_gt 
    mu = np.repeat(pm.math.matmul(state_assignment,expected_mean.reshape(-1,1))[:,0][:, np.newaxis], M, axis=1).eval()

    obs_data = pm.NegativeBinomial('ind',alpha=1/alpha, mu=mu) 
    obs_draw = pm.draw(obs_data)

```

Here is a heatmap of the simulated data, which looks as expected in terms of size, and values. Each value on one line seems to be generated from the same negative binomial distribution.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/e/eef4d4185fc7557c6bb6e34e8f3999d564a5bfee.png)

## Sampling

We can now perform the inference

```auto
with pm.Model() as inference_model:
    w = pm.Dirichlet("w", a=np.ones(max_state)) 
    latent_z = pm.Multinomial('z',n=1, p=w,shape=(N,max_state))

    alpha = pm.Uniform("alpha", lower=0, upper=0.1) 
    mu = np.repeat(pm.math.matmul(latent_z,expected_mean.reshape(-1,1))[:,0][:, np.newaxis], M, axis=1)
    obs_distrib = pm.NegativeBinomial('obs',alpha=1/alpha, mu=mu, observed=obs_draw)

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

```

We give the \mu parameters of the Negative Binomial in the most explicit way in the form of a matrix that matches the shape of the observation (N,M) to avoid any ambiguity.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/f/f7ebbd8082224eba4bbc3a5db8111abd06f70d58.png)

I would have expected to find the z and the observation in a big shared plate of size N and the observation nested in a smaller plate of size M.

## Inference results

We can first have a look at the trace

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/b/b266397955d2625c1857af927848edfcc104bbba.jpeg)  
Then, if we look to the posteriors for the latent variables we observe that all the observations were weirdly assigned to the first state  
 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/b/b9c4a521679be29b7ed5c611b50d8c65fe058f90.png)  
The state probabilities are also off, with the first state having the higher probability  
 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/9/9f85441ed64bbbb78f722fd15cab8c1424625a2b.png)  
The shape parameter is also reaching the upper bound of the prior, allowing wide distributions

## With categorical distribution

I then tried to model the latent variables using categorical distribution resulting in a new inference model

```auto
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)

```

The model is again not displaying nested plates

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/d/d60ac0a8c8f0ddc39a593593619826b34fdb182b.png)

However, now the inference works fine

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/9/9dd08e237570374ce70030c1de857b0649b44edd.png)  
 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/3/3b57f845f3eb59c1f577d5bd5011a53efad4cfc4.jpeg)

Can someone please help me to understand what is happening here ?

---

<div class="post-metadata">

### Author: ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)
#### Post date: [October 25, 2024, 3:10am UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022/2 "2024-10-25T03:10:11Z")

</div>

PyMC doesn’t have a step sampler that works for Multinomial distribution whereas you’ll notice there’s a specialized sampler for the Categorical distribution.

No clue about the nesting of the labels. There’s an heuristic used for it. Maybe if you add dims it will figure it out?

---

<div class="post-metadata">

### Author: ![Sopchi](https://avatars.discourse-cdn.com/v4/letter/s/e19adc/32.png) [@Sopchi](https://discourse.pymc.io/u/Sopchi)
#### Post date: [October 25, 2024, 12:01pm UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022/3 "2024-10-25T12:01:34Z")

</div>

Thank you very much for your help  
This is the outputs I get when using the categorical model and the multinomial model respectively

### Categorical

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/8/8e6db068b5ae79b4269db59652ad413a43b1be37.png)

### Multinomial

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pymc3/original/2X/9/98c7a6a1e41fd9e45562180e7e443f328a108f4b.png)

So if I understand correctly a sampler is chosen to sample from the multinomial but it does not work ? I am not sure I get the issue here. So, using pymc you can only draw from multinomial but not sample from it ? If that is the case then shouldn’t I receive a warning when trying to sample from such a model ?

Also, I can I know from which distribution supported by pymc I can sample from ?

For the nested plates I guess it’s just a displaying issue as it works fine with the model using categorical distributions. Fyi I tried with dims too and still it does not display nested plates.

Also, I wonder if this is the best, most efficient way to index a numpy array storing my negative binomial means of each states (exepected\_mean) with the categorical pytensor object ? :

```auto
latent_z = pm.Categorical('z', p=w,shape=N)
mu = pt.shared(expected_mean)[latent_z]
pm.NegativeBinomial('obs',alpha=1/alpha, mu=mu, observed=obs_draw.T)

```

Thank you for your help

---

<div class="post-metadata">

### Author: ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)
#### Post date: [October 25, 2024, 12:26pm UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022/4 "2024-10-25T12:26:35Z")

</div>

> [@Sopchi](#):
>
> Also, I can I know from which distribution supported by pymc I can sample from ?

In general you can sample any distribution, but sometimes the sampler will fail. This cannot be known in advance, and you will notice it by checking that `r_hat` is very high across chains or the samples seem “stuck”.

Furthermore, non-NUTS samples (which discrete variables must necessarily use) can fail to sample more quietly than the NUTS sampler (which continuous variables are automatically assigned to if we can get gradients).

The only case I know fails all the time is `Multinomial` because it’s the only discrete multivariate variable we have, and it has a very specific constraint (all values must add up to `n`). It would need a sampler that is aware of this fact. We could perhaps raise an error.

Note you can still use Multinomial just fine as a likelihood.

---

<div class="post-metadata">

### Author: ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)
#### Post date: [October 25, 2024, 12:28pm UTC](https://discourse.pymc.io/t/inference-issue-using-a-model-with-multinomial-latent-variables-on-multidimensional-observations/16022/5 "2024-10-25T12:28:27Z")

</div>

I opened an issue here: [Warn or raise when trying to sample a Multinomial variable · Issue #7548 · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/issues/7548)
