real - Cómo trazar un mapa de densidad 3D en python con matplotlib

matplotlib python 3

Question

Tengo un gran conjunto de datos de (x, y, z) posiciones de proteínas y me gustaría trazar áreas de alta ocupación como mapa de calor. Lo ideal es que la salida se vea similar a la siguiente visualización volumétrica, pero no estoy seguro de cómo lograrlo con matplotlib.

Mi idea inicial fue mostrar mis posiciones como un diagrama de dispersión 3D y colorear su densidad a través de un KDE. Codifiqué esto de la siguiente manera con datos de prueba:

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

¡Esto funciona bien! Sin embargo, mi información real contiene muchos miles de puntos de datos y el cálculo del kde y el diagrama de dispersión se vuelven extremadamente lentos.

Una pequeña muestra de mis datos reales:

Mi investigación sugiere que una mejor opción es evaluar el kde gaussiano en una cuadrícula. No estoy seguro de cómo hacerlo en 3D:

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

Answer 1

Gracias a mwaskon, por sugerir la biblioteca de mayavi.

Recreé el diagrama de dispersión de densidad en mayavi de la siguiente manera:

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

Establecer el modo de escala en ''ninguno'' evita que los glifos se escalan en proporción al vector de densidad. Además de los grandes conjuntos de datos, deshabilité la representación de escenas y usé una máscara para reducir el número de puntos.

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

A continuación, para evaluar el kde gaussiano en una cuadrícula:

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

Como mejora final, aceleré la evaluación de la función de densidad de densidad llamando a la función kde en paralelo.

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