Updates ShedTrace to inherit from Trace class

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