EM.R

EM.R

#!/usr/bin/Rscript

expectation <- function(y, params) {
        pi <- params$pi
        mu1 <- params$mu1
        mu2 <- params$mu2
        s1 <- params$s1
        s2 <- params$s2

        pi * dnorm(y, mean=mu2, sd=s2) / ((1-pi)*dnorm(y, mean=mu1, sd=s1) + pi*dnorm(y, mean=mu2, sd=s2))
}

maximization <- function(y, gamma, params) {
        pi <- params$pi
        mu1 <- params$mu1
        mu2 <- params$mu2
        s1 <- params$s1
        s2 <- params$s2

        mu1_new <- weighted.mean(y, -(gamma-1))
        mu2_new <- weighted.mean(y, gamma)
        s1 <- sqrt( weighted.mean((y-mu1_new)**2, -(gamma-1)) )
        s2 <- sqrt( weighted.mean((y-mu2_new)**2, gamma) )
        pi_new <- mean(gamma)
        list(pi=pi_new,mu1=mu1_new,mu2=mu2_new,s1=s1,s2=s2)
}

x1 <- rnorm(10000, mean=-3, sd=1)
x2 <- rnorm(10000, mean=3,sd=1)
x <- x1
l <- as.logical(rbinom(length(x),size=1,p=0.55))
x[l] <- x2[l]

init_params <-  list(pi=0.5, mu1=x[1], mu2=x[2], s1=sd(x), s2=sd(x))
params <- init_params
while(1) {
        old_params <- params
        gamma <- expectation(x, params) 
        params <- maximization(x, gamma, params)
        print(unlist(params))
        if( all( abs(unlist(params) - unlist(old_params)) < max(unlist(old_params) * 1e-6, rep(1e-6,5) ) ) ) {
                break
        }
}
print(unlist(params))
plot(density(x))
r <- range(x)
grid <- seq(r[1],r[2],0.01)
lines( grid, (1-params$pi) * dnorm(grid, mean=params$mu1, sd=params$s1) + params$pi * dnorm(grid, mean=params$mu2, sd=params$s2), col=2)