Type inference for mixture distributions

[kaandocal]

I am currently using Distributions.jl with mixture distributions (example: mixture of Gaussians) via MixtureModel. I have noticed that the compiler seems unable to perform type inference for basic functions such as mean and logpdf even when the MixtureModel struct is fully instantiated. Checking the source code these functions are implemented on the UnivariateMixture. I am currently specialising the functions I need for the specialised types I am working with but I was wondering if there isn't a better way of doing that....

[devmotion]

If you only work with mixtures of univariate Gaussians, you might want to use UnivariateGMM which is optimized for this use case.

Generally, if there are type inference problems with MixtureModel that are caused by the implementation in Distributions, it would be great to fix them. Do you have a concrete example?

[kaandocal]

Thanks for the recommendation! Unfortunately right now I'm using different mixture models including Poisson, NB etc.

Any example would be

using Distributions

mixture = MixtureModel([ Poisson(2.0), Poisson(10.0) ], Categorical(2))

@code_warntype mean(mixture)
@code_warntype pdf(mixture, 5)

Here several of the intermediate variables (including within _mixpdf1) don't have inferred types which leads to a significant performance decrease at the moment. I'd be happy to experiment with solutions but I don't want to break anything and Distributions is a very heavily used package...

[devmotion]

I made a PR with some fixes: #1308

[kaandocal]

Upon further inspection the problem seems to be that MixtureModel does not fully instantiate Categorical in its prior field. The compiler is then unable to infer the return value of probs(d); after annotating it all the type problems I've encountered in mean and logpdf go away.

[kaandocal]

This seems to be exactly the fix that's needed! Thanks a lot for submitting this.