混合正規分布のEMアルゴリズム

EMTest.scala

package jp.yukoba

import org.apache.commons.math3.distribution.NormalDistribution
import org.scalatest.FunSuite

import scala.util.Random

class EMTest extends FunSuite {
  /** 統計的学習の基礎のp.312より */
  test("混合正規分布のEMアルゴリズム") {
    val data = Vector(
      -0.39, 0.12, 0.94, 1.67, 1.76, 2.44, 3.72, 4.28, 4.92, 5.53,
      0.06, 0.48, 1.01, 1.68, 1.80, 3.25, 4.12, 4.60, 5.28, 6.22)
    val dataMean = data.sum / data.size

    val shuffledData = Random.shuffle(data)
    val u1_i = shuffledData(0)
    val u2_i = shuffledData(1)
    val s1_i = data.map(x => (x - dataMean) * (x - dataMean)).sum / data.size
    val s2_i = s1_i
    val pi_i = 0.5

    ((u1_i, u2_i, s1_i, s2_i, pi_i) /: (0 until 100)) { case ((u1, u2, s1, s2, pi), _) =>
      // 期待値ステップ
      val gamma = data.map(y => {
        def normDist(u: Double, s: Double, x: Double) = new NormalDistribution(u, math.sqrt(s)).density(x)

        val a = pi * normDist(u2, s2, y)
        val b = (1 - pi) * normDist(u1, s1, y)
        a / (a + b)
      })

      // 最大化ステップ
      val gammaSum2 = gamma.sum
      val gammaSum1 = data.size - gammaSum2
      val u1_n = data.indices.map(i => (1 - gamma(i)) * data(i)).sum / gammaSum1
      val u2_n = data.indices.map(i => gamma(i) * data(i)).sum / gammaSum2
      val s1_n = data.indices.map(i => (1 - gamma(i)) * (data(i) - u1) * (data(i) - u1)).sum / gammaSum1
      val s2_n = data.indices.map(i => gamma(i) * (data(i) - u2) * (data(i) - u2)).sum / gammaSum2
      val pi_n = gammaSum2 / data.size

      println(s"$u1_n, $u2_n, $s1_n, $s2_n, $pi_n")
      (u1_n, u2_n, s1_n, s2_n, pi_n)
    }
  }
}