cluster_covid_infection_estimates.py

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

households = 10000
infected_households_prop = 0.01
persons_per_household = 100
sample_prop = 0.001
samples = int(sample_prop * households * persons_per_household)

experiments = 1000
infected_prop_diffs = []
for experiment in range(experiments):
    infected_households = np.random.rand(households) < infected_households_prop
    infected_persons = np.repeat(infected_households, persons_per_household)
    true_infected_prop = infected_persons.sum() / len(infected_persons)

    sampled = np.random.choice(infected_persons, samples, replace=False)
    sampled_infected_prop = sampled.sum() / samples

    infected_prop_diffs.append(sampled_infected_prop - true_infected_prop)

print(np.mean(infected_prop_diffs))
sns.histplot(infected_prop_diffs)
plt.show()

infected_prop_diffs = []
for experiment in range(experiments):
    infected_households = np.random.rand(households) < infected_households_prop
    true_infected_prop = infected_households.sum() / len(infected_households)

    sampled = np.random.choice(infected_households, samples, replace=False)
    sampled_infected_prop = sampled.sum() / samples

    infected_prop_diffs.append(sampled_infected_prop - true_infected_prop)

print(np.mean(infected_prop_diffs))
sns.histplot(infected_prop_diffs)
plt.show()