import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.spatial.transform import Rotation as R # Used for rotating vectors
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 15
# --- Define Constants and Initial (Un-rotated) Vectors in the Jz basis ---
j = 3
m_arr = np.arange(-j, j + 1, 1)
# J0 is along the z-axis, with magnitude sqrt(j(j+1))
# sqrt(3 * 4)
J0 = np.array([0, 0, np.sqrt(j * (j + 1))])
# Apply a fixed Euler rotation to move the entire system to an arbitrary
# position for a more engaging 3D visualization.
def r_euler(jvec,m):
if (np.abs(m) <= j and np.linalg.norm(jvec) != 0):
return R.from_euler('xyz', [0, -np.arccos(m/np.linalg.norm(jvec))*180/np.pi, 0], degrees=True).apply(jvec)
else:
return np.array([0,0,0])
# Apply the rotation to all vectors
#L = r_euler.apply(L0)
#S = r_euler.apply(S0)
def unit(v):
"""Returns the unit vector."""
norm = np.linalg.norm(v)
if norm < 1e-10:
return np.zeros_like(v)
return v / norm
# --- Plot setup ---
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection="3d")
def draw_vector(r, v, color, label):
"""Draws a vector starting at position r with components v."""
ax.quiver(r[0], r[1], r[2], v[0], v[1], v[2],
color=color, lw=1.2, arrow_length_ratio=0.1)
# Add text label slightly offset from the vector tip
ax.text(r[0]+v[0]*1.1, r[1]+v[1]*1.1, r[2]+v[2]*1.1, label, color=color, fontsize=12)
# z axis
ax.quiver(0, 0, -4, 0, 0, 8, color='k', lw=.6, arrow_length_ratio=0.03)
ax.text(0.2, 0, 4, s=f'$z$', ha='right', color='k', fontsize=12)
# --- Precession cones ---
def plot_cone(axis, angle, height, color, alpha=0.3, resolution=40):
"""
Plots a cone whose axis is aligned with the given vector 'axis'.
The cone is centered at the origin.
"""
axis = unit(axis)
# Generate points for a cone aligned with the global z-axis
theta = np.linspace(0, 2 * np.pi, resolution)
z = np.linspace(0, height, resolution)
theta, z = np.meshgrid(theta, z)
r = z * np.tan(angle)
x = r * np.cos(theta)
y = r * np.sin(theta)
# Rotation matrix: align z-axis to the custom 'axis'
def rotation_matrix(u):
u = unit(u)
z = np.array([0, 0, 1])
v = np.cross(z, u)
c = np.dot(z, u)
# Handle parallel/anti-parallel case
if np.linalg.norm(v) < 1e-10:
if c < 0:
return np.diag([1, -1, -1]) # 180 deg flip for anti-parallel
return np.eye(3) # No rotation needed
# Rodrigues' formula for arbitrary axis rotation
vx = np.array([[0, -v[2], v[1]],
[v[2], 0, -v[0]],
[-v[1], v[0], 0]])
R = np.eye(3) + vx + vx @ vx * (1 / (1 + c))
return R
R_mat = rotation_matrix(axis)
# Rotate the cone coordinates
xyz = R_mat @ np.vstack([x.flatten(), y.flatten(), z.flatten()])
x_rot = xyz[0].reshape(x.shape)
y_rot = xyz[1].reshape(x.shape)
z_rot = xyz[2].reshape(x.shape)
ax.plot_surface(x_rot, y_rot, z_rot, color=color, alpha=alpha, linewidth=0, shade=True)
def plot_disk(axis, radius, color, alpha=0.3, resolution=40):
"""
Plots a disk in the xy-plane of a 3D axes object.
"""
# Generate polar coordinates for the disk
phi = np.linspace(0, 2 * np.pi, resolution)
r = np.linspace(0, radius, resolution)
phi, r = np.meshgrid(phi, r)
# Convert polar to Cartesian coordinates
x = r * np.cos(phi)
y = r * np.sin(phi)
z = np.zeros_like(x) # Z-coordinate is constant at 0 for the xy-plane
# Plot the surface
ax.plot_surface(x, y, z, color=color, alpha=alpha, linewidth=0, shade=True)
for m in m_arr:
J = r_euler(J0,m)
draw_vector(np.zeros(3), J, "green", r"$\vec{j}$") # J starts at origin
ax.plot([J[0],0], [J[1],0], [J[2],J[2]], linewidth=0.5, linestyle='--', color='k')
ax.text(0.2, 0, J[2]-0.1, s=f'${m} \hbar$', ha='right', color='k', fontsize=12)
angle_Jz = np.arccos(J[2] / np.linalg.norm(J))
if (m==0):
plot_disk([0, 0, 1], np.linalg.norm(J), "green", alpha=0.15)
else:
plot_cone([0, 0, 1], angle_Jz, m, "green", alpha=0.15)
max_extent = 2.5
ax.set_xlim([-max_extent, max_extent])
ax.set_ylim([-max_extent, max_extent])
ax.set_zlim([-max_extent, max_extent])
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.grid(False)
ax.set_axis_off()
ax.set_box_aspect([1, 1, 1])
# Set initial viewpoint (optional)
ax.view_init(elev=20, azim=120)
plt.tight_layout()
plt.show()
