# ExpectationMaximization

This package provides a simple implementation of the Expectation Maximization (EM) algorithm used to fit mixture models. Due to Julia amazing dispatch systems, generic and reusable code spirit, and the Distributions.jl package, the code while being very generic is both very expressive and fast! (Have a look at the Benchmark section)

In particular, it works on a lot of mixtures:

- Mixture of Univariate continuous distributions
- Mixture of Univariate discrete distributions
- Mixture of Multivariate distributions (continuous or discrete)
- Mixture of mixtures (univariate or multivariate and continuous or discrete)
- More?

Note that Distributions *currently* does not allow `MixtureModel`

to both have discrete and continuous components (but what does that? Rain).

Just define a `mix::MixtureModel`

and do `fit_mle(mix, y)`

with your data `y`

and that's it!
**Have a look at the Examples section**.

To work, the only requirements are that the components of the mixture `dist ∈ dists = components(mix)`

considered (custom or coming from an existing package)

- Are a subtype of
`Distribution`

i.e.`dist<:Distribution`

. - The
`logpdf(dist, y)`

is defined (it is used in the E-step) - The
`fit_mle(dist, y, weigths)`

returns the distribution with parameters equals to MLE. This is used in the M-step of the`ClassicalEM`

algorithm. For the`StocasticEM`

version, only`fit_mle(dist, y)`

is needed. Type or instance version of`fit_mle`

for your`dist`

are accepted thanks to this conversion line.

## TODO (feel free to contribute)

[] Add more variants to of the EM algorithm (so far there are the classic and stochastic version).

[] Better benchmark against other EM implementations

[] Speed up code (always!). So far, I focused on readable code.

## Example

Also have a look at the [examples](@ref Examples) section.

using Distributions
using ExpectationMaximization

### Model

N = 50_000
θ₁ = 10
θ₂ = 5
α = 0.2
β = 0.3
# Mixture Model here one can put any classical distributions
mix_true = MixtureModel([Exponential(θ₁), Gamma(α, θ₂)], [β, 1 - β])
# Generate N samples from the mixture
y = rand(mix_true, N)

### Inference

# Initial guess
mix_guess = MixtureModel([Exponential(1), Gamma(0.5, 1)], [0.5, 1 - 0.5])
# Fit the MLE with the EM algorithm
mix_mle = fit_mle(mix_guess, y; display = :iter, atol = 1e-3, robust = false, infos = false)

### Verify results

rtol = 5e-2
p = params(mix_mle)[1] # (θ₁, (α, θ₂))
isapprox(β, probs(mix_mle)[1]; rtol = rtol)
isapprox(θ₁, p[1]...; rtol = rtol)
isapprox(α, p[2][1]; rtol = rtol)
isapprox(θ₂, p[2][2]; rtol = rtol)