plot_trisurface with custom color array

Question

I basically want to "imshow" the pdf of a three-dimensional Dirichlet distribution on its support. Function simplex below computes regular points on that support, which are stored in the array sim. The array pdf holds a scalar density for each row in sim.

First thing I thought of was to use a triangulation. However, the color argument of plot_trisurface supports only one single color for all triangles. Setting cmap colors the triangles based on the z-coordinate values (See Fig. 1). Also plot_trisurface ignores the facecolors kwarg. What I want, however, is to color the surface based on pdf.

As a workaround I found, that I could interpolated the surface as 3d scatter plot. This generally gives the desired visualization, yet I ist clearly visible that it's a scatter plot; especially on the borders. (See Fig 2.)

Is there a way to plot the projection of the pdf onto the simplex?

import itertools
import matplotlib.pyplot as plt
import numpy as np
from scipy import stats


def simplex(n_vals):
    base = np.linspace(0, 1, n_vals, endpoint=False)
    coords = np.asarray(list(itertools.product(base, repeat=3)))
    return coords[np.isclose(coords.sum(axis=-1), 1.0)]


sim = simplex(20)
pdf = stats.dirichlet([1.1, 1.5, 1.3]).pdf(sim.T)

fig1 = plt.figure()
ax1 = fig1.add_subplot(1, 2, 1, projection='3d', azim=20)
ax2 = fig1.add_subplot(1, 2, 2, projection='3d', azim=20)
ax1.plot_trisurf(x, y, z, color='k')
ax2.plot_trisurf(x, y, z, cmap='Spectral')

fig2 = plt.figure()
ax21 = fig2.add_subplot(projection='3d', azim=20)
ax21.scatter3D(*sim.T, s=50, alpha=.5, c=pdf, cmap='Spectral')

@gboffi The answer you posted is not relevant, since I don't have a rectangular grid. And thank you for your general advice: googling, yeah, awesome idea ... — MaxPowers
– MaxPowers, Commented Aug 7, 2020 at 10:15

gboffi · Accepted Answer · 2020-08-09 15:57:28Z

This figure, complete with a colorbar,

was produced by the following script — the function map_colors, defined at the end of the script, could interest the general reader.

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
from itertools import product as Π

# the distribution that we want to study
dirichlet = stats.dirichlet([1.1, 1.5, 1.3])

# generate the "mesh"
N = 30 # no. of triangles along an edge
s = np.linspace(0, 1, N+1)
x, y, z = np.array([(x,y,1-x-y) for x,y in Π(s,s) if x+y<1+1E-6]).T

# plot as usual
fig = plt.figure() 
ax = fig.add_subplot(1, 1, 1, projection='3d', azim=20) 
p3dc = ax.plot_trisurf(x, y, z)

########## change the face colors ####################
mappable = map_colors(p3dc, dirichlet.pdf, 'Spectral')
# ####################################################

# possibly add a colormap
plt.colorbar(mappable, shrink=0.67, aspect=16.7)

# we are done
plt.show()

def map_colors(p3dc, func, cmap='viridis'):
    """
Color a tri-mesh according to a function evaluated in each barycentre.

    p3dc: a Poly3DCollection, as returned e.g. by ax.plot_trisurf
    func: a single-valued function of 3 arrays: x, y, z
    cmap: a colormap NAME, as a string

    Returns a ScalarMappable that can be used to instantiate a colorbar.
    """
    
    from matplotlib.cm import ScalarMappable, get_cmap
    from matplotlib.colors import Normalize
    from numpy import array

    # reconstruct the triangles from internal data
    x, y, z, _ = p3dc._vec
    slices = p3dc._segslices
    triangles = array([array((x[s],y[s],z[s])).T for s in slices])

    # compute the barycentres for each triangle
    xb, yb, zb = triangles.mean(axis=1).T
    
    # compute the function in the barycentres
    values = func(xb, yb, zb)

    # usual stuff
    norm = Normalize()
    colors = get_cmap(cmap)(norm(values))

    # set the face colors of the Poly3DCollection
    p3dc.set_fc(colors)

    # if the caller wants a colorbar, they need this
    return ScalarMappable(cmap=cmap, norm=norm)

Aguy · Accepted Answer · 2020-08-07 12:01:43Z

3

This is a way to do so by coloring each triangle in a triangulation object with the right color. Is this what you were looking for? The only thing is that each patch has a uniform color which make the patches somewhat visible.

# Setup is the same

import itertools
import matplotlib.pyplot as plt
from pylab import get_cmap
from matplotlib.tri import Triangulation, LinearTriInterpolator
import numpy as np
from scipy import stats
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def simplex(n_vals):
    base = np.linspace(0, 1, n_vals, endpoint=False)
    coords = np.asarray(list(itertools.product(base, repeat=3)))
    return coords[np.isclose(coords.sum(axis=-1), 1.0)]

sim = simplex(20)
pdf = stats.dirichlet([1.1, 1.5, 1.3]).pdf(sim.T)

# For shorter notation we define x, y and z:

x = sim[:, 0]
y = sim[:, 1]
z = sim[:, 2]

# Creating a triangulation object and using it to extract the actual triangles. 
# Note if it is necessary that no patch will be vertical (i.e. along the z direction)

tri = Triangulation(x, y)

triangle_vertices = np.array([np.array([[x[T[0]], y[T[0]], z[T[0]]],
                                        [x[T[1]], y[T[1]], z[T[1]]], 
                                        [x[T[2]], y[T[2]], z[T[2]]]]) for T in tri.triangles])

# Finding coordinate for the midpoints of each triangle. 
# This will be used to extract the color

midpoints = np.average(triangle_vertices, axis = 1)
midx = midpoints[:, 0]
midy = midpoints[:, 1]

# Interpolating the pdf and using it with the selected cmap to produce the color RGB vector for each face. 
# Some roundoff and normalization are needed

face_color_function = LinearTriInterpolator(tri, pdf)
face_color_index = face_color_function(midx, midy)
face_color_index[face_color_index < 0] = 0
face_color_index /= np.max(pdf)

cmap = get_cmap('Spectral')

# Creating the patches and plotting

collection = Poly3DCollection(triangle_vertices, facecolors=cmap(face_color_index), edgecolors=None)

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.add_collection(collection)
plt.show()

Obviously increasing the resolution would make the plot smoother.

edited Aug 7, 2020 at 12:01

answered Aug 7, 2020 at 11:54

Aguy

8,1196 gold badges35 silver badges62 bronze badges

2 Comments

MaxPowers Over a year ago

Just for simplification you could do triangle_vertices = sim[tri.triangles], which is a one-liner equivalent to your approach.

Nicolas Rougier Over a year ago

In your Poly3dCollection, you can use antialiased=False to get rid of the artifacts.

Collectives™ on Stack Overflow

plot_trisurface with custom color array

2 Answers 2

Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related