-
Notifications
You must be signed in to change notification settings - Fork 0
/
condprob.py
35 lines (28 loc) · 970 Bytes
/
condprob.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import survey
import pmf
import utils
import matplotlib.pyplot as plot
def condProb(cond, pmf) :
result = pmf.Copy()
for i,j in pmf.Items() :
if i < cond :
result.Mult(i,0)
result.Renormalize()
return result.Prob(cond)
table = survey.Pregnancies()
table.ReadRecords()
liveBirths = [x for x in table.records if x.outcome == 1]
pmf = pmf.MakePmfFromList([p.prglength for p in liveBirths])
print 'Probability of being born at 39 weeks: ',pmf.Prob(39)
print 'Conditional Probability of being born at 39 weeks if at week 38: ',condProb(39,pmf)
firstBirths,otherBirths = utils.partition(liveBirths,
lambda p : p.outcome == 1 and p.birthord == 1,
lambda p : p.prglength)
pmf = pmf.MakePmfFromList(firstBirths)
xs,ys = pmf.Render()
plot.bar(xs, ys, width=0.4, color='white')
pmf = pmf.MakePmfFromList(otherBirths)
xs,ys = pmf.Render()
plot.bar([x + 0.5 for x in xs], ys, width=0.4, color='blue')
#plot.bar(valsN,probsN, width=0.4, color='blue')
plot.show()