matplotlib 3D中的条形图

zfycwa2u  于 2023-05-01  发布在  其他
关注(0)|答案(1)|浏览(137)

我想画一个三维条形图。我知道如何使用以下代码来实现:

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111, projection='3d')
nbins = 50
# for c, z in zip(['r', 'g', 'b', 'y'], [30, 20, 10, 0]):
ys = np.random.normal(loc=10, scale=10, size=2000)

hist, bins = np.histogram(ys, bins=nbins)
xs = (bins[:-1] + bins[1:])/2

ax.bar(xs, hist, zs=30, zdir='y', color='r', ec='r', alpha=0.8)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

plt.show()

这将呈现如下内容:https://i.stack.imgur.com/KK2If.png
然而,我的目标是使条形图遵循我作为参数给予的线。例如,参数zdir='y'使图具有其当前方向。理想情况下,我想传递一个参数,使图遵循给定的直线,例如y=2x+1。
有人能帮助我们达到预期的结果吗?

kxe2p93d

kxe2p93d1#

一种方法是使用Poly3DCollection:其思想是计算每个条形的坐标和方向,然后将其添加到图中。
可以从3D空间中的矩形开始并应用适当的变换矩阵来计算每个条的位置和取向。
如果要更改curve,则还需要更改条width

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.patches import Rectangle

################
# Generates data
################
nbins = 50
ys = np.random.normal(loc=10, scale=10, size=2000)
hist, bins = np.histogram(ys, bins=nbins)
xs = (bins[:-1] + bins[1:])/2

#################################################
# Create a single bar and a transformation matrix
#################################################

# rectangle of width=height=1, centered at x,y=0
# covering the z range [0, height]
rect = np.array([
    [-0.5, 0, 0, 1],
    [0.5, 0, 0, 1],
    [0.5, 0, 1, 1],
    [-0.5, 0, 1, 1],
])
def translate(x, y, z):
    d = np.eye(4, dtype=float)
    d[:, -1] = [x, y, z, 1]
    return d
def scale(sx, sy, sz):
    d = np.eye(4, dtype=float)
    d[np.diag_indices(4)] = [sx, sy, sz, 1]
    return d
def rotate(t):
    d = np.eye(4, dtype=float)
    d[:2, :2] = np.array([
    [np.cos(t), -np.sin(t)],
    [np.sin(t), np.cos(t)]])
    return d
def transformation_matrix(t, x, y, z, w, h):
    return translate(x, y, z) @ rotate(t) @ scale(w, 1, h)
def apply_transform(t, x, y, z, w, h):
    """Apply the transformation matrix to the rectangle"""
    verts = transformation_matrix(t, x, y, z, w, h) @ rect.T
    return verts.T

#################
# Create the plot
#################
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

curve = lambda x: 2 * x + 1
# curve = lambda x: np.sin(0.05 * x)

xstep = abs(xs[0] - xs[1])
# NOTE: chose an appropriate bar width
width = xstep * 1.5

ys = curve(xs)
# previous bar coordinates
xp = np.roll(xs, 1)
yp = np.roll(ys, 1)
xp[0] = xs[0] - xstep
yp[0] = curve(xp[0])
# compute the orientation of the bars
theta = np.arctan2((ys - yp), (xs - xp))

# customize the appearance of the bar
facecolor = "tab:red"
edgecolor = "k"
linewidth = 0
# loop to add each bar
for x, y, t, h in zip(xs, ys, theta, hist):
    verts_matrix = apply_transform(t, x, y, 0, width, h)
    x, y, z = verts_matrix[:, 0], verts_matrix[:, 1], verts_matrix[:, 2]
    verts = [list(zip(x, y, z))]
    c = Poly3DCollection(verts, facecolor=facecolor, edgecolor=edgecolor, linewidth=linewidth)
    ax.add_collection3d(c)

# eventually show a legend
ax.legend([Rectangle((0, 0), 1, 1, facecolor=facecolor, edgecolor=edgecolor, linewidth=linewidth)], ["Bar Plot"])
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

ax.set_xlim(xs.min(), xs.max())
ax.set_ylim(ys.min(), ys.max())
ax.set_zlim(0, 100)

plt.show()

编辑以解释正在发生的事情:

考虑具有4个顶点的通用矩形:左下右下右上左上为了简单起见,让我们固定宽度=高度=1。然后我们考虑一个参考系x,y,z,我们画这个矩形。顶点的坐标为:左下角(-0.5,0,0)、右下角(0.5,0,0)、右上方(0.5、0、1)和左上角(-0.注意,该矩形在x方向上以零为中心。如果我们把它移动到x=2,那么它将以该位置为中心。你可以在rect中看到上面的坐标:为什么这个变量的第四列都是1?这是一个数学技巧,能够将平移矩阵应用于顶点。
让我们来谈谈transformation matrices (wikipedia has a nice page about it)。再次考虑我们的通用矩形:我们可以缩放它,旋转它,平移它,得到一个新的矩形在我们想要的位置和方向。
因此,上面的代码为每个转换定义了一个函数translate, scale, rotate。事实证明,我们可以将多个变换矩阵相乘以获得整体变换:这就是transformation_matrix所做的,它将前面提到的变换组合成一个矩阵。
最后,我使用apply_transform将变换矩阵应用于通用矩形:这将计算具有指定尺寸(宽度、高度)的指定位置/方向上的新矩形的顶点的坐标。

相关问题