Python code for reproduce figure 8.3, chapter 8, book: Intro to Probability for Data Science

python_figure_chapter8_83.txt

# python code for reproduce figure 8.3, chapter 8, book: Intro to Probability for Data Science
# https://probability4datascience.com/python08.html
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Bernoulli log-likelihood sequence
N = 50
S = np.linspace(0, N)
theta = np.linspace(0.1, 0.9, 100)
S_grid, theta_grid = np.meshgrid(S, theta)

def log_likelihood_bernoulli(S, theta):
  return S * np.log(theta) + (S.shape[1] - S) * np.log(1 - theta)


fig = plt.figure(figsize=(16, 6))
grid = plt.GridSpec(2, 4, wspace=0.4, hspace=0.3)

# surface plot
surface_ax = fig.add_subplot(grid[0:, 0], projection='3d')
surface_ax.plot_surface(S_grid, theta_grid, log_likelihood_bernoulli(S_grid, theta_grid), rstride=1, cstride=1,
                cmap='jet', edgecolor='none')
surface_ax.set_xlabel('S')
surface_ax.set_ylabel(r'$\theta$')
surface_ax.set_title(r'$log \mathcal{L}(\theta|S)$')
surface_ax.view_init(20, 690)

# contour plot
contour_ax = fig.add_subplot(grid[0:, 1:3])
cmap = contour_ax.imshow(log_likelihood_bernoulli(S_grid, theta_grid), extent=[0, 50, 0.1, 0.9], 
               cmap='jet', aspect = 'auto', alpha=1, zorder=1)
contour_ax.set_xlabel(r'S')
contour_ax.set_ylabel(r'$\theta$')
contour_ax.axvline(x=12, ymin=0, ymax=1, color='black', linestyle='dotted', linewidth=1.75)
contour_ax.axvline(x=25, ymin=0, ymax=1, color='black', linestyle='dotted', linewidth=1.75)
divider = make_axes_locatable(contour_ax)
cax = divider.append_axes('bottom', size='1.5%', pad=0.55)
fig.colorbar(cmap ,cax=cax, orientation='horizontal', label= 'log-likelihoood')

# plot slice at S=25
slice_s25 = fig.add_subplot(grid[0, 3])
slice_s25.plot(theta_grid, log_likelihood_bernoulli(np.ones(S_grid.shape) * 25, theta_grid), '-k')
slice_s25.annotate(r'$log \mathcal{L}(\theta | S=25)$', xy=(0.3, -55), size=14)
slice_s25.grid(True)

# plot slice at S=12
slice_s12 = fig.add_subplot(grid[1, 3])
slice_s12.plot(theta_grid, log_likelihood_bernoulli(np.ones(S_grid.shape) * 12, theta_grid), '-k')
slice_s12.annotate(r'$log \mathcal{L}(\theta | S=12)$', xy=(0.2, -80), size=14)
slice_s12.grid(True)
slice_s12.set_xlabel(r'$\theta$')
fig.suptitle('Bernoulli log-likelihood sequence:  ' + r'$log \mathcal{L}(\theta | x)=S log(\theta) + (N-S) log(1 - \theta)$' + ' , with ' + r'$S=\sum_{n=1}^{N}x_i$')