# Untitled

a guest
Mar 14th, 2020
479
0
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",
13. "recovered",
14. ])
15. EpiParams = recordtype("EpiParams", [
16. "days_infected",
17. "infection_rate",
18. "frac_critical",
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)
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,
84. recovered=0)
85. params = EpiParams(
86. days_infected=12,
87. infection_rate=3.5,
88. frac_critical=0.05,
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,
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
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)