import networkx as nx
from mayavi import mlab
import numpy as np
import random as rd
from IPython.display import Image
from IPython.core.display import HTML 

#definicija funkcij
def draw(G,pos):
    draw_networkx_nodes(G,pos[:,0:2],node_size=20*pos[:,3],alpha=0.9,linewidths=1.0)
    draw_networkx_edges(G,pos[:,0:2])
    #plt.xlim(4,5)
    #plt.ylim(4,11)
    plt.show()

def draw_3d(pos,graph,size,weight):
    xyz=np.array([pos[v] for v in sorted(graph)])
    pts = mlab.points3d(pos[:,0], pos[:,1], pos[:,2],size[:],scale_factor=0.2,scale_mode='none',colormap='jet',resolution=20)#
    pts.mlab_source.dataset.lines = np.array(graph.edges())
    tube = mlab.pipeline.tube(pts,tube_radius=0.01)#weight[:]
    mlab.pipeline.surface(tube, color=(1.0,0.0,0.0))

N=10
G=nx.Graph()
size=np.zeros([N],float)
pos=[]
for i in range(N):
    G.add_node(i)
    size[i]=rd.random()*10
    pos.append([rd.randint(0,9),rd.randint(0,9),rd.randint(0,9)])

for i in range(N):
    j=rd.randint(0,9)
    while j==i:
        j=rd.randint(0,9)
    G.add_edge(i,j,weight=rd.random()*10)

edges = G.edges()
weight = [G[u][v]['weight'] for u,v in edges]

pos=np.array(pos)
draw_3d(pos,G,size,weight)
mlab.savefig("3D_graph.png",magnification=2)

Image("3D_graph.png", width=500, height=500)