Calculate network wiring cost based on a position network using Python and NetworkX

cost.py

# -*- coding: utf-8 -*-
import networkx as nx
import numpy as np
__author__ = """\n""".join(['Erik Ziegler <erik.ziegler@ulg.ac.be>'])

__all__ = ['cost']

def cost(G, edge_key='weight', position_key='dn_position', H=None, minimum_weight=0):
	if G.is_multigraph() or G.is_directed():
		raise Exception('network cost is not implemented for ',
						'directed or multiedge graphs.')
	C, notcounted, weighted_edge_cost, sum_of_edge_weights = _compute_cost(G, edge_key, position_key, H, minimum_weight)
	return C, notcounted, weighted_edge_cost, sum_of_edge_weights


def _compute_cost(G, edge_key, position_key, H=None, minimum_weight=0):
	nodes = list(G.nodes_iter())
	if H is not None:
		del G
		G = H
		print 'Using given position network'
	N = nx.number_of_nodes(G)
	length = np.zeros((N,N))
	print 'Calculating distance between nodes'
	for node_i in range(0,N):
		for node_j in range(0,N):
			if node_j > node_i:
				x_i = float(G.node[nodes[node_i]][position_key][0])
				y_i = float(G.node[nodes[node_i]][position_key][1])
				z_i = float(G.node[nodes[node_i]][position_key][2])
				x_j = float(G.node[nodes[node_j]][position_key][0])
				y_j = float(G.node[nodes[node_j]][position_key][1])
				z_j = float(G.node[nodes[node_j]][position_key][2])
				length[node_i][node_j] = np.sqrt(np.power((x_j-x_i),2) + np.power((y_j-y_i),2) + np.power((z_j-z_i),2))
	length = length + length.T

	print 'Calculating wiring cost'
	numerator = 0.0
	notcounted = 0
	weighted_edge_cost = 0
	sum_of_edge_weights = 0
	print 'using minimum weight'
	print minimum_weight
	for u, v, data in G.edges_iter(data=True):
		if float(data[edge_key]) >= minimum_weight:
			numerator = numerator + length[u-1][v-1]
			weighted_edge_cost = weighted_edge_cost + data[edge_key]*length[u-1][v-1]
			sum_of_edge_weights = sum_of_edge_weights + data[edge_key]
			#print 'not counting edge'
			notcounted = notcounted + 1
			
	print 'not counted'
	print notcounted
	denominator = np.nansum(length)
	assert np.isfinite(numerator)
	assert np.isfinite(denominator)
	cost = numerator/denominator
	print numerator
	print denominator
	print cost
	print weighted_edge_cost
	return cost, notcounted, weighted_edge_cost, sum_of_edge_weights