feat: made colors configurable for scatter plots

array based on healthy,infected, immune, dead

Author: chorsnell

plot.py

@@ -3,15 +3,21 @@
 
 mpl.style.use('style.mplstyle')
 
+# healthy, infected, immune, dead
+personStateColors = ['#404040', '#ff0000', '#00ff00', '#000000']
+
+# sets graph title
 def ax2Title(ax2):
 ax2.set_title('number of infected')
 
+# updates figure with common values across all simulations
 def figUpdate(ax1, ax2, x_plot, y_plot):
 ax1.set_xlim(x_plot[0], x_plot[1])
 ax1.set_ylim(y_plot[0], y_plot[1])
 ax1.axis('off')
 ax2Title(ax2)
 
+# initialises figure with common values across all simulations
 def figInit(xbounds, ybounds, pop_size):
 fig = plt.figure(figsize=(5,7))
 spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5,2])

simple_simulation.py

@@ -5,7 +5,7 @@
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 
-from plot import figInit, figUpdate
+from plot import figInit, figUpdate, personStateColors
 
 #set seed for reproducibility
 np.random.seed(100)
@@ -318,16 +318,16 @@ def update(frame, population, infection_range=0.01, infection_chance=0.03,
  figUpdate(ax1, ax2, xbounds, ybounds)
 
  healthy = population[population[:,6] == 0][:,1:3]
- ax1.scatter(healthy[:,0], healthy[:,1], color='gray', s = 2, label='healthy')
+ ax1.scatter(healthy[:,0], healthy[:,1], color=personStateColors[0], s = 2, label='healthy')
 
  infected = population[population[:,6] == 1][:,1:3]
- ax1.scatter(infected[:,0], infected[:,1], color='red', s = 2, label='infected')
+ ax1.scatter(infected[:,0], infected[:,1], color=personStateColors[1], s = 2, label='infected')
 
  immune = population[population[:,6] == 2][:,1:3]
- ax1.scatter(immune[:,0], immune[:,1], color='green', s = 2, label='immune')
+ ax1.scatter(immune[:,0], immune[:,1], color=personStateColors[2], s = 2, label='immune')
 
  fatalities = population[population[:,6] == 3][:,1:3]
- ax1.scatter(fatalities[:,0], fatalities[:,1], color='black', s = 2, label='fatalities')
+ ax1.scatter(fatalities[:,0], fatalities[:,1], color=personStateColors[3], s = 2, label='fatalities')
 
  #add text descriptors
  ax1.text(xbounds[0], 

simulation.py

@@ -12,7 +12,7 @@
 from population import initialize_population, initialize_destination_matrix,\
 set_destination_bounds, save_data
 
-from plot import figInit, figUpdate
+from plot import figInit, figUpdate, personStateColors
 
 #set seed for reproducibility
 np.random.seed(100)
@@ -118,16 +118,16 @@ def update(frame, population, destinations, pop_size, infection_range=0.01,
 
  #plot population segments
  healthy = population[population[:,6] == 0][:,1:3]
- ax1.scatter(healthy[:,0], healthy[:,1], color='gray', s = 2, label='healthy')
+ ax1.scatter(healthy[:,0], healthy[:,1], color=personStateColors[0], s = 2, label='healthy')
 
  infected = population[population[:,6] == 1][:,1:3]
- ax1.scatter(infected[:,0], infected[:,1], color='red', s = 2, label='infected')
+ ax1.scatter(infected[:,0], infected[:,1], color=personStateColors[1], s = 2, label='infected')
 
  immune = population[population[:,6] == 2][:,1:3]
- ax1.scatter(immune[:,0], immune[:,1], color='green', s = 2, label='immune')
+ ax1.scatter(immune[:,0], immune[:,1], color=personStateColors[2], s = 2, label='immune')
 
  fatalities = population[population[:,6] == 3][:,1:3]
- ax1.scatter(fatalities[:,0], fatalities[:,1], color='black', s = 2, label='dead')
+ ax1.scatter(fatalities[:,0], fatalities[:,1], color=personStateColors[3], s = 2, label='dead')
 
 
  #add text descriptors

simulation_hospital.py

@@ -12,7 +12,7 @@
 from population import initialize_population, initialize_destination_matrix,\
 set_destination_bounds, save_data
 
-from plot import figInit, figUpdate
+from plot import figInit, figUpdate, personStateColors
 
 
 def update(frame, population, destinations, pop_size, infection_range=0.01, 
@@ -97,21 +97,21 @@ def update(frame, population, destinations, pop_size, infection_range=0.01,
  healthy = population[population[:,6] == 0][:,1:3]
  ax1.scatter(healthy[:healthcare_workers][:,0], 
  healthy[:healthcare_workers][:,1], 
- marker= 'P', s = 2, color='green', 
+ marker= 'P', s = 2, color=personStateColors[2], 
  label='healthy')
 
  ax1.scatter(healthy[healthcare_workers:][:,0], 
  healthy[healthcare_workers:][:,1], 
- color='gray', s = 2, label='healthy')
+ color=personStateColors[0], s = 2, label='healthy')
 
  infected = population[population[:,6] == 1][:,1:3]
- ax1.scatter(infected[:,0], infected[:,1], color='red', s = 2, label='infected')
+ ax1.scatter(infected[:,0], infected[:,1], color=personStateColors[1], s = 2, label='infected')
 
  immune = population[population[:,6] == 2][:,1:3]
- ax1.scatter(immune[:,0], immune[:,1], color='green', s = 2, label='immune')
+ ax1.scatter(immune[:,0], immune[:,1], color=personStateColors[2], s = 2, label='immune')
 
  fatalities = population[population[:,6] == 3][:,1:3]
- ax1.scatter(fatalities[:,0], fatalities[:,1], color='black', s = 2, label='dead')
+ ax1.scatter(fatalities[:,0], fatalities[:,1], color=personStateColors[3], s = 2, label='dead')
 
 
  #add text descriptors