-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
148 lines (122 loc) · 4.47 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import streamlit as st
import matplotlib.pyplot as plt
from simulation import Simulation
if "running" not in st.session_state:
st.session_state.running = False
if "data" not in st.session_state:
st.session_state.data = None
def run_simulation():
epidemic_simulation = Simulation(
population=population_size,
initially_infected=initially_infected,
recovery_chance=recovery_chance / 100,
transmission_rate=transmission_rate / 100,
)
days = epidemic_simulation.run()
unaffected, infected, recovered = zip(*epidemic_simulation.history)
days = range(len(epidemic_simulation.history))
st.session_state.data = days, unaffected, infected, recovered
st.title("Pandemic Simulation")
with st.expander(
label="Simulation Description",
icon=":material/info:",
):
st.markdown(
body="""
This is a simple simulation of a pandemic, with a few assumptions:
- Population of given size starts with the specified amount of sick people
- Every single day all the people meet each other
- For every contact with a healthy person, a sick person has a probability of spreading the disease specified by the user
- After that, every sick person has a chance of recovery specified by the user
- This cycle continues until there is no sick person left
- People who recovered from the disease cannot become sick again
- The disease is not lethal and everyone infected survives until the end of the simulation
- The simulation may either end when every single person gets infected and everyone recovers, or the disease spreads too slowly and initial population of sick people recovers before spreading the disease further
- For the same input parameters, the simulation may have a different outcomes
"""
)
st.header("Simulation parameters")
disabled = st.session_state.running
with st.form(key="form"):
left, right = st.columns(2)
with left:
population_size = st.number_input(
label="Population size:",
min_value=1,
max_value=10_000,
value=1000,
step=1,
disabled=disabled,
)
initially_infected = st.number_input(
label="Number of initially infected people:",
min_value=1,
max_value=population_size,
value=5,
step=1,
disabled=disabled,
)
with right:
transmission_rate = st.number_input(
label="Transmission probability [%]:",
min_value=0.1,
max_value=100.0,
value=5.0,
step=0.1,
disabled=disabled,
help="Determines how probable is it that a sick person will spread the virus to a healthy one",
)
recovery_chance = st.number_input(
label="Recovery chance [%]:",
min_value=0.1,
max_value=100.0,
value=2.5,
step=0.1,
disabled=disabled,
help="Determines how probable is it that a sick person will recover from the disease on a given day",
)
submitted = st.form_submit_button(
label="Run simulation",
type="primary",
icon=":material/play_arrow:",
on_click=lambda: st.session_state.update(running=True),
disabled=disabled,
)
if submitted:
with st.spinner("Running simulation..."):
run_simulation()
st.session_state.update(running=False)
st.rerun()
st.divider()
if not st.session_state.data or st.session_state.running:
# Show placeholder message
c = st.container(
border=True,
height=200,
)
c.markdown(
body=f"""
<div style="text-align: center;">
Results will appear here.
</div>""",
unsafe_allow_html=True,
)
else:
days, unaffected, infected, recovered = st.session_state.data
plt.figure(figsize=(10, 6))
plt.plot(days, unaffected, label="unaffected")
plt.plot(days, infected, label="infected")
plt.plot(days, recovered, label="recovered")
plt.xlabel("days")
plt.ylabel("number of people")
plt.title("Epidemic simulation results")
plt.legend()
st.pyplot(plt)
st.code(
f"""
Population size: {population_size}
Number of initially infected people: {initially_infected}
Transmission probability: {transmission_rate}%
Recovery chance: {recovery_chance}%
"""
)