kclique.py

import networkx as nx

from .base import EpochCommunity

# TODO: this is probably bugged, performs terribly for bubble

class KCliqueCommunity(EpochCommunity):
    def __init__(self, k=3, threshold=300, epoch=604800, **kwargs):
        super().__init__(epoch, **kwargs)
        self.k = k
        self.threshold = threshold

    def process(self, network):
        while True:
            yield self.tick()
            g = self.next_epoch(self.env.now)

            G = nx.Graph()
            G.add_nodes_from(g.nodes())
            for a, b, duration in g.edges(data='duration'):
                if duration > self.threshold:
                    G.add_edge(a, b)

            for community in nx.k_clique_communities(G, self.k):
                for node in community:
                    self.community[node] = community