pythonmatplotlibmatplotlib-3d

How do I color 3d triangular mesh using another scalar in matplotlib?


I have 3d vertex position data (x,y,z), scalar values at each vertex v, and vertex indices for each triangle tri.

If I were to draw the surface geometry itself, I could use plot_trisurf. If I were to draw a 2D projected plot with each triangle colored according to the vertex scalar, than I could use tripcolor with one of the coordinates omitted. However, it seems oddly unclear how I draw the 3D surface with each face colored according to the vertex scalar, not the z coordinate value.

How do I do this? If it is not within matplotlib's ability, is there a python tool that can do this?


Solution

  • You can manually create the faces from triangle indices and color them based on your scalar values.

    import numpy as np
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
    from matplotlib import cm
    
    # Example data
    # Vertex coordinates
    x = np.array([0, 1, 0, 1])
    y = np.array([0, 0, 1, 1])
    z = np.array([0, 0, 0, 1])
    v = np.array([0.1, 0.5, 0.9, 0.6])  # Scalar values at each vertex
    
    # Triangle definitions (vertex indices)
    tri = np.array([
        [0, 1, 2],
        [1, 2, 3]
    ])
    
    # Create triangle vertices for plotting
    verts = [list(zip(x[t], y[t], z[t])) for t in tri]
    
    # Compute average scalar for each face for coloring
    face_colors = np.mean(v[tri], axis=1)
    
    # Normalize for colormap
    norm = plt.Normalize(v.min(), v.max())
    colors = cm.viridis(norm(face_colors))
    
    # Plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    mesh = Poly3DCollection(verts, facecolors=colors, edgecolors='k', linewidths=0.5)
    ax.add_collection3d(mesh)
    
    ax.auto_scale_xyz(x, y, z)
    plt.colorbar(cm.ScalarMappable(norm=norm, cmap='viridis'), ax=ax, label='Scalar value')
    plt.show()