add method to plot SIR plot only, for example if visualisation is turned off

Author: paulvangentcom

simulation.py

@@ -15,7 +15,7 @@
 keep_at_destination, reset_destinations
 frompopulationimportinitialize_population, initialize_destination_matrix,\
 set_destination_bounds, save_data, save_population, Population_trackers
-fromvisualiserimportbuild_fig, draw_tstep, set_style
+fromvisualiserimportbuild_fig, draw_tstep, set_style, plot_sir
 
 #set seed for reproducibility
 #np.random.seed(100)
@@ -33,11 +33,8 @@ def __init__(self, *args, **kwargs):
 self.pop_tracker=Population_trackers()
 
 #initalise destinations vector
-self.destinations=initialize_destination_matrix(self.Config.pop_size, 1)
+self.destinations=initialize_destination_matrix(self.Config.pop_size, 1)
 
-self.fig, self.spec, self.ax1, self.ax2=build_fig(self.Config)
-
-#set_style(self.Config)
 
 
 defpopulation_init(self):
@@ -50,6 +47,10 @@ def tstep(self):
 '''
  takes a time step in the simulation
  '''
+
+ifself.frame==0andself.Config.visualise:
+#initialize figure
+self.fig, self.spec, self.ax1, self.ax2=build_fig(self.Config)
 
 #check destinations if active
 #define motion vectors if destinations active and not everybody is at destination
@@ -187,6 +188,12 @@ def run(self):
 print('total infectious: %i'%len(self.population[(self.population[:,6] ==1) |
  (self.population[:,6] ==4)]))
 print('total unaffected: %i'%len(self.population[self.population[:,6] ==0]))
+
+
+defplot_sir(self, size=(6,3), include_fatalities=False, 
+title='S-I-R plot of simulation'):
+plot_sir(self.Config, self.pop_tracker, size, include_fatalities,
+title)
 
 
 

visualiser.py

@@ -119,4 +119,52 @@ def draw_tstep(Config, population, pop_tracker, frame,
 plt.savefig('%s/%i.png'%(Config.plot_path, frame))
 except:
 check_folder(Config.plot_path)
-plt.savefig('%s/%i.png'%(Config.plot_path, frame))
+plt.savefig('%s/%i.png'%(Config.plot_path, frame))
+
+
+defplot_sir(Config, pop_tracker, size=(6,3), include_fatalities=False,
+title='S-I-R plot of simulation'):
+'''plots S-I-R parameters in the population tracker
+
+ Keyword arguments
+ -----------------
+ Config : class
+ the configuration class
+
+ pop_tracker : ndarray
+ the population tracker, containing
+
+ size : tuple
+ size at which the plot will be initialised (default: (6,3))
+
+ include_fatalities : bool
+ whether to plot the fatalities as well (default: False) 
+ '''
+
+#set plot style
+set_style(Config)
+
+#get color palettes
+palette=Config.get_palette()
+
+#plot the thing
+plt.figure(figsize=size)
+plt.title(title) 
+plt.plot(pop_tracker.susceptible, color=palette[0], label='susceptible')
+plt.plot(pop_tracker.infectious, color=palette[1], label='infectious')
+plt.plot(pop_tracker.recovered, color=palette[2], label='recovered')
+ifinclude_fatalities:
+plt.plot(pop_tracker.fatalities, color=palette[3], label='fatalities')
+
+#add axis labels
+plt.xlabel('time in hours')
+plt.ylabel('population')
+
+#add legend
+plt.legend()
+
+#beautify
+plt.tight_layout()
+
+#initialise
+plt.show()