Tutorial 6: Multinomial Inference with Dirichlet Priors

"""Tutorial 6: Use Dirichlet inference to test if specific balls deviate from uniformity"""

import pandas as pd
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

# --- Load cleaned data ---

print("Loading cleaned Powerball data from Tutorial 1...")
draws = pd.read_parquet('data/processed/powerball_clean.parquet')
print(f"Loaded {len(draws)} draws\n")

# --- Extract ball frequencies ---

# Collect all white balls from all draws
all_balls = []
for _, row in draws.iterrows():
    all_balls.extend([row['ball1'], row['ball2'], row['ball3'], row['ball4'], row['ball5']])

all_balls = np.array(all_balls)
n_total_picks = len(all_balls)  # Should be 5 * n_draws

print(f"Total ball picks: {n_total_picks}")
print(f"  (5 balls per draw × {len(draws)} draws)\n")

# Count frequency of each ball (1-69)
ball_counts = np.zeros(69)
for ball_num in range(1, 70):
    ball_counts[ball_num - 1] = np.sum(all_balls == ball_num)

# Expected frequency if uniform
expected_frequency = n_total_picks / 69

print(f"Expected frequency per ball (if uniform): {expected_frequency:.2f}")
print(f"\nObserved frequencies (first 10 balls):")
for i in range(10):
    print(f"  Ball {i+1}: {ball_counts[i]:.0f} (expected: {expected_frequency:.2f})")
print("  ...\n")

# --- Define prior ---

# Uniform Dirichlet prior: all alpha = 1
# This is the 69-dimensional generalization of Beta(1,1)
alpha_prior = np.ones(69)

print(f"Prior: Dirichlet(α₁=1, α₂=1, ..., α₆₉=1)")
print(f"  This is a flat prior over all 69 balls")
print(f"  Sum of alphas: {np.sum(alpha_prior):.0f} (weak prior, easy to update)\n")

# --- Bayesian update ---

# The Dirichlet-Multinomial conjugate update:
# Posterior = Dirichlet(alpha_prior + counts)
alpha_posterior = alpha_prior + ball_counts

print(f"Posterior: Dirichlet(α₁={alpha_posterior[0]:.0f}, α₂={alpha_posterior[1]:.0f}, ..., α₆₉={alpha_posterior[68]:.0f})")
print(f"  Updated belief after observing {len(draws)} draws")
print(f"  Sum of alphas: {np.sum(alpha_posterior):.0f} (strong posterior, hard to change)\n")

# --- Compute posterior statistics for each ball ---

# Posterior mean for each ball
posterior_means = alpha_posterior / np.sum(alpha_posterior)

# Expected probability if uniform
expected_prob = 1.0 / 69

# Compute credible intervals for each ball
# For Dirichlet, we use the fact that each marginal is approximately normal for large counts
# This is justified by the Central Limit Theorem
lower_bounds = []
upper_bounds = []

for i in range(69):
    # Use normal approximation for the marginal
    # Mean: alpha_i / sum(alpha)
    # Variance: (alpha_i * (sum(alpha) - alpha_i)) / (sum(alpha)^2 * (sum(alpha) + 1))
    alpha_sum = np.sum(alpha_posterior)
    mean = alpha_posterior[i] / alpha_sum
    variance = (alpha_posterior[i] * (alpha_sum - alpha_posterior[i])) / (alpha_sum**2 * (alpha_sum + 1))
    std = np.sqrt(variance)
    
    # 95% credible interval using ±1.96 standard deviations
    lower = mean - 1.96 * std
    upper = mean + 1.96 * std
    
    lower_bounds.append(lower)
    upper_bounds.append(upper)

lower_bounds = np.array(lower_bounds)
upper_bounds = np.array(upper_bounds)

# --- Identify suspicious balls ---

# A ball is "suspicious" if the expected uniform probability (1/69)
# falls outside its 95% credible interval
suspicious_balls = []

for i in range(69):
    if expected_prob < lower_bounds[i] or expected_prob > upper_bounds[i]:
        suspicious_balls.append(i + 1)  # Ball numbers are 1-indexed

print(f"Suspicious balls (expected value outside 95% CI):")
if len(suspicious_balls) == 0:
    print("  None! All balls are consistent with uniformity.")
else:
    print(f"  Found {len(suspicious_balls)} suspicious balls:")
    for ball_num in suspicious_balls[:10]:  # Show first 10
        idx = ball_num - 1
        print(f"    Ball {ball_num}: observed {ball_counts[idx]:.0f}, " 
              f"expected {expected_frequency:.2f}, "
              f"95% CI: [{lower_bounds[idx]*n_total_picks:.1f}, {upper_bounds[idx]*n_total_picks:.1f}]")
    if len(suspicious_balls) > 10:
        print(f"    ... and {len(suspicious_balls) - 10} more")

print()

# --- Visualize posterior as heatmap ---

print("Creating heatmap visualization...")

# Reshape ball numbers into a 7x10 grid for visualization
# (69 balls don't fit perfectly, so last row will have 9 balls)
grid_data = np.zeros((7, 10))
grid_data[:] = np.nan  # Initialize with NaN

ball_idx = 0
for row in range(7):
    for col in range(10):
        if ball_idx < 69:
            grid_data[row, col] = posterior_means[ball_idx]
            ball_idx += 1

# Create heatmap
fig, ax = plt.subplots(figsize=(14, 8))

# Plot heatmap
# vmin=0.013 and vmax=0.016 are chosen to "zoom in" around the expected value of 1/69 ≈ 0.01449
# This makes small deviations from uniformity more visually apparent in the color scale
im = ax.imshow(grid_data, cmap='RdYlGn_r', aspect='auto', vmin=0.013, vmax=0.016)

# Add colorbar
cbar = plt.colorbar(im, ax=ax)
cbar.set_label('Posterior Probability', fontsize=12)

# Mark expected uniform probability
expected_line = ax.axhline(y=-0.5, color='blue', linestyle='--', linewidth=2, alpha=0)  # Invisible line for legend
ax.plot([], [], color='blue', linestyle='--', linewidth=2, label=f'Expected (1/69 = {expected_prob:.4f})')

# Annotate each cell with ball number and mark suspicious ones
ball_idx = 0
for row in range(7):
    for col in range(10):
        if ball_idx < 69:
            ball_num = ball_idx + 1
            
            # Check if suspicious
            if ball_num in suspicious_balls:
                color = 'red'
                weight = 'bold'
                marker = '*'
            else:
                color = 'black'
                weight = 'normal'
                marker = ''
            
            ax.text(col, row, f"{ball_num}{marker}", 
                   ha='center', va='center', fontsize=10, 
                   color=color, fontweight=weight)
            ball_idx += 1

ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Posterior Probabilities for Each Ball (Heatmap)\n' + 
             '(Red/starred balls have 95% CI excluding expected uniform value)',
             fontsize=14, fontweight='bold')
ax.legend(loc='upper right', fontsize=10)

plt.tight_layout()
plt.savefig('outputs/tutorial6/ball_posterior_heatmap.png', dpi=150)
plt.close()

print("  Saved: ball_posterior_heatmap.png\n")

# --- Summary ---

print("Summary")
print("-" * 40)
print(f"We tested all 69 balls for deviations from uniformity.")
print(f"Using a flat Dirichlet prior and {len(draws)} draws:")
print(f"  Expected frequency per ball: {expected_frequency:.2f}")
print(f"  Suspicious balls: {len(suspicious_balls)}")
print()
if len(suspicious_balls) == 0:
    print("Conclusion: All balls are consistent with uniform randomness.")
else:
    print("Conclusion: Some balls deviate from expected frequency.")
    print("However, with 69 balls and 95% CIs, we expect ~3-4 false positives.")
    print("These deviations may just be noise, not evidence of bias.\n")