numpy fanciness

vertex_lists_to_vertex_component_coloring.py

#!/usr/bin/env python

# calling : python snap_to_comp_coloring.py <snap_components> <snap_singletons> <comp_coloring>
# in: [size] [node ID] [node ID] [...] and [singleton ID]
# out: [node ID]  [comp ID]
#      [node ID]  [comp ID]
#      ...
#      [singleton ID]  [comp ID]

# IMPORT MODULES

import numpy as np
import time
import sys
import os

# PROCESS COMMAND LINE ARGUMENTS

for i, word in enumerate(sys.argv):
	if word[0:2] == './':
		sys.argv[i] = os.getcwd() + word[1:]

snap_components_f = open(sys.argv[1],'rb')
comp_coloring_f = open(sys.argv[3],'wb')

# SNAP OUTPUT TO COMPONENT COLORING

a = time.time()

sizes = []
nodes = []
for line in snap_components_f:
	if not line.startswith('#'):
		words = line.split()
		sizes.append(int(words[0]))
		nodes.append(words[1:])

index = np.array(sizes).argsort()[::-1]
sizes = np.take(sizes,index,axis=0)
nodes = np.take(nodes,index,axis=0)

cn = 1
for comp in nodes:
	comp_coloring_f.write(color.join(comp) + '\t'+str(cn)+'\n')
	cn = cn + 1

if sys.argv[2] != 'none':
	snap_singletons_f = open(sys.argv[2],'rb')
	sn = 1
	for line in snap_singletons_f:
		comp_coloring_f.write(line.strip() + '\t' + str(sn + cn - 1) + '\n')
		sn = sn + 1

	b = time.time()

snap_components_f.close()
snap_singletons_f.close()
comp_coloring_f.close()

# OUTPUT RUNTIME STATISTICS TO LOG FILE

print "[{:s}]".format(sys.argv[0])
print "snap connected components = {:d}".format(cn - 1)
print "singletons = {:d}".format(sn - 1)
print "sizes = [",
for c in range(0,cn - 1): print "{:d}".format(sizes[c]),
for s in range(0,sn - 1): print "{:d}".format(1),
print "]"
print "time = {:.10f}".format(b - a)
print ""