import numpy as np
import matplotlib.pyplot as plt
from scipy.special import ellipk, ellipe
# --- Set up plot styling ---
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 14
# Constants
mu0 = 4 * np.pi * 1e-7
# --- FIELD CALCULATION FUNCTIONS (re-used for context) ---
def B_field_single_loop(x_point, z_point, R, current, z_offset):
# Vectorized Biot-Savart function
rho = np.abs(x_point)
z_prime = z_point - z_offset
k_squared_numerator = 4 * R * rho
k_squared_denominator = (R + rho)**2 + z_prime**2
k_squared = np.divide(k_squared_numerator, k_squared_denominator,
out=np.zeros_like(x_point),
where=k_squared_denominator != 0)
k_squared = np.where(k_squared < 0, 0, k_squared)
k_squared = np.where(k_squared > 1, 1.0 - 1e-9, k_squared)
K_val = ellipk(k_squared)
E_val = ellipe(k_squared)
term1 = (mu0 * current) / (2 * np.pi)
denom_sqrt_K = np.sqrt(k_squared_denominator)
# Bx Component
m1 = (rho != 0)
Bx_numerator = z_prime * (E_val * ( (R**2 + rho**2 + z_prime**2) / ((R - rho)**2 + z_prime**2) ) - K_val)
Bx_factor = term1 / denom_sqrt_K
Bx = np.zeros_like(x_point)
Bx[m1] = Bx_factor[m1] * (Bx_numerator[m1] / rho[m1])
Bx *= np.sign(x_point)
# Bz Component
Bz_factor = term1 / denom_sqrt_K
Bz_term = K_val + E_val * ( (R**2 - rho**2 - z_prime**2) / ((R - rho)**2 + z_prime**2) )
Bz = Bz_factor * Bz_term
return Bx, Bz
def B_anti_helmholtz_2D(x, z, R_coil, d_coil, current):
Bx1, Bz1 = B_field_single_loop(x, z, R_coil, current, -d_coil/2)
Bx2, Bz2 = B_field_single_loop(x, z, R_coil, -current, d_coil/2)
return Bx1 + Bx2, Bz1 + Bz2
# --- Plotting Setup ---
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 3), width_ratios=[2,3])
# Coil parameters
R_coil = 1.0
d_coil = R_coil
current = 1.0
# Plotting area
xlim = (-2.5 * R_coil, 2.5 * R_coil)
zlim = (-2.5 * R_coil, 2.5 * R_coil)
ax1.set_xlim(xlim)
ax1.set_ylim(zlim)
ax1.set_aspect('equal')
# Create a grid for streamplot
num_points = 100
x_grid = np.linspace(xlim[0], xlim[1], num_points)
z_grid = np.linspace(zlim[0], zlim[1], num_points)
X, Z = np.meshgrid(x_grid, z_grid)
# Calculate the magnetic field at each grid point
Bx_field, Bz_field = B_anti_helmholtz_2D(X, Z, R_coil, d_coil, current)
# --- Define Explicit Symmetric Seed Points (Avoiding Axes) ---
# 1. Define seeds only for the first quadrant (x > 0, z > 0)
num_radial_points = 5
num_angular_points = 6
# Radial distances: Start well away from the origin (0.4)
radial_distances = np.linspace(0.4, 2.0, num_radial_points)
# Angular range: from 15 degrees (pi/12) to 75 degrees (5*pi/12), avoiding the axes (0 and 90)
angles = np.linspace(np.pi / 12, 5 * np.pi / 12, num_angular_points)
# Generate seed points in Quadrant 1
seed_x_q1 = []
seed_z_q1 = []
for r in radial_distances:
for angle in angles:
seed_x_q1.append(r * np.cos(angle))
seed_z_q1.append(r * np.sin(angle))
seed_x_q1 = np.array(seed_x_q1)
seed_z_q1 = np.array(seed_z_q1)
# 2. Reflect Q1 seeds to the other three quadrants to ensure symmetry
# Quadrant 2 (x < 0, z > 0) -> Reflection across z-axis (x -> -x)
seed_x_q2 = -seed_x_q1
seed_z_q2 = seed_z_q1
# Quadrant 3 (x < 0, z < 0) -> Reflection across origin (x -> -x, z -> -z)
seed_x_q3 = -seed_x_q1
seed_z_q3 = -seed_z_q1
# Quadrant 4 (x > 0, z < 0) -> Reflection across x-axis (z -> -z)
seed_x_q4 = seed_x_q1
seed_z_q4 = -seed_z_q1
# 3. Combine all seeds
seed_x_all = np.concatenate([seed_x_q1, seed_x_q2, seed_x_q3, seed_x_q4])
seed_z_all = np.concatenate([seed_z_q1, seed_z_q2, seed_z_q3, seed_z_q4])
seed_points = np.vstack([seed_x_all, seed_z_all])
# --- Streamline Visualization with Fixed Symmetric Seed Points ---
ax1.streamplot(X, Z, Bx_field, Bz_field,
color='black',
linewidth=1,
start_points=seed_points.T,
arrowstyle='->',
arrowsize=1.5)
# --- Plot Coil Cross-sections ---
circle_radius = 0.1
# Coil 1 (z = -d_coil/2)
ax1.add_patch(plt.Circle((R_coil, -d_coil/2), circle_radius,color='lightgray', zorder=5))
ax1.text(R_coil+0.025, -d_coil/2+0.021, r'$\times$', color='black', fontsize=8, ha='center', va='center', weight='bold', zorder=6)
ax1.add_patch(plt.Circle((-R_coil, -d_coil/2), circle_radius, color='lightgray', zorder=5))
ax1.text(-R_coil+0.025, -d_coil/2+0.045, r'$\cdot$', color='black', fontsize=12, ha='center', va='center', weight='bold', zorder=6)
# Coil 2 (z = d_coil/2)
ax1.add_patch(plt.Circle((R_coil, d_coil/2), circle_radius, color='lightgray', zorder=5))
ax1.text(R_coil+0.025, d_coil/2+0.045, r'$\cdot$', color='black', fontsize=12, ha='center', va='center', weight='bold', zorder=6)
ax1.add_patch(plt.Circle((-R_coil, d_coil/2), circle_radius, color='lightgray', zorder=5))
ax1.text(-R_coil+0.025, d_coil/2+0.021, r'$\times$', color='black', fontsize=8, ha='center', va='center', weight='bold', zorder=6)
ax1.spines['top'].set_visible(False)
ax1.spines['right'].set_visible(False)
ax1.spines['bottom'].set_visible(False)
ax1.spines['left'].set_visible(False)
ax1.set_xticks([])
ax1.set_yticks([])
ax1.grid(False)
# --- Parameters for the Atomic System and MOT ---
# Unperturbed transition frequency (arbitrary units for plotting, but represents E/hbar)
omega0_level = 1.0 # Energy of the unperturbed excited state (l=1, ml=0)
# Zeeman splitting coefficient (related to Bohr magneton and g-factor)
# This determines the slope of the Zeeman-split levels.
mu_B_prime_g = 0.2 # (mu_B * g_factor * B_gradient_strength) in E/hbar units
# Laser detuning
delta0 = 0.2 # Amount of red-detuning (omega0 - omega)
laser_omega_level = omega0_level - delta0
# --- Plotting Range ---
z_min, z_max = -2.0, 2.0
z_axis = np.linspace(z_min, z_max, 200)
# --- Zeeman Level Calculations ---
# l=0 ground state (ml=0)
E_ground = 0.0 * z_axis
# l=1 excited states
# In an anti-Helmholtz field, B(z) is roughly B'z. So Zeeman splitting is proportional to z.
# Assuming B(z) = B' * z (linear field with gradient B' at origin)
# Energy split: E_ml = E_0 + ml * (mu_B * g) * B(z)
# For simplicity, let mu_B_prime_g * z represent the ml=1 splitting from ml=0 level
# The sign depends on the convention and actual field direction.
# For the plot, ml=+1 increases with z, ml=-1 decreases with z.
E_l1_ml_1 = omega0_level + mu_B_prime_g * z_axis # m_l = +1
E_l1_ml_0 = omega0_level * np.ones_like(z_axis) # m_l = 0
E_l1_ml_neg1 = omega0_level - mu_B_prime_g * z_axis # m_l = -1
# --- Plot Energy Levels ---
ax2.plot(z_axis, E_ground, 'k-', linewidth=1.5, label=r'$l=0, m_l=0$')
ax2.plot(z_axis, E_l1_ml_0, 'k-', linewidth=1.5, label=r'$l=1, m_l=0$')
ax2.plot(z_axis, E_l1_ml_1, 'k-', linewidth=1.5, label=r'$l=1, m_l=+1$')
ax2.plot(z_axis, E_l1_ml_neg1, 'k-', linewidth=1.5, label=r'$l=1, m_l=-1$')
ax2.quiver(z_min, -0.2, 4.5, 0, scale=1, scale_units='xy', angles='xy', width=0.003, headwidth=10, headlength=10, headaxislength=9)
# --- Add the fake break here ---
# Coordinates for the rectangle to hide the middle section
break_y_min = 0.5
break_y_max = 0.55
ax2.quiver(0, break_y_max, 0, 2-break_y_max, scale=1, scale_units='xy', angles='xy', width=0.003, headwidth=10, headlength=10, headaxislength=9)
ax2.quiver(0, -0.25, 0, break_y_min+0.22, scale=1, scale_units='xy', angles='xy', width=0.003, headwidth=0, headlength=0, headaxislength=0)
# Add the break symbol "//"
ax2.text(0, (break_y_min + break_y_max) / 2-0.02, r'$//$',
ha='center', va='center', fontsize=10, color='black', rotation=90)
# --- Plot Laser Frequency (Red-Detuned) ---
ax2.plot([-2,2],[laser_omega_level,laser_omega_level], color='red', linestyle='--', linewidth=1)
# --- Add Transition Arrows and Labels ---
# Arrow for omega0 (unperturbed transition at z=0)
ax2.annotate("", xy=(0.6, omega0_level), xytext=(0.6, E_ground[0]),
arrowprops=dict(arrowstyle="<->", color="black", lw=1))
ax2.text(0.7, (omega0_level + E_ground[0])/2, r'$\hbar\omega_0$', va='center', ha='left', fontsize=12)
# Arrow for omegaL
ax2.annotate("", xy=(z_min+1, laser_omega_level), xytext=(z_min+1, E_ground[0]),
arrowprops=dict(arrowstyle="<->", color="black", lw=1))
ax2.text(z_min+1.1, (laser_omega_level + E_ground[0])/2, r'$\hbar\omega_L$', va='center', ha='left', fontsize=12)
# --- Add Labels for Energy Levels and m_l states ---
# Ground state label
ax2.text(z_max + 0.6, -0.25, r'$z, B_z$', va='center', ha='left', fontsize=12)
ax2.text(0.1, 2.05, r'$E$', va='center', ha='left', fontsize=12)
ax2.text(z_max + 0.1, E_l1_ml_1[-1]+0.3, r'$m_\ell$', va='center', ha='left', fontsize=12)
ax2.text(z_max + 0.1, E_ground[0], r'$0$', va='center', ha='left', fontsize=12)
ax2.text(z_max + 0.1, E_l1_ml_0[-1], r'$0$', va='center', ha='left', fontsize=12)
ax2.text(z_max + 0.1, E_l1_ml_1[-1], r'$1$', va='center', ha='left', fontsize=12)
ax2.text(z_max + 0.1, E_l1_ml_neg1[-1], r'$-1$', va='center', ha='left', fontsize=12)
ax2.text(z_min - 0.3, E_ground[0]+0.125, r'$|g\rangle$', va='center', ha='center', fontsize=12)
ax2.text(z_min - 0.3, E_l1_ml_0[-1]+0.125, r'$|e\rangle$', va='center', ha='center', fontsize=12)
ax2.text(z_min - 0.3, E_ground[0]-0.125, r'$\ell=0$', va='center', ha='center', fontsize=12)
ax2.text(z_min - 0.3, E_l1_ml_0[-1]-0.125, r'$\ell=1$', va='center', ha='center', fontsize=12)
# Photon arrows
# --- Parameters for the Sine Wave Arrow ---
start_x = z_min
end_x = z_min+0.6
mid_y = 0.35 # Central y-position for the wave
amplitude = 0.05 # Height of the wave
frequency = 10 # Number of oscillations along the arrow's length
num_points = 100 # Number of points to define the wave (smoother is better)
# 1. Generate the sine wave path
x_path = np.linspace(start_x, end_x, num_points)
y_path = mid_y + amplitude * np.sin(frequency * np.pi * (x_path - start_x) / (end_x - start_x))
y_path[-20:] = y_path[-21]
# 2. Plot the sine wave line
ax2.plot(x_path, y_path, color='red', linewidth=1)
# 3. Add an arrowhead at the end of the sine wave
# We need the last point and a point just before it to calculate the tangent for the arrow direction.
arrow_end_x = x_path[-1]
arrow_end_y = y_path[-1]
# Get a point slightly before the end for tangent calculation
# Using the second to last point
arrow_base_x = x_path[-2]
arrow_base_y = y_path[-2]
# Calculate the direction vector for the arrowhead
dx = arrow_end_x - arrow_base_x
dy = arrow_end_y - arrow_base_y
# Use ax.annotate for the arrowhead, with xytext very close to xy
# The arrow will point from xytext to xy
ax2.annotate("",
xy=(arrow_end_x, arrow_end_y),
xytext=(arrow_end_x - dx*0.1, arrow_end_y - dy*0.1), # A short segment just before the end
arrowprops=dict(arrowstyle="simple, head_length=0.2, head_width=0.15", # <-- INCREASE THESE VALUES
color="red",
lw=2,
shrinkA=0,
shrinkB=0))
ax2.text((start_x+end_x)/2+0.3, mid_y+0.05, r'$\sigma^+$', color='red', va='bottom', ha='center', fontsize=12)
# Photon arrows2
# --- Parameters for the Sine Wave Arrow ---
start2_x = z_max
end2_x = z_max-0.6
# 1. Generate the sine wave path
x2_path = np.linspace(start2_x, end2_x, num_points)
y2_path = mid_y + amplitude * np.sin(frequency * np.pi * (x_path - start_x) / (end_x - start_x))
y2_path[-20:] = y_path[-21]
# 2. Plot the sine wave line
ax2.plot(x2_path, y2_path, color='red', linewidth=1)
# 3. Add an arrowhead at the end of the sine wave
# We need the last point and a point just before it to calculate the tangent for the arrow direction.
arrow2_end_x = x2_path[-1]
arrow2_end_y = y2_path[-1]
# Get a point slightly before the end for tangent calculation
# Using the second to last point
arrow2_base_x = x2_path[-2]
arrow2_base_y = y2_path[-2]
# Calculate the direction vector for the arrowhead
dx2 = arrow2_end_x - arrow2_base_x
dy2 = arrow2_end_y - arrow2_base_y
# Use ax.annotate for the arrowhead, with xytext very close to xy
# The arrow will point from xytext to xy
ax2.annotate("",
xy=(arrow2_end_x, arrow2_end_y),
xytext=(arrow2_end_x - dx2*0.1, arrow2_end_y - dy2*0.1), # A short segment just before the end
arrowprops=dict(arrowstyle="simple, head_length=0.2, head_width=0.15", # <-- INCREASE THESE VALUES
color="red",
lw=2,
shrinkA=0,
shrinkB=0))
ax2.text((start2_x+end2_x)/2-0.2, mid_y+0.05, r'$\sigma^-$', color='red', va='bottom', ha='center', fontsize=12)
# General settings
ax2.set_xlim([-2.5,2.5])
ax2.set_ylim([-0.5,omega0_level+1])
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(False)
ax2.spines['bottom'].set_visible(False)
ax2.spines['left'].set_visible(False)
ax2.set_xticks([])
ax2.set_yticks([])
plt.show()