covid-party.py

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

sns.set_context('talk')
sns.set_style('darkgrid')
ps = [0.001, 0.005, 0.01, 0.02, 0.035]
cities = ['Rural Montana', 'Seattle', 'Minneapolis', 'Chicago', 'Miami']
cp = sns.cubehelix_palette(len(ps), start=0.5, rot=-0.75)
sns.set_palette(cp)
x = np.linspace(2,100, 1000)
dpi = 40
plt.figure(figsize=(506 / dpi, 253 / dpi))
for i,p in enumerate(ps):
    y = 1 - (1 - p)**x
    plt.plot(x, y, label=f'$P_i$={p:.2%} ({cities[i]})')
plt.gca().set_xscale('log')
plt.gca().get_xaxis().set_major_formatter(mpl.ticker.ScalarFormatter())
plt.xticks([2, 5, 10, 20, 50, 100])
plt.title('Are you attending a COVID-19 Party?')
plt.xlabel('Party Attendees')
plt.ylabel('Probability someone has COVID')
plt.legend(bbox_to_anchor=(1.0, 1.05))
sns.despine()
plt.tight_layout()
plt.savefig('covid.png', dpi=300)
plt.show()