* One-dimensional Gaussian Mixture Model (Expectation Maximization)

 *

 * Keith O'Hara <kohara@swarthmore.edu>

 * ported to p5.js in March 2024

 *

 * click to add 1-D data points

 * 'r': reset

 * 'c' classify points with cluters

 * ' ': take a EM step

 * 

 */

​

const NUM_MODELS = 3;

​

let mean = [];

let variance;

let mixture = []; // clusters weights

let x = [];  // data points

let z = [];  // likelihoood that point i is in cluster j

let classify = false;

​

function setup() {

  createCanvas(800, 200);

  noStroke();

  colorMode(HSB, NUM_MODELS, 1, 1, .25);

  reset();

}

​

function mouseDragged() {

  x.push(mouseX);

  reset();

}

​

function keyPressed() {

​

  if (key == 'c') {

    classify = !classify;

  }

  else  if (key == 'r') {

    reset();

  }

  else if (key == 'i'){

    for (let m = 0; m < NUM_MODELS; m++) {

      print("mixture: " + mixture[m] + " mean: " + mean[m] + " variane: " + variance[m]);

    }

  }

  else {

    EM();

  }

}

​

function draw() {

  background(0);

  text(NUM_MODELS + " clusters", 25, 25);

  for (let i = 0; i < x.length; i++) {  

    if (classify) {

      for (let j = 0; j < NUM_MODELS; j++) {

        fill(j, 1, 1, z[i][j]); 

        ellipse(x[i], height-20, 10, 10);

      }

    }

    else {

      fill(0, 0, 1);

      ellipse(x[i], height - 20, 10, 10);

    }

  }

​

  for (let i = 0; i < width; i++) {

    for (let m = 0; m < NUM_MODELS; m++) {

      fill(m, 1, 1);

      ellipse(i, height-30 -70*height*mixture[m]*evalGaussianPDF(i, mean[m], variance[m]), 5, 5);

    }

  }

​

  fill(0, 0, 1);

  for (let i = 0; i < width; i++) {

    let s = 0;

    for (let m = 0; m < NUM_MODELS; m++) {

      s += mixture[m]*evalGaussianPDF(i, mean[m], variance[m]);

    }

    ellipse(i, height-30 -70*height*s, 2, 2);

  }

}

​

function reset() {

  mean = new Array(NUM_MODELS).fill(0);

  variance = new Array(NUM_MODELS).fill(0);

  mixture = new Array(NUM_MODELS).fill(0);

  z = new Array(x.length).fill().map((x) => new Array(NUM_MODELS).fill(0));

​

  for (let i = 0; i < NUM_MODELS; i++) {

    mean[i] = random(width);

    variance[i] = random(width, sq(width));

    mixture[i] = 1.0/NUM_MODELS;

  }

}

​

function evalGaussianPDF(x, mean, variance) {

  return (1.0/sqrt(variance*2*PI)) * exp(-sq(x - mean) / (2 * variance));

}

​

function expectation() {

  for (let n = 0; n < x.length; n++) {

    let normalizer = 0;

    for (let m = 0; m < NUM_MODELS; m++) {

      let p =  evalGaussianPDF(x[n], mean[m], variance[m]);

      z[n][m] = p * mixture[m];

      normalizer += z[n][m];

    }

    for (let m = 0; m < NUM_MODELS; m++) {

      if (normalizer > 0) z[n][m] /= normalizer;

    }

  }

}

​

function maximization() {

  for (let m = 0; m < NUM_MODELS; m++) {

    mean[m] = 0;

    variance[m] = 0;

​

    let zsum = 0;

    for (let n = 0; n < x.length; n++) {

      mean[m] += x[n]*z[n][m];

      zsum += z[n][m];

    }

    if (zsum > 0) mean[m] /= zsum;

​

    mixture[m] = zsum/x.length;

    for (let n = 0; n < x.length; n++) {

      variance[m]  += sq((x[n] - mean[m]))*z[n][m];

    }

    if (zsum > 0) variance[m] /= zsum;

  }

}

​

function EM() {

  expectation();

  maximization();

}