Wald Test

wald_test-shift.py

import numpy as np
import matplotlib.pyplot as plt

def simulate_bernoulli_trials(p, n):
    return np.random.binomial(1, p, n)

def wald_test_statistic(data, p_null):
    p_hat = np.mean(data)
    se = np.sqrt(p_hat * (1 - p_hat) / len(data))
    return (p_hat - p_null) / se

def rejection_probability(lambda_value, n_trials, p_null=0.5, threshold=1.645):
    p_values = [p_null - lambda_value / np.sqrt(n) for n in n_trials]
    print(p_values)
    rejection_rates = []
    wald_stats = []

    for n, p in zip(n_trials, p_values):
        data = simulate_bernoulli_trials(p, n)
        wald_stat = wald_test_statistic(data, p_null)
        rejected = 1 if wald_stat > threshold else 0
        wald_stats.append(wald_stat)
        rejection_rates.append(rejected)

    return rejection_rates, wald_stats


plt.figure(figsize=(10, 6))

n_trials = list(range(100, 10000001, 50000))

# Parameters
lambda_value = 0.01
rejection_rates, wald_stats = rejection_probability(lambda_value, n_trials, p_null=0.5)
print(lambda_value, rejection_rates, wald_stats)
# plt.plot(n_trials, rejection_rates, marker='v', color='orange')
plt.plot(n_trials, wald_stats, marker='v', color='orange', label=r'$\lambda='+str(lambda_value)+'$')

# Parameters
lambda_value = 0.1
rejection_rates, wald_stats = rejection_probability(lambda_value, n_trials, p_null=0.5)
print(lambda_value, rejection_rates, wald_stats)
# plt.plot(n_trials, rejection_rates, marker='o', color='blue')
plt.plot(n_trials, wald_stats, marker='o', color='blue', label=r'$\lambda='+str(lambda_value)+'$')

# Parameters
lambda_value = 1
rejection_rates, wald_stats = rejection_probability(lambda_value, n_trials, p_null=0.5)
print(lambda_value, rejection_rates, wald_stats)
# plt.plot(n_trials, rejection_rates, marker='x', color='red')
plt.plot(n_trials, wald_stats, marker='x', color='red', label=r'$\lambda='+str(lambda_value)+'$')

# Parameters
lambda_value = 2
rejection_rates, wald_stats = rejection_probability(lambda_value, n_trials, p_null=0.5)
print(lambda_value, rejection_rates, wald_stats)
# plt.plot(n_trials, rejection_rates, marker='+', color='green')
plt.plot(n_trials, wald_stats, marker='+', color='green', label=r'$\lambda='+str(lambda_value)+'$')
# Plotting

plt.hlines(1.6445, 0, n_trials[-1], label=r'Rejection $Z>1.6445$')

# plt.plot(n_trials, rejection_rates, marker='o')
plt.xlabel('Sample Size n')
plt.ylabel('Wald Test value')
plt.title('Rejection Probability of $H_0$ vs. Sample Size')
plt.grid(True)
plt.legend()
plt.show()