A maximização da expectativa (EM) é um tipo de método probabilístico para classificar dados. Por favor, corrija-me se eu estiver errado se não for um classificador.

O que é uma explicação intuitiva desta técnica EM? O que está expectationaqui e o que está sendo maximized?

Nota: o código por trás desta resposta pode ser encontrado aqui .

Suponha que temos alguns dados amostrados de dois grupos diferentes, vermelho e azul:

Aqui, podemos ver qual ponto de dados pertence ao grupo vermelho ou azul. Isso facilita encontrar os parâmetros que caracterizam cada grupo. Por exemplo, a média do grupo vermelho é cerca de 3, a média do grupo azul é cerca de 7 (e poderíamos encontrar a média exata se quiséssemos).

Geralmente, isso é conhecido como estimativa de máxima verossimilhança . Com alguns dados, calculamos o valor de um parâmetro (ou parâmetros) que melhor explica esses dados.

Agora imagine que não podemos ver qual valor foi amostrado de qual grupo. Tudo parece roxo para nós:

Aqui, sabemos que existem dois grupos de valores, mas não sabemos a qual grupo pertence um determinado valor.

Ainda podemos estimar as médias para o grupo vermelho e o grupo azul que melhor se ajustam a esses dados?

Sim, muitas vezes podemos! A maximização da expectativa nos dá uma maneira de fazer isso. A ideia geral por trás do algoritmo é esta:

1. Comece com uma estimativa inicial do que cada parâmetro pode ser.
2. Calcule a probabilidade de que cada parâmetro produza o ponto de dados.
3. Calcule pesos para cada ponto de dados indicando se é mais vermelho ou mais azul com base na probabilidade de ser produzido por um parâmetro. Combine os pesos com os dados ( expectativa ).
4. Calcule uma estimativa melhor para os parâmetros usando os dados ajustados por peso ( maximização ).
5. Repita as etapas 2 a 4 até que a estimativa do parâmetro convirja (o processo para de produzir uma estimativa diferente).

Essas etapas precisam de mais explicações, portanto, examinarei o problema descrito acima.

Exemplo: estimativa de média e desvio padrão

Usarei Python neste exemplo, mas o código deve ser bem fácil de entender se você não estiver familiarizado com essa linguagem.

Suponha que temos dois grupos, vermelho e azul, com os valores distribuídos como na imagem acima. Especificamente, cada grupo contém um valor retirado de uma distribuição normal com os seguintes parâmetros:

import numpy as np
from scipy import stats

np.random.seed(110) # for reproducible results

# set parameters
red_mean = 3
red_std = 0.8

blue_mean = 7
blue_std = 2

# draw 20 samples from normal distributions with red/blue parameters
red = np.random.normal(red_mean, red_std, size=20)
blue = np.random.normal(blue_mean, blue_std, size=20)

both_colours = np.sort(np.concatenate((red, blue))) # for later use...

Aqui está uma imagem desses grupos vermelho e azul novamente (para evitar que você tenha que rolar para cima):

Quando podemos ver a cor de cada ponto (ou seja, a qual grupo ele pertence), é muito fácil estimar a média e o desvio padrão para cada grupo. Apenas passamos os valores de vermelho e azul para as funções embutidas no NumPy. Por exemplo:

>>> np.mean(red)
2.802
>>> np.std(red)
0.871
>>> np.mean(blue)
6.932
>>> np.std(blue)
2.195

Mas e se não pudermos ver as cores dos pontos? Ou seja, em vez de vermelho ou azul, cada ponto foi colorido de roxo.

Para tentar recuperar os parâmetros de média e desvio padrão para os grupos vermelho e azul, podemos usar a maximização da expectativa.

Nossa primeira etapa ( etapa 1 acima) é adivinhar os valores dos parâmetros para a média e o desvio padrão de cada grupo. Não precisamos adivinhar de forma inteligente; podemos escolher qualquer número que quisermos:

# estimates for the mean
red_mean_guess = 1.1
blue_mean_guess = 9

# estimates for the standard deviation
red_std_guess = 2
blue_std_guess = 1.7

Essas estimativas de parâmetro produzem curvas de sino que se parecem com isto:

Essas são estimativas ruins. Ambos os meios (as linhas pontilhadas verticais) parecem distantes de qualquer tipo de "meio" para grupos de pontos sensíveis, por exemplo. Queremos melhorar essas estimativas.

A próxima etapa ( etapa 2 ) é calcular a probabilidade de cada ponto de dados aparecer sob as estimativas dos parâmetros atuais:

likelihood_of_red = stats.norm(red_mean_guess, red_std_guess).pdf(both_colours)
likelihood_of_blue = stats.norm(blue_mean_guess, blue_std_guess).pdf(both_colours)

Aqui, simplesmente colocamos cada ponto de dados na função de densidade de probabilidade para uma distribuição normal usando nossas estimativas atuais na média e no desvio padrão para vermelho e azul. Isso nos diz, por exemplo, que com nossas estimativas atuais, o ponto de dados em 1,761 tem muito mais probabilidade de ser vermelho (0,189) do que azul (0,00003).

Para cada ponto de dados, podemos transformar esses dois valores de verossimilhança em pesos ( etapa 3 ) para que somam 1 da seguinte forma:

likelihood_total = likelihood_of_red + likelihood_of_blue

red_weight = likelihood_of_red / likelihood_total
blue_weight = likelihood_of_blue / likelihood_total

Com nossas estimativas atuais e nossos pesos recém-calculados, agora podemos calcular novas estimativas para a média e o desvio padrão dos grupos vermelho e azul ( etapa 4 ).

Calculamos duas vezes a média e o desvio padrão usando todos os pontos de dados, mas com os pesos diferentes: uma vez para os pesos vermelhos e uma vez para os pesos azuis.

O ponto chave da intuição é que quanto maior o peso de uma cor em um ponto de dados, mais o ponto de dados influencia as próximas estimativas para os parâmetros dessa cor. Isso tem o efeito de "puxar" os parâmetros na direção certa.

def estimate_mean(data, weight):
    """
    For each data point, multiply the point by the probability it
    was drawn from the colour's distribution (its "weight").

    Divide by the total weight: essentially, we're finding where 
    the weight is centred among our data points.
    """
    return np.sum(data * weight) / np.sum(weight)

def estimate_std(data, weight, mean):
    """
    For each data point, multiply the point's squared difference
    from a mean value by the probability it was drawn from
    that distribution (its "weight").

    Divide by the total weight: essentially, we're finding where 
    the weight is centred among the values for the difference of
    each data point from the mean.

    This is the estimate of the variance, take the positive square
    root to find the standard deviation.
    """
    variance = np.sum(weight * (data - mean)**2) / np.sum(weight)
    return np.sqrt(variance)

# new estimates for standard deviation
blue_std_guess = estimate_std(both_colours, blue_weight, blue_mean_guess)
red_std_guess = estimate_std(both_colours, red_weight, red_mean_guess)

# new estimates for mean
red_mean_guess = estimate_mean(both_colours, red_weight)
blue_mean_guess = estimate_mean(both_colours, blue_weight)

Temos novas estimativas para os parâmetros. Para melhorá-los novamente, podemos voltar para a etapa 2 e repetir o processo. Fazemos isso até que as estimativas convirjam ou depois que algum número de iterações tiver sido executado ( etapa 5 ).

Para nossos dados, as primeiras cinco iterações desse processo são assim (as iterações recentes têm uma aparência mais forte):

Vemos que as médias já estão convergindo em alguns valores, e as formas das curvas (governadas pelo desvio padrão) também estão se tornando mais estáveis.

Se continuarmos por 20 iterações, terminaremos com o seguinte:

O processo EM convergiu para os seguintes valores, que se revelaram muito próximos dos valores reais (onde podemos ver as cores - sem variáveis ​​ocultas):

          | EM guess | Actual |  Delta
----------+----------+--------+-------
Red mean  |    2.910 |  2.802 |  0.108
Red std   |    0.854 |  0.871 | -0.017
Blue mean |    6.838 |  6.932 | -0.094
Blue std  |    2.227 |  2.195 |  0.032

No código acima, você deve ter notado que a nova estimativa para o desvio padrão foi calculada usando a estimativa da iteração anterior para a média. Em última análise, não importa se calcularmos primeiro um novo valor para a média, pois estamos apenas encontrando a variância (ponderada) dos valores em torno de algum ponto central. Ainda veremos as estimativas para os parâmetros convergirem.

EM é um algoritmo para maximizar uma função de verossimilhança quando algumas das variáveis ​​em seu modelo não são observadas (ou seja, quando você tem variáveis ​​latentes).

Você pode simplesmente perguntar, se estamos apenas tentando maximizar uma função, por que não usamos apenas o mecanismo existente para maximizar uma função. Bem, se você tentar maximizar isso pegando derivadas e definindo-as como zero, descobrirá que, em muitos casos, as condições de primeira ordem não têm solução. Há um problema da galinha e do ovo em que, para resolver os parâmetros do modelo, você precisa saber a distribuição dos dados não observados; mas a distribuição de seus dados não observados é uma função dos parâmetros do seu modelo.

EM tenta contornar isso adivinhando iterativamente uma distribuição para os dados não observados, em seguida, estimando os parâmetros do modelo, maximizando algo que é um limite inferior na função de verossimilhança real e repetindo até a convergência:

O algoritmo EM

Comece com suposições para os valores dos parâmetros do seu modelo

E-step: Para cada ponto de dados com valores ausentes, use a equação do modelo para resolver a distribuição dos dados ausentes, considerando sua estimativa atual dos parâmetros do modelo e dados observados (observe que você está resolvendo uma distribuição para cada um valor, não para o valor esperado). Agora que temos uma distribuição para cada valor ausente, podemos calcular a expectativa da função de verossimilhança com relação às variáveis ​​não observadas. Se nossa estimativa para o parâmetro do modelo estava correta, essa probabilidade esperada será a probabilidade real de nossos dados observados; se os parâmetros não estiverem corretos, será apenas um limite inferior.

Passo M: Agora que temos uma função de verossimilhança esperada sem variáveis ​​não observadas, maximize a função como faria no caso totalmente observado, para obter uma nova estimativa dos parâmetros do modelo.

Repita até a convergência.

Aqui está uma receita simples para entender o algoritmo de maximização da expectativa:

1- Leia este tutorial de EM por Do e Batzoglou.

2- Você pode ter pontos de interrogação em sua cabeça, dê uma olhada nas explicações nesta página de troca de pilha de matemática .

3- Olhe para este código que escrevi em Python que explica o exemplo no artigo do tutorial EM do item 1:

Aviso: o código pode ser confuso / abaixo do ideal, já que não sou um desenvolvedor Python. Mas isso faz o trabalho.

import numpy as np
import math

#### E-M Coin Toss Example as given in the EM tutorial paper by Do and Batzoglou* #### 

def get_mn_log_likelihood(obs,probs):
    """ Return the (log)likelihood of obs, given the probs"""
    # Multinomial Distribution Log PMF
    # ln (pdf)      =             multinomial coeff            *   product of probabilities
    # ln[f(x|n, p)] = [ln(n!) - (ln(x1!)+ln(x2!)+...+ln(xk!))] + [x1*ln(p1)+x2*ln(p2)+...+xk*ln(pk)]     

    multinomial_coeff_denom= 0
    prod_probs = 0
    for x in range(0,len(obs)): # loop through state counts in each observation
        multinomial_coeff_denom = multinomial_coeff_denom + math.log(math.factorial(obs[x]))
        prod_probs = prod_probs + obs[x]*math.log(probs[x])

    multinomial_coeff = math.log(math.factorial(sum(obs))) -  multinomial_coeff_denom
    likelihood = multinomial_coeff + prod_probs
    return likelihood

# 1st:  Coin B, {HTTTHHTHTH}, 5H,5T
# 2nd:  Coin A, {HHHHTHHHHH}, 9H,1T
# 3rd:  Coin A, {HTHHHHHTHH}, 8H,2T
# 4th:  Coin B, {HTHTTTHHTT}, 4H,6T
# 5th:  Coin A, {THHHTHHHTH}, 7H,3T
# so, from MLE: pA(heads) = 0.80 and pB(heads)=0.45

# represent the experiments
head_counts = np.array([5,9,8,4,7])
tail_counts = 10-head_counts
experiments = zip(head_counts,tail_counts)

# initialise the pA(heads) and pB(heads)
pA_heads = np.zeros(100); pA_heads[0] = 0.60
pB_heads = np.zeros(100); pB_heads[0] = 0.50

# E-M begins!
delta = 0.001  
j = 0 # iteration counter
improvement = float('inf')
while (improvement>delta):
    expectation_A = np.zeros((5,2), dtype=float) 
    expectation_B = np.zeros((5,2), dtype=float)
    for i in range(0,len(experiments)):
        e = experiments[i] # i'th experiment
        ll_A = get_mn_log_likelihood(e,np.array([pA_heads[j],1-pA_heads[j]])) # loglikelihood of e given coin A
        ll_B = get_mn_log_likelihood(e,np.array([pB_heads[j],1-pB_heads[j]])) # loglikelihood of e given coin B

        weightA = math.exp(ll_A) / ( math.exp(ll_A) + math.exp(ll_B) ) # corresponding weight of A proportional to likelihood of A 
        weightB = math.exp(ll_B) / ( math.exp(ll_A) + math.exp(ll_B) ) # corresponding weight of B proportional to likelihood of B                            

        expectation_A[i] = np.dot(weightA, e) 
        expectation_B[i] = np.dot(weightB, e)

    pA_heads[j+1] = sum(expectation_A)[0] / sum(sum(expectation_A)); 
    pB_heads[j+1] = sum(expectation_B)[0] / sum(sum(expectation_B)); 

    improvement = max( abs(np.array([pA_heads[j+1],pB_heads[j+1]]) - np.array([pA_heads[j],pB_heads[j]]) ))
    j = j+1

Tecnicamente, o termo "EM" é um pouco subespecificado, mas suponho que você se refere à técnica de análise de agrupamento de Modelagem de Mistura Gaussiana, que é uma instância do princípio EM geral.

Na verdade, a análise de cluster EM não é um classificador . Eu sei que algumas pessoas consideram o agrupamento como uma "classificação não supervisionada", mas na verdade a análise de agrupamento é algo bem diferente.

A principal diferença, e o grande mal-entendido de classificação que as pessoas sempre têm com a análise de cluster, é que: na análise de cluster, não existe uma "solução correta" . É um método de descoberta de conhecimento , na verdade serve para descobrir algo novo ! Isso torna a avaliação muito complicada. Muitas vezes é avaliada usando uma classificação conhecida como referência, mas isso nem sempre é apropriado: a classificação que você tem pode ou não refletir o que está nos dados.

Deixe-me dar um exemplo: você tem um grande conjunto de dados de clientes, incluindo dados de gênero. Um método que divide esse conjunto de dados em "masculino" e "feminino" é ideal quando você o compara com as classes existentes. Em uma forma de "previsão" de pensar isso é bom, pois para novos usuários agora você pode prever o sexo. Em uma maneira de pensar de "descoberta de conhecimento", isso é realmente ruim, porque você queria descobrir alguma estrutura nova nos dados. Um método que iria, por exemplo, dividir os dados em idosos e crianças, entretanto, teria uma pontuação tão pior quanto possível em relação à classe masculina / feminina. No entanto, esse seria um excelente resultado de agrupamento (se a idade não for fornecida).

Agora de volta ao EM. Basicamente, ele assume que seus dados são compostos de várias distribuições normais multivariadas (observe que esta é uma suposição muito forte, em particular quando você fixa o número de clusters!). Em seguida, ele tenta encontrar um modelo ideal local para isso, melhorando alternadamente o modelo e a atribuição de objeto ao modelo .

Para melhores resultados em um contexto de classificação, escolha o número de clusters maiores do que o número de classes, ou mesmo aplicar o clustering para solteiros classes somente (para descobrir se há alguma estrutura dentro da classe!).

Digamos que você queira treinar um classificador para distinguir "carros", "bicicletas" e "caminhões". É pouco útil assumir que os dados consistem em exatamente 3 distribuições normais. No entanto, você pode presumir que há mais de um tipo de carro (e caminhões e bicicletas). Então, em vez de treinar um classificador para essas três classes, você agrupa carros, caminhões e bicicletas em 10 grupos cada (ou talvez 10 carros, 3 caminhões e 3 bicicletas, o que for), em seguida, treina um classificador para distinguir essas 30 classes e, em seguida, mesclar o resultado da classe de volta às classes originais. Você também pode descobrir que existe um cluster que é particularmente difícil de classificar, por exemplo, Trikes. Eles são um pouco carros e um pouco bicicletas. Ou caminhões de entrega, que são mais carros superdimensionados do que caminhões.

Se outras respostas forem boas, tentarei fornecer outra perspectiva e abordar a parte intuitiva da pergunta.

O algoritmo EM (Expectation-Maximization) é uma variante de uma classe de algoritmos iterativos usando dualidade

Trecho (ênfase minha):

Em matemática, uma dualidade, de um modo geral, traduz conceitos, teoremas ou estruturas matemáticas em outros conceitos, teoremas ou estruturas, de maneira um a um, muitas vezes (mas nem sempre) por meio de uma operação de involução: se o dual de A é B, então o dual de B é A. Tais involuções às vezes têm pontos fixos , de modo que o dual de A é o próprio A

Normalmente, um B dual de um objeto A está relacionado a A de alguma forma que preserva alguma simetria ou compatibilidade . Por exemplo AB = const

Exemplos de algoritmos iterativos, empregando dualidade (no sentido anterior) são:

1. Algoritmo euclidiano para o maior divisor comum e suas variantes
2. Algoritmo e variantes Gram – Schmidt Vector Base
3. Média aritmética - Desigualdade média geométrica e suas variantes
4. Algoritmo de maximização de expectativa e suas variantes (veja também aqui para uma visão geométrica de informações )
5. (.. outros algoritmos semelhantes ..)

De forma semelhante, o algoritmo EM também pode ser visto como duas etapas de maximização dupla :

.. [EM] é visto como maximizando uma função conjunta dos parâmetros e da distribuição sobre as variáveis ​​não observadas. O E-step maximiza esta função com respeito à distribuição sobre as variáveis ​​não observadas; o passo M em relação aos parâmetros ..

Em um algoritmo iterativo usando dualidade, há a suposição explícita (ou implícita) de um ponto de convergência de equilíbrio (ou fixo) (para EM, isso é provado usando a desigualdade de Jensen)

Portanto, o esboço de tais algoritmos é:

1. Etapa semelhante a E: Encontre a melhor solução x em relação a y dado ser mantido constante.
2. Etapa semelhante a M (dual): Encontre a melhor solução y em relação ax (conforme calculado na etapa anterior) sendo mantida constante.
3. Critério de Rescisão / Etapa de Convergência: Repita as etapas 1, 2 com os valores atualizados de x , y até a convergência (ou o número especificado de iterações seja alcançado)

Observe que quando tal algoritmo converge para um ótimo (global), ele encontrou uma configuração que é a melhor em ambos os sentidos (isto é, tanto no domínio / parâmetros x quanto no domínio / parâmetros y ). No entanto, o algoritmo pode apenas encontrar um ótimo local e não o ótimo global .

eu diria que esta é a descrição intuitiva do esboço do algoritmo

Para os argumentos estatísticos e aplicações, outras respostas deram boas explicações (verifique também as referências nesta resposta)

A resposta aceita faz referência ao Papel EM da Chuong , que faz um trabalho decente ao explicar o EM. Há também um vídeo do youtube que explica o jornal com mais detalhes.

Para recapitular, aqui está o cenário:

1st:  {H,T,T,T,H,H,T,H,T,H} 5 Heads, 5 Tails; Did coin A or B generate me?
2nd:  {H,H,H,H,T,H,H,H,H,H} 9 Heads, 1 Tails
3rd:  {H,T,H,H,H,H,H,T,H,H} 8 Heads, 2 Tails
4th:  {H,T,H,T,T,T,H,H,T,T} 4 Heads, 6 Tails
5th:  {T,H,H,H,T,H,H,H,T,H} 7 Heads, 3 Tails

Two possible coins, A & B are used to generate these distributions.
A & B have an unknown parameter: their bias towards heads.

We don't know the biases, but we can simply start with a guess: A=60% heads, B=50% heads.

No caso da pergunta da primeira tentativa, intuitivamente pensaríamos que B a gerou, já que a proporção de caras corresponde muito bem ao viés de B ... mas esse valor foi apenas um palpite, então não podemos ter certeza.

Com isso em mente, gosto de pensar na solução EM assim:

• Cada tentativa de virar pode 'votar' na moeda que mais gosta
  • Isso se baseia em quão bem cada moeda se ajusta à sua distribuição
  • OU, do ponto de vista da moeda, há uma grande expectativa de ver essa tentativa em relação à outra moeda (com base nas probabilidades logarítmicas ).
• Dependendo de quanto cada tentativa gosta de cada moeda, pode atualizar a estimativa do parâmetro da moeda (viés).
  • Quanto mais uma tentativa gosta de uma moeda, mais ela consegue atualizar a tendência da moeda para refletir a sua!
  • Essencialmente, os vieses da moeda são atualizados pela combinação dessas atualizações ponderadas em todos os testes, um processo denominado ( maximazation ), que se refere a tentar obter as melhores estimativas para o viés de cada moeda, dado um conjunto de testes.

Isso pode ser uma simplificação exagerada (ou até mesmo fundamentalmente errado em alguns níveis), mas espero que isso ajude em um nível intuitivo!

EM é usado para maximizar a probabilidade de um modelo Q com variáveis ​​latentes Z.

É uma otimização iterativa.

theta <- initial guess for hidden parameters
while not converged:
    #e-step
    Q(theta'|theta) = E[log L(theta|Z)]
    #m-step
    theta <- argmax_theta' Q(theta'|theta)

e-step: dada a estimativa atual de Z, calcule a função de log-verossimilhança esperada

m-step: encontre theta que maximiza este Q

Exemplo de GMM:

e-step: estimar atribuições de rótulo para cada ponto de dados, dada a estimativa do parâmetro gmm atual

passo m: maximizar um novo theta de acordo com as novas atribuições de rótulo

K-means também é um algoritmo EM e há muitas animações explicando no K-means.

Usando o mesmo artigo de Do e Batzoglou citado na resposta de Zhubarb, implementei o EM para esse problema em Java . Os comentários à sua resposta mostram que o algoritmo fica preso em um ótimo local, o que também ocorre com minha implementação se os parâmetros thetaA e thetaB forem iguais.

Abaixo está a saída padrão do meu código, mostrando a convergência dos parâmetros.

thetaA = 0.71301, thetaB = 0.58134
thetaA = 0.74529, thetaB = 0.56926
thetaA = 0.76810, thetaB = 0.54954
thetaA = 0.78316, thetaB = 0.53462
thetaA = 0.79106, thetaB = 0.52628
thetaA = 0.79453, thetaB = 0.52239
thetaA = 0.79593, thetaB = 0.52073
thetaA = 0.79647, thetaB = 0.52005
thetaA = 0.79667, thetaB = 0.51977
thetaA = 0.79674, thetaB = 0.51966
thetaA = 0.79677, thetaB = 0.51961
thetaA = 0.79678, thetaB = 0.51960
thetaA = 0.79679, thetaB = 0.51959
Final result:
thetaA = 0.79678, thetaB = 0.51960

Abaixo está minha implementação Java do EM para resolver o problema em (Do e Batzoglou, 2008). A parte central da implementação é o loop para executar EM até que os parâmetros convirjam.

private Parameters _parameters;

public Parameters run()
{
    while (true)
    {
        expectation();

        Parameters estimatedParameters = maximization();

        if (_parameters.converged(estimatedParameters)) {
            break;
        }

        _parameters = estimatedParameters;
    }

    return _parameters;
}

Abaixo está o código completo.

import java.util.*;

/*****************************************************************************
This class encapsulates the parameters of the problem. For this problem posed
in the article by (Do and Batzoglou, 2008), the parameters are thetaA and
thetaB, the probability of a coin coming up heads for the two coins A and B,
respectively.
*****************************************************************************/
class Parameters
{
    double _thetaA = 0.0; // Probability of heads for coin A.
    double _thetaB = 0.0; // Probability of heads for coin B.

    double _delta = 0.00001;

    public Parameters(double thetaA, double thetaB)
    {
        _thetaA = thetaA;
        _thetaB = thetaB;
    }

    /*************************************************************************
    Returns true if this parameter is close enough to another parameter
    (typically the estimated parameter coming from the maximization step).
    *************************************************************************/
    public boolean converged(Parameters other)
    {
        if (Math.abs(_thetaA - other._thetaA) < _delta &&
            Math.abs(_thetaB - other._thetaB) < _delta)
        {
            return true;
        }

        return false;
    }

    public double getThetaA()
    {
        return _thetaA;
    }

    public double getThetaB()
    {
        return _thetaB;
    }

    public String toString()
    {
        return String.format("thetaA = %.5f, thetaB = %.5f", _thetaA, _thetaB);
    }

}


/*****************************************************************************
This class encapsulates an observation, that is the number of heads
and tails in a trial. The observation can be either (1) one of the
experimental observations, or (2) an estimated observation resulting from
the expectation step.
*****************************************************************************/
class Observation
{
    double _numHeads = 0;
    double _numTails = 0;

    public Observation(String s)
    {
        for (int i = 0; i < s.length(); i++)
        {
            char c = s.charAt(i);

            if (c == 'H')
            {
                _numHeads++;
            }
            else if (c == 'T')
            {
                _numTails++;
            }
            else
            {
                throw new RuntimeException("Unknown character: " + c);
            }
        }
    }

    public Observation(double numHeads, double numTails)
    {
        _numHeads = numHeads;
        _numTails = numTails;
    }

    public double getNumHeads()
    {
        return _numHeads;
    }

    public double getNumTails()
    {
        return _numTails;
    }

    public String toString()
    {
        return String.format("heads: %.1f, tails: %.1f", _numHeads, _numTails);
    }

}

/*****************************************************************************
This class runs expectation-maximization for the problem posed by the article
from (Do and Batzoglou, 2008).
*****************************************************************************/
public class EM
{
    // Current estimated parameters.
    private Parameters _parameters;

    // Observations from the trials. These observations are set once.
    private final List<Observation> _observations;

    // Estimated observations per coin. These observations are the output
    // of the expectation step.
    private List<Observation> _expectedObservationsForCoinA;
    private List<Observation> _expectedObservationsForCoinB;

    private static java.io.PrintStream o = System.out;

    /*************************************************************************
    Principal constructor.
    @param observations The observations from the trial.
    @param parameters The initial guessed parameters.
    *************************************************************************/
    public EM(List<Observation> observations, Parameters parameters)
    {
        _observations = observations;
        _parameters = parameters;
    }

    /*************************************************************************
    Run EM until parameters converge.
    *************************************************************************/
    public Parameters run()
    {

        while (true)
        {
            expectation();

            Parameters estimatedParameters = maximization();

            o.printf("%s\n", estimatedParameters);

            if (_parameters.converged(estimatedParameters)) {
                break;
            }

            _parameters = estimatedParameters;
        }

        return _parameters;

    }

    /*************************************************************************
    Given the observations and current estimated parameters, compute new
    estimated completions (distribution over the classes) and observations.
    *************************************************************************/
    private void expectation()
    {

        _expectedObservationsForCoinA = new ArrayList<Observation>();
        _expectedObservationsForCoinB = new ArrayList<Observation>();

        for (Observation observation : _observations)
        {
            int numHeads = (int)observation.getNumHeads();
            int numTails = (int)observation.getNumTails();

            double probabilityOfObservationForCoinA=
                binomialProbability(10, numHeads, _parameters.getThetaA());

            double probabilityOfObservationForCoinB=
                binomialProbability(10, numHeads, _parameters.getThetaB());

            double normalizer = probabilityOfObservationForCoinA +
                                probabilityOfObservationForCoinB;

            // Compute the completions for coin A and B (i.e. the probability
            // distribution of the two classes, summed to 1.0).

            double completionCoinA = probabilityOfObservationForCoinA /
                                     normalizer;
            double completionCoinB = probabilityOfObservationForCoinB /
                                     normalizer;

            // Compute new expected observations for the two coins.

            Observation expectedObservationForCoinA =
                new Observation(numHeads * completionCoinA,
                                numTails * completionCoinA);

            Observation expectedObservationForCoinB =
                new Observation(numHeads * completionCoinB,
                                numTails * completionCoinB);

            _expectedObservationsForCoinA.add(expectedObservationForCoinA);
            _expectedObservationsForCoinB.add(expectedObservationForCoinB);
        }
    }

    /*************************************************************************
    Given new estimated observations, compute new estimated parameters.
    *************************************************************************/
    private Parameters maximization()
    {

        double sumCoinAHeads = 0.0;
        double sumCoinATails = 0.0;
        double sumCoinBHeads = 0.0;
        double sumCoinBTails = 0.0;

        for (Observation observation : _expectedObservationsForCoinA)
        {
            sumCoinAHeads += observation.getNumHeads();
            sumCoinATails += observation.getNumTails();
        }

        for (Observation observation : _expectedObservationsForCoinB)
        {
            sumCoinBHeads += observation.getNumHeads();
            sumCoinBTails += observation.getNumTails();
        }

        return new Parameters(sumCoinAHeads / (sumCoinAHeads + sumCoinATails),
                              sumCoinBHeads / (sumCoinBHeads + sumCoinBTails));

        //o.printf("parameters: %s\n", _parameters);

    }

    /*************************************************************************
    Since the coin-toss experiment posed in this article is a Bernoulli trial,
    use a binomial probability Pr(X=k; n,p) = (n choose k) * p^k * (1-p)^(n-k).
    *************************************************************************/
    private static double binomialProbability(int n, int k, double p)
    {
        double q = 1.0 - p;
        return nChooseK(n, k) * Math.pow(p, k) * Math.pow(q, n-k);
    }

    private static long nChooseK(int n, int k)
    {
        long numerator = 1;

        for (int i = 0; i < k; i++)
        {
            numerator = numerator * n;
            n--;
        }

        long denominator = factorial(k);

        return (long)(numerator / denominator);
    }

    private static long factorial(int n)
    {
        long result = 1;
        for (; n >0; n--)
        {
            result = result * n;
        }

        return result;
    }

    /*************************************************************************
    Entry point into the program.
    *************************************************************************/
    public static void main(String argv[])
    {
        // Create the observations and initial parameter guess
        // from the (Do and Batzoglou, 2008) article.

        List<Observation> observations = new ArrayList<Observation>();
        observations.add(new Observation("HTTTHHTHTH"));
        observations.add(new Observation("HHHHTHHHHH"));
        observations.add(new Observation("HTHHHHHTHH"));
        observations.add(new Observation("HTHTTTHHTT"));
        observations.add(new Observation("THHHTHHHTH"));

        Parameters initialParameters = new Parameters(0.6, 0.5);

        EM em = new EM(observations, initialParameters);

        Parameters finalParameters = em.run();

        o.printf("Final result:\n%s\n", finalParameters);
    }
}