Create SHED dataset tools

pydtn/shed.py

+"""pydtn module for SHED dataset specific tools."""
+
+__all__ = [
+    'write_meta_file',
+    'read_meta_file',
+    'shed_trace',
+]
+
+__author__ = 'Jarrod Pas <j.pas@usask.ca>'
+
+
+import csv
+import json
+from collections import defaultdict
+from itertools import groupby, count
+from os import path
+
+from pydtn import Contact
+
+
+def write_meta_file(meta_path, csv_path, duty_cycle_length=300):
+    """Return metadata for a data set, from the dataset."""
+    nodes = set()
+    last = -1
+    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
+            nodes.add(source)
+            nodes.add(target)
+            last = max(last, int(slot))
+
+    common = path.commonprefix([meta_path, csv_path])
+    csv_path = path.relpath(csv_path, common)
+
+    meta_data = {
+        'data': csv_path,
+        'nodes': len(nodes),
+        'duration': last * duty_cycle_length,
+        'duty_cycle_length': duty_cycle_length,
+    }
+
+    with open(meta_path, 'w') as meta_file:
+        json.dump(meta_data, meta_file, sort_keys=True, indent=2)
+        meta_file.write('\n')
+
+
+def read_meta_file(meta_path):
+    """Return metadata for a data set, from a metadata file."""
+    with open(meta_path) as meta_file:
+        return json.load(meta_file)
+    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)