Merge branch 'feature/trace-class' into 'develop'

Feature/trace class

See merge request !7

examples/random_node.py

@@ -4,7 +4,7 @@ __author__ = "Jarrod Pas <j.pas@usask.ca>"
 
 from random import Random
 
-from pydtn import Network, Node, random_trace, RandomTraffic
+from pydtn import Network, Node, RandomTrace, RandomTraffic
 
 
 class RandomNode(Node):
@@ -53,7 +53,7 @@ def main():
     trace_options = {
         'seed': seed,
     }
-    trace = random_trace(nodes, **trace_options)
+    trace = RandomTrace(nodes, **trace_options)
 
     network = Network(nodes, traffic=traffic, trace=trace)
     network.run(until=duration)

examples/shed.py

@@ -9,7 +9,7 @@ from multiprocessing import Pool
 
 from pydtn import Network, RandomTraffic, Node, EpidemicNode
 from pydtn.community import BubbleNode, HCBFNode, LouvainCommunity
-import pydtn.shed as shed
+from pydtn.shed import ShedTrace
 
 
 Simulation = namedtuple('Simulation', ['trace', 'node_type', 'seed'])
@@ -19,8 +19,7 @@ def run_simulation(simulation):
     """Run a simulation."""
     seed = simulation.seed
 
-    trace_metadata = shed.read_meta_file(simulation.trace)
-    trace = shed.shed_trace(simulation.trace)
+    trace = ShedTrace(simulation.trace)
 
     epoch = 7*24*60*60  # 7 days
 
@@ -31,13 +30,13 @@ def run_simulation(simulation):
     }
     nodes = {
         node_id: simulation.node_type(**node_options)
-        for node_id in range(trace_metadata['nodes'])
+        for node_id in range(trace.nodes)
     }
 
     traffic_options = {
         'seed': seed,
         'start': epoch,
-        'speed': 30 * 60,  # 1 packet every 30 mins
+        'step': 30 * 60,  # 1 packet every 30 mins
     }
     traffic = RandomTraffic(nodes, **traffic_options)
 
@@ -45,7 +44,7 @@ def run_simulation(simulation):
     network.run()
 
     stats = {
-        'trace': simulation.trace,
+        'trace': trace.path,
         'node_type': node_type.__name__,
         'seed': seed,
     }

pydtn/__init__.py

@@ -19,8 +19,9 @@ __all__ = [
     'FloodingNode',
 
     'Contact',
-    'csv_trace',
-    'random_trace',
+    'Trace',
+    'CSVTrace',
+    'RandomTrace',
 
     'Traffic',
     'RandomManyToManyTraffic',
@@ -540,31 +541,87 @@ class FloodingNode(Node):
 Contact = namedtuple('Contact', ['time', 'a', 'b', 'join'])
 
 
-def random_trace(nodes, seed=None, step=1):
-    """Generate a random contact trace."""
-    random = Random(seed)
+class Trace:
+    """Contact trace base class."""
 
-    if isinstance(nodes, dict):
-        nodes = list(nodes)
+    def __init__(self, nodes):
+        """Create a contact trace."""
+        self._nodes = nodes
 
-    for now in count(step=step):
-        join = bool(random.getrandbits(1))
-        node_a, node_b = random.sample(nodes, 2)
-        yield Contact(now, node_a, node_b, join)
+    @property
+    def nodes(self):
+        """Return number of nodes in trace."""
+        return self._nodes
+
+    def __iter__(self):
+        """Yield no contacts."""
+        yield
+
+    @staticmethod
+    def create_contact(now, node_a, node_b, join):
+        """Create a contact."""
+        return Contact(time=now, a=node_a, b=node_b, join=join)
+
+
+class CSVTrace(Trace):
+    """
+    Contact trace from csv file.
+
+    The first row is used as headers.
 
+    The second row is the following information about the trace:
+    [ int(duration), int(node_count), ignored, ignored ]
+
+    The csv uses the following columns:
+    [ int(time), int(node_a), int(node_b), int(join) ]
+    time -- simulation time that the contact even happens
+    node_{a,b} -- nodes involved in contact
+    join -- whether the contact is going up (1) or going down (0)
+    """
+
+    def __init__(self, path):
+        """Create a csv trace generator."""
+        with open(path) as csv_file:
+            csv_file = csv.reader(csv_file)
+            # skip header
+            next(csv_file)
+            # get node count from first row of csv. it is additional stats
+            # about the trace
+            # duration, placeholder, placehold, node_count
+            nodes = int(next(csv_file)[1])
+
+        super().__init__(nodes)
+        self.path = path
+
+    def __iter__(self):
+        """Yield contacts from csv file."""
+        with open(self.path) as csv_file:
+            csv_file = csv.reader(csv_file)
+            # skip header and first row
+            next(csv_file)
+            next(csv_file)
+            for row in csv_file:
+                yield self.create_contact(*map(int, row))
+
+
+class RandomTrace(Trace):
+    """Contact trace from complete randomness."""
+
+    def __init__(self, nodes, seed=None, step=1):
+        """Create a random trace generator."""
+        super().__init__(len(nodes))
+        self._node_list = list(nodes)
+        self.step = step
+        self.seed = seed
+
+    def __iter__(self):
+        """Yield random trace."""
+        random = Random(self.seed)
 
-def csv_trace(path):
-    """Generate contact trace from csv file."""
-    with open(path) as trace_file:
-        reader = csv.reader(trace_file)
-        # skip header
-        next(reader)
-        # skip first line, it is additional stats about the trace:
-        # duration, placeholder, placehold, node_count
-        next(reader)
-        for row in reader:
-            now, source, destination, join = map(int, row)
-            yield Contact(now, source, destination, join)
+        for now in count(step=self.step):
+            join = bool(random.getrandbits(1))
+            node_a, node_b = random.sample(self._node_list, 2)
+            yield self.create_contact(now, node_a, node_b, join)
 
 
 class Traffic:

pydtn/shed.py

@@ -3,7 +3,7 @@
 __all__ = [
     'write_meta_file',
     'read_meta_file',
-    'shed_trace',
+    'ShedTrace',
 ]
 
 __author__ = 'Jarrod Pas <j.pas@usask.ca>'
@@ -15,7 +15,7 @@ from collections import defaultdict
 from itertools import groupby, count
 from os import path
 
-from pydtn import Contact
+from pydtn import Trace
 
 
 def write_meta_file(meta_path, csv_path, duty_cycle_length=300):
@@ -56,53 +56,88 @@ def read_meta_file(meta_path):
     raise RuntimeError('Should not get here...')
 
 
-def _get_contact_pairs(csv_path):
-    pairs = defaultdict(set)
-    with open(csv_path) as csv_file:
-        csv_file = csv.reader(csv_file)
-        next(csv_file)
-        for row in csv_file:
-            _, source, _, target, _, slot = row
-            pair = min(source, target), max(source, target)
-            slot = int(slot)
-            pairs[pair].add(slot)
-    return dict(pairs)
-
-
-def shed_trace(meta_path):
-    """
-    Generate contact trace for a duty cycle based SHED dataset.
-
-    Keyword Arguments:
-    duty_cycle_length -- duration of each duty cycle (default 300)
-    """
-    meta = read_meta_file(meta_path)
-    pairs = _get_contact_pairs(meta['data'])
-    node = count()
-    nodes = {}
-    contacts = []
-    for (source, target), slots in pairs.items():
-        # get canonical node id for source
-        if source not in nodes:
-            nodes[source] = next(node)
-        source = nodes[source]
-
-        # get canonical node id for source
-        if target not in nodes:
-            nodes[target] = next(node)
-        target = nodes[target]
-
-        slots = sorted(slots)
-
-        # groups consecutive slots
-        # if the lambda is mapped it will return:
-        # [1, 2, 3, 6, 7, 9] -> [-1, -1, -1, -3, -3, -4]
-        for _, group in groupby(enumerate(slots), lambda p: p[0]-p[1]):
-            times = list(map(lambda g: g[1], group))
-            start = times[0] * meta['duty_cycle_length']
-            end = (times[-1] + 1) * meta['duty_cycle_length']
-
-            contacts.append(Contact(start, source, target, True))
-            contacts.append(Contact(end, source, target, False))
-
-    yield from sorted(contacts)
+class ShedTrace(Trace):
+    """Generator for contact traces from duty cycle based SHED datasets."""
+
+    def __init__(self, meta_path):
+        """
+        Create SHED trace generator.
+
+        Arguments:
+        meta_path -- path to a metadata file created with `write_meta_file`
+        """
+        self.path = meta_path
+        self.meta = read_meta_file(meta_path)
+
+        super().__init__(self.meta['nodes'])
+
+        self._pairs = None
+        self._contacts = None
+
+    @property
+    def contact_pairs(self):
+        """Return all pairs and their duty cycles that they are in contact."""
+        if self._pairs is not None:
+            return self._pairs
+
+        pairs = defaultdict(set)
+        with open(self.meta['data']) as csv_file:
+            csv_file = csv.reader(csv_file)
+            next(csv_file)
+            for row in csv_file:
+                _, source, _, target, _, slot = row
+                pair = min(source, target), max(source, target)
+                slot = int(slot)
+                pairs[pair].add(slot)
+
+        self._pairs = dict(pairs)
+        return self._pairs
+
+    @property
+    def contacts(self):
+        """
+        Calculate contacts in the SHED dataset.
+
+        Computes contacts by finding consecutive duty cycles that any give pair
+        of nodes saw each other. The end of each contact is at the start of
+        the duty cycle that they did not see each other.
+        """
+        if self._contacts is not None:
+            return self._contacts
+
+        node = count()
+        nodes = {}
+        contacts = []
+        for (source, target), slots in self.contact_pairs.items():
+            # get canonical node id for source
+            if source not in nodes:
+                nodes[source] = next(node)
+            source = nodes[source]
+
+            # get canonical node id for source
+            if target not in nodes:
+                nodes[target] = next(node)
+            target = nodes[target]
+
+            slots = sorted(slots)
+
+            # groups consecutive slots
+            # if the lambda is mapped it will return:
+            # [1, 2, 3, 6, 7, 9] -> [-1, -1, -1, -3, -3, -4]
+            for _, group in groupby(enumerate(slots), lambda p: p[0]-p[1]):
+                times = list(map(lambda g: g[1], group))
+                start = times[0] * self.meta['duty_cycle_length']
+                end = (times[-1] + 1) * self.meta['duty_cycle_length']
+
+                contacts.extend([
+                    self.create_contact(start, source, target, True),
+                    self.create_contact(end, source, target, False),
+                ])
+
+        contacts.sort()
+        self._contacts = contacts
+        return self._contacts
+
+    def __iter__(self):
+        """Yield contacts in SHED dataset."""
+        yield from self.contacts