refactoring
[r15k.git] / r15k / server / graph.py
1 # -*- coding: utf-8 -*-
2
3 import random
4 import networkx as nx
5 import logging
6 import time
7
8 logger = logging.getLogger('r15k.graph')
9
10 def network(nodes, seed):
11 start = time.clock()
12 logger.info("Generating network graph with seed=%d, size=%d", seed, len(nodes))
13 numEdges = len(nodes)/10
14 edgesPerNode = {}
15
16 def add_random_edges(graph):
17 while len(graph.edges()) < (len(nodes) + numEdges - 1):
18 first = random.choice(nodes)
19 second = random.choice(nodes)
20 if first != second and (edgesPerNode[first] < 3) and (edgesPerNode[second] <3):
21 graph.add_edge(first, second)
22
23 graph = nx.Graph()
24 S, T = set(nodes), set()
25 # Randomly select a first node, and place it in T.
26 node_s = random.sample(S, 1).pop()
27 S.remove(node_s)
28 T.add(node_s)
29
30 # Create a random connected graph.
31 while S:
32 # Select random node from S, and another in T.
33 node_s, node_t = random.sample(S, 1).pop(), random.sample(T, 1).pop()
34 # Create an edge between the nodes, and move the node from S to T.
35 graph.add_edge(node_s, node_t)
36 edgesPerNode[node_s] = edgesPerNode.get(node_s, 0) + 1
37 edgesPerNode[node_t] = edgesPerNode.get(node_t, 0) + 1
38 S.remove(node_s)
39 T.add(node_s)
40
41 # Add random edges until the number of desired edges is reached.
42 add_random_edges(graph)
43 graph.positions = nx.pygraphviz_layout(graph, prog = 'neato')
44 elapsed = time.clock() - start
45 logger.info("Generating network graph took: %d (%f)", elapsed, elapsed)
46 return graph
47