WIP Mixture Models

[AustinRochford]

Re: #1401

This is my first pass at a flexible mixture model class, and definitely needs a lot of work. As this notebook shows, this code can support two use cases:

1. comp_dists is a PyMC3 distribution; that is, each of the mixture components are from the same distributional family, differing only in their parameters (e.g. a mixture of normals)
2. comp_dists is an iterable of PyMC3 distributions (e.g. a zero inflated Poisson)

There are a few issues to address here, that I'd love to get some feedback on:

• How to subclass Discrete or Continuous as appropriate, based on the component distributions (it seems like there is some Python metaprogramming magic that should work here, but I am not very well versed in that sort of thing)
• Intuitive broadcasting
• NUTS seems slow for these mixture models; this may be an initialization/scaling problem

@twiecki @springcoil @fonnesbeck any feedback/guidance you could give would be much appreciated :)

[twiecki]

Nice, this is a really clear implementation.

How to subclass Discrete or Continuous as appropriate, based on the component distributions (it seems like there is some Python metaprogramming magic that should work here, but I am not very well versed in that sort of thing)

I assume you mean for discrete mixtures? If you look at the code, it's mostly about default dtypes:

class Discrete(Distribution):
    """Base class for discrete distributions"""

    def __init__(self, shape=(), dtype='int64', defaults=['mode'], *args, **kwargs):
        super(Discrete, self).__init__(
            shape, dtype, defaults=defaults, *args, **kwargs)


class Continuous(Distribution):
    """Base class for continuous distributions"""

    def __init__(self, shape=(), dtype='float64', defaults=['median', 'mean', 'mode'], *args, **kwargs):
        super(Continuous, self).__init__(
            shape, dtype, defaults=defaults, *args, **kwargs)

As such, maybe we could inherit from Distribution instead and take the dtype from the mixture?

Intuitive broadcasting

Can you specify an example case?

NUTS seems slow for these mixture models; this may be an initialization/scaling problem

Yes, that's the most common reason. If you post an example it would help with testing.

[AustinRochford]

@twiecki thanks for the feedback.

The Continuous vs. Discrete approach makes complete sense.

For broadcasting, the notebook I posted shows two different cases. As the code is now, when the component distributions all have the same type (as in the normal mixture case comp_dists=pm.Normal.dist(mu, sd)), then the observations for the Mixture need to be broadcastable with mu (hence observed=x[:, np.newaxis]), but when the components are specified as a list of distributions (as in the zero-inflated Poisson case comp_dists=[pm.ConstantDist.dist(0), pm.Poissoin.dist(lam)), the observations should be one-dimensional. As I write this out, I think this is easier to handle than I originally thought.

[AustinRochford]

I added a NUTS example to the linked notebook; 6 samples/sec seems quite low from my experience.

[twiecki]

I think this looks good. The NUTS issue I bet is initialization, can you try with ADVI init (https://gist.github.com/jonsedar/cd4985bbfafdba61b3c8d077dd91f237)?

In any case, I wouldn't block on the NUTS issue, we can resolve it later. Instead, I would focus on:

• tests
• docs (the example notebook you have is close, just needs a bit of text).

[AustinRochford]

@twiecki yup, those are my priorities. I would also add random value generation to support posterior predictive sampling to the list of necessary additions.

One more question: do we want to replace the current ZeroInflated* implementations with mixtures?

[twiecki]

One more question: do we want to replace the current ZeroInflated* implementations with mixtures?

Yes, that's probably cleaner.

[AustinRochford]

I also need to better understand/resolve the issues in #1449 before I can make this code work with dependent weights.

[AustinRochford]

This should perhaps use tt.shapepadright instead of np.newaxis

[AustinRochford]

I am not going to port the ZeroInflated* models to be subclasses in this pull request. Due to #1452, we no longer broadcast a certain way in ConstantDist's logp that makes that a bit trickier. To keep things simpler, I will separate it into a subsquent PR unless anyone has objections.

[twiecki]

That sounds right. Let me know when you want us to take another look.

[AustinRochford]

@twiecki I think I have gotten this into a pretty good state and would love your feedback on it again. Happy to make any changes you think are necessary to merge.

[twiecki]

The simplicity to create a Normal Mixture here is really validating, nicely done.

[twiecki]

This is a really high-quality PR, thanks!