概率质量函数：怀孕周期的PMF

__author__ = 'dell'

import survey
import Pmf
import matplotlib.pyplot as pyplot

table = survey.Pregnancies()
table.ReadRecords()
print "Number of pregnancies ", len(table.records)

firsts = survey.Pregnancies()
others = survey.Pregnancies()

for p in table.records:
    if p.outcome != 1:
        continue
    if p.birthord == 1:
        firsts.AddRecord(p)
    else:
        others.AddRecord(p)

print 'Num of the first babies :', len(firsts)
print 'Num of others :', len(others)

firsts.length = [r.prglength for r in firsts.records]
others.length = [r.prglength for r in others.records]

firsts.hist = Pmf.MakeHistFromList(firsts.length, 'firsts')
others.hist = Pmf.MakeHistFromList(others.length, 'others')

hists = [firsts.hist, others.hist]

width = 0.4
shifts = [-width, 0.0]

option_list = [dict(color='0.9'), dict(color='blue')]
pyplot.clf()


def Shift(xs, shift):
    return [x + shift for x in xs]

for i, hist in enumerate(hists):
        xs, fs = hist.Render()
        xs = Shift(xs, shifts[i])
        pyplot.bar(xs, fs, label=hist.name, width=width, **option_list[i])

pyplot.show()

运行结果：