#!/usr/bin/env python # -*- coding: utf-8 -*- # # Copyright (C) 2013 W. Trevor King # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as # published by the Free Software Foundation, either version 3 of the # License, or (at your option) any later version. # # This program is distributed in the hope that it will be useful, but # WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU # Lesser General Public License for more details. # # You should have received a copy of the GNU Lesser General Public # License along with this program. If not, see # . """Understanding expectation maximization You have a bag of red, green, and blue balls, from which you draw N times with replacement and get n1 red, n2 green, and n3 blue balls. The probability of any one combination of n1, n2, and n3 is given by the multinomial distribution: p(n1, n2, n3) = N! / (n1! n2! n3!) p1^n1 p2^n2 p3^n3 From some outside information, we can parameterize this model in terms of a single hidden variable p: p1 = 1/4 p2 = 1/4 + p/4 p3 = 1/2 - p/4 If we are red/green colorblind, we only measure m1 = n1 + n2 m2 = n3 What is p (the hidden variable)? What were n1 and n2? """ import numpy as _numpy class BallBag (object): """Color-blind ball drawings """ def __init__(self, p=1): self._p = p self.pvals = [0.25, 0.25 + p/4., 0.5 - p/4.] def draw(self, n=10): """Draw `n` balls from the bag with replacement Return (m1, m2), where m1 is the number of red or green balls and m2 is the number of blue balls. """ nvals = _numpy.random.multinomial(n=n, pvals=self.pvals) m1 = sum(nvals[:2]) # red and green m2 = nvals[2] return (m1, m2) class Analyzer (object): def __init__(self, m1, m2): self.m1 = m1 self.m2 = m2 self.p = self.E_n1 = self.E_n2 = None def __call__(self): pass def print_results(self): print('Results for {}:'.format(type(self).__name__)) for name,attr in [ ('p', 'p'), ('E[n1|m1,m2]', 'E_n1'), ('E[n2|m1,m2]', 'E_n2'), ]: print(' {}: {}'.format(name, getattr(self, attr))) class Naive (Analyzer): """Simple analysis With a large enough sample, the measured m1 and m2 give good estimates for (p1 + p2) and p3. You can use either of these to solve for the unknown p, and then solve for E[n1] and E[n2]. While this is an easy approach for the colored ball example, it doesn't generalize well to more complicated models ;). """ def __call__(self): N = self.m1 + self.m2 p1 = 0.25 p3 = self.m2 / float(N) p2 = 1 - p1 - p3 self.p = 4*p2 - 1 self.E_n1 = p1 * N self.E_n2 = p2 * N class MaximumLikelihood (Analyzer): """Analytical ML estimation The likelihood of a general model θ looks like: L(x^N,θ) = sum_{n=1}^N log[p(x_n|θ)] + log[p(θ)] dropping the p(θ) term and applying to this situation, we have: L([m1,m2], p) = N! / (m1! m2!) (p1 + p2)^m1 p3^m2 = N! / (m1! m2!) (1/2 + p/4)^m1 (1/2 - p/4)^m2 which comes from recognizing that to a color-blind experimenter the three ball colors are effectively two ball colors, so they'll have a binomial distribution. Maximizing the log-likelihood: log[L(m1,m2)] = log[N!…] + m1 log[1/2 + p/4] + m2 log[1/2 - p/4] d/dp log[L] = m1/(1/2 + p/4)/4 + m2/(1/2 - p/4)/(-4) = 0 m1 (2 - p) = m2 (2 + p) 2 m1 - m1 p = 2 m2 + m2 p (m1 + m2) p = 2 (m1 - m2) p = 2 (m1 - m2) / (m1 + m2) Given this value of p, the the expected values of n1 and n2 are: E[n1|m1,m2] = p1 / (p1 + p2) * m1 # from the relative probabilities = (1/4) / (1/2 + p/4) * m1 = m1 / (2 + p) = m1 / [2 + 2 (m1 - m2) / (m1 + m2)] = m1/2 * (m1 + m2) / (m1 + m2 + m1 - m2) = m1 * (m1 + m2) / (4 m1) = (m1 + m2) / 4 E[n2|m1,m2] = p2 / (p1 + p2) * m1 = (1/4 + p/4) / (1/2 + p/4) * m1 = m1 (1 + p) / (2 + p) = m1 [1 + 2 ((m1 - m2) / (m1 + m2)] / [2 + 2 (m1 - m2) / (m1 + m2)] = m1/2 * (m1 + m2 + 2m1 - 2m2) / (m1 + m2 + m1 - m2) = m1 * (3m1 - m2) / (4 m1) = (3m1 - m2) / 4 So with a draw of m1 = 61 and m2 = 39, the ML estimates are: p = 0.44 E[n1|m1,m2] = 25.0 E[n2|m1,m2] = 36.0 """ def __call__(self): N = self.m1 + self.m2 self.p = 2 * (self.m1 - self.m2) / float(N) self.E_n1 = N / 4. self.E_n2 = (3*self.m1 - self.m2) / 4. class ExpectationMaximizer (Analyzer): """Expectation maximization Sometimes analytical ML is hard, so instead we iteratively optimize: Q(θ_t|θ_{t-1}) = E[log[p(x,y,θ_t)]|x,θ_{t-1}] Applying to this situation, we have: Q(p_t|p_{t-1}) = E[log[p([m1,m2],[n1,n2,n3],p_t)]|[m1,m2],p_{t-1}] where: p(m1,m2,n1,n2,n3,p) = δ(m1-n1-n2)δ(m2-n3)p(n1,n2,n3) Plugging in and expanding the log: Q(p_t|p_{t-1}) = E[log[δ(m1…)] + log[δ(m2…)] + log[p(n1,n2,n3,p_t)] |m1,m2,p_{t-1}] ≈ E[log[p(n1,n2,n3,p_t)]|m1,m2,p_{t-1}] # drop boring δ terms ≈ E[log[N!…] + n1 log[1/4] + n2 log[1/4 + p_t/4] + n3 log[1/2 - p_t/4] |m1,m2,p_{t-1}] ≈ E[n2 log[1/4 + p_t/4] + n3 log[1/2 - p_t/4] |m1,m2,p_{t-1}] # drop non-p_t terms ≈ E[n2|m1,m2,p_{t-1}] log[1/4 + p_t/4] + m2 log[1/2 - p_t/4] Maximizing (the M step): d/dp_t Q(p_t|p_{t-1}) ≈ E[n2|m1,m2,p_{t-1}] / (1/4 + p_t/4)/4 + m2 / (1/2 - p_t/4)/(-4) = 0 E[n2|m1,m2,p_{t-1}] / (1 + p_t) = m2 / (2 - p_t) E[n2|m1,m2,p_{t-1}] (2 - p_t) = m2 (1 + p_t) p_t (E[n2|m1,m2,p_{t-1}] + m2) = 2 E[n2|m1,m2,p_{t-1}] - m2 p_t = (2 E[n2|m1,m2,p_{t-1}] - m2) / (E[n2|m1,m2,p_{t-1}] + m2) To get a value for p_t, we need to evaluate those expectations (the E step). Using a subset of the ML analysis: E[n2|m1,m2,p_{t-1}] = m1 (1 + p_{t-1}) / (2 + p_{t-1}) """ def __init__(self, p=0, **kwargs): super(ExpectationMaximizer, self).__init__(**kwargs) self.p = 0 # prior belief def _E_step(self): """Caculate E[ni|m1,m2,p_{t-1}] given the prior parameter p_{t-1} """ return { 'E_n1': m1 / (2. + self.p), 'E_n2': m1 * (1. + self.p) / (2. + self.p), 'E_n3': m2, } def _M_step(self, E_n1, E_n2, E_n3): "Maximize Q(p_t|p_{t-1}) over p_t" self.p = (2.*E_n2 - self.m2) / (E_n2 + self.m2) def __call__(self, n=10): for i in range(n): print(' estimated p{}: {}'.format(i, self.p)) Es = self._E_step() self._M_step(**Es) for key,value in self._E_step().items(): setattr(self, key, value) if __name__ == '__main__': p = 0.6 bag = BallBag(p=p) m1,m2 = bag.draw(n=100) print('estimate p using m1 = {} and m2 = {}'.format(m1, m2)) for analyzer in [ ExpectationMaximizer(m1=m1, m2=m2), MaximumLikelihood(m1=m1, m2=m2), Naive(m1=m1, m2=m2), ]: analyzer() analyzer.print_results()