如何使用matplotlib在python中绘制3D密度图

Question

我有一个包含(x,y,z)蛋白质位置的大型数据集，想要绘制高密度区域的热力图。理想情况下，输出应该看起来类似于下面体积可视化的样子，但我不确定如何在matplotlib中实现这一点。

我的初始想法是将我的位置显示为3D散点图，并通过KDE对它们的密度进行颜色编码。我使用测试数据编写了以下代码：

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D 

mu, sigma = 0, 0.1 
x = np.random.normal(mu, sigma, 1000)
y = np.random.normal(mu, sigma, 1000)
z = np.random.normal(mu, sigma, 1000)

xyz = np.vstack([x,y,z])
density = stats.gaussian_kde(xyz)(xyz) 

idx = density.argsort()
x, y, z, density = x[idx], y[idx], z[idx], density[idx]

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x, y, z, c=density)
plt.show()

这很好用！但是，我的实际数据包含许多数千个数据点，计算KDE和散点图变得极其缓慢。

这是我实际数据的一个小样本：

我的研究表明，在网格上评估高斯KDE是更好的选择。只是我不确定如何在三维中进行：

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D 

mu, sigma = 0, 0.1 
x = np.random.normal(mu, sigma, 1000)
y = np.random.normal(mu, sigma, 1000)

nbins = 50

xy = np.vstack([x,y])
density = stats.gaussian_kde(xy) 

xi, yi = np.mgrid[x.min():x.max():nbins*1j, y.min():y.max():nbins*1j]
di = density(np.vstack([xi.flatten(), yi.flatten()]))

fig = plt.figure()
ax = fig.add_subplot(111)
ax.pcolormesh(xi, yi, di.reshape(xi.shape))
plt.show() 

- nv_wu

---

针对这个应用程序，我认为你最好使用Mayavi来进行三维可视化，因为它更适合于此类应用。以下是文档中的一个示例，可以帮助你入门。 - mwaskom

Answer 1

感谢mwaskon推荐mayavi库。

我按如下方式在mayavi中重新创建了密度散点图：

import numpy as np
from scipy import stats
from mayavi import mlab

mu, sigma = 0, 0.1 
x = 10*np.random.normal(mu, sigma, 5000)
y = 10*np.random.normal(mu, sigma, 5000)
z = 10*np.random.normal(mu, sigma, 5000)

xyz = np.vstack([x,y,z])
kde = stats.gaussian_kde(xyz)
density = kde(xyz)

# Plot scatter with mayavi
figure = mlab.figure('DensityPlot')
pts = mlab.points3d(x, y, z, density, scale_mode='none', scale_factor=0.07)
mlab.axes()
mlab.show()

将scale_mode设置为“none”可以防止图形按照密度向量成比例缩放。此外，对于大型数据集，我禁用了场景渲染，并使用掩码减少了点的数量。

# Plot scatter with mayavi
figure = mlab.figure('DensityPlot')
figure.scene.disable_render = True

pts = mlab.points3d(x, y, z, density, scale_mode='none', scale_factor=0.07) 
mask = pts.glyph.mask_points
mask.maximum_number_of_points = x.size
mask.on_ratio = 1
pts.glyph.mask_input_points = True

figure.scene.disable_render = False 
mlab.axes()
mlab.show()

接下来，要在网格上评估高斯核密度估计：

import numpy as np
from scipy import stats
from mayavi import mlab

mu, sigma = 0, 0.1 
x = 10*np.random.normal(mu, sigma, 5000)
y = 10*np.random.normal(mu, sigma, 5000)    
z = 10*np.random.normal(mu, sigma, 5000)

xyz = np.vstack([x,y,z])
kde = stats.gaussian_kde(xyz)

# Evaluate kde on a grid
xmin, ymin, zmin = x.min(), y.min(), z.min()
xmax, ymax, zmax = x.max(), y.max(), z.max()
xi, yi, zi = np.mgrid[xmin:xmax:30j, ymin:ymax:30j, zmin:zmax:30j]
coords = np.vstack([item.ravel() for item in [xi, yi, zi]]) 
density = kde(coords).reshape(xi.shape)

# Plot scatter with mayavi
figure = mlab.figure('DensityPlot')

grid = mlab.pipeline.scalar_field(xi, yi, zi, density)
min = density.min()
max=density.max()
mlab.pipeline.volume(grid, vmin=min, vmax=min + .5*(max-min))

mlab.axes()
mlab.show()

为了最终改进，我通过并行调用kde函数加快了核密度函数的评估。

import numpy as np
from scipy import stats
from mayavi import mlab
import multiprocessing

def calc_kde(data):
    return kde(data.T)

mu, sigma = 0, 0.1 
x = 10*np.random.normal(mu, sigma, 5000)
y = 10*np.random.normal(mu, sigma, 5000)
z = 10*np.random.normal(mu, sigma, 5000)

xyz = np.vstack([x,y,z])
kde = stats.gaussian_kde(xyz)

# Evaluate kde on a grid
xmin, ymin, zmin = x.min(), y.min(), z.min()
xmax, ymax, zmax = x.max(), y.max(), z.max()
xi, yi, zi = np.mgrid[xmin:xmax:30j, ymin:ymax:30j, zmin:zmax:30j]
coords = np.vstack([item.ravel() for item in [xi, yi, zi]]) 

# Multiprocessing
cores = multiprocessing.cpu_count()
pool = multiprocessing.Pool(processes=cores)
results = pool.map(calc_kde, np.array_split(coords.T, 2))
density = np.concatenate(results).reshape(xi.shape)

# Plot scatter with mayavi
figure = mlab.figure('DensityPlot')

grid = mlab.pipeline.scalar_field(xi, yi, zi, density)
min = density.min()
max=density.max()
mlab.pipeline.volume(grid, vmin=min, vmax=min + .5*(max-min))

mlab.axes()
mlab.show()

- nv_wu

---

太棒了！你知道如何使用不是Jet的颜色映射绘制它吗？ - crypdick

通过在调用volume时使用关键字参数color=，您可以将其更改为单一颜色。更高级的颜色映射可以通过创建ColorTransferFunction来完成，例如http://docs.enthought.com/mayavi/mayavi/auto/mlab_pipeline_other_functions.html#volume中的示例。 - coderforlife

我无法在Python 3.5.4中使其工作，有没有办法让它与matplotlib、seaborn或bokeh一起工作？ - O.rka

Mayavi在内部使用VTK C++绑定，但后者目前不支持Python 3。我发现matplotlib在等值面或任何复杂的3D绘图方面都很差。我相信可以使用bokeh对其进行重写。大多数基于JavaScript/浏览器的绘图库具有令人印象深刻的渲染能力。 - nv_wu

当您仅考虑图中的一个幻灯片时，您是否有此类图的示例？ （例如，z = 0的绘图） - Cruz