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.

community.py

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