| 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.
Project: Support and Testing
Views: 7138License: GPL3
"""This file contains code for use with "Think Stats",1by Allen B. Downey, available from greenteapress.com23Copyright 2014 Allen B. Downey4License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html5"""67from __future__ import print_function89import numpy as np10import sys1112import nsfg13import first14import thinkstats215import thinkplot161718def PmfMean(pmf):19"""Computes the mean of a PMF.2021Returns:22float mean23"""24mean = 0.025for x, p in pmf.d.items():26mean += p * x27return mean282930def PmfVar(pmf, mu=None):31"""Computes the variance of a PMF.3233Args:34mu: the point around which the variance is computed;35if omitted, computes the mean3637Returns:38float variance39"""40if mu is None:41mu = pmf.Mean()4243var = 0.044for x, p in pmf.d.items():45var += p * (x - mu) ** 246return var474849def Diffs(t):50"""List of differences between the first elements and others.5152t: list of numbers5354returns: list of numbers55"""56first = t[0]57rest = t[1:]58diffs = [first - x for x in rest]59return diffs606162def PairWiseDifferences(live):63"""Summarize pairwise differences for children of the same mother.6465live: DataFrame of pregnancy records for live births66"""67live = live[live.prglngth >= 37]68preg_map = nsfg.MakePregMap(live)6970diffs = []71for caseid, indices in preg_map.items():72lengths = live.loc[indices].prglngth.values73if len(lengths) >= 2:74diffs.extend(Diffs(lengths))7576mean = thinkstats2.Mean(diffs)77print('Mean difference between pairs', mean)7879pmf = thinkstats2.Pmf(diffs)80thinkplot.Hist(pmf, align='center')81thinkplot.Show(xlabel='Difference in weeks',82ylabel='PMF')838485def main(script):86"""Tests the functions in this module.8788script: string script name89"""90live, firsts, others = first.MakeFrames()91PairWiseDifferences(live)9293# test PmfMean and PmfVar94prglngth = live.prglngth95pmf = thinkstats2.Pmf(prglngth)96mean = PmfMean(pmf)97var = PmfVar(pmf)9899assert(mean == pmf.Mean())100assert(var == pmf.Var())101print('mean/var preg length', mean, var)102103print('%s: All tests passed.' % script)104105106if __name__ == '__main__':107main(*sys.argv)108109110