Contact
CoCalc Logo Icon
StoreFeaturesDocsShareSupport News AboutSign UpSign In
| Download

Think Stats by Allen B. Downey Think Stats is an introduction to Probability and Statistics for Python programmers.

This is the accompanying code for this book.

Website: http://greenteapress.com/wp/think-stats-2e/

Views: 7089
License: GPL3
1
"""This file contains code for use with "Think Stats",
2
by Allen B. Downey, available from greenteapress.com
3
4
Copyright 2014 Allen B. Downey
5
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
6
"""
7
8
from __future__ import print_function
9
10
import numpy as np
11
import sys
12
13
import nsfg
14
import first
15
import thinkstats2
16
import thinkplot
17
18
19
def PmfMean(pmf):
20
"""Computes the mean of a PMF.
21
22
Returns:
23
float mean
24
"""
25
mean = 0.0
26
for x, p in pmf.d.items():
27
mean += p * x
28
return mean
29
30
31
def PmfVar(pmf, mu=None):
32
"""Computes the variance of a PMF.
33
34
Args:
35
mu: the point around which the variance is computed;
36
if omitted, computes the mean
37
38
Returns:
39
float variance
40
"""
41
if mu is None:
42
mu = pmf.Mean()
43
44
var = 0.0
45
for x, p in pmf.d.items():
46
var += p * (x - mu) ** 2
47
return var
48
49
50
def Diffs(t):
51
"""List of differences between the first elements and others.
52
53
t: list of numbers
54
55
returns: list of numbers
56
"""
57
first = t[0]
58
rest = t[1:]
59
diffs = [first - x for x in rest]
60
return diffs
61
62
63
def PairWiseDifferences(live):
64
"""Summarize pairwise differences for children of the same mother.
65
66
live: DataFrame of pregnancy records for live births
67
"""
68
live = live[live.prglngth >= 37]
69
preg_map = nsfg.MakePregMap(live)
70
71
diffs = []
72
for caseid, indices in preg_map.items():
73
lengths = live.loc[indices].prglngth.values
74
if len(lengths) >= 2:
75
diffs.extend(Diffs(lengths))
76
77
mean = thinkstats2.Mean(diffs)
78
print('Mean difference between pairs', mean)
79
80
pmf = thinkstats2.Pmf(diffs)
81
thinkplot.Hist(pmf, align='center')
82
thinkplot.Show(xlabel='Difference in weeks',
83
ylabel='PMF')
84
85
86
def main(script):
87
"""Tests the functions in this module.
88
89
script: string script name
90
"""
91
live, firsts, others = first.MakeFrames()
92
PairWiseDifferences(live)
93
94
# test PmfMean and PmfVar
95
prglngth = live.prglngth
96
pmf = thinkstats2.Pmf(prglngth)
97
mean = PmfMean(pmf)
98
var = PmfVar(pmf)
99
100
assert(mean == pmf.Mean())
101
assert(var == pmf.Var())
102
print('mean/var preg length', mean, var)
103
104
print('%s: All tests passed.' % script)
105
106
107
if __name__ == '__main__':
108
main(*sys.argv)
109
110