def update(sim, scat, linhas1, linhas2, linhas3, statistics):
"""
Execute an iteration of the simulation and update the animation graphics
:param sim:
:param scat:
:param linhas1:
:param linhas2:
:param statistics:
:return:
"""
sim.execute()
scat.set_facecolor([color2(a) for a in sim.get_population()])
scat.set_offsets(sim.get_positions())
df1, df2 = update_statistics(sim, statistics)
for col in linhas1.keys():
linhas1[col].set_data(df1.index.values, df1[col].values)
for col in linhas2.keys():
linhas2[col].set_data(df2.index.values, df2[col].values)
for col in linhas3.keys():
linhas3[col].set_data(df1.index.values, df1[col].values)
ret = [scat]
for l in linhas1.values():
ret.append(l)
for l in linhas2.values():
ret.append(l)
for l in linhas3.values():
ret.append(l)
return tuple(ret)
def draw_graph2(sim, ax=None, edges=False):
import networkx as nx
from covid_abs.graphics import color2
G = nx.Graph()
pos = {}
G.add_node(sim.healthcare.id, type='healthcare')
pos[sim.healthcare.id] = [sim.healthcare.x, sim.healthcare.y]
houses = []
for house in sim.houses:
G.add_node(house.id, type='house')
pos[house.id] = [house.x, house.y]
houses.append(house.id)
buss = []
for bus in sim.business:
G.add_node(bus.id, type='business')
pos[bus.id] = [bus.x, bus.y]
buss.append(bus.id)
colors = {}
for person in sim.population:
G.add_node(person.id, type='person')
col = color2(person)
if col not in colors:
colors[col] = {'status': person.status, 'severity': person.infected_status, 'id':[]}
colors[col]['id'].append(person.id)
pos[person.id] = [person.x, person.y]
if edges:
for house in sim.houses:
for person in house.homemates:
G.add_edge(house.id, person.id)
for bus in sim.business:
for person in bus.employees:
G.add_edge(bus.id, person.id)
#nx.draw(G, ax=ax, pos=pos, node_color=colors, node_size=sizes)
nx.draw_networkx_nodes(G, pos=pos, nodelist=[sim.healthcare.id], node_color='darkseagreen', label='Hospital')
nx.draw_networkx_nodes(G, pos=pos, nodelist=houses, node_color='cyan', label='Houses')
nx.draw_networkx_nodes(G, pos=pos, nodelist=buss, node_color='darkviolet', label='Business')
for key in colors.keys():
nx.draw_networkx_nodes(G, pos=pos, nodelist=[sim.healthcare.id], node_color='darkseagreen', label='Hospital')
def draw_graph(sim, ax=None, edges=False):
import networkx as nx
from covid_abs.graphics import color2
G = nx.Graph()
colors = []
pos = {}
sizes = []
G.add_node(sim.healthcare.id, type='healthcare')
colors.append('darkseagreen')
pos[sim.healthcare.id] = [sim.healthcare.x, sim.healthcare.y]
sizes.append(100)
for house in sim.houses:
G.add_node(house.id, type='house')
colors.append('cyan')
pos[house.id] = [house.x, house.y]
sizes.append(100)
for bus in sim.business:
G.add_node(bus.id, type='business')
colors.append('darkviolet')
pos[bus.id] = [bus.x, bus.y]
sizes.append(100)
for person in sim.population:
G.add_node(person.id, type='person')
colors.append(color2(person))
pos[person.id] = [person.x, person.y]
sizes.append(10)
if edges:
for house in sim.houses:
for person in house.homemates:
G.add_edge(house.id, person.id)
for bus in sim.business:
for person in bus.employees:
G.add_edge(bus.id, person.id)
nx.draw(G, ax=ax, pos=pos, node_color=colors, node_size=sizes)
if ax is not None:
ax.set_xlim((0, sim.length))
ax.set_ylim((0, sim.height))
def execute_simulation(sim, **kwargs):
"""
Execute a simulation and plot its results
:param sim: a Simulation or MultiopulationSimulation object
:param iterations: number of interations of the simulation
:param iteration_time: time (in miliseconds) between each iteration
:return: an animation object
"""
statistics = {'info': [], 'ecom': []}
fig, ax = plt.subplots(nrows=2, ncols=2, figsize=[20, 10])
# fig, ax = plt.subplots(nrows=1, ncols=3, figsize=[20, 5])
# plt.close()
frames = kwargs.get('iterations', 100)
iteration_time = kwargs.get('iteration_time', 250)
sim.initialize()
ax[0][0].set_title('Simulation Environment')
ax[0][0].set_xlim((0, sim.length))
ax[0][0].set_ylim((0, sim.height))
pos = np.array(sim.get_positions())
scat = ax[0][0].scatter(
pos[:, 0],
pos[:, 1],
c=[color2(a) for a in sim.get_population()],
edgecolors=[edgecolor(a) for a in sim.get_population()],
s=[nodesize(a) for a in sim.get_population()],
linewidths=0.1)
df1, df2 = update_statistics(sim, statistics)
ax[0][1].set_title('Contagion Evolution')
ax[0][1].set_xlim((0, frames))
ax[0][1].set_ylim((0, 1))
linhas1 = {}
ax[0][1].axhline(y=sim.critical_limit,
c="black",
ls='--',
label='Critical limit')
for col in df1.columns.values:
if col not in ["R0"]:
linhas1[col], = ax[0][1].plot(df1.index.values,
df1[col].values,
c=color1(col),
label=col)
ax[0][1].set_xlabel("Nº of Days")
ax[0][1].set_ylabel("% of Population")
handles, labels = ax[0][1].get_legend_handles_labels()
lgd = ax[0][1].legend(handles, labels,
loc='top right') #2, bbox_to_anchor=(0, 0))
linhas2 = {}
ax[1][1].set_title('Economical Impact')
ax[1][1].set_xlim((0, frames))
for col in df2.columns.values:
linhas2[col], = ax[1][1].plot(df2.index.values,
df2[col].values,
c=color3(col),
label=legend_ecom[col])
ax[1][1].set_xlabel("Nº of Days")
ax[1][1].set_ylabel("Wealth")
handles, labels = ax[1][1].get_legend_handles_labels()
lgd = ax[1][1].legend(handles, labels,
loc='top right') #2, bbox_to_anchor=(1, 1))
linhas3 = {}
ax[1][0].set_title("R0")
ax[1][0].set_xlim((0, frames))
ax[1][0].set_ylim((0, 10))
for col in df1.columns.values:
if col == "R0":
linhas3[col], = ax[1][0].plot(df1.index.values,
df1[col].values,
c=color1(col),
label=col)
ax[1][0].set_xlabel("Nº of Days")
ax[1][0].set_ylabel("R0")
animate = lambda i: update(sim, scat, linhas1, linhas2, linhas3, statistics
)
init = lambda: clear(scat, linhas1, linhas2, linhas3)
# animation function. This is called sequentially
anim = animation.FuncAnimation(fig,
animate,
init_func=init,
frames=frames,
interval=iteration_time,
blit=True)
if kwargs.get("filename", None) is not None:
anim.save(kwargs.get("filename", None) + ".gif", writer='imagemagick')
fig.savefig(kwargs.get("filename", None) + ".png", dpi=300)
return anim
def execute_simulation(sim, **kwargs):
"""
Execute a simulation and plot its results
:param sim: a Simulation or MultiopulationSimulation object
:param iterations: number of interations of the simulation
:param iteration_time: time (in miliseconds) between each iteration
:return: an animation object
"""
third_plot = kwargs.get('third_plot', 'ecom')
statistics = {'info': [], third_plot: []}
fig, ax = plt.subplots(nrows=1, ncols=3, figsize=[20, 5])
# plt.close()
frames = kwargs.get('iterations', 100)
iteration_time = kwargs.get('iteration_time', 250)
sim.initialize()
ax[0].set_title('Simulation Environment')
ax[0].set_xlim((0, sim.length))
ax[0].set_ylim((0, sim.height))
pos = np.array(sim.get_positions())
scat = ax[0].scatter(pos[:, 0],
pos[:, 1],
c=[color2(a) for a in sim.get_population()])
df1, df2 = update_statistics(sim, statistics, third_plot)
ax[1].set_title('Contagion Evolution')
ax[1].set_xlim((0, frames))
linhas1 = {}
ax[1].axhline(y=sim.critical_limit,
c="black",
ls='--',
label='Critical limit')
for col in df1.columns.values:
if col != 'Asymptomatic':
linhas1[col], = ax[1].plot(df1.index.values,
df1[col].values,
c=color1(col),
label=col)
ax[1].set_xlabel("Nº of Days")
ax[1].set_ylabel("% of Population")
handles, labels = ax[1].get_legend_handles_labels()
lgd = ax[1].legend(handles, labels,
loc='upper right') #2, bbox_to_anchor=(0, 0))
if third_plot == 'ecom':
linhas2 = {}
ax[2].set_title('Economical Impact')
ax[2].set_xlim((0, frames))
for col in df2.columns.values:
linhas2[col], = ax[2].plot(df2.index.values,
df2[col].values,
c=color3(col),
label=legend_ecom[col])
ax[2].set_xlabel("Nº of Days")
ax[2].set_ylabel("Wealth")
handles, labels = ax[2].get_legend_handles_labels()
lgd = ax[2].legend(handles, labels,
loc='upper right') #2, bbox_to_anchor=(1, 1))
elif third_plot == 'R':
linhas2 = {}
ax[2].set_title('Average effective Reproduction Number')
ax[2].set_xlim((0, frames))
ax[2].set_ylim((0, 3))
for col in df2.columns.values:
linhas2[col], = ax[2].plot(df2.index.values,
df2[col].values,
c="c",
label='R-Value')
ax[2].set_xlabel("Nº of Days")
ax[2].set_ylabel("R")
handles, labels = ax[2].get_legend_handles_labels()
lgd = ax[2].legend(handles, labels,
loc='upper right') # 2, bbox_to_anchor=(1, 1))
animate = lambda i: update(sim, scat, linhas1, linhas2, statistics,
third_plot)
init = lambda: clear(scat, linhas1, linhas2)
# animation function. This is called sequentially
anim = animation.FuncAnimation(fig,
animate,
init_func=init,
frames=frames,
interval=iteration_time,
blit=True)
return anim
