RでEMアルゴリズムによる混合ガウスモデル最尤推定

PRML 9.2.2に記載の通り、EMアルゴリズムによって混合ガウスモデルの最尤推定が行われる過程と、対数尤度関数の収束の様子を示します。

library(mvtnorm)
library(plotrix)
frame()
set.seed(0)
par(mfrow=c(4, 4))
par(mar=c(2.5, 2.5, 1, 0.1))
par(mgp=c(1.5, .5, 0))
xrange <- c(-2, 2)
yrange <- c(-2, 2)
D <- 2
K <- 2
data(faithful)
x <- as.matrix(faithful)
N <- nrow(x)
#N <- 100
#x <- rmvnorm(N / 2, c(-1, -1), matrix(c(.4, 0, 0, .4), D))
#x <- rbind(x, rmvnorm(N / 2, c(1, 1), matrix(c(.4, 0, 0, .4), D)))
x <- t((t(x) - apply(x, 2, mean)) / apply(x, 2, sd))  # normalize
mu <- matrix(c(-1.5, 1, 1.5, -1), K, byrow=T)
sigma <- rep(list(diag(0.5, D)), K)
pz <- rep(1 / K, K)
gamma <- matrix(NA, nrow=N, ncol=K)
likelihood <- numeric()

iteration <- 0
repeat {
    cat("mu\n");print(mu)
    cat("sigma\n");print(sigma)
    cat("pi\n");print(pz)

    plot(x, xlim=xrange, ylim=yrange, col=ifelse(is.na(gamma[, 2]), 1, hsv(gamma[, 2] * .4)), pch=1)
    for (k in 1:K) {
        e <- eigen(sigma[[k]])
        draw.ellipse(mu[k, 1], mu[k, 2], sqrt(e$values[1]), sqrt(e$values[2]), 
            atan2(e$vectors[2, 1], e$vectors[1, 1]) / pi * 180,
            border=hsv((k - 1) * .4, 1, 0.8), lwd=2)
    }
    title(paste0("EM step#", iteration))

    # E step
    for (n in 1:N) {
        pzx <- sapply(
            1:K, 
            function(k) pz[k] * dmvnorm(x[n, ], mean=mu[k, ], sigma=sigma[[k]])
            )
        gamma[n, ] <- pzx / sum(pzx)
    }

    # M step
    nk <- colSums(gamma)
    for (k in 1:K) {
        mu[k, ] <- colSums(x * gamma[, k]) / nk[k]
        sigma[[k]] <- matrix(
            rowSums(sapply(1:N, function(n) gamma[n, k] * outer(x[n, ] - mu[k, ], x[n, ] - mu[k, ]))),
            D) / nk[k]
        pz[k] <- nk[k] / N
    }

    # likelihood
    likelihood <- c(likelihood, sum(sapply(1:N, function(n)
            log(sum(sapply(
                1:K, 
                function(k)
                    pz[k] * dmvnorm(x[n, ], mean=mu[k, ], sigma=sigma[[k]])
                )))
            )))

    if (length(likelihood) > 1 
        && likelihood[length(likelihood)] - likelihood[length(likelihood) - 1] < 1.0E-2) {
        break
    }
    iteration <- iteration + 1
}

plot(likelihood, type="l", xlab="iteration", ylab="ln p(X)")
title("ln p(X)")