Code
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def profile(l, m, theta, phi):
"""
Calculates the absolute square of the normalized complex spherical harmonic Y_l^m(theta, phi).
Args:
l (int): The orbital quantum number.
m (int): The magnetic quantum number.
theta (numpy.ndarray): The polar angle (from 0 to pi).
phi (numpy.ndarray): The azimuthal angle (from 0 to 2*pi).
Returns:
numpy.ndarray: The value of |Y_l^m(theta, phi)|^2.
"""
if l == 1:
if m == 0:
return (3.0 / (8.0 * np.pi)) * np.sin(theta)**2
elif abs(m) == 1:
# The e^(i*phi) term squared has a magnitude of 1, so we only
# need to square the sine term and the prefactor.
return (3.0 / (16.0 * np.pi)) * (1 + np.cos(theta)**2)
else:
# Return zeros for any other (l,m) pairs not explicitly handled.
return np.zeros_like(theta)
# Set up the spherical coordinate grid
theta_phi_samples = 100
theta = np.linspace(0, np.pi, theta_phi_samples)
phi = np.linspace(0, 2 * np.pi, theta_phi_samples)
theta, phi = np.meshgrid(theta, phi)
# Define the (l, m) pairs to be plotted, avoiding redundant plots
# We only need to plot for m >= 0
lm_pairs = [
(1, 0),
(1, 1)
]
# Calculate the squared magnitude for all spherical harmonics and find the global maximum
all_R_values = []
for l, m in lm_pairs:
R = profile(l, m, theta, phi)
all_R_values.append(R)
# Find the global maximum to normalize all plots to the same scale
R_max = max(np.max(r) for r in all_R_values)
# Set up the plot with a grid of subplots (2 rows, 3 columns)
fig = plt.figure(figsize=(15, 10))
# Loop through the pairs and create a subplot for each
for i, (l, m) in enumerate(lm_pairs):
ax = fig.add_subplot(2, 3, i + 1, projection='3d')
# Set equal aspect ratio for the plot box
ax.set_box_aspect([1, 1, 1])
# Get the pre-calculated R and normalize it
R = all_R_values[i] / np.max(all_R_values[i])#R_max
if (l,m)==(0,0): R = R/1.5
# Convert to Cartesian coordinates for plotting
x = R * np.sin(theta) * np.cos(phi)
y = R * np.sin(theta) * np.sin(phi)
z = R * np.cos(theta)
# Normalize the R values for color mapping and apply alpha
norm = plt.Normalize(vmin=R.min(), vmax=R.max())
if (l,m)==(0,0): norm = plt.Normalize(vmin=R.min()-.1, vmax=R.max()+.1)
colors = plt.cm.jet(norm(R))
# Plot the surface with colors based on the radius and alpha set directly
ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=colors,
edgecolor='none', alpha=0.8)
# Manually set the axis limits to be identical across all plots
ax.set_xlim([-1, 1])
ax.set_ylim([-1, 1])
ax.set_zlim([-1, 1])
# Set titles for each subplot, using the pm symbol where appropriate
if m == 0:
title = r'$\pi$-transition' +'\n' + r'$P(\Omega) = \frac{3}{8\pi} \sin(\theta)^2$'
ax.quiver(
0, 0, 1, # Position of the vector's tail
0, 0, 0.7, # Components of the vector
color='red',
arrow_length_ratio=0.2, # Customize arrow appearance
)
ax.text( 0, 0.2, 1.3, s=r'$\vec{\varepsilon}$', color='red', fontsize=14)
else:
title = r'$\sigma^\pm$-transitions' +'\n' + r'$P(\Omega) = \frac{3}{16\pi} (1+\cos(\theta)^2)$'
# 2. Define the path of the circular arc
# Define start and end angles in degrees
start_angle = 30
end_angle = 350
# Convert angles to radians for calculation
theta = np.linspace(np.deg2rad(start_angle), np.deg2rad(end_angle), 100)
radius = 0.5
# Calculate the x, y, z coordinates of the curve
x_curve = radius * np.cos(theta)
y_curve = radius * np.sin(theta)
z_curve = np.full_like(theta, 1.4) # Keep the arc on the z=0 plane
# 3. Plot the curve
ax.plot(x_curve, y_curve, z_curve, color='red')
# 4. Add the arrowhead using ax.quiver()
# Get the last two points of the curve to define the arrow direction
x_end, y_end, z_end = x_curve[-1], y_curve[-1], z_curve[-1]
x_tail, y_tail, z_tail = x_curve[-2], y_curve[-2], z_curve[-2]
# Calculate the vector components for the arrow from the tail to the tip
u_arrow = x_end - x_tail
v_arrow = y_end - y_tail
w_arrow = z_end - z_tail
# Plot the arrowhead. The `ax.quiver()` function plots a straight arrow.
ax.quiver(
x_tail, y_tail, z_tail, # Tail of the arrow
u_arrow, v_arrow, w_arrow, # Direction components
color='red',
length=0.1, # Control the length of the arrowhead
normalize=True, # Normalize the vector to control arrow length
arrow_length_ratio=2.2 # Customize arrow appearance
)
ax.text( 0.1, 0.1, 1.3, s=r'$\vec{\varepsilon}$', color='red', fontsize=14)
ax.set_title(title, fontsize=14)
# Hide the axes to make the plot cleaner
ax.set_axis_off()
plt.tight_layout()
plt.show()