Skip to content

Instantly share code, notes, and snippets.

@correlator correlator/community.py
Last active Aug 29, 2015

Embed
What would you like to do?
A class that simulates the spread of community on a graph based on the SIR model https://en.wikipedia.org/wiki/Epidemic_model#The_SIR_model.
import scipy.stats
from scipy import random
import heapq
import networkx as nx
from copy import deepcopy
import matplotlib.pyplot as plt
matplotlib.rcParams.update({'font.size': 22})
class Community():
"""
example usage: > graph = nx.barabasi_albert_graph(20, 1)
> model = Community(graph)
> results = model.run('mov_try/')
"""
TIME = 0
NODE_ID = 1
COLOR_MAP = { 'prospect': 'blue',
'supporter': 'green',
'churn': 'orange'
}
def __init__(self, networkx_graph, support_prob = 0.2, churn_rate = 0.11):
self.community = deepcopy(networkx_graph)
self.community_size = self.community.number_of_nodes()
self.pos = nx.graphviz_layout(self.community, prog='fdp')
self.support_prob = support_prob
self.churn_rate = churn_rate
def run(self, save_dir = None, start_node = None, max_iter = 500):
"""
Returns a dict with number of prospects, supporters,
and churned people at each step.
"""
self._prepare_for_run(start_node, save_dir)
while self.t < max_iter and self._has_current_supporters():
self.t += 1
self._run_next_step()
return self.summary
def _prepare_for_run(self, start_node, save_dir):
self._initialize_constants(save_dir)
self._initialize_graph(start_node)
def _initialize_constants(self, save_dir):
self.save_dir = save_dir
self.t = 0
self.supporter_count = 0
self.churned_count = 0
self.churn_events = []
self.summary = { 'prospect': [], 'supporter': [], 'churn': [] }
def _initialize_graph(self, start_node = None):
nx.set_node_attributes(self.community, 'engagement', 'prospect')
first_supporter = start_node or random.randint(0, self.community_size)
self._make_supporter(first_supporter)
def _run_next_step(self):
self._find_new_supporters()
self._find_churn()
self._update_summary()
if self.save_dir is not None:
self._save_plot()
def _find_new_supporters(self):
for supporter in self._current_supporters():
for prospect in self._prospect_neighbors(supporter):
if (random.random() < self.support_prob):
self._make_supporter(prospect)
def _make_supporter(self, node_id):
self.supporter_count += 1
self.community.node[node_id]['engagement'] = 'supporter'
churn_time = self.t + scipy.random.exponential(1/self.churn_rate)
churn_event = self._build_churn_event(churn_time, node_id)
heapq.heappush(self.churn_events, churn_event)
def _build_churn_event(self, churn_time, node_id):
churn_event = [None] * 2
churn_event[self.TIME] = churn_time
churn_event[self.NODE_ID] = node_id
return tuple(churn_event)
def _find_churn(self):
if len(self.churn_events) > 0:
while self.churn_events[0][self.TIME] <= self.t:
node_to_churn = heapq.heappop(self.churn_events)
self._churn_node(node_to_churn[self.NODE_ID])
if len(self.churn_events) == 0:
break
def _churn_node(self, node_id):
self.churned_count += 1
self.supporter_count -= 1
self.community.node[node_id]['engagement'] = 'churn'
def _current_supporters(self):
return [node_id for node_id, attr_dict in self.community.node.items()
if attr_dict ['engagement'] == 'supporter' ]
def _prospect_neighbors(self, node_id):
neighbors = self.community.neighbors(node_id)
return [neighbor for neighbor in neighbors
if self.community.node[neighbor]['engagement'] == 'prospect']
def _update_summary(self):
churns = self.churned_count
supporters = self.supporter_count
prostects = self.community_size - churns - supporters
self.summary['prospect'].append(prostects)
self.summary['supporter'].append(supporters)
self.summary['churn'].append(churns)
def _color_values(self):
return [self.COLOR_MAP[attributes['engagement']]
for node, attributes in self.community.nodes_iter(data=True)]
def _save_plot(self):
plt.figure(figsize = (15,15))
plt.suptitle('Supporter Prob: %.2f, Churn Rate: %.2f'
% (self.support_prob, self.churn_rate))
plt.subplot(2,1,1)
nx.draw(self.community, self.pos, node_color = self._color_values())
plt.subplot(2,1,2)
plt.plot(self.summary['prospect'], c = self.COLOR_MAP['prospect'], linewidth=4)
plt.plot(self.summary['supporter'], c= self.COLOR_MAP['supporter'], linewidth=4)
plt.plot(self.summary['churn'], c= self.COLOR_MAP['churn'], linewidth=4)
plt.xlabel('Week number')
plt.ylabel('Number of people')
plt.legend(('prospects', 'supporters', 'churned'))
plt.savefig(self.save_dir + '/graph_%03d' % self.t)
plt.close('all')
def _has_current_supporters(self):
return len(self._current_supporters()) > 0
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.