How to color the edges of a 3d graph in Mayavi?

Question

I have a 3d graph created using Mayavi and the edges have to be colored by a scalar value.

The following code creates the graph but I am not sure how to color the edges

import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from mayavi import mlab

def main(edge_color=(0.8, 0.8, 0.8), edge_size=0.02):

    t = [1, 2, 4, 4, 5, 3, 5]
    h = [2, 3, 6, 5, 6, 4, 1]


    ed_ls = [(x, y) for x, y in zip(t, h)]
    G = nx.OrderedGraph()
    G.add_edges_from(ed_ls)
    nx.draw(G)
    plt.show()

    graph_pos = nx.spring_layout(G, dim=3)

    # numpy array of x,y,z positions in sorted node order
    xyz = np.array([graph_pos[v] for v in sorted(G)])
    mlab.figure(1)
    mlab.clf()
    pts = mlab.points3d(xyz[:, 0], xyz[:, 1], xyz[:, 2])
    pts.mlab_source.dataset.lines = np.array(G.edges())
    tube = mlab.pipeline.tube(pts, tube_radius=edge_size)
    mlab.pipeline.surface(tube, color=edge_color)

    mlab.show()  # interactive window

main()

Scalar values to be used for coloring the edges

scalar = [0.1, 0.7, 0.3, 0.5, 0.9, 0.8, 0.2]

Any suggestions on how to do this will be really helpful.

I also see another problem in the 3d graph that has been created. One of the edges is not connected to a node.

EDIT: From what I understand, mlab.pipeline.surface(tube, color=edge_color)
is used to color the edge/tube .

Updated code:

def main(edge_color=(0.8, 0.2, 0.8), edge_size=0.02, graph_colormap='winter'):
    t = [1, 2, 4, 4, 5, 3, 5]
    h = [2, 3, 6, 5, 6, 4, 1]

    ed_ls = [(x, y) for x, y in zip(t, h)]
    G = nx.OrderedGraph()
    G.add_edges_from(ed_ls)
    nx.draw(G)
    plt.show()
    scalars = np.array(G.nodes())+5
    pprint(scalars)
    e_color = [(0.8, 0.2, 0.8), (0.8, 0.2, 0.8), (0.8, 0.2, 0.8),
               (0.8, 0.2, 0.8), (0.8, 0.2, 0.8), (0.8, 0.2, 0.8),
               (0.8, 0.2, 0.8)]
    graph_pos = nx.spring_layout(G, dim=3)

    # numpy array of x,y,z positions in sorted node order
    xyz = np.array([graph_pos[v] for v in sorted(G)])
    mlab.figure(1)
    mlab.clf()
    pts = mlab.points3d(xyz[:, 0], xyz[:, 1], xyz[:, 2],
                        scalars,
                        colormap=graph_colormap
                        )
    pts.mlab_source.dataset.lines = np.array(G.edges())
    tube = mlab.pipeline.tube(pts, tube_radius=edge_size)
    #mlab.pipeline.surface(tube, color=e_color)  # doesn't work
    mlab.pipeline.surface(tube, color=edge_color)  # doesn't work

    mlab.show()  # interactive window

But the problems is I am no able to assign different color for different edge/tube

Answer 1

A possible solution, not at all automated, but sufficient for a proof of concept.

import networkx as nx
import numpy as np
from mayavi import mlab

t = [1, 2, 4, 4, 5, 3, 5]
h = [2, 3, 6, 5, 6, 4, 1]
ed_ls = [(x, y) for x, y in zip(t, h)]
G = nx.OrderedGraph()
G.add_edges_from(ed_ls)
graph_pos = nx.spring_layout(G, dim=3)
xyz = np.array([graph_pos[v] for v in G])
print(xyz.shape)
mlab.points3d(xyz[:, 0], xyz[:, 1], xyz[:, 2],
              np.linspace(1, 2, xyz.shape[0]),
              colormap='winter', resolution=100, scale_factor=0.3)
smallTri = np.tile(xyz[-3:, :], (2, 1))[:4, :]
remEdges = np.vstack((xyz[-1, :], xyz[:-2, :]))
allEdges = np.vstack((smallTri, remEdges))
for i in range(allEdges.shape[0] - 1):
    mlab.plot3d(allEdges[i:i + 2, 0], allEdges[i:i + 2, 1],
                allEdges[i:i + 2, 2], color=(0.2, 1 - 0.1 * i, 0.8))
mlab.show()