create methods to dump simulation data to disk

Author: paulvangentcom

config.py

@@ -36,6 +36,7 @@ def __init__(self, *args, **kwargs):
 self.x_plot= [0, 1] 
 self.y_plot= [0, 1]
 self.save_plot=False
+self.plot_path='render/'#folder where plots are saved to
 self.plot_style='default'#can be default, dark, ...
 self.colorblind_mode=False
 #if colorblind is enabled, set type of colorblindness

population.py

@@ -9,6 +9,7 @@
 importnumpyasnp
 
 frommotionimportget_motion_parameters
+fromutilsimportcheck_folder
 
 definitialize_population(Config, mean_age=45, max_age=105,
 xbounds=[0, 1], ybounds=[0, 1]):
@@ -157,7 +158,7 @@ def set_destination_bounds(population, destinations, xmin, ymin,
 returnpopulation, destinations
 
 
-defsave_data(population, infected, fatalities):
+defsave_data(population, pop_tracker):
 '''dumps simulation data to disk
 
  Function that dumps the simulation data to specific files on the disk.
@@ -176,10 +177,11 @@ def save_data(population, infected, fatalities):
  the array containing data of fatalities over time
  '''
 num_files=len(glob('data/*'))
-os.makedirs('data/%i'%num_files)
+check_folder('data/%i'%num_files)
 np.save('data/%i/population.npy'%num_files, population)
-np.save('data/%i/infected.npy'%num_files, infected)
-np.save('data/%i/fatalities.npy'%num_files, fatalities)
+np.save('data/%i/infected.npy'%num_files, pop_tacker.infectious)
+np.save('data/%i/recovered.npy'%num_files, pop_tracker.recovered)
+np.save('data/%i/fatalities.npy'%num_files, pop_tracker.fatalities)
 
 
 classPopulation_trackers():

simulation.py

@@ -154,6 +154,9 @@ def run(self):
 print('\nCTRL-C caught, exiting')
 sys.exit(1)
 
+ifself.Config.save_data:
+save_data(self.population, self.pop_tracker)
+
 
 if__name__=='__main__':
 

utils.py

@@ -0,0 +1,10 @@
+'''
+collection of utility methods shared across files
+'''
+
+importos
+
+defcheck_folder(folder='render/'):
+'''check if folder exists, make if not present'''
+ifnotos.path.exists(folder):
+os.makedirs(folder)

visualiser.py

@@ -7,6 +7,7 @@
 importnumpyasnp
 
 fromenvironmentimportbuild_hospital
+fromutilsimportcheck_folder
 
 defset_style(Config):
 '''sets the plot style
@@ -15,6 +16,7 @@ def set_style(Config):
 ifConfig.plot_style.lower() =='dark':
 mpl.style.use('plot_styles/dark.mplstyle')
 
+
 defbuild_fig(Config, figsize=(5,7)):
 set_style(Config)
 fig=plt.figure(figsize=(5,7))
@@ -30,6 +32,8 @@ def build_fig(Config, figsize=(5,7)):
 #ax2.set_xlim(0, simulation_steps)
 ax2.set_ylim(0, Config.pop_size+100)
 
+#if 
+
 returnfig, spec, ax1, ax2
 
 
@@ -111,4 +115,8 @@ def draw_tstep(Config, population, pop_tracker, frame,
 plt.pause(0.0001)
 
 ifConfig.save_plot:
-plt.savefig('render/%i.png'%frame)
+try:
+plt.savefig('%s/%i.png'%(Config.plot_path, frame))
+except:
+check_folder(Config.plot_path)
+plt.savefig('%s/%i.png'%(Config.plot_path, frame))