带有自定义颜色数组的 plot_trisurface

Question

我基本上想在其支持上“显示”三维 Dirichlet 分布的 pdf。下面的函数simplex 计算该支撑上的常规点，这些点存储在数组sim 中。数组pdf 保存sim 中每一行的标量密度。

我想到的第一件事是使用三角测量。但是，plot_trisurface 的 color 参数仅支持所有三角形的一种颜色。设置cmap 根据 z 坐标值为三角形着色（参见图 1）。 plot_trisurface 也忽略了 facecolors kwarg。然而，我想要的是根据pdf 为表面着色。

作为我发现的一种解决方法，我可以将曲面内插为 3d 散点图。这通常给出了所需的可视化，但我清楚地看到它是一个散点图；特别是在边界上。 （见图2）

有没有办法将 pdf 投影到单纯形上？

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')

Answer 1

这是一种通过使用正确颜色为三角剖分对象中的每个三角形着色的方法。这是你要找的吗？唯一的问题是每个补丁都有统一的颜色，这使得补丁有些可见。

# 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()

显然增加分辨率会使情节更流畅。

Answer 2

这个图，加上colorbar，

由以下脚本生成 — 函数 map_colors，定义在脚本末尾，普通读者可能会感兴趣。

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)