A/B testing, bandits, and hierarchical models

A/B testing, bandits, and hierarchical models#

This notebook is based on the first session of the Applied Bayesian Modeling course offered by PyMC Labs.

Click here to run this notebook on Colab. For the exercises, use the version with blanks.

# Install PreliZ if needed (pinned to work with the PyMC on Colab)
try:
    import preliz
except ImportError:
    %pip install -q "preliz<0.23"

import arviz as az
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import pymc as pm
import preliz as pz

# Make the figures a little smaller so we fit more on the screen
plt.rcParams["figure.dpi"] = 75
plt.rcParams["figure.figsize"] = (6.4, 3.2)

# Change the way NumPy 2.0 displays floats
np.set_printoptions(legacy="1.25")

Part 1: The beta-binomial model#

In digital marketing, A/B testing is used to evaluate the effect of different versions of an ad, web page, etc. More generally, it has the same structure as tests of medical treatments, public policy interventions, product design, and more.

As an example, suppose we test an email campaign that includes a promotional offer. We send version A to 100 people, and 8 of them accept the offer. So the observed conversion rate in this sample is 8%.

Here’s a dictionary that represents the results.

A = dict(n=100, k=8)

Some questions we’d like to answer:

  • What is the probability someone responds to this email?

  • How confident are we about that estimate?

  • If we send the same email to 100 people, how many will respond?

And if we have more than one version of the email, we want to compare them.

To answer these questions, we’ll use the beta-binomial model:

  • A beta distribution to represent what we believe about the conversion rate.

  • A binomial to represent the distribution of outcomes (number of responses).

The beta distribution#

The first step is to define a prior distribution that represents what we believe about conversion rates of ads like this one, in general, before we consider the data.

We’ll use a beta distribution because it’s defined from 0 to 1, which is the right domain for a probability, and it has parameters we can choose to specify the location and width of the prior.

Here’s the beta distribution we’ll use for this example.

from preliz import Beta

dist = Beta(2, 5)
dist.mean(), dist.std()

(0.2857142857142857, 0.15971914124998499)

dist.plot_pdf(moments="mv")
plt.xlabel("Conversion rate")
plt.ylabel("Density")
plt.title("Beta distribution");
_images/1bb952a9611b8d9f7cc17c1666c33a260f9a99fa566526ea3ffea21469005466.png

With these parameters, the mean is about 0.28, which means we think a typical response rate is about 28%. But the width of the distribution indicates that we are very unsure about that.

The binomial distribution#

Now let’s pick an arbitrary single value from the prior, 0.2, and suppose we send 50 emails with that response rate. The binomial distribution shows how many responses we should expect.

from preliz import Binomial

dist = Binomial(50, 0.2)
dist.mean()

10.0

dist.plot_pdf(moments="mv", color="C1")
plt.xlabel("Number of conversions")
plt.ylabel("PMF")
plt.title("Binomial distribution");
_images/97ff7c6538c063dd03a1269fcd7a3c0422c48ae00f7816bd4b899715fbd06a53.png

The mean is 10, but values as low as 2 and as high as 20 have non-negligible probability.

In general we don’t know the response rate – we have to estimate it from data.

The Bayesian update#

The fundamental idea in Bayesian statistics is that we can use data to update our beliefs.

  • We start with a prior distribution that represents what we believe before we see the data,

  • We compute a posterior distribution that represents what we believe after we see the data.

There are several ways to do this computation:

  • Mathematically, or

  • Computationally, using a grid approximation, variational inference, or MCMC.

MCMC is Markov chain Monte Carlo sampling, which takes a prior distribution and data, and generates a random sample from an approximate posterior distribution.

PyMC provides several MCMC algorithms along with distributions we can use to assemble a model of the data-generating process.

In this example, we model the data-generating process as:

  • Choose a conversion rate from the prior, which is a beta distribution

  • Choose a number of conversions from the likelihood, which is a binomial distribution

Exercise: Fill in this model with a beta prior and binomial likelihood.

# Fill this in

with pm.Model() as model:
    # Prior
    conversion_rate = 0
    
    # Likelihood
    pm.Binomial("obs", p=0, n=0, observed=0)

# Solution

with pm.Model() as model:
    # Prior
    conversion_rate = pm.Beta("conversion_rate", 2, 5)

    # Likelihood
    pm.Binomial("obs", p=conversion_rate, n=A["n"], observed=A["k"])

Here’s a graphical representation of the model.

pm.model_to_graphviz(model)
_images/d8ff0f6eacbfa64b6fddfd3c3c12098dd2ea9be77217cbcd4d7f6ba527f7e23f.svg

First, let’s draw a sample from the prior distribution.

with model:
    prior_idata1 = pm.sample_prior_predictive()

Sampling: [conversion_rate, obs]

Here’s what it looks like.

def plot_prior(idata, **options):
    return az.plot_posterior(idata, group="prior", hdi_prob="hide", **options)

plot_prior(prior_idata1)
plt.xlabel("Conversion rate")
plt.ylabel("Density")
plt.title("Prior distribution");
_images/f2cdca941d69f5ad63447342254ed1f41bf0e18bbc64f15295580b4bd62b2835.png

The result looks like the PDF of the beta distribution we computed.

Now let’s sample from the posterior.

with model:
    idata1 = pm.sample()

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [conversion_rate]



Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.









Here’s what the posterior distribution looks like.






az.plot_posterior(idata1);








_images/d27cb5b2182a8e45d6a08265ce5f5590ac75170d1261922f74ce7b2ba76c9701.png





The posterior distribution represents what we believe after seeing the data.
Its mean is about 9%, a little above the observed conversion rate, 8%, because the prior pulls it up.


The 94% HDI extends from about 4% to 15%.
Which means there is a 94% chance that the true conversion rate falls in this range, subject to the assumptions of the model (for example, that the rate does not change over time).


To compare the prior and posterior distributions, we’ll combine idata1 and prior_idata1.






with model:
    idata1.extend(prior_idata1)









And then we can use plot_dist_comparison.






az.plot_dist_comparison(idata1);








_images/52e782aa25709fb569a50ac7a963404a9c8d2dcad8aa39a38ef5d48927014d88.png





In summary:


    

  • The prior distribution represents what we believe about the conversion rate before seeing the data.
    • 

  • The posterior distribution represents what we believe about the conversion rate after seeing the data.
    • 



As we get more data, the posterior distribution gets narrower.






What about that prior?#


The previous example uses a weakly informative prior – based on what we know about email campaigns, we think the response rate could fall anywhere in a wide range.


Weakly informative priors “let the data speak for itself”.


But when you have only a small dataset, the prior makes a big difference, so you should use as much information as you have.


Here’s a problem that illustrates the point, from Chapter 4 of Think Bayes.




In Major League Baseball, most players have a batting average between .170 and .310, which means that their probability of getting a hit is between 0.17 and 0.31.


Suppose a player appearing in their first game gets 3 hits out of 3 attempts. What is the posterior distribution for their probability of getting a hit?




Here’s a model that estimates a proportion after observing k=3 successes out of n=3 attempts.
It uses a weakly informative prior (the same one we used for the email campaign).






with pm.Model() as model_rookie:
    batting_average = pm.Beta("p", 2, 5)  
    pm.Binomial("obs", p=batting_average, n=3, observed=3)
    
    prior_idata2 = pm.sample_prior_predictive()
    idata2 = pm.sample()









Sampling: [obs, p]





Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [p]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.









Here’s what that prior looks like.






plot_prior(prior_idata2);








_images/67cde69ee01faed59a79110623cb2c7d74cb7fd63f14ece1d4f078ec589364f1.png





Here’s the posterior.






az.plot_posterior(idata2);








_images/c1beaa4a76eece11ce5386ae8e9032d25085b663639cce6ebfed1b420262a544.png





In this case, “letting the data speak for itself” yields an unrealistic result.


We have a lot of information about batting averages in the major leagues, which we can use to choose a more reasonable prior.


We can specify the prior in terms of mean and dispersion, and convert to the parameters of the beta distribution.






def compute_params(mean, std):
    alpha = mean * ((mean * (1 - mean)) / std**2 - 1)
    beta = (1 - mean) * ((mean * (1 - mean)) / std**2 - 1)
    return alpha, beta













alpha, beta = compute_params(mean=0.240, std=0.070)
alpha, beta









(8.693877551020407, 27.53061224489796)









Exercise: Go back and modify the prior of model_rookie.
Compute and plot the posterior, and see what effect the new prior has on the results.
Given what we know about the context, do the results seem more reasonable?


This example demonstrates the Bayesian version of shrinkage, so called because an estimate based on small data “shrinks” toward the prior mean.








Part 2: Bayesian bandits#


Let’s get back to the first example, estimating the conversion rate of an email campaign.


For purposes of simulation, let’s assume that the actual long run conversion rate for version A is 10%.
We’ll use a dictionary to keep track of this probability, p, the number of times we have tested this version, n, and the number of successful conversions, k.






A = dict(p=0.1, n=0, k=0)









When we run a campaign, we’ll update this dictionary with the number of people we contact and the number of conversions, drawn from a binomial distribution with parameters n and p.


This function uses a PyMC distribution as a random number generator.
This usage is not common, but it can be useful with distributions that have different parameters in different libraries.
If you use PyMC for everything, compatibility is ensured.






rng = np.random.default_rng(17)


def run_campaign(d, n=100):
    # Draw the number of successes for this campaign
    dist = pm.Binomial.dist(n=n, p=d["p"])
    k = pm.draw(dist, random_seed=rng)
    
    # Update totals
    d["n"] += n
    d["k"] += k









As an example, suppose we send version A to 100 people.






run_campaign(A, 100)
A









{'p': 0.1, 'n': 100, 'k': 12}









We can use this model to estimate the conversion rate.






with pm.Model() as model_A:
    conversion_rate = pm.Beta("conversion_rate", 2, 5)
    pm.Binomial("obs", p=conversion_rate, n=A["n"], observed=A["k"])
    
    idata3 = pm.sample()









Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [conversion_rate]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.









Here’s the posterior distribution.






az.plot_posterior(idata3);








_images/ec09f109ae6a1035681602d5d966681c17786e0569446acdd9f464c97ce7e738.png





Now suppose we introduce a new promotion that we think is more enticing, and for the sake of the example, we’ll assume it is in fact more effective, so p=0.15.






B = dict(p=0.15, n=0, k=0)









But we don’t know yet that B is more effective, so we’ll do an uneven A/B test where we send the new email to 20 people in the next batch of 100.






run_campaign(A, 80)
A









{'p': 0.1, 'n': 180, 'k': 18}













run_campaign(B, 20)
B









{'p': 0.15, 'n': 20, 'k': 6}









Now let’s update our beliefs.




Loop back to here#


Here’s the previous model extended to estimate conversion rates for A and B.






with pm.Model() as model_AB:
    
    # Priors for conversion rates of A and B
    conversion_rate_A = pm.Beta("conversion_rate_A", 2, 5)
    conversion_rate_B = pm.Beta("conversion_rate_B", 2, 5)
    
    # Likelihoods for observed data
    pm.Binomial("obs_A", p=conversion_rate_A, n=A["n"], observed=A["k"])
    pm.Binomial("obs_B", p=conversion_rate_B, n=B["n"], observed=B["k"])

    idata4 = pm.sample()









Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [conversion_rate_A, conversion_rate_B]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.









Here are the posterior distributions.






az.plot_posterior(idata4);








_images/9148aa04dbf2d7575d23cc2b069ca18c82878dd9cd702f5f5006c10ed000f2ca.png





Here’s how we can plot them on the same axes.






posterior = az.extract(idata4)
samples_A = posterior["conversion_rate_A"].values
samples_B = posterior["conversion_rate_B"].values

az.plot_kde(samples_A, label="conversion_rate_A")
az.plot_kde(samples_B, label="conversion_rate_B", plot_kwargs={"color": "C1"})

plt.xlabel("Conversion Rate")
plt.title("Posterior Distributions")
plt.legend();








_images/f699dd66aadc4041987be1a4509ea5e3a570c6b96ee181564307d12b94b5a013.png





Based on these results, it looks like B is better, but the distributions overlap, so we might be tempted to run a hypothesis test.


But rather than merely distinguishing between “yes, there’s a difference” or “no, there isn’t”, it is more useful to compute the distribution of possible differences. We want to know how different the conversion rates are, not just whether they are different.






diff = samples_B - samples_A









Here’s what it looks like.






az.plot_posterior(diff);








_images/f39a6da463d89f84591cd55214ce528d1d8be7da567dd7c76d482696825ce09c.png





On average, we think B is better (positive difference), but it’s still possible that A is better (negative difference).


Here are the probabilities of superiority for A and B.






p_superior_A = (samples_A > samples_B).mean()
p_superior_B = (samples_B > samples_A).mean()

p_superior_A, p_superior_B









(0.0115, 0.9885)









At any point in time, these posterior distributions represent what we believe about A and B, but how do we put those beliefs into action?






Thompson sampling#


One option is Thompson sampling, where we use A and B in proportion to their probability of superiority.
As long as we are unsure, we use them both about equally.
As we get confident that one is better, we use it more often.


For the next 100 emails, we’ll use the probability of superiority to choose the proportions.






n = 100
n_A = (n * p_superior_A).round().astype(int)
n_B = n - n_A

n_A, n_B









(1, 99)









And we’ll simulate the campaign.






run_campaign(A, n_A)
A









{'p': 0.1, 'n': 181, 'k': 18}













run_campaign(B, n_B)
B









{'p': 0.15, 'n': 119, 'k': 22}









Now go back and run the model again.
If you loop around a few times, you should see:


    

  • The estimated conversion rates for A and B converge to their true values.
    • 

  • The probability of superiority of B converges to 1.
    • 

  • We gradually use A less and B more.
    • 



This strategy, also known as the Bayesian bandit strategy, is asymptotically optimal: as the experiment goes on, its regret (the conversions lost by using inferior versions) grows as slowly as possible.








Part 3: Hierarchical models#


Our next step is to make a hierarchical version of the model, where the parameters of the prior distribution are drawn from yet another distribution.


To represent the distributions of these parameters, we’ll use a gamma distribution.
It’s a continuous distribution, defined for non-negative numbers, so it is often used to represent a distribution of possible rates, like events per unit of time.


There are different ways to parameterize the gamma distribution – we’ll use the shape-rate form, with parameters called alpha and beta.
(Sorry about all the Greek letters!)






alpha = 16
beta = 8









We’ll use PreliZ again to represent this distribution.






from preliz import Gamma

dist = Gamma(alpha, beta)
dist.mean(), dist.std()









(2.0, 0.5)









Here’s what it looks like.






dist.plot_pdf(moments="mv")
plt.xlabel("Quantity")
plt.ylabel("Density")
plt.title("Gamma distribution");








_images/c5f09dae48f06f2b6a8fb1f56e58fd2c5db8305921905e7ad2e2a1064c2092ef.png





Now we’re ready to build a hierarchical model.




Bandits revisited#


Continuing the A/B testing example, suppose we have an old email we have sent 200 times, and we get 20 conversions.






A = dict(p=0.1, n=200, k=20)
A









{'p': 0.1, 'n': 200, 'k': 20}













obs_proportion_A = A["k"] / A["n"]
obs_proportion_A









0.1









And a new version we have sent only 10 times, with two conversions.






D = dict(p=0.2, n=10, k=2)
D









{'p': 0.2, 'n': 10, 'k': 2}













obs_proportion_D = D["k"] / D["n"]
obs_proportion_D









0.2









Here’s the model that estimates conversion rates for A and D.






with pm.Model() as model_AD:
    # Priors for conversion rates of A and D
    conversion_rate_A = pm.Beta("conversion_rate_A", 2, 5)
    conversion_rate_D = pm.Beta("conversion_rate_D", 2, 5)
    
    # Likelihoods for observed data
    pm.Binomial("obs_A", p=conversion_rate_A, n=A["n"], observed=A["k"])
    pm.Binomial("obs_D", p=conversion_rate_D, n=D["n"], observed=D["k"])
    
    idata5 = pm.sample()









Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [conversion_rate_A, conversion_rate_D]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.









Here are the posterior distributions, showing the observed proportions as reference values.






parameters = ["conversion_rate_A", "conversion_rate_D"]
obs_proportions = [obs_proportion_A, obs_proportion_D]













axes = az.plot_posterior(idata5, var_names=parameters)

for ax, x in zip(axes, obs_proportions):
    ax.axvline(x=x, color="gray", linestyle="--", label="Observed")
    ax.legend()








_images/baf53a53aaaae41ee8b8b4ad4244bd35e27160d0f6e98dcbd248300c9b192e6d.png





In this model, the conversion rates are independent – the data from A doesn’t affect the posterior distribution for D, and vice versa.






pm.model_to_graphviz(model_AD)








_images/3236e1d47b5c7f4f83956a56ba2380b6189413266e9bc181ad815cfaf4436411.svg





And that means we are leaving information on the table: assuming A and D are similar, what we learn about A gives us information about the effectiveness of campaigns like this in general, which influences what we should believe about D.


We can take advantage of this additional information by making the model hierarchical.






# Modify this

with pm.Model() as model_hierarchical:
    # Priors for conversion rates of A and D
    conversion_rate_A = pm.Beta("conversion_rate_A", 2, 5)
    conversion_rate_D = pm.Beta("conversion_rate_D", 2, 5)
    
    # Likelihoods for observed data
    pm.Binomial("obs_A", p=conversion_rate_A, n=A["n"], observed=A["k"])
    pm.Binomial("obs_D", p=conversion_rate_D, n=D["n"], observed=D["k"])













# Solution

with pm.Model() as model_hierarchical:
    # Hyperpriors for the Beta distributions of A and D
    alpha = pm.Gamma("alpha", 16, 8)
    beta = pm.Gamma("beta", 40, 8)
    
    # Priors for conversion rates of A and D, sharing hyperparameters
    conversion_rate_A = pm.Beta("conversion_rate_A", alpha, beta)
    conversion_rate_D = pm.Beta("conversion_rate_D", alpha, beta)

    # come back and add this later
    conversion_rate_new = pm.Beta("conversion_rate_new", alpha, beta)
    
    # Likelihoods for observed data
    pm.Binomial("obs_A", p=conversion_rate_A, n=A["n"], observed=A["k"])
    pm.Binomial("obs_D", p=conversion_rate_D, n=D["n"], observed=D["k"])









The graphical representation of the model shows how the campaigns interact with each other through the hyperparameters.






pm.model_to_graphviz(model_hierarchical)








_images/0ff85d487dee8ccc8e230d64e5234f13470c0c94d3a2935496615a3ad91fe62b.svg





Let’s sample the priors.






with model_hierarchical:
    prior_idata6 = pm.sample_prior_predictive()









Sampling: [alpha, beta, conversion_rate_A, conversion_rate_D, conversion_rate_new, obs_A, obs_D]









Here’s what the hyperpriors look like.






hyperparameters = ["alpha", "beta"]
plot_prior(prior_idata6, var_names=hyperparameters);








_images/897907961972811f37db5af770dd8091ad0ef3eeb3ef2604ee7c2af43d5ebe7c.png





The means are close to 2 and 5, which were the fixed parameters we used in the previous version.
So the priors for the conversion rates are similar to what we had before.






parameters = ["conversion_rate_A", "conversion_rate_D"]
plot_prior(prior_idata6, var_names=parameters);








_images/00f78aa66ac112f7a570dff0f8d148431a76696bd7a0820f077fbb9695d5f273.png





Now when we sample, we simultaneously update beliefs about the conversion rates and the hyperparameters.






with model_hierarchical:
    idata6 = pm.sample()









Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [alpha, beta, conversion_rate_A, conversion_rate_D, conversion_rate_new]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 2 seconds.









Here are the posteriors for the hyperparameters.






az.plot_posterior(idata6, var_names=hyperparameters);








_images/5d49ef7312525896e05de470edbad22f14b70fc152c4b8cb4c73421874963106.png





Finally, here are the posteriors for the conversion rates, showing the observed proportions as reference values.






axes = az.plot_posterior(idata6, var_names=parameters)

for ax, x in zip(axes, obs_proportions):
    ax.axvline(x=x, color="gray", linestyle="--", label="Observed")
    ax.legend()








_images/dcf7aab0353a48a780c84926a67fe20c9c6841955630e50eda33de1f7ca49ad7.png





In this example, the effect of the hierarchical model is small – the posterior means are about the same as the ones we got from the non-hierarchical model.
If we had more versions of the email, we would see more interaction.


But even in this case, the hierarchical model has a useful feature: the posterior distributions of alpha and beta represent what we learned about campaigns in general, based on the results from A and D.
Together they imply a posterior belief about the conversion rate of a new, untested version.


To generate samples for a new, untested version, we can add a new variable to the model


    conversion_rate_new = pm.Beta("conversion_rate_new", alpha, beta)





Because we have no observed data for this variable, its posterior distribution is based on the posterior distributions of the hyperparameters alpha and beta, which were updated based on the data from A and D.


Here’s what it looks like.






az.plot_posterior(idata6, var_names=["conversion_rate_new"]);








_images/4390485c9321b50c5b1bd1d17c103fe1dfa708576128fbe2d7b5f1d9f506a10a.png





Because the response rates of A and D are below the prior mean, the posterior mean is lower than the prior mean (0.29).






Posterior predictive distributions#


So far we have run


    

  • pm.sample_prior_predictive, which samples from the prior distribution (and the prior predictive distribution, which we won’t talk about today), and
    • 

  • pm.sample, which samples from the posterior distribution.
    • 



Now we’ll run sample_posterior_predictive, which samples from the posterior predictive distribution, which is the distribution of the data we expect if we run the same experiment again.






with model_hierarchical:
    posterior_pred = pm.sample_posterior_predictive(idata6)
    idata6.extend(posterior_pred)









Sampling: [obs_A, obs_D]











We’ll use this function to plot it.






def plot_posterior_predictive(idata, **options):
    return az.plot_posterior(idata, group="posterior_predictive", **options);













plot_posterior_predictive(idata6);








_images/6a576b602fd9ea75666ab1e3087315ee92502809c80fdf3651566a7210701bc6.png





These distributions predict the number of conversions we expect if we run the same campaign again – with the same values of n – given our posterior beliefs about the conversion rates.


These distributions represent two sources of uncertainty:


    

  • We don’t know exactly what the conversion rates are, and
    • 

  • Even if we did, we don’t know how many conversions we would get.
    • 



The first is epistemic, because it relates to what we know.
The second is aleatoric, because it relates to randomness (literally “dice”).


As we learn more, the first source of uncertainty gets smaller; the second does not.






Hierarchical baseball#


In the previous example, with only two campaigns, the hierarchical model doesn’t learn very much about campaigns in general, so it doesn’t impose a lot of shrinkage.


Now let’s look at an example with more data – baseball!
Here’s a dataset borrowed from Bill Engels’s tutorial on Gaussian processes.






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)

download("https://github.com/AllenDowney/BDAwithPyMC/raw/main/data/batting_2023.csv")









The dataset includes every batter who had at least one “at bat” in Major League Baseball (MLB).






df = pd.read_csv("batting_2023.csv")









Here are the batters with the lowest batting averages.






cols = ["Name", "Tm", "AB", "H", "BA"]
df.sort_values(by="BA")[cols].head(10)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      388
      Shelby Miller
      Los Angeles
      1
      0
      0.0
    

    

      390
      Calvin Mitchell
      Pittsburgh
      4
      0
      0.0
    

    

      291
      Connor Kaiser
      Colorado
      4
      0
      0.0
    

    

      530
      Braden Shewmake
      Atlanta
      4
      0
      0.0
    

    

      328
      Irving Lopez
      St. Louis
      11
      0
      0.0
    

    

      517
      Josh Sborz
      Texas
      1
      0
      0.0
    

    

      626
      Adam Wainwright
      St. Louis
      2
      0
      0.0
    

    

      618
      Meibrys Viloria
      Cleveland
      3
      0
      0.0
    

    

      524
      Chandler Seagle
      San Diego
      1
      0
      0.0
    

    

      112
      Diego Castillo
      Arizona
      1
      0
      0.0
    

  







And the highest.






df.sort_values(by="BA")[cols].tail(10)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      629
      Forrest Wall
      Atlanta
      13
      6
      0.462
    

    

      452
      Michael Pérez
      New York
      8
      4
      0.500
    

    

      327
      Alejo López
      Cincinnati
      2
      1
      0.500
    

    

      125
      Zack Collins
      Cleveland
      4
      2
      0.500
    

    

      597
      Cole Tucker
      Colorado
      8
      4
      0.500
    

    

      142
      Charlie Culberson
      Atlanta
      1
      1
      1.000
    

    

      135
      Kutter Crawford
      Boston
      0
      0
      NaN
    

    

      307
      Mark Kolozsvary
      Baltimore
      0
      0
      NaN
    

    

      495
      José Rodriguez
      Chicago
      0
      0
      NaN
    

    

      543
      Nick Solak
      Atlanta,Detroit
      0
      0
      NaN
    

  







And here’s the distribution.






az.plot_dist(df["BA"], kind="kde")
plt.xlabel("Batting average")
plt.ylabel("Density");








_images/81381945156f231933a75a7906f109f23fc5bebacb492e4cc499d2ff532b6946.png





What do you notice about the players with the lowest and highest batting averages?


Now let’s make a hierarchical model to estimate batting averages for everyone in the league.






at_bats = df["AB"].values
hits = df["H"].values
n_players = len(df)
n_players









660













# Modify this (starting with the previous example)

with pm.Model() as model_baseball:
    # Hyperpriors for the Beta distributions of A and D
    alpha = pm.Gamma("alpha", 16, 8)
    beta = pm.Gamma("beta", 40, 8)
    
    # Priors for conversion rates of A and D, sharing hyperparameters
    conversion_rate_A = pm.Beta("conversion_rate_A", alpha, beta)
    conversion_rate_D = pm.Beta("conversion_rate_D", alpha, beta)
    
    # Likelihoods for observed data
    pm.Binomial("obs_A", p=conversion_rate_A, n=A["n"], observed=A["k"])
    pm.Binomial("obs_D", p=conversion_rate_D, n=D["n"], observed=D["k"])













# Solution

with pm.Model() as model_baseball:
    # Hyperpriors for the Beta distribution
    alpha = pm.Gamma("alpha", 16, 8)
    beta = pm.Gamma("beta", 40, 8)
    
    # Player-level batting averages sharing the same Beta prior
    player_ba = pm.Beta("player_ba", alpha=alpha, beta=beta, shape=n_players)
    player_new = pm.Beta("player_new", alpha, beta)
    
    # Likelihood for observed hits
    pm.Binomial("hits", p=player_ba, n=at_bats, observed=hits)













with model_baseball:
    idata7 = pm.sample()









Initializing NUTS using jitter+adapt_diag...





Multiprocess sampling (4 chains in 4 jobs)





NUTS: [alpha, beta, player_ba, player_new]







Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 6 seconds.













az.plot_posterior(idata7, var_names=hyperparameters);








_images/a9918c11d3a76257abf2123fd1c49aab6fd90406fdf10344c88d74f08b6a4b28.png





The posterior distribution of the hyperparameters implies this posterior distribution for a new player.






az.plot_posterior(idata7, var_names=["player_new"]);








_images/9074f5881c64a1fbcc9299bf5ed9c426f16cf82bceca6fc39f5f0fbc68eece5e.png





Now let’s see what the model inferred about the players.
We’ll use this function to summarize the posterior for a given player.






def plot_posterior_player(df, idata, name):
    row = df.query('Name == @name')
    display(row[cols])

    # Get the index for the player you want
    idx = row.index[0]

    # Extract posterior samples for that specific player
    player_ba_samples = idata.posterior["player_ba"].sel(player_ba_dim_0=idx)

    # Plot
    az.plot_posterior(player_ba_samples)
    plt.axvline(x=row["BA"].values, color="gray", linestyle="--", label="Observed BA")
    plt.title(name)
    plt.legend()













name = "Chadwick Tromp"
plot_posterior_player(df, idata7, name)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      595
      Chadwick Tromp
      Atlanta
      16
      2
      0.125
    

  



_images/566beaec60d1182de87e2f378d071b3eff712fa092f3e184f79fa4d494914aee.png









name = "Jared Walsh"
plot_posterior_player(df, idata7, name)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      633
      Jared Walsh
      Los Angeles
      104
      13
      0.125
    

  



_images/acd82017a026c0e8d4ab7020e2625d886ebf43a4e16fbe025860310f9ab747bc.png









name = "Forrest Wall"
plot_posterior_player(df, idata7, name)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      629
      Forrest Wall
      Atlanta
      13
      6
      0.462
    

  



_images/04e09e80918eaaf17fe4a98187aeaf3de69acdb8a573c4eb7e956d1f22402ba5.png









name = "Kutter Crawford"
plot_posterior_player(df, idata7, name)













  
    

      
      Name
      Tm
      AB
      H
      BA
    

  
  
    

      135
      Kutter Crawford
      Boston
      0
      0
      NaN
    

  



_images/0c9fe73b6820d2246220b65b22de237d694ae41140591f6ff66b206ec3883c23.png





Baseball mystery: how do you score a run without a plate appearance?






row = df.query("Name == 'Kutter Crawford'")
row.iloc[0]









Unnamed: 0                141
Name          Kutter Crawford
Age                        27
#days                     592
Lev                    Maj-AL
Tm                     Boston
G                           1
PA                          0
AB                          0
R                           1
H                           0
2B                          0
3B                          0
HR                          0
RBI                         0
BB                          0
IBB                         0
SO                          0
HBP                         0
SH                          0
SF                          0
GDP                         0
SB                          0
CS                          0
BA                        NaN
OBP                       NaN
SLG                       NaN
OPS                       NaN
mlbID                  676710
Name: 135, dtype: object









Answer


Copyright 2025–2026 Allen Downey


Code license: MIT License


Text license: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International