The gist of my program is to simulate the growth of some virus. I used (attempted) OOP concepts to break down the problem into chunks and have the chunks talk to each other.
Not sure if my implementation is effective but it seems to work pretty good. Currently, the bottleneck seems to be in the plotting. I'm still learning about matplotlib so I'm not surprised.
The program has five classes. The first class just keeps track of simulation details, nothing too fancy.
class Details(): def __init__(self,num_people_1d,fig_size): self.num_people_1d = num_people_1d self.total_people = self.num_people_1d**2 self.fig_size = fig_size self.x_length = self.fig_size/num_people_1d self.y_length = self.fig_size/num_people_1d
The second class is the display. At each iteration, the 2D grid is updated with details about individuals that are either infected or have perished. The display class is updated with the information
import matplotlib.pyplot as plt import matplotlib.patches as patches class Display(): def __init__(self,details_instance): self.size_x,self.size_y = details_instance.fig_size self.length_1d = details_instance.num_people_1d self.x_length = details_instance.x_length self.y_length = details_instance.y_length def create_plot(self,plot_size = (5,5)): self.fig = plt.figure(figsize = plot_size) self.ax = self.fig.subplots() canvas = patches.Rectangle((0,0),1,1,fill=True, edgecolor='none',facecolor='g') self.ax.add_patch(canvas) def update_plot(self,infected_table=None,kill_table = None): #Transposing tables infected_table = list(map(list, zip(*infected_table))) kill_table = list(map(list, zip(*kill_table))) for i,row in enumerate(infected_table): for j,col in enumerate(row): infected_person = col dead_person = kill_table[i][j] if dead_person: coord = i*self.x_length,j*self.y_length square = patches.Rectangle(coord, self.x_length,self.y_length, fill=True, edgecolor = 'none', facecolor = 'r') self.ax.add_patch(square) if infected_person and not dead_person: coord = i*self.x_length,j*self.y_length square = patches.Rectangle(coord, self.x_length,self.y_length, fill=True, edgecolor = 'none', facecolor = 'y') self.ax.add_patch(square) plt.show() plt.pause(0.1)
The next class is the virus class. Not too much going on, just the infect and mortality rate.
class Virus(): def __init__(self,infectionRate = 0.1,mortalityRate = 0.01): self.IR = infectionRate self.MR = mortalityRate
Then is the person class. This class just keeps some basic information. If the person is infected or dead, and some simple methods.
import random class Person(): def __init__(self,id = None,discrete_location = None,infected = False): self.id = id if discrete_location: self.dl_x,self.dl_y = discrete_location else: raise Exception() self.infected = infected self.neighbors =  self.dead = False def become_infected(self,virus): if not self.dead: self.infected = True self.virus = virus def do_i_live(self): return random.random()>self.virus.MR def kill(self): self.dead = True self.infected = False
The last class is the Population class. This class holds most of the code that actually does stuff because it is what updates all of the individuals.
from person import Person import random class Population(): def __init__(self,persons=,details_instance =None,virus_strain=None): if len(persons)<1: print('There is no population! Adding a member') self.persons = persons self.count = 0 self.add_person() else: self.persons = persons self.details_instance = details_instance self.virus_strain = virus_strain self.dead_persons = [] def add_person(self): if len(self.persons)<1: self.persons.append(Person(id=self.count, discrete_location = (0,0), infected = False) ) self.count +=1 else: loc_x = self.details_instance.x_length*(self.count%self.details_instance.num_people_1d) loc_y = self.details_instance.y_length*((self.count - self.count%self.details_instance.num_people_1d)/self.details_instance.num_people_1d) person = Person(id = self.count, discrete_location = (loc_x,loc_y), infected = False) self.count +=1 self.persons.append(person) def get_infected_table(self): truth_table =  current_list =  i = 0 while i < self.count: current_list.append(self.persons[i].infected) i+=1 if (i)%(self.details_instance.num_people_1d) ==0: truth_table.append(current_list) current_list =  if self.count%(self.details_instance.num_people_1d) !=0: truth_table.append(current_list) return truth_table def get_dead_table(self): truth_table =  current_list =  i = 0 while i < self.count: current_list.append(self.persons[i].dead) i+=1 if (i)%(self.details_instance.num_people_1d) ==0: truth_table.append(current_list) current_list =  if self.count%(self.details_instance.num_people_1d) !=0: truth_table.append(current_list) return truth_table def kill_infected(self,infected_table): linear_indices = self.get_infected_indices(infected_table) for index in linear_indices: still_living = self.persons[index].do_i_live() if not still_living: self.persons[index].kill() self.dead_persons.append(index) def add_neighbors(self): #Currently returns the linear index! Compatible with persons!! if len(self.persons)<=1: return #One method: #Use self.count and modulos to identify neighbors #Possibly a better method that I do not follow: #Using discrete location to identify neighbors #Using first method for i in range(self.count): #at left boundary if i%self.details_instance.num_people_1d==0: left = -1 else: left = i-1 #at right boundary if (i+1)%self.details_instance.num_people_1d==0: right = -1 else: right = i+1 up = i+self.details_instance.num_people_1d down = i - self.details_instance.num_people_1d #First build potential neighbors potential_neighbors = [left,right,up,down] #Second identify if any potential neighbors don't exist neighbor_list =  for j in potential_neighbors: if (j >= 0) and (j<self.count): neighbor_list.append(j) #Third update the person with neighbors self.persons[i].neighbors = neighbor_list def spread_infection(self,infected_table): linear_indices = self.get_infected_indices(infected_table) for index in linear_indices: current_infected_person = self.persons[index] neighbors = current_infected_person.neighbors for neighbor in neighbors: if random.random()<current_infected_person.virus.IR: self.persons[neighbor].become_infected(self.virus_strain) def get_infected_count(self): infected_people = 0 for person in self.persons: if person.infected: infected_people+=1 return infected_people def get_dead_count(self): dead_people = 0 for person in self.persons: if person.dead: dead_people+=1 return dead_people def get_infected_indices(self,infected_table): #returns the linear indices of those infected linear_indices= for i,row in enumerate(infected_table): for j,col in enumerate(row): if col: linear_indices.append(j+i*self.details_instance.num_people_1d) return linear_indices
To run all of this code I wrote the following script:
from person import Person from virus import Virus from display import Display from details import Details from population import Population import random num_people_1d = 10 simul_details = Details(num_people_1d = num_people_1d,fig_size = (1,1)) virus_strain1 = Virus() pop = Population(details_instance = simul_details,virus_strain=virus_strain1) number_people = num_people_1d**2-1 for i in range(number_people): pop.add_person() pop.add_neighbors() starting_person = random.randint(0,number_people-1) print('The starting person is %d' % starting_person) pop.persons[starting_person].become_infected(virus_strain1) current_infected = pop.get_infected_table() current_dead = pop.get_dead_table() simul_display = Display(details_instance=simul_details) simul_display.create_plot() total = 100 for iter in range(total): infected_people = pop.get_infected_count() dead_people = pop.get_dead_count() print('The iteration we are on is %d with %d infected' %(iter,infected_people)) simul_display.update_plot(current_infected,current_dead) pop.spread_infection(current_infected) current_infected=pop.get_infected_table() pop.kill_infected(current_infected) current_dead = pop.get_dead_table() if infected_people+dead_people > number_people: print('All individuals are infected or dead!') break
This is all of the code. Any comments would be gratefully appreciated.