Simulation of Vaccine Efficacy

vaccine.py

import random

VACCINE_EFFICACY = 0.504
VACCINE_WAIT = 14
SAMPLE_SIZE = 20_000
TOTAL_POPULATION = 10_000_000
INFECTION_RATE = 0.0002

class SimulateVaccine(object):
    def __init__(self):
        self.populations = list(range(TOTAL_POPULATION))
        self.sample_group()
        self.vaccinated_infected = set()
        self.placebo_infected = set()

    def sample_group(self):
        allsamples = random.sample(self.populations, SAMPLE_SIZE)
        self.vaccinated = set(allsamples[:SAMPLE_SIZE//2])
        self.placebo = set(allsamples[SAMPLE_SIZE//2:])

    def is_infected(self, rate):
        if random.random() < rate:
            return True
        return False

    def infect(self, day):
        efficacy = VACCINE_EFFICACY * day / VACCINE_WAIT if day < VACCINE_WAIT else VACCINE_EFFICACY
        for person_id in self.populations:
            if self.is_infected(INFECTION_RATE):
                if person_id in self.vaccinated:
                    if self.is_infected(1 - efficacy):
                        self.vaccinated_infected.add(person_id)
                if person_id in self.placebo:
                    self.placebo_infected.add(person_id)

    def sample_infected(self):
        return len(self.vaccinated_infected) + len(self.placebo_infected)

    def run(self, target):
        day = 0
        vaccinated_prior = 0
        placebo_prior = 0
        while self.sample_infected() < target:
            self.infect(day)
            if day == VACCINE_WAIT:
                vaccine_prior = len(self.vaccinated_infected)
                placebo_prior = len(self.placebo_infected)
            print(day, len(self.vaccinated_infected), len(self.placebo_infected))
            day += 1
        attack_vaccinated = (len(self.vaccinated_infected) - vaccine_prior) / (SAMPLE_SIZE//2)
        attack_placebo = (len(self.placebo_infected) - placebo_prior) / (SAMPLE_SIZE//2)
        efficacy = (attack_placebo - attack_vaccinated) * 100 / attack_placebo
        print("Efficacy = ", efficacy)

if __name__ == "__main__":
    sim = SimulateVaccine()
    sim.run(150)