shed.py

"""Example to run a batch of simlations on SHED data."""

__author__ = 'Jarrod Pas <j.pas@usask.ca>'

import sys
from argparse import ArgumentParser
from collections import namedtuple
from multiprocessing import Pool
from os import path
from pprint import pprint

<<<<<<< HEAD
from pydtnsim import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
from pydtnsim.community import BubbleKCliqueNode, BubbleLouvainNode
from pydtnsim.community import HCBFKCliqueNode, HCBFLouvainNode
=======
from pydtn import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
from pydtn.community import BubbleKCliqueNode, BubbleLouvainNode
from pydtn.community import HCBFKCliqueNode, HCBFLouvainNode
from pydtn.community import  HCBFAGKmeansNode, BubbleAGKmeansNode
>>>>>>> d960e624440fdfe57ae489eb5130565794d246d0


Simulation = namedtuple('Simulation', ['trace', 'node_type', 'seed'])


def run_simulation(simulation):
    """Run a simulation."""
    seed = simulation.seed

    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

    node_type = simulation.node_type
    node_options = {
        'tick_rate': 5 * 60,  # 5 mins
        'epoch': epoch,
        'k': 3,
        'agk':3,
        'context': {}
    }
    nodes = {
        node_id: simulation.node_type(**node_options)
        for node_id in range(trace.nodes)
    }

    traffic_options = {
        'seed': seed,
        'start': epoch,
        'step': 60 * 60,  # 1 packet every 1 mins
    }
    traffic = RandomTraffic(nodes, **traffic_options)

    network = Network(nodes, traffic=traffic, trace=trace)
    network.run()

    stats = {
        'trace': simulation.trace,
        'node_type': node_type.__name__,
        'seed': seed,
    }
    stats.update(network.stats_summary)

    # return stats because we can't pickle the network as it is a generator.
    return stats


def main(args):
    """Run simulation for each seed in args."""
    log = pprint if args['pretty'] else print
    pool = Pool()
    simulations = []

    trace = args['shed']
    node_types = [
        Node,
        EpidemicNode,
        BubbleAGKmeansNode,
        HCBFAGKmeansNode,
        BubbleLouvainNode,
        HCBFLouvainNode,
        BubbleKCliqueNode,
        HCBFKCliqueNode,
    ]

    for seed in args['seeds']:
        for node_type in node_types:
            sim = Simulation(trace=trace, node_type=node_type, seed=seed)
            simulations.append(sim)
    for stats in pool.imap_unordered(run_simulation, simulations):
        log(stats)


def parse_args(args):
    """Parse arguments."""
    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])

    args = parser.parse_args(args)
    return vars(args)


if __name__ == '__main__':
    sys.exit(main(parse_args(sys.argv[1:])))