plot_trisurface с настраиваемым массивом цветов

Question

Я в основном хочу "показать" pdf трехмерного распределения Дирихле на его поддержке. Функция simplex ниже вычисляет обычные точки на этой опоре, которые хранятся в массиве sim. Массив pdf содержит скалярную плотность для каждой строки в sim.

Первое, о чем я подумал, - это использовать триангуляцию. Однако аргумент color для plot_trisurface поддерживает только один единственный цвет для всех треугольников. Установка cmap раскрашивает треугольники на основе значений z-координаты (см. Рис. 1). Также plot_trisurface игнорирует facecolors kwarg. Однако я хочу раскрасить поверхность на основе 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.)

Поверхность, интерполированная диаграммой рассеяния

Есть ли способ построить проекцию 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

Эта цифра, в комплекте с 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)

Answer 2

Это способ сделать это, раскрасив каждый треугольник в объекте триангуляции правильным цветом. Это то, что ты искал? Единственное, что каждый патч имеет однородный цвет, что делает патчи несколько заметными.

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

введите описание изображения здесь

Очевидно, что увеличение разрешения сделает график более плавным.