pythonmatplotlibcoordinatesarrowsqubit

How to draw sphere with arrow pointing from sphere in matplotlib


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

enter image description here

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

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


Solution

  • 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:

    enter image description here