I was (and still am) very confused by the varying benchmark results for `Real`

. Tested by replacing all `Real`

with `Float64`

```
import Pkg
Pkg.activate(temp=true, io=devnull)
V_OLD, V_NEW = "0.12.99", "0.12.102"
Pkg.add(name="LoopVectorization", version=V_NEW, io=devnull)
Pkg.add(["Distributions", "ProgressMeter", "BenchmarkTools"], io=devnull)
Pkg.status()
module EM
import Distributions: Dirichlet
using LoopVectorization
const AV = AbstractVector{T} where T
const AM = AbstractMatrix{T} where T
abstract type AbstractRegularization end
const MaybeReg = Union{AbstractRegularization, Nothing}
struct InvGamma <: AbstractRegularization
Ξ±::Float64
Ξ²::Float64
end
function init_G(K, N, a::Float64)
distr = Dirichlet(K, a)
rand(distr, N)
end
normal(x::Float64, ΞΌ::Float64, var::Float64) = exp(-(x-ΞΌ)^2 / (2var)) / sqrt(2Ο * var)
function ELBO(G, p, mu, var, data, reg::Nothing)::Float64
K, N = size(G)
ret = 0.0
@tturbo for n β 1:N, k β 1:K
q = G[k, n]
ret += q * (
log(p[k]) - (log(2Ο) + log(var[k]) + (data[n] - mu[k])^2 / var[k]) / 2
- log(q + 1e-100)
)
end
ret
end
@inline ELBO(G, p, mu, var, data, reg::InvGamma) = ELBO(G, p, mu, var, data, nothing)
# ===== E step =====
function finalize_posteriors!(G, norm::AV, reg::Nothing)
K, N = size(G)
@tturbo for n β 1:N, k β 1:K
G[k, n] /= norm[n]
end
end
@inline finalize_posteriors!(G, norm::AV, reg::InvGamma) =
finalize_posteriors!(G, norm, nothing)
function step_E!(G, norm::AV{<:Float64}, p, mu, var, data, reg::MaybeReg)
K, N = size(G)
@assert length(norm) == N
norm .= 0
@tturbo for n β 1:N, k β 1:K
G[k, n] = p[k] * exp(-(data[n] - mu[k])^2 / (2var[k])) / sqrt(2Ο * var[k])
norm[n] += G[k, n]
end
finalize_posteriors!(G, norm, reg)
end
# ===== M step =====
M_weights!(p, G, norm, reg::Nothing) = p .= norm ./ sum(norm)
@inline M_weights!(p, G, norm, reg::InvGamma) = M_weights!(p, G, norm, nothing)
function M_means!(mu, G, norm, data, reg::Nothing)
K, N = size(G)
mu .= 0
@tturbo for n β 1:N,k β 1:K
mu[k] += G[k, n] / norm[k] * data[n]
end
end
@inline M_means!(mu, G, norm, data, reg::InvGamma) =
M_means!(mu, G, norm, data, nothing)
finalize_variances!(var::AV, norm::AV, reg::Nothing) = @turbo var .= var ./ norm .+ 1e-6
function finalize_variances!(var::AV, norm::AV, reg::InvGamma)
Ξ±, Ξ² = reg.Ξ±, reg.Ξ²
@turbo @. var = (2Ξ² + var) / (2Ξ± + norm)
nothing
end
function M_variances!(var, G, norm, data, mu, reg::MaybeReg)
K, N = size(G)
var .= 0
@tturbo for n β 1:N, k β 1:K
var[k] += G[k, n] * (data[n] - mu[k])^2
end
finalize_variances!(var, norm, reg)
end
function step_M!(G, norm, p, mu, var, data, reg::MaybeReg)
K, N = size(G)
@assert K < N
evidences = @view norm[1:K]
evidences .= 0
@tturbo for n β 1:N, k β 1:K
evidences[k] += G[k, n]
end
M_weights!(p, G, evidences, reg)
M_means!(mu, G, evidences, data, reg)
M_variances!(var, G, evidences, data, mu, reg)
end
end # module
using ProgressMeter, Random, BenchmarkTools # , Revise
import Distributions: UnivariateGMM, Categorical
const AV = AbstractVector{T} where T
const AM = AbstractMatrix{T} where T
function sample_GMM(p::AV, mu::AV, var::AV; N::Integer)
distr = UnivariateGMM(mu, sqrt.(var), Categorical(p))
rand(distr, N)
end
mutable struct GaussianMixture{T<:Float64}
K::Integer
G::Union{Matrix{T}, Nothing}
norm::Union{Vector{T}, Nothing}
p::Vector{T}
mu::Vector{T}
var::Vector{T}
ELBO_hist::Vector{<:Float64}
end
function GaussianMixture{T}(K::Integer) where T<:Float64
@assert K > 0
GaussianMixture{T}(
K, nothing, nothing,
zeros(T, K), zeros(T, K), zeros(T, K),
Float64[]
)
end
@inline GaussianMixture(K::Integer) = GaussianMixture{Float64}(K)
function fit!(
gmm::GaussianMixture{T}, data::AV{<:Float64};
maxiter::Integer=10_000,
reg::EM.MaybeReg=nothing, a::Float64=100., reinit::Bool=false
) where T<:Float64
K, N = gmm.K, length(data)
has_reinit = false
if gmm.G === nothing || size(gmm.G) β (K, N) || reinit
gmm.G = EM.init_G(K, N, a)
gmm.norm = zeros(T, N)
has_reinit = true
end
@assert size(gmm.G) == (K, N)
if has_reinit
EM.step_M!(gmm.G, gmm.norm, gmm.p, gmm.mu, gmm.var, data, reg)
else
# Keep old parameters
EM.step_E!(gmm.G, gmm.norm, gmm.p, gmm.mu, gmm.var, data, reg)
end
gmm.ELBO_hist = zeros(T, maxiter)
for i β 1:maxiter
EM.step_E!(gmm.G, gmm.norm, gmm.p, gmm.mu, gmm.var, data, reg)
EM.step_M!(gmm.G, gmm.norm, gmm.p, gmm.mu, gmm.var, data, reg)
gmm.ELBO_hist[i] = EM.ELBO(gmm.G, gmm.p, gmm.mu, gmm.var, data, reg)
end
(; p=gmm.p, mu=gmm.mu, var=gmm.var)
end
function fit_rolling!(
gmm::GaussianMixture{T}, data::AV{<:Float64}, win_size::Integer;
kwargs...
) where T<:Float64
the_range = win_size:length(data)
P = Matrix{T}(undef, gmm.K, length(the_range))
M, V = similar(P), similar(P)
@showprogress for (i, off) β enumerate(the_range)
win = @view data[off-the_range[1]+1:off]
p, mu, var = fit!(gmm, win; kwargs...)
P[:, i] .= p
M[:, i] .= mu
V[:, i] .= var
end
P, M, V
end
@info "Generating data..."
const data = sample_GMM([.6, .4], [0, 0], [.3, .2]; N=500)
const K, WIN = 2, 300
@info "Parameters" K WIN
@info "Fitting... (no regularization)"
@benchmark fit_rolling!(mix, $data, $WIN) setup=(Random.seed!(85); mix = GaussianMixture(K))
@info "Fitting... (with regularization)"
@benchmark fit_rolling!(mix, $data, $WIN; reg=reg) setup=(Random.seed!(85); mix = GaussianMixture(K); reg = EM.InvGamma(.1, .1))
```

I see

which looks good to me.