Network lab python source code

main.py

import random, json
import pandas as pd

def get_orig_edges(min_time, max_time):
    edges = set()
    for line in open('./contacts.txt'):
        source, target, time = line.strip().split()
        source, target, time = int(source), int(target), int(time)
        if time < min_time or time > max_time:
            continue
        edges.add((source, target))
    return edges

def get_neigh_from_edges(edges):
    neigh = {}
    for source, target in edges:
        if source not in neigh:
            neigh[source] = set()
        neigh[source].add(target)
        if target not in neigh:
            neigh[target] = set()
        neigh[target].add(source)
    return neigh

def random_walk(neigh, start, trials, prob_restart=0.15):
    res = set()
    cur_node = start
    for i in range(trials):
        neighs = neigh[cur_node]
        if random.random() < prob_restart:
            cur_node = start
        else:
            next_node = random.sample(neighs, 1)[0]
            res.add((cur_node, next_node))
            cur_node = next_node
    return res

def get_layers(edges, start):
    layers = {start: 0}
    changed = True
    while changed:
        changed = False
        for source, target in edges:
            sl = layers.get(source)
            tl = layers.get(target)
            if sl is None and tl is None:
                continue
            elif sl is None and tl is not None:
                layers[source] = tl + 1
                changed = True
            elif tl is None and sl is not None:
                layers[target] = sl + 1
                changed = True
            elif sl + 1 < tl:
                layers[target] = sl + 1
                changed = True
            elif tl + 1 < sl:
                layers[source] = tl + 1
                changed = True
    return layers

def group_layers(layers):
    return [[
        node for node, layer in layers.items()
        if layer == cur_layer
    ] for cur_layer in range(max(layers.values()) + 1)]

def get_json_data(min_time, max_time, start, iter_count):
    print('Extracting edges...')
    edges = get_orig_edges(min_time, max_time)
    print('%d edges found' % len(edges))
    neigh = get_neigh_from_edges(edges)
    print('%d vertices found' % len(neigh))
    walk = random_walk(neigh, start, iter_count)
    print('%d edges sampled' % len(walk))
    walk_neigh = get_neigh_from_edges(walk)
    print('%d vertices sampled' % len(walk_neigh))
    layers = get_layers(walk, start)
    grouped_layers = group_layers(layers)
    print('%d layers found: %s' % (len(grouped_layers), repr(list(map(len, grouped_layers)))))
    return json.dumps({
        'neigh': {vertex: list(nodes) for vertex, nodes in walk_neigh.items()},
        'layers': layers,
    })