add methods to save population data every 'n' timesteps

Author: paulvangentcom

config.py

@@ -14,8 +14,10 @@ def __init__(self, *args, **kwargs):
  self.verbose = False #whether to print infections, recoveries and fatalities to the terminal
  self.simulation_steps = 10000 #total simulation steps performed
  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.save_data = False #whether to dump data at end of simulation
+ self.save_pop = False #whether to save population matrix every 'save_pop_freq' timesteps
+ self.save_pop_freq = 10 #population data will be saved every 'n' timesteps. Default: 10
+ self.save_pop_folder = 'pop_data/' #folder to write population timestep data to
  self.endif_no_infections = True #whether to stop simulation if no infections remain
 
  #scenario flags

population.py

@@ -179,11 +179,30 @@ def save_data(population, pop_tracker):
  num_files = len(glob('data/*'))
  check_folder('data/%i' %num_files)
  np.save('data/%i/population.npy' %num_files, population)
- np.save('data/%i/infected.npy' %num_files, pop_tacker.infectious)
+ np.save('data/%i/infected.npy' %num_files, pop_tracker.infectious)
  np.save('data/%i/recovered.npy' %num_files, pop_tracker.recovered)
  np.save('data/%i/fatalities.npy' %num_files, pop_tracker.fatalities)
 
 
+def save_population(population, tstep=0, folder='data_tstep'):
+ '''dumps population data at given timestep to disk
+
+ Function that dumps the simulation data to specific files on the disk.
+ Saves final state of the population matrix, the array of infected over time,
+ and the array of fatalities over time
+
+ Keyword arguments
+ -----------------
+ population : ndarray
+ the array containing all the population information
+
+ tstep : int
+ the timestep that will be saved
+ ''' 
+ check_folder('%s/' %(folder))
+ np.save('%s/population_%i.npy' %(folder, tstep), population)
+
+
 class Population_trackers():
  '''class used to track population parameters