Austin Rochford (Posts about Bayesian Statistics)

Playing Battleship with Bayesian Search Theory, Thompson Sampling, and Approximate Bayesian Computation

Austin Rochford — Thu, 02 Sep 2021 04:00:00 GMT

As a child, I spent many hours playing Battleship against my Mom or brother to pass the time.

Battleship is a classic game of incomplete information that involves placing ships of different lengths on a square grid ($10 \times 10$ in the most popular version) and trying to guess the location of all of the opponent’s ships (sinking them) before the opponent does the same.

Many years after I stopped playing Battleship with my family, I developed an interest in Bayesian search theory after learning about a firm that specializes in search consulting for the US government at jobs-in-industry panel in graduate school. I had already happily lined up a job in e-commerce optimization, but the idea of using statistics to guide the potentially expensive searches piqued my intellectual curiosity. I e-mailed the presenter afterward for more information and he recommended [Theory of Optimal Search]^(Stone, Lawrence D. Theory of optimal search. Elsevier, 1976.) as the sessential reference in the field. The book is interesting if a bit dry, it is about mathematical optimization after all, and I was always curious about the connection with the search theory it develops and Battleship, which has always stood out in my mind as a basic search problem. Though most of the theory in this book is not directly applicable to Battleship for various reasons, a continuing theme is allocating search in areas with the largest posterior probability of containing the target object given the search results so far. A 1985 result of [Assaf and Zamir]^(Assaf, David, and Shmuel Zamir. “Optimal sequential search: a Bayesian approach.” The Annals of Statistics 13, no. 3 (1985): 1213-1221.) shows the optimality of this strategy in a situation much closer to that of Battleship.

Inspired by these forays into search theory I have thought idly for many years about constructing a near-optimal Bayesian approach to playing Battleship. This post is the culmination of these thoughts, showing how to use ideas from Bayesian search theory, approximate Bayesian computation (ABC), and Thompson sampling to construct an easily tractable, near optimal strategy for Battleship that only requires simulating random Battleship boards and not reasoning through hard-coded special cases to achieve strong play.

Playing Battleship

First we make the necessary Python imports and do a bit of housekeeping.

%matplotlib inline

from abc import ABC, abstractmethod
import datetime
from IPython.display import HTML
from itertools import product
import multiprocessing as mp
from tqdm import tqdm, trange

from empiricaldist import Pmf
import matplotlib as mpl
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.ticker import IndexLocator, StrMethodFormatter
import numpy as np
from scipy import linalg, stats
import seaborn as sns

FIG_WIDTH = 8
FIG_HEIGHT = 6
mpl.rcParams['figure.figsize'] = (FIG_WIDTH, FIG_HEIGHT)

CMAP = LinearSegmentedColormap.from_list(
    'battleship',
    [
        (0.0, 'blue'),
        (0.2, 'orange'),
        (1., 'red')
    ]
)
CMAP.set_bad('w')

sns.set(color_codes=True)

pct_formatter = StrMethodFormatter('{x:.1%}')

SEED = 123456789 # for reproducibility

rng = np.random.default_rng(SEED)

We represent a Battleship board as a sequence of rectangular grids, one per ship. For standard Battleship, that means a game consists of five $10 \times 10$ grids, the first grid containing a ship of length five, the seconding containing a ship of length four, and so on.

The following constants are useful for defining a standard game of Battleship. (Note that we are using Hasbro’s names for the ships, revised in 2002.)

GRID_LENGTH = 10
SHIP_NAMES = [
    "Carrier",
    "Battleship",
    "Destroyer",
    "Submarine",
    "Patrol boat"
]
SHIP_SIZES = np.array([
    5, 
    4,
    3,
    3,
    2
])

n_ships = SHIP_SIZES.size

The following plots show how a Battleship board is represented as a sequence of five grids, one per ship.

ships = np.zeros((n_ships, GRID_LENGTH, GRID_LENGTH))
ships[0, 0, :5] = 1
ships[1, 4:8, 7] = 1
ships[2, 9, 3:6] = 1
ships[3, 3:6, 1] = 1
ships[4, 5, 8:10] = 1

def plot_board(board, cmap=CMAP,
               cbar=False, cbar_kwargs=None,
               ax=None, **heatmap_kwargs):
    if ax is None:
        _, ax = plt.subplots(figsize=(FIG_WIDTH, FIG_WIDTH))
        
    heatmap_kwargs.setdefault('vmin', 0)
    heatmap_kwargs.setdefault('vmax', 1)
    
    board_ = np.atleast_2d(board)
    mask = board_.mask if np.ma.is_masked(board_) else None
    
    if cbar_kwargs is None:
        cbar_kwargs = {}
        
    cbar_kwargs.setdefault('format', pct_formatter)
    cbar_kwargs.setdefault('label', "Posterior probability of a hit")
    
    sns.heatmap(board_, mask=mask, cmap=cmap,
                cbar=cbar, square=True,
                linewidths=1.5, linecolor='k',
                cbar_kws=cbar_kwargs,
                ax=ax, **heatmap_kwargs)
    
    if board_.shape[0] == 1:
        ax.set_yticklabels([])
        
    if cbar:
        cbar = ax.figure.axes[-1]
        cbar.set_ylabel(cbar.get_ylabel(), rotation=270)
        
    return ax

fig, axes = plt.subplots(ncols=n_ships, sharex=True, sharey=True,
                         figsize=(n_ships * 2.5, 2.5))

for (ship, name, ax) in zip(ships, SHIP_NAMES, axes):
    plot_board(ship, ax=ax);
    
    ax.set_xticklabels([]);
    ax.set_yticklabels([]);
    ax.set_title(name);
    
fig.tight_layout();

Adding these ship grids gives us the two-dimensional board familiar to anyone who has played battleship.

def to_board(ships, ship_axis=0):
    return ships.sum(axis=ship_axis)

plot_board(to_board(ships));

The following class implements most of the basics necessary to track the state of a game of Battleship, make guesses, and reveal the contents of guessed spaces.

class Battleship:
    def __init__(self, ships):
        self._ships = ships
        self._turn_revealed = []
        
    @property
    def _board(self):
        return to_board(self._ships)
        
    @property
    def grid_length(self):
        return self._board.shape[0]
    
    @property
    def is_solved(self):
        return self.revealed.sum() == self._board.sum()
    
    @property
    def revealed(self):
        return to_board(self._revealed_ships)
        
    @property
    def _revealed_ships(self):
        if self.turns > 0:
            return self._turn_revealed[-1]
        else:
            return np.ma.masked_all_like(self._ships)
    
    @property
    def ship_sizes(self):
        return self._ships.sum(axis=(1, 2))
    
    @property
    def sunk(self):
        ship_sizes = self._ships.sum(axis=(1, 2))
        revealed_sizes = (self._revealed_ships
                              .sum(axis=(1, 2))
                              .filled(0))
        
        return ship_sizes == revealed_sizes
    
    @property
    def turn_revealed(self):
        return [np.ma.masked_all_like(self._board)] \
                + [to_board(revealed) for revealed in self._turn_revealed]
    
    @property
    def turns(self):
        return len(self._turn_revealed)

These bookkeeping methods are fairly self-explanatory. An instance of Battleship tracks which squares have been revealed to the opponent using numpy’s masked arrays. The revealed grid is masked wherever the opponent has not yet guessed.

The heart of our Battleship class is the guess method, which takes the coordinates to guess, and returns a tuple containing whether or not the guess resulted in a hit or miss, and the index of the ship sunk as a result of that guess (None if no ship was sunk). (Here we use a dirty trick to add a method to the Battleship class after it has already been defined for expository clarity. Don’t do this in practice!)

class Battleship(Battleship):
    def guess(self, i, j):
        if not self.revealed.mask[i, j]:
            raise ValueError(f"{i}, {j} already guessed")
        else:
            prev_sunk = self.sunk

            next_ships = self._revealed_ships.copy()
            next_ships[:, i, j] = self._ships[:, i, j]
            self._turn_revealed.append(next_ships)
            
            curr_sunk = self.sunk
            
            if (curr_sunk == prev_sunk).all():
                sunk = None
            else:
                sunk = (curr_sunk & ~prev_sunk).argmax()
            
            return self._board[i, j], sunk

In addition to this class representing game state and history, we define an abstract class Strategy that generates the guesses necessary to play a game.

class Strategy(ABC):
    @abstractmethod
    def next_guess(self, revealed):
        pass

    def reveal(self, i, j, hit_or_miss, sunk):
        pass

The next_guess method takes a masked array of the board as revealed so far and should return the coordinates of the next spot to guess. The reveal method takes guessed coordinates, an indicator of whether the guess resulted in a hit or miss, and the index of the ship sunk as a result of that guess (or None, as with guess above). This method allows a Strategy to update internal state based on the result of guesses. While some simple strategies can infer the next guess purely based on the state of the revealed board, we will see that it is useful when taking a Bayesian approach to maintain state inside a strategy.

The play function take a configuration of ships and a Strategy and plays the corresponding game of Battleship.

def play(ships, strategy, progress_bar=False):
    game = Battleship(ships)
    
    if progress_bar:
        pbar = tqdm(total=int(ships.sum()))
    
    while not game.is_solved:
        i, j = strategy.next_guess(game.revealed)
        hit_or_miss, sunk = game.guess(i, j)
        strategy.reveal(i, j, hit_or_miss, sunk)
        
        if progress_bar and hit_or_miss == 1:
            pbar.update()
            
    if progress_bar:
        pbar.close()
        
    return game

The rest of this post progresses from strategies for simplified versions of Battleship to finally showing how to use ABC and Thompson sampling to achieve near optimal play.

Manual guessing

This framework allows us to manually play Battleship by inputting user guesses and watching how the board evolves.

class ManualStrategy(Strategy):
    def next_guess(self, revealed):
        ax = plot_board(revealed)
        ax.set_title("Currently revealed")
        plt.show()

        i = int(input("Enter row to guess: "))
        j = int(input("Enter column to guess: "))
        
        return i, j

try:
    play(ships, ManualStrategy())
except KeyboardInterrupt:
    print("Game ended")

Enter row to guess:  5
Enter column to guess:  4

Enter row to guess:  0
Enter column to guess:  3

Enter row to guess:  1
Enter column to guess:  3

Enter row to guess:  0
Enter column to guess:  2

Enter row to guess:  0
Enter column to guess:  4

Game ended

This strategy is not really useful for our purposes, but it provides a good illustration of how the game is usually played. The white squares here indicate the cells whose contents are unknown because we have not yet guessed them.

Row with one ship

We begin by simplifying the idea of Battleship to the most basic possible setting of a single row containing a single ship.

The following function generates all possible rows of a given length containing a ship of a given size.

def get_all_ship_rows(grid_length, ship_size):
    i, j = np.indices((grid_length, grid_length))[:, :-ship_size + 1]
    
    return 1 * (i <= j) & (j < i + ship_size)

For a row of length three and a ship of size two (3/2), there are only two such rows.

all_rows_3_2 = get_all_ship_rows(3, 2)

ax = plot_board(all_rows_3_2)
ax.set_title("All rows 3/2");

For a row of length ten and a ship of size four (10/4), there are seven such rows.

all_rows_10_4 = get_all_ship_rows(10, 4)

ax = plot_board(all_rows_10_4)
ax.set_title("All rows 10/4");

We will compare our Bayesian strategies to several benchmark strategies, starting with the simplest possible strategy: random guessing.

Random guessing

The random guessing approach to Battleship is equivalent to the following urn problem from classic probability theory. Imagine an urn containing 100 balls (corresponding to cells on the board), 17 of which are red (corresponding to cells covered by ships), and 83 of which are blue. The number of turns required to solve Battleship through random guessing has the same distribution as the number of balls of any color that are drawn before all 17 red balls have been drawn. It is well known (in the right circles, at least) that this corresponds to the negative hypergeometric distribution.

Original image credit John D. Cook

The following function returns the probability mass function (pmf) of the appropriate negative hypergeometric distribution given a grid and ship size.

def get_random_guess_dist(grid_size, ship_size, n_hit=None):
    if n_hit is None:
        n_hit = ship_size

    support = np.arange(ship_size, grid_size + 1)
    nhg = stats.nhypergeom(grid_size,
                           grid_size - ship_size,
                           n_hit)
    
    return Pmf(nhg.pmf(support - ship_size), support)

We see, unsurprisingly, that random guessing will usually take near the maximum number of turns to solve the game.

std_random_pmf = get_random_guess_dist(GRID_LENGTH**2, SHIP_SIZES.sum())

def plot_turn_dist(pmf, kind='bar', mean=False, ax=None, mean_kwargs=None, **kwargs):
    if ax is None:
        _, ax = plt.subplots()
        
    if kind == 'bar':
        # pandas barplot uses odd indexing, making it hard to
        # mix bar and line plots if we use pandas's versions
        kwargs.setdefault('width', 1)
        kwargs.setdefault('alpha', 0.75)

        ax.bar(pmf.index, pmf, **kwargs)
    elif kind == 'line':
        pmf.plot(ax=ax, **kwargs)
    else:
        raise ValueError("kind must be one of 'bar' or 'line'")    
        
    if mean:
        if mean_kwargs is None:
            mean_kwargs = {}
            
        mean_kwargs.setdefault('ls', '--')
        
        ax.axvline(pmf.mean(), **mean_kwargs);
    
    ax.xaxis.grid(False)
    ax.set_xlabel("Turns")
    
    return ax

def make_pct_yaxis(ax):
    ax.yaxis.set_major_formatter(pct_formatter)
    ax.set_ylabel("Probability")
    
    return ax

std_turn_ax = plot_turn_dist(std_random_pmf, kind='line', mean=True)
    
make_pct_yaxis(std_turn_ax);
std_turn_ax.set_title("Random guessing\nStandard Battleship");

The dashed vertical line shows the expected number of turns to solve Battleship by randomly guessing,

std_random_pmf.mean()

95.38888888888883

Returning the the case of a row with a single ship, we visualize the turn distribution of random guessin strategies for various row lengths and a ship of size two.

RANDOM_ROW_LENGTHS = [4, 5, 8, 10]

fig, axes = plt.subplots(nrows=len(RANDOM_ROW_LENGTHS) // 2, ncols=2,
                         sharey=True,
                         figsize=(FIG_WIDTH, 1.25 * FIG_HEIGHT))

for (row_length, ax) in zip(RANDOM_ROW_LENGTHS, axes.flat):
    pmf = get_random_guess_dist(row_length, 2)
    plot_turn_dist(pmf, ax=ax);
    
    make_pct_yaxis(ax);
    ax.set_title(f"{row_length}/2");

fig.suptitle("Random guessing");
fig.tight_layout();

We do the same for a ship of size three.

fig, axes = plt.subplots(nrows=len(RANDOM_ROW_LENGTHS) // 2, ncols=2,
                         sharey=True,
                         figsize=(FIG_WIDTH, 1.25 * FIG_HEIGHT))

for (row_length, ax) in zip(RANDOM_ROW_LENGTHS, axes.flat):
    pmf = get_random_guess_dist(row_length, 3)
    plot_turn_dist(pmf, ax=ax);
    
    make_pct_yaxis(ax);
    ax.set_title(f"{row_length}/3");

fig.suptitle("Random guessing");
fig.tight_layout();

As with the case of standard Battleship, the most likely outcome is to use all the turns, and the expected value will be slightly below that.

Optimal search

In the simple setting of a row with a single ship, it is not too hard to work out the optimal search strategy, which consists of two phases. First we must locate the ship by getting at least one hit. Once the ship has been located, we must sink it by attacking spaces next to known hits until all of the ship has been found.

After at least one hit

We first implement optimal search after at least one hit has occurred. In this case, we guess the spot immediately to the left of the leftmost hit, unless that hit is in the first cell, or the cell to its left has already been shown to contain water. In those two cases, we guess the cell to the right of the rightmost hit until all cells containing the ship have been hit. The following function implements this strategy.

def next_guess_1d_with_hit(row):
    first = row.argmax()

    # found the left edge, fill out the ship to the right
    if first == 0 or row[first - 1] == 0:
        return first + row[first:].mask.argmax()
    # find the left edge
    else:
        return first - 1

We can show how this strategy works through a simple animated example. Suppose the position of a length three ship on a length ten row is as shown below.

board = np.zeros(10)
board[6:9] = 1

plot_board(board);

Also suppose that after three guesses, the following spots have been revealed.

revealed = np.ma.masked_all_like(board)
revealed[[1, 4, 7]] = board[[1, 4, 7]]

plot_board(revealed);

The above strategy will take three guesses to sink the ship, as shown in the animation below.

turn_revealed = [revealed]

next_j = next_guess_1d_with_hit(revealed)
next_revealed = revealed.copy()
next_revealed[next_j] = board[next_j]
turn_revealed.append(next_revealed)
revealed = next_revealed

next_j = next_guess_1d_with_hit(revealed)
next_revealed = revealed.copy()
next_revealed[next_j] = board[next_j]
turn_revealed.append(next_revealed)
revealed = next_revealed

next_j = next_guess_1d_with_hit(revealed)
next_revealed = revealed.copy()
next_revealed[next_j] = board[next_j]
turn_revealed.append(next_revealed)
revealed = next_revealed

def animate_boards(boards,
                   cmap=CMAP, cbar=False, ax=None,
                   heatmap_kwargs=None, **ani_kwargs):
    if ax is None:
        fig, ax = plt.subplots(figsize=(FIG_WIDTH, FIG_WIDTH))
    else:
        fig = ax.figure
        
    if heatmap_kwargs is None:
        heatmap_kwargs = {}
        
    plot_board(boards[0],
               cmap=cmap, cbar=cbar, ax=ax,
               **heatmap_kwargs)
    
    quadmesh, *_ = ax.get_children()
    
    def ani_func(i):
        quadmesh.set_array(boards[i])
        
        return quadmesh,
    
    ani_kwargs.setdefault('blit', True)
    ani_kwargs.setdefault('frames', len(boards))
    
    return FuncAnimation(fig, ani_func, **ani_kwargs)

%%capture
ani = animate_boards(turn_revealed, interval=300)

HTML(ani.to_html5_video())

Your browser does not support the video tag.

Getting the first hit

Immediately it is clear that we only have to search one of the conrguence classes of the cell index mod the ship size to guarantee a hit. In the following plot of a row of length ten, each congruence class mod three is colored differently.

plot_board(np.arange(10) % 3, cmap='plasma', vmin=0, vmax=2);

A moment’s consideration shows that the congruence class of -1 modulo the ship size will always have the fewest elements of any of the congruence class, so we choose to search along these cells. The following plot shows the cells to be searched until a hit is found in red.

def get_search_cells_1d(grid_length, ship_size):
    i = np.arange(grid_length)
    
    return i[i % ship_size == -1 % ship_size]

row = np.zeros(10)
row[get_search_cells_1d(10, 3)] = 1

plot_board(row);

The following subclass of Strategy implements the optimal strategy of searching along this grid until a hit is found, then greedily filling out the ship.

class Single1DOptimalStrategy(Strategy):
    def __init__(self, grid_length, ship_size):
        self._grid_length = grid_length
        self._ship_size = ship_size
        
    def next_guess(self, revealed):
        if revealed.mask.all() or revealed.sum() == 0:
            search_j = get_search_cells_1d(self._grid_length,
                                           self._ship_size)            
            next_j = search_j[revealed.mask[0, search_j].argmax()]
        else:
            return 0, next_guess_1d_with_hit(revealed[0])
                
        return 0, next_j

We now use this strategy to play all three games with a length three ship on a length five row.

all_rows_5_3 = get_all_ship_rows(5, 3)

opt_row_games_5_3 = [
    play(ship_row[np.newaxis, np.newaxis],
         Single1DOptimalStrategy(5, 3))
        for ship_row in all_rows_5_3
]
opt_row_turns_5_3 = Pmf.from_seq([game.turns for game in opt_row_games_5_3])

ax = plot_turn_dist(opt_row_turns_5_3)

make_pct_yaxis(ax);
ax.set_title("Optimal play\nShip size 3, row length 5");

Two of the games take four turns and one takes three turns. Notably, no games take five turns. This is because we know that the third cell is a guaranteed hit in this configuration.

We can repeat this exercise for a row of length ten.

all_rows_10_3 = get_all_ship_rows(10, 3)

opt_row_games_10_3 = [
    play(ship_row[np.newaxis, np.newaxis],
         Single1DOptimalStrategy(10, 3))
        for ship_row in all_rows_10_3
]
opt_row_turns_pmf_10_3 = Pmf.from_seq([game.turns for game in opt_row_games_10_3])

ax = plot_turn_dist(opt_row_turns_pmf_10_3)

make_pct_yaxis(ax);
ax.set_title("Optimal play\nShip size 3, row length 10");

We see that no games take more than six turns with this configuration. Since the equivalence class of -1 modulo 3 contains three members (2, 5, and 8), we are guaranteed to get our first hit in at most three guesses, and once we have a first hit it will take at most three guesses to sink the ship (two hits and potentially one miss if the first hit was not on the rightmost cell of the ship).

Bayesian (Thompson sampling)

The above thought exercise is interesting, but it involves several domain-specific insights (the optimal search grid, how to optimally sink a ship that has been located) that do not generalize immediately to two-dimensional grids that allow ships to be oriented in either rows or columns. (Note that if we restrict our ship in an $m \times n$ grid to be oriented along a row, we’re essentially playing a one-dimensional game on a row of length $m \cdot n$.) In this section we will show how to reproduce the above results using a Bayesian approach to this game of simplified Battleship. This approach will generalize to a near-optimal strategy for standard Battleship with many fewer hard coded rules than the above optimal one-dimensional strategy requires to generalize.

We begin by enumerating all possible boards for a given row/ship configuration (in this case 5/3).

ax = plot_board(all_rows_5_3)
ax.set_yticklabels([]);

We consider this set of ships as a uniform prior distribution on the set of possible board layouts. With this perspective, when we take the average in each columns, we get the prior probability that each cell contains a ship.

We already see that this Bayesian approach thhat we can conclude that the third cell must contain a ship without doing any modular arithmetic.

For a configuration with a longer row (10/3), we see that the middle six cells have the higest prior probability of containing a ship. This makes sense as there are three configurations in which each of these cells contains a ship, but only one configuration in which the first cell contains a ship, and similarly for the other cells near the edges.

all_rows_10_3 = get_all_ship_rows(10, 3)

plot_board(all_rows_10_3.mean(axis=0), cbar=True);

Looking at these probabilities, a resonable strategy would seem to be to guess the next cell with the highest probability of containing a hit. In both the 5/3 and 10/3 cases, this would result in guessing the third cell next. (By convention when many cells have the same hit probability, we will choose the one wit smallest coordinates.)

Let’s focus on the 5/3 configuration for the moment and return to 10/3 after. Suppose the true 5/3 board is as follows.

board_5_3 = np.zeros(5)
board_5_3[1:4] = 1

plot_board(board_5_3);

Guessing the third cell must result in a hit in the 5/3 configuration, so after one guess the following board has been revealed.

revealed_5_3 = np.ma.masked_all_like(board_5_3)
revealed_5_3[2] = board_5_3[2]

plot_board(revealed_5_3);

With this information, what should our next guess be? For the first guess we chose the cell with the highest prior probability of yielding a hit, so it seems intuitive that we would choose the cell with the highest posterior probability of yielding a hit, given the observation from our first guess.

We can get this posterior distribution by eliminating the boards that do not match the cells that have been revealed to us.

def is_compat(ships, revealed, board_axis=(-2, -1)):
    return (ships == revealed).all(axis=board_axis)

compat_5_3 = all_rows_5_3[
    is_compat(all_rows_5_3, revealed_5_3, board_axis=-1)
]

ax = plot_board(compat_5_3)
ax.set_yticklabels([]);

Since the third cell is guaranteed to be a hit in the 5/3 configuration, all ship layouts are compatible with the result of our first guess, and the posterior is the same as the prior restricted to unknown cells. With this restriction our posterior is as follows.

plot_board(
    np.ma.masked_array(compat_5_3.mean(axis=0),
                       mask=~revealed_5_3.mask),
    cbar=True
);

We see that out of all the unknown cells, the second and third are tied for the highest posterior probability of being a hit. By our convention, we will choose the second cell as the next guess, which we know will reveal a hit.

revealed_5_3[1] = board_5_3[1]

plot_board(revealed_5_3);

With this observations we are able to eliminate one of the potential boards.

compat_5_3 = all_rows_5_3[
    is_compat(all_rows_5_3, revealed_5_3, board_axis=-1)
]

ax = plot_board(compat_5_3)
ax.set_yticklabels([]);

Combining these two compatible boards to form the posterior distribution, we get the following.

ax = plot_board(
    np.ma.masked_array(compat_5_3.mean(axis=0),
                       mask=~revealed_5_3.mask),
    cbar=True
)
ax.set_yticklabels([]);

We see that it is impossible for the fourth cell to contain a hit, and there’s a 50% chance that the first and fourth cells will be a hit. By our convention above, we guess the first cell, resulting in a miss.

revealed_5_3[0] = board_5_3[0]

plot_board(revealed_5_3);

Finally we see that there is only one board compatible with the information that has been revealed, so we guess the fourth cell to end the game.

compat_5_3 = all_rows_5_3[
    is_compat(all_rows_5_3, revealed_5_3, board_axis=-1)
]

ax = plot_board(compat_5_3)
ax.set_yticklabels([]);

We now return to the 10/3 case. Suppose we are working with the following board.

board_10_3 = np.zeros(10)
board_10_3[5:8] = 1

plot_board(board_10_3);

Recalling the prior distribution on hits in the 10/3 configuration, we see first guess the third cell, revealing a miss.

plot_board(all_rows_10_3.mean(axis=0), cbar=True);

revealed_10_3 = np.ma.masked_all_like(board_10_3)
revealed_10_3[2] = board_10_3[2]

plot_board(revealed_10_3);

Excluding boards that are incompatible with a miss in the third cell, we see that we have already ruled out a ship in the first two cells as well.

compat_10_3 = all_rows_10_3[
    is_compat(all_rows_10_3, revealed_10_3, board_axis=-1)
]

ax = plot_board(compat_10_3)
ax.set_yticklabels([]);

Unsurprisingly, the sixth through eight cells are tied for the highest probability of yielding a hit, so we guess the six cell, revealing a hit.

plot_board(compat_10_3.mean(axis=0), cbar=True);

revealed_10_3[5] = board_10_3[5]

plot_board(revealed_10_3);

Hitting the rest of the ship will proceed similarly to the 5/3 case, so we stop playing here.

This approach of calculating the posterior distribution by enumerating the possible grids and eliminating those that are incompatible with the observed hits and misses is a very simple type of approximate Bayesian computation (ABC). ABC is a fascinating and deep field of research in its own right, and we are really only scratching its surface here.

Choosing to guess the cell with the maximum a posteriori (MAP) is a simplified (greedy) form of Thompson sampling, a Bayesian approach to playing games of incomplete information that chooses the next action according to the probability that it maximizes the expected reward. In the simple situation of a single ship (in both one and two dimensions) we can calculate the cell(s) that has the MAP probability chance of yielding a hit (the expected reward) exactly, so this strategy becomes greedy. (I have a soft spot for Thompsons sampling in general as we use it in a number of our machine-learning based products at Kibo, where I work. For a discussion of the applications of Thompson sampling to e-commerce optimization, take a look at my talks.)

def argmax_2d(x):
    max_i = x.max(axis=1).argmax()
    max_j = x[max_i].argmax()
    
    return max_i, max_j

class SingleThompsonStrategy(Strategy):
    def __init__(self, poss_ships):
        self._poss_ships = poss_ships
    
    def compat_ships(self, revealed):
        if revealed.mask.all():
            return self._poss_ships
        else:
            is_compat_ = is_compat(self._poss_ships, revealed)

            return self._poss_ships[is_compat_]
    
    def next_guess(self, revealed):
        post = np.ma.masked_array(
            self.compat_ships(revealed).mean(axis=0),
            mask=~revealed.mask
        )
        
        return argmax_2d(post)

We watch this strategy play the game we started above.

ts = SingleThompsonStrategy(all_rows_10_3[:, np.newaxis, :])
game_10_3 = play(board_10_3[np.newaxis, np.newaxis, :], ts,
                 progress_bar=True)

100%|██████████| 3/3 [00:00<00:00, 391.81it/s]

%%capture
ani = animate_boards(game_10_3.turn_revealed)

HTML(ani.to_html5_video())

Your browser does not support the video tag.

As with our optimal row strategy, we can play every possible 10/3 game with the Thompson sampling strategy and compare the distributions of the number of turns.

ts_row_games_10_3 = [
    play(row[np.newaxis, np.newaxis],
         SingleThompsonStrategy(all_rows_10_3[:, np.newaxis, :]))
        for row in all_rows_10_3
]
ts_row_turns_pmf_10_3 = Pmf.from_seq([
    game.turns for game in ts_row_games_10_3
])

fig, (opt_ax, ts_ax) = plt.subplots(nrows=2, sharex=True, sharey=True,
                                    figsize=(FIG_WIDTH, 2 * FIG_HEIGHT))

plot_turn_dist(opt_row_turns_pmf_10_3, color='C0', ax=opt_ax);

make_pct_yaxis(opt_ax);
opt_ax.set_title("Optimal play");

plot_turn_dist(ts_row_turns_pmf_10_3, color='C1', ax=ts_ax);

make_pct_yaxis(ts_ax);
ts_ax.set_title("Thompson sampling");

fig.suptitle("Ship size 3, row length 10");
fig.tight_layout();

It turns out that for a single row with a single ship, the Thompson sampling strategy is optimal!

Image credit SINMANTYX

It is extremely cool to me that we can reproduce the optimal strategy without too much domain-specific thought, but just by enumerating all possible boards and sequentially eliminating the ones that are no longer compatible with the observed (revealed) data! This result is particularly exciting because, as we will see, the generalization of the row-optimal strategy to a two-dimensional grid is not straightforward (even for just one ship), but the Thompson sampling approach generalizes much more readily.

Square, one ship

We now turn to the case of a square grid with one ship on it to see how these strategies generalize to something closer to standard Battleship.

Random guessing

Once again the number of turns it takes to sink the ship follows a negative hypergeometric distribution. For the 5/5/3 configuration, this distribution is shown below.

random_5_5_3_pmf = get_random_guess_dist(5**2, 3)

ax = plot_turn_dist(random_5_5_3_pmf, kind='line', mean=True)

make_pct_yaxis(ax);
ax.set_title(r"5/5/3");

There are not really any surprises here; the geometry of the board and the ability to have ships orient along rows or columns doesn’t affect the random guessing strategy at all.

Near optimal

We now construct a near optimal strategy for one ship on a two-dimensional grid based on trying to reduce the problem to search on a one-dimensional row/column as quickly as possible. We call this strategy “near optimal” because we will see that it takes slightly more turns, on average, than a Thompson sampling strategy to sink the ship.

It is tempting to treat an $m \times n$ two-dimensional grid as a $m \cdot n$ one-dimensional grid and use the search grid from the above strategy, but this grid will only minimize the turns-to-first-hit when $m$ and $n$ are relatively prime. To illustrate this fact, we show the results of wrapping the search grid for the 100/2 configuration to the 10/10/2 configuration below.

bad_search_j_1d = get_search_cells_1d(100, 2)
bad_search_i, bad_search_j = np.divmod(bad_search_j_1d, 10)
bad_search_grid = np.zeros((10, 10))
bad_search_grid[bad_search_i, bad_search_j] = 1

ax = plot_board(bad_search_grid)
ax.set_title(r"Don't do this!", fontweight='bold');

Clearly this search strategy can fail to locate a column-oriented ship in an even-numbered column even after fifty guesses!

Instead, the optimal grid to search for the initial hit (at least for a square grid) is to cycle through the equivalence classes modulo the ship size in each row, as shown below for both the 10/10/2 and the 5/5/3 configurations.

def get_search_cells_2d(grid_shape, ship_size):
    i, j = np.indices(grid_shape)
    in_grid = (i - 1) % ship_size == j % ship_size
    
    return i[in_grid], j[in_grid]

search_i, search_j = get_search_cells_2d((10, 10), 2)
search_grid = np.zeros((10, 10))
search_grid[search_i, search_j] = 1

ax = plot_board(search_grid)
ax.set_title(r"Ship size 3, $5 \times 5$ grid");

search_i, search_j = get_search_cells_2d((5, 5), 3)
search_grid = np.zeros((5, 5))
search_grid[search_i, search_j] = 1

ax = plot_board(search_grid)
ax.set_title(r"Ship size 3, $5 \times 5$ grid");

We can tell that these search grids are optimal because each cell highlighted in red is one cell less than the ship shize away from any other hightlighted cell.

Once a hit has been found by searching these cells, we have to determine if the ship is oriented along the row or column passing through that cell. For hits near a boundary (in cells less than a ship’s length away from either the top/bottom or left/right edge) we only have to test one cell to decide which way the ship is oriented. For hits at the interior we have to, in general, test two adjacent cells (either left and right or above and below) to determine the orientation of the ship.

The following plot highlights the cells that only require one test guess to determine the row/column orientation following a hit for the 10/10/3 configuration.

near_edge = np.zeros((10, 10))
i, j = np.indices((10, 10))
near_edge[(i < 2) | (7 < i) | (j < 2) | (7 < j)] = 1

plot_board(near_edge);

The strategy outlined here is implemented below.

class SingleGridStrategy(Strategy):
    def __init__(self, ship_size):
        self._ship_size = ship_size

    def next_guess(self, revealed):
        if revealed.mask.all() or revealed.sum() == 0:
            next_unmasked = (revealed.mask
                                     [search_i, search_j]
                                     .argmax())

            return search_i[next_unmasked], search_j[next_unmasked]
        else:
            grid_length, _ = revealed.shape
            hit_i, hit_j = argmax_2d(revealed)
            ship_size = self._ship_size
            
            if revealed.sum() == 1:                
                if hit_i < ship_size - 1 or grid_length - ship_size < hit_i:
                    next_i = hit_i + (1 if hit_i < ship_size - 1 else -1)
                    
                    if revealed.mask[next_i, hit_j]:
                        return next_i, hit_j
                    else:
                        is_row = revealed[next_i, hit_j] == 0
                elif hit_j < ship_size - 1 or grid_length - ship_size < hit_j:
                    next_j = hit_j + (1 if hit_j < ship_size - 1 else -1)
                    
                    if revealed.mask[hit_i, next_j]:
                        return hit_i, next_j
                    else:
                        is_row = revealed[hit_i, next_j] == 1
                elif revealed.mask[hit_i, hit_j - 1]:
                    return hit_i, hit_j - 1
                elif revealed.mask[hit_i, hit_j + 1]:
                    return hit_i, hit_j + 1
                else:
                    is_row = False
            else:                
                is_row = revealed[hit_i].sum() > 1
                
            if is_row:
                return hit_i, next_guess_1d_with_hit(revealed[hit_i])
            else:
                return next_guess_1d_with_hit(revealed[:, hit_j]), hit_j

This strategy involves quite a bit of branching logic and more than a few magic numbers. While I do not doubt that this implementation can be simplified, avoiding complex branching logic is one of the motivations for a simulation-based Bayesian approach to Battleship.

To understand the distribution of turns required by this strategy to sink the ship, we define a function that generates all square grids of a given shape with a single ship of a given size.

def get_all_ships(grid_length, ship_size):
    rows = get_all_ship_rows(grid_length, ship_size)
    n_rows, _ = rows.shape
    
    i = np.arange(grid_length)

    boards = np.zeros((
        n_rows, grid_length,
        grid_length, grid_length,
    ))
    boards[:, i, i, :] = rows[:, np.newaxis]
    long_boards = boards.reshape((-1, grid_length, grid_length))
    
    return np.concatenate((
        long_boards,
        long_boards.transpose(0, 2, 1)
    ))

We generate all possible boards for the 5/5/3 configuration and visualize a few of them.

all_ships_5_5_3 = get_all_ships(5, 3)

fig, axes = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True,
                         figsize=(FIG_WIDTH, FIG_WIDTH))

plt_boards = rng.choice(all_ships_5_5_3, size=axes.size, replace=False)

for (board, ax) in zip(plt_boards, axes.flat):
    plot_board(board, ax=ax);

fig.suptitle("Four boards, 5/5/3");
fig.tight_layout();

We now play all 5/5/3 games using this strategy.

grid_games_5_5_3 = [
    play(ship[np.newaxis], SingleGridStrategy(3))
        for ship in all_ships_5_5_3
]
grid_turns_pmf_5_5_3 = Pmf.from_seq([
    game.turns for game in grid_games_5_5_3
])

ax = plot_turn_dist(grid_turns_pmf_5_5_3,
                    color='C1', label="Near optimal",
                    mean=True, mean_kwargs={'c': 'C1'});
plot_turn_dist(random_5_5_3_pmf,
               kind='line', c='C0',
               mean=True, mean_kwargs={'c': 'C0'},
               label="Random",
               ax=ax)

make_pct_yaxis(ax);

ax.legend();

We see that this near optimal strategy is an significant improvement over random guessing (which is expected). Intuitively, it seems odd that the distribution is bimodal. In fact, we will see after implementing Thompson sampling in the two-dimensional case, that this strategy is definitely not optimal.

Bayesian (Thompson sampling)

The generalization of one-dimensional Thompson sampling to a two-dimensional grid is much simpler than the above generalization of the one-dimensional optimal strategy. The implementation of Thompson sampling for a single ship in two dimensions is shown below.

class SingleThompsonStrategy(Strategy):
    def __init__(self, poss_ships):
        self._poss_ships = poss_ships
    
    def compat_ships(self, revealed):
        if revealed.mask.all():
            return self._poss_ships
        else:
            is_compat_ = is_compat(self._poss_ships, revealed)

            return self._poss_ships[is_compat_]
    
    def next_guess(self, revealed):
        post = np.ma.masked_array(
            self.compat_ships(revealed).mean(axis=0),
            mask=~revealed.mask
        )
        
        return argmax_2d(post)

We now play all 5/5/3 games with this strategy to understand its turn distribution.

ts_games_5_5_3 = [
    play(ship[np.newaxis], SingleThompsonStrategy(all_ships_5_5_3))
        for ship in all_ships_5_5_3
]
ts_turns_pmf_5_5_3 = Pmf.from_seq([
    game.turns for game in ts_games_5_5_3
])

As alluded to above, we see that the turn distribution for Thompson sampling is concentrated to the left of the turn distribution for our near optimal strategy.

ax = plot_turn_dist(grid_turns_pmf_5_5_3,
                    color='C1', label="Near optimal",
                    mean=True, mean_kwargs={'c': 'C1'});
plot_turn_dist(random_5_5_3_pmf,
               kind='line', c='C0',
               mean=True, mean_kwargs={'c': 'C0'},
               label="Random",
               ax=ax)
plot_turn_dist(ts_turns_pmf_5_5_3,
               color='C2', label="Thompson sampling",
               mean=True, mean_kwargs={'c': 'C2'},
               ax=ax);

make_pct_yaxis(ax);

ax.legend(loc='upper right');

Note that while the turn distribution for Thompson sampling is concentrated further to the left, the expected numbers of turns for Thompson sampling and our near optimal strategy, indicated by the dashed vertical lines, are quite close.

No doubt we could improve the logic of our near optimal strategy to match (or perhaps even beat) Thompson sampling, but the generalization of Thompson sampling to the two-dimensional grid is so much more elegant that it seems hardly worthwhile.

Standard battleship

Now that we understand these approaches to simplified versions of Battleship, we are ready to tackle the Thompson sampling approach to standard Battleship.

Random

From above we recall the turn distribution for random guessing at standard Battleship.

std_turn_ax.figure

Optimistic expected values

Generalizing the optimal one-dimensional strategy to two dimensions was sufficiently difficult that we do not attempt to define even a near optimal strategy for standard Battleship. Configurations where several ships are adjacent, such as the one shown below, will certainly lead to complex and bug-prone branching logic.

ships = np.zeros((n_ships, GRID_LENGTH, GRID_LENGTH))
ships[0, 0, :5] = 1
ships[1, 1, :4] = 1
ships[2, 2, :3] = 1
ships[3, 3, :3] = 1
ships[4, 4, :2] = 1

fig, axes = plt.subplots(ncols=n_ships, sharex=True, sharey=True,
                         figsize=(n_ships * 2.5, 2.5))

for (ship, name, ax) in zip(ships, SHIP_NAMES, axes):
    plot_board(ship, ax=ax);
    
    ax.set_xticklabels([]);
    ax.set_yticklabels([]);
    ax.set_title(name);
    
fig.tight_layout();

plot_board(to_board(ships));

In spite of these challenges, we will want to at least be able to benchmark the average number of turns Thompson sampling requires against other approaches. Fortunately we can propose several scenarios for which the expected number of turns is straightforward to calculate.

Extremely optimistic

We can make two assumptions to get a reasonable floor on the expected number of turns.

The first five hits we get are each from a different ship.
As soon as we hit a ship, we know the complete location of that ship.

Together with a search strategy for the first five hits, we can calculate the expected value in this unrealistic but instructive scenario. Weakening these assumptions also leads to more realistic expected values to which we can compare the performance of Thompson sampling.

A fairly efficient search strategy for the first five hits in standard Battle ship is to search a mod-two grid (since the smallest ship has size two).

search_i, search_j = get_search_cells_2d((10, 10), 2)
search_grid = np.zeros((10, 10))
search_grid[search_i, search_j] = 1

ax = plot_board(search_grid)
ax.set_title(r"Standard Battleship");

The distribution of the number of turns it takes to get five hits in these fifty cells is given below.

extreme_hit_pmf = get_random_guess_dist(search_grid.sum(),
                                        np.ceil(SHIP_SIZES / 2).sum(),
                                        n_hit=SHIP_SIZES.size)

ax = plot_turn_dist(extreme_hit_pmf, mean=True)

make_pct_yaxis(ax);
ax.set_title("Finding five ships on a\n"
             r"$10 \times 10$ grid (optimistically)");

Under the second assumption, it will take us 12 more guesses to sink all the ships, so the expected number of turns to solve the game (optimistically) is

extreme_ev = extreme_hit_pmf.mean() + (SHIP_SIZES - 1).sum()

extreme_ev

40.18181818181814

Under the slightly more conservative (but still optimistic) assumption that it takes on average one and a half misses to determine the location and orientation of a ship, the expected value becomes

optim_ev = extreme_hit_pmf.mean() + 1.5 * SHIP_SIZES.size + (SHIP_SIZES - 2).sum()

optim_ev

42.68181818181814

We will use these two figures to benchmark the performance of Thompson sampling for standard Battleship in the next section.

Bayesian (Thompson sampling)

Generalizing the Thompson sampling strategy from one ship on a two-dimensional grid to standard Battleship is slighly more conceptually involved than generalize from one- to two-dimensional grids with one ship, but not too much more.

The generalization also requires some computational optimizations, as we can no longer expect to enumerate and store all possible configurations of the board on reasonably sized hardware. To illustrate this and facilitate our eventual solution to this issue, we generate all possible configurations of each ship.

all_ships = [get_all_ships(GRID_LENGTH, ship_size) for ship_size in SHIP_SIZES]

Even though we can enumerate the possibilities for each ship easily,

n_ship_grids = np.array([len(ships) for ships in all_ships])

n_ship_grids

array([120, 140, 160, 160, 180])

there are more than $10^{10}$ possible combinations in the cartesian product of these configurations

np.log10(n_ship_grids).sum()

10.88882175214102

Now, not all elements of the Cartesian product produce valid boards (the ships may overlap). An example of such an element is given below.

ships = np.zeros((n_ships, GRID_LENGTH, GRID_LENGTH))
ships[0, 0, :5] = 1
ships[1, 0, :4] = 1
ships[2, 0, :3] = 1
ships[3, 0, :3] = 1
ships[4, 0, :2] = 1

fig, axes = plt.subplots(ncols=n_ships, sharex=True, sharey=True,
                         figsize=(n_ships * 2.5, 2.5))

for (ship, name, ax) in zip(ships, SHIP_NAMES, axes):
    plot_board(ship, ax=ax);
    
    ax.set_xticklabels([]);
    ax.set_yticklabels([]);
    ax.set_title(name);
    
fig.tight_layout();

plot_board(to_board(ships), vmax=5.,
           cbar=True,
           cbar_kwargs={
               'format': '%d',
               'label': "Number of ships"
           });

The following function tests to see if an array of ship grids constitutes a valid board (has no overlap).

def has_no_overlap(ships, ship_axis=0, board_axis=(-2, -1)):
    return (to_board(ships, ship_axis=ship_axis)
                .max(axis=board_axis) \
                == 1)

has_no_overlap(ships)

False

Even eliminating sets of ships that do overlap, there are still too many possible boards to enumerate. To address this challenge we will switch from Thompson enumeration, which is a more proper description of what we’ve been doing so far to Thompson sampling. Instead of enumerating all boards compatible with the currently revealed information and choosing the unknown cell with the highest posterior probability of yielding a hit, we will draw a sample from the posterior distribution of compatible boards, calculate the cell most likely to yield a hit based on this sample to use as the next guess.

This sampling approach means we only need to be able to simulate random games of Battleship and check if they are compatible with the currently revealed information in order to generalize Thompson sampling to standard Battleship. Simulating random games of Battleship is not particularly hard, but simulating games compatible with boards that have many spots revealed can take a prohibitively long time if done naively by sampling from each ship’s possible configurations and then testing to see if the product of the samples is compatible. (I have tried this and had it take over an hour to generate a compatible sample with approximately 30 cells revealed.)

Fortunately we can address this issue by intelligently propagating hits and misses from the aggregate board to each ship’s grid and restricting to a subset of each ship’s position that is more likely to produce a compatible board when combined with all the other ship samples.

Before walking through this process, we introduce the machinery necessary to sample random (compatible) boards.

def _sample_compat_ships_block(poss_ships, revealed,
                               rng=None, block_size=10_000):
    if rng is None:
        rng = np.random.default_rng()

    samples = np.stack([
            rng.choice(ships, size=block_size, shuffle=False)
                for ships in poss_ships
        ],
        axis=1
    )
    
    has_no_overlap_ = has_no_overlap(samples, ship_axis=1)
    valid_samples = samples[has_no_overlap_]
    
    if revealed.mask.all():
        return valid_samples
    else:
        is_compat_ = is_compat(to_board(valid_samples, ship_axis=1), revealed)
    
        return valid_samples[is_compat_]
    
def sample_compat_ships_block(poss_ships, revealed,
                              rng=None, block_size=10_000):
    block = []
    
    while len(block) == 0:
        block = _sample_compat_ships_block(
            poss_ships, revealed,
            rng=rng, block_size=block_size
        )
        
    return block


def generate_compat_ships(poss_ships, revealed,
                          rng=None, block_size=10_000):
    while True:
        yield from sample_compat_ships_block(
            poss_ships, revealed,
            rng=rng, block_size=block_size
        )

def take(n, gen, progress_bar=False):
    n_range = trange(n) if progress_bar else range(n)
    
    return [x for _, x in zip(n_range, gen)]

def sample_compat_ships(poss_ships, revealed, n,
                        rng=None, block_size=10_000,
                        progress_bar=False):
    compat_gen = generate_compat_ships(poss_ships, revealed,
                                       rng=rng)
    
    return np.array(take(n, compat_gen,
                         progress_bar=progress_bar))

def sample_ships(poss_ships, n,
                 rng=None, block_size=10_000,
                 progress_bar=False):
    empty = np.ma.masked_all_like(all_ships[0][0])
    
    return sample_compat_ships(poss_ships, empty, n,
                               rng=rng, block_size=block_size,
                               progress_bar=progress_bar)

We plot four randomly sampled boards below.

ships = sample_ships(all_ships, 4, rng=rng)

fig, axes = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True,
                         figsize=(FIG_WIDTH, FIG_WIDTH))

for (ships_, ax) in zip(ships, axes.flat):
    plot_board(to_board(ships_), ax=ax);

fig.suptitle("Four boards, standard Battleship");
fig.tight_layout();

To illustrate how long it can take to sample compatible boards from the full Cartesian product, we randomly mask one third of the cells in the first board above and attempt to sample a compatible board.

masked_i, masked_j = rng.choice(
    np.indices((GRID_LENGTH, GRID_LENGTH)).reshape(2, -1),
    size=GRID_LENGTH**2 // 3, axis=1,
    replace=False, shuffle=False
)
mask = np.full((GRID_LENGTH, GRID_LENGTH), False)
mask[masked_i, masked_j] = True

revealed = np.ma.masked_array(to_board(ships[0]), mask=mask)

plot_board(revealed);

start = datetime.datetime.now()

try:
    sample_compat_ships(all_ships, revealed, 1, rng=rng)
except KeyboardInterrupt:
    end = datetime.datetime.now()
    
    print("You gave up trying to sample a compatible board"
          f"after {(end - start).total_seconds():.1f} seconds")

You gave up trying to sample a compatible boardafter 59.4 seconds

In order to implement the strategy for sampling compatible boards outlined above, we need to think about how to propagate hit/miss information about the combined board to each ship. For a miss, this propagation is easy, as all ships must have been missed. Propagating the information from a hit is much trickier, since in general we don’t know which ship was hit. We do however, know which ship was hit when our opponent announces that we sank one of their ships, because they have to tell us which one.

To implement Thompson sampling for standard Battleship then, we track not just what has been revealed, but what we can conclude has been revealed per ship. We use this per-ship revealed information to reduce our possibilities the compatible ship grids before sampling, then sample from the Cartesian product of these reduced sets, which produces compatible samples in at most seconds, even when many cells have been revealed.

Of course there are many situations where we can infer which ship has been hit. Consider the following board.

revealed = np.ma.masked_all((10, 10))
revealed[0, :2] = 0
revealed[1, 2] = 0
revealed[2, :2] = 0

plot_board(revealed);

If guessing (1, 0) results in a hit, we can obviously conclude that the patrol boat was hit. We do not account for situations such as this when propagating hit information from the combined board to individual ships to avoid cumbersome and error-prone conditional logic in our strategy. (In fact, Thompson sampling accounts for this automatically without a special rule, as there will be only one possible position for the patrol boat compatible with this revealed information.)

We implement Thompson sampling for standard Battleship as described above.

class ThompsonStrategy(Strategy):
    def __init__(self, all_ships, rng=None, block_size=10_000):
        self._ship_strats = [
            SingleThompsonStrategy(ships) for ships in all_ships
        ]
        
        n_ship = len(all_ships)
        grid_shape = all_ships[0][0].shape
        self._ships_revealed = np.ma.masked_all((n_ship,) + grid_shape)

        self._rng = rng if rng is not None else np.random.default_rng()
        self._block_size = block_size
    
    def next_guess(self, revealed, n=None):
        post = np.ma.masked_array(
            self.sample_post(revealed, n=n),
            mask=~revealed.mask
        )
        
        return argmax_2d(post)
    
    def sample_post(self, revealed, n=None):
        all_compat_ships = [
            ship_strat.compat_ships(ship_revealed)
                for (ship_strat, ship_revealed)
                    in zip(self._ship_strats, self._ships_revealed)
        ]
        
        if n is None:
            compat_samples = sample_compat_ships_block(
                all_compat_ships, revealed,
                rng=self._rng, block_size=self._block_size
            )
        else:
            compat_samples = sample_compat_ships(
                all_compat_ships, revealed, n=n,
                rng=self._rng, block_size=self._block_size
            )
        
        return to_board(compat_samples, ship_axis=1).mean(axis=0)

    def reveal(self, i, j, hit_or_miss, sunk):
        if hit_or_miss == 0 or sunk is not None:
            self._ships_revealed[:, i, j] = 0
            
            if sunk is not None:
                self._ships_revealed[sunk, i, j] = 1

We use this strategy to play a game and visualize the results.

strat = ThompsonStrategy(all_ships, rng=rng)
game = play(ships[0], strat,
            progress_bar=True)

100%|██████████| 17/17 [00:02<00:00, 7.06it/s]

game.turns

%%capture
ani = animate_boards(game.turn_revealed)

HTML(ani.to_html5_video())

Your browser does not support the video tag.

With a bit of extra bookeeping, we can also watch the posterior distribution used for Thompson sampling evolve after each turn.

game = Battleship(ships[0])
strat = ThompsonStrategy(all_ships, rng=rng)

posts = []

with tqdm(total=SHIP_SIZES.sum()) as pbar:
    while not game.is_solved:
        post = strat.sample_post(game.revealed, n=1_000)
        i, j = argmax_2d(np.ma.masked_array(
            post, mask=~game.revealed.mask
        ))

        hit_or_miss, sunk = game.guess(i, j)
        strat.reveal(i, j, hit_or_miss, sunk)
        
        posts.append(post)
        
        if hit_or_miss == 1:
            pbar.update()

100%|██████████| 17/17 [00:09<00:00, 1.83it/s]

%%capture
ani = animate_boards(posts, cbar=True)

HTML(ani.to_html5_video())

Your browser does not support the video tag.

I’m not really sure how interpretable this is, but it is kind of mesmerizing.

Finally, we play 1,000 random games of standard Battleship using this strategy to understand the distributions of the number of turns required.

N_GAME = 1_000

def pickleable_play(args):
    i, ships = args

    return play(
        ships,
        ThompsonStrategy(all_ships,
                         rng=np.random.default_rng(SEED + i))
    )

with mp.Pool(mp.cpu_count()) as pool:
    ts_games = [game for game in tqdm(
        pool.imap_unordered(
            pickleable_play,
            zip(range(N_GAME),
                sample_ships(all_ships, N_GAME, rng=rng))
        ),
        total=N_GAME
    )]

100%|██████████| 1000/1000 [15:32<00:00, 1.07it/s]

ts_turns = np.array([game.turns for game in ts_games])
ts_turns_pmf = Pmf.from_seq(ts_turns)

ts_turns_pmf.to_csv('./ts_turns_pmf.csv')

ax = plot_turn_dist(std_random_pmf, kind='line', mean=True,
                    label="Random guessing")
plot_turn_dist(ts_turns_pmf, color='C1',
               mean=True, mean_kwargs={'c': 'C1'},
               label="Thompson sampling", ax=ax);
ax.axvline(extreme_ev, c='C2', ls='--',
           label="Extremely optimistic\nexpected value");
ax.axvline(optim_ev, c='C3', ls='--',
           label="Optimistic\nexpected value");

ax.xaxis.set_major_locator(IndexLocator(base=5, offset=-2));

ax.set_ylim(0, 0.0825);
make_pct_yaxis(ax);

ax.legend(loc='upper right', bbox_to_anchor=(0.825, 1));
ax.set_title("Standard Battleship");

print(f"""
Extremely optimistic expected value = {extreme_ev:.2f}
Optimistic expected value = {optim_ev:.2f}
Thompson sampling expected value = {ts_turns_pmf.mean():.2f}
""")

Extremely optimistic expected value = 40.18 Optimistic expected value = 42.68 Thompson sampling expected value = 45.89

It’s very exciting to see that Thompson sampling is not too much less efficient (on average) than our (extremely) optimistic baseline expected values!

For fun, we visualize the simulated games with the fewest and most turns, to see what it looks like when we get very lucky or unlucky.

min_ts_game = ts_games[(ts_turns == ts_turns.min()).argmax()]
max_ts_game = ts_games[(ts_turns == ts_turns.max()).argmax()]

%%capture
min_ani = animate_boards(min_ts_game.turn_revealed)
max_ani = animate_boards(max_ts_game.turn_revealed)

HTML(min_ani.to_html5_video())

Your browser does not support the video tag.

HTML(max_ani.to_html5_video())

Your browser does not support the video tag.

Interestingly, we see that the fastest game three ships adjacent to each other. We likely got a bit of a speedup from finding two of those while attempting to sink the first one hit.

These simulated games naturally invite the question of whether Thompson sampling is optimal for standard Battleship (in the appropriate sense) as we have implemented it. I strongly suspect that it is near-optimal, but not quite optimal. Likely Thompson sampling combined with a few simple rules (like a restricted initial search grid) will produce slightly better (and perhaps even optimal) results. Searching the internet for prior work on near optimal Battleship strategies yields a few articles about strategies that take, on average, 42-44 guesses to complete the game. Our simple Thompson sampling approach is quite close to that range, and with a few heuristics bolted on would likely be able to shave off a few turns to reach that level of performance. We can also improve the performance of Thompson sampling at the cost of computational time by increasing the block_size used for sampling, resulting in a more accurate approximation to the posterior distribution. Another interesting feature of the strategies in those articles is that they try to balance searching (exploration) and sinking known ships (exploitation) with rules, which we do not have to do explicitly when using Thompson sampling.

Two other interesting extensions of this work would be to change the process of randomly generating ships (which is equivalent to changing the prior used in ABC) and to explore which strategies for initially placing ships take the longest for Thompson sampling to solve.

This framework is generalizable to variations of Battleship such as salvo or or limited lying is allowed, and we may explore those generalizations in future posts.

Many thanks to Kiril Zvezdarov and Meenal Jhajharia for providing helpful feedback on early drafts of this post.

This post is available as a Jupyter notebook here.

Bayesian Splines with Heteroskedastic Noise in Python with PyMC3

Austin Rochford — Wed, 11 Aug 2021 04:00:00 GMT

Splines are a powerful tool when modeling nonlinear relationships. This post shows how to include splines in a Bayesian model in Python using pymc3. In addition, we will show how to use a second spline component to handle heteroskedastic data, that is, data where the noise scale is not constant.

Image credit Wikipedia

To illustrate these concepts, we will use Lidar data from Larry Wasserman’s excellent book All of Nonparametric Statistics.

Load the Data

First we make the necessary Python imports and do some light housekeeping.

%matplotlib inline

from warnings import filterwarnings

from aesara import shared, tensor as at
import arviz as az
import matplotlib as mpl
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
import scipy as sp
import seaborn as sns

You are running the v4 development version of PyMC3 which currently still lacks key features. You probably want to use the stable v3 instead which you can either install via conda or find on the v3 GitHub branch: https://github.com/pymc-devs/pymc3/tree/v3

filterwarnings('ignore', category=UserWarning, module='arviz')

mpl.rcParams['figure.figsize'] = (8, 6)

sns.set(color_codes=True)

We are now ready to load the data.

DATA_URL = 'http://www.stat.cmu.edu/~larry/all-of-nonpar/=data/lidar.dat'

df = pd.read_csv(DATA_URL, sep=' +', engine='python')

df.head()

	range	logratio
0	390	-0.050356
1	391	-0.060097
2	393	-0.041901
3	394	-0.050985
4	396	-0.059913

df.describe()

	range	logratio
count	221.000000	221.000000
mean	554.751131	-0.291156
std	95.912396	0.282475
min	390.000000	-0.949554
25%	472.000000	-0.542305
50%	555.000000	-0.108043
75%	637.000000	-0.054825
max	720.000000	0.026907

We standardize both range and logratio to make it easier to specify priors once we begin building our spline models.

std_df = (df - df.mean()) / df.std()

std_df.head()

	range	logratio
0	-1.717725	0.852467
1	-1.707299	0.817981
2	-1.686447	0.882398
3	-1.676020	0.850240
4	-1.655168	0.818631

Exploratory Data Analysis

The task at hand is to model (standardized) logratio as a function of (standardized) range.

fig, (std_ax, joint_ax) = plt.subplots(
    nrows=2, sharex=True, sharey=False,
    gridspec_kw={'height_ratios': [1, 4]}
)

sns.scatterplot(x="range", y="logratio", data=std_df,
                alpha=0.5, ax=joint_ax);

(std_df.groupby(std_df["range"].round(1))
       ["logratio"]
       .std()
       .rolling(5)
       .mean()
       .plot(ax=std_ax));

std_ax.set_ylabel("Standard\ndeviation\n(binned)");

fig.tight_layout();

The scatter plot shows that the relationship is definitely nonlinear, and there is no obvious (to me at least) transform of logratio that will make the relationship linear. The top plot shows how the (binned, smoothed) standard deviation of logratio varies with range. As is evident from both plots, as range increases, so does the scale of variation of logratio.

Introduction to Splines

Regression splines are a type of of generalized additive model (GAM) that use linear combinations of (generally low-degree) polynomials to introduce nonlinear relationships between covariates and responses.

To begin constructing our spline model, we must choose a number of knots (also known as anchors or control points) in the domain of our co variate. In this post we will use twenty splines in the interval $[-1.75, 1.75]$, which comfortably contains the observed values of range.

N_KNOT = 20

knots = np.linspace(-1.75, 1.75, N_KNOT)

The following plot shows the location of the knots in the Lidar data.

fig, (std_ax, joint_ax) = plt.subplots(
    nrows=2, sharex=True, sharey=False,
    gridspec_kw={'height_ratios': [1, 4]}
)

sns.scatterplot(x="range", y="logratio", data=std_df,
                alpha=0.5, ax=joint_ax);
sns.rugplot(knots, height=0.075,
            c='k', label="Knots",
            ax=joint_ax);

(std_df.groupby(std_df["range"].round(1))
       ["logratio"]
       .std()
       .rolling(5)
       .mean()
       .plot(ax=std_ax));
sns.rugplot(knots, height=0.15,
            c='k', label="Knots",
            ax=std_ax);

std_ax.set_ylabel("Standard\ndeviation\n(binned)");

joint_ax.legend();
fig.tight_layout();

Let $x^*_i$, $i = 1, 2, \ldots, 20$ be the location of the $j$-th knot. The spline model we will use is given by

\[E(Y\ |\ X) = \sum_{j = 1}^{20} \beta_j \cdot B_{j, k; \mathbf{x}^*}(X)\].

(If one applies a link function to the conditional expectation of the left hand side, this becomes a generalized additive model.) For spline regression, $B_{j, k; \mathbf{x}^*}(\cdot)$ is a $k$-th-degree polynomial in $x$ and $x^*$. There are many possible choices for these functions. We will use scipy’s cubic B-spline implementation. For more information splines, consult Simon Wood’s excellent book Generalized Additive Models.

basis = sp.interpolate.BSpline(knots, np.eye(N_KNOT), 3)

We see that basis is a callable function function that will give the design matrix for spline regression at a given set of points.

hasattr(basis, '__call__')

True

We build this design matrix at the (standardized) value of range.

dmat = shared(basis(std_df["range"]))

With dmat in hand, we are ready to build our model with pymc3. We follow the model specified in Milad Kharratzadeh’s excellent short paper Splines in Stan.

The model for the conditional mean is given above,

\[\mu\ |\ X = \sum_{j = 1}^{20} \beta_j \cdot B_{j, k; \mathbf{x}^*}(X).\]

We put a Gaussian random walk prior (GRW) on these coefficients $\beta_j$, under the intuition that the coefficients for adjacent knots should be similar. We parameterize our GRW as follows:

\[ \begin{align*} \mu_{\beta} & \sim N(0, 2.5^2) \\ \Delta_{\beta, j} & \sim N(0, 1) \\ \sigma_{\beta} & \sim \mathrm{Half}-N(2.5^2) \\ \beta_j & = \mu_{\beta} + \sigma_{\beta} \cdot \sum_{i = 1}^j \Delta_{\beta, i}. \end{align*} \]

# the scale necessary to make a halfnormal distribution
# have unit variance
HALFNORMAL_SCALE = 1. / np.sqrt(1. - 2. / np.pi)

coords = {"knot": np.arange(N_KNOT)}

with pm.Model(coords=coords) as model:
    μ_β = pm.Normal("μ_β", 0., 2.5)
    Δ_β = pm.Normal("Δ_β", 0., 1., dims="knot")
    σ_β = pm.HalfNormal("σ_β", 2.5 * HALFNORMAL_SCALE)
    β = pm.Deterministic("β", μ_β + σ_β * Δ_β.cumsum(),
                         dims="knot")
    μ = at.dot(dmat, β)

Our observational model here is normal, with unknown variance $\sigma \sim \mathrm{Half-}N(2.5^2)$.

with model:
    σ = pm.HalfNormal("σ", 2.5 * HALFNORMAL_SCALE)
    obs = pm.Normal("obs", μ, σ, observed=std_df["logratio"])

We now sample from the posterior distribution of this model.

SEED = 123456789
CORES = 3

SAMPLE_KWARGS = {
    'cores': CORES,
    'random_seed': [SEED + i for i in range(CORES)],
    'return_inferencedata': True,
    'target_accept': 0.95
}

with model:
    trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [μ_β, Δ_β, σ_β, σ]

100.00% [6000/6000 02:21<00:00 Sampling 3 chains, 0 divergences]

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

None of the standard sampling diagnostics show cause for concern.

az.plot_energy(trace);

az.rhat(trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    μ_β      float64 1.002
    Δ_β      float64 1.003
    σ_β      float64 1.003
    σ        float64 1.001
    β        float64 1.002

To visualize our predictions, we sample from the posterior predictive distribution along a grid of reasonable values for range.

pp_range = np.linspace(-1.75, 1.75, 100)
dmat.set_value(basis(pp_range))

with model:
    pp_trace = pm.sample_posterior_predictive(trace)

100.00% [3000/3000 00:00<00:00]

We now plot the posterior predictions.

fig, (std_ax, joint_ax) = plt.subplots(
    nrows=2, sharex=True, sharey=False,
    gridspec_kw={'height_ratios': [1, 4]}
)

low, high = np.percentile(pp_trace["obs"], [2.5, 97.5], axis=0)
joint_ax.fill_between(pp_range, low, high,
                      color='k', alpha=0.25,
                      label="95% credible interval");

sns.scatterplot(x="range", y="logratio", data=std_df,
                alpha=0.5, ax=joint_ax);

(std_df.groupby(std_df["range"].round(1))
       ["logratio"]
       .std()
       .rolling(5)
       .mean()
       .plot(ax=std_ax, label="Actual"));

std_ax.plot(pp_range, pp_trace["obs"].std(axis=0),
            c='k', label="Posterior predictive");

joint_ax.plot(pp_range, pp_trace["obs"].mean(axis=0),
              c='k', label="Posterior expected value");

std_ax.set_ylabel("Standard\ndeviation\n(binned)");

std_ax.legend(loc='upper left', bbox_to_anchor=(0., 1.65));
joint_ax.legend(loc='lower left');
fig.tight_layout();

Visually, we appear to have captured the relationship between range and the expected value of logratio reasonably well. The credible interval and the standard deviation above are a bit odd though. We have built a homoskedastic (same-variance) observational model, so the credible interval has roughly the same width, even though the data show a small variance for small values of range, and variance increases as range does.

Accounting for heteroskedasticity

In order to remedy this issue, we will build a heteroskedastic model that allows the variance of logratio to vary with ratio. In fact, we will use a spline to model the changing variance as well.

Let $\gamma_j$ come from a GRW similar to $\beta_j$. We define

\[ \begin{align*} \eta_{\sigma}\ |\ X & = \sum_{j = 1}^{20} \gamma_j \cdot B_{j, k; \mathbf{x}^*}(X) \\ \sigma\ |\ X & = 0.05 + \exp(\eta_{\sigma}). \end{align*} \]

Note that the $0.05$ factor in the definition of $\sigma\ |\ X$ sets a lower bound on the variance, which is necessary for computational stability.

The model is a straightforward adaptation of the homoskedastic one.

dmat.set_value(basis(std_df["range"]))

with pm.Model(coords=coords) as var_model:
    β0 = pm.Normal("β0", 0., 2.5)
    Δ_β = pm.Normal("Δ_β", 0., 1., dims="knot")
    σ_β = pm.HalfNormal("σ_β", 2.5 * HALFNORMAL_SCALE)
    β = pm.Deterministic("β", β0 + Δ_β.cumsum() * σ_β,
                         dims="knot")
    μ = at.dot(dmat, β)
    
    γ0 = pm.Normal("γ0", 0., 2.5)
    Δ_γ = pm.Normal("Δ_γ", 0., 1., dims="knot")
    σ_γ = pm.HalfNormal("σ_γ", 2.5 * HALFNORMAL_SCALE)
    γ = pm.Deterministic("γ", γ0 + Δ_γ.cumsum() * σ_γ,
                         dims="knot")
    η_σ = at.dot(dmat, γ)
    σ = 0.05 + at.exp(η_σ)

    obs = pm.Normal("obs", μ, σ, observed=std_df["logratio"])

We now sample from this model.

with var_model:
    var_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β0, Δ_β, σ_β, γ0, Δ_γ, σ_γ]

100.00% [6000/6000 04:42<00:00 Sampling 3 chains, 0 divergences]

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

Again the sampling diagnostics show no cause for concern.

az.plot_energy(var_trace);

az.rhat(var_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    β0       float64 0.9998
    Δ_β      float64 1.006
    σ_β      float64 1.0
    γ0       float64 1.001
    Δ_γ      float64 1.005
    σ_γ      float64 1.0
    β        float64 1.002
    γ        float64 1.002

We do see that the values of $\gamma_j$ that correspond to small values of range have small coefficients, and the coefficients grow as range gets larger.

ax, = az.plot_forest(var_trace, var_names=["γ"])

ax.set_xlabel(r"$\gamma_j$");

ax.set_yticklabels(np.arange(N_KNOT)[::-1]);
ax.set_ylabel("$j$");

Again we sample from the posterior predictive distribution of this model.

dmat.set_value(basis(pp_range))

with var_model:
    pp_var_trace = pm.sample_posterior_predictive(var_trace)

100.00% [3000/3000 00:00<00:00]

We plot these predictions in order to compare them to those of the homoskedastic model.

fig, (std_ax, joint_ax) = plt.subplots(
    nrows=2, sharex=True, sharey=False,
    gridspec_kw={'height_ratios': [1, 4]}
)

low, high = np.percentile(pp_var_trace["obs"], [2.5, 97.5], axis=0)
joint_ax.fill_between(pp_range, low, high,
                      color='k', alpha=0.25,
                      label="95% credible interval");

sns.scatterplot(x="range", y="logratio", data=std_df,
                alpha=0.5, ax=joint_ax);

(std_df.groupby(std_df["range"].round(1))
       ["logratio"]
       .std()
       .rolling(5)
       .mean()
       .plot(ax=std_ax, label="Actual"));

std_ax.plot(pp_range, pp_trace["obs"].std(axis=0),
            c='k', label="Posterior predictive\n(homoskedastic)");
std_ax.plot(pp_range, pp_var_trace["obs"].std(axis=0),
            c='r', ls='--',
            label="Posterior predictive\n(heteroskedastic)");

joint_ax.plot(pp_range, pp_trace["obs"].mean(axis=0),
              c='k', label="Posterior predictive\n(homoskedastic)");
joint_ax.plot(pp_range, pp_var_trace["obs"].mean(axis=0),
              c='r', ls='--',
              label="Posterior predictive\n(heteroskedastic)");

std_ax.set_ylabel("Standard\ndeviation\n(binned)");

std_ax.legend(loc='upper left', ncol=3,
              bbox_to_anchor=(0., 1.6));
joint_ax.legend(loc='lower left');
fig.tight_layout();

We see that the homo- and heteroskedastic models produce essentially the same estimate of the expected value of logratio, but that the heteroskedastic model comes closer to capturing the true change in the variance.

We now compare these two models using Pareto-smoothed importance sampling leave-one-out cross-validation (PSIS-LOO).

traces = {
    'Homoskedastic': trace,
    'Heteroskedastic': var_trace
}

comp_df = az.compare(traces)

comp_df.loc[:, :"weight"]

	rank	loo	p_loo	d_loo	weight
Heteroskedastic	0	41.369658	17.759399	0.000000	1.000000e+00
Homoskedastic	1	-41.703944	14.188435	83.073602	8.058976e-11

fig, ax = plt.subplots()

az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

Interestingly, the PSIS-LOO score for the heteroskedastic model is significantly higher than that of the homoskedastic model, even though these two models predict essentially the same conditional mean.

This post is available as a Jupyter notebook here.

A Bayesian Alternative to Synthetic Control Comparative Case Studies in Python with PyMC3

Austin Rochford — Sun, 08 Aug 2021 04:00:00 GMT

Recently I had cause to be perusing Yiqing Xu’s research and dove deep into the 2020 paper A Bayesian Alternative to Synthetic Control for Comparative Case Studies¹ (BASC-CCS from now on). In an effort to better understand the method introduced in this paper, I decided to replicate (most of) the results form the simulation study described in Sections 4.1 and A.4.1 in Python using PyMC3.

BASC-CCS uses time series cross sectional data (TSCS) to infer causal effects from observational data where direct control of the treatment mechanism is not possible. As such it falls under the umbrella of causal inference methods. More specifically, it uses the Bayesian approach that treats causal inference from observational data as a problem of imputing missing (control) data for treated units.

Generate the Data

We begin by generating the data to be modeled. First we make the necessary Python imports and do some light housekeeping.

%matplotlib inline

from warnings import filterwarnings

from aesara import tensor as at
import arviz as az
import matplotlib as mpl
from matplotlib import pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
import pymc3 as pm
import scipy as sp
import seaborn as sns
import xarray as xr

You are running the v4 development version of PyMC3 which currently still lacks key features. You probably want to use the stable v3 instead which you can either install via conda or find on the v3 GitHub branch: https://github.com/pymc-devs/pymc3/tree/v3

filterwarnings('ignore', category=pm.ImputationWarning, module='pymc3')

mpl.rcParams['figure.figsize'] = (8, 6)

sns.set(color_codes=True)

BASC-CCS is designed for TSCS where a number of units are observed over a period of time. Following the BASC-CCS paper, we simulate data from $N = 50$ units over $T = 30$ time periods.

N = 50
T = 30

t = np.arange(T)

In this simulation 10% of the units will be treated starting at time $T_{\mathrm{treat}} = 21$.

T_treat = 21
P_treat = 0.1

Which units are treated will be determined by a linear combination of two unobserved latent factors, drawn from a standard normal distribution for each unit, $\Gamma_{i, 1}, \Gamma_{i, 2} \sim N(0, 1)$ (note that we use capital letters for simulated quantities and the corresponding lowercase letters for the corresponding inferred parameters).

SEED = 12345 # for reproducibility

rng = np.random.default_rng(SEED)

Γ = rng.normal(size=(N, 2))

Treatment is determined according to the variable

\[\mathrm{tr}_i = 0.7 \cdot \Gamma_{i, 1} + 0.3 \cdot \Gamma_{i, 2} + \varepsilon^{\Gamma}_i,\]

where $\varepsilon^{\Gamma}_i \sim N(0, 0.25).$

tr = Γ.dot([0.7, 0.3]) + rng.normal(0., 0.5, size=N)

The five units with the largest values of $tr_i$ are treated.

treat_crit = np.percentile(tr, 100 * (1. - P_treat))
treated = tr > treat_crit

n_treated = treated.sum()

n_treated

The following plot shows the relationship between the latent values $\Gamma_i$ and which units are treated.

CMAP = 'winter'

fig, ax = plt.subplots(figsize=(8, 8))
ax.set_aspect('equal');


norm = plt.Normalize(tr.min(), tr.max())
sns.scatterplot(x=Γ[~treated, 0], y=Γ[~treated, 1], hue=tr[~treated],
                hue_norm=norm, palette=CMAP,
                label="Untreated", legend=False,
                ax=ax);
sns.scatterplot(x=Γ[treated, 0], y=Γ[treated, 1], hue=tr[treated],
                s=100, marker='s',
                hue_norm=norm, palette=CMAP,
                label="Treated", legend=False,
                ax=ax);

Γ_mult = 1.05
Γ_min, Γ_max = Γ_mult * Γ.min(), Γ_mult * Γ.max()
plot_Γ = np.linspace(Γ_min, Γ_max, 100)

ax.plot(plot_Γ, (treat_crit - 0.7 * plot_Γ) / 0.3,
        c='k', ls='--', label="Treatment threshold\n(noiseless)");

sm = plt.cm.ScalarMappable(norm=norm, cmap=CMAP)
cbar = fig.colorbar(sm)

ax.set_xlim(Γ_min, Γ_max);
ax.set_xlabel(r"$\Gamma_{i, 1}$");

ax.set_ylim(Γ_min, Γ_max);
ax.set_ylabel(r"$\Gamma_{i, 2}$");

cbar.set_label(r"$\mathrm{tr}_i$");
ax.legend(loc='upper left');

The inclusion of the noise term $\varepsilon^{\Gamma}_i$ causes some of the treated units to be below the treatment threshold line in this plot and some of the points above the treatment threshold line to not be treated.

We now generate $K = 10$ covariates for each unit, one of which is a constant intercept and the rest of which follow a standard normal distribution.

K = 10

X = np.empty((N, K))
X[:, 0] = 1.
X[:, 1:] = rng.normal(size=(N, K - 1))

Following the paper, the first $K_* = 4$ covariates influence the outcome, and the rest are uncorrelated with it. The true regression coefficients, $B_j$ are taken from the BASC-CCS paper.

K_star = 4

B = np.zeros(K)
B[:K_star] = 3., 6., 4., 2.

The influence of each of the first $K_*$ covariates varies by unit, according to $A_{i, j} \sim N\left(0, \left(\frac{B_j}{2}\right)^2\right)$.

A = np.zeros((N, K))
A[:, :K_star] = rng.normal(0., 0.5 * B[:K_star], size=(N, K_star))

The influence of these first $K_*$ covariates also varies over time according to an $AR(1)$ process $\Xi_{j, t}$ with autocorrelation of $0.6$ and standard normal innovations.

def ar1(k, innov):
    t = np.arange(innov.shape[-1])
    expon = sp.linalg.toeplitz(t)
        
    return np.dot(innov, np.triu(np.power(k, expon)))

innov_Ξ = rng.normal(size=(K_star, T))

Ξ = np.zeros((K, T))
Ξ[:K_star] = ar1(0.6, innov_Ξ)

fig, ax = plt.subplots(figsize=(8, 6))

ax.plot(t, Ξ.T, c='k', alpha=0.75);

ax.set_xlabel("$t$");
ax.set_ylabel(r"$\Xi_{i, t}$");

The horizontal lines at zero here correspond to the $K - K_*$ covariates that have no influence on the outcome and do not vary over time.

The noise in outcomes is related to the latent parameters $\Gamma_i$ through factor loadings $F_t$, which also follow an $AR(1)$ process with autocorrelation of $0.7$ and standard normal innovations.

innov_F = rng.normal(size=(2, T))

F = ar1(0.7, innov_F)

fig, ax = plt.subplots(figsize=(8, 6))

ax.plot(t, F.T, c='k', alpha=0.75);

ax.set_xlabel("$t$");
ax.set_ylabel(r"$F_{i, t}$");

Finally, we simulate the treament effects, which are

\[ \Delta_{i, t} = \begin{cases} t - T_{\mathrm{treat}} + \varepsilon^{\mathrm{treat}}_{i, t} & \mathrm{if}\ t > T_{\mathrm{treat}} \\ 0 & \mathrm{if}\ t \leq T_{\mathrm{treat}} \end{cases} \]

where $\varepsilon^{\mathrm{treat}}_{i, t} \sim N(0, 0.25)$.

Δ = np.zeros((N, T))
Δ[:, T_treat:] = t[T_treat:] - T_treat \
                    + rng.normal(0., 0.5, size=(N, T - T_treat))

The array w indicates which units where treated at a given time, with

\[ w_{i, t} = \begin{cases} 1 & \mathrm{if\ unit}\ i\ \mathrm{is\ treated\ at\ time}\ t \\ 0 & \mathrm{otherwise} \end{cases}. \]

w = np.zeros((N, T))
w[treated, T_treat:] = 1

We plot the treatment effects that our model will attempt to recover.

fig, ax = plt.subplots(figsize=(8, 6))

ax.plot(t, (Δ * w)[~treated][0],
        c='k', label="Untreated");
ax.plot(t, (Δ * w)[treated][0],
        c='r', alpha=0.75,
        label="Treated");
ax.plot(t, (Δ * w)[treated][1:].T,
        c='r', alpha=0.75);

ax.set_xlabel("$t$");
ax.set_ylabel(r"$\Delta_{i, t} \cdot w_{i, t}$");
ax.legend(loc='upper left');

We combine each of these components to generate the data to be modeled.

Θ = (B + A)[..., np.newaxis] + Ξ[np.newaxis]
XΘ = (X[..., np.newaxis] * Θ).sum(axis=1)

We combine the reatment effects, the effects of the covariates, the effects of the latent factors, and standard normal noise to get the observed outcomes, $y_{i, t}$.

y = Δ * w + XΘ + Γ.dot(F) + rng.normal(size=(N, T))

fig, ax = plt.subplots(figsize=(8, 6))

ax.axvline(T_treat, c='k', ls='--',
           label="$T_{\mathrm{treat}}$");

ax.plot(t, y[~treated].T, c='k', alpha=0.5);
ax.plot([], [], c='k', alpha=0.5,
        label="Untreated");

ax.plot(t, y[treated].T, c='r', lw=3, alpha=0.75);
ax.plot([], [], c='r', alpha=0.75,
        label="Treated");

ax.set_xlabel("$t$");
ax.set_ylabel("$y_{i, t}$");
ax.legend(loc="upper left");

We see that the treated units have outcomes that generally trend up after $T_{\mathrm{treat}}$, but there is a lot of visual noise to interpret when comparing those to the untreated units. Thee BASC-CCS model w will build in the next section will help to quantify the difference and cut through this noise.

Modeling

Unidentified latent factors

We begin with a model that has unidentified latent factors in order to determine which factors we should constrain to best identify our model. For more details about implementing factor analysis models in PyMC see my previous post on the topic.

Since we are using the Bayesian causal-inference-as-missing-data paradigm, we define the control observations as a masked array, with entries masked when $w_{i, t} = 1$, indicating that the $i$-th unit was treated at time $t$.

y_ctrl = np.ma.array(y, mask=w)

For simplicity our model will use two latent factors, even though in general we do not know the true number of latent factors. For a more rigorous discussion of how to choose the number of latent factors, see Machine Learning, A Probabilistic Perspective², §12.3.

N_factor = 2

We define the coordinates for our parameters. (For more information on how to get PyMC3 to interact nicely with xarray via coordinates, see Oriol Abril’s excellent post on the subject.)

coords = {
    "unit": np.arange(N),                        # units
    "fact": np.arange(N_factor),                 # latent factors
    "feat": np.arange(K),                        # features
    "time": t,                                   # time
    "time_block": np.arange(T - (N_factor + 1))  # block of time for identifying factors
}

We put centered normal priors on our shared regression coefficients, $\beta_j$. This prior differs from the one in the paper, which uses a sparse prior for these coefficients. Fortunately since we are using a modular probabilistic programming library, this prior can be changed relatively easily.

with pm.Model(coords=coords) as ref_model:
    β = pm.Normal("β", 0., 5., dims="feat")

We put hierarchical normally distributed priors with mean zero (for identifiability) on the per-unit random effects $\alpha_{i, j}$ and the per-time random effects $\xi_{j, t}$.

# the scale necessary to make a halfnormal distribution
# have unit variance
HALFNORMAL_SCALE = 1. / np.sqrt(1. - 2. / np.pi)

def noncentered_normal(name, dims, μ=None,):
    if μ is None:
        μ = pm.Normal(f"μ_{name}", 0., 5.)

    Δ = pm.Normal(f"Δ_{name}", 0., 1., dims=dims)
    σ = pm.HalfNormal(f"σ_{name}", 5. * HALFNORMAL_SCALE)
    
    return pm.Deterministic(name, μ + Δ * σ)

with ref_model:
    α = noncentered_normal("α", ("unit", "feat"), μ=0.)
    ξ = noncentered_normal("ξ", ("feat", "time"), μ=0.)

Note that here we do not use the fact that the true values of $\Xi_{j, t}$ form an AR(1) process.

We combine $\beta_j$, $\alpha_{i, j}$, and $\xi_{j, t}$ to form the full coefficient cube, $\theta_{i, j, t}$.

θ = β[np.newaxis, :, np.newaxis] \
        + α[..., np.newaxis] \
        + ξ[np.newaxis, ...]

We build latent factor component of the model as in the previous post, knowing that this model specificiation is only identified up to reflections of f_unid and γ_unid.

with ref_model:
    f_pos_row = pm.HalfNormal("f_pos_row", HALFNORMAL_SCALE,
                              dims="fact")
    f_block_unid = pm.Normal("f_block_unid", 0., 1.,
                             dims=("time_block", "fact"))
    f_unid = at.concatenate((
        at.eye(N_factor),
        at.shape_padleft(f_pos_row),
        f_block_unid
    ))
    γ_unid = pm.Normal("γ_unid", 0., 1., dims=("unit", "fact"))

Finally, we specify our observational likelihood.

with ref_model:
    μ_ctrl = (X[..., np.newaxis] * θ).sum(axis=1) + γ_unid.dot(f_unid.T)
    
    σ = pm.HalfNormal("σ", 5. * HALFNORMAL_SCALE)
    obs_ctrl = pm.Normal("obs_ctrl", μ_ctrl, σ, observed=y_ctrl)

We now draw 100 samples from this unidentified model.

CORES = 3

SAMPLE_KWARGS = {
    "cores": CORES,
    "random_seed": [SEED + i for i in range(CORES)],
    "return_inferencedata": True
}

with ref_model:
    ref_trace = pm.sample(tune=100, draws=100, **SAMPLE_KWARGS)

Only 100 samples in chain.
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β, Δ_α, σ_α, Δ_ξ, σ_ξ, f_pos_row, f_block_unid, γ_unid, σ, obs_ctrl_missing]

100.00% [600/600 02:14<00:00 Sampling 3 chains, 18 divergences]

Sampling 3 chains for 100 tune and 100 draw iterations (300 + 300 draws total) took 136 seconds.
There were 18 divergences after tuning. Increase `target_accept` or reparameterize.
The acceptance probability does not match the target. It is 1.0, but should be close to 0.8. Try to increase the number of tuning steps.
The rhat statistic is larger than 1.4 for some parameters. The sampler did not converge.
The number of effective samples is smaller than 10% for some parameters.

az.rhat(ref_trace).max()

<xarray.Dataset>
Dimensions:           ()
Data variables:
    β                 float64 1.813
    Δ_α               float64 1.885
    σ_α               float64 1.814
    Δ_ξ               float64 1.86
    σ_ξ               float64 1.805
    f_pos_row         float64 1.727
    f_block_unid      float64 1.818
    γ_unid            float64 1.808
    σ                 float64 1.832
    obs_ctrl_missing  float64 1.82
    α                 float64 1.85
    ξ                 float64 1.848

As expected, the $\hat{R}$ parameters are quite high, indiciating that the chains have not converged well. This behavior is expected, since we know that this model is reflection-invariant and therefore has a multimodal posterior.

As in the previous post on identification in factor analysis, we choose a row of f_block_unid that we constrain to have all positive entries and change the signs of the entries f_block_unid and γ_unid so that they have the same relationship with the signs of this row. We choose the row that has the chain whose latent factors have the largest posterior expected distance from the origin.

block_mean = (ref_trace.posterior["f_block_unid"]
                       .mean(dim="draw"))
block_mean_norm = np.square(block_mean).sum(dim="fact")
sign_row = (block_mean_norm.max(dim="chain")
                           .argmax()
                           .item())
sign_chain = (block_mean_norm.sel({"time_block": sign_row})
                             .argmax()
                             .item())

target_sign = (np.sign(block_mean.sel({"chain": sign_chain,
                                       "time_block": sign_row}))
                 .data)

We can now specify the identified model, which is largely the same as the previous model.

with pm.Model(coords=coords) as model:
    β = pm.Normal("β", 0., 5., dims="feat")
    α = noncentered_normal("α", ("unit", "feat"), μ=0.)
    ξ = noncentered_normal("ξ", ("feat", "time"), μ=0.)
    θ = β[np.newaxis, :, np.newaxis] \
            + α[..., np.newaxis] \
            + ξ[np.newaxis, ...]
    
    f_pos_row = pm.HalfNormal("f_pos_row", HALFNORMAL_SCALE,
                              dims="fact")
    f_block_unid = pm.Normal("f_block_unid", 0., 1.,
                             dims=("time_block", "fact"))
    
    γ_unid = pm.Normal("γ_unid", 0., 1., dims=("unit", "fact"))

We now enforce the sign constraints that will break the reflectional invariance of the previous model and specify the observational likelihood.

with model:
    unid_sign = at.sgn(f_block_unid[sign_row])
    
    f_block = pm.Deterministic(
        "f_block", target_sign * unid_sign * f_block_unid,
        dims=("time_block", "fact")
    )
    f = at.concatenate((
        at.eye(N_factor),
        at.shape_padleft(f_pos_row),
        f_block
    ))
    
    γ = pm.Deterministic(
        "γ", target_sign * unid_sign * γ_unid,
        dims=("unit", "fact")
    )

    μ_ctrl = (X[..., np.newaxis] * θ).sum(axis=1) + γ.dot(f.T)
    
    σ = pm.HalfNormal("σ", 5. * HALFNORMAL_SCALE)
    obs_ctrl = pm.Normal("obs_ctrl", μ_ctrl, σ, observed=y_ctrl)

We now sample from this identified model.

with model:
    trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β, Δ_α, σ_α, Δ_ξ, σ_ξ, f_pos_row, f_block_unid, γ_unid, σ, obs_ctrl_missing]

100.00% [6000/6000 13:54<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 1_000 draw iterations (3_000 + 3_000 draws total) took 835 seconds.
The rhat statistic is larger than 1.4 for some parameters. The sampler did not converge.
The estimated number of effective samples is smaller than 200 for some parameters.

For the moment we ignore the convergence warnings, because we expect the $\hat{R}$ values for f_block_unid and γ_unid to be high.

The energy plot shows no cause for concern.

az.plot_energy(trace);

Ignoring the variables we know are not identified, all of our $\hat{R}$ statistics are reasonable as well.

rhat_var_names = [
    var_name for var_name in trace.posterior.data_vars
        if not var_name.endswith("_unid")
]

az.rhat(trace, var_names=rhat_var_names).max()

<xarray.Dataset>
Dimensions:           ()
Data variables:
    β                 float64 1.004
    Δ_α               float64 1.007
    σ_α               float64 1.003
    Δ_ξ               float64 1.008
    σ_ξ               float64 1.0
    f_pos_row         float64 1.002
    σ                 float64 1.002
    obs_ctrl_missing  float64 1.004
    α                 float64 1.008
    ξ                 float64 1.008
    f_block           float64 1.008
    γ                 float64 1.008

The model has recovered the true regression coefficients reasonably well.

ax, = az.plot_forest(trace, var_names=["β"],
                     combined=True, hdi_prob=0.95)

ax.scatter(B[::-1], ax.get_yticks(),
           c='k', zorder=5, label="Actual");

ax.set_yticklabels([]);
ax.set_ylabel(r"$\beta_j$");

ax.legend(loc='upper left');

We now turn to estimating the causal effect. First we build an array of the posterior imputed control outcomes for the treated units.

post_treat_ctrl = xr.DataArray(
    trace.posterior["obs_ctrl_missing"]
          .data
          .reshape((CORES, 1000, n_treated, T - T_treat)),
    dims=(
        "chain", "draw",
        "treat_unit", "treat_time"
    )
)

post_treat_ctrl.head()

<xarray.DataArray (chain: 3, draw: 5, treat_unit: 5, treat_time: 5)>
array([[[[ 1.25759892e+01,  1.30612264e+01,  1.57087369e+01,
           1.49536257e+01,  1.64677110e+01],
         [-4.12810485e-02,  6.26047750e+00,  7.79874263e+00,
           1.23859034e+01,  1.61490831e+01],
         [ 7.27955473e+00,  1.04034349e+01,  1.12521199e+01,
           1.11063292e+01,  1.13530200e+01],
         [-5.36650155e+00, -5.83322477e+00, -3.35144524e+00,
           3.44000735e+00,  1.88207349e+00],
         [-1.31260800e+01, -1.22691019e+01, -1.33695535e+01,
          -5.94977044e+00, -8.61557509e+00]],

        [[ 1.30278045e+01,  1.42052058e+01,  1.40630688e+01,
           1.47218969e+01,  1.93233792e+01],
         [ 2.45746861e+00,  3.67240422e+00,  6.87289651e+00,
           8.61550996e+00,  1.48857372e+01],
         [ 8.68996670e+00,  1.22336995e+01,  1.01608424e+01,
           9.78275733e+00,  1.42336802e+01],
         [-6.32208820e+00, -3.22852886e+00, -1.89722169e+00,
           3.30217994e-01,  3.78225547e+00],
         [-1.53244972e+01, -1.35265466e+01, -1.36205170e+01,
...
           1.35873648e+01,  1.50553915e+01],
         [ 3.01241871e+00,  5.42217922e+00,  6.50204080e+00,
           1.15586759e+01,  1.33643026e+01],
         [ 7.65560768e+00,  1.29871011e+01,  9.75434111e+00,
           1.09479191e+01,  1.26587440e+01],
         [-3.59855020e+00, -4.76761066e+00, -4.10285725e+00,
           2.14369279e+00,  3.28661610e-01],
         [-1.39193859e+01, -1.79066255e+01, -1.26329862e+01,
          -4.73557646e+00, -7.94774028e+00]],

        [[ 1.38132791e+01,  1.24831973e+01,  1.20970456e+01,
           1.37274249e+01,  1.51289416e+01],
         [ 3.95173624e+00,  5.69156061e+00,  6.21128104e+00,
           1.20260114e+01,  1.34918240e+01],
         [ 7.22084256e+00,  1.30855164e+01,  9.73445859e+00,
           1.10445040e+01,  1.27090074e+01],
         [-3.85094179e+00, -3.93828746e+00, -3.88376440e+00,
           1.53955903e+00,  8.58625045e-01],
         [-1.39993793e+01, -1.79798655e+01, -1.25980525e+01,
          -4.70834640e+00, -7.87338772e+00]]]])
Dimensions without coordinates: chain, draw, treat_unit, treat_time

xarray.DataArray

chain: 3
draw: 5
treat_unit: 5
treat_time: 5

12.58 13.06 15.71 14.95 16.47 … -14.0 -17.98 -12.6 -4.708 -7.873

array([[[[ 1.25759892e+01,  1.30612264e+01,  1.57087369e+01,
           1.49536257e+01,  1.64677110e+01],
         [-4.12810485e-02,  6.26047750e+00,  7.79874263e+00,
           1.23859034e+01,  1.61490831e+01],
         [ 7.27955473e+00,  1.04034349e+01,  1.12521199e+01,
           1.11063292e+01,  1.13530200e+01],
         [-5.36650155e+00, -5.83322477e+00, -3.35144524e+00,
           3.44000735e+00,  1.88207349e+00],
         [-1.31260800e+01, -1.22691019e+01, -1.33695535e+01,
          -5.94977044e+00, -8.61557509e+00]],

        [[ 1.30278045e+01,  1.42052058e+01,  1.40630688e+01,
           1.47218969e+01,  1.93233792e+01],
         [ 2.45746861e+00,  3.67240422e+00,  6.87289651e+00,
           8.61550996e+00,  1.48857372e+01],
         [ 8.68996670e+00,  1.22336995e+01,  1.01608424e+01,
           9.78275733e+00,  1.42336802e+01],
         [-6.32208820e+00, -3.22852886e+00, -1.89722169e+00,
           3.30217994e-01,  3.78225547e+00],
         [-1.53244972e+01, -1.35265466e+01, -1.36205170e+01,
...
           1.35873648e+01,  1.50553915e+01],
         [ 3.01241871e+00,  5.42217922e+00,  6.50204080e+00,
           1.15586759e+01,  1.33643026e+01],
         [ 7.65560768e+00,  1.29871011e+01,  9.75434111e+00,
           1.09479191e+01,  1.26587440e+01],
         [-3.59855020e+00, -4.76761066e+00, -4.10285725e+00,
           2.14369279e+00,  3.28661610e-01],
         [-1.39193859e+01, -1.79066255e+01, -1.26329862e+01,
          -4.73557646e+00, -7.94774028e+00]],

        [[ 1.38132791e+01,  1.24831973e+01,  1.20970456e+01,
           1.37274249e+01,  1.51289416e+01],
         [ 3.95173624e+00,  5.69156061e+00,  6.21128104e+00,
           1.20260114e+01,  1.34918240e+01],
         [ 7.22084256e+00,  1.30855164e+01,  9.73445859e+00,
           1.10445040e+01,  1.27090074e+01],
         [-3.85094179e+00, -3.93828746e+00, -3.88376440e+00,
           1.53955903e+00,  8.58625045e-01],
         [-1.39993793e+01, -1.79798655e+01, -1.25980525e+01,
          -4.70834640e+00, -7.87338772e+00]]]])

Coordinates: (0)
Attributes: (0)

We plot these against the observed treated values to visualize the treatment effect.

fig, ax = plt.subplots(figsize=(8, 6))

treated_c = [f"C{i}" for i in range(n_treated)]

ax.plot(post_treat_ctrl.mean(dim=("chain", "draw")).T,
        ls='--');
ax.set_prop_cycle(None);
ax.plot(y[treated, T_treat:].T);

ax.set_xlabel(r"$t - T_{\mathrm{treated}}$");
ax.set_ylabel("y_{i, t}");

handles = [
    Line2D([0], [0], c='k',
           label="Treated (observed)"),
    Line2D([0], [0], c='k', ls='--',
           label="Control (posterior expected value)")
]
ax.legend(loc='upper left', handles=handles);

ax.set_title("Treated units");

Here we see that the posterior estimates of the treatment effect are quite close to the actual effect.

fig, ax = plt.subplots(figsize=(8, 6))

low, high = (
    (y[treated, T_treat:] - post_treat_ctrl)
        .quantile([0.025, 0.975],
                  dim=("chain", "draw", "treat_unit"))
)
ax.fill_between(np.arange(T - T_treat), low, high,
                color='C0', alpha=0.25,
                label="95% credible interval");
ax.plot([0, T - T_treat - 1], [0, T - T_treat - 1], 
        c='k', label="Actual");
ax.plot(
    y[treated, T_treat:].T \
        - post_treat_ctrl.mean(dim=("chain", "draw")).T,
    c='C0'
);

handles, _ = ax.get_legend_handles_labels()
handles.insert(
    1, Line2D([0], [0],c='C0', label="Posterior expected")
)
ax.legend(loc='upper left', handles=handles);

ax.set_xlabel(r"$t - T_{\mathrm{treated}}$");
ax.set_ylabel("Treatment effect");

This post is available as a Jupyter notebook here.

Pang, Xun, Licheng Liu, and Yiqing Xu. “A Bayesian alternative to synthetic control for comparative case studies.” Available at SSRN (2020).↩︎
Murphy, Kevin P. Machine learning: a probabilistic perspective. MIT press, 2012.↩︎

Modeling Shark Attacks in Python with PyMC3

Austin Rochford — Sun, 27 Jun 2021 04:00:00 GMT

On a recent visit to Cape May, New Jersey I bought and read Shark Attacks of the Jersey Shore: A History, which is an interesting account of verified shark attacks in New Jersey since the nineteenth century.

While I was reading the book, I grew curious about modeling the frequency of shark attacks and went on the hunt for relevant data. There have not been many recent shark attacks in New Jersey, so I broadened my search and found the Global Shark Attack File (GSAF). The GSAF provides a fairly detailed incident log which appears to be updated every few days. This post presents an analysis of the GSAF data in Python using PyMC3. (It is worth mentioning that I am not a shark biologist, and I am sure specialists have produced much more useful and accurate models of shark attacks. Still, this seems like a fun small project as long as we don’t take the results too seriously.)

First we make some Python imports and do a bit of housekeeping.

%matplotlib inline

from functools import reduce
from warnings import filterwarnings

from aesara import tensor as at
import arviz as az
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
import seaborn as sns
import us

You are running the v4 development version of PyMC3 which currently still lacks key features. You probably want to use the stable v3 instead which you can either install via conda or find on the v3 GitHub branch: https://github.com/pymc-devs/pymc3/tree/v3

filterwarnings('ignore', category=UserWarning, module='aesara')
filterwarnings('ignore', category=UserWarning, module='arviz')
filterwarnings('ignore', category=UserWarning, module='pandas')

plt.rcParams['figure.figsize'] = (8, 6)
sns.set(color_codes=True)

Load and Transform the Data

We begin by downloading the data from the GSAF and loading the relevant columns into a Pandas data frame.

%%bash
DATA_URL=http://www.sharkattackfile.net/spreadsheets/GSAF5.xls
DATA_DEST=./data/GSAF5.xls

if [[ ! -e $DATA_DEST ]];
then
    wget -q -O $DATA_DEST $DATA_URL
fi

full_df = pd.read_excel('./data/GSAF5.xls',
                        usecols=[
                           "Case Number", "Date", "Type",
                           "Country", "Area", "Location",
                           "Fatal (Y/N)"
                        ],
                        index_col="Case Number")
full_df["Date"] = full_df["Date"].apply(pd.to_datetime, errors='coerce')

full_df.head()

	Date	Type	Country	Area	Location	Fatal (Y/N)
Case Number
2021.06.14.b	2021-06-14	Unprovoked	USA	Florida	New Smyrna Beach, Volusia County	N
2021.06.14.a	2021-06-14	Unprovoked	USA	Florida	New Smyrna Beach, Volusia County	N
2021.06.12	2021-06-12	Provoked	ENGLAND	West Sussex	Littlehampton	N
2021.06.11	2021-06-11	Unprovoked	AUSTRALIA	Western Australia	Five Fingers Reef	N
2021.05.23	2021-05-23	Unprovoked	USA	South Carolina	Burkes Beach, Hilton Head, Beaufort County	N

full_df.tail()

	Date	Type	Country	Area	Location	Fatal (Y/N)
Case Number
NaN	NaT	NaN	NaN	NaN	NaN	NaN
NaN	NaT	NaN	NaN	NaN	NaN	NaN
NaN	NaT	NaN	NaN	NaN	NaN	NaN
NaN	NaT	NaN	NaN	NaN	NaN	NaN
xx	NaT	NaN	NaN	NaN	NaN	NaN

We can already see by inspecting the data frame that there will be quite a bit of missing data to handle.

full_df.index.isnull().mean()

0.6595373706846449

full_df.isnull().mean()

Date           0.775621
Type           0.742880
Country        0.744391
Area           0.760587
Location       0.763726
Fatal (Y/N)    0.763842
dtype: float64

We begin by filtering out rows with missing data in important columns.

FILTERS = [
    ~full_df.index.isnull(),
    ~full_df["Date"].isnull(),
    ~full_df["Type"].isnull()
]

(full_df[reduce(np.logical_and, FILTERS)]
        .isnull()
        .mean())

Date           0.000000
Type           0.000000
Country        0.004843
Area           0.052413
Location       0.061927
Fatal (Y/N)    0.079398
dtype: float64

We see now that there is significantly less missing data, so we proceed to examine various aspects of the data.

ax = (full_df[reduce(np.logical_and, FILTERS)]
             ["Type"]
             .value_counts()
             .plot(kind='barh'))

ax.set_xscale('log');
ax.set_xlabel("Number of attacks");

ax.invert_yaxis();
ax.set_ylabel("Type of attack");

Unprovoked attacks are by far the most common. We will focus our analysis on this type of attack.

FILTERS.append(full_df["Type"] == "Unprovoked")

ax = (full_df[reduce(np.logical_and, FILTERS)]
             ["Country"]
             .value_counts()
             .plot(kind='barh', figsize=(8, 30)))

ax.set_xscale('log');
ax.set_xlabel("Number of unprovoked attacks");

ax.invert_yaxis();
ax.set_ylabel("Country");

While the data contain information about attacks in many countries, we will focus on the United States. This focus is due partially to the fact that I was in New Jersey when considering this problem, and partially due to the fact that finding standardized data across countries can be more challenging than within different regions of a single country.

FILTERS.append(full_df["Country"] == "USA")

ax = (full_df[reduce(np.logical_and, FILTERS)]
["Area"]
.value_counts()
.plot(kind='barh', figsize=(8, 12)))

ax.set_xscale('log');
ax.set_xlabel("Number of unprovoked attacks");

ax.invert_yaxis();
ax.set_ylabel("Country");

We see that for unprovoked attacks in the United States, Area roughly corresponds to state, with some territories included as well. Since most of the territories appear very rarely, we restrict our analysis to states so that it is easier to find information about them.

FILTERS.append(
full_df["Area"].isin([
state.name for state in us.states.STATES
])
)

Finally we look at how the number of shark attacks has changed over time.

ax = (full_df.assign(Year=full_df["Date"].dt.year)
             [reduce(np.logical_and, FILTERS)]
             ["Year"]
             .value_counts()
             .sort_index()
             .plot())

ax.set_xlabel("Year");
ax.set_ylabel("Unprovoked shark attacks\nin the United States");

We see that the number of shark attacks has increased over time. This phenomenon is likely partially due to population growth and partially due to improved reporting mechanisms for shark attacks. We willk keep a relatively modern focus an analysze shark attacks between 1960 and 2020.

YEAR_RANGE = (1960, 2020)

FILTERS.append(full_df["Date"].dt.year.between(*YEAR_RANGE))

Now that we have defined the set of attacks we will model, we produce another data frame including only these attacks and make some light transformations of the data.

df = (full_df[reduce(np.logical_and, FILTERS)]
             .copy()
             .rename(columns={"Area": "State"}))
df["Year"] = df["Date"].dt.year

df.head()

	Date	Type	Country	State	Location	Fatal (Y/N)	Year
Case Number
2020.12.30	2020-12-30	Unprovoked	USA	California	Coronado, San Diego County	N	2020
2020.12.08	2020-12-08	Unprovoked	USA	Hawaii	Honolua Bay	Y	2020
2020.12.06.b	2020-12-06	Unprovoked	USA	Oregon	Seaside Cove, Clatsop County	N	2020
2020.11.26	2020-11-26	Unprovoked	USA	Hawaii	Maui	N	2020
2020.10.31	2020-10-31	Unprovoked	USA	Florida	Ormond Beach	N	2020

df.tail()

	Date	Type	Country	State	Location	Fatal (Y/N)	Year
Case Number
1904.00.00.a	1970-01-01 00:00:00.000001904	Unprovoked	USA	Hawaii	Off Diamond Head, Honolulu, O’ahu	Y	1970
1896.00.00.b	1970-01-01 00:00:00.000001896	Unprovoked	USA	Florida	NaN	Y	1970
1883.00.00.a	1970-01-01 00:00:00.000001883	Unprovoked	USA	South Carolina	NaN	Y	1970
1882.00.00.b	1970-01-01 00:00:00.000001882	Unprovoked	USA	Florida	In the bay near the naval yard at Pensacola, E…	N	1970
1852.00.00	1970-01-01 00:00:00.000001852	Unprovoked	USA	South Carolina	Mount Pleasant, Charleston County	Y	1970

df.shape

(1514, 7)

After applying these filters, just over 1,500 attacks remain. This analysis will focus on the number of shark attacks in a state in a given year. Subsequent posts may analyze other aspects of this data. First we count the number of attacks in a given state in a given year. (We will shorten the phrase “unprovoked shark attacks in the United States between 1960 and 2020” to “attacks” for the remainder of the post.)

attacks_nz = (df.groupby(["State", "Year"])
                .size()
                .rename("Attacks"))

attacks_nz.describe()

count    333.000000
mean       4.546547
std        6.527174
min        1.000000
25%        1.000000
50%        2.000000
75%        5.000000
max       38.000000
Name: Attacks, dtype: float64

The series attacks_nz includes a row only when there was at least one attack in that state in that year. We also want to include zero entries for state/year combinations that saw now attacks, which we do now by reindexing attacks_nz.

attacks_index = (pd.MultiIndex.from_product((
                        attacks_nz.index
                                  .get_level_values("State")
                                  .unique(),
                        YEAR_RANGE[0] + np.arange(attacks_nz.index
                                                            .get_level_values("Year")
                                                            .values
                                                            .ptp())))
                   .rename("Year", level=1))
attacks_df = (attacks_nz.reindex(attacks_index, fill_value=0)
                        .astype(np.int64)
                        .to_frame())

attacks_df.head()

		Attacks
State	Year
Alabama	1960	0
	1961	0
	1962	0
	1963	0
	1964	0

attacks_df.tail()

		Attacks
State	Year
Washington	2015	0
	2016	0
	2017	1
	2018	0
	2019	0

Modeling

We now turn to modeling the data.

ax = attacks_df["Attacks"].hist(bins=attacks_df["Attacks"].max() + 1)

ax.set_xlabel("Number of attacks");

ax.set_yscale('log');
ax.set_ylabel("Number of state-years");

We see that the vast majority of state-years have no shark attacks, and that when there is at least one attack, there are rarely very many. We see that the index of dispersion is significantly larger than one, so the data shows overdispersion.

attacks_df.var() / attacks_df.mean()

Attacks    12.729978
dtype: float64

Negative Binomial

Due to this overdispersion, we begin with a Negative Binomial model.

Let $y_{i, t}$ represent the number of attacks in the $i$-th state in year $t$. We use the priors

\[ \begin{align*} \mu & \sim \operatorname{Lognormal}(0, 2.5^2) \\ \alpha & \sim \operatorname{Half}-N(2.5^2). \end{align*} \]

with pm.Model() as nb_model:
    μ = pm.Lognormal("μ", 0., 2.5)
    α = pm.HalfNormal("α", 2.5)

We now let $y_{i, t} \sim NB(\mu, \alpha)$

y = attacks_df["Attacks"].values

with nb_model:
    obs = pm.NegativeBinomial("obs", μ, α, observed=y)

We use pymc3 to sample from the posterior distribution of this model.

CHAINS = 3
SEED = 12345

SAMPLE_KWARGS = {
    'cores': CHAINS,
    'random_seed': [SEED + i for i in range(CHAINS)],
    'return_inferencedata': True
}

with nb_model:
    nb_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [μ, α]

100.00% [6000/6000 00:15<00:00 Sampling 3 chains, 0 divergences]

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

The standard sampling diagnostics (energy plots, BFMI, and $\hat{R}$) show no cause for concern.

def make_diagnostic_plots(trace, axes=None, min_mult=0.995, max_mult=1.005):
    if axes is None:
        fig, axes = plt.subplots(ncols=2,
                                 sharex=False, sharey=False,
                                 figsize=(16, 6))
        
    az.plot_energy(trace, ax=axes[0])
    
    
    rhat = az.rhat(trace).max()
    axes[1].barh(np.arange(len(rhat.variables)), rhat.to_array(),
                 tick_label=list(rhat.variables.keys()))
    axes[1].axvline(1, c='k', ls='--')

    axes[1].set_xlim(
        min_mult * min(rhat.min().to_array().min(), 1),
        max_mult * max(rhat.max().to_array().max(), 1)
    )
    axes[1].set_xlabel(r"$\hat{R}$")

    axes[1].set_ylabel("Variable")
    
    return fig, axes

make_diagnostic_plots(nb_trace);

Examining the posterior distribution for $\mu$ shows that the model will predicts more than one shark attack per year on average.

az.plot_posterior(nb_trace, var_names=["μ"]);

This prediction of about one shark attack per state per year captures nationwide average well enough, but is clearly not very useful at the state level. We see specifically how bad it is by examining the posterior predictions of the model.

with nb_model:
    nb_pp_trace = pm.sample_posterior_predictive(nb_trace)

100.00% [3000/3000 00:05<00:00]

First we produce a long data frame where each row represents a single posterior predictive sample from the distribution of the number of attacks in a given state-year. This data frame will be the basis for all of our posterior predictive plots.

def make_pp_full_df(pp_obs):
    return (pd.DataFrame(
                {i: samples for i, samples in enumerate(pp_obs)},
                index=attacks_df.index
              )
              .reset_index()
              .melt(id_vars=["State", "Year"],
                    var_name="Sample")
              .set_index(["State", "Year", "Sample"]))

nb_pp_full_df = make_pp_full_df(nb_pp_trace["obs"])

nb_pp_full_df.head()

			value
State	Year	Sample
Alabama	1960	0	1
	1961	0	0
	1962	0	0
	1963	0	0
	1964	0	0

ALPHA = 0.05

def summarize_pp_df(pp_full_df, level=None, alpha=ALPHA):
    if level is None:
        level = ["State", "Year"]
    
    return (pp_full_df.groupby(level=level)
                      ["value"]
                      .agg(
                          mean='mean',
                          sum='sum',
                          low=lambda s: s.quantile(alpha / 2.),
                          high=lambda s: s.quantile(1. - alpha / 2.),
                      )
                      .assign(attacks=attacks_df["Attacks"]
                                                .groupby(level=level)
                                                .sum()))

nb_pp_df = summarize_pp_df(nb_pp_full_df)

nb_pp_df.head()

		mean	sum	low	high	attacks
State	Year
Alabama	1960	1.231000	3693	0	12.000	0
	1961	1.212667	3638	0	12.000	0
	1962	1.206000	3618	0	11.025	0
	1963	1.232000	3696	0	12.000	0
	1964	1.177667	3533	0	11.000	0

The data frame nb_pp_df contains the posterior predictive mean and posterior predictive quantiles for each state-year combination. We now plot the predictions and actual data, broken down by state.

def plot_pp_over_time(data=None, y="mean", *args, **kwargs):
    pp_df = data
    
    ax = plt.gca()
    (pp_df.plot("Year", y,
                c='k', label="Posterior expected value",
                ax=ax))
    (pp_df.reset_index()
          .plot("Year", "attacks",
                kind='scatter', c='k', zorder=5,
                label="Actual", ax=ax))

    ax.fill_between(pp_df["Year"], pp_df["low"], pp_df["high"],
                    color='C0', alpha=0.5,
                    label="95% posterior credible interval")

    ax.set_ylabel("Number of shark attacks")

    return ax

grid = sns.FacetGrid(nb_pp_df.reset_index(), col="State",
                     col_wrap=2, sharey=False, aspect=1.5)
grid.map_dataframe(plot_pp_over_time);

grid.axes[0].legend(loc="upper left");
grid.set_titles("{col_name}");
grid.fig.tight_layout();

We see that the predictions for each state are indeed the same (ignoring some slight Monte Carlo variation), and therefore vastly overpredict shark attacks for most states while massively underpredicting shark attacks for some states.

The plot below which does not include a time axis further reinforces this point.

nb_pp_state_df = summarize_pp_df(nb_pp_full_df, level="State")
ax = (nb_pp_state_df.reset_index()
                    .plot.scatter("mean", "State", color='C0',
                                  xerr=nb_pp_state_df[["low", "high"]]
                                                     .sub(nb_pp_state_df["mean"],
                                                          axis=0)
                                                     .abs()
                                                     .values.T,
                                  zorder=5,
                                  label="Posterior predictive mean"))
(attacks_df.reset_index()
           .plot.scatter("Attacks", "State",
                         color='k', alpha=0.5,
                         label="Actual", ax=ax));

ax.set_xlabel("Number of shark Attacks");
ax.invert_yaxis();

The predictions from this model are particularly bad for states like California, Florida, Hawaii, and the Carolinas, where years with many shark attacks are relatively more common than they are for other states.

Negative Binomial Regression

California, Florida, and to a lesser extent the Carolinas are similar in that they have relatively large populations. All of these states are similar in that their population is relatively concentrated on the coast.

The table below shows that overdispersion is still quite prevalent when we condition on state, so will use negative binomial regression to account for the effect of state-level factors on the number of attacks.

(attacks_df.groupby(level="State")
           .var()
           .div(attacks_df.groupby("State")
                          .mean()))

	Attacks
State
Alabama	1.090395
California	1.408464
Connecticut	1.000000
Delaware	0.966102
Florida	8.373881
Georgia	1.090395
Hawaii	3.457544
Louisiana	0.915254
Maine	1.000000
Massachusetts	1.254237
Mississippi	1.000000
New Jersey	1.457627
New York	2.000000
North Carolina	2.342930
Oregon	1.156634
Rhode Island	1.000000
South Carolina	2.382694
Texas	1.917465
Virginia	0.881356
Washington	0.983051

First we load state-leve population and coastline data from Wikipedia.

POP_URL = "https://en.wikipedia.org/wiki/List_of_U.S._states_and_territories_by_historical_population"

population = (pd.read_html(POP_URL)[3]
                .iloc[:-1]
                .melt(id_vars="Name",
                      var_name="Year",
                      value_name="Population")
                .rename(columns={"Name": "State"})
                .fillna(0)
                .astype({
                    "Year": np.int64,
                    "Population": np.float64
                })
                .set_index(["State", "Year"])
                .sort_index()
                ["Population"])

population.head()

State    Year
Alabama  1960    3266740.0
         1970    3444165.0
         1980    3893888.0
         1990    4040587.0
         2000    4447100.0
Name: Population, dtype: float64

population.tail()

State    Year
Wyoming  1980    469557.0
         1990    453588.0
         2000    493782.0
         2010    563626.0
         2020    576851.0
Name: Population, dtype: float64

The series population contains the population of each state according to the United States census conducted every ten years.

COAST_URL = "https://en.wikipedia.org/wiki/List_of_U.S._states_and_territories_by_coastline"

coast_df, _ = pd.read_html(COAST_URL)
coast_df = (coast_df[["State or territory", "Method 1 (CRS)", "Coast/area ratio (ft/mi2)"]]
                    .iloc[:-1])
coast_df.columns = coast_df.columns.droplevel(0)
coast_df = coast_df.drop(["Rank", "Method 2"], axis=1)
coast_df = (coast_df.rename(columns={
                        "State or territory": "State",
                        "Method 1": "Coastline to area"
                    })
                    .set_index("State")
                    .sort_index())

coast_df.head()

	Coastline	Coastline to area
State
Alabama	53 mi (85 km)	5.3
Alaska	6,640 mi (10,690 km)	53
American Samoa	–	–
California	840 mi (1,350 km)	27
Connecticut	96 mi (154 km)	91

coast_df["Coastline"] = (
    coast_df["Coastline"]
            .str.split(expand=True)
            .iloc[:, 0]
            .str.replace(",", "")
            .str.replace("[–—]", "0", regex=True)
            .astype(np.float64)
)
coast_df["Coastline to area"] = (
    coast_df["Coastline to area"]
            .astype(str)
            .str.replace("[–—]", "-1", regex=True)
            .astype(np.float64)
            .replace(-1, np.nan)
)

coast_df.head()

	Coastline	Coastline to area
State
Alabama	53.0	5.3
Alaska	6640.0	53.0
American Samoa	0.0	NaN
California	840.0	27.0
Connecticut	96.0	91.0

coast_df.tail()

	Coastline	Coastline to area
State
U.S. Minor Outlying Islands	0.0	NaN
U.S. Virgin Islands	0.0	NaN
Virginia	112.0	14.0
Washington	157.0	12.0
Wisconsin	0.0	NaN

The data frame coast_df contains the length of a state’s coastline (in miles) as well as the ratio of its coastline to area (in square miles).

We now combine attacks, population, and coast_df into a single data frame.

attacks_df = (attacks_df.merge(coast_df,
                               left_index=True, right_index=True)
                        .merge(population, how='left',
                               left_index=True, right_index=True)
                        .fillna(method='ffill'))

attacks_df["Population to coastline"] = attacks_df["Population"] / attacks_df["Coastline"]

attacks_df.head()

		Attacks	Coastline	Coastline to area	Population	Population to coastline
State	Year
Alabama	1960	0	53.0	5.3	3266740.0	61636.603774
	1961	0	53.0	5.3	3266740.0	61636.603774
	1962	0	53.0	5.3	3266740.0	61636.603774
	1963	0	53.0	5.3	3266740.0	61636.603774
	1964	0	53.0	5.3	3266740.0	61636.603774

attacks_df.tail()

		Attacks	Coastline	Coastline to area	Population	Population to coastline
State	Year
Washington	2015	0	157.0	12.0	6724540.0	42831.464968
	2016	0	157.0	12.0	6724540.0	42831.464968
	2017	1	157.0	12.0	6724540.0	42831.464968
	2018	0	157.0	12.0	6724540.0	42831.464968
	2019	0	157.0	12.0	6724540.0	42831.464968

Here the population data is from the most recent United States census prior to the year in question (thanks to fillna(method='ffill'). Below we plot the relationship between the four explanatory variables and the number of attacks.

fig, axes = plt.subplots(ncols=2, nrows=2, sharey=True,
                         figsize=(12, 9))

for col, ax in zip(attacks_df.columns[1:], axes.ravel()):
    attacks_df.plot.scatter(col, "Attacks",
                            color='C0', alpha=0.5,
                            ax=ax);
    
    ax.set_xscale('log');
    ax.set_yscale('log');

fig.tight_layout();

The top two and bottom two plots are very similar. We choose to use coastline length (in miles) and population as our predictors since those relationships seem to be a bit more linear on the log-log scale.

We standardize the logarithms of coastline length and population before using them as predictors.

def standardize(x):
    return (x - x.mean()) / x.std()

log_coast = np.log(attacks_df["Coastline"].values)
x_coast = at.constant(standardize(log_coast))

log_pop = np.log(attacks_df["Population"].values)
x_pop = at.constant(standardize(log_pop))

We use the priors $\beta_0, \beta_{\text{coast}}, \beta_{\text{pop}} \sim N(0, 2.5^2)$ on the regression coefficients and set

\[\eta_{i, t} = \beta_0 + \beta_{\text{coast}} \cdot x_{\text{coast}, i} + \beta_{\text{pop}} \cdot x_{\text{pop}, i, t}.\] The mean is then $\mu_{i, t} = \exp \eta_{i, t}.$

with pm.Model() as nb_reg_model:
    β0 = pm.Normal("β0", 0., 2.5)
    β_coast = pm.Normal("β_coast", 0., 2.5)
    β_pop = pm.Normal("β_pop", 0., 2.5)
    η = β0 + β_coast * x_coast + β_pop * x_pop
    μ = at.exp(η)

As in the previous model, $\alpha \sim \operatorname{Half}-N(2.5^2)$ and $y_{i, t} \sim NB(\mu_{i, t}, \alpha)$.

with nb_reg_model:
    α = pm.HalfNormal("α", 2.5)
    obs = pm.NegativeBinomial("obs", μ, α, observed=y)

We again sample from the posterior distribution of the model.

with nb_reg_model:
    nb_reg_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β0, β_coast, β_pop, α]

100.00% [6000/6000 00:24<00:00 Sampling 3 chains, 0 divergences]

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

As before, the standard sampling diagnostics show no cause for concern.

make_diagnostic_plots(nb_reg_trace);

To see if the inclusion of these predictors has improved the model fit, we sample from and plot the posterior predictive distributions per state-year.

with nb_reg_model:
    nb_reg_pp_trace = pm.sample_posterior_predictive(nb_reg_trace)

100.00% [3000/3000 00:02<00:00]

nb_reg_pp_full_df = make_pp_full_df(nb_reg_pp_trace["obs"])

nb_reg_pp_df = summarize_pp_df(nb_reg_pp_full_df)
grid = sns.FacetGrid(nb_reg_pp_df.reset_index(), col="State",
                     col_wrap=2, sharey=False, aspect=1.5)
grid.map_dataframe(plot_pp_over_time);

grid.axes[0].legend(loc="upper left");
grid.set_titles("{col_name}");
grid.fig.tight_layout();

nb_reg_pp_state_df = summarize_pp_df(nb_reg_pp_full_df, level="State")
ax = (nb_reg_pp_state_df.reset_index()
                    .plot.scatter("mean", "State", color='C0',
                                  xerr=nb_reg_pp_state_df[["low", "high"]]
                                                     .sub(nb_reg_pp_state_df["mean"],
                                                          axis=0)
                                                     .abs()
                                                     .values.T,
                                  zorder=5,
                                  label="Posterior predictive mean"))
(attacks_df.reset_index()
           .plot.scatter("Attacks", "State",
                         color='k', alpha=0.5,
                         label="Actual", ax=ax));

ax.set_xlabel("Number of shark Attacks");
ax.invert_yaxis();

These predictions certainly appear visually better than those from the regular negative binomial model without predictors, but still leave significant room for improvement, espescially in states like Florida, California, Hawaii, and the Carolinas.

We use Pareto-smoothed importance sampling leave-one-out cross-validation (PSIS-LOO) to quantify the improvement of this model over the basic negative binomial model.

traces = {
    "NB": nb_trace,
    "NB regression": nb_reg_trace
}

fig, ax = plt.subplots()

comp_df = az.compare(traces)
az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

We see that the PSIS-LOO score of the regression model is significantly higher than that of the simple negative binomial model.

Mixed effects

There are many factors about each state that we have not included as predictors (average water temperature, prevalence of beach tourism, etc.). To account for this these and other unmeasured differences between states, we now add state-level mixed effects to our model.

Let $j = 1, 2, \ldots, n_{\text{state}}$ denote the index of a state and $j(i)$ denoth the index of the state of the $i$-th observation.

We place hierarchical normal priors on $\beta_{0, j},$ $\beta_{\text{coast}, j},$ and $\beta_{\text{pop}, j}$ and set

\[\eta_{i, t} = \beta_{0, j(i)} + \beta_{\text{coast}, j(i)} \cdot x_{\text{coast}, i} + \beta_{\text{pop}, j(i)} \cdot x_{\text{pop}, i, t}.\] Again the mean is then $\mu_{i, t} = \exp \eta_{i, t}.$

The hierarchical prior for $\beta_0$ is equivalent to

\[ \begin{align*} \mu_{\beta_0}, \gamma_{0, \text{pop}} & \sim N(0, 2.5^2) \\ \sigma_{\beta_0} & \sim \operatorname{Half}-N(2.5^2) \\ \beta_{0, j} & \sim N\left(\mu_{\beta_0} + \gamma_{0, \text{pop}} \cdot \bar{x}_{\text{pop}, j}, \sigma_{\beta_0}^2\right). \end{align*} \]

In reality, we implement this system of priors using a non-centered parametrization that is mathematically equivalent but often more computationally efficient. Note the inclusion of a term for $\bar{x}_{\text{pop}, j}$, which is the average population for the $j$-th state. We include this term to avoid violating the Gauss-Markov assumptions.

j_, state_map = (attacks_df.index
                          .get_level_values("State")
                          .factorize(sort=True))
n_state = state_map.size
j = at.constant(j_)

x_pop_bar = at.constant(standardize(
    attacks_df.assign(x_pop=x_pop.eval())
              .groupby(level="State")
              ["x_pop"]
              .mean()
              .values
))

def noncentered_normal(name, shape, μ=None):
    if μ is None:
        μ = pm.Normal(f"μ_{name}", 0., 2.5)

    Δ = pm.Normal(f"Δ_{name}", 0., 1., shape=shape)
    σ = pm.HalfNormal(f"σ_{name}", 2.5)
    
    return pm.Deterministic(name, μ + Δ * σ)

def noncentered_normal_with_avg(name, avg_map, shape, μ=None):
    if μ is None:
        μ = pm.Normal(f"μ_{name}", 0., 2.5)

    avg_terms_sum = 0.
    
    for term_name, x_bar in avg_map.items():
        avg_terms_sum += pm.Normal(f"γ_{name}_{term_name}_bar", 0., 2.5) * x_bar
    
    return noncentered_normal(name, shape, μ=μ + avg_terms_sum)

with pm.Model() as mixed_model:
    β0 = noncentered_normal_with_avg(
        "β0", {"pop": x_pop_bar}, n_state
    )

The prior for $\beta_{\text{coast}, j}$ is

\[ \begin{align*} \mu_{\beta_\text{coast}}, \gamma_{\text{coast}^2}, \gamma_{\text{coast}, \text{pop}} & \sim N(0, 2.5^2) \\ \sigma_{\beta_{\text{coast}}} & \sim \operatorname{Half}-N(2.5^2) \\ \beta_{\text{coast}, j} & \sim N\left(\mu_{\beta_\text{coast}} + \gamma_{\text{coast}^2} \cdot x_{\text{coast}, j} + \gamma_{\text{coast}, \text{pop}} \cdot \bar{x}_{\text{pop}, j}, \sigma_{\beta_0}^2\right) \end{align*} \]

Here $\bar{x}_{\text{coast}, j} = x_{\text{coast}, j}$ since the a state’s coastline does not change over time (we ignoring the effects of global warming here). Due to this fact, we do not include the term $\gamma_{0, \text{coast}} \cdot x_{\text{coast}, j}$ in the definition of $\beta_0$, because this term would lead to $\mu_{\beta_{\text{coast}}}$ becoming unidentified. We use the analagous system of priors for $\beta_{\text{pop}}.$

log_state_coast = np.log(coast_df["Coastline"]
                                 .loc[state_map]
                                 .values)
x_coast_bar = at.constant(standardize(log_state_coast))

with mixed_model:
    β_coast = noncentered_normal_with_avg(
        "β_coast",
        {"pop": x_pop_bar, "coast": x_coast_bar},
        n_state
    )
    β_pop = noncentered_normal_with_avg(
        "β_pop",
        {"pop": x_pop_bar, "coast": x_coast_bar},
        n_state)

The definitions of $\eta_{i, t}$, $\mu_{i, t}$, and $\alpha$ are similar to those in the previous model.

with mixed_model:
    η = β0[j] + β_coast[j] * x_coast + β_pop[j] * x_pop
    μ = at.exp(η)
    
    α = pm.HalfNormal("α", 2.5)
    obs = pm.NegativeBinomial("obs", μ, α, observed=y)

We again sample from the posterior distribution of the model.

with mixed_model:
    mixed_trace = pm.sample(**SAMPLE_KWARGS, target_accept=0.95)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [μ_β0, γ_β0_pop_bar, Δ_β0, σ_β0, μ_β_coast, γ_β_coast_pop_bar, γ_β_coast_coast_bar, Δ_β_coast, σ_β_coast, μ_β_pop, γ_β_pop_pop_bar, γ_β_pop_coast_bar, Δ_β_pop, σ_β_pop, α]

100.00% [6000/6000 24:00<00:00 Sampling 3 chains, 3 divergences]

Sampling 3 chains for 1_000 tune and 1_000 draw iterations (3_000 + 3_000 draws total) took 1441 seconds.
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.
There were 2 divergences after tuning. Increase `target_accept` or reparameterize.

The few divergences here could be eliminated by increasing target_accept and waiting longer, but we will proceed with our analysis (I’m impatient). As before, the standard sampling diagnostics show no cause for concern.

make_diagnostic_plots(mixed_trace);

To see if the inclusion of mixed effects has improved the model fit, we sample from and plot the posterior predictive distributions per state-year.

with mixed_model:
    mixed_pp_trace = pm.sample_posterior_predictive(mixed_trace)

100.00% [3000/3000 00:05<00:00]

mixed_pp_full_df = make_pp_full_df(mixed_pp_trace["obs"])

mixed_pp_df = summarize_pp_df(mixed_pp_full_df)
grid = sns.FacetGrid(mixed_pp_df.reset_index(), col="State",
                     col_wrap=2, sharey=False, aspect=1.5)
grid.map_dataframe(plot_pp_over_time);

grid.axes[0].legend(loc="upper left");
grid.set_titles("{col_name}");
grid.fig.tight_layout();

mixed_pp_state_df = summarize_pp_df(mixed_pp_full_df, level="State")
ax = (mixed_pp_state_df.reset_index()
                    .plot.scatter("mean", "State", color='C0',
                                  xerr=mixed_pp_state_df[["low", "high"]]
                                                     .sub(mixed_pp_state_df["mean"],
                                                          axis=0)
                                                     .abs()
                                                     .values.T,
                                  zorder=5,
                                  label="Posterior predictive mean"))
(attacks_df.reset_index()
           .plot.scatter("Attacks", "State",
                         color='k', alpha=0.5,
                         label="Actual", ax=ax));

ax.set_xlabel("Number of shark Attacks");
ax.invert_yaxis();

We see that the predictions for Florida, California, Hawaii, and the Carolinas are much better in this model.

traces["NB mixed"] = mixed_trace

fig, ax = plt.subplots()

comp_df = az.compare(traces)
az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

As expected the PSIS-LOO score for this model is significantly better than the fixed effects negative binomial regression model.

It is interesting but not surprising that mixed effects are required to produce reasonable predictions of the number of shark attacks in each state over time. Future posts may investigate further explanatory variables and other aspects of this data set.

This post is available as a Jupyter notebook here.

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sun Jun 27 2021

Python implementation: CPython
Python version       : 3.8.8
IPython version      : 7.22.0

numpy     : 1.20.2
matplotlib: 3.4.1
pandas    : 1.2.3
us        : 2.0.2
seaborn   : 0.11.1
aesara    : 2.0.12
arviz     : 0.11.2
pymc3     : 4.0

A Bayesian Analysis of Lego Prices in Python with PyMC3

Austin Rochford — Thu, 10 Jun 2021 04:00:00 GMT

This post is the second in a series of posts analyzing Lego pricing data scraped from brickset.com. The first post gave an empirical analysis of the fairness of the price for Darth Vader’s Meditation Chamber (75296), disproving my initial impression that the set was overpriced for its size and prompting me to preorder it. In this post we will build several more formal models of the price of Lego sets based on their size and use them to answer a few related questions.

The newest addition to my collection (75244).

First we make the necessary Python imports and do some light housekeeping.

%matplotlib inline

import datetime
from functools import reduce
import itertools as itl
from warnings import filterwarnings

from aesara import shared, tensor as at
import arviz as az
from matplotlib import MatplotlibDeprecationWarning, pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.ticker import FuncFormatter, MultipleLocator, StrMethodFormatter
import numpy as np
import pandas as pd
import pymc3 as pm
import scipy as sp
from sklearn.preprocessing import StandardScaler
import seaborn as sns

filterwarnings('ignore', category=UserWarning, message=".*FixedFormatter.*")
filterwarnings('ignore', category=MatplotlibDeprecationWarning, module='pandas')
filterwarnings('ignore', category=UserWarning, module='aesara')
filterwarnings('ignore', category=UserWarning, module='arviz')
filterwarnings('ignore', category=UserWarning, module='pandas')

plt.rcParams['figure.figsize'] = (8, 6)

sns.set(color_codes=True)

dollar_formatter = StrMethodFormatter("${x:,.2f}")

Load the Data

We begin the real work by loading the data scraped from Brickset. See the first post in this series for more background on the data.

DATA_URL = 'https://austinrochford.com/resources/lego/brickset_01011980_06012021.csv.gz'

def to_datetime(year):
    return np.datetime64(f"{round(year)}-01-01")

full_df = (pd.read_csv(DATA_URL,
                       usecols=[
                           "Year released", "Set number",
                           "Name", "Set type", "Theme", "Subtheme",
                           "Pieces", "RRP"
                       ])
             .dropna(subset=[
                 "Year released", "Set number",
                 "Name", "Set type", "Theme",
                 "Pieces", "RRP"
             ]))
full_df["Year released"] = full_df["Year released"].apply(to_datetime)
full_df = (full_df.set_index(["Year released", "Set number"])
                  .sort_index())

We see that the data set contains information on approximately 8,500 Lego sets produced between 1980 and June 2021.

full_df.head()

		Name	Set type	Theme	Pieces	RRP	Subtheme
Year released	Set number
1980-01-01	1041-2	Educational Duplo Building Set	Normal	Dacta	68.0	36.50	NaN
	1075-1	LEGO People Supplementary Set	Normal	Dacta	304.0	14.50	NaN
	1101-1	Replacement 4.5V Motor	Normal	Service Packs	1.0	5.65	NaN
	1123-1	Ball and Socket Couplings & One Articulated Joint	Normal	Service Packs	8.0	16.00	NaN
	1130-1	Plastic Folder for Building Instructions	Normal	Service Packs	1.0	14.00	NaN

full_df.tail()

		Name	Set type	Theme	Pieces	RRP	Subtheme
Year released	Set number
2021-01-01	80022-1	Spider Queen’s Arachnoid Base	Normal	Monkie Kid	1170.0	119.99	Season 2
	80023-1	Monkie Kid’s Team Dronecopter	Normal	Monkie Kid	1462.0	149.99	Season 2
	80024-1	The Legendary Flower Fruit Mountain	Normal	Monkie Kid	1949.0	169.99	Season 2
	80106-1	Story of Nian	Normal	Seasonal	1067.0	79.99	Chinese Traditional Festivals
	80107-1	Spring Lantern Festival	Normal	Seasonal	1793.0	119.99	Chinese Traditional Festivals

full_df.describe()

	Pieces	RRP
count	8502.000000	8502.000000
mean	258.739591	31.230506
std	481.627846	44.993559
min	0.000000	0.000000
25%	32.000000	7.000000
50%	98.000000	17.000000
75%	300.000000	39.990000
max	11695.000000	799.990000

We also load data about my personal collection and add a column to full_df indicating whether or not I own each set.

MY_COLLECTION_URL = 'https://austinrochford.com/resources/lego/Brickset-MySets-owned-20210602.csv'

my_df = pd.read_csv(MY_COLLECTION_URL)

my_df.index

Index(['8092-1', '10221-1', '10266-1', '10281-1', '10283-1', '21309-1',
       '21312-1', '21320-1', '21321-1', '31091-1', '40174-1', '40268-1',
       '40391-1', '40431-1', '40440-1', '41602-1', '41608-1', '41609-1',
       '41628-1', '75030-1', '75049-1', '75074-1', '75075-1', '75093-1',
       '75099-1', '75136-1', '75137-1', '75138-1', '75162-1', '75176-1',
       '75187-1', '75229-1', '75243-1', '75244-1', '75248-1', '75254-1',
       '75255-1', '75263-1', '75264-1', '75266-1', '75267-1', '75269-1',
       '75273-1', '75277-1', '75278-1', '75283-1', '75292-1', '75297-1',
       '75302-1', '75306-1', '75308-1', '75317-1', '75318-1'],
      dtype='object')

full_df["austin"] = (full_df.index
                            .get_level_values("Set number")
                            .isin(my_df.index))

Since the data spans more than 40 years, it is important to adjust for inflation. We use the Consumer Price Index for All Urban Consumers: All Items in U.S. City Average from the U.S. Federal Reserve to adjust for inflation.

CPI_URL = 'https://austinrochford.com/resources/lego/CPIAUCNS202100401.csv'

years = pd.date_range('1979-01-01', '2021-01-01', freq='Y') \
            + datetime.timedelta(days=1)
cpi_df = (pd.read_csv(CPI_URL, index_col="DATE", parse_dates=["DATE"])
            .loc[years])
cpi_df["to2021"] = cpi_df.loc["2021-01-01"] / cpi_df

We now add a column RRP2021, which is RRP adjusted to 2021 dollars.

full_df["RRP2021"] = (pd.merge(full_df, cpi_df,
                               left_on=["Year released"],
                               right_index=True)
                        .apply(lambda df: df["RRP"] * df["to2021"],
                               axis=1))

Based on the exploratory data analysis in the previous post, we filter full_df down to approximately 6,300 sets to be included in our analysis.

FILTERS = [
    full_df["Set type"] == "Normal",
    full_df["Pieces"] > 10,
    full_df["Theme"] != "Duplo",
    full_df["Theme"] != "Service Packs",
    full_df["Theme"] != "Bulk Bricks",
    full_df["RRP"] > 0
]

df = full_df[reduce(np.logical_and, FILTERS)].copy()

df.head()

		Name	Set type	Theme	Pieces	RRP	Subtheme	austin	RRP2021
Year released	Set number
1980-01-01	1041-2	Educational Duplo Building Set	Normal	Dacta	68.0	36.50	NaN	False	122.721632
	1075-1	LEGO People Supplementary Set	Normal	Dacta	304.0	14.50	NaN	False	48.752429
	5233-1	Bedroom	Normal	Homemaker	26.0	4.50	NaN	False	15.130064
	6305-1	Trees and Flowers	Normal	Town	12.0	3.75	Accessories	False	12.608387
	6306-1	Road Signs	Normal	Town	12.0	2.50	Accessories	False	8.405591

df.describe()

	Pieces	RRP	RRP2021
count	6312.000000	6312.000000	6312.000000
mean	337.177440	37.038246	45.795998
std	535.619271	49.657704	58.952563
min	11.000000	0.600000	0.971220
25%	68.000000	9.990000	11.866173
50%	177.000000	19.990000	26.712319
75%	400.000000	49.500000	55.952471
max	11695.000000	799.990000	897.373477

Modeling

Upon plotting the data, a linear relationship between the logarithm of pieces and the logarithm of reccomended retail price (in 2021 dollars) becomes immediately apparent.

ax = sns.scatterplot(x="Pieces", y="RRP2021", data=df,
                alpha=0.05);
sns.scatterplot(x="Pieces", y="RRP2021",
                data=df[df["austin"] == True],
                label="My sets", ax=ax); 

ax.set_xscale('log');

ax.set_yscale('log');
ax.set_ylabel("Retail price (2021 $)");

ax.legend(loc='upper left');

Pooled model

We start with a very simple pooled linear model of log price versus log pieces. Let $Y_i$ denote the recommended retail price of the $i$-th set and $y_i = \log Y_i$. Similarly, let $X_i$ denote the number of pieces in the $i$-th set and $x_i = \log X_i$. Due to the log-log linear relationship between $Y_i$ and $X_i$, our models will take the form

\[Y_i = f_i(X_i) \cdot \varepsilon_i,\]

where $\varepsilon_i$ accounts for the unmeasured noise in the price of the $i$-th set. We place a log-normal prior on $\varepsilon_i$ so that $ (_i) = 1$ and $\textrm{Var}\left(\log(\varepsilon_j)\right) = \sigma_i^2.$ standard distributional math shows that

\[y_i \sim N\left(\log\left(f_i(X_i))\right) - \frac{\sigma_i^2}{2}, \sigma_i^2\right).\]

Each model is of the form

\[f_i(X_i) = b_{0, i} \cdot X_i^{\beta_i},\]

\[\log\left(f_i(X_i)\right) = \beta_{0, i} + \beta_i x_i,\]

where $\beta_{0, i} = \log b_{0, i}$.

We define log_pieces and log_rrp2021 which correspond to $x_i$ and $y_i$ respectively.

pieces = df["Pieces"].values
log_pieces = np.log(df["Pieces"].values)

rrp2021 = df["RRP2021"].values
log_rrp2021 = np.log(rrp2021)

For computational efficiency, we will standardize $x_i$ to produce $\tilde{x}_i$ which has mean zero and variance one.

scaler = StandardScaler().fit(log_pieces[:, np.newaxis])

def scale_log_pieces(log_pieces, scaler=scaler):
    return scaler.transform(log_pieces[:, np.newaxis])[:, 0]

std_log_pieces = scale_log_pieces(log_pieces)
std_log_pieces_ = shared(std_log_pieces)

Here definining std_log_pieces_ as an aesara shared variable will allow us to change its value later for posterior predictive sampling.

Our models are therefore given by

\[\log\left(f_i(X_i)\right) = \beta_{0, i} + \beta_i \tilde{x}_i.\]

For the pooled model, $\beta_{0, i} = \beta_0$, $\beta_i = \beta$, and $\sigma_i = \sigma$ are constant across all sets. We give them the priors

\[ \begin{align*} \beta_0 & \sim N(0, 2.5^2) \\ \beta & \sim N(0, 2.5^2) \\ \sigma & \sim \textrm{HalfNormal}(5^2). \end{align*} \]

with pm.Model() as pooled_model:
    β0 = pm.Normal("β0", 0., 2.5)
    β_pieces = pm.Normal("β_pieces", 0., 2.5)
    
    σ = pm.HalfNormal("σ", 5.)

We define the mean and likelihood of the observed data as above.

with pooled_model:
    μ = β0 + β_pieces * std_log_pieces_ - 0.5 * σ**2
    obs = pm.Normal("obs", μ, σ, observed=log_rrp2021)

We now sample from posterior distribution of this model.

CHAINS = 3
SEED = 123456789

SAMPLE_KWARGS = {
    'cores': CHAINS,
    'random_seed': [SEED + i for i in range(CHAINS)],
    'return_inferencedata': True
}

with pooled_model:
    pooled_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β0, β_pieces, σ]

100.00% [6000/6000 00:15<00:00 Sampling 3 chains, 0 divergences]

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

The energy plot, BFMI, and $\hat{R}$ statistics for these samples show no cause for concern.

def make_diagnostic_plots(trace, axes=None, min_mult=0.995, max_mult=1.005):
    if axes is None:
        fig, axes = plt.subplots(ncols=2, sharex=False, sharey=False,
                                 figsize=(16, 6))
        
    az.plot_energy(trace, ax=axes[0])
    
    
    rhat = az.rhat(trace).max()
    axes[1].barh(np.arange(len(rhat.variables)), rhat.to_array(),
                 tick_label=list(rhat.variables.keys()))
    axes[1].axvline(1, c='k', ls='--')

    axes[1].set_xlim(
        min_mult * min(rhat.min().to_array().min(), 1),
        max_mult * max(rhat.max().to_array().max(), 1)
    )
    axes[1].set_xlabel(r"$\hat{R}$")

    axes[1].set_ylabel("Variable")
    
    return fig, axes

make_diagnostic_plots(pooled_trace);

We sample from the posterior predictive distribution at the observed data points and plot the residuals.

with pooled_model:
    pp_trace_data = pm.sample_posterior_predictive(pooled_trace)
    
df["pooled_resid"] = (log_rrp2021 - pp_trace_data["obs"]).mean(axis=0)

100.00% [3000/3000 00:02<00:00]

ax = sns.scatterplot(x="Pieces", y="pooled_resid", data=df,
                     alpha=0.1)

ax.set_xscale('log');
ax.set_xlabel("Pieces");

ax.set_ylabel("Log RRP residual (2021 $)");

ax.set_title("Pooled model");

We will investigate these residuals further shortly. First we sample from the posterior predictive distribution along a grid spanning reasonable values of pieces and plot the resulting predictions.

pp_pieces = np.logspace(1, np.log10(pieces.max()))
pp_log_pieces = np.log(pp_pieces)
pp_std_log_pieces = scale_log_pieces(pp_log_pieces)

std_log_pieces_.set_value(pp_std_log_pieces)

with pooled_model:
    pooled_pp_trace = pm.sample_posterior_predictive(pooled_trace)

100.00% [3000/3000 00:01<00:00]

ax = sns.scatterplot(x="Pieces", y="RRP2021", data=df,
                alpha=0.05);
sns.scatterplot(x="Pieces", y="RRP2021",
                data=df[df["austin"] == True],
                alpha=0.5, label="My sets",
                ax=ax);

pp_low, pp_high = np.percentile(
    np.exp(pooled_pp_trace["obs"]),
    [2.5, 97.4],
    axis=0
)
ax.plot(pp_pieces, np.exp(pooled_pp_trace["obs"]).mean(axis=0),
        c='k', label="Posterior predicted mean");
ax.fill_between(pp_pieces, pp_low, pp_high,
                color='k', alpha=0.15, zorder=-1,
                label="95% posterior interval");

ax.set_xscale('log');

ax.set_yscale('log');
ax.set_ylabel("Retail price (2021 $)");

ax.legend(loc='upper left');

The predictions seem reasonable. We now plot the posterior distribution of the regression coefficients.

az.plot_posterior(pooled_trace, var_names=["β0", "β_pieces"]);

When we break down the residuals for a few notable themes, we see that Star Wars and Disney sets have larger residuals than Creator sets with the corresponding number of pieces.

ax = sns.scatterplot(x="Pieces", y="pooled_resid", data=df,
                     alpha=0.1)

for theme in ["Star Wars", "Creator", "Disney"]:
    sns.scatterplot(x="Pieces", y="pooled_resid",
                    data=df[df["Theme"] == theme],
                    alpha=0.5, label=f"{theme} sets",
                    ax=ax);

ax.set_xscale('log');
ax.set_xlabel("Pieces");

ax.set_ylabel("Log RRP residual (2021 $)");

ax.set_title("Pooled model");

This difference in the distribution of residuals across themes suggests that we introduce theme-specific factors into our model.

Unpooled model

In order to introduce theme-specific factors, we create indicator variables for each theme.

theme_id, theme_map = df["Theme"].factorize(sort=True)
n_theme = theme_map.size

theme_id_ = shared(theme_id)

The simplest model that includes theme-specific factors is an unpooled model that consists of independent regressions per theme. With $j(i)$ denoting the index of the theme of the $i$-th set, we have the model

\[\log\left(f_i(X_i)\right) = \beta_{0, j(i)} + \beta_{j(i)} \tilde{x}_i\]

with the i.i.d. priors

\[ \begin{align*} \beta_{0, 1}, \ldots \beta_{0, n_j} & \sim N(0, 2.5^2) \\ \beta_1, \ldots \beta_{n_j} & \sim N(0, 2.5^2) \\ \sigma_1, \ldots, \sigma_{n_j} & \sim \textrm{HalfNormal}(5^2). \end{align*} \]

std_log_pieces_.set_value(std_log_pieces)

with pm.Model() as unpooled_model:
    β0 = pm.Normal("β0", 0., 2.5, shape=n_theme)
    β_pieces = pm.Normal("β_pieces", 0., 2.5, shape=n_theme)
    
    σ = pm.HalfNormal("σ", 5., shape=n_theme)
    μ = β0[theme_id_] + β_pieces[theme_id_] * std_log_pieces_ \
            - 0.5 * σ[theme_id_]**2
    
    obs = pm.Normal("obs", μ, σ[theme_id_], observed=log_rrp2021)

We don’t actually sample from this model because it is underdetermined for many themes that have few sets associated with them, as the following plot shows.

ax = (df["Theme"]
        .value_counts()
        .nsmallest(n=20)
        .plot.barh())

ax.set_xlabel("Number of sets");
ax.set_ylabel("Theme");

There are five themes with only a single set, and a number of themes with fewer than five sets total. Regression models with only a single data point are underdetermined, and those with few observations are likely to have wide error bars. Instead of imposing a minimum number of sets per theme and then using an unpooled model, we will use a partially pooled model to share some information across themes/regularize while allowing themes with many sets to deviate from the average if the data indicate that they should.

Partially pooled model

One model that shares information across themes is called “partially pooled” (it may also be called multilevel or hierarchical). In the partially pooled model, the intercept and piece coefficient for each theme are drawn not from i.i.d normal distributions but from a normal distribution with a common mean and variance. This shared hyperprior distribution regularizes the theme intercepts and coefficients, shrinking those for themes with fewer sets towards the pooled values.

Conceptually, we use the priors

\[ \begin{align*} \mu_{\beta_0} & \sim N(0, 2.5^2) \\ \sigma_{\beta_0} & \sim \textrm{HalfNormal}(2.5^2) \\ \beta_{0, 1}, \ldots, \beta_{0, n_j} & \sim N\left(\mu_{\beta_0}, \sigma_{\beta_0}^2\right) \end{align*} \]

and similarly for $\beta$. In reality, we use a non-centered parametrization parameterization that is mathematically equivalent but often more computationally efficient.

def hierarchical_normal(name, shape, μ=None):
    if μ is None:
        μ = pm.Normal(f"μ_{name}", 0., 2.5)

    Δ = pm.Normal(f"Δ_{name}", 0., 1., shape=shape)
    σ = pm.HalfNormal(f"σ_{name}", 2.5)
    
    return pm.Deterministic(name, μ + Δ * σ)

with pm.Model() as partial_model:
    β0 = hierarchical_normal("β0", n_theme)
    β_pieces = hierarchical_normal("β_pieces", n_theme)
    
    σ = pm.HalfNormal("σ", 5.)
    μ = β0[theme_id_] + β_pieces[theme_id_] * std_log_pieces_ \
            - 0.5 * σ**2
    
    obs = pm.Normal("obs", μ, σ, observed=log_rrp2021)

with partial_model:
    partial_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [μ_β0, Δ_β0, σ_β0, μ_β_pieces, Δ_β_pieces, σ_β_pieces, σ]

100.00% [6000/6000 03:15<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 1_000 draw iterations (3_000 + 3_000 draws total) took 196 seconds.
The estimated number of effective samples is smaller than 200 for some parameters.

Again our diagnostics show no cause for concern.

make_diagnostic_plots(partial_trace);

After calculating the residuals for the partially pooled model we see that they tend to be closer to zero and the diagonal bands are less pronounced.

with partial_model:
    pp_trace_data = pm.sample_posterior_predictive(partial_trace)
    
df["partial_resid"] = (log_rrp2021 - pp_trace_data["obs"]).mean(axis=0)

100.00% [3000/3000 00:02<00:00]

fig, axes = plt.subplots(ncols=2, sharex=True, sharey=True,
                         figsize=(14, 6))

sns.scatterplot(x="Pieces", y="pooled_resid", data=df,
                alpha=0.1, ax=axes[0]);

axes[0].set_xscale('log');
axes[0].set_xlabel("Pieces");

axes[0].set_ylabel("Log RRP residual (2021 $)");

axes[0].set_title("Pooled model");

sns.scatterplot(x="Pieces", y="partial_resid", data=df,
                alpha=0.1, ax=axes[1]);

axes[1].set_xscale('log');
axes[1].set_xlabel("Pieces");

axes[1].set_ylabel("Log RRP residual (2021 $)");

axes[1].set_title("Partially pooled model");

fig.tight_layout();

We also see that the residuals for the example themes (Star Wars, Creator, and Disney) are more comingled than for the pooled model.

fig, axes = plt.subplots(ncols=2, sharex=True, sharey=True,
                         figsize=(14, 6))

sns.scatterplot(x="Pieces", y="pooled_resid", data=df,
                alpha=0.1, ax=axes[0]);

for theme in ["Star Wars", "Creator", "Disney"]:
    sns.scatterplot(x="Pieces", y="pooled_resid",
                    data=df[df["Theme"] == theme],
                    alpha=0.5, label=f"{theme} sets",
                    ax=axes[0]);

axes[0].set_xscale('log');
axes[0].set_xlabel("Pieces");

axes[0].set_ylabel("Log RRP residual (2021 $)");

axes[0].set_title("Pooled model");

sns.scatterplot(x="Pieces", y="partial_resid", data=df,
                alpha=0.1, ax=axes[1]);

for theme in ["Star Wars", "Creator", "Disney"]:
    sns.scatterplot(x="Pieces", y="partial_resid",
                    data=df[df["Theme"] == theme],
                    alpha=0.5, label=f"{theme} sets",
                    ax=axes[1]);

axes[1].set_xscale('log');
axes[1].set_xlabel("Pieces");

axes[1].set_ylabel("Log RRP residual (2021 $)");

axes[1].set_title("Partially pooloed model");

fig.tight_layout();

Now that we have posterior samples from two models, we use Pareto-smoothed importance sampling leave-one-out cross validation¹

traces = {
    "Unpooled": pooled_trace,
    "Partially pooled": partial_trace
}

comp_df = az.compare(traces)

comp_df.loc[:, :"weight"]

	rank	loo	p_loo	d_loo	weight
Partially pooled	0	-2315.578245	190.594364	0.000000	0.99951
Unpooled	1	-5133.311893	4.871287	2817.733648	0.00049

fig, ax = plt.subplots(figsize=(12, 6))

az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

The table and chart show that by sharing some information across themes while also incorporating their varying effects the partially pooled model has substantially better predictive accuracy than the pooled model (higher values of loo correspond to better predictive ability).

In order to compare the posterior predictions of the pooled and partially pooled model directly, we need to define a grid over theme/standardized log piece combinations.

pp_df = pd.DataFrame.from_records(
    itl.product(range(n_theme), pp_std_log_pieces),
    columns=["theme_id", "std_log_pieces"]
)
pp_df["theme"] = theme_map[pp_df["theme_id"].values]
pp_df["pieces"] = np.exp(
    scaler.inverse_transform(pp_df["std_log_pieces"]
                                  .values[:, np.newaxis])
          .ravel()
)

pp_df.head()

	std_log_pieces	theme	pieces
0	-2.373432	4 Juniors	10.000000
1	-2.252260	4 Juniors	11.550805
2	-2.131088	4 Juniors	13.342109
3	-2.009916	4 Juniors	15.411210
4	-1.888744	4 Juniors	17.801188

pp_df.tail()

	theme_id	std_log_pieces	theme	pieces
6545	130	3.079305	Znap	6569.786072
6546	130	3.200477	Znap	7588.631734
6547	130	3.321649	Znap	8765.480484
6548	130	3.442820	Znap	10124.835517
6549	130	3.563992	Znap	11695.000000

std_log_pieces_.set_value(pp_df["std_log_pieces"].values)
theme_id_.set_value(pp_df["theme_id"].values)

We now sample from the posterior predictive distribution and compare the predictions to those from the pooled model.

with partial_model:
    partial_pp_trace = pm.sample_posterior_predictive(partial_trace)

100.00% [3000/3000 00:02<00:00]

pp_df["pp_partial_mean"] = np.exp(partial_pp_trace["obs"]).mean(axis=0)
pp_df["pp_partial_low"], pp_df["pp_partial_high"] = np.percentile(
    np.exp(partial_pp_trace["obs"]),
    [2.5, 97.5],
    axis=0
)

pp_df.head()

	std_log_pieces	theme	pieces	pp_partial_mean	pp_partial_low	pp_partial_high
0	-2.373432	4 Juniors	10.000000	5.214650	2.380664	10.484612
1	-2.252260	4 Juniors	11.550805	5.917360	2.668467	11.413480
2	-2.131088	4 Juniors	13.342109	6.864103	3.045283	13.610948
3	-2.009916	4 Juniors	15.411210	7.653877	3.515344	14.966362
4	-1.888744	4 Juniors	17.801188	8.679962	4.009006	16.935274

ax = sns.scatterplot(x="Pieces", y="RRP2021", data=df,
                alpha=0.05);

(pp_df.groupby("theme_id")
      .plot("pieces", "pp_partial_mean",
            c='k', alpha=0.1, label='_nolegend_',
            ax=ax));
ax.plot(pp_pieces, np.exp(pooled_pp_trace["obs"]).mean(axis=0),
        c='r', lw=2, label="Pooled");

ax.set_xscale('log');

ax.set_yscale('log');
ax.set_ylabel("Retail price (2021 $)");

handles, _ = ax.get_legend_handles_labels()
partial_line = Line2D([0], [0], color='k',
                      label="Partially pooled")
handles.append(partial_line)

ax.legend(handles=handles,
          title="Posterior predicted mean",
          loc='upper left');

Note that the partial pooling model produces one line per theme, with varying slopes and intercepts. We can see that while the predictions for some themes are quite far away from the pooled predictions most are close to them, as shown by the stacking of the semi-transparent lines. Below we highlight the partially pooled predictions for a few themes of interest among all partially pooled predictions.

PLOT_THEMES = [
    "Creator Expert",
    "Disney",
    "Star Wars",
    "Harry Potter",
    "Marvel Super Heroes",
    "Ninjago",
    "City",
    "Space",
    "Jurassic World"
]

ax = sns.scatterplot(x="Pieces", y="RRP2021", data=df,
                alpha=0.05);

(pp_df.groupby("theme_id")
      .plot("pieces", "pp_partial_mean",
            c='k', alpha=0.1, label='_nolegend_',
            ax=ax));
(pp_df[pp_df["theme"].isin(PLOT_THEMES)]
      .groupby("theme_id")
      .plot("pieces", "pp_partial_mean",
            c='r', label='_nolegend_',
            ax=ax));

ax.set_xscale('log');

ax.set_yscale('log');
ax.set_ylabel("Retail price (2021 $)");

title = "Highlighted: " \
            + ", ".join(PLOT_THEMES[:4]) + ",\n" \
            + ", ".join(PLOT_THEMES[4:-1]) \
            + ", and " + PLOT_THEMES[-1]
ax.set_title(title);

These notable themes have preditions close to the global average, which makes sense since they all have a fair number of sets, so the global mean is pulled towards their theme effects.

It has long been known in the quantitative social sciences that when naively modeled (as we have) that group-level effects (in this case theme-level effects), can strongly correlate with item-level covariates (in this case standardized log pieces). We now investigate whether our partially pooled model results in such a correlation.

First we calculate the average value of std_log_pieces within each theme.

theme_mean_std_log_pieces = (pd.Series(std_log_pieces, index=df.index)
                               .groupby(df["Theme"])
                               .mean())

ax = theme_mean_std_log_pieces.plot.hist(bins=20, lw=0., alpha=0.5)

ax.set_xlabel("Average standardized log pieces\nper set within theme");

We see that these theme averages vary fairly widely. The following plots show the relationship between each theme’s average standardized log pieces and the posterior expected values of $\beta_{0, j}$ and $\beta_j$ respectively.

post_β0 = partial_trace["posterior"]["β0"].mean(dim=('chain', 'draw'))
post_β_pieces = partial_trace["posterior"]["β_pieces"].mean(dim=('chain', 'draw'))

fig, (β0_ax, β_pieces_ax) = plt.subplots( 
    ncols=2, sharex=True, sharey=False,
    figsize=(15, 6)
)

β0_ax.scatter(theme_mean_std_log_pieces, post_β0);

β0_ax.set_xlabel("Average standardized log pieces\nper set within theme");
β0_ax.set_ylabel("Posterior expected\nvalue of $\\beta_0$");
 
β0_corr, β0_corr_p_val = sp.stats.pearsonr(
    theme_mean_std_log_pieces.values,
    post_β0
)
β0_ax.set_title(f"Correlation coefficient = {β0_corr:0.2f}\n"
                f"Two-tailed p-value {β0_corr_p_val:.2e}");

β_pieces_ax.scatter(theme_mean_std_log_pieces, post_β_pieces);

β_pieces_ax.set_xlabel("Average standardized log pieces\nper set within theme");
β_pieces_ax.set_ylabel("Posterior expected\nvalue of $\\beta_{pieces}$");

β_pieces_corr, β_pieces_corr_p_val = sp.stats.pearsonr(
    theme_mean_std_log_pieces.values,
    post_β_pieces
)
β_pieces_ax.set_title(f"Correlation coefficient = {β_pieces_corr:0.2f}\n"
                      f"Two-tailed p-value {β_pieces_corr_p_val:.2e}");

Both plots show highly significant correlations, with a stronger correlation between the theme average and the posterior expected average of $\beta_{0, j}$. These correlations lead to violations of the Gauss-Markov assumptions in our regression design, so we remedy them with an expanded model in the next section.

Partially pooled mean model

Fortunately Bafumi and Gelman² show how to remedy this violation by including the theme-level average standardized log pieces as a covariate in $\beta_{0, j}$ and $\beta_j$. if we denote the average standardized log pieces for sets in the $j$-th theme as $\bar{\tilde{x}}_j$, in this extended model, $\beta_{0, j}$ is defined as

\[ \begin{align*} \mu_{\beta_0}, \gamma_{\beta, 0} & \sim N(0, 2.5^2) \\ \sigma_{\beta_0} & \sim \textrm{HalfNormal}(2.5^2) \\ \beta_{0, j} & \sim N\left(\mu_{\beta_0} + \gamma_{\beta, 0} \cdot \bar{\tilde{x}}_j, \sigma_{\beta_0}^2\right), \end{align*} \]

and similarly for $\beta_j$.

def hierarchical_normal_with_mean(name, x_mean, shape, μ=None):
    γ = pm.Normal(f"γ_{name}", 0., 2.5)
        
    return hierarchical_normal(name, shape, μ=μ) + γ * x_mean

with pm.Model() as mean_model:
    σ = pm.HalfNormal("σ", 5.)
    
    β0 = hierarchical_normal_with_mean(
        "β0", at.constant(theme_mean_std_log_pieces),
        n_theme
    )
    β_pieces = hierarchical_normal_with_mean(
        "β_pieces", at.constant(theme_mean_std_log_pieces),
        n_theme
    )
    
    μ = β0[theme_id_] + β_pieces[theme_id_] * std_log_pieces_ \
            - 0.5 * σ**2
    
    obs = pm.Normal("obs", μ, σ, observed=log_rrp2021)

We now sample from this model.

std_log_pieces_.set_value(std_log_pieces)
theme_id_.set_value(theme_id)

with mean_model:
    mean_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ, γ_β0, μ_β0, Δ_β0, σ_β0, γ_β_pieces, μ_β_pieces, Δ_β_pieces, σ_β_pieces]

100.00% [6000/6000 03:57<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 1_000 draw iterations (3_000 + 3_000 draws total) took 239 seconds.
The number of effective samples is smaller than 10% for some parameters.

Once again the sampling diagnostics show no cause for concern.

make_diagnostic_plots(mean_trace);

The correlations from above have disappeared in this model.

post_β0 = mean_trace["posterior"]["β0"].mean(dim=('chain', 'draw'))
post_β_pieces = mean_trace["posterior"]["β_pieces"].mean(dim=('chain', 'draw'))

fig, (β0_ax, β_pieces_ax) = plt.subplots( 
    ncols=2, sharex=True, sharey=False,
    figsize=(15, 6)
)

β0_ax.scatter(theme_mean_std_log_pieces, post_β0);

β0_ax.set_xlabel("Average standardized log pieces\nper set within theme");
β0_ax.set_ylabel("Posterior expected\nvalue of $\\beta_0$");
 
β0_corr, β0_corr_p_val = sp.stats.pearsonr(
    theme_mean_std_log_pieces.values,
    post_β0
)
β0_ax.set_title(f"Correlation coefficient = {β0_corr:0.2f}\n"
                f"Two-tailed p-value {β0_corr_p_val:.2f}");

β_pieces_ax.scatter(theme_mean_std_log_pieces, post_β_pieces);

β_pieces_ax.set_xlabel("Average standardized log pieces\nper set within theme");
β_pieces_ax.set_ylabel("Posterior expected\nvalue of $\\beta_{pieces}$");

β_pieces_corr, β_pieces_corr_p_val = sp.stats.pearsonr(
    theme_mean_std_log_pieces.values,
    post_β_pieces
)
β_pieces_ax.set_title(f"Correlation coefficient = {β_pieces_corr:0.2f}\n"
                      f"Two-tailed p-value {β_pieces_corr_p_val:.2f}");

We add this model to our comparison.

traces["Partially pooled with mean"] = mean_trace

comp_df = az.compare(traces)

comp_df.loc[:, :"weight"]

	rank	loo	p_loo	d_loo	weight
Partially pooled	0	-2315.578245	190.594364	0.000000	0.958583
Partially pooled with mean	1	-2316.840059	187.512592	1.261814	0.000000
Unpooled	2	-5133.311893	4.871287	2817.733648	0.041417

fig, ax = plt.subplots(figsize=(12, 6))

az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

We see that the the PSIS-LOO scores of both partially pooled model are better than the pooled model and statistically indistinguishable from each other. Given these scores, we will work with the model that accounts for the theme’s average standardized log pieces going forward, since it satisfies the Gauss-Markov assumptions better.

The residuals from the two partially pooled models are also visually quite similar, though the second partial pooling model does have less pronounced diagonal bands.

with mean_model:
    pp_trace_data = pm.sample_posterior_predictive(mean_trace)

df["mean_resid"] = (log_rrp2021 - pp_trace_data["obs"]).mean(axis=0)

100.00% [3000/3000 00:01<00:00]

fig, axes = plt.subplots(ncols=2, sharex=True, sharey=True,
                         figsize=(14, 6))

sns.scatterplot(x="Pieces", y="pooled_resid", data=df,
                alpha=0.1, ax=axes[0]);

axes[0].set_xscale('log');
axes[0].set_xlabel("Pieces");

axes[0].set_ylabel("Log RRP residual (2021 $)");

axes[0].set_title("Partially pooled model");

sns.scatterplot(x="Pieces", y="mean_resid", data=df,
                alpha=0.1, ax=axes[1]);

axes[1].set_xscale('log');
axes[1].set_xlabel("Pieces");

axes[1].set_ylabel("Log RRP residual (2021 $)");

axes[1].set_title("Partially pooled mean model");

fig.tight_layout();

Interestingly, there is more variation in the varying intercepts ($\beta_{0, j}$) than the varying slopes ($\beta_j$), since $\sigma_{\beta_0}$ has a larger posterior expected value than $\sigma_{\beta}$.

az.plot_posterior(mean_trace,
                  var_names=[
                      "μ_β0", "μ_β_pieces",
                      "σ_β0", "σ_β_pieces"
                  ],
                  grid=(2, 2));

Plotting the residuals of log RRP in 2021 dollars versus time shows that the residuals for this model are significantly larger for sets between 1980 and 1995. This phenomenon makes sense as Lego has released many more sets per year after 1995 than in the initial period, so the estimates are dominated by the data in this later period.

ax = (df.groupby(level="Year released")
        ["mean_resid"]
        .mean()
        .plot(c='k', label="Year average"));

strip_ax = ax.twiny()
sns.stripplot(x="Year released", y="mean_resid",
              data=df.reset_index(),
              jitter=1.5,
              color='C0', alpha=0.1,
              ax=strip_ax);

strip_ax.set_axis_off();

ax.set_ylim(-2, 2);
ax.set_ylabel("Log RRP residual (2021 $)");

ax.legend(loc='upper left');

Time-varying model

We address this issue with the residuals by adding a time-varying component to $\beta_0$. First we build t, which indicates the number of years after the data begins in 1980 that the set was released.

year = df.index.get_level_values("Year released").year.values
t = year - year.min()
t_ = shared(t)
n_year = int(t.max() + 1)

The prior on $\beta_j$ stays the same as in the mean model, but we add a Gaussian random walk term, $\beta_{0, t}$ to the intercept, $\beta_0$. We use the priors

\[ \begin{align*} \mu_{\beta_{0, pieces}}, \gamma_{\beta_{0, pieces}} & \sim N(0, 2.5^2) \\ \sigma_{\beta_{0, pieces}} & \sim \textrm{HalfNormal}(2.5^2) \\ \Delta_{\beta_{0, i}} & \sim N(0, 0.25^2) \\ \beta_{0, t} & \sim \sum_{i = 1}^t \Delta_{\beta_{0, i}} \\ \beta_{0, j, t} & \sim N\left(\mu_{\beta_0} + \beta_{0, t} + \gamma_{\beta, 0} \cdot \bar{\tilde{x}}_j, \sigma_{\beta_0}^2\right). \end{align*} \]

We see that the time-varying factor $\beta_{0, t}$ is constant across themes in order to reduce the computational burden when sampling from this model.

with pm.Model() as time_model:
    Δ_β0_t = pm.Normal("Δ_β0_t", 0., 0.25, shape=n_year)
    β0_t = pm.Deterministic("β0_t", Δ_β0_t.cumsum())
    
    β0_theme = hierarchical_normal_with_mean(
        "β0_theme", at.constant(theme_mean_std_log_pieces),
        n_theme
    )
    β0_i = β0_t[t_] + β0_theme[theme_id_]

    β_pieces = hierarchical_normal_with_mean(
        "β_pieces", at.constant(theme_mean_std_log_pieces),
        n_theme
    )
    
    σ = pm.HalfNormal("σ", 5.)
    μ = β0_i + β_pieces[theme_id_] * std_log_pieces_ - 0.5 * σ**2
    
    obs = pm.Normal("obs", μ, σ, observed=log_rrp2021)

We now sample from this model.

with time_model:
    time_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [Δ_β0_t, γ_β0_theme, μ_β0_theme, Δ_β0_theme, σ_β0_theme, γ_β_pieces, μ_β_pieces, Δ_β_pieces, σ_β_pieces, σ]

100.00% [6000/6000 17:36<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 1_000 draw iterations (3_000 + 3_000 draws total) took 1056 seconds.
The number of effective samples is smaller than 25% for some parameters.

Once again the sampling diagnostics show no cause for concern.

make_diagnostic_plots(time_trace);

We see that the time-varying model performs noticeably better than the other models in terms of PSIS-LOO score.

traces["Partially pooled with mean and time"] = time_trace

comp_df = az.compare(traces)

comp_df.loc[:, :"weight"]

	rank	loo	p_loo	d_loo	weight
Partially pooled with mean and time	0	-2058.373203	245.239325	0.000000	9.205189e-01
Partially pooled	1	-2315.578245	190.594364	257.205042	4.565283e-02
Partially pooled with mean	2	-2316.840059	187.512592	258.466855	1.857221e-12
Unpooled	3	-5133.311893	4.871287	3074.938690	3.382830e-02

fig, ax = plt.subplots(figsize=(12, 6))

az.plot_compare(comp_df, plot_ic_diff=False, ax=ax);

The residuals for this model are similar to those of the previous one when plotted against the number of pieces in the set.

with time_model:
    pp_time_data = pm.sample_posterior_predictive(time_trace)

df["time_resid"] = (log_rrp2021 - pp_time_data["obs"]).mean(axis=0)

100.00% [3000/3000 00:02<00:00]

fig, axes = plt.subplots(ncols=2, sharex=True, sharey=True,
                         figsize=(14, 6))

sns.scatterplot(x="Pieces", y="mean_resid", data=df,
                alpha=0.1, ax=axes[0]);

axes[0].set_xscale('log');
axes[0].set_xlabel("Pieces");

axes[0].set_ylabel("Log RRP residual (2021 $)");

axes[0].set_title("Partially pooled mean model");

sns.scatterplot(x="Pieces", y="time_resid", data=df,
                alpha=0.1, ax=axes[1]);

axes[1].set_xscale('log');
axes[1].set_xlabel("Pieces");

axes[1].set_ylabel("Log RRP residual (2021 $)");

axes[1].set_title("Partially pooled mean model with time");

fig.tight_layout();

The following plot shows that this model has significantly reduced the amount of time-dependence in the residuals.

ax = (df.groupby(level="Year released")
        ["time_resid"]
        .mean()
        .plot(c='k', label="Year average"));

strip_ax = ax.twiny()
sns.stripplot(x="Year released", y="time_resid",
              data=df.reset_index(),
              jitter=1.5,
              color='C0', alpha=0.1,
              ax=strip_ax);

strip_ax.set_axis_off();

ax.set_ylim(-2, 2);
ax.set_ylabel("Log RRP residual (2021 $)");

ax.legend(loc='upper left');

Again, there is more variation in the varying intercepts ($\beta_{0, j}$) than the varying slopes ($\beta_j$), since $\sigma_{\beta_0}$ has a larger posterior expected value than $\sigma_{\beta}$.

az.plot_posterior(time_trace,
                  var_names=[
                      "μ_β0_theme", "μ_β_pieces",
                      "σ_β0_theme", "σ_β_pieces"
                  ],
                  grid=(2, 2));

We now revisit the question of whether or not Darth Vader’s Meditation Chamber is fairly priced. In the previous post in this series, we compared the price per piece of this set to all sets, all Star Wars sets, and all sets currently in my collection. We found that it was in or near the first quartile of price per piece among all sets and Star Wars sets, and near the median priece per piece in my collection. With this model in hand, we can ask where the adjusted recommended retail price falls in the predicted distribution of posterior predicted prices for a Star Wars set with 663 pieces.

VADER_MEDITATION = "75296-1"

vader_info = df.xs(VADER_MEDITATION, level="Set number")
vader_iloc = (df.index
                .get_level_values("Set number")
                .isin([VADER_MEDITATION])
                .argmax())

(vader_info["RRP2021"].values[0] \
    >= np.exp(pp_trace_data["obs"][:, vader_iloc])).mean()

0.2816666666666667

The set falls in towards the bottom of the second quartile of predicted prices for such a set, which is consistent with the empirical price per piece analysis of the previous post. This model can also answer some counterfactual questions about my collection. We see that my sets are highly concentrated in the Star Wars theme, with a smattering of other themes represented as well.

(df[df["austin"] == True]
   ["Theme"]
   .value_counts(ascending=True)
   .plot.barh());

This preponderance of Star Wars sets is even stronger than it initially appears because all but one of my BrickHeadz sets are Star Wars characters. The single exception are the adorable German Shepherds (40440).

In particular, we can compare the actual cost of my collection to the predicted expected cost of my collection conditioned on the piece count and theme of the sets.

coll_value = (df[df["austin"] == True]
                ["RRP2021"]
                .sum())

coll_value

3206.3933684401236

(np.exp(pp_trace_data["obs"])
                     [:, df["austin"] == True]
                     .mean(axis=0)
                     .sum())

3826.537853977332

So my collection cost about 20% less than expected conditioned on piece count and theme. (This figure actually understates my frugality as I have received several sets as gifts, notably the massive Super Star Destroyer (10221).) We can easily check where my collection’s actual value falls in the posterior predictive distribution of its possible values.

pp_coll_value = (np.exp(pp_trace_data["obs"])
                   [:, df["austin"] == True]
                   .sum(axis=1))

ax, = (sns.displot(pp_coll_value, kind='kde',
                   height=6, aspect=8 / 6,
                   label="Posterior predictive distribution")
          .axes.flatten())

ax.axvline(coll_value, c='k', ls='--',
           label="Actual value");

ax.xaxis.set_major_formatter(dollar_formatter);
ax.set_xlabel("Collection value (2021 $)");

ax.set_yticks([]);

ax.legend(loc='upper right');

(coll_value > pp_coll_value).mean()

0.007

My collection’s actual value is in the first percentile of the posterior predictive distribution of possible collection values; according to this model, I really am quite cheap!

It is worth noting that in this post we have modeled log price directly, not price per piece as in the last model. Those familiar with offset terms in offsets in generalizedl linear models will notice that the model we have specified can be equivalent to an model of (log) price per piece if we fix $\beta_j = 1$, since $\log\left(\frac{x}{y}\right) = \log x - \log y.$ The following plot shows how far outside the high density area of the posterior distribution of $\mu_{\beta_0}$ the value one is.

az.plot_posterior(time_trace, var_names=["μ_β0_theme"], ref_val=1.);

Due to this fact, we prefer to model (log) price directly over using an offset model that includes (log) piece count on both sides of the equal sign.

Future posts will explore further inferences from this model and develop extended models using other attributes of each set available via Brickset.

This post is available as a Jupyter notebook here.

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Thu Jun 10 2021

Python implementation: CPython
Python version       : 3.8.8
IPython version      : 7.22.0

pandas    : 1.2.3
aesara    : 2.0.10
pymc3     : 4.0
scipy     : 1.6.2
seaborn   : 0.11.1
matplotlib: 3.4.1
arviz     : 0.11.2
numpy     : 1.20.2

Vehtari, A., Gelman, A., & Gabry, J. (2017). Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC. (PSIS-LOO) to compare them.) Statistics and computing, 27(5), 1413-1432.↩︎
Bafumi, J., & Gelman, A. E. (2006). Fitting multilevel models when predictors and group effects correlate.↩︎

The Hierarchical Regularized Horseshoe Prior in PyMC3

Austin Rochford — Sat, 29 May 2021 04:00:00 GMT

For some time I have been interested in better understanding the horseshoe prior ¹ by implementing it in PyMC3. The horsehoe prior is a continuous alternative to the spike-and-slab prior for sparse Bayesian estimation. The continuity of the horeshoe prior allows for simpler estimation using general purpose Bayesian computation packages such as pymc3 and Stan. Piironen and Vehtari ² have proposed a hierarchical regularized horseshoe prior that has advantages over the original horseshoe prior when it comes to specifying the hyperprior distributions on the regularization parameters. In the appendix of their paper, Piironen and Vehtari give a sample implementation in Stan. This post gives a corresponding implementation in pymc3 and tests it on simulated data.

First we import the necessary Python packages and do some light housekeeping.

%matplotlib inline

from warnings import filterwarnings

from aesara import tensor as at
import arviz as az
from matplotlib import pyplot as plt
import numpy as np
import pymc3 as pm
import seaborn as sns

filterwarnings('ignore', category=UserWarning, module='aesara')
filterwarnings('ignore', category=UserWarning, module='arviz')

sns.set(color_codes=True)

Simulated Data

Since the horseshoe is designed for sparse estimation of regression parameters, we simulate data from a model with $D = 50$ variables, $D_0 = 5$ of which are nonzero.

D = 50
D0 = 5

SEED = 123456789 # for reproducibility

rng = np.random.default_rng(SEED)

INTERCEPT = rng.uniform(-3, 3)

COEF = np.zeros(D)
COEF[:D0] = rng.choice([-1, 1], size=D0) * rng.normal(5, 1, size=D0)

Now we draw $N = 100$ random observations with $x_i \sim N(0, 1).$

N = 100

X = rng.normal(size=(N, D))

We now simulate responses from this regression model with $N(0, 1)$-distributed noise.

SIGMA = 1.

y = INTERCEPT + X.dot(COEF) + rng.normal(0, SIGMA, size=N)

We see that there is a linear relationship between $y$ and the first $D_0$ variables and no relationship between $y$ and a random sample from the subsequent $D - D0$ variables.

fig, axes = plt.subplots(nrows=2, ncols=D0,
                         sharex=True, sharey=True,
                         figsize=(16, 6))

for i, (ax, coef) in enumerate(zip(axes[0], COEF)):
    ax.scatter(X[:, i], y, alpha=0.75);
    
    ax.set_xlabel(f"$x_{{ {i} }}$");
    ax.set_title(f"$\\beta_{{ {i} }} \\approx {coef:.2f}$");


zero_coef_ix = rng.choice(range(D0, D), replace=False, size=D0)
zero_coef_ix.sort()

for ax, i in zip(axes[1], zero_coef_ix):
    ax.scatter(X[:, i], y, alpha=0.75);
    
    ax.set_xlabel(f"$x_{{ {i} }}$");
    ax.set_title(f"$\\beta_{{ {i} }} = 0$");
    
axes[0, 0].set_ylabel("$y$");
axes[1, 0].set_ylabel("$y$");

fig.tight_layout();

The Regularized Horseshoe

First we place a reasonable half-normal prior on the error variance $\sigma^2$.

with pm.Model() as model:
    σ = pm.HalfNormal("σ", 2.5)

The hierarchical regularized horseshoe uses two levels of regularization, global and local. There is a global parameter, $\tau$, that will shrink all parameters towards zero (similarly to ridge regression) along with local parameters $\lambda_i$ for each coefficient. A long-tailed prior on the $\lambda_i$s allows some of them to be nonzero, with the scale of $\tau$ setting the prior expected number of nonzero parameters.

An important calculation (§3.3 of Piironen and Vehtari) shows that if we believe there are actually $D_*$ nonzero coefficients, the prior on the global shrinkage parameter $\tau$ should be

\[ \tau \sim \textrm{Half-StudentT}_{2} \left(\frac{D_*}{D - D_*} \cdot \frac{\sigma}{\sqrt{N}}\right). \]

Since we know the generating process for this data, we set $D_* = D_0$, although Piironen and Vehtari show that this guess only needs to be order-of-magnitude correct in practice.

with model:
    τ = pm.HalfStudentT("τ", 2, D0 / (D - D0) * σ / np.sqrt(N))

The regularized horseshoe uses the prior

\[\beta_i \sim N\left(0, \tau^2 \cdot \tilde{\lambda}_i^2\right)\]

for the coefficients, where

\[\tilde{\lambda}_i^2 = \frac{c^2 \lambda_i^2}{c^2 + \tau^2 \lambda_i^2}.\]

Given these definitions, it only remains to specify priors on $\lambda_i$ and $c$. As indicated above, we use a long-tailed prior $\lambda_i \sim \textrm{Half-StudentT}_5(1)$ on the local parameters

with model:
    λ = pm.HalfStudentT("λ", 5, shape=D)

Following §2.3 of Piironen and Vehtari, we place an inverse gamma prior on $c^2$, $c^2 \sim \textrm{InverseGamma}(1, 1).$

with model:
    c2 = pm.InverseGamma("c2", 1, 1)

With these priors in place, we define $\tilde{\lambda}_i$

with model:
    λ_ = λ * at.sqrt(c2 / (c2 + τ**2 * λ**2))

To implement $\beta_i \sim N\left(0, \tau^2 \cdot \tilde{\lambda}_i^2\right)$ more efficiently, we use the parameterization

\[ \begin{align*} z_i & \sim N(0, 1), \\ \beta_i & = z_i \cdot \tau \cdot \tilde{\lambda_i}. \end{align*} \]

with model:
    z = pm.Normal("z", 0., 1., shape=D)
    β = pm.Deterministic("β", z * τ * λ_)

Note that it is important to constrain two of the three factors $z_i$, $\tau$, and $\tilde{\lambda}_i$ to be positive, otherwise this model will not be fully identified. Unlike in Piironen and Vehtari, we enforce identifiability by using half-Student T priors on $\tau$ and $\lambda_i$.

Using a relatively flat $N(0, 10^2)$ prior for the intercept, the likelihood of the observed data is

with model:
    β0 = pm.Normal("β0", 0, 10.)
    
    obs = pm.Normal("obs", β0 + at.dot(X, β), σ, observed=y)

We proceed to sample from this model

CHAINS = 3

SAMPLE_KWARGS = {
    'cores': CHAINS,
    'target_accept': 0.99,
    'max_treedepth': 15,
    'random_seed': [SEED + i for i in range(CHAINS)],
    'return_inferencedata': True
}

with model:
    trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β0, z, c2, λ, τ, σ]

100.00% [6000/6000 05:20<00:00 Sampling 3 chains, 0 divergences]

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

The energy plot, BFMIs, and $\hat{R}$ statistics for these samples show no cause for concern.

az.plot_energy(trace);

az.rhat(trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    z        float64 1.006
    β0       float64 1.001
    σ        float64 1.001
    τ        float64 1.002
    λ        float64 1.004
    c2       float64 1.0
    β        float64 1.002

The relative error in recovering the true nonzero parameters is quite good (the largest relative error is approximately 3%).

def post_mean(trace, var_name):
    return trace["posterior"][var_name].mean(dim=("chain", "draw"))

np.abs((post_mean(trace, "β")[:D0] - COEF[:D0]) / COEF[:D0]).max()

<xarray.DataArray 'β' ()>
array(0.03201216)

Visually, the true parameter values lie comfortably inside the 95% high posterior density intervals for these nonzero parameters.

ax, = az.plot_forest(
    trace, var_names=["β"], coords={"β_dim_0": range(D0)},
    kind='ridgeplot', ridgeplot_truncate=False, ridgeplot_alpha=0.5,
    hdi_prob=0.95, combined=True,
    figsize=(8, 6)
)
ax.scatter(COEF[:D0][::-1], ax.get_yticks(),
           c='C1',
           label="Actual value");

ax.set_xlabel(r"$\beta_i$");

ax.set_ylim(bottom=None, top=1.55 * ax.get_yticks().max())
ax.set_yticklabels(range(D0)[::-1]);
ax.set_ylabel(r"$i$");

ax.legend(loc='upper center');
ax.set_title("Posterior distribution of nonzero coefficients");

The (absolute) largest posterior expected value for a coefficient that is actually zero is relatively small compared to the scale of the nonzero parameters.

np.abs(post_mean(trace, "β")[D0:]).max()

<xarray.DataArray 'β' ()>
array(0.2778543)

The following plot shows the posterior distributions of the $D - D_0 = 45$ coefficients. Most of the distributions show a pronounced peak around zero, as we would expect for a sparse estimator.

fig, ax = plt.subplots(figsize=(8, 6))

for i in range(D0, D):
    ax.hist(trace["posterior"]["β"][..., i].data.ravel(),
            histtype='step', bins=100,
            color='C0', alpha=0.25,);

ax.set_xlabel(f"$\\beta_i$\n$i = {D0},...,{D}$");

ax.set_yticks([]);
ax.set_ylabel("Density");

ax.set_title("Posterior distributions");

And surely enough, all but one of the 95% high posterior density intervals for the coefficients that are actually zero contain zero.

post_β_low, post_β_high = trace["posterior"]["β"][..., D0:].quantile([0.025, 0.975], dim=('chain', 'draw'))

((post_β_low <= 0) & (0 <= post_β_high)).mean()

<xarray.DataArray 'β' ()>
array(0.97777778)

This post is available as a Jupyter notebook here.

Many thanks to Sue Márquez for pointing out that the original post spelled Aki’s name wrong, and sincere apologies to Aki.

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sat May 29 2021

Python implementation: CPython
Python version       : 3.8.8
IPython version      : 7.22.0

matplotlib: 3.4.1
arviz     : 0.11.2
numpy     : 1.20.2
pymc3     : 3.11.1
aesara    : 2.0.6
seaborn   : 0.11.1

Carvalho, C. M., Polson, N. G., & Scott, J. G. (2009, April). Handling sparsity via the horseshoe. In Artificial Intelligence and Statistics (pp. 73-80). PMLR.↩︎
Piironen, J., & Vehtari, A. (2017). Sparsity information and regularization in the horseshoe and other shrinkage priors. Electronic Journal of Statistics, 11(2), 5018-5051.↩︎

A PyMC3 Analysis of Tyrannosaurid Growth Curves

Austin Rochford — Mon, 17 May 2021 04:00:00 GMT

In Chapter 7 of their excellent book Bayesian Statistical Methods (BSM from now on), Reich and Ghosh work an example analysis of the growth rates of Tyrannosaurids in R using JAGS. This dataset and analysis are quite fun, so this post shows how to execute (almost) the same analysis in Python using PyCM3.

The Dataset

In addition to statistics, I am a Lego nerd.

To quote BSM,

We analyze the data from 20 fossils to estimate the growth curves of four tyrannosaurid species: Albertosaurus, Daspletosaurus, Gorgosaurus and Tyrannosaurus. The data are taken from Table 1 of Erickson, GM et al (2004).

First we make some standard Python imports and load the dataset from the author’s website.

%matplotlib inline

from warnings import filterwarnings

from aesara import shared, tensor as at
import arviz as az
from matplotlib import MatplotlibDeprecationWarning, pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
import seaborn as sns

filterwarnings('ignore', category=UserWarning, module='aesara')
filterwarnings('ignore', category=UserWarning, module='arviz')
filterwarnings('ignore', category=MatplotlibDeprecationWarning, module='pandas')

sns.set(color_codes=True)

%%bash
DATA_URI=https://www4.stat.ncsu.edu/~bjreich/BSMdata/Erickson.csv
DATA_DEST=/tmp/Erickson.csv

if [[ ! -e $DATA_DEST ]];
then
    wget -q -O $DATA_DEST $DATA_URI
fi

df = (pd.read_csv('/tmp/Erickson.csv')
        .rename(columns=str.lower)
        .sort_values('age')
        .reset_index(drop=True))

df.head()

	taxon	spec. #	age	mass
0	Albertosaurus	RTMP 2002.45.46	2	50.3
1	Tyrannosaurus	LACM 28471	2	29.9
2	Gorgosaurus	FMNH PR2211	5	127.0
3	Gorgosaurus	RTMP 86.144.1	7	229.0
4	Daspletosaurus	RTMP 94.143.1	10	496.0

Exploratory data analysis

First we summarize and plot the data.

df.describe()

	age	mass
count	20.00000	20.000000
mean	15.10000	1568.760000
std	7.05542	1562.465702
min	2.00000	29.900000
25%	13.00000	579.250000
50%	15.50000	1123.500000
75%	18.75000	1795.000000
max	28.00000	5654.000000

grid = sns.catplot(x='age', y='mass', hue='taxon', data=df)

grid.ax.set_xlim(left=1);
grid.ax.set_xscale('log');

grid.ax.set_ylim(bottom=100);
grid.ax.set_yscale('log');

Note the logarithmic scale on both of the axes here. A few points are immediately apparent.

There are only 20 samples, so the data is fairly small.
Mass varies fairly widely across taxa.
The relationship between age and mass appears linear on the log-log scale.
Variance appears to increase with age.

Modeling

Given the third and fourth observations about the data,

[w]e use multiplicative error rather than additive error because variation in the population likely increases with mass/age.

Throughout this post $Y_i$ and $X_i$ will correspond to the mass and age of the $i$-th sample respectively. Each of our four models will take the form

\[Y_i = f_{j(i)}(X_i) \cdot \varepsilon_{j(i)}.\]

Here $\varepsilon_j$ corresponds to the multiplicative error for the $j$-th taxon and $j(i)$ is the taxon identifier of the $i$-th sample.

We place a log-normal prior on $\varepsilon_j$ so that $ (_j) = 1$ and $\textrm{Var}\left(\log(\varepsilon_j)\right) = \sigma_j^2$. With the notation $y_i = \log Y_i$ and $x_i = \log X_i$, standard distributional math shows that

\[y_i \sim N\left(\log\left(f_{j(i)}(X_{j(i)})\right) - \frac{\sigma_{j(i)}^2}{2}, \sigma_{j(i)}^2\right).\]

Each of the four models we consider arises from a choice of the form of $f_j(X)$ (along with a prior for any of its unknown parameters) and a prior on $\sigma_j$.

Unpooled log-linear model

Given the third observation above, we begin with a log-linear model for the relationship between age and mass. If

\[f_j(X) = A_j \cdot X^{b_j}\]

then

\[\log\left(f_j(X)\right) = a_j + b_j \log X,\]

where $a_j = \log A_j$.

We define log_age and log_mass which correspond to $x_i$ and $y_i$ respectively.

log_age = np.log(df['age'].values)
log_mass = np.log(df['mass'].values)

We turn the taxon names into numeric identifiers and set up shared containers for log_age and taxon_id. These shared containers are aesara’s way to facilitate posterior predictive sampling. aesara is the PyMC team’s fork of theano. Since active development on theano has ceased, pymc3 now uses aesara for tensor calculations.

taxon_id, taxon_map = df['taxon'].factorize(sort=True)
n_taxon = taxon_map.size

taxon_id_ = shared(taxon_id)
log_age_ = shared(log_age)

Given the second observation above, we will actually specify one model per taxon (hence the adjective “unpooled”). This model has regression coefficients

\[ \begin{align*} a_j & \sim N(0, 2.5^2) \\ b_j & \sim N(0, 2.5^2). \end{align*} \]

with pm.Model() as loglin_model:
    a = pm.Normal("a", 0., 2.5, shape=n_taxon)
    b = pm.Normal("b", 0., 2.5, shape=n_taxon)

Note that these prior differs differ slightly (in their variance) from those in the BSM. Nevertheless, the results will be quite similar. With the prior on $\sigma_j \sim N(0, 2.5^2)$, we can finish specifying the model in code.

with loglin_model:
    σ = pm.HalfNormal("σ", 2.5, shape=n_taxon)
    μ = a[taxon_id_] + b[taxon_id_] * log_age_ - 0.5 * σ[taxon_id_]**2

    obs = pm.Normal("obs", μ, σ[taxon_id_], observed=log_mass)

We now sample from the posterior distribution of the model.

CHAINS = 3
DRAWS = 2000
SEED = 12345 # for reproducibility

SAMPLE_KWARGS = {
    'cores': CHAINS,
    'draws': DRAWS,
    'random_seed': [SEED + i for i in range(CHAINS)],
    'return_inferencedata': True,
    'target_accept': 0.995
}

with loglin_model:
    loglin_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ, b, a]

100.00% [9000/9000 03:15<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 2_000 draw iterations (3_000 + 6_000 draws total) took 198 seconds.
The number of effective samples is smaller than 25% for some parameters.

To check that there is no obvious problem with our samples we check the energy plot, Bayesian fraction of missing information (BFMI), and Gelman-Rubin statistics.

az.plot_energy(loglin_trace);

Since the marginal and transitional energy distributions are similar and each BFMI is greater than 0.2, these diagnostics show no obvious issues.

az.rhat(loglin_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    a        float64 1.008
    b        float64 1.006
    σ        float64 1.003

The Gelman-Rubin statistics for each parameter are acceptable (smaller than 1.05) as well.

We now draw samples from the posterior predictive distribution of $Y$. First we build a grid for $x$ (pp_log_age) and $j$ (pp_taxon_id) that span the observed range of these quantities in the data.

pp_log_age = np.repeat(np.linspace(0, np.log(30)), n_taxon)
pp_taxon_id = np.tile(np.arange(n_taxon), pp_log_age.size // n_taxon)

Now we sample from the posterior predictive distribution.

ALPHA = 0.05

def get_pp_df(model, trace, alpha=ALPHA):
    taxon_id_.set_value(pp_taxon_id)
    log_age_.set_value(pp_log_age)
    
    with model:
        pp_trace = pm.sample_posterior_predictive(trace)
        
    taxon_id_.set_value(taxon_id)
    log_age_.set_value(log_age)
        
    return pd.DataFrame({
        'age': np.exp(pp_log_age),
        'taxon_id': pp_taxon_id,
        'pp_mean': np.exp(pp_trace['obs'].mean(axis=0)),
        'pp_low': np.exp(np.percentile(pp_trace['obs'], 100 * alpha / 2., axis=0)),
        'pp_high': np.exp(np.percentile(pp_trace['obs'], 100 * (1 - alpha / 2.), axis=0))
    })

Note the lines in get_pp_df where we change the values of the shared variables taxon_id_ and log_age_ to their posterior predictive grid values before sampling and back to the observed values after sampling.

pp_loglin_df = get_pp_df(loglin_model, loglin_trace)

100.00% [6000/6000 00:31<00:00]

pp_loglin_df.head()

	age	taxon_id	pp_mean	pp_low	pp_high
0	1.000000	0	19.767345	7.513964	44.358809
1	1.000000	1	3.751587	0.029910	96.423621
2	1.000000	2	8.544977	2.598875	22.385467
3	1.000000	3	7.263371	2.957169	17.696855
4	1.071878	0	21.720457	8.950703	48.788518

Finally, after all of this preparatory work, we get to the payoff and plot the posterior predictive means and intervals for each taxon across our grid.

def make_posterior_mean_label(name):
    return "Posterior mean" if name is None else f"{name}\nposterior mean"

def make_posterior_interval_label(name):
    return "95% interval" if name is None else f"{name}\n95% interval"

def plot_taxon_posterior(taxon, df, pp_df,
                         name=None, plot_data=True,
                         legend=True, ax=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 6))

    ax.fill_between(pp_df['age'], pp_df['pp_low'], pp_df['pp_high'],
                    alpha=0.25, label=make_posterior_interval_label(name));
    pp_df.plot('age', 'pp_mean', label=make_posterior_mean_label(name), ax=ax)

    if plot_data:
        df.plot.scatter('age', 'mass', c='C0', label="Data", ax=ax)
    
    ax.set_title(taxon)

    if legend:
        ax.legend(loc=2)
    else:
        ax.get_legend().remove()
        
    return ax

FIGSIZE = (12, 9)

def plot_pp_by_taxon(pp_df,
                     name=None, plot_data=True,
                     fig=None, axes=None, figsize=FIGSIZE):
    if fig is None or axes is None:
        fig, axes = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True, figsize=figsize)

    for plt_taxon_id, ax in enumerate(axes.flatten()):
        taxon = taxon_map[plt_taxon_id]
        taxon_df = df[taxon_id == plt_taxon_id]
        taxon_pp_df = pp_df[pp_df['taxon_id'] == plt_taxon_id]

        plot_taxon_posterior(taxon, taxon_df, taxon_pp_df,
                             name=name, plot_data=plot_data,
                             legend=(plt_taxon_id == 0), ax=ax)
        ax.set_xlim(right=30)
        ax.set_ylim(top=6000)
        
    return fig, axes

plot_pp_by_taxon(pp_loglin_df);

Visually comparing these plots to the corresponding ones in the worked JAGS example shows reasonable agreement in spite of our slightly different prior specification.

Pooled log-linear model

It is quite noticeable that in the above posterior predictive plots that the credible intervals are wider for taxa that have fewer data points. In particular there are only three samples of taxon Daspletosaurus. A hierarchical model will allow us to share information from taxa with more observations when inferring the posterior distributions of parameters for taxa with fewer observations. This information sharing leads to the use of the term “pooled” for this type of model. The pooled log linear model still takes the form

\[\log\left(f_j(X)\right) = a_j + b_j \log X,\]

but uses different priors on $a_j$ and $b_j$. Conceptually, the pooled model uses the priors

\[ \begin{align*} \mu_a & \sim N(0, 2.5^2) \\ \sigma_a & \sim \textrm{HalfNormal}(2.5^2) \\ a_j & \sim N(\mu_a, \sigma_a^2), \end{align*} \]

and similarly for $b_j$.

In reality, this intuitive parameterization can often present computational challenges, so we use the following non-centered parameterization instead.

\[ \begin{align*} \mu_a & \sim N(0, 2.5^2) \\ \Delta_{a_j} & \sim N(0, 1) \\ \sigma_a & \sim \textrm{HalfNormal}(2.5^2) \\ a_j & = \mu_a + \Delta_{a_j} \cdot \sigma_a. \end{align*} \]

def hierarchical_normal(name, shape):
    μ = pm.Normal(f"μ_{name}", 0., 2.5)
    Δ = pm.Normal(f"Δ_{name}", 0., 1., shape=shape)
    σ = pm.HalfNormal(f"σ_{name}", 2.5)
    
    return pm.Deterministic(name, μ + Δ * σ)

with pm.Model() as pooled_loglin_model:
    a = hierarchical_normal("a", n_taxon)
    b = hierarchical_normal("b", n_taxon)

To further share information across taxa, we enforce the constraint that the noise parameters are identical across taxa, $\sigma_1 = \sigma_2 = \sigma_3 = \sigma_4 = \sigma$, with the prior $\sigma \sim \textrm{HalfNormal}(2.5^2)$.

with pooled_loglin_model:
    σ = pm.HalfNormal("σ", 2.5)
    μ = a[taxon_id_] + b[taxon_id_] * log_age_ - 0.5 * σ**2

    obs = pm.Normal("obs", μ, σ, observed=log_mass)

We now sample from the pooled log-linear model.

with pooled_loglin_model:
    pooled_loglin_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ, σ_b, Δ_b, μ_b, σ_a, Δ_a, μ_a]

100.00% [9000/9000 08:33<00:00 Sampling 3 chains, 0 divergences]

Sampling 3 chains for 1_000 tune and 2_000 draw iterations (3_000 + 6_000 draws total) took 514 seconds.
The number of effective samples is smaller than 25% for some parameters.

Again our diagnostics show no obvious cause for concern.

az.plot_energy(pooled_loglin_trace);

az.rhat(pooled_loglin_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    μ_a      float64 1.002
    Δ_a      float64 1.002
    μ_b      float64 1.0
    Δ_b      float64 1.001
    σ_a      float64 1.002
    a        float64 1.002
    σ_b      float64 1.004
    b        float64 1.002
    σ        float64 1.006

Again we sample from the posterior predictive distribution of this model and plot the results.

pp_pooled_loglin_df = get_pp_df(pooled_loglin_model, pooled_loglin_trace)

100.00% [6000/6000 00:52<00:00]

pp_pooled_loglin_df.head()

	age	taxon_id	pp_mean	pp_low	pp_high
0	1.000000	0	17.405405	9.378821	30.602004
1	1.000000	1	9.940230	3.453810	27.524651
2	1.000000	2	9.484354	4.619197	18.747739
3	1.000000	3	8.394668	4.823675	15.338217
4	1.071878	0	19.095916	10.369680	33.178229

plot_pp_by_taxon(pp_pooled_loglin_df);

Visually, each of the credible intervals appears smaller than for the unpooled model, espescially for Daspletosaurus. We can overlay the posterior predictive plots for the unpooled and pooled log-linear models to confirm this observation.

fig, axes = plot_pp_by_taxon(pp_loglin_df, plot_data=False, name="Unpooled log-linear")
plot_pp_by_taxon(pp_pooled_loglin_df,
                 name="Pooled log-linear",
                 fig=fig, axes=axes);

Unpooled logistic model

While the log-linear models are reasonable given our exploratory data analysis above (and produce visually plausible results), they are physically unrealistic. With a log-linear model, as the dinosaur gets older its mass will increase without bound. In reality, living creatures have a maximum lifespan, so the log-linear model may be reasonable for realistic ages. Nevertheless, we can incorporate upper bounds on mass by using a logistic function for $f_j$,

\[f_j(x) = a_j + \frac{b_j}{1 + \exp\left(-(x - c_j)\ /\ d_j\right)}\]

def logistic(x, a, b, c, d, tensor=False):
    exp = at.exp if tensor else np.exp
        
    return a + b / (1. + exp(-(x - c) / d))

def inv_logistic(y, a, b, c, d, tensor=False):
    log = at.log if tensor else np.log
    
    return c - d * log(b / (y - a) - 1)

Note that logistic and its inverse function, inv_logistic, switch between the aesera and numpy implementations of exp and log depending on the tensor parameter.

fig, ax = plt.subplots(figsize=(8, 6))

A, B, C, D = 1, 5, 10, 3

f = lambda x: logistic(x, A, B, C, D)
f_inv = lambda y: inv_logistic(y, A, B, C, D)

X_MAX = f_inv(A + B - 0.01)
logistic_x = np.linspace(0, X_MAX)

ax.plot(logistic_x, f(logistic_x));

ax.set_xlim(0, X_MAX);

ax.set_yticks([0, f(0), f(2 * X_MAX)]);
ax.set_yticklabels(["0", r"$f(0)$",
                    r"$\lim_{x \to \infty}\ f(x)$"]);

ax.set_title(f"Logistic function $f$ with\n$a =${A}, $b =${B}, $c =${C}, $d =${D}");

We require $a, b, d > 0$, so we set

\[ \begin{align*} a_j & = \exp \alpha_j \\ b_j & = \exp \beta_j \\ d_j & = \exp \delta_j. \end{align*} \]

For the unpooled logistic model we use $\alpha_j \sim N(0, 2.5^2)$ and similarly for $\beta_j$ and $\delta_j$. We use the same prior for $c_j$ as in the unpooled log-linear model.

with pm.Model() as logistic_model:
    α = pm.Normal("α", 0., 2.5, shape=n_taxon)
    a = pm.Deterministic("a", at.exp(α))
    
    β = pm.Normal("β", 0., 2.5, shape=n_taxon)
    b = pm.Deterministic("b", at.exp(β))
    
    c = pm.Normal("c", 0., 2.5, shape=n_taxon)
    
    δ = pm.Normal("δ", 0., 2.5, shape=n_taxon)
    d = pm.Deterministic("d", at.exp(δ))

The noise variance is also the same as in the unpooled log-linear model. The likelihood is quite similar as well.

with logistic_model:
    σ = pm.HalfNormal("σ", 2.5, shape=n_taxon)
    μ = at.log(
            logistic(log_age_,
                     a[taxon_id_], b[taxon_id_], c[taxon_id_], d[taxon_id_],
                     tensor=True)
        ) - 0.5 * σ[taxon_id_]**2

    obs = pm.Normal("obs", μ, σ[taxon_id_], observed=log_mass)

We now sample from the unpooled logistic model.

with logistic_model:
    logistic_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ, δ, c, β, α]

100.00% [9000/9000 09:42<00:00 Sampling 3 chains, 13 divergences]

Sampling 3 chains for 1_000 tune and 2_000 draw iterations (3_000 + 6_000 draws total) took 583 seconds.
There were 2 divergences after tuning. Increase `target_accept` or reparameterize.
The chain reached the maximum tree depth. Increase max_treedepth, increase target_accept or reparameterize.
There were 6 divergences after tuning. Increase `target_accept` or reparameterize.
There were 5 divergences after tuning. Increase `target_accept` or reparameterize.
The estimated number of effective samples is smaller than 200 for some parameters.

Again our diagnostics show no obvious cause for concern.

az.plot_energy(logistic_trace);

az.rhat(logistic_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    α        float64 1.017
    β        float64 1.02
    c        float64 1.008
    δ        float64 1.008
    a        float64 1.017
    b        float64 1.02
    d        float64 1.014
    σ        float64 1.024

Again we sample from the posterior predictive distribution of this model and plot the results.

pp_logistic_df = get_pp_df(logistic_model, logistic_trace)

100.00% [6000/6000 00:35<00:00]

pp_logistic_df.head()

	age	taxon_id	pp_mean	pp_low	pp_high
0	1.000000	0	20.229437	0.655100	809.706349
1	1.000000	1	372.903045	0.370631	5551.674728
2	1.000000	2	60.408886	0.477719	1983.882017
3	1.000000	3	11.421099	1.119319	53.975799
4	1.071878	0	21.909702	0.812396	642.989663

fig, axes = plot_pp_by_taxon(pp_loglin_df, name="Unpooled log-linear")
plot_pp_by_taxon(pp_logistic_df,
                 name="Unpooled logistic",
                 fig=fig, axes=axes);

Interestingly, the posterior predictive mean for Daspletosaurus doesn’t seem to fit the observed data as well as the log-linear model, and the credible intervals are quite wide. This observation makes sense given that the logistic model has twice as many parameters, but the number of observations has remained the same.

Pooled logistic

We can once again share information across taxa by specifying a pooled logistic model. This model places hierarchical normal priors on $\alpha$, $\beta$, $c$, and $\delta$.

with pm.Model() as pooled_logistic_model:
    α = hierarchical_normal("α", n_taxon)
    a = pm.Deterministic("a", at.exp(α))
    
    β = hierarchical_normal("β", n_taxon)
    b = pm.Deterministic("b", at.exp(β))
    
    c = hierarchical_normal("c", n_taxon)
    
    δ = hierarchical_normal("δ", n_taxon)
    d = pm.Deterministic("d", at.exp(δ))

The noise variance and likelihood are similar to those of the unpooled logistic model, except we constrain all taxa to share the same noise variance, $\sigma$, as in the pooled log-linear model.

with pooled_logistic_model:
    σ = pm.HalfNormal("σ", 2.5)
    μ = at.log(
            logistic(log_age_,
                     a[taxon_id_], b[taxon_id_], c[taxon_id_], d[taxon_id_],
                     tensor=True)
        ) - 0.5 * σ**2

    obs = pm.Normal("obs", μ, σ, observed=log_mass)

with pooled_logistic_model:
    pooled_logistic_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ, σ_δ, Δ_δ, μ_δ, σ_c, Δ_c, μ_c, σ_β, Δ_β, μ_β, σ_α, Δ_α, μ_α]

100.00% [9000/9000 19:02<00:00 Sampling 3 chains, 3 divergences]

Sampling 3 chains for 1_000 tune and 2_000 draw iterations (3_000 + 6_000 draws total) took 1143 seconds.
There were 2 divergences after tuning. Increase `target_accept` or reparameterize.
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.

Again our diagnostics show no obvious cause for concern.

az.plot_energy(pooled_logistic_trace);

(az.rhat(pooled_logistic_trace['posterior'])
   .max()
   .to_array()
   .max())

<xarray.DataArray ()>
array(1.00208561)

Again we sample from the posterior predictive distribution of this model and plot the results.

pp_pooled_logistic_df = get_pp_df(pooled_logistic_model, pooled_logistic_trace)

100.00% [6000/6000 01:08<00:00]

pp_pooled_logistic_df.head()

	age	taxon_id	pp_mean	pp_low	pp_high
0	1.000000	0	22.852762	10.267422	56.497339
1	1.000000	1	9.839328	0.906120	125.299677
2	1.000000	2	12.589540	3.182979	91.279583
3	1.000000	3	11.080350	4.344195	36.627854
4	1.071878	0	24.465512	11.729286	56.373575

fig, axes = plot_pp_by_taxon(pp_logistic_df, plot_data=False, name="Unpooled logistic")
plot_pp_by_taxon(pp_pooled_logistic_df,
                 name="Pooled logistic",
                 fig=fig, axes=axes);

We see that the posterior predictive means for the pooled logistic model fit better than those from the unpooled model, and that the credible intervals are again significantly smaller.

We now compare the posterior predictions for the pooled log-linear and logistic models.

fig, axes = plot_pp_by_taxon(pp_pooled_loglin_df, plot_data=False, name="Pooled log-linear")
plot_pp_by_taxon(pp_pooled_logistic_df,
                 name="Pooled logistic",
                 fig=fig, axes=axes);

These models agree with each other fairly well, espescially in age ranges where the taxon data is concentrated, with the logistic model making slightly more plausible predictions as age increases.

Model comparison

We now compare these four models using two of the information criteria available in arviz. While BSM uses the deviance information criterion (DIC) to compare models, DIC is no longer supported in pymc3 or arviz, so we use Pareto-smoothed importance sampling leave-one-out cross validation (PSIS-LOO) and the widely applicable information criterion (WAIC) to compare models. These information criteria are available in arviz. For details on these methods, consult the pymc3 documentation.

traces = {
    'Unpooled log-linear': loglin_trace,
    'Pooled log-linear': pooled_loglin_trace,
    'Unpooled logistic': logistic_trace,
    'Pooled logistic': pooled_logistic_trace
}

def plot_compare(comp_df, ic, height=6, fig=None, axes=None):
    if fig is None or axes is None:
        fig, axes = plt.subplots(
            ncols=2, figsize=(16 * (height / 6), height)
        )
        
    score_ax, weight_ax = axes

    az.plot_compare(comp_df, insample_dev=False, plot_ic_diff=False, ax=score_ax)
    
    score_ax.set_xlabel("Log score")
    score_ax.set_ylabel("Model")

    comp_df['weight'].plot.barh(ax=weight_ax)

    weight_ax.set_xscale('log')
    weight_ax.set_xlim(right=2)
    weight_ax.set_xlabel("Stacking weight")

    weight_ax.set_yticklabels([])
    weight_ax.invert_yaxis()

    fig.suptitle(f"{ic} comparison")
    fig.tight_layout()
    
    return fig, axes

The following plots show the PSIS-LOO comparison for these models.

loo_df = az.compare(traces, ic='loo', seed=SEED)
plot_compare(loo_df, "LOO");

Note that larger scores are better, so PSIS-LOO favors the pooled models over the unpooled models and favors the pooled logistic model most highly of all.

waic_df = az.compare(traces, ic='waic', seed=SEED)
plot_compare(waic_df, "WAIC");

Using WAIC produces quite similar results.

Interestingly, the results in BSM using DIC favor the logistic model slightly.

Many thanks to Reich and Ghosh for their excellent textbook which brought this fun dataset to my attention and to Kiril Zvezdarov for his comments on early drafts.

The notebook this post was generated from is available here.

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sun May 16 2021

Python implementation: CPython
Python version       : 3.8.8
IPython version      : 7.22.0

pymc3     : 3.11.1
seaborn   : 0.11.1
pandas    : 1.2.3
arviz     : 0.11.2
matplotlib: 3.4.1
numpy     : 1.20.2
aesara    : 2.0.6

Evaluating Precision in Sports Analytics

Austin Rochford — Thu, 20 Dec 2018 05:00:00 GMT

One of my data anaylsis pet peeves is false precision. Just because it is possible calculate a quantity to three decimal places doesn’t mean all of those decimal places are meaningful. This post explores how much precision is justified in the context of two common sports statistics: batting average in Major League Baseball and save percentage in the National Hockey League. Using Bayesian hierarchical models, we find that though these quantities are conventionally calculated to three decimal places, only the first two decimal places of precision are justified.

%matplotlib inline

import arviz as az
from matplotlib import pyplot as plt
from matplotlib.ticker import StrMethodFormatter
import numpy as np
import pandas as pd
import pymc3 as pm
import scipy as sp
import seaborn as sns

sns.set(color_codes=True)

svpct_formatter = ba_formatter = StrMethodFormatter("{x:.3f}")

Batting average

We begin by loading hitting data for the 2018 MLB season from Baseball Reference.

def get_data_url(filename):
    return f"https://www.austinrochford.com/resources/sports_precision/{filename}"

def load_data(filepath, player_col, usecols):
    df = pd.read_csv(filepath, usecols=[player_col] + usecols)
    
    return (pd.concat((df[player_col]
                         .str.split('\\', expand=True)
                         .rename(columns={0: 'name', 1: 'player_id'}),
                       df.drop(player_col, axis=1)),
                      axis=1)
              .rename(columns=str.lower)
              .groupby(['player_id', 'name'])
              .first() # players that switched teams have their entire season stats listed once per team
              .reset_index('name'))

mlb_df = load_data(get_data_url('2018_batting.csv'), 'Name', ['AB', 'H'])

mlb_df.head()

	name	ab	h
player_id
abreujo02	Jose Abreu	499	132
acunaro01	Ronald Acuna	433	127
adamewi01	Willy Adames	288	80
adamja01	Jason Adam	0	0
adamsau02	Austin L. Adams	0	0

batter_df = mlb_df[mlb_df['ab'] > 0]
n_player, _ = batter_df.shape

This data set covers nearly 1,000 MLB players.

n_player

Batting average(https://en.wikipedia.org/wiki/Batting_average_(baseball%29) is the most basic summary of a player’s batting performance and is defined as their number of hits divided by their number of at bats. In order to assess the amount of precision that is justified when calculating batting average, we build a hierarchical logistic model. Let $n_i$ be the number of at bats for the $i$-th player and let $y_i$ be their number of hits. Our model is

\[ \begin{align*} \mu_{\eta} & \sim N(0, 5^2) \\ \sigma_{\eta} & \sim \textrm{Half-}N(2.5^2) \\ \eta_i & \sim N(\mu, \sigma_{\eta}^2) \\ \textrm{ba}_i & = \textrm{sigm}(\eta_i) \\ y_i\ |\ n_i & \sim \textrm{Binomial}(n_i, \textrm{ba}_i). \end{align*} \]

We specify this model in pymc3 below using a non-centered parametrization.

def hierarchical_normal(name, shape, μ=None):
    if μ is None:
        μ = pm.Normal(f"μ_{name}", 0., 5.)
    
    Δ = pm.Normal(f"Δ_{name}", shape=shape)
    σ = pm.HalfNormal(f"σ_{name}", 2.5)
    
    return pm.Deterministic(name, μ + Δ * σ)

with pm.Model() as mlb_model:
    η = hierarchical_normal("η", n_player)
    ba = pm.Deterministic("ba", pm.math.sigmoid(η))
    
    hits = pm.Binomial("hits", batter_df['ab'], ba, observed=batter_df['h'])

We proceeed to sample from the model’s posterior distribution.

CHAINS = 3
SEED = 88564 # from random.org, for reproducibility

SAMPLE_KWARGS = {
    'draws': 1000,
    'tune': 1000,
    'chains': CHAINS,
    'cores': CHAINS,
    'random_seed': list(SEED + np.arange(CHAINS))
}

with mlb_model:
    mlb_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ_η, Δ_η, μ_η]
Sampling 3 chains: 100%|██████████| 6000/6000 [00:44<00:00, 133.94draws/s]

Before drawing conclusions from the posterior samples, we use arviz to verify that there are no obvious problems with the sampler diagnostics.

az.plot_energy(mlb_trace);

az.gelman_rubin(mlb_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    μ_η      float64 1.01
    Δ_η      float64 1.0
    σ_η      float64 1.0
    η        float64 1.0
    ba       float64 1.0

First we’ll examine the posterior distribution of Mike Trout’s batting average.

fig, ax = plt.subplots(figsize=(8, 6))

trout_ix = (batter_df.index == 'troutmi01').argmax()
ax.hist(mlb_trace['ba'][:, trout_ix], bins=30, alpha=0.5);

ax.vlines(batter_df['h']
                   .div(batter_df['ab'])
                   .loc['troutmi01'],
          0, 275,
          linestyles='--',
          label="Actual batting average");

ax.xaxis.set_major_formatter(ba_formatter);
ax.set_xlabel("Batting average");

ax.set_ylabel("Posterior density");

ax.legend();
ax.set_title("Mike Trout");

We see that the posterior places significant mass between .260 and .320, quite a wide range of batting averages. This range roughly corresponds to the 95% credible interval for his 2018 batting average.

np.percentile(mlb_trace['ba'][:, trout_ix], [2.5, 97.5])

array([ 0.25516468,  0.32704036])

We will use the width of the 95% credible interval for each player’s batting average to determine how many digits of precision are justified.

mlb_df = batter_df.assign(
    width_95=sp.stats.iqr(mlb_trace["ba"], axis=0, rng=(2.5, 97.5))
)

The following plot shows the width of these intervals, grouped by the number of at bats the player had in 2018.

def plot_ci_width(grouped, width):
    fig, ax = plt.subplots(figsize=(8, 6))

    low = grouped.quantile(0.025)
    high = grouped.quantile(0.975)

    ax.fill_between(low.index, low, high,
                    alpha=0.25,
                    label=f"{width:.0%} interval");

    grouped.mean().plot(ax=ax, label="Average")

    ax.set_ylabel("Width of 95% credible interval");
    ax.legend(loc=0);
    
    return ax

ax = plot_ci_width(mlb_df['width_95'].groupby(mlb_df['ab'].round(-2)), 0.95)

ax.set_xlim(0, mlb_df['ab'].max());
ax.set_xlabel("At bats");

ax.set_ylim(bottom=0.);
ax.yaxis.set_major_formatter(ba_formatter);

ax.set_title("Batting average");

We see that, on average, about 100 at bats are required to justify a single digit of precision in a player’s batting average. Even in the limit of very many at bats (600 at bats corresponds to just under four at bats per game across a 162 game season) the 95% credible interval has an average width approaching 0.060. This limit indicates that batting average is at most meaningful to the second digit, and even the second digit has a fair bit of uncertainty. This result is not surprising; calculating batting average to three decimal places is a historical convention, but I don’t think many analysts rely on the third digit for their decisions/arguments. While intuitive, it is pleasant to have a somewhat rigorous justification for this practice.

Save percentage

We apply a similar analysis to save percentage in the NHL. First we load 2018 goaltending data from Hockey Reference.

nhl_df = load_data(get_data_url('2017_2018_goalies.csv'), 'Player', ['SA', 'SV'])

nhl_df.head()

	name	sa	sv
player_id
allenja01	Jake Allen	1614	1462
andercr01	Craig Anderson	1768	1588
anderfr01	Frederik Andersen	2211	2029
appleke01	Ken Appleby	55	52
bernijo01	Jonathan Bernier	1092	997

n_goalie, _ = nhl_df.shape

This data set consists of the goaltending performance of just under 100 players.

n_goalie

Our save percentage model is almost identical to the batting average model. Let $n_i$ be the number of at shots the $i$-th goalie faced and let $y_i$ be the number of saves they made. The model is

\[ \begin{align*} \mu_{\eta} & \sim N(0, 5^2) \\ \sigma_{\eta} & \sim \textrm{Half-}N(2.5^2) \\ \eta_i & \sim N(\mu, \sigma_{\eta}^2) \\ \textrm{svp}_i & = \textrm{sigm}(\eta_i) \\ y_i\ |\ n_i & \sim \textrm{Binomial}(n_i, \textrm{svp}_i). \end{align*} \]

with pm.Model() as nhl_model:
    η = hierarchical_normal("η", n_goalie)
    svp = pm.Deterministic("svp", pm.math.sigmoid(η))
    
    saves = pm.Binomial("saves", nhl_df['sa'], svp, observed=nhl_df['sv'])

with nhl_model:
    nhl_trace = pm.sample(nuts_kwargs={'target_accept': 0.9}, **SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ_η, Δ_η, μ_η]
Sampling 3 chains: 100%|██████████| 6000/6000 [00:17<00:00, 338.38draws/s]

Once again, the convergence diagnostics show no cause for concern.

az.plot_energy(nhl_trace);

az.gelman_rubin(nhl_trace).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    μ_η      float64 1.0
    Δ_η      float64 1.0
    σ_η      float64 1.0
    η        float64 1.0
    svp      float64 1.0

We examine the posterior distribution of Sergei Bobrovsky’s save percentage.

fig, ax = plt.subplots(figsize=(8, 6))

bobs_ix = (nhl_df.index == 'bobrose01').argmax()
ax.hist(nhl_trace['svp'][:, bobs_ix], bins=30, alpha=0.5);

ax.vlines(nhl_df['sv']
                .div(nhl_df['sa'])
                .loc['bobrose01'],
          0, 325,
          linestyles='--',
          label="Actual save percentage");

ax.xaxis.set_major_formatter(ba_formatter);
ax.set_xlabel("Save percentage");

ax.set_ylabel("Posterior density");

ax.legend(loc=2);
ax.set_title("Sergei Bobrovsky");

We see that the posterior places significant mass between .905 and .925. We see below that the best and worst save percentages (for goalies that faced at least 200 shots in 2018) are separated by about 0.070.

(nhl_df['sv']
       .div(nhl_df['sa'])
       [nhl_df['sa'] > 200]
       .quantile([0., 1.]))

0.0    0.866995
1.0    0.933712
dtype: float64

Sergei Bobrovsky’s 0.020-wide credible interval is a significant proportion of this 0.070 total range.

np.percentile(nhl_trace['svp'][:, bobs_ix], [2.5, 97.5])

array([ 0.90683748,  0.92526507])

As with batting average, we plot the width of each goalie’s interval, grouped by the number of shots they faced.

nhl_df = nhl_df.assign(
    width_95=sp.stats.iqr(nhl_trace["svp"], axis=0, rng=(2.5, 97.5))
)

ax = plot_ci_width(nhl_df['width_95'].groupby(nhl_df['sa'].round(-2)), 0.95)

ax.set_xlim(0, nhl_df['sa'].max());
ax.set_xlabel("Shots against");

ax.set_ylim(bottom=0.);
ax.yaxis.set_major_formatter(svpct_formatter);

ax.set_title("Save percentage");

This plot shows that even goalies that face many (2000+) shots have credible intervals wider that 0.010, a signifcant proportion of the total variation between goalies.

This post is available as a Jupyter notebook here.

Monotonic Effects in PyMC3

Austin Rochford — Sat, 10 Nov 2018 05:00:00 GMT

Last week I came across the following tweet from Paul Bürkner about a paper he coauthored about including ordinal predictors in Bayesian regression models, and I thought the approach was very clever.

Have you ever wondered how to handle ordinal predictors in your regression models? We propose a simple and intuitive method that is readily available via #brms and @mcmc_stan: https://t.co/dKg4AphvsG
— Paul Bürkner (@paulbuerkner) November 2, 2018

The code in the paper uses brms and Stan to illustrate these concepts. In this post I’ll be replicating some of the paper’s analysis in Python and PyMC3, mostly for my own edification.

%matplotlib inline

from itertools import product

import arviz as az
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
from rpy2.robjects import pandas2ri, r
import seaborn as sns
from theano import shared

sns.set()

The paper uses data from the ordPens R package, which we download and load into a Pandas DataFrame.

%%bash
if [[ ! -e ~/data/ordPens ]];
then
    mkdir -p data
    wget -q -O data/ordPens_0.3-1.tar.gz ~/data/ https://cran.r-project.org/src/contrib/ordPens_0.3-1.tar.gz
    tar xzf data/ordPens_0.3-1.tar.gz -C data
fi

!ls data/ordPens/data/

ICFCoreSetCWP.RData

pandas2ri.activate()

r.load('data/ordPens/data/ICFCoreSetCWP.RData');
all_df = r['ICFCoreSetCWP']

all_df.head()

	b1602	b122	b126	b130	b134	b140	b147	b152	b164	b180	…	e450	e455	e460	e465	e570	e575	e580	e590	s770	phcs
1	0	1	2	1	1	0	0	2	0	0	…	4	0	0	0	3	3	4	3	0	44.33
2	3	2	2	3	3	2	3	3	3	1	…	3	3	2	2	2	2	2	2	2	21.09
3	0	1	2	1	1	0	1	2	0	0	…	4	0	0	0	3	3	4	0	0	41.74
4	0	0	0	2	1	0	0	0	0	0	…	2	2	-1	0	0	2	2	1	1	33.96
5	0	0	0	0	0	0	0	0	0	0	…	1	0	0	0	0	0	0	0	0	46.29

5 rows × 68 columns

The variable of interest here is phcs, which is a subjective physical health score. The predictors we are interested in are d450 and d455.

df = all_df[['d450', 'd455', 'phcs']]

df.head()

	d450	d455	phcs
1	0	2	44.33
2	3	3	21.09
3	0	2	41.74
4	3	2	33.96
5	0	0	46.29

These predictors are ratings on a five-point scale (0-4) of the patient’s impairment while walking (d450) and moving around (d455). For more information on this data, consult the ordPens documentation.

The following plots show a fairly strong monotonic relationship between d450, d455, and phcs.

fig, (d450_ax, d455_ax) = plt.subplots(ncols=2, sharey=True, figsize=(16, 6))

sns.stripplot('d450', 'phcs', data=df,
              jitter=0.1, color='C0', alpha=0.75,
              ax=d450_ax);
sns.stripplot('d455', 'phcs', data=df,
              jitter=0.1, color='C0', alpha=0.75,
              ax=d455_ax);

d455_ax.set_ylabel("");
fig.tight_layout();

The big idea of the paper is to include monotonic effects due to these ordinal predictors as follows. A scalar $b \sim N(0, 10^2)$ parameterizes the overall strength and direction of the relationship, and a Dirichlet vector $\xi \sim \textrm{Dirichlet}(1, \ldots, 1)$ encodes how much of $b$ is gained at each level. The parameters $b$ and $\xi$ are combined into

\[mo(i) = b \sum_{k = 0}^i \xi_k\]

which can be included as a term in a regression model. It is evident that if $i < j$ then $mo(i) \leq mo(j)$ since

\[mo(j) - mo(i) = b \sum_{k = i + 1}^j \xi_k \geq 0\]

and therefore the effect of this term will be monotonic as desired.

The following function constructs this distribution in PyMC3.

def monotonic_prior(name, n_cat):
    b = pm.Normal(f'b_{name}', 0., 10.)
    ξ = pm.Dirichlet(f'ξ_{name}', np.ones(n_cat))

    return pm.Deterministic(f'mo_{name}', b * ξ.cumsum())

With this notation in hand, our model is

\[ \begin{align*} \mu_i & = \beta_0 + mo_{\textrm{d450}}(j_i) + mo_{\textrm{d455}}(k_i) \\ \beta_0 & \sim N(0, 10^2) \\ y_i & \sim N(\mu_i, \sigma^2) \\ \sigma & \sim \textrm{HalfNormal}(5^2) \end{align*} \]

where $j_i$ and $k_i$ are the level of d450 and d455 for the $i$-th patient respectively and $y_i$ is that patient’s phcs score.

We now express this model in PyMC3.

d450 = df['d450'].values
d450_cats = np.unique(d450)
d450_n_cat = d450_cats.size
d450_ = shared(d450)

d455 = df['d455'].values
d455_cats = np.unique(d455)
d455_n_cat = d455_cats.size
d455_ = shared(d455)

phcs = df['phcs'].values

with pm.Model() as model:
    β0 = pm.Normal('β0', 0., 10.)

    mo_d450 = monotonic_prior('d450', d450_n_cat)
    mo_d455 = monotonic_prior('d455', d455_n_cat)
    
    μ = β0 + mo_d450[d450_] + mo_d455[d455_]
    σ = pm.HalfNormal('σ', 5.)
    phcs_obs = pm.Normal('phcs', μ, σ, observed=phcs)

We now sample from the model’s posterior distribution.

CHAINS = 3
SEED = 934520 # from random.org

SAMPLE_KWARGS = {
    'draws': 1000,
    'tune': 1000,
    'chains': CHAINS,
    'random_seed': list(SEED + np.arange(CHAINS))
}

with model:
    trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 2 jobs)
NUTS: [σ, ξ_d455, b_d455, ξ_d450, b_d450, β0]
Sampling 3 chains: 100%|██████████| 6000/6000 [00:41<00:00, 145.59draws/s]
The number of effective samples is smaller than 25% for some parameters.

We use arviz to check the performance of our sampler.

inf_data = az.convert_to_inference_data(trace)

The energy plot, BFMI, and Gelman-Rubin statistics show no cause for concern.

az.plot_energy(inf_data);

az.gelman_rubin(inf_data).max()

<xarray.Dataset>
Dimensions:  ()
Data variables:
    β0       float64 1.0
    b_d450   float64 1.0
    b_d455   float64 1.0
    ξ_d450   float64 1.0
    mo_d450  float64 1.0
    ξ_d455   float64 1.0
    mo_d455  float64 1.0
    σ        float64 1.0

We now sample from the model’s posterior predictive distribution and visualize the results.

pp_d450, pp_d455 = np.asarray(list(zip(*product(d450_cats, d455_cats))))

d450_.set_value(pp_d450)
d455_.set_value(pp_d455)

with model:
    pp_trace = pm.sample_posterior_predictive(trace)

100%|██████████| 3000/3000 [00:07<00:00, 388.49it/s]

pp_df = pd.DataFrame({
    'd450': pp_d450,
    'd455': pp_d455,
    'pp_phcs_mean': pp_trace['phcs'].mean(axis=0)
})

The important feature of this encoding of ordinal predictors is that the $\xi$ parameters allow different levels of the predictor to contribute to result in a different change in the effect, which is in contrast to what happens when these are included as linear predictors, which is quite common in the literature.

REF_CAT = 1

fig, (d450_ax, d455_ax) = plt.subplots(ncols=2, sharey=True, figsize=(16, 6))

(pp_df.pivot_table('pp_phcs_mean', 'd450', 'd455')
      .plot(marker='o', ax=d450_ax));

d450_ax.set_xticks(d450_cats);
d450_ax.set_ylabel("Posterior predictive phcs");

(pp_df.pivot_table('pp_phcs_mean', 'd455', 'd450')
      .plot(marker='o', ax=d455_ax));

d455_ax.set_xticks(d455_cats);

fig.tight_layout();

The following plot corresponds to Figure 3 in the original paper, and the dark lines agree with the mean in that figure quite well.

fig, (d450_ax, d455_ax) = plt.subplots(ncols=2, sharey=True, figsize=(16, 6))

(pp_df[pp_df['d455'] != REF_CAT]
      .pivot_table('pp_phcs_mean', 'd450', 'd455')
      .plot(marker='o', c='k', alpha=0.5, legend=False,
            ax=d450_ax));
(pp_df[pp_df['d455'] == REF_CAT]
      .plot('d450', 'pp_phcs_mean',
            marker='o', c='k',
            label=f"Refernce category (d455 = {REF_CAT})",
            ax=d450_ax));

d450_ax.set_xticks(d450_cats);
d450_ax.set_ylabel("Posterior excpected phcs");

(pp_df[pp_df['d450'] != REF_CAT]
      .pivot_table('pp_phcs_mean', 'd455', 'd450')
      .plot(marker='o', c='k', alpha=0.5, legend=False,
            ax=d455_ax));
(pp_df[pp_df['d450'] == REF_CAT]
      .plot('d455', 'pp_phcs_mean',
            marker='o', c='k',
            label=f"Refernce category (d450 = {REF_CAT})",
            ax=d455_ax));

d455_ax.set_xticks(d455_cats);

fig.tight_layout();

For reference, we compare this model to a model that includes d450 and d455 as linear predictors. This model is given by

\[ \begin{align*} \mu_i & = \beta_0 + \beta_{\textrm{d450}} \cdot j(i) + \beta_{\textrm{d455}} \cdot k(i) \\ \beta_0, \beta_{\textrm{d450}}, \beta_{\textrm{d455}} & \sim N(0, 10^2) \\ y_i & \sim N(\mu_i, \sigma^2) \\ \sigma & \sim \textrm{HalfNormal}(5^2) \end{align*} \]

d450_.set_value(d450)
d455_.set_value(d455)

with pm.Model() as linear_model:
    β0 = pm.Normal('β0', 0., 10.)
    β_d450 = pm.Normal('β_d450', 0., 10.)
    β_d455 = pm.Normal('β_d455', 0., 10.)
    
    μ = β0 + β_d450 * d450_ + β_d455 * d455_
    σ = pm.HalfNormal('σ', 5.)
    phcs_obs = pm.Normal('phcs', μ, σ, observed=phcs)

with linear_model:
    linear_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 2 jobs)
NUTS: [σ, β_d455, β_d450, β0]
Sampling 3 chains: 100%|██████████| 6000/6000 [00:07<00:00, 771.92draws/s]

As in the paper, compare these models by stacking their posterioir predictive distributions.

comp_df = (pm.compare({
                model: trace,
                linear_model: linear_trace
             })
             .rename({
                 0: "Paper",
                 1: "Linear"
             }))

comp_df

	WAIC	pWAIC	dWAIC	weight	SE	dSE	var_warn
Paper	2825.24	6.22	0	1	29.01	0	0
Linear	2830.2	3.7	4.97	0	29.09	4.42	0

We see that the model from the paper has a lower WAIC and gets 100% of the weight, a strong sign that it is surperior to the linear model.

This post is available as a Jupyter notebook here.

An Improved Analysis of NBA Foul Calls with Python

Austin Rochford — Sun, 04 Feb 2018 05:00:00 GMT

Last April, I wrote a post that used Bayesian item-response theory models to analyze NBA foul call data. Last November, I spoke about a greatly improved version of these models at PyData NYC. This post is a write-up of the models from that talk.

Last Two-minute Report

Since late in the 2014-2015 season, the NBA has issued last two minute reports. These reports give the league’s assessment of the correctness of foul calls and non-calls in the last two minutes of any game where the score difference was three or fewer points at any point in the last two minutes.

These reports are notably different from play-by-play logs, in that they include information on non-calls for notable on-court interactions. This non-call information presents a unique opportunity to study the factors that impact foul calls. There is a level of subjectivity inherent in the the NBA’s definition of notable on-court interactions which we attempt to mitigate later using season-specific factors.

Loading the data

Russel Goldenberg of The Pudding has been scraping the PDFs that the NBA publishes and transforming them into a CSV for some time. I am grateful for his work, which has enabled this analysis.

We download the data locally to be kind to GitHub.

%matplotlib inline

import datetime
from itertools import product
import logging
import pickle

from matplotlib import pyplot as plt
from matplotlib.offsetbox import AnchoredText
from matplotlib.ticker import FuncFormatter, StrMethodFormatter
import numpy as np
import pandas as pd
import scipy as sp
import seaborn as sns
from sklearn.preprocessing import LabelEncoder
from theano import tensor as tt

pct_formatter = StrMethodFormatter('{x:.1%}')

sns.set()
blue, green, *_ = sns.color_palette()

plt.rc('figure', figsize=(8, 6))

LABELSIZE = 14
plt.rc('axes', labelsize=LABELSIZE)
plt.rc('axes', titlesize=LABELSIZE)
plt.rc('figure', titlesize=LABELSIZE)
plt.rc('legend', fontsize=LABELSIZE)
plt.rc('xtick', labelsize=LABELSIZE)
plt.rc('ytick', labelsize=LABELSIZE)

SEED = 207183 # from random.org, for reproducibility

# keep theano from complaining about compile locks for small models
(logging.getLogger('theano.gof.compilelock')
        .setLevel(logging.CRITICAL))

%%bash
DATA_URI=https://raw.githubusercontent.com/polygraph-cool/last-two-minute-report/32f1c43dfa06c2e7652cc51ea65758007f2a1a01/output/all_games.csv
DATA_DEST=/tmp/all_games.csv

if [[ ! -e $DATA_DEST ]];
then
    wget -q -O $DATA_DEST $DATA_URI
fi

We use only a subset of the columns in the source data set.

USECOLS = [
    'period',
    'seconds_left',
    'call_type',
    'committing_player',
    'disadvantaged_player',
    'review_decision',
    'play_id',
    'away',
    'home',
    'date',
    'score_away',
    'score_home',
    'disadvantaged_team',
    'committing_team'
]

orig_df = pd.read_csv(
    '/tmp/all_games.csv',
    usecols=USECOLS,
    index_col='play_id',
    parse_dates=['date']
)

The data set contains more than 16,000 plays.

orig_df.shape[0]

Each row of the DataFrame represents a play and each column describes an attrbiute of the play:

period is the period of the game,
seconds_left is the number of seconds remaining in the game,
call_type is the type of call,
committing_player and disadvantaged_player are the names of the players involved in the play,
review_decision is the opinion of the league reviewer on whether or not the play was called correctly:
- review_decision = "INC" means the call was an incorrect noncall,
- review_decision = "CNC" means the call was an correct noncall,
- review_decision = "IC" means the call was an incorrect call, and
- review_decision = "CC" means the call was an correct call,
away and home are the abbreviations of the teams involved in the game,
date is the date on which the game was played,
score_away and score_home are the scores of the away and home team during the play, respectively, and
disadvantaged_team and committing_team indicate how each team is involved in the play.

orig_df.head(n=2).T

play_id	20150301CLEHOU-0	20150301CLEHOU-1
period	Q4	Q4
seconds_left	112	103
call_type	Foul: Shooting	Foul: Shooting
committing_player	Josh Smith	J.R. Smith
disadvantaged_player	Kevin Love	James Harden
review_decision	CNC	CC
away	CLE	CLE
home	HOU	HOU
date	2015-03-01 00:00:00	2015-03-01 00:00:00
score_away	103	103
score_home	105	105
disadvantaged_team	CLE	HOU
committing_team	HOU	CLE

Research questions

In this post, we answer two questions:

How does game context impact foul calls?
Is (not) committing and/or drawing fouls a measurable player skill?

The previous post focused on the second question, and gave the first question only a cursory treatment. This post enhances our treatment of the first question, in order to control for non-skill factors influencing foul calls (namely intentional fouls). Controlling for these factors makes our estimates of player skill more realistic.

Exploratory Data Analysis

First we examine the types of calls present in the data set.

(orig_df['call_type']
        .value_counts()
        .head(n=15))

Foul: Personal                     4736
Foul: Shooting                     4201
Foul: Offensive                    2846
Foul: Loose Ball                   1316
Turnover: Traveling                 779
Instant Replay: Support Ruling      607
Foul: Defense 3 Second              277
Instant Replay: Overturn Ruling     191
Foul: Personal Take                 172
Turnover: 3 Second Violation        139
Turnover: 24 Second Violation       126
Turnover: 5 Second Inbound           99
Stoppage: Out-of-Bounds              96
Violation: Lane                      84
Foul: Away from Play                 82
Name: call_type, dtype: int64

The portion of call_type before the colon is the general category of the call. We count the occurence of these categories below.

(orig_df['call_type']
        .str.split(':', expand=True)
        .iloc[:, 0]
        .value_counts()
        .plot(
            kind='bar',
            color=blue, logy=True, 
            title="Call types"
        )
        .set_ylabel("Frequency"));

We restrict our attention to foul calls, though other call types would be interesting to study in the future.

foul_df = orig_df[
    orig_df['call_type']
           .fillna("UNKNOWN")
           .str.startswith("Foul")
]

We count the foul call types below.

(foul_df['call_type']
        .str.split(': ', expand=True)
        .iloc[:, 1]
        .value_counts()
        .plot(
            kind='bar',
            color=blue, logy=True,
            title="Foul Types"
        )
        .set_ylabel("Frequency"));

We restrict our attention to the five foul types below, which generally involve two players. This subset of fouls allows us to pursue our second research question in the most direct manner.

FOULS = [
    f"Foul: {foul_type}"
        for foul_type in [
            "Personal",
            "Shooting",
            "Offensive",
            "Loose Ball",
            "Away from Play"
        ]
]

Data transformation

There are a number of misspelled team names in the data, which we correct.

TEAM_MAP = {
    "NKY": "NYK",
    "COS": "BOS",
    "SAT": "SAS",
    "CHi": "CHI",
    "LA)": "LAC",
    "AT)": "ATL",
    "ARL": "ATL"
}

def correct_team_name(col):
    def _correct_team_name(df):
        return df[col].apply(lambda team_name: TEAM_MAP.get(team_name, team_name))
    
    return _correct_team_name

We also convert each game date to an NBA season.

def date_to_season(date):
    if date >= datetime.datetime(2017, 10, 17):
        return '2017-2018'
    elif date >= datetime.datetime(2016, 10, 25):
        return '2016-2017'
    elif date >= datetime.datetime(2015, 10, 27):
        return '2015-2016'
    else:
        return '2014-2015'

We clean the data by

restricting to plays that occured during the last two minutes of regulation,
imputing incorrect noncalls when review_decision is missing,
correcting team names,
converting game dates to seasons,
restricting to the foul types discussed above,
restricting to the plays that happened during the 2015-2016 and 2016-2017 regular seasons (those are the only full seasons in the data set as of February 2018), and
dropping unneeded rows and columns.

clean_df = (foul_df.where(lambda df: df['period'] == "Q4")
                   .where(lambda df: (df['date'].between(datetime.datetime(2016, 10, 25),
                                                         datetime.datetime(2017, 4, 12))
                                     | df['date'].between(datetime.datetime(2015, 10, 27),
                                                          datetime.datetime(2016, 5, 30)))
                   )
                   .assign(
                       review_decision=lambda df: df['review_decision'].fillna("INC"),
                       committing_team=correct_team_name('committing_team'),
                       disadvantged_team=correct_team_name('disadvantaged_team'),
                       away=correct_team_name('away'),
                       home=correct_team_name('home'),
                       season=lambda df: df['date'].apply(date_to_season)
                   )
                   .where(lambda df: df['call_type'].isin(FOULS))
                   .dropna()
                   .drop('period', axis=1)
                   .assign(call_type=lambda df: (df['call_type']
                                                   .str.split(': ', expand=True)  
                                                   .iloc[:, 1])))

About 55% of the rows in the original data set remain.

clean_df.shape[0] / orig_df.shape[0]

0.5516564417177914

clean_df.head(n=2).T

play_id	20151028INDTOR-1	20151028INDTOR-2
seconds_left	89	73
call_type	Shooting	Shooting
committing_player	Ian Mahinmi	Bismack Biyombo
disadvantaged_player	DeMar DeRozan	Paul George
review_decision	CC	IC
away	IND	IND
home	TOR	TOR
date	2015-10-28 00:00:00	2015-10-28 00:00:00
score_away	99	99
score_home	106	106
disadvantaged_team	TOR	IND
committing_team	IND	TOR
disadvantged_team	TOR	IND
season	2015-2016	2015-2016

We use scikit-learn’s LabelEncoder to transform categorical features (call type, player, and season) to integers.

call_type_enc = LabelEncoder().fit(
    clean_df['call_type']
)
n_call_type = call_type_enc.classes_.size

player_enc = LabelEncoder().fit(
    np.concatenate((
        clean_df['committing_player'],
        clean_df['disadvantaged_player']
    ))
)
n_player = player_enc.classes_.size

season_enc = LabelEncoder().fit(
    clean_df['season']
)
n_season = season_enc.classes_.size

We transform the data by

rounding seconds_left to the nearest second (purely for convenience),
transforming categorical features to integer ids,
setting foul_called equal to one or zero depending on whether or not a foul was called, and
setting score_committing and score_disadvantaged to the score of the committing and disadvantaged teams, respectively.

df = (clean_df[['seconds_left']]
              .round(0)
              .assign(
                call_type=call_type_enc.transform(clean_df['call_type']),
                foul_called=1. * clean_df['review_decision'].isin(['CC', 'INC']),
                player_committing=player_enc.transform(clean_df['committing_player']),
                player_disadvantaged=player_enc.transform(clean_df['disadvantaged_player']),
                score_committing=clean_df['score_home'].where(
                    clean_df['committing_team'] == clean_df['home'],
                    clean_df['score_away']
                ),
                score_disadvantaged=clean_df['score_home'].where(
                    clean_df['disadvantaged_team'] == clean_df['home'],
                    clean_df['score_away']
                ),
                season=season_enc.transform(clean_df['season'])
              ))

The resulting data is ready for analysis.

df.head(n=2).T

play_id	20151028INDTOR-1	20151028INDTOR-2
seconds_left	89.0	73.0
call_type	4.0	4.0
foul_called	1.0	0.0
player_committing	162.0	36.0
player_disadvantaged	98.0	358.0
score_committing	99.0	106.0
score_disadvantaged	106.0	99.0
season	0.0	0.0

Modeling

We follow George Box’s modeling workflow, as summarized by Dustin Tran:

build a model of the science,
infer the model given data, and
criticize the model given data.

Baseline model

Build a model of the science

Below we examine the foul call rate by season.

def make_foul_rate_yaxis(ax, label="Observed foul call rate"):
    ax.yaxis.set_major_formatter(pct_formatter)
    ax.set_ylabel(label)
    
    return ax

make_foul_rate_yaxis(
    df.pivot_table('foul_called', 'season')
      .rename(index=season_enc.inverse_transform)
      .rename_axis("Season")
      .plot(kind='bar', rot=0, legend=False)
);

There is a pronounced difference between the foul call rate in the 2015-2016 and 2016-2017 NBA seasons; our first model accounts for this difference.

We use pymc3 to specify our models. Our first model is given by

\[ \begin{align*} \beta^{\textrm{season}}_s & \sim N(0, 5) \\ \eta^{\textrm{game}}_k & = \beta^{\textrm{season}}_{s(k)} \\ p_k & = \textrm{sigm}\left(\eta^{\textrm{game}}_k\right). \end{align*} \]

We use a logistic regression model with different factors for each season.

import pymc3 as pm

with pm.Model() as base_model:
    β_season = pm.Normal('β_season', 0., 5., shape=n_season)
    p = pm.Deterministic('p', pm.math.sigmoid(β_season))

Foul calls are Bernoulli trials, $y_k \sim \textrm{Bernoulli}(p_k).$

season = df['season'].values

with base_model:
    y = pm.Bernoulli(
        'y', p[season],
        observed=df['foul_called']
    )

Infer the model given data

We now sample from the model’s posterior distribution.

NJOBS = 3

SAMPLE_KWARGS = {
    'draws': 1000,
    'njobs': NJOBS,
    'random_seed': [
        SEED + i for i in range(NJOBS)
    ],
    'nuts_kwargs': {
        'target_accept': 0.9
    }
}

with base_model:
    base_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [β_season]
100%|██████████| 1500/1500 [00:07<00:00, 198.65it/s]

Convergence diagnostics

We rely on three diagnostics to ensure that our samples have converged to the posterior distribution:

Energy plots: if the two distributions in the energy plot differ significantly (espescially in the tails), the sampling was not very efficient.
Bayesian fraction of missing information (BFMI): BFMI quantifies this difference with a number between zero and one. A BFMI close to (or exceeding) one is preferable, and a BFMI lower than 0.2 is indicative of efficiency issues.
Gelman-Rubin statistics: Gelman-Rubin statistics near one are preferable, and values less than 1.1 are generally taken to indicate convergence.

For more information on energy plots and BFMI consult Robust Statistical Workflow with PyStan.

bfmi = pm.bfmi(base_trace)

max_gr = max(
    np.max(gr_stats) for gr_stats in pm.gelman_rubin(base_trace).values()
)

CONVERGENCE_TITLE = lambda: f"BFMI = {bfmi:.2f}\nGelman-Rubin = {max_gr:.3f}"

(pm.energyplot(base_trace, legend=False, figsize=(6, 4))
   .set_title(CONVERGENCE_TITLE()));

Criticize the model given data

We use the samples from p’s posterior distribution to calculate residuals, which we use to criticize our models. These residuals allow us to assess how well our model describes the data-generation process and to discover unmodeled sources of variation.

base_trace['p']

array([[ 0.4052151 ,  0.30696232],
       [ 0.3937377 ,  0.30995026],
       [ 0.39881138,  0.29866616],
       ..., 
       [ 0.40279887,  0.31166828],
       [ 0.4077945 ,  0.30299785],
       [ 0.40207901,  0.29991789]])

resid_df = (df.assign(p_hat=base_trace['p'][:, df['season']].mean(axis=0))
              .assign(resid=lambda df: df['foul_called'] - df['p_hat']))

resid_df[['foul_called', 'p_hat', 'resid']].head()

	foul_called	p_hat	resid
play_id
20151028INDTOR-1	1.0	0.403875	0.596125
20151028INDTOR-2	0.0	0.403875	-0.403875
20151028INDTOR-3	1.0	0.403875	0.596125
20151028INDTOR-4	0.0	0.403875	-0.403875
20151028INDTOR-6	0.0	0.403875	-0.403875

The per-season residuals are quite small, which is to be expected.

(resid_df.pivot_table('resid', 'season')
         .rename(index=season_enc.inverse_transform))

	resid
season
2015-2016	-0.000162
2016-2017	-0.000219

Anyone who has watched a close basketball game will realize that we have neglected an important factor in late game foul calls — intentional fouls. Near the end of the game, intentional fouls are used by the losing team when they are on defense to end the leading team’s possession as quickly as possible.

The influence of intentional fouls in the plot below is shown by the rapidly increasing of the residuals as the number of seconds left in the game decreases.

def make_time_axes(ax,
                   xlabel="Seconds remaining in game",
                   ylabel="Observed foul call rate"):
    ax.invert_xaxis()
    ax.set_xlabel(xlabel)
    
    return make_foul_rate_yaxis(ax, label=ylabel)

make_time_axes(
    resid_df.pivot_table('resid', 'seconds_left')
            .reset_index()
            .plot('seconds_left', 'resid', kind='scatter'),
    ylabel="Residual"
);

Possession model

Build a model of the science

The following plot illustrates the fact that only the trailing team has any incentive to committ intentional fouls.

df['trailing_committing'] = (df['score_committing']
                               .lt(df['score_disadvantaged'])
                               .mul(1.)
                               .astype(np.int64))

make_time_axes(
    df.pivot_table(
        'foul_called',
        'seconds_left',
        'trailing_committing'
      )
      .rolling(20)
      .mean()
      .rename(columns={
          0: "No", 1: "Yes"
      })
      .rename_axis(
          "Committing team is trailing",
          axis=1
      )
      .plot()
);

Intentional fouls are only useful when the trailing (and committing) team is on defense. The plot below reflects this fact; shooting and personal fouls are almost always called against the defensive player; we see that they are called at a much higher rate than offensive fouls.

ax = (df.pivot_table('foul_called', 'call_type')
        .rename(index=call_type_enc.inverse_transform)
        .rename_axis("Call type", axis=0)
        .plot(kind='barh', legend=False))

ax.xaxis.set_major_formatter(pct_formatter);
ax.set_xlabel("Observed foul call rate");

We continue to model the differnce in foul call rates between seasons.

with pm.Model() as poss_model:
    β_season = pm.Normal('β_season', 0., 5., shape=2)

Throughout this post, we will use hierarchical distributions to model the variation of foul call rates. For much more information on hierarchical models, consult Data Analysis Using Regression and Multilevel/Hierarchical Models. We use the priors

\[ \begin{align*} \sigma_{\textrm{call}} & \sim \operatorname{HalfNormal}(5) \\ \beta^{\textrm{call}}_{c} & \sim \operatorname{Hierarchical-Normal}(0, \sigma_{\textrm{call}}^2). \end{align*} \]

For sampling efficiency, we use an [non-centered parametrization](http://twiecki.github.io/blog/2017/02/08/bayesian-hierchical-non-centered/#The-Funnel-of-Hell-(and-how-to-escape-it%29) of the hierarchical normal distribution.

def hierarchical_normal(name, shape, σ_shape=1):
    Δ = pm.Normal(f'Δ_{name}', 0., 1., shape=shape)
    σ = pm.HalfNormal(f'σ_{name}', 5., shape=σ_shape)
    
    return pm.Deterministic(name, Δ * σ)

Each call type has a different foul call rate.

with poss_model:
    β_call = hierarchical_normal('β_call', n_call_type)

We add score difference and the number of possessions by which the committing team is trailing to the DataFrame.

df['score_diff'] = (df['score_disadvantaged']
                      .sub(df['score_committing']))

df['trailing_poss'] = (df['score_diff']
                         .div(3)
                         .apply(np.ceil))

trailing_poss_enc = LabelEncoder().fit(df['trailing_poss'])
trailing_poss = trailing_poss_enc.transform(df['trailing_poss'])
n_trailing_poss = trailing_poss_enc.classes_.size

The plot below shows that the foul call rate (over time) varies based on the score difference (quantized into possessions) between the disadvanted team and the committing team. We assume that at most three points can be scored in a single possession (while this is not quite correct, four-point plays are rare enough that we do not account for them in our analysis).

make_time_axes(
    df.pivot_table(
        'foul_called',
        'seconds_left',
        'trailing_poss'
      )
      .loc[:, 1:3]
      .rolling(20).mean()
      .rename_axis(
          "Trailing possessions\n(committing team)",
          axis=1
      )
      .plot()
);

The plot below reflects the fact that intentional fouls are disproportionately personal fouls; the rate at which personal fouls are called increases drastically as the game nears its end.

make_time_axes(
    df.pivot_table('foul_called', 'seconds_left', 'call_type')
      .rolling(20).mean()
      .rename(columns=call_type_enc.inverse_transform)
      .rename_axis(None, axis=1)
      .plot()
);

Due to the NBA’s shot clock, the natural timescale of a basketball game is possessions, not seconds, remaining.

df['remaining_poss'] = (df['seconds_left']
                          .floordiv(25)
                          .add(1))

remaining_poss_enc = LabelEncoder().fit(df['remaining_poss'])
remaining_poss = remaining_poss_enc.transform(df['remaining_poss'])
n_remaining_poss = remaining_poss_enc.classes_.size

Below we plot the foul call rate across trailing possession/remaining posession pairs. Note that we always calculate trailing possessions (trailing_poss) from the perspective of the committing team. For instance, trailing_poss = 1 indicates that the committing team is trailing by 1-3 points, whereas trailing_poss = -1 indicates that the committing team is leading by 1-3 points.

ax = sns.heatmap(
    df.pivot_table(
        'foul_called',
        'trailing_poss',
        'remaining_poss'
      )
      .rename_axis(
          "Trailing possessions\n(committing team)",
          axis=0
      )
      .rename_axis("Remaining possessions", axis=1),
    cmap='seismic',
    cbar_kws={'format': pct_formatter}
)

ax.invert_yaxis();
ax.set_title("Observed foul call rate");

The heatmap above shows that the foul call rate increases significantly when the committing team is trailing by more than the number of possessions remaining in the game. That is, teams resort to intentional fouls only when the opposing team can run out the clock and guarantee a win. (Since we have quantized the score difference and time into posessions, this conclusion is not entirely correct; it is, however, correct enough for our purposes.)

call_name_df = df.assign(
    call_type=lambda df: call_type_enc.inverse_transform(
        df['call_type'].values
    )
)

diff_df = (pd.merge(
                call_name_df,
                call_name_df.groupby('call_type')
                            ['foul_called']
                            .mean()
                            .rename('avg_foul_called')
                            .reset_index()
             )
             .assign(diff=lambda df: df['foul_called'] - df['avg_foul_called']))

The heatmaps below are broken out by call type, and show the difference between the foul call rate for each trailing/remaining possession combination and the overall foul call rate for the call type in question

def plot_foul_diff_heatmap(*_, data=None, **kwargs):
    ax = plt.gca()

    sns.heatmap(
        data.pivot_table(
            'diff',
            'trailing_poss',
            'remaining_poss'
        ),
        cmap='seismic', robust=True,
        cbar_kws={'format': pct_formatter}
    )
    
    ax.invert_yaxis()
    ax.set_title("Observed foul call rate")

(sns.FacetGrid(diff_df, col='call_type', col_wrap=3, aspect=1.5)
    .map_dataframe(plot_foul_diff_heatmap)
    .set_axis_labels(
        "Remaining possessions",
        "Trailing possessions\n(committing team)"
    )
    .set_titles("{col_name}"));

These plots confirm that most intentional fouls are personal fouls. They also show that the three-way interaction between trailing possesions, remaining possessions, and call type are important to model foul call rates.

\[ \begin{align*} \sigma_{\textrm{poss}, c} & \sim \operatorname{HalfNormal}(5) \\ \beta^{\textrm{poss}}_{t, r, c} & \sim \operatorname{Hierarchical-Normal}(0, \sigma_{\textrm{poss}, c}^2) \end{align*} \]

with poss_model:
    β_poss = hierarchical_normal(
        'β_poss',
        (n_trailing_poss, n_remaining_poss, n_call_type),
        σ_shape=(1, 1, n_call_type)
    )

The foul call rate is a combination of season, call type, and possession factors.

\[\eta^{\textrm{game}}_k = \beta^{\textrm{season}}_{s(k)} + \beta^{\textrm{call}}_{c(k)} + \beta^{\textrm{poss}}_{t(k),r(k),c(k)}\]

call_type = df['call_type'].values

with poss_model:
    η_game = β_season[season] \
                + β_call[call_type] \
                + β_poss[
                    trailing_poss,
                    remaining_poss,
                    call_type
                ]

\[ \begin{align*} p_k & = \operatorname{sigm}\left(\eta^{\textrm{game}}_k\right) \end{align*} \]

with poss_model:
    p = pm.Deterministic('p', pm.math.sigmoid(η_game))
    y = pm.Bernoulli('y', p, observed=df['foul_called'])

Infer the model given data

Again, we sample from the model’s posterior distribution.

with poss_model:
    poss_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ_β_poss_log__, Δ_β_poss, σ_β_call_log__, Δ_β_call, β_season]
100%|██████████| 1500/1500 [07:56<00:00,  3.15it/s]
There were 5 divergences after tuning. Increase `target_accept` or reparameterize.
There were 10 divergences after tuning. Increase `target_accept` or reparameterize.
There were 9 divergences after tuning. Increase `target_accept` or reparameterize.
The number of effective samples is smaller than 25% for some parameters.

The BFMI and Gelman-Rubin statistics for this model indicate no problems with sampling and good convergence.

bfmi = pm.bfmi(poss_trace)

max_gr = max(
    np.max(gr_stats) for gr_stats in pm.gelman_rubin(poss_trace).values()
)

(pm.energyplot(poss_trace, legend=False, figsize=(6, 4))
   .set_title(CONVERGENCE_TITLE()));

Criticize the model given data

Again, we calculate residuals.

resid_df = (df.assign(p_hat=poss_trace['p'].mean(axis=0))
              .assign(resid=lambda df: df.foul_called - df.p_hat))

The following plots show that, grouped various ways, the residuals for this model are relatively well-distributed.

ax = sns.heatmap(
    resid_df.pivot_table(
                'resid',
                'trailing_poss',
                'remaining_poss'
            )
            .rename_axis(
                "Trailing possessions\n(committing team)",
                axis=0
            )
            .rename_axis(
                "Remaining possessions",
                axis=1
            )
            .loc[-3:3],
    cmap='seismic', 
    cbar_kws={'format': pct_formatter}
)

ax.invert_yaxis();
ax.set_title("Observed foul call rate");

N_BIN = 20

bin_ix, bins = pd.qcut(
    resid_df.p_hat, N_BIN,
    labels=np.arange(N_BIN),
    retbins=True
)

ax = (resid_df.groupby(bins[bin_ix])
              .resid.mean()
              .rename_axis('p_hat', axis=0)
              .reset_index()
              .plot('p_hat', 'resid', kind='scatter'))

ax.xaxis.set_major_formatter(pct_formatter);
ax.set_xlabel(r"Binned $\hat{p}$");

make_foul_rate_yaxis(ax, label="Residual");

ax = (resid_df.groupby('seconds_left')
              .resid.mean()
              .reset_index()
              .plot('seconds_left', 'resid', kind='scatter'))
make_time_axes(ax, ylabel="Residual");

Model selection

Now that we have two models, we can engage in model selection. We use the widely applicable Bayesian information criterion (WAIC) for model selection.

MODEL_NAME_MAP = {
    0: "Base",
    1: "Possession"
}

comp_df = (pm.compare(
                (base_trace, poss_trace),
                (base_model, poss_model)
             )
             .rename(index=MODEL_NAME_MAP)
             .loc[MODEL_NAME_MAP.values()])

Since smaller WAICs are better, the possession model clearly outperforms the base model.

comp_df

	WAIC	pWAIC	dWAIC	weight	SE	dSE	var_warn
Base	11610.1	2.11	1541.98	0	56.9	73.43	0
Possession	10068.1	82.93	0	1	88.05	0	0

fig, ax = plt.subplots()

ax.errorbar(
    np.arange(len(MODEL_NAME_MAP)),
    comp_df.WAIC,
    yerr=comp_df.SE, fmt='o'
);

ax.set_xticks(np.arange(len(MODEL_NAME_MAP)));
ax.set_xticklabels(comp_df.index);
ax.set_xlabel("Model");

ax.set_ylabel("WAIC");

Player item-response theory model

Build a model of the science

We now turn to the question of whether or not committing and/or drawing fouls is a measurable skill. We use an item-response theory (IRT) model to study this question. For more information on Bayesian item-response models, consult the following references.

Practical Issues in Implementing and Understanding Bayesian Ideal Point Estimation is an excellent introduction to applied Bayesian IRT models and has inspired much of this work.
Bayesian Item Response Modeling — Theory and Applications is a comprehensive mathematical overview of Bayesien IRT modeling.

The item-response theory model includes the season, call type, and possession terms of the previous models.

with pm.Model() as irt_model:
    β_season = pm.Normal('β_season', 0., 5., shape=n_season)
    β_call = hierarchical_normal('β_call', n_call_type)
    β_poss = hierarchical_normal(
        'β_poss',
        (n_trailing_poss, n_remaining_poss, n_call_type),
        σ_shape=(1, 1, n_call_type)
    )
    
    η_game = β_season[season] \
                + β_call[call_type] \
                + β_poss[
                    trailing_poss,
                    remaining_poss,
                    call_type
                ]

Each disadvantaged player has an ideal point (per season).

\[ \begin{align*} \sigma_{\theta} & \sim \operatorname{HalfNormal}(5) \\ \theta^{\textrm{player}}_{i, s} & \sim \operatorname{Hierarchical-Normal}(0, \sigma_{\theta}^2) \end{align*} \]

player_disadvantaged = df['player_disadvantaged'].values
n_player = player_enc.classes_.size

with irt_model:
    θ_player = hierarchical_normal(
        'θ_player', (n_player, n_season)
    )
    θ = θ_player[player_disadvantaged, season]

Each committing player has an ideal point (per season).

\[ \begin{align*} \sigma_{b} & \sim \operatorname{HalfNormal}(5) \\ b^{\textrm{player}}_{j, s} & \sim \operatorname{Hierarchical-Normal}(0, \sigma_{b}^2) \end{align*} \]

player_committing = df['player_committing'].values

with irt_model:
    b_player = hierarchical_normal(
        'b_player', (n_player, n_season)
    )
    b = b_player[player_committing, season]

Players affect the foul call rate through the difference in their ideal points.

\[\eta^{\textrm{player}}_k = \theta_k - b_k\]

with irt_model:
    η_player = θ - b

The sum of the game and player effects determines the foul call probability.

\[\eta_k = \eta^{\textrm{game}}_k + \eta^{\textrm{player}}_k\]

with irt_model:
    η = η_game + η_player

with irt_model:
    p = pm.Deterministic('p', pm.math.sigmoid(η))
    y = pm.Bernoulli(
        'y', p,
        observed=df['foul_called']
    )

Infer the model given data

Again, we sample from the model’s posterior distribution.

with irt_model:
    irt_trace = pm.sample(**SAMPLE_KWARGS)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [σ_b_player_log__, Δ_b_player, σ_θ_player_log__, Δ_θ_player, σ_β_poss_log__, Δ_β_poss, σ_β_call_log__, Δ_β_call, β_season]
100%|██████████| 1500/1500 [13:55<00:00,  1.80it/s]
There were 3 divergences after tuning. Increase `target_accept` or reparameterize.
There were 1 divergences after tuning. Increase `target_accept` or reparameterize.
There were 4 divergences after tuning. Increase `target_accept` or reparameterize.
The estimated number of effective samples is smaller than 200 for some parameters.

None of the sampling diagnostics indicate problems with convergence.

bfmi = pm.bfmi(irt_trace)

max_gr = max(
    np.max(gr_stats) for gr_stats in pm.gelman_rubin(irt_trace).values()
)

(pm.energyplot(irt_trace, legend=False, figsize=(6, 4))
   .set_title(CONVERGENCE_TITLE()));

Criticize the model given data, take three

The binned residuals for this model are more asymmetric than for the previous models, but still not too bad.

resid_df = (df.assign(p_hat=irt_trace['p'].mean(axis=0))
              .assign(resid=lambda df: df['foul_called'] - df['p_hat']))

N_BIN = 50

bin_ix, bins = pd.qcut(
    resid_df.p_hat, N_BIN,
    labels=np.arange(N_BIN),
    retbins=True
)

ax = (resid_df.groupby(bins[bin_ix])
              .resid.mean()
              .rename_axis('p_hat', axis=0)
              .reset_index()
              .plot('p_hat', 'resid', kind='scatter'))

ax.xaxis.set_major_formatter(pct_formatter);
ax.set_xlabel(r"Binned $\hat{p}$");

make_foul_rate_yaxis(ax, label="Residual");

ax = (resid_df.groupby('seconds_left')
              .resid.mean()
              .reset_index()
              .plot('seconds_left', 'resid', kind='scatter'))
make_time_axes(ax, ylabel="Residual");

Model selection

The IRT model is a marginal improvement over the possession model in terms of WAIC.

MODEL_NAME_MAP[2] = "IRT"

comp_df = (pm.compare(
                (base_trace, poss_trace, irt_trace),
                (base_model, poss_model, irt_model)
             )
             .rename(index=MODEL_NAME_MAP)
             .loc[MODEL_NAME_MAP.values()])

comp_df

	WAIC	pWAIC	dWAIC	weight	SE	dSE
Base	11610.1	2.11	1566.92	0	56.9	74.03
Possession	10068.1	82.93	24.94	0.08	88.05	10.99
IRT	10043.2	216.6	0	0.91	88.47	0

fig, ax = plt.subplots()
ax.errorbar(
    np.arange(len(MODEL_NAME_MAP)), comp_df.WAIC,
    yerr=comp_df.SE, fmt='o'
);
ax.set_xticks(np.arange(len(MODEL_NAME_MAP)));
ax.set_xticklabels(comp_df.index);
ax.set_xlabel("Model");
ax.set_ylabel("WAIC");

We now produce two DataFrames containing the estimated player ideal points per season.

def varname_to_param(varname):
    return varname[0]

def varname_to_player(varname):
    return int(varname[3:-2])

def varname_to_season(varname):
    return int(varname[-1])

irt_df = (pm.trace_to_dataframe(
                irt_trace, varnames=['θ_player', 'b_player']
            )
            .rename(columns=lambda col: col.replace('_player', ''))
            .T
            .apply(
               lambda s: pd.Series.describe(
                   s, percentiles=[0.055, 0.945]
               ),
               axis=1
            )
            [['mean', '5.5%', '94.5%']]
            .rename(columns={
                 '5.5%': 'low',
                 '94.5%': 'high'
            })
            .rename_axis('varname')
            .reset_index()
            .assign(
                param=lambda df: df['varname'].apply(varname_to_param),
                player=lambda df: df['varname'].apply(varname_to_player),
                season=lambda df: df['varname'].apply(varname_to_season)
            )
            .drop('varname', axis=1))

irt_df.head()

	mean	low	high	param	player	season
0	-0.016516	-0.323127	0.272022	θ	0	0
1	0.003845	-0.289842	0.290248	θ	0	1
2	0.008519	-0.289715	0.304574	θ	1	0
3	0.031367	-0.240853	0.339297	θ	1	1
4	-0.037530	-0.320010	0.226110	θ	2	0

player_irt_df = irt_df.pivot_table(
    index='player',
    columns=['param', 'season'],
    values='mean'
)

player_irt_df.head()

param	b		θ
season	0	1	0	1
player
0	-0.069235	0.013339	-0.016516	0.003845
1	-0.003003	0.001853	0.008519	0.031367
2	0.084515	0.089058	-0.037530	-0.002373
3	-0.028946	0.004360	0.003514	-0.000334
4	-0.001976	0.280380	0.072932	0.005571

The following plot shows that the committing skill appears to be somewhat larger than the disadvantaged skill. This difference seems reasonable because most fouls are committed by the player on defense; committing skill is quite likely to be correlated with defensive ability.

def plot_latent_params(df):
    fig, ax = plt.subplots()
    
    n, _ = df.shape
    y = np.arange(n)

    ax.errorbar(
        df['mean'], y,
        xerr=(df[['high', 'low']]
                .sub(df['mean'], axis=0)
                .abs()
                .values.T),
        fmt='o'
    )

    ax.set_yticks(y)
    ax.set_yticklabels(
        player_enc.inverse_transform(df.player)
    )
    ax.set_ylabel("Player")
    
    return fig, ax

fig, axes = plt.subplots(
    ncols=2, nrows=2, sharex=True,
    figsize=(16, 8)
)
(θ0_ax, θ1_ax), (b0_ax, b1_ax) = axes

bins = np.linspace(
    0.9 * irt_df['mean'].min(),
    1.1 * irt_df['mean'].max(),
    75
)

θ0_ax.hist(
    player_irt_df['θ', 0],
    bins=bins, normed=True
);
θ1_ax.hist(
    player_irt_df['θ', 1],
    bins=bins, normed=True
);

θ0_ax.set_yticks([]);
θ0_ax.set_title(
    r"$\hat{\theta}$ (" + season_enc.inverse_transform(0) + ")"
);

θ1_ax.set_yticks([]);
θ1_ax.set_title(
    r"$\hat{\theta}$ (" + season_enc.inverse_transform(1) + ")"
);

b0_ax.hist(
    player_irt_df['b', 0],
    bins=bins, normed=True, color=green
);
b1_ax.hist(
    player_irt_df['b', 1],
    bins=bins, normed=True, color=green
);

b0_ax.set_xlabel(
    r"$\hat{b}$ (" + season_enc.inverse_transform(0) + ")"
);

b0_ax.invert_yaxis();
b0_ax.xaxis.tick_top();
b0_ax.set_yticks([]);

b1_ax.set_xlabel(
    r"$\hat{b}$ (" + season_enc.inverse_transform(1) + ")"
);

b1_ax.invert_yaxis();
b1_ax.xaxis.tick_top();
b1_ax.set_yticks([]);

fig.suptitle("Disadvantaged skill", size=18);
fig.text(0.45, 0.02, "Committing skill", size=18)
fig.tight_layout();

We now examine the top and bottom ten players in each ability, across both seasons.

The top players in terms of disadvantaged ability tend to be good scorers (Jimmy Butler, Ricky Rubio, John Wall, Andre Iguodala). The presence of DeAndre Jordan in the top ten may to be due to the hack-a-Shaq phenomenon. In future work, it would be interesting to control for the disavantage player’s free throw percentage in order to mitigate the influence of the hack-a-Shaq effect on the measurement of latent skill.

Interestingly, the bottom players (in terms of disadvantaged ability) include many stars (Pau Gasol, Carmelo Anthony, Kevin Durant, Kawhi Leonard). The presence of these stars in the bottom may somewhat counteract the pervasive narrative that referees favor stars in their foul calls.

top_bot_irt_df = (irt_df.groupby('param')
                    .apply(
                        lambda df: pd.concat((
                                df.nlargest(10, 'mean'),
                                df.nsmallest(10, 'mean')
                            ),
                            axis=0, ignore_index=True
                        )
                    )
                    .reset_index(drop=True))

top_bot_irt_df.head()

	mean	low	high	param	player	season
0	0.351946	-0.026786	0.762273	b	86	0
1	0.320737	-0.027064	0.713128	b	23	0
2	0.280380	-0.071020	0.695970	b	4	1
3	0.279678	-0.057249	0.647667	b	462	1
4	0.271735	-0.106795	0.676231	b	78	0

fig, ax = plot_latent_params(
    top_bot_irt_df[top_bot_irt_df['param'] == 'θ']
                  .sort_values('mean')
)
ax.set_xlabel(r"$\hat{\theta}$");
ax.set_title("Top and bottom ten");

The top ten players in terms of committing skill include many defensive standouts (Danny Green — twice, Gordon Hayward, Paul George).

The bottom ten players include many that are known to be defensively challenged (Ricky Rubio and James Harden). Dwight Howard was, at one point, a fierce defender of the rim, but was well past his prime in 2015, when our data set begins. Chris Paul’s presence in the bottom is somewhat surprising.

fig, ax = plot_latent_params(
    top_bot_irt_df[top_bot_irt_df['param'] == 'b']
                  .sort_values('mean')
)
ax.set_xlabel(r"$\hat{b}$");
ax.set_title("Top and bottom ten");

In the sports analytics community, year-over-year correlation of latent parameters is the test of whether or not a latent quantity truly measures a skill. The following plots show a slight year-over-year correlation in the committing skill, but not much correlation in the disadvantaged skill.

def p_val_to_asterisks(p_val):
    if p_val < 0.0001:
        return "****"
    elif p_val < 0.001:
        return "***"
    elif p_val < 0.01:
        return "**"
    elif p_val < 0.05:
        return "*"
    else:
        return ""

def plot_corr(x, y, **kwargs):
    corrcoeff, p_val = sp.stats.pearsonr(x, y)
    asterisks = p_val_to_asterisks(p_val)
    
    artist = AnchoredText(
        f'{corrcoeff:.2f}{asterisks}',
        loc=10, frameon=False,
        prop=dict(size=LABELSIZE)
    )
    plt.gca().add_artist(artist)
    plt.grid(b=False)

PARAM_MAP = {
    'θ': r"$\hat{\theta}$",
    'b': r"$\hat{b}$"
}

def replace_label(label):
    param, season = eval(label)

    return "{param}\n({season})".format(
        param=PARAM_MAP[param],
        season=season_enc.inverse_transform(season)
    )
    
def style_grid(grid):
    for ax in grid.axes.flat:
        ax.grid(False)
        ax.set_xticklabels([]);
        ax.set_yticklabels([]);
        
        if ax.get_xlabel():
            ax.set_xlabel(replace_label(ax.get_xlabel()))

        if ax.get_ylabel():
            ax.set_ylabel(replace_label(ax.get_ylabel()))
            
    return grid

player_all_season = set(df.groupby('player_disadvantaged')
                          .filter(lambda df: df['season'].nunique() == n_season)
                          ['player_committing']) \
                        & set(df.groupby('player_committing')
                                .filter(lambda df: df['season'].nunique() == n_season)
                                ['player_committing'])

style_grid(
    sns.PairGrid(
        player_irt_df.loc[player_all_season],
        size=1.75
       )
       .map_upper(plt.scatter, alpha=0.5)
       .map_diag(plt.hist)
       .map_lower(plot_corr)
);

Since we can only reasonably estimate the skills of players for which we have sufficient foul call data, we plot the correlations below for players that appeared in at least ten plays in each season.

MIN = 10

player_has_min = set(df.groupby('player_disadvantaged')
                       .filter(lambda df: (df['season']
                                             .value_counts()
                                             .gt(MIN)
                                             .all()))
                       ['player_committing']) \
                    & set(df.groupby('player_committing')
                            .filter(lambda df: (df['season']
                                                  .value_counts()
                                                  .gt(MIN)
                                                  .all()))
                            ['player_committing'])

grid = style_grid(
    sns.PairGrid(
        player_irt_df.loc[player_has_min],
        size=1.75
       )
       .map_upper(plt.scatter, alpha=0.5)
       .map_diag(plt.hist)
       .map_lower(plot_corr)
)

As expected, the season-over-season latent skill correlations are higher for this subset of players.

From this figure, it seems that committing skill ($\hat{b}$) exists and is measurable, but is fairly small. It also seems that disadvantaged skill ($\hat{\theta}$), if it exists, is difficult to measure from this data set. Ideally, the NBA would release a foul report for the entirety of every game, but that seems quite unlikely.

In the future it would be useful to include a correction for the probability that a given game appears in the data set. This correction would help with the sample bias introduced by the fact that only games that are close in the last two minutes are included in the NBA’s reports.

This post is available as a Jupyter notebook here.