Guest User

Untitled

a guest
Mar 14th, 2020
322
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from __future__ import division
  2. import numpy as np
  3. import matplotlib.pyplot as plt
  4. from collections import namedtuple
  5. import matplotlib.widgets as wid
  6. from recordtype import recordtype #changed because namedtuple is immutable (can't be modified)
  7.  
  8. PopState = namedtuple("PopState", [
  9. "day",
  10. "susceptible",
  11. "infected",
  12. "dead",
  13. "recovered",
  14. ])
  15. EpiParams = recordtype("EpiParams", [
  16. "days_infected",
  17. "infection_rate",
  18. "frac_critical",
  19. "frac_dead_icu",
  20. "frac_dead_hospital",
  21. "frac_dead_home",
  22. "start_day",
  23. "controlled_infection_rate",
  24. ])
  25.  
  26. ICU_BEDS = 90000
  27. HOSPITAL_BEDS = 900000
  28.  
  29. def update_state(state, params):
  30. day = state.day
  31. S = state.susceptible
  32. I = state.infected
  33. D = state.dead
  34. R = state.recovered
  35. N = S + I + D + R
  36.  
  37. days_betw_contacts = params.days_infected / params.infection_rate
  38. if day >= params.start_day:
  39. days_betw_contacts = params.days_infected / params.controlled_infection_rate
  40.  
  41. new_infections = (I * S) / (N * days_betw_contacts)
  42. recovered_infections = (I / params.days_infected) * (1 - params.frac_critical)
  43.  
  44. # For critical cases, we split into ICU, hospital, and home care, in descending order
  45. C = I * params.frac_critical;
  46. C_icu = min(C, ICU_BEDS);
  47. C_hospital = 0 if C < ICU_BEDS else min(C - ICU_BEDS, HOSPITAL_BEDS)
  48. C_home = 0 if C < (ICU_BEDS + HOSPITAL_BEDS) else C - (ICU_BEDS + HOSPITAL_BEDS)
  49.  
  50. recovered_critical = (
  51. ((C_icu / params.days_infected) * (1 - params.frac_dead_icu)) +
  52. ((C_hospital / params.days_infected) * (1 - params.frac_dead_hospital)) +
  53. ((C_home / params.days_infected) * (1 - params.frac_dead_home))
  54. )
  55.  
  56. dead_critical = (
  57. ((C_icu / params.days_infected) * (params.frac_dead_icu)) +
  58. ((C_hospital / params.days_infected) * (params.frac_dead_hospital)) +
  59. ((C_home / params.days_infected) * (params.frac_dead_home))
  60. )
  61.  
  62. newS = S - new_infections
  63. newI = I + new_infections - recovered_infections - recovered_critical - dead_critical
  64. newR = R + recovered_infections + recovered_critical
  65. newD = D + dead_critical
  66.  
  67. return PopState(day + 1, newS, newI, newD, newR)
  68.  
  69. def run_model(state, params):
  70. infections = [state.infected]
  71. deaths = [state.dead]
  72. while state.day<1000: #(state.infected) >= 1: # changed because, as implemented, the chart needs a constant number of values on x-axis
  73. state = update_state(state, params)
  74. infections.append(state.infected)
  75. deaths.append(state.dead)
  76. return infections, deaths, state
  77.  
  78. total_population = 327 * 1e6
  79. init_state = PopState(
  80. day=1,
  81. susceptible=total_population - 1,
  82. infected=1,
  83. dead=0,
  84. recovered=0)
  85. params = EpiParams(
  86. days_infected=12,
  87. infection_rate=3.5,
  88. frac_critical=0.05,
  89. frac_dead_icu=0.02,
  90. frac_dead_hospital=0.05,
  91. frac_dead_home=0.5,
  92. start_day=79,
  93. # Modify this parameter to vary the strength of infection controls
  94. controlled_infection_rate=3.5,
  95. )
  96.  
  97. infections_over_time, deaths_over_time, final_state = run_model(init_state, params)
  98.  
  99.  
  100. fig1, (ax1,ax2) = plt.subplots(2, 1, sharex='col')
  101. fig2 = plt.figure()
  102.  
  103. line1,=ax1.plot(infections_over_time)
  104. line2,=ax2.plot(deaths_over_time)
  105.  
  106.  
  107. ax1.set_title('Active Infections')
  108. ax2.set_title('Cumulative Deaths')
  109.  
  110.  
  111. axcolor = 'lightgoldenrodyellow'
  112.  
  113. #axis objects (for sliders)
  114. axIR = fig2.add_axes([0.3, .9, 0.5, 0.03], facecolor=axcolor)
  115. axdays_infected = fig2.add_axes([0.3, .8, 0.5, 0.03], facecolor=axcolor)
  116. axfrac_critical = fig2.add_axes([0.3, .7, 0.5, 0.03], facecolor=axcolor)
  117. axfrac_dead_icu = fig2.add_axes([0.3, .6, 0.5, 0.03], facecolor=axcolor)
  118. axfrac_dead_hospital = fig2.add_axes([0.3, .5, 0.5, 0.03], facecolor=axcolor)
  119. axfrac_dead_home = fig2.add_axes([0.3, .4, 0.5, 0.03], facecolor=axcolor)
  120. axstart_day = fig2.add_axes([0.3, .3, 0.5, 0.03], facecolor=axcolor)
  121.  
  122.  
  123.  
  124. #slider objects
  125. sIR = wid.Slider(axIR, 'Infection Rate', 0, 5.0, valinit=3.5, valstep=.02)
  126. sdays_infected = wid.Slider(axdays_infected, 'Days Infected', 1, 21, valinit=12, valstep=1)
  127. sfrac_critical = wid.Slider(axfrac_critical, 'Rate Critical Inf', 0, .2, valinit=.05, valstep=.01)
  128. sfrac_dead_icu = wid.Slider(axfrac_dead_icu, 'Mort Rate - ICU', 0, .1, valinit=.02, valstep=.01)
  129. sfrac_dead_hospital = wid.Slider(axfrac_dead_hospital, 'Mort Rate - Hospital', 0, .2, valinit=.05, valstep=.01)
  130. sfrac_dead_home = wid.Slider(axfrac_dead_home, 'Mort Rate - Home', 0, 1, valinit=.5, valstep=.01)
  131. sstart_day = wid.Slider(axstart_day, 'Start Day', 1, 999, valinit=79, valstep=1)
  132.  
  133. # days_infected=12,
  134. # frac_critical=0.05,
  135. # frac_dead_icu=0.02,
  136. # frac_dead_hospital=0.05,
  137. # frac_dead_home=0.5,
  138. # start_day=79,
  139.  
  140.  
  141.  
  142.  
  143. def update(val):
  144. params.controlled_infection_rate = sIR.val
  145. params.days_infected = sdays_infected.val
  146. params.frac_critical = sfrac_critical.val
  147. params.frac_dead_icu = sfrac_dead_icu.val
  148. params.frac_dead_hospital = sfrac_dead_hospital.val
  149. params.frac_dead_home = sfrac_dead_home.val
  150. params.start_day = sstart_day.val
  151.  
  152.  
  153.  
  154.  
  155. infections_over_time, deaths_over_time, final_state = run_model(init_state, params)
  156.  
  157. line1.set_ydata(infections_over_time)
  158. line2.set_ydata(deaths_over_time)
  159.  
  160. fig1.canvas.draw_idle()
  161.  
  162.  
  163. sIR.on_changed(update)
  164. sdays_infected.on_changed(update)
  165. sfrac_critical.on_changed(update)
  166. sfrac_dead_icu.on_changed(update)
  167. sfrac_dead_hospital.on_changed(update)
  168. sfrac_dead_home.on_changed(update)
  169. sstart_day.on_changed(update)
  170.  
  171.  
  172. plt.show()
RAW Paste Data