zika.py

# Biology IA Simulations
# By Daniel Watson

import numpy
import matplotlib.pyplot as plt

# Controlled variables
h = .1 # days (step size)
population = 3.9e6 # remains constant
mosquitoes = 1e8 # remains constant
initial_mosquitoes_infected = 1000
recovery_time = 4.5 # days
mosquito_lifetime = 14. # days

# Simulation
def zika(num_steps, bite_factor):
  bites = 3. / mosquito_lifetime * bite_factor
  # Arrays for data
  s = numpy.zeros(num_steps + 1)
  i = numpy.zeros(num_steps + 1)
  r = numpy.zeros(num_steps + 1)
  m_s = numpy.zeros(num_steps + 1)
  m_i = numpy.zeros(num_steps + 1)
  # Distribution of humans
  s[0] = population
  # Distribution of mosquitoes
  m_i[0] = initial_mosquitoes_infected
  m_s[0] = mosquitoes - m_i[0]
  # Variables for storing maximum values
  max_h = i[0]
  max_m = m_i[0]
  time_max_h = 0
  time_max_m = 0

  # Forward Euler Method
  for k in range(num_steps):
    # Dynamic quantities used in calculations
    i2r = i[k] / recovery_time
    s2i_h = bites * s[k] * m_i[k] / mosquitoes
    s2i_m = bites * m_s[k] * i[k] / population
    mosquito_births = (m_i[k] + m_s[k]) / mosquito_lifetime # mosquito population is constant
    # Update all the system according to the five differential equations
    r[k+1] = r[k] + h * i2r
    i[k+1] = i[k] + h * (s2i_h - i2r)
    s[k+1] = s[k] - h * s2i_h
    m_i[k+1] = m_i[k] + h * (s2i_m - m_i[k] / mosquito_lifetime)
    m_s[k+1] = m_s[k] + h * (mosquito_births - s2i_m - m_s[k] / mosquito_lifetime)
    # Record maximum values for humans and mosquitoes
    if i[k+1] > max_h:
      max_h = i[k+1]
      time_max_h = k
    if m_i[k+1] > max_m:
      max_m = m_i[k+1]
      time_max_m = k

  return {'humans': i,
          'mosquitoes': m_i,
          'recovered': r,
          'max_humans': max_h,
          'max_mosquitoes': max_m,
          'max_humans_step': time_max_h,
          'max_mosquitoes_step': time_max_m
         }

# Run the simulation 100 times with a range of bite factors
# to plot it against maximum infection percentages or occurrence
def plot_bite_factor(option, bite_factor_start_value, bite_factor_end_value, show_legend=True):
  num_bite_factor_tests = 100
  # Arrays for plotting
  hum = []
  mos = []
  # Variables that depend on arguments
  bite_factor_step_size = (bite_factor_end_value - bite_factor_start_value) / float(num_bite_factor_tests)
  tests = bite_factor_step_size * numpy.array(range(num_bite_factor_tests + 1)) + bite_factor_start_value
  
  # Simulation loop
  for bite_factor in tests:
    total_time = 365*10
    num_steps = int(total_time / h)
    simulation = zika(num_steps, bite_factor)
    if option == 'maxima':
      y_h = 100 * simulation['max_humans'] / population
      y_m = 100 * simulation['max_mosquitoes'] / mosquitoes
      ylabel = "Maximum perentage infected"
    elif option == 'time':
      y_h = simulation['max_humans_step'] * total_time / num_steps
      y_m = simulation['max_mosquitoes_step'] * total_time / num_steps
      ylabel = "Maximum infection percentage time (days)"
    else:
      raise RuntimeError('"%s" is not a valid option!' % option)
    hum.append(y_h)
    mos.append(y_m)
    print "b={b}: ({h}, {m})".format(b=bite_factor, h=y_h, m=y_m)
  plt.plot(tests, hum, label="Humans")
  plt.plot(tests, mos, label="Mosquitoes")
  if show_legend:
    plt.legend(("Humans", "Mosquitoes"), loc='upper right')
  axes = plt.gca()
  axes.set_xlabel("Bite factor")
  axes.set_ylabel(ylabel)
  plt.xlim(xmin=bite_factor_start_value, xmax=bite_factor_start_value + \
           num_bite_factor_tests * bite_factor_step_size)
  plt.ylim(ymin=0.)
  plt.show()

# Plot percentages of infected human and mosquito
# populations over time, for a single bite factor
def plot_simulation(bite_factor, total_time=365*10, show_legend=True, show_recovered=True):
  # Argument dependents
  num_steps = int(total_time / h)
  times = h * numpy.array(range(num_steps + 1))
  # Simulation results
  simulation = zika(num_steps, bite_factor)
  print "max_h: " + str(100 * simulation['max_humans'] / population)
  print "max_h time: " + str(simulation['max_humans_step'] * total_time / num_steps)
  print "max_m: " + str(100 * simulation['max_mosquitoes'] / mosquitoes)
  print "max_m time: " + str(simulation['max_mosquitoes_step'] * total_time / num_steps)
  print "R: " + str(100 * simulation['recovered'][-1] / population)
  # Plot arrays of simulation values
  plt.plot(times, 100 * simulation['humans'] / population, label="Humans")
  plt.plot(times, 100 * simulation['mosquitoes'] / mosquitoes, label="Mosquitoes")
  if show_recovered:
    plt.plot(times, 100 * simulation['recovered'] / population, label="Recovered humans")
  if show_legend:
    plt.legend(("Humans", "Mosquitoes", "Recovered humans"), loc='upper right')
  axes = plt.gca()
  axes.set_xlabel("Time (days)")
  axes.set_ylabel("Percentage infected")
  plt.title("Zika simulation for b=" + str(bite_factor))
  plt.xlim(xmin=0.,xmax=total_time)
  plt.ylim(ymin=0.)
  plt.show()

### Simulation tests. Only use one at a time; feel
### free to play with these or to use your own.
### The tests below are those used to generate the
### graphs presented on the actual report.

#plot_bite_factor('time', 0., 1.)
#plot_bite_factor('time', .5879, .589481, show_legend=False)
#plot_bite_factor('time', .58892, .58898, show_legend=False)
#plot_bite_factor('time', .58943357, .59288, show_legend=False)
#plot_bite_factor('maxima', 0., 1.)
#plot_bite_factor('maxima', 0., .5889386)
#plot_simulation(.1, total_time=365)
#plot_simulation(.3, total_time=365, show_legend=False)
#plot_simulation(.5, total_time=365*2, show_legend=False)
#plot_simulation(.5889386, show_legend=False, show_recovered=False)
#plot_simulation(.5894681, show_legend=False, show_recovered=False)
#plot_simulation(.5911396, show_legend=False, show_recovered=False)
#plot_simulation(.5928111, show_legend=False, show_recovered=False)
#plot_simulation(.595, show_legend=False, show_recovered=False)
#plot_simulation(.6, show_legend=False, show_recovered=False)
plot_simulation(.7, show_legend=False, show_recovered=False)