How to calculate the mixed normal distribution multiple times simultaneously

Dear pyro experts

I am not sure how to implement this.
I am trying to estimate with a 1D mixed normal distribution.
Following your tutorial, I was able to build the model as follows.

Def model(data):
    Weights = pyro.sample(‘weights’, dust.Dirichlet(0.5*torch.ones(K)))
    With pyro.plate(‘components’, M):
        Locs = pyro.sample(‘locs’, dust.Normal(0., 10.))
        Scale = pyro.sample(‘scale’, dust.LogNormal(0., 2.))
    With pyro.plate(‘data’, N):
        Assignment = pyro.sample(‘assignment’, dust.Categorical(Weights))
        Pyro.sample(‘obs’, dust.Normal(locs[assignment], scale[assignment]), obs=data)

However, there are K data in one dimension, and running this code K times will take a long time.
Since the number of each data is the same N, we can consider it as N data in K dimensions.
So, I want to extend the above implementation to K dimensions, but it does not work. How should I implement this?

Def model(data):
    With pyro.plate(‘dim’, K):
        Weights = pyro.sample(‘weights’, dust.Dirichlet(0.5*torch.ones(K)))
        With pyro.plate(‘components’, M):
            Locs = pyro.sample(‘locs’, dust.Normal(0., 10.))
            Scale = pyro.sample(‘scale’, dust.LogNormal(0., 2.))
        With pyro.plate(‘data’, N):
            Assignment = pyro.sample(‘assignment’, dust.Categorical(Weights))
            Pyro.sample(‘obs’, dust.Normal(locs[assignment], scale[assignment]), obs=data)

Also, is it faster to do MCMC at the same time?

Furthermore, given that there is no correlation between the two, I think it can be implemented as a multivariate mixed normal distribution, but which is the correct implementation?

Hi, I’m a little confused by your example code, since the dimensionality of Weights and hence the number of components is specified by K rather than M. What is K in your notation?

To clarify, do you mean that data now has shape (N, K) rather than shape (N,) and you want Locs and Scale to have shape (M, K) rather than (M,) where M is the number of components, but with each entry drawn from the same 1-dimensional prior?

If so, as discussed in the “Distribution Shape” section of the tensor shape tutorial, you should remove the dim plate and instead .expand the distributions at those sites and declare the final dimension dependent with .to_event:

def model(data):
    # M components
    weights = pyro.sample('weights', dist.Dirichlet(0.5*torch.ones(M)))
    with pyro.plate('components', M):
        locs = pyro.sample('locs', dist.Normal(0, 10).expand([K]).to_event(1))
    # ... and similarly for Scale and obs

See the tensor shape and enumeration tutorials for more background.

If this isn’t what you want, please provide some math or pseudocode for the generative model you have in mind.

Also, is it faster to do MCMC at the same time?

Faster than what?

Dear eb8680_2

Thanks for replying. I am studying the link you gave me.

The code I wrote was very bad. Sorry about that. Let me explain that I want to do.

I have N data I know is mixed normal distribution of order M. If I want to find the parameters (weights, locs, scales) of the mixed normal distribution from these N data, the code should be as follows, right?

def model(data):
    weights = pyro.sample("weights", dist.Dirichlet(0.5*torch.ones(M)))
    with pyro.plate("components", M):
        locs = pyro.sample("locs", dist.Normal(0., 10.))
        scale = pyro.sample("scale", dist.LogNormal(0., 2.))
    with pyro.plate("data", N):
        assignment = pyro.sample("assignment", dist.Categorical(weights))
        pyro.sample("obs", dist.Normal(locs[assignment], scale[assignment]), obs=data)

In addition, I have K data of these N data, and I want to estimate the mixture normal distribution for all K data. Therefore, I thought of the following process.

for K number of data.
    Parameter estimation by MAP estimation or MCMC with the above model

I think i can do this by performing the estimation process multiple times.
However, if K is large, the “for loop” is too slow.

So, I thought of solving the data as a multivariate mixture of uncorrelated normal distributions with the from (K, N). I think this can be faster than the “for loop” by using the GPU.

However, I don’t know how to write the code. How can I estimate the parameters of a multivariate normal distribution by using the information that there is no correlation?

one way to do this would be to use the lower level hmc interface

http://num.pyro.ai/en/stable/mcmc.html?highlight=hmc(#numpyro.infer.hmc.hmc

to define a method for fitting a single dataset and then using vmap to do many such fits in parallel. care needs to be taken to split the rng keys appropriately