def _animation_update(self, t, drawn_nodes, drawn_I, time_markers): status = EoN.get_statuses(self.G, self.node_history, t) Inodelist = [node for node in self.nodelist if status[node] == 'I'] drawn_nodes.set_color([self.colordict[status[node]] for node in self.nodelist]) drawn_I[0].remove() drawn_I[0] = nx.draw_networkx_nodes(self.G, pos=self.pos, nodelist=Inodelist, color = self.colordict['I'], ax = self.network_axes, **self.kwargs) for tm in time_markers: tm.remove() self._highlight_time(t, time_markers)
def _initialize(self): initial_status = EoN.get_statuses(self.G, self.node_history, self.frame_times[0]) colorlist = [self.colordict[initial_status[node]] for node in self.nodelist] nodeset = {node for node in self.nodelist} edgelist = [edge for edge in self.G.edges() if edge[0] in nodeset and edge[1] in nodeset] nx.draw_networkx_edges(self.G, pos=self.pos, edgelist=edgelist, ax = self.network_axes) drawn_nodes = nx.draw_networkx_nodes(self.G, pos=self.pos, ax = self.network_axes, nodelist = self.nodelist, color=colorlist, **self.kwargs) Inodelist = [node for node in self.nodelist if initial_status[node] == 'I'] drawn_I = [nx.draw_networkx_nodes(self.G, pos=self.pos, nodelist=Inodelist, color = self.colordict['I'], ax = self.network_axes, **self.kwargs)] self.network_axes.set_xticks([]) self.network_axes.set_yticks([]) time_markers = [None for ax in self.timeseries_axi] self._highlight_time(self.frame_times[0], time_markers) return drawn_nodes, drawn_I, time_markers
def visualize(G, node_history, time, colordict={'S':'#009a80','I':'#ff2020', 'R':'gray'}, t=None, S=None, I=None, R=None, pos = None, SIR = True, figname=None): r''' ''' if SIR: if t is None: visualization = EoN_Plot(timeseries=[]) else: visualization = EoN_plot() else: if t is None: visualization = EoN_Plot(timeseries = [], SIR = SIR) else: visualization = EoN_Plot(timeseries = ['S', 'I'], SIR = SIR) status = EoN.get_statuses(G, node_history, time) visualization.plot_network(G, status, pos=pos, colordict=colordict) if t is not None: visualization.plot_timeseries(t, S, I, R, label = 'simulation') visualization.legend() if figname is not None: visualization.save(figname) return visualization