Introduction to PMFs

Copyright 2015–2026 Allen Downey

MIT License

Click here to run this notebook on Colab

from os.path import basename, exists


def download(url):
    filename = basename(url)
    if not exists(filename):
        from urllib.request import urlretrieve

        local, _ = urlretrieve(url, filename)
        print("Downloaded " + local)


BASE = "https://github.com/AllenDowney/CompStats/raw/master/"
download(BASE + "2002FemPreg.dct")
download(BASE + "2002FemPreg.dat.gz")
download(BASE + "pg2591.txt")

# this line makes Jupyter show figures in the notebook
%matplotlib inline

Hist objects

A histogram is a map from each possible value to the number of times it appears. A map can be a mathematical function or, as in the examples below, a Python data structure that provides the ability to look up a value and get its probability.

Counter is a data structure provided by Python; I am defining a new data structure, called a Hist, that has all the features of a Counter, plus a few more that I define.

import random
import matplotlib.pyplot as plt
from collections import Counter

class Hist(Counter):
    
    def __add__(self, other):
        """Returns the Pmf of the sum of elements from self and other."""
        return Hist(x + y for x, y in product(self.elements(), other.elements()))
    
    def choice(self):
        """Chooses a random element."""
        return random.choice(list(self.elements()))
    
    def plot(self, **options):
        """Plots the Pmf."""
        plt.bar(*zip(*self.items()), **options)
        plt.xlabel('Values')
        plt.ylabel('Counts')
    
    def ranks(self):
        """Returns ranks and counts as lists."""
        return zip(*enumerate(sorted(self.values(), reverse=True)))

As an example, I’ll make a Hist of the letters in my name:

hist = Hist('allen')
hist

Hist({'l': 2, 'a': 1, 'e': 1, 'n': 1})

We can look up a letter and get the corresponding count:

hist['l']

2

Or loop through all the letters and print their counts:

for letter in hist:
    print(letter, hist[letter])

a 1
l 2
e 1
n 1

Counter provides most_common, which makes a list of (element, count) pairs:

hist.most_common()

[('l', 2), ('a', 1), ('e', 1), ('n', 1)]

Here they are in a more readable form:

for letter, count in hist.most_common():
    print(letter, count)

l 2
a 1
e 1
n 1

I defined choice, which returns a random element from the Hist. On average, ‘l’ should appear twice as often as the other letters.

for i in range(10):
    print(hist.choice())

l
e
l
l
l
e
l
n
l
l

One (perhaps surprising) thing you can use Hists for: checking whether two words are anagrams of each other. If two words are anagrams, they have the same Hist:

def is_anagram(word1, word2):
    return Hist(word1) == Hist(word2)

Here’s a simple test:

is_anagram('allen', 'nella')

True

And my favorite anagram pair:

is_anagram('tachymetric', 'mccarthyite')

True

And here’s a false one, just to make sure:

is_anagram('abcd', 'abccd')

False

So far the elements in the Hists have been letters (actually strings), but in statistics it is more common to work with numerical elements. Here’s a Hist that represents the possible outcomes of a six-sided die:

d6 = Hist([1,2,3,4,5,6])
d6

Hist({1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1})

Hist provides a plot function:

import seaborn as sns
COLORS = sns.color_palette()

d6.plot(color=COLORS[3])

elements returns an iterator

d6.elements()

<itertools.chain at 0x7de5f2830400>

Which is easier to see if you convert to a list:

list(d6.elements())

[1, 2, 3, 4, 5, 6]

The product of two iterators is an iterator that enumerates all pairs:

from itertools import product

product(d6.elements(), d6.elements())

<itertools.product at 0x7de5f283a180>

Here are the elements of the product:

list(product(d6.elements(), d6.elements()))

[(1, 1),
 (1, 2),
 (1, 3),
 (1, 4),
 (1, 5),
 (1, 6),
 (2, 1),
 (2, 2),
 (2, 3),
 (2, 4),
 (2, 5),
 (2, 6),
 (3, 1),
 (3, 2),
 (3, 3),
 (3, 4),
 (3, 5),
 (3, 6),
 (4, 1),
 (4, 2),
 (4, 3),
 (4, 4),
 (4, 5),
 (4, 6),
 (5, 1),
 (5, 2),
 (5, 3),
 (5, 4),
 (5, 5),
 (5, 6),
 (6, 1),
 (6, 2),
 (6, 3),
 (6, 4),
 (6, 5),
 (6, 6)]

Now we can compute the sum of all pairs:

list(x + y for x, y in product(d6.elements(), d6.elements()))

[2,
 3,
 4,
 5,
 6,
 7,
 3,
 4,
 5,
 6,
 7,
 8,
 4,
 5,
 6,
 7,
 8,
 9,
 5,
 6,
 7,
 8,
 9,
 10,
 6,
 7,
 8,
 9,
 10,
 11,
 7,
 8,
 9,
 10,
 11,
 12]

And finally make a Hist of the sums:

Hist(x + y for x, y in product(d6.elements(), d6.elements()))

Hist({7: 6, 6: 5, 8: 5, 5: 4, 9: 4, 4: 3, 10: 3, 3: 2, 11: 2, 2: 1, 12: 1})

But all of that is provided by __add__, which we can call using the + operator:

twice = d6 + d6
twice

Hist({7: 6, 6: 5, 8: 5, 5: 4, 9: 4, 4: 3, 10: 3, 3: 2, 11: 2, 2: 1, 12: 1})

Now we can plot the histogram of outcomes from rolling two dice:

twice.plot(color=COLORS[2])

Or three dice:

thrice = twice + d6
thrice

Hist({10: 27,
      11: 27,
      9: 25,
      12: 25,
      8: 21,
      13: 21,
      7: 15,
      14: 15,
      6: 10,
      15: 10,
      5: 6,
      16: 6,
      4: 3,
      17: 3,
      3: 1,
      18: 1})

Notice that this is looking more and more like a bell curve:

thrice.plot(color=COLORS[1])

As the number of dice increases, the result converges to a normal distribution, also known as a Gaussian distribution.

Are first babies more likely to be late?

This is one of the first topics I wrote about in my blog, and still the most popular, with more than 100,000 page views:

http://allendowney.blogspot.com/2011/02/are-first-babies-more-likely-to-be-late.html

I used data from the National Survey of Family Growth (NSFG):

import re

import numpy as np
import pandas as pd


def read_stata_dct(dct_file):
    """Reads a Stata dictionary file.

    dct_file: string filename

    returns: list of (start, end) column specs, list of variable names
    """
    starts, names = [], []
    with open(dct_file) as f:
        for line in f:
            match = re.search(r'_column\(([^)]*)\)', line)
            if match:
                starts.append(int(match.group(1)) - 1)
                names.append(line.split()[2].lower())
    colspecs = list(zip(starts, starts[1:] + [None]))
    return colspecs, names


def read_fem_preg(dct_file='2002FemPreg.dct', dat_file='2002FemPreg.dat.gz'):
    """Reads the NSFG pregnancy data.

    returns: DataFrame
    """
    colspecs, names = read_stata_dct(dct_file)
    df = pd.read_fwf(dat_file, colspecs=colspecs, names=names, compression='gzip')
    df = df.copy()  # defragment the columns read_fwf creates

    # replace 'not ascertained', 'refused', 'don't know' with NaN
    na_vals = [97, 98, 99]
    df['birthwgt_lb'] = df['birthwgt_lb'].replace(na_vals, np.nan)
    df['birthwgt_oz'] = df['birthwgt_oz'].replace(na_vals, np.nan)

    # birthwgt_lb contains at least one bogus value (51 lbs)
    df.loc[df['birthwgt_lb'] > 20, 'birthwgt_lb'] = np.nan

    # birth weight is stored in two columns, lbs and oz;
    # convert to a single column in lb
    df['totalwgt_lb'] = df['birthwgt_lb'] + df['birthwgt_oz'] / 16
    return df

preg = read_fem_preg()

preg.head()

	caseid	pregordr	howpreg_n	howpreg_p	moscurrp	nowprgdk	pregend1	pregend2	nbrnaliv	multbrth	...	laborfor_i	religion_i	metro_i	basewgt	adj_mod_basewgt	finalwgt	secu_p	sest	cmintvw	totalwgt_lb
0	1	1	NaN	NaN	NaN	NaN	6.0	NaN	1.0	NaN	...	0	0	0	3410.389399	3869.349602	6448.271112	2	9	1231	8.8125
1	1	2	NaN	NaN	NaN	NaN	6.0	NaN	1.0	NaN	...	0	0	0	3410.389399	3869.349602	6448.271112	2	9	1231	7.8750
2	2	1	NaN	NaN	NaN	NaN	5.0	NaN	3.0	5.0	...	0	0	0	7226.301740	8567.549110	12999.542264	2	12	1231	9.1250
3	2	2	NaN	NaN	NaN	NaN	6.0	NaN	1.0	NaN	...	0	0	0	7226.301740	8567.549110	12999.542264	2	12	1231	7.0000
4	2	3	NaN	NaN	NaN	NaN	6.0	NaN	1.0	NaN	...	0	0	0	7226.301740	8567.549110	12999.542264	2	12	1231	6.1875

5 rows × 244 columns

The variable outcome encodes the outcome of the pregnancy. Outcome 1 is a live birth.

preg.outcome.value_counts().sort_index()

outcome
1    9148
2    1862
3     120
4    1921
5     190
6     352
Name: count, dtype: int64

pregordr is 1 for first pregnancies, 2 for others.

preg.pregordr.value_counts().sort_index()

pregordr
1     5033
2     3766
3     2334
4     1224
5      613
6      308
7      158
8       78
9       38
10      17
11       8
12       5
13       3
14       3
15       1
16       1
17       1
18       1
19       1
Name: count, dtype: int64

I selected live births, then split into first babies and others.

live = preg[preg.outcome == 1]
firsts = live[live.birthord == 1]
others = live[live.birthord != 1]
len(firsts), len(others)

(4413, 4735)

The mean pregnancy lengths are slightly different:

firsts.prglngth.mean(), others.prglngth.mean()

(np.float64(38.60095173351461), np.float64(38.52291446673706))

The difference is 0.078 weeks:

diff = firsts.prglngth.mean() - others.prglngth.mean()
diff

np.float64(0.07803726677754952)

Which is 13 hours. Note: the best units to report are often not the units you computed.

diff * 7 * 24

np.float64(13.11026081862832)

Let’s see if we can visualize the difference in the histograms:

first_hist = Hist(firsts.prglngth)
other_hist = Hist(others.prglngth)

I used some plotting options to put two bar charts side-by-side:

def plot_distributions(dist1, dist2):
    dist1.plot(width=-0.45, align='edge', color=COLORS[3], label='firsts')
    dist2.plot(width=0.45, align='edge', color=COLORS[4], label='others')
    plt.xlim(33.5, 43.5)
    plt.legend()

Here are the two histograms:

plot_distributions(first_hist, other_hist)

Remember that the vertical axis is counts. In this case, we are comparing counts with different totals, which might be misleading.

An alternative is to compute a probability mass function (PMF), which divides the counts by the totals, yielding a map from each element to its probability.

The probabilities are “normalized” to add up to 1.

import numpy as np

class Pmf(Hist):
    
    def normalize(self):
        total = sum(self.values())
        for element in self:
            self[element] /= total
        return self
    
    def plot_cumulative(self, **options):
        xs, ps = zip(*sorted(self.items()))
        cs = np.cumsum(ps, dtype=float)
        cs /= cs[-1]
        plt.plot(xs, cs, **options)

Now we can compare PMFs fairly.

first_pmf = Pmf(firsts.prglngth).normalize()
other_pmf = Pmf(others.prglngth).normalize()

Using PMFs, we see that some of the difference at 39 weeks was an artifact of the different sample sizes:

plot_distributions(first_pmf, other_pmf)

Even so, it is not easy to compare PMFs. One more alternative is the cumulative mass function (CMF), which shows, for each \(t\), the total probability up to and including \(t\).

first_pmf.plot_cumulative(linewidth=4, color=COLORS[3], label='firsts')
other_pmf.plot_cumulative(linewidth=4, color=COLORS[4], label='others')
plt.xlim(23.5, 44.5)
plt.legend(loc='upper left')

<matplotlib.legend.Legend at 0x7de5e92ba7e0>

The CDFs are similar up to week 38. After that, first babies are more likely to be born late.

Note: don’t be afraid of thick lines. Differences that are only visible with thin lines are unlikely to be real.

Word Frequencies

Next topic: let’s look at histograms of words, bigrams and trigrams.

The following function reads lines from a file or URL and splits them into words:

def iterate_words(filename):
    """Read lines from a file and split them into words."""
    for line in open(filename):
        for word in line.split():
            yield word.strip()

Here’s an example using a book from Project Gutenberg. wc is a histogram of word counts:

# FAIRY TALES
# By The Brothers Grimm
# https://www.gutenberg.org/cache/epub/2591/pg2591.txt
wc = Hist(iterate_words('pg2591.txt'))

Here are the 20 most common words:

wc.most_common(20)

[('the', 6507),
 ('and', 5250),
 ('to', 2707),
 ('a', 1932),
 ('he', 1817),
 ('of', 1450),
 ('was', 1337),
 ('in', 1080),
 ('she', 1049),
 ('that', 1021),
 ('his', 1014),
 ('you', 941),
 ('it', 881),
 ('her', 880),
 ('had', 827),
 ('I', 755),
 ('they', 751),
 ('for', 721),
 ('with', 720),
 ('as', 718)]

Word frequencies in natural languages follow a predictable pattern called Zipf’s law (which is an instance of Stigler’s law, which is also an instance of Stigler’s law).

We can see the pattern by lining up the words in descending order of frequency and plotting their ranks (1st, 2nd, 3rd, …) versus counts (6507, 5250, 2707):

ranks, counts = wc.ranks()
plt.plot(ranks, counts, linewidth=10, color=COLORS[5])
plt.xlabel('rank')
plt.ylabel('count')

Text(0, 0.5, 'count')

Huh. Maybe that’s not so clear after all. The problem is that the counts drop off very quickly. If we use the highest count to scale the figure, most of the other counts are indistinguishable from zero.

Also, there are more than 10,000 words, but most of them appear only a few times, so we are wasting most of the space in the figure in a regime where nothing is happening.

This kind of thing happens a lot. A common way to deal with it is to compute the log of the quantities or to plot them on a log scale:

ranks, counts = wc.ranks()
plt.plot(ranks, counts, linewidth=4, color=COLORS[5])
plt.xlabel('rank')
plt.ylabel('count')
plt.xscale('log')
plt.yscale('log')

This (approximately) straight line is characteristic of Zipf’s law.

n-grams

On to the next topic: bigrams and trigrams.

from itertools import tee

def pairwise(iterator):
    """Iterates through a sequence in overlapping pairs.
    
    If the sequence is 1, 2, 3, the result is (1, 2), (2, 3), (3, 4), etc.
    """
    a, b = tee(iterator)
    next(b, None)
    return zip(a, b)

bigrams is the histogram of word pairs:

bigrams = Hist(pairwise(iterate_words('pg2591.txt')))

And here are the 20 most common:

bigrams.most_common(20)

[(('to', 'the'), 444),
 (('in', 'the'), 399),
 (('of', 'the'), 369),
 (('and', 'the'), 349),
 (('into', 'the'), 294),
 (('said', 'the'), 251),
 (('on', 'the'), 199),
 (('and', 'when'), 168),
 (('he', 'had'), 164),
 (('he', 'was'), 164),
 (('to', 'be'), 163),
 (('it', 'was'), 152),
 (('Then', 'the'), 151),
 (('I', 'will'), 149),
 (('that', 'he'), 143),
 (('at', 'the'), 142),
 (('came', 'to'), 138),
 (('and', 'he'), 135),
 (('she', 'was'), 129),
 (('all', 'the'), 125)]

Similarly, we can iterate the trigrams:

def triplewise(iterator):
    a, b, c = tee(iterator, 3)
    next(b)
    next(c)
    next(c)
    return zip(a, b, c)

And make a histogram:

trigrams = Hist(triplewise(iterate_words('pg2591.txt')))

Here are the 20 most common:

trigrams.most_common(20)

[(('came', 'to', 'the'), 65),
 (('and', 'when', 'he'), 50),
 (('out', 'of', 'the'), 50),
 (('said', 'to', 'the'), 34),
 (('he', 'came', 'to'), 33),
 (('and', 'when', 'she'), 33),
 (('went', 'into', 'the'), 32),
 (('went', 'to', 'the'), 31),
 (('and', 'said', 'to'), 31),
 (('one', 'of', 'the'), 30),
 (('came', 'to', 'a'), 30),
 (('and', 'as', 'he'), 29),
 (('they', 'came', 'to'), 29),
 (('he', 'did', 'not'), 28),
 (('there', 'was', 'a'), 28),
 (('that', 'he', 'had'), 28),
 (('and', 'I', 'will'), 27),
 (('that', 'it', 'was'), 25),
 (('and', 'at', 'last'), 24),
 (('and', 'when', 'the'), 24)]

And now for a little fun. I’ll make a dictionary that maps from each word pair to a Hist of the words that can follow.

from collections import defaultdict

d = defaultdict(Hist)
for a, b, c in trigrams:
    d[a, b][c] += trigrams[a, b, c]

Now we can look up a pair and see what might come next:

d['the', 'blood']

Hist({'ran': 2,
      'of': 2,
      'on': 1,
      'that': 1,
      'came,': 1,
      'streamed': 1,
      'fell': 1,
      'might': 1,
      'ran.': 1})

Here are the most common words that follow “into the”:

d['into', 'the'].most_common(10)

[('forest', 15),
 ('forest,', 13),
 ('garden', 9),
 ('kitchen,', 8),
 ('cellar', 8),
 ('wide', 7),
 ('room,', 7),
 ('water,', 7),
 ('wood', 6),
 ('kitchen', 6)]

Here are the words that follow “said the”:

d['said', 'the'].most_common(10)

[('old', 13),
 ('man,', 12),
 ('little', 10),
 ('fisherman,', 8),
 ('father,', 7),
 ('ass,', 6),
 ('other;', 5),
 ('wife,', 5),
 ('fish;', 5),
 ('fish.', 5)]

Hist provides choice, which chooses a random word with probability proportional to count:

d['said', 'the'].choice()

'bird,'

Given a prefix, we can choose a random suffix:

prefix = 'said', 'the'
suffix = d[prefix].choice()
suffix

'Sultan;'

Then we can shift the words and compute the next prefix:

prefix = prefix[1], suffix
prefix

('the', 'Sultan;')

Repeating this process, we can generate random new text that has the same correlation structure between words as the original:

for i in range(100):
    suffix = d[prefix].choice()
    print(suffix, end=' ')
    prefix = prefix[1], suffix

'I will give you your own life.' When he had eaten and drunk, he showed them the mark, and they put out his needle and goose, will never be able to get my ball again, I would willingly have been all over with soot, but left one of the sea! Hearken to me! My wife saw your rampion from the spot if he comes out of which there was so hard and fast that with all his people thither, and want to go away.' 'Ah,' said Elsie, 'if we marry each other they were proud and ill-behaved, and scorned to 

With a prefix of two words, we typically get text that flirts with sensibility.