Maximum Likelihood estimator for k heads in n trials. https://www.udacity.com/course/viewer#!/c-st101/l-48727700/e-48726376/m-48724282

mle.csv

mle.py

import itertools
import math

def combo(n, r):
    """nCr"""
    
    return math.factorial(n) / math.factorial(n - r) / math.factorial(r)


def prob_exact(heads, tails, prob_h):
    """Binomial distribution of `heads` successes in `heads + tails` trials"""
    
    n = heads + tails
    return combo(n, heads) * pow(prob_h, heads) * pow(1 - prob_h, tails)


def seq(start, end, step):
    """Like range, for floating point steps"""

    sample_count = (end - start) / step
    return itertools.islice(itertools.count(start, step), sample_count)


def prob_seq(prob_head, seq):
    """Maximum likelihood estimator given a sequence of heads(1) and tails(0) and 
    probabilility of heads"""
    
    heads =  len(filter(lambda i: i == "1", seq))
    tails = len(seq) - heads
    return prob_exact(heads, tails, prob_head)


def plot_distribution():
    """Plot the MLE curve"""

    probs = ((i, prob_seq(i, "101")) for i in seq(0, 1, 0.01))

    def prob_to_str(k, v):
        return "{},{}".format(k, v)
    
    prob_strs = (prob_to_str(k, v) for k, v in probs)
    s = "\n".join(prob_strs)    
    
    with open("mle.csv", 'w') as f:
        f.write(s)