Максимізація очікувань (ЕМ) - це свого роду ймовірнісний метод класифікації даних. Будь ласка, виправте мене, якщо я помиляюся, якщо це не класифікатор.

Що таке інтуїтивне пояснення цієї методики ЕМ? Що expectationтут і що буває maximized?

Примітка: код за цією відповіддю можна знайти тут .

Припустимо, у нас є деякі дані, відібрані з двох різних груп, червоної та синьої:

Тут ми можемо побачити, яка точка даних належить до групи червоного чи синього. Це дозволяє легко знайти параметри, що характеризують кожну групу. Наприклад, середнє значення червоної групи становить близько 3, середнє значення синьої групи - близько 7 (і ми могли б знайти точні засоби, якщо б хотіли).

Це, загалом кажучи, відоме як максимальна оцінка ймовірності . З огляду на деякі дані, ми обчислюємо значення параметра (або параметрів), що найкраще пояснює ці дані.

А тепер уявіть, що ми не можемо побачити, яке значення було відібране з якої групи. Для нас все виглядає фіолетово:

Тут ми знаємо, що є дві групи цінностей, але ми не знаємо, до якої групи належить якесь значення.

Чи можемо ми все-таки оцінити засоби для червоної та синьої групи, які найкраще відповідають цим даним?

Так, часто ми можемо! Максималізація очікування дає нам спосіб це зробити. Сама загальна ідея алгоритму така:

1. Почніть з початкової оцінки того, яким може бути кожен параметр.
2. Обчисліть ймовірність того, що кожен параметр виробляє точку даних.
3. Обчисліть ваги для кожної точки даних із зазначенням того, вона більше червоного або синього кольору, виходячи з ймовірності її отримання параметром. Поєднайте ваги з даними ( очікування ).
4. Обчисліть кращу оцінку параметрів, використовуючи дані з коригуванням ваги ( максимізація ).
5. Повторюйте кроки 2 - 4, поки оцінка параметра не сходиться (процес припиняє видавати іншу оцінку).

Ці кроки потребують додаткового пояснення, тому я перегляну описану вище проблему.

Приклад: оцінка середнього та стандартного відхилень

Я буду використовувати Python в цьому прикладі, але код повинен бути досить легким для розуміння, якщо ви не знайомі з цією мовою.

Припустимо, у нас є дві групи, червона та синя, зі значеннями, розподіленими як на зображенні вище. Зокрема, кожна група містить значення, отримане з звичайного розподілу з такими параметрами:

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...

Ось зображення цих червоно-синіх груп знову (щоб уберегти вас від необхідності прокручування вгору):

Коли ми можемо побачити колір кожної точки (тобто до якої групи вона належить), дуже легко оцінити середнє та стандартне відхилення для кожної групи. Ми просто передаємо червоні та сині значення вбудованим функціям у NumPy. Наприклад:

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

Але що робити, якщо ми не бачимо кольори точок? Тобто замість червоного чи синього кольори кожна точка була пофарбована у фіолетовий колір.

Щоб спробувати відновити середні та стандартні параметри відхилення для червоної та синьої груп, ми можемо використовувати Максималізацію очікування.

Наш перший крок ( етап 1 вище) - вгадати значення параметрів для середньої групи та середнього відхилення кожної групи. Нам не треба розумно здогадуватися; ми можемо вибрати будь-які цифри, які нам подобаються:

# 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

Ці оцінки параметрів створюють криві дзвіночка, які виглядають так:

Це погані оцінки. Обидва засоби (вертикальні пунктирні лінії) виглядають далеко не будь-якими «серединами» для розумних груп точок, наприклад. Ми хочемо покращити ці оцінки.

Наступним кроком ( крок 2 ) є обчислення ймовірності появи кожної точки даних під поточними здогадами параметрів:

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)

Тут ми просто помістили кожну точку даних у функцію щільності ймовірності для нормального розподілу, використовуючи наші поточні здогадки на середнє та стандартне відхилення для червоного та синього. Це говорить нам, наприклад, що за нашими теперішніми здогадами точка даних у 1.761 набагато частіше буде червоною (0,189), ніж синьою (0,00003).

Для кожної точки даних ми можемо перетворити ці два значення правдоподібності у ваги ( крок 3 ) так, щоб вони дорівнювали 1 таким чином:

likelihood_total = likelihood_of_red + likelihood_of_blue

red_weight = likelihood_of_red / likelihood_total
blue_weight = likelihood_of_blue / likelihood_total

За нашими поточними оцінками та нашими щойно обчисленими вагами ми тепер можемо обчислити нові оцінки середнього та стандартного відхилення червоної та синьої груп ( крок 4 ).

Ми двічі обчислюємо середнє та стандартне відхилення, використовуючи всі точки даних, але з різною вагою: один раз для червоних ваг і один раз для синіх ваг.

Ключовий біт інтуїції полягає в тому, що чим більша вага кольору в точці даних, тим більше точка даних впливає на наступні оцінки параметрів цього кольору. Це спричиняє «потягування» параметрів у правильному напрямку.

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)

У нас є нові оцінки параметрів. Щоб знову їх покращити, ми можемо перейти до кроку 2 і повторити процес. Ми робимо це, поки оцінки не збігаються або після того, як буде виконано деяку кількість ітерацій ( крок 5 ).

За нашими даними, перші п’ять ітерацій цього процесу виглядають приблизно так (останні ітерації мають сильніший вигляд):

Ми бачимо, що засоби вже сходяться на деяких значеннях, а форми кривих (керуються стандартним відхиленням) також стають більш стійкими.

Якщо ми продовжимо 20 ітерацій, ми закінчимось наступним:

Процес ЕМ конвергується до таких значень, які виявляються дуже близькими до фактичних значень (де ми можемо бачити кольори - відсутність прихованих змінних):

          | 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

У наведеному вище коді, можливо, ви помітили, що нову оцінку для стандартного відхилення було обчислено, використовуючи оцінку попередньої ітерації для середнього. Зрештою, не має значення, чи спочатку обчислимо нове значення для середнього значення, оскільки ми лише знаходимо (зважену) дисперсію значень навколо якоїсь центральної точки. Ми ще побачимо оцінки зближення параметрів.

EM - це алгоритм для максимізації функції ймовірності, коли деякі зі змінних у вашій моделі не спостерігаються (тобто коли у вас є приховані змінні).

Ви можете справедливо запитати, якщо ми просто намагаємося максимально функціонувати, чому б ми просто не використали існуючу техніку для максимізації функції. Що ж, якщо ви спробуєте це максимально використати, беручи похідні та встановлюючи їх до нуля, ви виявите, що в багатьох випадках умови першого порядку не мають рішення. У цьому є проблема курки-яєць, щоб вирішити для вашої моделі параметри потрібно знати розподіл ваших незабезпечених даних; але розподіл ваших незабезпечених даних є функцією параметрів вашої моделі.

ЕМ намагається обійти це, ітеративно відгадуючи розподіл для непомічених даних, потім оцінюючи параметри моделі, максимізуючи щось, що є нижньою межею фактичної функції ймовірності, і повторюючи до конвергенції:

Алгоритм ЕМ

Почніть з здогадки щодо значень параметрів моделі

Е-крок: для кожної точки даних, у якої відсутні значення, використовуйте рівняння моделі для вирішення розподілу відсутніх даних, враховуючи поточну здогадку параметрів моделі та дані, що спостерігаються (зауважте, що ви вирішуєте розподіл для кожного відсутнього значення, а не для очікуваного значення). Тепер, коли у нас є розподіл на кожне відсутнє значення, ми можемо обчислити очікування функції ймовірності відносно непомічених змінних. Якщо наші здогадки щодо параметра моделі були правильними, ця очікувана ймовірність буде фактичною ймовірністю наших спостережуваних даних; якщо параметри були невірними, це буде просто нижня межа.

М-крок: Тепер, коли у нас є функція очікуваної ймовірності, в якій відсутні незаміщені змінні, максимізуйте функцію, як це було б у випадку, що спостерігається, щоб отримати нову оцінку параметрів вашої моделі.

Повторюйте до зближення.

Ось прямий рецепт для розуміння алгоритму максимізації очікування:

1- Прочитайте цей підручник з Е. До та Бацоглу.

2- У вас можуть стояти запитання в голові, подивіться пояснення на цій сторінці обміну стеками математики .

3- Подивіться на цей код, який я написав на Python, що пояснює приклад у документі з підручника EM статті 1:

Попередження: Код може бути брудним / субоптимальним, оскільки я не розробник Python. Але це робить свою роботу.

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

У технічному плані термін "ЕМ" трохи не визначений, але я припускаю, що ви посилаєтесь на методику кластерного аналізу моделей Гаусса, що є прикладом загального принципу ЕМ.

Власне, кластерний аналіз EM не є класифікатором . Я знаю, що деякі люди вважають кластеризацію "непідконтрольною класифікацією", але насправді кластерний аналіз - це щось зовсім інше.

Ключова відмінність, і велика класифікація непорозумінь, яку люди завжди мають з кластерним аналізом, полягає в тому, що: в аналізах кластерів немає "правильного рішення" . Це метод виявлення знань , він насправді призначений для пошуку чогось нового ! Це робить оцінку дуже хитрою. Його часто оцінюють, використовуючи відому класифікацію як еталонну, але це не завжди доречно: класифікація, яку ви маєте, може відображати або не відображати дані, що містяться в даних.

Дозвольте навести приклад: у вас є великий набір клієнтів, включаючи гендерні дані. Метод, який розбиває ці дані, встановлені на "чоловічі" та "жіночі", є оптимальним при порівнянні їх із існуючими класами. У "передбачуваному" способі мислення це добре, оскільки для нових користувачів тепер можна передбачити їхню стать. У способі мислення "відкриття знань" це насправді погано, тому що ви хотіли відкрити для себе якусь нову структуру . Метод, який б, наприклад, розділив дані на людей похилого віку та дітей, однак став би настільки ж гіршим, як це може бути досягнуто стосовно чоловічого / жіночого класу. Однак це було б чудовим результатом кластеризації (якщо б не було вказано вік).

Тепер повернемось до ЕМ. По суті, це передбачає, що ваші дані складаються з декількох багатоваріантних нормальних розподілів (зауважте, що це дуже вагоме припущення, зокрема, коли ви фіксуєте кількість кластерів!). Потім він намагається знайти локальну оптимальну модель для цього шляхом почергового вдосконалення моделі та призначення об'єкта моделі .

Для досягнення найкращих результатів в контексті класифікації, виберіть кількість кластерів великих , ніж число класів, або навіть застосувати кластеризацию для окремих класів тільки (щоб з'ясувати , чи є якась - то структура в класі!).

Скажіть, ви хочете навчити класифікатора, щоб розказувати "машини", "велосипеди" та "вантажівки". Мало користі, якщо припустити, що дані складаються з 3-х звичайних розподілів. Однак ви можете припустити, що існує більше одного типу автомобілів (і вантажівок, і велосипедів). Тож замість того, щоб навчати класифікатора для цих трьох класів, ви згрупуєте машини, вантажівки та велосипеди в 10 кластерів (або, можливо, 10 машин, 3 вантажівки та 3 велосипеди, що завгодно), а потім навчите класифікатора, щоб розділити ці 30 класів, а потім об'єднати результат класу назад до початкових класів. Ви також можете виявити, що є один кластер, який важко класифікувати, наприклад, Трікс. Вони дещо машини, а трохи велосипеди. Або вантажні автомобілі, які більше нагадують негабаритні машини, ніж вантажні автомобілі.

Інші відповіді добрі, я спробую надати іншу точку зору і вирішити інтуїтивну частину питання.

Алгоритм EM (Expectation-Maximization) - це варіант класу ітеративних алгоритмів з використанням подвійності

Витяг (моє наголос):

У математиці подвійність, взагалі кажучи, перекладає поняття, теореми чи математичні структури в інші поняття, теореми чи структури, один на один, часто (але не завжди) за допомогою операції інволюції: якщо подвійний A є B, тоді дуал B є A. Такі інволюції іноді мають фіксовані точки , так що подвійний A є A сам

Зазвичай подвійний B об'єкта A пов'язаний з A певним чином, що зберігає певну симетрію або сумісність . Наприклад AB = const

Прикладами ітеративних алгоритмів, що використовують подвійність (у попередньому розумінні), є:

1. Евклідовий алгоритм для найбільшого спільного дільника та його варіанти
2. Алгоритм та варіанти векторного грам-Шмідта та варіанти
3. Середнє арифметичне - середня геометрична нерівність та її варіанти
4. Алгоритм очікування-Максимізація та його варіанти (див. Тут також інформаційно-геометричний вигляд )
5. (.. інші подібні алгоритми ..)

Аналогічним чином алгоритм ЕМ також можна розглядати як два подвійні етапи максимізації :

.. [ЕМ] розглядається як максимізація спільної функції параметрів та розподілу по незабезпеченим змінним. Е-крок максимізує цю функцію щодо розподілу по непоміченим змінним; крок М по відношенню до параметрів ..

В ітеративному алгоритмі з використанням подвійності існує явне (або неявне) припущення про рівноважну (або фіксовану) точку зближення (для ЕМ це доводиться, використовуючи нерівність Дженсена)

Отже, контур таких алгоритмів такий:

1. Е-подібний крок: Знайдіть найкраще рішення x стосовно того, що y утримується постійним.
2. М-подібний крок (подвійний): знайти найкраще рішення y щодо x (як було обчислено на попередньому кроці), яке буде постійним.
3. Критерій припинення / крок конвергенції: Повторюйте кроки 1, 2 із оновленими значеннями x , y до досягнення конвергенції (або визначеної кількості ітерацій)

Зауважте, що коли такий алгоритм конвергується до (глобального) оптимуму, він знайшов конфігурацію, яка найкраща в обох сенсах (тобто в x домені / параметрах та y домені / параметрах). Однак алгоритм може просто знайти локальний оптимум, а не глобальний оптимум.

Я б сказав, що це інтуїтивний опис контуру алгоритму

Щодо статистичних аргументів та додатків, інші відповіді дали хороші пояснення (перевірте також посилання в цій відповіді)

Прийнята відповідь посилається на документ Chuong EM , який робить гідну роботу з поясненням EM. Також є відео на YouTube, яке детальніше пояснює статтю.

Для резюме, ось такий сценарій:

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.

Що стосується питання першого судового розгляду, інтуїтивно ми могли б подумати, що B створив це, оскільки частка головок дуже добре відповідає упередженню Б. ... але це значення було лише здогадкою, тому ми не можемо бути впевнені.

Зважаючи на це, мені подобається думати про рішення ЕМ так:

• Кожен пробний фліп отримує "голосування", за яку монету вона найбільше подобається
  • Це ґрунтується на тому, наскільки кожна монета відповідає її розповсюдженню
  • АБО, з точки зору монети, великі сподівання побачити цей випробування відносно іншої монети (виходячи з ймовірності журналу ).
• Залежно від того, наскільки кожна проба любить кожну монету, вона може оновити здогадки про параметр цієї монети (зміщення).
  • Чим більше пробі подобається монета, тим більше отримує оновлення упередженості монети, щоб відобразити її власну!
  • По суті, ухили монети оновлюються шляхом комбінування цих зважених оновлень у всіх випробуваннях, процес, який називається ( максимізація ), який стосується намагання отримати найкращі здогадки щодо ухилу кожної монети з набором випробувань.

Це може бути надмірне спрощення (або навіть принципово неправильне на деяких рівнях), але я сподіваюся, що це допомагає на інтуїтивному рівні!

ЕМ використовується для максимізації ймовірності моделі Q з прихованими змінними Z.

Це ітеративна оптимізація.

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-крок: задана поточна оцінка Z обчислює очікувану функцію вірогідності

m-крок: знайдіть тету, яка максимізує цей Q

Приклад GMM:

e-крок: присвоєння міток оцінки для кожної точки даних з урахуванням поточної оцінки параметра gmm

m-крок: максимізуйте нову тету з урахуванням нових призначень етикетки

K-засоби - це також алгоритм ЕМ, і на K-засобах є багато пояснювальних анімацій.

Використовуючи ту саму статтю До і Бацоглоу, цитовану у відповіді Жубарба, я реалізував ЕМ для цієї проблеми на Java . У коментарях до його відповіді видно, що алгоритм зациклюється на локальному оптимальному, що також відбувається з моєю реалізацією, якщо параметри thetaA і thetaB однакові.

Нижче наведено стандартний вихід мого коду, що показує збіжність параметрів.

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

Нижче наведено мою реалізацію ЕМ на Java для вирішення проблеми в (Do and Batzoglou, 2008). Основна частина реалізації - це цикл для запуску ЕМ до тих пір, поки параметри не зберуться.

private Parameters _parameters;

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

        Parameters estimatedParameters = maximization();

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

        _parameters = estimatedParameters;
    }

    return _parameters;
}

Нижче наведено весь код.

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);
    }
}