EM MVN: decouple from SemObsMissing

so EM MVN could be done when SemObsMissing is constructed

Author: Alexey Stukalov

src/StructuralEquationModels.jl

@@ -42,8 +42,9 @@ include("frontend/pretty_printing.jl")
 # observed
 include("observed/data.jl")
 include("observed/covariance.jl")
-include("observed/missing.jl")
+include("observed/missing_pattern.jl")
 include("observed/EM.jl")
+include("observed/missing.jl")
 # constructor
 include("frontend/specification/Sem.jl")
 include("frontend/specification/documentation.jl")

src/additional_functions/start_val/start_fabin3.jl

@@ -17,19 +17,10 @@ function start_fabin3(model::AbstractSemSingle; kwargs...)
     )
 end
 
-function start_fabin3(observed, imply, optimizer, args...; kwargs...)
+function start_fabin3(observed::SemObserved, imply, optimizer, args...; kwargs...)
     return start_fabin3(imply.ram_matrices, obs_cov(observed), obs_mean(observed))
 end
 
-# SemObservedMissing
-function start_fabin3(observed::SemObservedMissing, imply, optimizer, args...; kwargs...)
-    if !observed.em_model.fitted
-        em_mvn(observed; kwargs...)
-    end
-
-    return start_fabin3(imply.ram_matrices, observed.em_model.Σ, observed.em_model.μ)
-end
-
 function start_fabin3(
     ram_matrices::RAMMatrices,
     Σ::AbstractMatrix,

src/frontend/fit/fitmeasures/minus2ll.jl

@@ -40,10 +40,7 @@ end
 # compute likelihood for missing data - H1 -------------------------------------------------
 # -2ll =  ∑ log(2π)*(nᵢ + mᵢ) + ln(Σᵢ) + (mᵢ - μᵢ)ᵀ Σᵢ⁻¹ (mᵢ - μᵢ)) + tr(SᵢΣᵢ)
 function minus2ll(observed::SemObservedMissing)
-    observed.em_model.fitted || em_mvn(observed)
-
-    μ = observed.em_model.μ
-    Σ = observed.em_model.Σ
+    Σ, μ = obs_cov(observed), obs_mean(observed)
 
     F = 0.0
     for pat in observed.patterns

src/observed/EM.jl

@@ -5,15 +5,17 @@
 # what about random restarts?
 
 """
-    em_mvn(;
-        observed::SemObservedMissing,
-        start_em = start_em_observed,
-        max_iter_em = 100,
-        rtol_em = 1e-4,
-        kwargs...)
+    em_mvn(patterns::AbstractVector{SemObservedMissingPattern};
+           start_em = start_em_observed,
+           max_iter_em = 100,
+           rtol_em = 1e-4,
+           kwargs...)
 
-Estimates the covariance matrix and mean vector of the normal distribution via expectation maximization for `observed`.
-Overwrites the statistics stored in `observed`.
+Estimates the covariance matrix and mean vector of the
+multivariate normal distribution (MVN)
+via expectation maximization (EM) for `observed`.
+
+Returns the tuple of the EM covariance matrix and the EM mean vector.
 
 Uses the EM algorithm for MVN-distributed data with missing values
 adapted from the supplementary material to the book *Machine Learning: A Probabilistic Perspective*,
@@ -22,23 +24,21 @@ copyright (2010) Kevin Murphy and Matt Dunham: see
 [*emAlgo.m*](https://github.com/probml/pmtk3/blob/master/toolbox/Algorithms/optimization/emAlgo.m) scripts.
 """
 function em_mvn(
-    observed::SemObservedMissing;
+    patterns::AbstractVector{<:SemObservedMissingPattern};
     start_em = start_em_observed,
-    max_iter_em = 100,
-    rtol_em = 1e-4,
+    max_iter_em::Integer = 100,
+    rtol_em::Number = 1e-4,
     kwargs...,
 )
-    n_man = SEM.n_man(observed)
-
-    # preallocate stuff?
-    𝔼x_pre = zeros(n_man)
-    𝔼xxᵀ_pre = zeros(n_man, n_man)
+    n_man = SEM.n_man(patterns[1])
 
     ### precompute for full cases
-    fullpat = observed.patterns[1]
-    if nmissed_vars(fullpat) == 0
-        sum!(reshape(𝔼x_pre, 1, n_man), fullpat.data)
-        mul!(𝔼xxᵀ_pre, fullpat.data', fullpat.data)
+    𝔼x_full = zeros(n_man)
+    𝔼xxᵀ_full = zeros(n_man, n_man)
+    if nmissed_vars(patterns[1]) == 0
+        fullpat = patterns[1]
+        sum!(reshape(𝔼x_full, 1, n_man), fullpat.data)
+        mul!(𝔼xxᵀ_full, fullpat.data', fullpat.data)
     else
         @warn "No full cases pattern found"
     end
@@ -47,72 +47,79 @@ function em_mvn(
     # estepFn = (em_model, data) -> estep(em_model, data, EXsum, EXXsum, ismissing, missingRows, n_obs)
 
     # initialize
-    em_model = start_em(observed; kwargs...)
-    em_model_prev = EmMVNModel(zeros(n_man, n_man), zeros(n_man), false)
-    iter = 1
-    done = false
-    𝔼x = zeros(n_man)
-    𝔼xxᵀ = zeros(n_man, n_man)
-
-    while !done
-        step!(em_model, observed, 𝔼x, 𝔼xxᵀ, 𝔼x_pre, 𝔼xxᵀ_pre)
-
-        if iter > max_iter_em
-            done = true
-            @warn "EM Algorithm for MVN missing data did not converge. Likelihood for FIML is not interpretable.
-            Maybe try passing different starting values via 'start_em = ...' "
-        elseif iter > 1
-            # done = isapprox(ll, ll_prev; rtol = rtol)
-            done =
-                isapprox(em_model_prev.μ, em_model.μ; rtol = rtol_em) &&
-                isapprox(em_model_prev.Σ, em_model.Σ; rtol = rtol_em)
+    Σ₀, μ = start_em(patterns; kwargs...)
+    Σ = convert(Matrix, Σ₀)
+    @assert all(isfinite, Σ) all(isfinite, μ)
+    Σ_prev, μ_prev = copy(Σ), copy(μ)
+
+    iter = 0
+    converged = false
+    while !converged && (iter < max_iter_em)
+        em_step!(Σ, μ, Σ_prev, μ_prev, patterns, 𝔼x_full, 𝔼xxᵀ_full)
+
+        if iter > 0
+            Δμ = norm(μ - μ_prev)
+            ΔΣ = norm(Σ - Σ_prev)
+            Δμ_rel = Δμ / max(norm(μ_prev), norm(μ))
+            ΔΣ_rel = ΔΣ / max(norm(Σ_prev), norm(Σ))
+            #@info "Iteration #$iter: ΔΣ=$(ΔΣ) ΔΣ/Σ=$(ΔΣ_rel) Δμ=$(Δμ) Δμ/μ=$(Δμ_rel)"
+            # converged = isapprox(ll, ll_prev; rtol = rtol)
+            converged = ΔΣ_rel <= rtol_em && Δμ_rel <= rtol_em
+        end
+        if !converged
+            Σ, Σ_prev = Σ_prev, Σ
+            μ, μ_prev = μ_prev, μ
         end
-
-        # print("$iter \n")
         iter += 1
-        copyto!(em_model_prev.μ, em_model.μ)
-        copyto!(em_model_prev.Σ, em_model.Σ)
+        #@info "$iter\n"
     end
 
-    # update EM Mode in observed
-    observed.em_model.Σ .= em_model.Σ
-    observed.em_model.μ .= em_model.μ
-    observed.em_model.fitted = true
+    if !converged
+        @warn "EM Algorithm for MVN missing data did not converge in $iter iterations.\n" *
+              "Likelihood for FIML is not interpretable.\n" *
+              "Maybe try passing different starting values via 'start_em = ...' "
+    else
+        @info "EM for MVN missing data converged in $iter iterations"
+    end
 
-    return nothing
+    return Σ, μ
 end
 
 # E and M steps -----------------------------------------------------------------------------
 
-# update em_model
-function step!(em_model::EmMVNModel, observed::SemObserved, 𝔼x, 𝔼xxᵀ, 𝔼x_pre, 𝔼xxᵀ_pre)
+function em_step!(
+    Σ::AbstractMatrix,
+    μ::AbstractVector,
+    Σ₀::AbstractMatrix,
+    μ₀::AbstractVector,
+    patterns::AbstractVector{<:SemObservedMissingPattern},
+    𝔼x_full,
+    𝔼xxᵀ_full,
+)
     # E step, update 𝔼x and 𝔼xxᵀ
-    fill!(𝔼x, 0)
-    fill!(𝔼xxᵀ, 0)
-
-    μ = em_model.μ
-    Σ = em_model.Σ
+    copy!(μ, 𝔼x_full)
+    copy!(Σ, 𝔼xxᵀ_full)
 
     # Compute the expected sufficient statistics
-    for pat in observed.patterns
+    for pat in patterns
         (nmissed_vars(pat) == 0) && continue # skip full cases
 
         # observed and unobserved vars
         u = pat.miss_mask
         o = pat.obs_mask
 
         # precompute for pattern
-        Σoo_chol = cholesky(Symmetric(Σ[o, o]))
-        Σuo = Σ[u, o]
-        μu = μ[u]
-        μo = μ[o]
+        Σoo_chol = cholesky(Symmetric(Σ₀[o, o]))
+        Σuo = Σ₀[u, o]
+        μu = μ₀[u]
+        μo = μ₀[o]
 
         𝔼xu = fill!(similar(μu), 0)
         𝔼xo = fill!(similar(μo), 0)
         𝔼xᵢu = similar(μu)
 
         𝔼xxᵀuo = fill!(similar(Σuo), 0)
-        𝔼xxᵀuu = n_obs(pat) * (Σ[u, u] - Σuo * (Σoo_chol \ Σuo'))
+        𝔼xxᵀuu = n_obs(pat) * (Σ₀[u, u] - Σuo * (Σoo_chol \ Σuo'))
 
         # loop trough data
         @inbounds for rowdata in eachrow(pat.data)
@@ -124,24 +131,21 @@ function step!(em_model::EmMVNModel, observed::SemObserved, 𝔼x, 𝔼xxᵀ, 
             𝔼xo .+= rowdata
         end
 
-        𝔼xxᵀ[o, o] .+= pat.data' * pat.data
-        𝔼xxᵀ[u, o] .+= 𝔼xxᵀuo
-        𝔼xxᵀ[o, u] .+= 𝔼xxᵀuo'
-        𝔼xxᵀ[u, u] .+= 𝔼xxᵀuu
+        Σ[o, o] .+= pat.data' * pat.data
+        Σ[u, o] .+= 𝔼xxᵀuo
+        Σ[o, u] .+= 𝔼xxᵀuo'
+        Σ[u, u] .+= 𝔼xxᵀuu
 
-        𝔼x[o] .+= 𝔼xo
-        𝔼x[u] .+= 𝔼xu
+        μ[o] .+= 𝔼xo
+        μ[u] .+= 𝔼xu
     end
 
-    𝔼x .+= 𝔼x_pre
-    𝔼xxᵀ .+= 𝔼xxᵀ_pre
-
     # M step, update em_model
-    em_model.μ .= 𝔼x ./ n_obs(observed)
-    em_model.Σ .= 𝔼xxᵀ ./ n_obs(observed)
-    mul!(em_model.Σ, em_model.μ, em_model.μ', -1, 1)
+    k = inv(sum(n_obs, patterns))
+    lmul!(k, Σ)
+    lmul!(k, μ)
+    mul!(Σ, μ, μ', -1, 1)
 
-    #Σ = em_model.Σ
     # ridge Σ
     # while !isposdef(Σ)
     #     Σ += 0.5I
@@ -150,42 +154,48 @@ function step!(em_model::EmMVNModel, observed::SemObserved, 𝔼x, 𝔼xxᵀ, 
     # diagonalization
     #if !isposdef(Σ)
     #    print("Matrix not positive definite")
-    #    em_model.Σ .= 0
-    #    em_model.Σ[diagind(em_model.Σ)] .= diag(Σ)
+    #    Σ .= 0
+    #    Σ[diagind(em_model.Σ)] .= diag(Σ)
     #else
-    # em_model.Σ = Σ
+    # Σ = Σ
     #end
 
-    return em_model
+    return Σ, μ
 end
 
 # generate starting values -----------------------------------------------------------------
 
 # use μ and Σ of full cases
-function start_em_observed(observed::SemObservedMissing; kwargs...)
-    fullpat = observed.patterns[1]
+function start_em_observed(patterns::AbstractVector{<:SemObservedMissingPattern}; kwargs...)
+    fullpat = patterns[1]
     if (nmissed_vars(fullpat) == 0) && (n_obs(fullpat) > 1)
         μ = copy(fullpat.obs_mean)
-        Σ = copy(fullpat.obs_cov)
+        Σ = copy(parent(fullpat.obs_cov))
         if !isposdef(Σ)
             Σ = Diagonal(Σ)
         end
-        return EmMVNModel(convert(Matrix, Σ), μ, false)
+        return Σ, μ
     else
         return start_em_simple(observed, kwargs...)
     end
 end
 
 # use μ = O and Σ = I
-function start_em_simple(observed::SemObservedMissing; kwargs...)
-    μ = zeros(n_man(observed))
-    Σ = rand(n_man(observed), n_man(observed))
+function start_em_simple(patterns::AbstractVector{<:SemObservedMissingPattern}; kwargs...)
+    nvars = n_man(first(patterns))
+    μ = zeros(nvars)
+    Σ = rand(nvars, nvars)
     Σ = Σ * Σ'
     # Σ = Matrix(1.0I, n_man, n_man)
-    return EmMVNModel(Σ, μ, false)
+    return Σ, μ
 end
 
 # set to passed values
-function start_em_set(observed::SemObservedMissing; model_em, kwargs...)
-    return em_model
+function start_em_set(
+    patterns::AbstractVector{<:SemObservedMissingPattern};
+    obs_cov::AbstractMatrix,
+    obs_mean::AbstractVector,
+    kwargs...,
+)
+    return copy(obs_cov), copy(obs_mean)
 end