rlab/example/emLearningEx.js

Javascript scientific library like R

// �ѦҤ��m�GNumerical example to understand Expectation-Maximization -- http://ai.stanford.edu/~chuongdo/papers/em_tutorial.pdf
// What is the expectation maximization algorithm? (PDF) -- http://stats.stackexchange.com/questions/72774/numerical-example-to-understand-expectation-maximization

var log=console.log;
var R = require("../rlab");

function logp(n) {
	return R.steps(1,n,1).log().sum();
}

// �Ǧ^�h�������� log �ȡI log( (n!)/(x1!x2!...xk!) p1^x1 p2^x2 ... pk^xk )
// = [log(n)+...+log(1)]-[log(x1)...]+....+x1*log(p1)+...+xk*log(pk)
function xplog(x, p) {
  var n = x.sum();
  var r=logp(n);
  for (var i in x) r-=logp(x[i]);
  return r+x.dot(p.log());
}

function EM() {
// 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
  var e = [ [5,5], [9,1], [8,2], [4,6], [7,3] ];
  var pA = [0.6,0.4], pB = [0.5,0.5];
  var gen=0, delta=9.9999;
  for (var gen=0; gen<1000 && delta > 0.001; gen++) {
    log("pA=%s pB=%s delta=%d", pA, pB, delta.toFixed(4));
    var sumA=[0,0], sumB=[0,0];
    for (var i in e) {
//      log("e[i]=%s", e[i]);
      var lA = xplog(e[i], pA);
      var lB = xplog(e[i], pB);
//      log("lA=%s lB=%s", lA, lB);
      var a  = lA.exp(), b = lB.exp();
//      log("a=%s b=%s", a, b);
      var wA = a/(a+b), wB=b/(a+b);
//      log("wA=%s wB=%s", wA, wB);
      var eA = wA.mul(e[i]);
      var eB = wB.mul(e[i]);
//      log("eA=%s eB=%s", eA, eB);
      sumA   = sumA.add(eA);
      sumB   = sumB.add(eB);
//      log("sumA=%s sumB=%s", sumA, sumB);
    }
//    log("sumA=%s sumB=%s", sumA, sumB);
    var npA = sumA.mul(1.0/sumA.sum());
    var npB = sumB.mul(1.0/sumB.sum());
    var dA  = npA.sub(pA);
    var dB  = npB.sub(pB);
//    log("dA=%s dB=%s", dA, dB);
    var delta = [dA, dB].max();
    pA = npA; pB=npB;
  }
}

EM();