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)
 
 
+defsave_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)
+
+
 classPopulation_trackers():
 '''class used to track population parameters