add visualisation loop to refactored code

Author: paulvangentcom

config.py

@@ -15,10 +15,11 @@ def __init__(self, *args, **kwargs):
 self.tstep=0#current simulation timestep
 self.save_data=True#whether to dump data
 self.save_timesteps=True#dumps population data every time step
+self.endif_no_infections=True#whether to stop simulation if no infections remain
 
 #scenario flags
 self.traveling_infects=False
-self.self_isolate=True
+self.self_isolate=False
 self.lockdown=False
 self.lockdown_percentage=0.1#after this proportion is infected, lock-down begins
 self.lockdown_compliance=0.95#fraction of the population that will obey the lockdown 
@@ -33,6 +34,7 @@ def __init__(self, *args, **kwargs):
 #size of the simulated world in coordinates
 self.x_plot= [0, 1] 
 self.y_plot= [0, 1]
+self.save_plot=False
 
 #population variables
 self.pop_size=2000

simulation.py

@@ -114,6 +114,11 @@ def tstep(self):
 #update population statistics
 self.pop_tracker.update_counts(self.population)
 
+#visualise
+draw_tstep(self.Config, self.population, self.pop_tracker, self.frame, 
+self.fig, self.spec, self.ax1, self.ax2)
+
+
 #run callback
 self.callback()
 
@@ -133,109 +138,24 @@ def callback(self):
 self.population[0][8] =50
 self.population[0][10] =1
 
-
-
-defupdate(frame, population, pop_tracker, destinations, pop_size, infection_range=0.01, 
-infection_chance=0.03, speed=0.01, recovery_duration=(200, 500), mortality_chance=0.02,
-xbounds=[0.02, 0.98], ybounds=[0.02, 0.98], x_plot=[0, 1], 
-y_plot=[0, 1], wander_range=0.05, risk_age=55, 
-critical_age=75, critical_mortality_chance=0.1,
-risk_increase='quadratic', no_treatment_factor=3, 
-treatment_factor=0.5, healthcare_capacity=250, age_dependent_risk=False, 
-treatment_dependent_risk=False, visualise=False, verbose=False,
-self_isolate=True, self_isolate_proportion=0.6, isolation_bounds=[0, 0, 0.1, 0.1],
-traveling_infects=False, lockdown=False, lockdown_percentage=0.1, 
-lockdown_vector=[], plot_style='default'):
-
-#add one infection to jumpstart
-ifframe==50:
-population[0][6] =1
-population[0][8] =50
-population[0][10] =1
-
-
-ifvisualise:
-#construct plot and visualise
-spec=fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5,2])
-ax1.clear()
-ax2.clear()
-
-ax1.set_xlim(x_plot[0], x_plot[1])
-ax1.set_ylim(y_plot[0], y_plot[1])
-
-ifself_isolateandisolation_bounds!=None:
-build_hospital(isolation_bounds[0], isolation_bounds[2],
-isolation_bounds[1], isolation_bounds[3], ax1,
-addcross=False)
-
-#plot population segments
-healthy=population[population[:,6] ==0][:,1:3]
-ax1.scatter(healthy[:,0], healthy[:,1], color='gray', s=2, label='healthy')
-
-infected=population[population[:,6] ==1][:,1:3]
-ax1.scatter(infected[:,0], infected[:,1], color='red', s=2, label='infected')
-
-immune=population[population[:,6] ==2][:,1:3]
-ax1.scatter(immune[:,0], immune[:,1], color='green', s=2, label='immune')
-
-fatalities=population[population[:,6] ==3][:,1:3]
-ax1.scatter(fatalities[:,0], fatalities[:,1], color='black', s=2, label='dead')
-
-
-#add text descriptors
-ax1.text(x_plot[0], 
-y_plot[1] + ((y_plot[1] -y_plot[0]) /100), 
-'timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i'%(frame,
-len(population),
-len(healthy), 
-len(infected), 
-len(immune), 
-len(fatalities)),
-fontsize=6)
-
-ax2.set_title('number of infected')
-ax2.text(0, pop_size*0.05, 
-'https://github.com/paulvangentcom/python-corona-simulation',
-fontsize=6, alpha=0.5)
-#ax2.set_xlim(0, simulation_steps)
-ax2.set_ylim(0, pop_size+200)
-
-iftreatment_dependent_risk:
-infected_arr=np.asarray(pop_tracker.infectious)
-indices=np.argwhere(infected_arr>=healthcare_capacity)
-
-ax2.plot([healthcare_capacityforxinrange(len(pop_tracker.infectious))], 
-color='red', label='healthcare capacity')
-
-ifplot_style.lower() =='default':
-ax2.plot(pop_tracker.infectious, color='gray')
-ax2.plot(pop_tracker.fatalities, color='black', label='fatalities')
-elifplot_style.lower() =='sir':
-ax2.plot(pop_tracker.infectious, color='gray')
-ax2.plot(pop_tracker.fatalities, color='black', label='fatalities')
-ax2.plot(pop_tracker.susceptible, color='blue', label='susceptible')
-ax2.plot(pop_tracker.recovered, color='green', label='recovered')
-else:
-raiseValueError('incorrect plot_style specified, use \'sir\' or \'default\'')
-
-iftreatment_dependent_risk:
-ax2.plot(indices, infected_arr[infected_arr>=healthcare_capacity], 
-color='red')
-
-ax2.legend(loc='best', fontsize=6)
-#plt.savefig('render/%i.png' %frame)
-
-returnpopulation
+defrun(self):
+fortinrange(self.Config.simulation_steps):
+try:
+sim.tstep()
+sys.stdout.write('\r')
+sys.stdout.write('%i / %i'%(t, self.Config.simulation_steps))
+exceptKeyboardInterrupt:
+print('\nCTRL-C caught, exciting')
+sys.exit(1)
 
 
 if__name__=='__main__':
 
-tstep=2000
-
+#initialize
 sim=Simulation()
-
 
-fortinrange(tstep):
-sim.tstep()
-sys.stdout.write('\r')
-sys.stdout.write('%i / %i'%(t, tstep))
+#set number of simulation steps
+sim.Config.simulation_steps=2000
+
+#run
+sim.run()

visualiser.py

@@ -95,4 +95,9 @@ def draw_tstep(Config, population, pop_tracker, frame,
 color='red')
 
 ax2.legend(loc='best', fontsize=6)
-#plt.savefig('render/%i.png' %frame)
+
+plt.draw()
+plt.pause(0.0001)
+
+ifConfig.save_plot:
+plt.savefig('render/%i.png'%frame)