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
 from population import initialize_population, initialize_destination_matrix,\
 set_destination_bounds, save_data, save_population, Population_trackers
-from visualiser import build_fig, draw_tstep, set_style
+from visualiser import build_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)
 
 
  def population_init(self):
@@ -50,6 +47,10 @@ def tstep(self):
  '''
  takes a time step in the simulation
  '''
+ 
+ if self.frame == 0 and self.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]))
+ 
+
+ def plot_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))
+ 
+ 
+def plot_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')
+ if include_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()