Update shed example

examples/shed.py

@@ -6,11 +6,12 @@ import sys
 from argparse import ArgumentParser
 from collections import namedtuple
 from multiprocessing import Pool
+from os import path
+from pprint import pprint
 
-from pydtn import Network, RandomTraffic, Node, EpidemicNode
+from pydtn import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
 from pydtn.community import BubbleKCliqueNode, BubbleLouvainNode
 from pydtn.community import HCBFKCliqueNode, HCBFLouvainNode
-from pydtn.shed import ShedTrace
 
 
 Simulation = namedtuple('Simulation', ['trace', 'node_type', 'seed'])
@@ -20,7 +21,9 @@ def run_simulation(simulation):
     """Run a simulation."""
     seed = simulation.seed
 
-    trace = ShedTrace(simulation.trace)
+    csv = path.join(simulation.trace, 'contact.csv')
+    metadata = path.join(simulation.trace, 'metadata.json')
+    trace = CSVTrace(csv, metadata=metadata)
 
     epoch = 7*24*60*60  # 7 days
 
@@ -39,7 +42,7 @@ def run_simulation(simulation):
     traffic_options = {
         'seed': seed,
         'start': epoch,
-        'step': 1 * 60,  # 1 packet every 1 mins
+        'step': 60 * 60,  # 1 packet every 1 mins
     }
     traffic = RandomTraffic(nodes, **traffic_options)
 
@@ -47,7 +50,7 @@ def run_simulation(simulation):
     network.run()
 
     stats = {
-        'trace': trace.path,
+        'trace': simulation.trace,
         'node_type': node_type.__name__,
         'seed': seed,
     }
@@ -59,9 +62,11 @@ def run_simulation(simulation):
 
 def main(args):
     """Run simulation for each seed in args."""
-    trace = args['shed']
+    log = pprint if args['pretty'] else print
     pool = Pool()
     simulations = []
+
+    trace = args['shed']
     node_types = [
         Node,
         EpidemicNode,
@@ -77,7 +82,7 @@ def main(args):
             simulations.append(sim)
 
     for stats in pool.imap_unordered(run_simulation, simulations):
-        print(stats)
+        log(stats)
 
 
 def parse_args(args):
@@ -85,6 +90,7 @@ def parse_args(args):
     parser = ArgumentParser()
 
     parser.add_argument('shed')
+    parser.add_argument('--pretty', action='store_true')
     parser.add_argument('--seeds', '-s',
                         metavar='SEED', type=int, nargs='+', default=[None])