Fitting a mixture model with a tidy EM

We simulate random data from two distributions- a normal centered at 0, and a normal centered at 4.

library(tidyverse)
theme_set(theme_bw())
set.seed(2016)
# setup: simulate data
observations <- data_frame(observation = seq_len(1000)) %>%
  mutate(oracle = sample(c("A", "B"), n(), TRUE)) %>%
  mutate(mu = ifelse(oracle == "A", 0, 4),
         x = rnorm(n(), mu, sd = 1))
ggplot(observations, aes(x)) +
  geom_histogram() +
  geom_vline(xintercept = c(0, 4), color = "red", lty = 2)

`stat_bin()` using `bins = 30`. Pick better value with `binwidth`.

How could we estimate assignments for this data, assuming we no longer knew them?

We perform expectation maximization:

set.seed(1337)
# we set up by assigning clusters randomly
maximized <- observations %>%
  mutate(assignment = sample(c("A", "B"), n(), TRUE))
for (i in 1:10) {
  # expectation: find the mean/sd of current clusters
  estimates <- maximized %>%
    group_by(assignment) %>%
    summarize(mu_hat = mean(x), sigma_hat = sd(x))
  
  # maximization: pick the best for each
  last <- maximized
  maximized <- observations %>%
    crossing(estimates) %>%
    mutate(log_likelihood = dnorm(x, mu_hat, sigma_hat, log = TRUE)) %>%
    group_by(observation) %>%
    top_n(1, log_likelihood) %>%
    ungroup()
  
  # have we converged yet?
  if (all(maximized$assignment == last$assignment)) {
    break
  }
}
ggplot(maximized, aes(x, fill = assignment)) +
  geom_histogram()

`stat_bin()` using `bins = 30`. Pick better value with `binwidth`.