In [1]:
# CREATE HIERARCHICAL SPATIAL TEMPORAL GRAPH

In [35]:
import networkx as nx
import numpy as np
import logging as log
import itertools
import random
import math
import pprint as pp
import matplotlib.pyplot as plt
log.basicConfig(level=log.DEBUG, format='%(asctime)s %(levelname)s: %(message)s')

In [11]:
def generate_forest_balanced_trees(r, h, n):
    graphs = []
    roots = []
    num_nodes = num_nodes_balanced_tree(r, h)
    starting_num = 0
    for i in range(0, n):
        #log.debug("building tree with starting root: %s", starting_num)
        g = nx.balanced_tree(r, h)
        g = nx.convert_node_labels_to_integers(g, first_label=starting_num)

        #log.debug("nodes: %s", pp.pformat(g.nodes()))

        graphs.append(g)
        roots.append(starting_num)
        starting_num += num_nodes
    trees = nx.union_all(graphs)
    return (trees, roots)

def num_nodes_balanced_tree(r,h):
    """
    Returns the number of nodes in a balanced tree, with branching factor R and height H.
    """
    total = 0
    for i in range(0,h+1):
        total += r ** i
        #log.debug("total: %s", total)
    return total

In [89]:
def generate_linked_roots_forest(graph, roots, frac = 0.3):
    """
    Given a forest of trees, and the list of root vertices, link 
    a specified fraction of the roots with edges to produce a 
    single connected graph from the disparate trees in the forest.
    
    Returns the input graph, after edge addition.  Warning - need to 
    check whether the result is connected.  With low fractions, it 
    may not be connected.
    """
    new_graph = graph.copy()
    edge_pairs = []
    # take combinations of the root nodes, filtering out self-combos
    nodes = []
    nodes.append(roots)
    nodes.append(roots)
    for tup in itertools.product(*nodes):
        if tup[0] == tup[1]:
            continue
        edge_pairs.append(tup)
    
    num_links = int(math.ceil(len(edge_pairs) * frac))
    log.debug("number of edges to create between random roots: %s", num_links)
    random_pairs = random.sample(edge_pairs, num_links)
    
    # now create the edges between the chosen pairs of roots
    for pair in random_pairs:
        new_graph.add_edge(pair[0], pair[1])
        
    return new_graph

In [22]:
%matplotlib inline

In [74]:
g = generate_forest_balanced_trees(2,2,6)
gr = g[0]
roots = g[1]

In [33]:
def draw_tree(graph):
    pos=nx.graphviz_layout(graph,prog='dot')
    nx.draw(graph,pos,with_labels=True,arrows=False)
    plt.show()

In [75]:
draw_tree(gr)



In [90]:
for i in range(0, 10):
    gr2 = None
    gr2 = generate_linked_roots_forest(gr,roots,frac=0.3)
    print nx.is_connected(gr2)
print roots


True
True
True
True
True
True
True
True
True
True
[0, 7, 14, 21, 28, 35]

In [87]:
draw_tree(gr2)



In [91]:
num_centroids = len(roots)

In [92]:
# choose random spatial locations for the roots, and then lay out the children
# randomly around the root centroids

In [94]:
x_max = 100
y_max = 100

In [96]:
centroids = []
for i in xrange(0,num_centroids):
    x = int(random.uniform(0, x_max))
    y = int(random.uniform(0, y_max))
    centroids.append((x,y))
print centroids


[(49, 17), (30, 69), (15, 7), (23, 5), (54, 44), (79, 92)]

In [97]:
# assign the centroid coordinates to the roots
for root in roots:
    label = "a-" + str(x) + "-" + str(y)
    gr2.node[root]['xcoord'] = x
    gr2.node[root]['ycoord'] = y
    gr2.node[root]['label'] = label
    gr2.node[root]['level'] = 'root'

In [101]:
for t in nx.generate_gml(gr2):
    print t


graph [
  name "(empty_graph(7))_with_int_labels"
  node [
    id 0
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 1
    label "1"
  ]
  node [
    id 2
    label "2"
  ]
  node [
    id 3
    label "3"
  ]
  node [
    id 4
    label "4"
  ]
  node [
    id 5
    label "5"
  ]
  node [
    id 6
    label "6"
  ]
  node [
    id 7
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 8
    label "8"
  ]
  node [
    id 9
    label "9"
  ]
  node [
    id 10
    label "10"
  ]
  node [
    id 11
    label "11"
  ]
  node [
    id 12
    label "12"
  ]
  node [
    id 13
    label "13"
  ]
  node [
    id 14
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 15
    label "15"
  ]
  node [
    id 16
    label "16"
  ]
  node [
    id 17
    label "17"
  ]
  node [
    id 18
    label "18"
  ]
  node [
    id 19
    label "19"
  ]
  node [
    id 20
    label "20"
  ]
  node [
    id 21
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 22
    label "22"
  ]
  node [
    id 23
    label "23"
  ]
  node [
    id 24
    label "24"
  ]
  node [
    id 25
    label "25"
  ]
  node [
    id 26
    label "26"
  ]
  node [
    id 27
    label "27"
  ]
  node [
    id 28
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 29
    label "29"
  ]
  node [
    id 30
    label "30"
  ]
  node [
    id 31
    label "31"
  ]
  node [
    id 32
    label "32"
  ]
  node [
    id 33
    label "33"
  ]
  node [
    id 34
    label "34"
  ]
  node [
    id 35
    label "a-79-92"
    level "root"
    xcoord 79
    ycoord 92
  ]
  node [
    id 36
    label "36"
  ]
  node [
    id 37
    label "37"
  ]
  node [
    id 38
    label "38"
  ]
  node [
    id 39
    label "39"
  ]
  node [
    id 40
    label "40"
  ]
  node [
    id 41
    label "41"
  ]
  edge [
    source 0
    target 1
  ]
  edge [
    source 0
    target 2
  ]
  edge [
    source 0
    target 35
  ]
  edge [
    source 0
    target 14
  ]
  edge [
    source 1
    target 3
  ]
  edge [
    source 1
    target 4
  ]
  edge [
    source 2
    target 5
  ]
  edge [
    source 2
    target 6
  ]
  edge [
    source 7
    target 8
  ]
  edge [
    source 7
    target 9
  ]
  edge [
    source 7
    target 35
  ]
  edge [
    source 7
    target 28
  ]
  edge [
    source 7
    target 21
  ]
  edge [
    source 8
    target 10
  ]
  edge [
    source 8
    target 11
  ]
  edge [
    source 9
    target 12
  ]
  edge [
    source 9
    target 13
  ]
  edge [
    source 14
    target 16
  ]
  edge [
    source 14
    target 35
  ]
  edge [
    source 14
    target 21
  ]
  edge [
    source 14
    target 15
  ]
  edge [
    source 15
    target 17
  ]
  edge [
    source 15
    target 18
  ]
  edge [
    source 16
    target 19
  ]
  edge [
    source 16
    target 20
  ]
  edge [
    source 21
    target 35
  ]
  edge [
    source 21
    target 22
  ]
  edge [
    source 21
    target 23
  ]
  edge [
    source 22
    target 24
  ]
  edge [
    source 22
    target 25
  ]
  edge [
    source 23
    target 26
  ]
  edge [
    source 23
    target 27
  ]
  edge [
    source 28
    target 35
  ]
  edge [
    source 28
    target 29
  ]
  edge [
    source 28
    target 30
  ]
  edge [
    source 29
    target 32
  ]
  edge [
    source 29
    target 31
  ]
  edge [
    source 30
    target 33
  ]
  edge [
    source 30
    target 34
  ]
  edge [
    source 35
    target 36
  ]
  edge [
    source 35
    target 37
  ]
  edge [
    source 36
    target 38
  ]
  edge [
    source 36
    target 39
  ]
  edge [
    source 37
    target 40
  ]
  edge [
    source 37
    target 41
  ]
]

In [ ]: