Untitled

import numpy as np
import matplotlib.pyplot as plt


class PandemicSimulation:
    def __init__(self, population_size, initial_infected_count, transmission_rate, recovery_chance, random_state=None):
        """
        Initialize the simulation parameters.
        
        Parameters
        ----------
        population_size : int
            Total number of individuals in the population.
        initial_infected_count : int
            Number of individuals initially infected at the start of the simulation.
        transmission_rate : float (0.0 to 1.0)
            Probability that an infection is transmitted from an infected to a healthy individual during a meeting.
        recovery_chance : float (0.0 to 1.0)
            Probability that an infected individual recovers at the end of a day.
        random_state : int or None
            Optional seed for reproducibility.
        """
        self.population_size = population_size
        self.initial_infected_count = initial_infected_count
        self.transmission_rate = transmission_rate
        self.recovery_chance = recovery_chance
        
        self.rng = np.random.default_rng(random_state)
        
        # Initialize population: True = infected, False = healthy
        self.population = np.zeros(population_size, dtype=bool)
        initially_infected_indices = self.rng.choice(population_size, initial_infected_count, replace=False)
        self.population[initially_infected_indices] = True
        
        self.days_elapsed = 0

    def simulate_day(self):
        """
        Simulate one day of meetings, infection spread, and recovery.
        """
        # Randomly permute the population indices
        indices = self.rng.permutation(self.population_size)
        
        # Pair individuals: (indices[0], indices[1]), (indices[2], indices[3]), ...
        # If odd number of individuals, the last one doesn't meet anyone.
        pairs = indices.reshape(-1, 2) if self.population_size % 2 == 0 else indices[:-1].reshape(-1, 2)
        
        # Meeting step: transmit infection
        for p1, p2 in pairs:
            infected1 = self.population[p1]
            infected2 = self.population[p2]
            
            # If one infected and one healthy, infection might spread
            if infected1 and not infected2:
                # p2 gets infected with probability transmission_rate
                if self.rng.random() < self.transmission_rate:
                    self.population[p2] = True
            elif infected2 and not infected1:
                # p1 gets infected with probability transmission_rate
                if self.rng.random() < self.transmission_rate:
                    self.population[p1] = True
        
        # Recovery step: each infected recovers with probability recovery_chance
        infected_indices = np.where(self.population == True)[0]
        recoveries = self.rng.random(len(infected_indices)) < self.recovery_chance
        # Set recovered individuals to healthy
        self.population[infected_indices[recoveries]] = False
        
        # Increment day counter
        self.days_elapsed += 1

    def run_until_cured(self, max_days=10_000):
        """
        Run the simulation until no one is infected or until max_days is reached.
        
        Returns
        -------
        int
            The number of days until the population is fully cured.
        """
        while np.any(self.population) and self.days_elapsed < max_days:
            self.simulate_day()
        return self.days_elapsed



# Set parameters for the simulation
population_size = 10_000
initial_infected_count = 10
transmission_rate = 0.05
recovery_chance = 0.10

# Run a single simulation
sim = PandemicSimulation(population_size, initial_infected_count, transmission_rate, recovery_chance, random_state=42)
days = sim.run_until_cured()
print(f"It took {days} days until the population was fully cured.")

# Run multiple simulations to get an average
n_simulations = 10_000
results = []
for i in range(n_simulations):
    sim = PandemicSimulation(population_size, initial_infected_count, transmission_rate, recovery_chance)
    days = sim.run_until_cured()
    results.append(days)

avg_days = np.mean(results)
print(f"Average days until cured over {n_simulations} simulations: {avg_days:.2f}")

plt.figure(figsize=(10,6))
plt.hist(results, bins=n_simulations, edgecolor='black', alpha=0.7)
plt.title('Distribution of Days Until Cured Over Multiple Simulations')
plt.xlabel('Days Until Cured')
plt.ylabel('Number of Simulations')
plt.grid(True, alpha=0.3)
plt.show()