Algorithms & Methods

ClassicEM<:AbstractEM

The EM algorithm was introduced by A. P. Dempster, N. M. Laird and D. B. Rubin in 1977 in the reference paper Maximum Likelihood from Incomplete Data Via the EM Algorithm.

Base.@kwdef struct StochasticEM<:AbstractEM
    rng::AbstractRNG = Random.GLOBAL_RNG
end

The Stochastic EM algorithm was introduced by G. Celeux, and J. Diebolt. in 1985 in The SEM Algorithm: A probabilistic teacher algorithm derived from the EM algorithm for the mixture problem.

The default random number generator is Random.GLOBAL_RNG. Pass another one with StochasticEM(rng) or StochasticEM(; rng = rng), e.g. StochasticEM(MersenneTwister(0)), to make a run reproducible. The argument must be an AbstractRNG: an integer seed such as StochasticEM(0) is a MethodError.

fit_mle(mix::MixtureModel, y::AbstractVecOrMat, weights...; method = ClassicEM(), display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false, infos=false)

Use an Expectation Maximization (EM) algorithm to maximize the Loglikelihood (fit) the mixture with an i.i.d sample y. The mix input is a mixture that is used to initialize the EM algorithm. When y is an AbstractMatrix, each column is one observation, i.e. there are size(y, 2) observations of dimension size(y, 1).

• weights at most one positional weight vector w, of length size_sample(y), may be given; it then computes a weighted version of the EM. (Useful for fitting mixture of mixtures)
• method determines the algorithm used.
• infos = true returns the tuple (mix_fitted, history) instead of just mix_fitted, where history::Dict{String,Any} holds "converged"::Bool, "iterations"::Int (number of EM iterations actually performed) and "logtots"::Vector (the loglikelihood after each of those iterations). The iteration-0 loglikelihood is not stored, so length(history["logtots"]) == history["iterations"], and it is empty when maxiter = 0.
• robust = true will prevent the (log)likelihood to overflow to -∞ or ∞.
• atol criteria determining the convergence of the algorithm. If the Loglikelihood difference between two iteration i and i+1 is smaller than atol i.e. |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<atol, the algorithm stops.
• rtol relative tolerance for convergence, |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<rtol*(|ℓ⁽ⁱ⁺¹⁾| + |ℓ⁽ⁱ⁾|)/2 (does not check if rtol is nothing)
• display value can be :none, :iter, :final to display Loglikelihood evolution at each iterations :iter or just the final one :final

fit_mle(mix::AbstractArray{<:MixtureModel}, y::AbstractVecOrMat, weights...; method = ClassicEM(), display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false, infos=false)

Do the same as fit_mle for each (initial) mixture in the mix array, then keep the fit with the largest final loglikelihood history["logtots"][end] (taken as -Inf when logtots is empty, e.g. maxiter = 0); ties keep the earliest initial condition.

Every initial condition runs inside a try/catch, so one singular solution does not abort the whole sweep (using robust = true should be enough to avoid most errors in the first place). A failing initial condition is reported with @debug and skipped.

• An InterruptException is never swallowed: it is rethrown immediately, so Ctrl-C still stops the sweep.
• If every initial condition fails, the error raised by the first failing one is rethrown.
• All keywords are forwarded unchanged to fit_mle(mix[j], y, weights...); see its docstring for their meaning.
• infos = true returns (mix_best, history_best) for the selected fit instead of just mix_best.

predict(mix::MixtureModel, y::AbstractVecOrMat; robust=false)

Evaluate the most likely category for each observation given a MixtureModel, i.e. the argmax over the row of predict_proba belonging to that observation. Returns a length-N Vector{Int} of component indices in 1:ncomponents(mix); ties go to the lowest index. When y is an AbstractMatrix, each column is one observation, so N = size(y, 2).

• robust = true will prevent the (log)likelihood to overflow to -∞ or ∞.

predict_proba(mix::MixtureModel, y::AbstractVecOrMat; robust=false)

Evaluate the probability for each observation to belong to a category given a MixtureModel. Returns a fresh N × K matrix whose row n is the posterior distribution of the component label of observation n, with K = ncomponents(mix) and N = size_sample(y) (length(y) for a vector, size(y, 2) for a matrix, where each column is one observation). Every row sums to 1 unless it is degenerate, i.e. the observation has zero density under every component (see robust).

• robust = true will prevent the (log)likelihood to under(overflow)flow to -∞ (or ∞).

fit_mle(g::D, args...) where {D<:Distribution}

In ExpectationMaximization.jl the "instance" version of fit_mle is supported (in addition of the current "type" version). Note this is not supported in Distributions.jl. Example: fit_mle(Bernoulli(0.2), x) is accepted in addition of fit_mle(Bernoulli, x) this allows compatibility with how fit_mle(g::Product) and fit_mle(g::MixtureModel) are written.

By default the instance is simply dropped in favour of typeof(g).name.wrapper. More specific methods are provided wherever the instance carries information the type does not: DiagNormal/IsoNormal/FullNormal (the covariance structure), Binomial (ntrials) and Categorical (ncategories).

fit_mle(g::Product, x::AbstractMatrix, args...)

The fit_mle for a multivariate Product distribution g is the product_distribution of the fit_mle of each of its components, marginal s being fitted on row s of x — i.e. each column of x is one observation and length(g) == size(x, 1) is required. args... is forwarded to every marginal fit_mle, so it is either empty or a single weight vector γ of length size(x, 2). Product is meant to be deprecated in the next versions of Distributions.jl. Use the analog VectorOfUnivariateDistribution type instead.

fit_mle(dists::VectorOfUnivariateDistribution, x::AbstractMatrix{<:Real}, args...)

The fit_mle for a VectorOfUnivariateDistribution dists is the product_distribution of the fit_mle of each of its components, marginal s being fitted on view(x, s, :) — i.e. each column of x is one observation and size(x, 1) == length(dists) is required. Because the row is passed as a view, component fit_mle methods must accept AbstractVector rather than Vector. args... is forwarded to every marginal fit_mle, so it is either empty or a single weight vector γ of length size(x, 2). VectorOfUnivariateDistribution should act like the old Product, while the sibling fit_mle(dists::ArrayOfUnivariateDistribution, x::AbstractArray, args...) (same idea, but x an array of arrays) is not really tested yet and is deliberately left undocumented.

fit_mle(::Type{<:Dirac}, x::AbstractArray{<:Real})
fit_mle(::Type{<:Dirac}, x::AbstractArray{<:Real}, w::AbstractArray{Float64})

fit_mle for Dirac distribution (weighted or not) data sets. Returns Dirac(first(x)) when all the observations carrying a non-zero weight are equal, and Dirac(NaN) otherwise. Note that the weighted method requires w::AbstractArray{Float64} exactly: another real element type (e.g. Vector{Float32}) is a MethodError, since there is no other weighted Dirac method to fall back on.

fit_mle(::Type{<:Laplace}, x::AbstractArray{<:Real}, w::AbstractArray{<:Real})

fit_mle for Laplace distribution weighted data sets.

fit_mle(::Type{<:Uniform}, x::AbstractArray{<:Real}, w::AbstractArray{<:Real})

fit_mle for Uniform distribution weighted data sets. It is the same as the unweighted fit applied to the observations carrying a non-zero weight: the MLE only depends on the extrema of the support, so the non-zero weight values themselves are irrelevant. Requires size(x) == size(w).

fit_mle!(α::AbstractVector, dists::AbstractVector{F} where {F<:Distribution}, y::AbstractVecOrMat, method::ClassicEM; display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false)
fit_mle!(α::AbstractVector, dists::AbstractVector{F} where {F<:Distribution}, y::AbstractVecOrMat, w::Union{Nothing,AbstractVector}, method::ClassicEM; display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false)

Use the EM algorithm to update in place the Distribution dists and weights α composing a mixture distribution. When y is an AbstractMatrix, each column is one observation. w, when given and not nothing, is a weight vector of length size_sample(y). Returns history::Dict{String,Any} with keys "converged"::Bool, "iterations"::Int and "logtots"::Vector (loglikelihood after each performed iteration, empty when maxiter = 0).

• robust = true will prevent the (log)likelihood to overflow to -∞ or ∞.
• atol criteria determining the convergence of the algorithm. If the Loglikelihood difference between two iteration i and i+1 is smaller than atol i.e. |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<atol, the algorithm stops.
• rtol relative tolerance for convergence, |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<rtol*(|ℓ⁽ⁱ⁺¹⁾| + |ℓ⁽ⁱ⁾|)/2 (does not check if rtol is nothing)
• display value can be :none, :iter, :final to display Loglikelihood evolution at each iterations :iter or just the final one :final

fit_mle!(α::AbstractVector, dists::AbstractVector{F} where {F<:Distribution}, y::AbstractVecOrMat, method::StochasticEM; display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false)
fit_mle!(α::AbstractVector, dists::AbstractVector{F} where {F<:Distribution}, y::AbstractVecOrMat, w::Union{Nothing,AbstractVector}, method::StochasticEM; display=:none, maxiter=1000, atol=1e-3, rtol=nothing, robust=false)

Use the stochastic EM algorithm to update in place the Distribution dists and weights α composing a mixture distribution. When y is an AbstractMatrix, each column is one observation. w, when given and not nothing, is a weight vector of length size_sample(y). Returns history::Dict{String,Any} with keys "converged"::Bool, "iterations"::Int and "logtots"::Vector (loglikelihood after each performed iteration, empty when maxiter = 0).

• Throws a DomainError if the loglikelihood is not finite: some observation then has zero density under every component, its posterior row is NaN, and the S-step would silently assign it to component 1. Try robust = true or another initial condition.
• robust = true will prevent the (log)likelihood to overflow to -∞ or ∞.
• atol criteria determining the convergence of the algorithm. If the Loglikelihood difference between two iteration i and i+1 is smaller than atol i.e. |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<atol, the algorithm stops.
• rtol relative tolerance for convergence, |ℓ⁽ⁱ⁺¹⁾ - ℓ⁽ⁱ⁾|<rtol*(|ℓ⁽ⁱ⁺¹⁾| + |ℓ⁽ⁱ⁾|)/2 (does not check if rtol is nothing)
• display value can be :none, :iter, :final to display Loglikelihood evolution at each iterations :iter or just the final one :final

E_step!(LL, c, γ, s, dists, α, y; robust=false)

E-step, in two stages: loglikelihoods! fills LL[n, k] = log(α[k]) + logpdf(dists[k], y[n]) (y[:, n] for multivariate samples), then _softmax_rows! turns every row into a posterior. Returns γ.

• LL the N × K log-likelihood matrix. Entirely overwritten, first with the log-likelihoods and then with the posteriors.
• c a length-N vector filled with c[n] = logsumexp(LL[n, :]) = log ℙ(y[n]). The fit_mle! drivers sum (or weight-sum) it to get the loglikelihood, so no extra pass over y is needed.
• γ the N × K posterior matrix, γ[n, k] = ℙ(zₙ = k ∣ yₙ). It may alias LL, and every caller in this package passes γ === LL: the log-likelihoods are not needed once the posteriors have been formed, so the posteriors are written over them. Aliasing is therefore only safe as long as nothing afterwards reads LL expecting log-likelihoods; pass a distinct γ if you need both matrices at once.
• s a length-N scratch vector. Its contents are meaningless on entry and on exit.
• dists, α the current components and mixing weights; y the sample (a vector, or a D × N matrix).
• robust = true clamps ±Inf log-likelihoods before normalizing, which is what prevents a degenerate, unnormalized row of γ.

loglikelihoods!(LL::AbstractMatrix, dists, α, y::AbstractVector)
loglikelihoods!(LL::AbstractMatrix, dists, α, y::AbstractMatrix)

Fill LL[n, k] = log(α[k]) + logpdf(dists[k], y[n]) and return LL. For the AbstractMatrix method each column of y is one observation, so the entry is log(α[k]) + logpdf(dists[k], y[:, n]).

This is the extension hook of the E-step: add a method for your component or sample type if it can score a whole sample at once. All such a method has to do is give LL[n, k] the value above; how it computes it is up to you.

M_step!(α, dists, y, γ, method::ClassicEM)
M_step!(α, dists, y, γ, w, method::ClassicEM)

For the ClassicEM the weights γ computed at the E-step for each observation in y are used to update α and dists in place: α[k] = mean(γ[:, k]) and dists[k] = fit_mle(dists[k], y, γ[:, k]).

The weighted variant folds the observation weights into the posteriors once (γ .*= w) and normalizes α by sum(w), rather than materializing w .* γ[:, k] for each of the K components. It therefore overwrites γ, which is sound inside fit_mle! only because the next E-step rewrites γ entirely — so γ must not be read after the M-step.

M_step!(α, dists, y, cat, method::StochasticEM)
M_step!(α, dists, y, cat, w, method::StochasticEM)

For the StochasticEM the cat drawn at S-step for each observation in y is used to update α and dists in place. cat[k] indexes the observations assigned to component k, so the subsample is passed as a view rather than copied — component fit_mle methods must therefore accept SubArrays. A component fit_mle that slices such a view by rows should gather it once first, as fit_mle(::Product, ::AbstractMatrix, args...) does through _gather_rows: row-slicing a view(y, :, cat[k]) re-copies its vector column index once per row. The weighted variant sets α[k] = sum(w[cat[k]]) / sum(w) and forwards view(w, cat[k]) to each component fit.

_loglikelihood_col!(LLₖ, d::MvNormal, logα, y::AbstractMatrix)

Distributions.sqmahal!, which the generic path reaches through logpdf!, materializes a D × N centered copy of the sample and then still solves one triangular system per observation. These methods evaluate the entire sample in one pass instead:

• isotropic or diagonal Σ: a weighted sum of squares, with no temporary at all;
• full Σ: one blocked in-place PDMats.whiten! (a BLAS-3 trsm) per MVNORMAL_BLOCKSIZE observations, instead of one BLAS-2 trsv per observation.

Measured 3.5x-11.4x faster than the generic path for D from 2 to 100 at N = 1e5, agreeing to a few units in the last place, with an allocation bounded by MVNORMAL_BLOCKSIZE rather than growing with N.

_softmax_rows!(c, γ, LL, s)

Fused allocation-free row-wise log-sum-exp and softmax: writes c[n] = logsumexp(LL[n, :]) and γ[n, :] = softmax(LL[n, :]) in a single set of column-major passes, where logsumexp! needed a 16N-byte temporary plus a second pass over LL.

• γ may alias LL: every element is a same-index read-then-write and no later pass re-reads LL.
• s is a length-N scratch vector.
• Rows whose maximum is not finite are left unnormalized so that c and γ reproduce the previous logsumexp! based implementation bit for bit (this is what robust = true is for).