bar plot in 3d following a given line

Question

I want to draw a bar plot in 3d. I know how to do that using the following code:

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

This will render something like this: https://i.stack.imgur.com/KK2If.png

However, my goal is to make the bar plot follows a line that I give as parameter. For example here, the parameter zdir='y' makes the plot have its current direction. Ideally I want to pass a parameter that makes the plot follows a given line for example y=2x+1.

Could someone help arrive at the desired result?

Answer 1

One way to achieve that is by using Poly3DCollection: the idea is to compute the coordinates and orientation of each bar, then add it to the plot.

The position and orientation of each bar can be computed starting from a rectangle in 3D space and applying the appropriate transformation matrix.

If you are going to change the curve, you will also need to change the bar 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()

EDIT to explain what is going on:

Consider a generic rectangle with 4 vertices: bottom left, bottom right, top right, top left. For simplicity, let's fix width=height=1. Then we consider a reference system x,y,z and we draw this rectangle. The coordinates of vertices are: bottom left (-0.5, 0, 0), bottom right (0.5, 0, 0), top right (0.5, 0, 1) and top left (-0.5, 0, 1). Note that this rectangle is centered around the zero in the x direction. If we move it to x=2, then it will be centered at that location. You can see the above coordinates in rect: why does this variable has a fourth column filled with ones? That's a mathematical trick to be able to apply a translation matrix to the vertices.

Let's talk about transformation matrices (wikipedia has a nice page about it). Consider again our generic rectangle: we can scale it, rotate it and translate it to get a new rectangle in the position and orientation we want.

So, the code above defines a function for each transformation, translate, scale, rotate. Turns out that we can multiply together multiple transformation matrices to get an overall transformation: that's what transformation_matrix does, it combines the aforementioned transformations into a single matrix.

Finally, I used apply_transform to apply the transformation matrix to the generic rectangle: this will compute the coordinates of the vertices of the new rectangle, in the specified position/orientation with the specified size (width, height).