Move more specific node types to pydtn.community

examples/shed.py

@@ -8,42 +8,14 @@ from collections import namedtuple
 from multiprocessing import Pool
 
 from pydtn import Network, RandomTraffic, Node, EpidemicNode
-from pydtn.community import BubbleNode, HCBFNode
-from pydtn.community import LouvainCommunity, KCliqueCommunity
+from pydtn.community import BubbleKCliqueNode, BubbleLouvainNode
+from pydtn.community import HCBFKCliqueNode, HCBFLouvainNode
 from pydtn.shed import ShedTrace
 
 
 Simulation = namedtuple('Simulation', ['trace', 'node_type', 'seed'])
 
 
-class LouvainNode(Node):
-    def __init__(self, **options):
-        options['community'] = LouvainCommunity(options['epoch'])
-        super().__init__(**options)
-
-
-class LouvainHCBFNode(LouvainNode, HCBFNode):
-    pass
-
-
-class LouvainBubbleNode(LouvainNode, BubbleNode):
-    pass
-
-
-class KCliqueNode(Node):
-    def __init__(self, **options):
-        options['community'] = KCliqueCommunity(options['epoch'], options['k'])
-        super().__init__(**options)
-
-
-class KCliqueHCBFNode(KCliqueNode, HCBFNode):
-    pass
-
-
-class KCliqueBubbleNode(KCliqueNode, BubbleNode):
-    pass
-
-
 def run_simulation(simulation):
     """Run a simulation."""
     seed = simulation.seed
@@ -92,10 +64,10 @@ def main(args):
     node_types = [
         Node,
         EpidemicNode,
-        KCliqueBubbleNode,
-        KCliqueHCBFNode,
-        LouvainBubbleNode,
-        LouvainHCBFNode,
+        BubbleKCliqueNode,
+        HCBFKCliqueNode,
+        BubbleLouvainNode,
+        HCBFLouvainNode,
     ]
 
     for seed in args['seeds']:

pydtn/community.py

@@ -16,8 +16,14 @@ __all__ = [
     "LouvainCommunity",
 
     "CommunityNode",
+    "KCliqueNode",
+    "LouvainNode",
     "BubbleNode",
+    "BubbleKCliqueNode",
+    "BubbleLouvainNode",
     "HCBFNode",
+    "HCBFKCliqueNode",
+    "HCBFLouvainNode",
 ]
 __author__ = 'Jarrod Pas <j.pas@usask.ca>'
 
@@ -301,6 +307,35 @@ class CommunityNode(Node):
         ]
 
 
+class LouvainNode(CommunityNode):
+    """Wrapper to add louvain community detection automatically to a node."""
+
+    def __init__(self, **options):
+        """
+        Create lovain node.
+
+        Keyword Arguments:
+        epoch -- how often to recalculate communities
+        """
+        options['community'] = LouvainCommunity(options['epoch'])
+        super().__init__(**options)
+
+
+class KCliqueNode(CommunityNode):
+    """Wrapper to add kclique community detection automatically to a node."""
+
+    def __init__(self, **options):
+        """
+        Create k-clique node.
+
+        Keyword Arguments:
+        epoch -- how often to recalculate communities
+        k -- initial community size
+        """
+        options['community'] = KCliqueCommunity(options['epoch'], options['k'])
+        super().__init__(**options)
+
+
 def _decide(node, others, key):
     """
     Make a decision on which node to send to best is decided based on key.
@@ -349,6 +384,18 @@ class BubbleNode(CommunityNode):
         return {}
 
 
+class BubbleKCliqueNode(KCliqueNode, BubbleNode):
+    """Bubble node with k-clique community detection."""
+
+    pass
+
+
+class BubbleLouvainNode(LouvainNode, BubbleNode):
+    """Bubble node with louvain community detection."""
+
+    pass
+
+
 class HCBFNode(CommunityNode):
     """Node with Hybrid Community Based forwarding."""
 
@@ -392,3 +439,15 @@ class HCBFNode(CommunityNode):
                 return {target: 'local-popularity'}
 
         return {}
+
+
+class HCBFKCliqueNode(KCliqueNode, HCBFNode):
+    """HCBF node with k-clique community detection."""
+
+    pass
+
+
+class HCBFLouvainNode(LouvainNode, HCBFNode):
+    """HCBF node with louvain community detection."""
+
+    pass