<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Think Bayes
Second Edition
Copyright 2020 Allen B. Downey
License
Step1:  Introduction
Three dimensions!
The Grizzly Bear Problem
In 1996 and 1997 Mowat and Strobeck deployed bear traps in locations in British Columbia and Alberta, in an effort to estimate the population of grizzly bears.  They describe the experiment in "Estimating Population Size of Grizzly Bears Using Hair Capture, DNA Profiling, and Mark-Recapture Analysis"
The "trap" consists of  a lure and several strands of barbed wire intended to capture samples of hair from bears that visit the lure.  Using the hair samples, the researchers use DNA analysis to identify individual bears.
During the first session, on June 29, 1996, the researchers deployed traps at 76 sites.  Returning 10 days later, they obtained 1043 hair samples and identified 23 different bears.  During a second 10-day session they obtained 1191 samples from 19 different bears, where 4 of the 19 were from bears they had identified in the first batch.
To estimate the population of bears from this data, we need a model for the probability that each bear will be observed during each session.  As a starting place, we'll make the simplest assumption, that every bear in the population has the same (unknown) probability of being sampled during each round.
With these assumptions we can compute the probability of the data for a range of possible populations.
As an example, let's suppose that the actual population of bears is 200.
After the first session, 23 of the 200 bears have been identified.
During the second session, if we choose 19 bears at random, what is the probability that 4 of them were previously identified?
I'll define
N
Step2:  So that's the distribution of k given N, K, and n.
Now let's go the other way
Step3:  So that's our prior.
To compute the likelihood of the data, we can use hypergeom with constants K and n, and a range of values of N.
Step4:  We can compute the posterior in the usual way.
Step5:  And here's what it looks like.
Step6:  The most likely value is 109.
Step7:  But the distribution is skewed to the right, so the posterior mean is substantially higher.
Step9:  Two parameter model
Step10:  Two parameters better than one?
Step11:  The Lincoln index problem
A few years ago my occasional correspondent John D. Cook wrote an excellent
blog post about the Lincoln index, which is a way to estimate the
number of errors in a document (or program) by comparing results from
two independent testers. 
http
Step13:  Unknown probabilities
Step14:  Chao et al
Step15:  Spina bifida
Step16:  Diabetes
<ASSISTANT_TASK:>
Python Code:
# If we're running on Colab, install empiricaldist
# https://pypi.org/project/empiricaldist/
import sys
IN_COLAB = 'google.colab' in sys.modules
if IN_COLAB:
    !pip install empiricaldist
# Get utils.py and create directories
import os
if not os.path.exists('utils.py'):
    !wget https://github.com/AllenDowney/ThinkBayes2/raw/master/code/soln/utils.py
        
if not os.path.exists('figs'):
    !mkdir figs
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from empiricaldist import Pmf, Cdf
from utils import decorate, savefig
Explanation: Think Bayes
Second Edition
Copyright 2020 Allen B. Downey
License: Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
End of explanation
from scipy.stats import hypergeom
N = 100
K = 23
n = 19
ks = np.arange(12)
ps = hypergeom(N, K, n).pmf(ks)
plt.bar(ks, ps)
decorate(xlabel='Number of bears observed twice',
         ylabel='PMF',
         title='Hypergeometric distribution of k (known population 200)')
Explanation: Introduction
Three dimensions!
The Grizzly Bear Problem
In 1996 and 1997 Mowat and Strobeck deployed bear traps in locations in British Columbia and Alberta, in an effort to estimate the population of grizzly bears.  They describe the experiment in "Estimating Population Size of Grizzly Bears Using Hair Capture, DNA Profiling, and Mark-Recapture Analysis"
The "trap" consists of  a lure and several strands of barbed wire intended to capture samples of hair from bears that visit the lure.  Using the hair samples, the researchers use DNA analysis to identify individual bears.
During the first session, on June 29, 1996, the researchers deployed traps at 76 sites.  Returning 10 days later, they obtained 1043 hair samples and identified 23 different bears.  During a second 10-day session they obtained 1191 samples from 19 different bears, where 4 of the 19 were from bears they had identified in the first batch.
To estimate the population of bears from this data, we need a model for the probability that each bear will be observed during each session.  As a starting place, we'll make the simplest assumption, that every bear in the population has the same (unknown) probability of being sampled during each round.
With these assumptions we can compute the probability of the data for a range of possible populations.
As an example, let's suppose that the actual population of bears is 200.
After the first session, 23 of the 200 bears have been identified.
During the second session, if we choose 19 bears at random, what is the probability that 4 of them were previously identified?
I'll define
N: actual (unknown) population size, 200.
K: number of bears identified in the first session, 23.
n: number of bears observed in the second session, 19 in the example.
k: the number of bears in the second session that had previously been identified, 4.
For given values of N, K, and n, the distribution of k is described by the hypergeometric distribution:
$PMF(k) = {K \choose k}{N-K \choose n-k}/{N \choose n}$
To understand why, consider: 
The denominator, ${ N \choose n}$, is the number of subsets of $n$ we could choose from a population of $N$ bears.
The numerator is the number of subsets that contain $k$ bears from the previously identified $K$ and $n-k$ from the previously unseen $N-K$.
SciPy provides hypergeom, which we can use to compute this PMF.
End of explanation
Ns = np.arange(50, 501)
prior_N = Pmf(1, Ns)
prior_N.index.name = 'N'
Explanation: So that's the distribution of k given N, K, and n.
Now let's go the other way: given K, n, and k, how can we estimate the total population, N?
As a starting place, let's suppose that, prior to this study, an expert in this domain would have estimated that the population is between 50 and 500, and equally likely to be any value in that range.
End of explanation
K = 23
n = 19
k = 4
likelihood = hypergeom(Ns, K, n).pmf(k)
Explanation: So that's our prior.
To compute the likelihood of the data, we can use hypergeom with constants K and n, and a range of values of N.
End of explanation
posterior_N = prior_N * likelihood
posterior_N.normalize()
Explanation: We can compute the posterior in the usual way.
End of explanation
posterior_N.plot()
decorate(xlabel='Population of bears (N)',
         ylabel='PDF',
         title='Posterior distribution of N')
Explanation: And here's what it looks like.
End of explanation
posterior_N.max_prob()
Explanation: The most likely value is 109.
End of explanation
posterior_N.mean()
posterior_N.credible_interval(0.9)
Explanation: But the distribution is skewed to the right, so the posterior mean is substantially higher.
End of explanation
ps = np.linspace(0, 1, 101)
prior_p = Pmf(1, ps)
prior_p.index.name = 'p'
from utils import make_joint
joint_prior = make_joint(prior_N, prior_p)
N_mesh, p_mesh = np.meshgrid(Ns, ps)
N_mesh.shape
from scipy.stats import binom
like1 = binom.pmf(K, N_mesh, p_mesh)
like1.sum()
like2 = binom.pmf(k, K, p_mesh) * binom.pmf(n-k, N_mesh-K, p_mesh)
like2.sum()
from utils import normalize
joint_posterior = joint_prior * like1 * like2
normalize(joint_posterior)
def plot_contour(joint, **options):
    Plot a joint distribution.
    
    joint: DataFrame representing a joint PMF
    
    cs = plt.contour(joint.columns, joint.index, joint, **options)
    decorate(xlabel=joint.columns.name, 
             ylabel=joint.index.name)
    return cs
plot_contour(joint_posterior)
decorate(title='Joint posterior distribution of N and p')
from utils import marginal
marginal_N = marginal(joint_posterior, 0)
posterior_N.plot(color='gray')
marginal_N.plot()
decorate(xlabel='Population of bears (N)',
         ylabel='PDF',
         title='Posterior marginal distribution of N')
marginal_N.mean(), marginal_N.credible_interval(0.9)
marginal_p = marginal(joint_posterior, 1)
marginal_p.plot()
decorate(xlabel='Probability of observing a bear',
         ylabel='PDF',
         title='Posterior marginal distribution of p')
from seaborn import JointGrid
def joint_plot(joint, **options):
    x = joint.columns.name
    x = 'x' if x is None else x
    y = joint.index.name
    y = 'y' if y is None else y
    # make a JointGrid with minimal data
    data = pd.DataFrame({x:[0], y:[0]})
    g = JointGrid(x, y, data, **options)
    # replace the contour plot
    g.ax_joint.contour(joint.columns, 
                   joint.index, 
                   joint, 
                   cmap='viridis')
    
    # replace the marginals
    marginal_x = marginal(joint, 0)
    g.ax_marg_x.plot(marginal_x.qs, marginal_x.ps)
    
    marginal_y = marginal(joint, 1)
    g.ax_marg_y.plot(marginal_y.ps, marginal_y.qs)
joint_plot(joint_posterior)
Explanation: Two parameter model
End of explanation
mean = (23 + 19) / 2
N1 = 138
p = mean/N1
p
from scipy.stats import binom
binom(N1, p).std()
binom(N1, p).pmf([23, 19]).prod()
N2 = 173
p = mean/N2
p
binom(N2, p).std()
binom(N2, p).pmf([23, 19]).prod()
Explanation: Two parameters better than one?
End of explanation
n0 = 20
n1 = 15
k11 = 3
k10 = n0 - k11
k01 = n1 - k11
k10, k01
Ns = np.arange(32, 350)
prior_N = Pmf(1, Ns)
prior_N.index.name = 'N'
p0, p1 = 0.2, 0.15
like0 = binom.pmf(n0, Ns, p0)
like0.sum()
like1 = binom.pmf(k01, Ns-n0, p1) * binom.pmf(k11, n0, p1)
like1.sum()
likelihood = like0 * like1
likelihood.shape
prior_N.shape
posterior_N = prior_N * likelihood
posterior_N.normalize()
posterior_N.plot()
decorate(xlabel='n',
         ylabel='PMF',
         title='Posterior marginal distribution of n with known p1, p2')
posterior_N.mean()
Explanation: The Lincoln index problem
A few years ago my occasional correspondent John D. Cook wrote an excellent
blog post about the Lincoln index, which is a way to estimate the
number of errors in a document (or program) by comparing results from
two independent testers. 
http://www.johndcook.com/blog/2010/07/13/lincoln-index/
Here's his presentation of the problem:
"Suppose you have a tester who finds 20 bugs in your program. You
want to estimate how many bugs are really in the program. You know
there are at least 20 bugs, and if you have supreme confidence in your
tester, you may suppose there are around 20 bugs. But maybe your
tester isn't very good. Maybe there are hundreds of bugs. How can you
have any idea how many bugs there are? There's no way to know with one
tester. But if you have two testers, you can get a good idea, even if
you don't know how skilled the testers are."
Suppose the first tester finds 20 bugs, the second finds 15, and they
find 3 in common; how can we estimate the number of bugs?
End of explanation
p0 = np.linspace(0, 1, 61)
prior_p0 = Pmf(1, p0)
prior_p0.index.name = 'p0'
p1 = np.linspace(0, 1, 51)
prior_p1 = Pmf(1, p1)
prior_p1.index.name = 'p1'
from utils import make_joint
joint = make_joint(prior_p0, prior_p1)
joint.shape
joint_pmf = Pmf(joint.transpose().stack())
joint_pmf.head()
joint_prior = make_joint(prior_N, joint_pmf)
joint_prior.shape
joint_prior.head()
likelihood = joint_prior.copy()
Ns = joint_prior.columns
for (p0, p1) in joint_prior.index:
    like0 = binom.pmf(n0, Ns, p0)
    like1 = binom.pmf(k01, Ns-n0, p1) * binom.pmf(k11, n0, p1)
    likelihood.loc[p0, p1] = like0 * like1
    
likelihood.to_numpy().sum()
from utils import normalize
joint_posterior = joint_prior * likelihood
normalize(joint_posterior)
joint_posterior.shape
from utils import marginal
posterior_N = marginal(joint_posterior, 0)
posterior_N.plot()
posterior_pmf = marginal(joint_posterior, 1)
posterior_pmf.shape
posterior_joint_ps = posterior_pmf.unstack().transpose()
posterior_joint_ps.head()
def plot_contour(joint, **options):
    Plot a joint distribution.
    
    joint: DataFrame representing a joint PMF
    
    cs = plt.contour(joint.columns, joint.index, joint, **options)
    decorate(xlabel=joint.columns.name, 
             ylabel=joint.index.name)
    return cs
plot_contour(posterior_joint_ps)
decorate(title='Posterior joint distribution for p1 and p2')
posterior_p1 = marginal(posterior_joint_ps, 0)
posterior_p2 = marginal(posterior_joint_ps, 1)
posterior_p1.plot(label='p1')
posterior_p2.plot(label='p2')
decorate(xlabel='Probability of finding a bug',
         ylabel='PDF',
         title='Posterior marginal distributions of p1 and p2')
posterior_p1.mean(), posterior_p1.credible_interval(0.9)
posterior_p2.mean(), posterior_p2.credible_interval(0.9)
joint_plot(posterior_joint_ps)
Explanation: Unknown probabilities
End of explanation
data = [-1000, 63, 55, 18, 69, 17, 21, 28]
index = pd.MultiIndex.from_product([[0, 1]]*3)
Kijk = pd.Series(data, index)
Kijk
Kijk.xs(1, level=0).sum()
Kijk.xs(1, level=1).sum()
Kijk.xs(1, level=2).sum()
n = pd.Series(1, range(3))
for level in n.index:
    n[level] = Kijk.xs(1, level=level).sum()
    
n
Kixx = Kijk.sum(level=0)
Kixx
Kijx = Kijk.sum(level=[0,1])
Kijx
def num_observed(K):
    return np.asarray(K)[1:].sum()
s0 = num_observed(Kixx)
s0
s1 = num_observed(Kijx)
s1
s2 = num_observed(Kijk)
s2
s = [s0, s1, s2]
s
k = pd.concat([Kixx, Kijx, Kijk])
k
k[1]
k[0,1]
k[0,0,1]
k[0]
N = 300
k0 = N - s[0]
k0
k00 = N - s[1]
k00
k000 = N - s[2]
k000
Ns = np.arange(s2, 500, 5)
prior_N = Pmf(1, Ns)
prior_N.index.name = 'N'
ps = np.linspace(0, 1, 101)
prior_p = Pmf(1, ps)
prior_p.index.name = 'p'
from utils import make_joint
joint_prior = make_joint(prior_N, prior_p)
def likelihood_round1(ps, n, k, s, joint_prior):
    like = joint_prior.copy()
    for N in joint_prior:
        like[N] = binom.pmf(n[0], N, ps)
    return like
like1 = likelihood_round1(ps, n, k, s, joint_prior)
def likelihood_round2(ps, n, k, s, joint_prior):
    like = joint_prior.copy()
    for N in joint_prior:
        k0 = N - s[0]
        like[N] = (binom.pmf(k[0,1], k0, ps) * 
                   binom.pmf(k[1,1], k[1], ps))
    return like
like2 = likelihood_round2(ps, n, k, s, joint_prior)
like2.to_numpy().sum()
joint_posterior2 = joint_prior * like1 * like2
normalize(joint_posterior2)
plot_contour(joint_posterior2)
decorate(title='Joint posterior of N and p')
marginal_N = marginal(joint_posterior2, 0)
marginal_N.plot()
decorate(xlabel='N',
         ylabel='PDF',
         title='Posterior marginal distribution of N')
marginal_N.mean(), marginal_N.credible_interval(0.9)
marginal_p = marginal(joint_posterior2, 1)
marginal_p.plot()
decorate(xlabel='p',
         ylabel='PDF',
         title='Posterior marginal distribution of p')
def likelihood_round3(ps, n, k, s, joint_prior):
    like = joint_prior.copy()
    for N in joint_prior:
        k00 = N - s[1]
        like[N] = (binom.pmf(k[0,0,1], k00, ps) * 
                   binom.pmf(k[0,1,1], k[0,1], ps) * 
                   binom.pmf(k[1,0,1], k[1,0], ps) * 
                   binom.pmf(k[1,1,1], k[1,1], ps))
    return like
like3 = likelihood_round3(ps, n, k, s, joint_prior)
like3.to_numpy().sum()
joint_posterior3 = joint_posterior2 * like3
normalize(joint_posterior3)
plot_contour(joint_posterior3)
decorate(title='Joint posterior of N and p')
marginal_N = marginal(joint_posterior3, 0)
marginal_N.plot()
decorate(xlabel='N',
         ylabel='PDF',
         title='Posterior marginal distribution of N')
marginal_N.mean(), marginal_N.credible_interval(0.9)
marginal_p = marginal(joint_posterior3, 1)
marginal_p.plot()
decorate(xlabel='p',
         ylabel='PDF',
         title='Posterior marginal distribution of p')
joint_plot(joint_posterior3)
Explanation: Chao et al
End of explanation
data_sb = [-100, 60, 49, 4, 247, 112, 142, 12]
num_observed(data_sb)
def make_stats(data, num_rounds):
    index = pd.MultiIndex.from_product([[0, 1]]*num_rounds)
    K = pd.Series(data, index)
    n = pd.Series(0, range(num_rounds))
    for level in n.index:
        n[level] = K.xs(1, level=level).sum()
    
    t = [K.sum(level=list(range(i+1)))
         for i in range(num_rounds)]
    
    s = [num_observed(Kx) for Kx in t]
        
    k = pd.concat(t)
    return n, k, s
n, k, s = make_stats(data_sb, 3)
n
k
s
Ns = np.arange(s[2], 1000, 5)
prior_N = Pmf(1, Ns)
prior_N.index.name = 'N'
prior_N.shape
probs0 = np.linspace(0.5, 1.0, 51)
prior_p0 = Pmf(1, probs0)
prior_p0.index.name = 'p0'
prior_p0.head()
probs1 = np.linspace(0.1, 0.5, 41)
prior_p1 = Pmf(1, probs1)
prior_p1.index.name = 'p1'
prior_p1.head()
probs2 = np.linspace(0.1, 0.4, 31)
prior_p2 = Pmf(1, probs2)
prior_p2.index.name = 'p2'
prior_p2.head()
def make_joint3(prior0, prior1, prior2):
    joint2 = make_joint(prior0, prior1)
    joint2_pmf = Pmf(joint2.transpose().stack())
    joint3 = make_joint(prior2, joint2_pmf)
    return joint3
joint_prior = make_joint3(prior_p0, prior_p1, prior_N)
joint_prior.head()
likelihood = joint_prior.copy()
Ns = joint_prior.columns
for (p0, p1) in joint_prior.index:
    like0 = binom.pmf(k[1], Ns, p0)
    k0 = Ns - s[0]
    like1 = binom.pmf(k[0,1], k0, p1) * binom.pmf(k[1,1], k[1], p1)
    likelihood.loc[p0, p1] = like0 * like1
    
likelihood.to_numpy().sum()
from utils import normalize
joint_posterior = joint_prior * likelihood
normalize(joint_posterior)
joint_posterior.shape
from utils import marginal
posterior_N = marginal(joint_posterior, 0)
posterior_N.plot()
decorate(xlabel='N',
         ylabel='PDF',
         title='Posterior for N after two rounds')
posterior_N.mean(), posterior_N.credible_interval(0.9)
posterior_pmf = marginal(joint_posterior, 1)
posterior_pmf.head()
posterior_joint_ps = posterior_pmf.unstack().transpose()
posterior_joint_ps.head()
plot_contour(posterior_joint_ps)
decorate(title='Posterior joint distribution for p1 and p2')
posterior_p1 = marginal(posterior_joint_ps, 0)
posterior_p2 = marginal(posterior_joint_ps, 1)
posterior_p1.plot(label='p1')
posterior_p2.plot(label='p2')
decorate(xlabel='Probability of finding a bug',
         ylabel='PDF',
         title='Posterior marginal distributions of p1 and p2')
joint3 = make_joint(prior_p2, posterior_pmf)
joint3.head()
joint3_pmf = Pmf(joint3.stack())
joint3_pmf.head()
prior4 = make_joint(posterior_N, joint3_pmf)
prior4.head()
prior4.shape
joint2 = make_joint(posterior_N, prior_p2)
joint2.head()
like2 = joint2.copy()
Ns = joint2.columns
for p2 in joint2.index:
    k00 = Ns - s[1]
    like = (binom.pmf(k[0,0,1], k00, p2) *
            binom.pmf(k[0,1,1], k[0,1], p2) *
            binom.pmf(k[1,0,1], k[1,0], p2) *
            binom.pmf(k[1,1,1], k[1,1], p2))
    
    like2.loc[p2] = like
like2.to_numpy().sum()
like2.head()
like4 = prior4.copy()
for (p0, p1, p2) in prior4.index:
    like4.loc[p0, p1, p2] = like2.loc[p2]
    
like4.to_numpy().sum()
posterior4 = prior4 * like4
normalize(posterior4)
marginal_N = marginal(posterior4, 0)
marginal_N.plot()
decorate(xlabel='N',
         ylabel='PDF',
         title='Posterior for N after three rounds')
marginal_N.mean(), marginal_N.credible_interval(0.9)
posterior_p012 = marginal(posterior4, 1)
posterior_p012.unstack().head()
posterior_p2 = marginal(posterior_p012.unstack(), 0)
posterior_p2.plot()
posterior_p01 = marginal(posterior_p012.unstack(), 1)
joint_plot(posterior_p01.unstack().transpose())
Explanation: Spina bifida
End of explanation
data4 = [-10000, 10, 182, 8, 74, 7, 20, 14, 709, 12, 650, 46, 104, 18, 157, 58]
num_observed(data4)
n, k, s = make_stats(data4, 4)
n
k
s
Explanation: Diabetes
End of explanation
<END_TASK>