How to draw sphere with arrow pointing from sphere in matplotlib

Question

I try to draw a sphere to represent a qubit using matplotlib

import numpy as np 
import matplotlib.pyplot as plt
theta = [0, np.pi]
phi = [0, 2* np.pi]
N=100
theta_array = np.linspace(theta[0], theta[1], N)
phi_array = np.linspace(phi[0], phi[1], N)
theta_grid, phi_grid = np.meshgrid(theta_array,phi_array)

x = np.sin(theta_grid) * np.cos(phi_grid)
y = np.sin(theta_grid) * np.sin(phi_grid)
z = np.cos(theta_grid)

fig = plt.figure(figsize=(6,6))
ax = fig.gca(projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, shade=False,linewidth=0)
plt.show()

I want to add tube arrows on the sphere with directions parallel with xyz axis, like this:

I am not an expert in matplotlib, so it's seem pretty tough to me. Can anyone help me? thanks in advance!

Answer 1

You can use the quiver function to to what you want. See code below:

import numpy as np 
import matplotlib.pyplot as plt

theta = [0, np.pi]
phi = [0, 2* np.pi]
N=100
theta_array = np.linspace(theta[0], theta[1], N)
phi_array = np.linspace(phi[0], phi[1], N)
theta_grid, phi_grid = np.meshgrid(theta_array,phi_array)

x = np.sin(theta_grid) * np.cos(phi_grid)
y = np.sin(theta_grid) * np.sin(phi_grid)
z = np.cos(theta_grid)

fig = plt.figure(figsize=(6,6))
ax = fig.gca(projection='3d')
ax.view_init(azim=60)
ax.plot_surface(x, y, z, rstride=1, cstride=1, shade=False,linewidth=0)

#Set up arrows

ax.quiver(1,0,0,1,0,0,color = 'k', alpha = .8, lw = 3) #x arrow
ax.text(2.4,0,0,'Sx',fontsize=20) 

ax.quiver(0,1,0,0,1,0,color = 'k', alpha = .8, lw = 3)#y arrow
ax.text(0,2.4,0,'Sy',fontsize=20)

ax.quiver(0,0,1,0,0,1,color = 'k', alpha = .8, lw = 3)#z arrow
ax.text(-0.3,0,1.8,'Sz',fontsize=20)

plt.show()

And the output gives: