Change a matplotlib 3D figure's frames into x,y and z arrows

Question

Can one can change the arrows of a figure into an arrow by superimposing arrows on top of the x, y and z axes to create the illusion of the axes being arrows or perhaps directly change the settings of the frames as Matplot lib framing in order to get the same outcome on a 3D plot, showing (x,y,z) with arrows?

Turning this

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# generate sample points and straight line
z = np.repeat(0, 100) 
x = np.repeat(1.0, 100) 
y = np.linspace(start=3.0, stop=6.0, num=100) 
ax.plot(x, y, z, c='red') # draw straight line
ax.view_init(45, -150) # angle to show

# set axes limits and labels
ax.set_xlabel(r"$x$"); ax.set_ylabel(r"$y$"); ax.set_zlabel(r"$z$")
ax.set_xlim(0,1.1) ;ax.set_ylim(6,3) ;ax.set_zlim(0,1.75)

#  Remove tick marks
ax.set_xticks([0,0.25,0.5,0.75,1]) ; ax.set_xticklabels(['0','1','2','4','T'])
ax.set_yticks([6.0,5.5,5,4.5,4.0,3.5,3]) ; ax.set_yticklabels(["","","","","","",""])
ax.set_zticks([1.75,1.25,0.75,0.25]) ax.set_zticklabels(['','','',''])

# change background colour to white
ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))

#plt.savefig("sample.png", type="png",dbi=400) # save image 
plt.tight_layout()
plt.show()

into something like this:

Answer 1

I don't usually use 3D graphs, and I did a lot of research to answer your question. Here's a great approach I found. I created a new Arrow 3D class and implemented it. In your code, I added the class and added arrows to the x-, y-, and z-axes. I manually shifted their positions to align them on the axes.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.proj3d import proj_transform
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

class Arrow3D(FancyArrowPatch):

    def __init__(self, x, y, z, dx, dy, dz, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._xyz = (x, y, z)
        self._dxdydz = (dx, dy, dz)

    def draw(self, renderer):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)

        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        super().draw(renderer)

def _arrow3D(ax, x, y, z, dx, dy, dz, *args, **kwargs):
    '''Add an 3d arrow to an `Axes3D` instance.'''

    arrow = Arrow3D(x, y, z, dx, dy, dz, *args, **kwargs)
    ax.add_artist(arrow)

setattr(Axes3D, 'arrow3D', _arrow3D)

fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection='3d')

# generate sample points and straight line
z = np.repeat(0, 100) 
x = np.repeat(1.0, 100) 
y = np.linspace(start=3.0, stop=6.0, num=100) 
ax.plot(x, y, z, c='red') # draw straight line
ax.view_init(45, -150) # angle to show

# set axes limits and labels
ax.set_xlabel(r"$x$"); ax.set_ylabel(r"$y$"); ax.set_zlabel(r"$z$")
ax.set_xlim(0,1.1) ;ax.set_ylim(6,3) ;ax.set_zlim(0,1.75)

#  Remove tick marks
ax.set_xticks([0,0.25,0.5,0.75,1])
ax.set_xticklabels(['0','1','2','4','T'])
ax.set_yticks([6.0,5.5,5,4.5,4.0,3.5,3])
ax.set_yticklabels(["","","","","","",""])
ax.set_zticks([1.75,1.25,0.75,0.25])
ax.set_zticklabels(['','','',''])

# change background colour to white
ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))

xlim = plt.gca().get_xlim()
ylim = plt.gca().get_ylim()
zlim = plt.gca().get_zlim()
# print(xlim,ylim,zlim)
# (0.0, 1.1) (6.0, 3.0) (0.0, 1.75)
ax.arrow3D(-0.03, ylim[0]+0.06, 0, xlim[1]+0.05, 0, 0, mutation_scale=20, arrowstyle='<|-|>',fc='k') # x axis
ax.arrow3D(-0.03, ylim[1], 0, 0, ylim[1]+0.1, 0, mutation_scale=20, arrowstyle='<|-|>', fc='k') # y axis
ax.arrow3D(-0.05, ylim[1], 0, 0, 0, zlim[1]+0.1, mutation_scale=20, arrowstyle='<|-|>', fc='k') # z axis

ax.text2D(0.05, 0.65,r'$\mathcal{Z}$', fontsize=18, ha='center', transform=ax.transAxes)
ax.text2D(0.60, -0.03,r'$\mathcal{Y}$', fontsize=18, ha='center', transform=ax.transAxes)
ax.text2D(0.95, 0.40,r'$\mathcal{X}$', fontsize=18, ha='center', transform=ax.transAxes)

plt.tick_params(axis='both', color='white')
#plt.savefig("sample.png", type="png",dbi=400) # save image 
# plt.tight_layout()
plt.show()

Answer 2

A little workaround to make it happen:

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

# copied from your code
fig = plt.Figure()
ax = plt.subplot(111, projection='3d')

z = np.repeat(0, 100) 
x = np.repeat(1.0, 100) 
y = np.linspace(start=3.0, stop=6.0, num=100)
ax.plot(x, y, z, c='red') # draw straight line
ax.view_init(45, -150) # angle to show
ax.set_xlim(0,1.1)
ax.set_ylim(6,3)
ax.set_zlim(0,1.75)

# my code starts here
xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
zmin, zmax = ax.get_zlim()

ax.quiver3D(xmin, ymin, zmin, (xmax-xmin), 0, 0, length=1, arrow_length_ratio=0.1, colors='k', linewidth=3)
ax.quiver3D(xmin, ymin, zmin, 0, (ymax-ymin), 0, length=1, arrow_length_ratio=0.1, colors='b', linewidth=3)
ax.quiver3D(xmin, ymax, zmin, 0, (ymin-ymax), 0, length=1, arrow_length_ratio=0.1, colors='b', linewidth=3)
ax.quiver3D(xmin, ymax, zmin, 0, 0, (zmax-zmin), length=1, arrow_length_ratio=0.1, colors='k', linewidth=3)

ax.quiver3D(xmax, ymin, zmin, 0, (ymax-ymin), 0, length=1, arrow_length_ratio=0, colors='k', linewidth=1, alpha=0.5)
ax.quiver3D(xmin, ymax, zmin, (xmax-xmin), 0, 0, length=1, arrow_length_ratio=0, colors='k', linewidth=1, alpha=0.5)
ax.quiver3D(xmax, ymax, zmin, 0, 0, (zmax-zmin), length=1, arrow_length_ratio=0, colors='k', linewidth=1, alpha=0.5)

ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
ax.set_zlim(zmin, zmax)

ax.set_xticks([xmax])
ax.set_yticks([ymin])
ax.set_zticks([zmax])

ax.set_xticklabels([r'$\mathcal{X}$'])
ax.set_yticklabels([r'$\mathcal{Y}$'])
ax.set_zticklabels([r'$\mathcal{Z}$'])
ax.grid(None)

for axis in [ax.w_xaxis, ax.w_yaxis, ax.w_zaxis]:
    axis.line.set_linewidth(0.01)
ax.tick_params(axis='x', colors='w', pad=-5, labelcolor='k', tick1On=False, tick2On=False)
ax.tick_params(axis='y', colors='w', pad=-5, labelcolor='k', tick1On=False, tick2On=False)
ax.tick_params(axis='z', colors='w', pad=-5, labelcolor='k', tick1On=False, tick2On=False)

The output is:

I'll try to summarize the code in points:

• The xyz labels can be set separately using the standard ax.set_xlabel('x-axis') and so on.
• The above code will also work for subplots. For example if you have ax = plt.subplot(121, projection='3d') instead of a fixed 111 position.
• The color of the y-axis can, of course, be changed to black. I purposely put blue so that you can find it in the code easier.

If anything is unclear leave comment, I'll try to explain more.