import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 10
# Energy levels (relative to the continuum at 0 eV) and their positions on the y-axis
# We use a dictionary to store the energy levels for easy lookup.
energy_levels_singlet = {
'1S': -6.587, # this is shifted by 18 eV for visibility
'2S': -3.97,
'2P': -3.37,
'3S': -1.67,
'3P': -1.5,
'3D': -1.51,
'4S': -0.913,
'4P': -0.845,
'4D': -0.851,
'4F': -0.850,
'5S': -0.576,
'5P': -0.541,
'5D': -0.544,
'5F': -0.544,
}
energy_levels_triplet = {
'2S': -4.769,
'2P': -3.62,
'3S': -1.67,
'3P': -1.58,
'3D': -1.51,
'4S': -0.993,
'4P': -0.879,
'4D': -0.851,
'4F': -0.850,
'5S': -0.615,
'5P': -0.559,
'5D': -0.544,
'5F': -0.544,
}
# x-coordinates for each term (S, P, D, F) to create columns
# Values are chosen to create a visual separation similar to the Grotrian diagram.
x_positions_singlet = {
'S': 0.5,
'P': 1.5,
'D': 2.5,
'F': 3.5
}
x_positions_triplet = {
'S': 5,
'P': 6,
'D': 7,
'F': 8
}
# --- Set up the plot ---
fig, ax = plt.subplots(figsize=(7, 5))
ax.set_ylim(-8, 1)
ax.set_xlim(0, 10.5)
# Hide the x-axis ticks and labels as we are creating custom annotations
y_tick_locations = [i for i in range(-8, 1)]
y_tick_labels = ['-26','-25','-24','-5','-4','-3','-2','-1','0']
ax.set_yticks(y_tick_locations)
ax.set_yticklabels(y_tick_labels)
ax.set_ylabel('energy (eV)')
ax.set_xticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
# --- Plot the energy level lines and labels ---
line_length = 0.4
for term in energy_levels_singlet:
key = term.strip('12345')
n = int(term.strip('SPDF'))
x=x_positions_singlet[key]
j=int(x-0.5)
ax.hlines(energy_levels_singlet[term], x - line_length, x + line_length, color='black', linewidth=2)
if (n < 4):
ax.text(x, energy_levels_singlet[term] - 0.1, f'${n} ^1{key}_{j}$', ha='center', va='top')
for term in energy_levels_triplet:
key = term.strip('12345')
n = int(term.strip('SPDF'))
x=x_positions_triplet[key]
j=int(x-4)
ax.hlines(energy_levels_triplet[term], x - line_length, x + line_length, color='black', linewidth=2)
if (n < 4):
if (key == 'S'):
ax.text(x, energy_levels_triplet[term] - 0.1, f'${n} ^3{key}_{j}$', ha='center', va='top')
else:
ax.text(x, energy_levels_triplet[term] - 0.1, f'${n}^3{key}_{{{j-2},{j-1},{j}}}$', ha='center', va='top')
# Custom labels for the columns and energy values
ax.text(2, 0.5, 'singlet', ha='center', va='center', fontsize=11)
ax.text(6.5, 0.5, 'triplet', ha='center', va='center', fontsize=11)
# Custom lines and labels for n=5, 4, 3, 2
#for n, y in [(5, energy_levels['5S']), (4, energy_levels['4S']), (3, energy_levels['3S']), (2, energy_levels['2S'])]:
for n in range(2,6):
y = - 13.6 / n**2
ax.text(9.5, y, f'$n={n}$', ha='left', va='center')
ax.hlines(y, 0, 9, color='black', ls=':', linewidth=0.3)
# He+ + e- continuum
ax.hlines(0, 0, 9, color='black', linewidth=1.5)
ax.fill_between(np.linspace(0, 9, 100), 0, 1, color='lightgrey', hatch='///', alpha=0.5)
ax.text(4.5, 0.2, r'He$^+ + e^{-}$', ha='center', va='bottom', fontsize=11)
# Hydrogen and custom boxes
ax.text(9.7, 0.5, 'H', ha='center', va='center', fontsize=11)
ax.hlines(0, 9.2, 10.2, color='black', linewidth=1.5)
ax.fill_between(np.linspace(9.2, 10.2, 10), 0, 1, color='lightgrey', hatch='///', alpha=0.5)
# --- Add the fake break here ---
# Coordinates for the rectangle to hide the middle section
break_y_min = -5.6
break_y_max = -5.4
#break_x_min = -0.1 # x position of left spine
#break_x_max = 0.3 # Width of the break patch, set slightly larger than the spine
# Add a white rectangle to hide the axis spine in the break region
#rect = Rectangle((break_x_min, break_y_min), break_x_max, break_y_max - break_y_min,
# facecolor='white', edgecolor='white', zorder=3, transform=ax.transData)
#ax.add_patch(rect)
ax.spines['left'].set_visible(False)
ax.vlines(0,ax.get_ylim()[0], break_y_min, color='black', linewidth=1)
ax.vlines(0,ax.get_ylim()[1], break_y_max, color='black', linewidth=1)
# Add the break symbol "//"
ax.text(0, (break_y_min + break_y_max) / 2+0.03, r'$//$',
ha='center', va='center', fontsize=14, color='black', rotation=90)
# --- Plot the transitions (red arrows) ---
# Format: (start_x, start_y), (end_x, end_y)
transitions = [
# Singlet transitions
(x_positions_singlet['P'], energy_levels_singlet['2P'], x_positions_singlet['S'], energy_levels_singlet['1S']),
(x_positions_singlet['P'], energy_levels_singlet['3P'], x_positions_singlet['S'], energy_levels_singlet['2S']),
(x_positions_singlet['P'], energy_levels_singlet['2P'], x_positions_singlet['S'], energy_levels_singlet['2S']),
(x_positions_singlet['P'], energy_levels_singlet['4P'], x_positions_singlet['S'], energy_levels_singlet['2S']),
(x_positions_singlet['P'], energy_levels_singlet['4P'], x_positions_singlet['S'], energy_levels_singlet['3S']),
(x_positions_singlet['S'], energy_levels_singlet['3S'], x_positions_singlet['P'], energy_levels_singlet['2P']),
(x_positions_singlet['S'], energy_levels_singlet['4S'], x_positions_singlet['P'], energy_levels_singlet['2P']),
(x_positions_singlet['S'], energy_levels_singlet['4S'], x_positions_singlet['P'], energy_levels_singlet['3P']),
(x_positions_singlet['D'], energy_levels_singlet['3D'], x_positions_singlet['P'], energy_levels_singlet['2P']),
(x_positions_singlet['D'], energy_levels_singlet['4D'], x_positions_singlet['P'], energy_levels_singlet['3P']),
(x_positions_singlet['P'], energy_levels_singlet['4P'], x_positions_singlet['D'], energy_levels_singlet['3D']),
(x_positions_singlet['F'], energy_levels_singlet['4F'], x_positions_singlet['D'], energy_levels_singlet['3D']),
(x_positions_singlet['P'], energy_levels_singlet['3P'], x_positions_singlet['S'], energy_levels_singlet['1S']),
# Triplet transitions
(x_positions_triplet['P'], energy_levels_triplet['3P'], x_positions_triplet['S'], energy_levels_triplet['2S']),
(x_positions_triplet['P'], energy_levels_triplet['2P'], x_positions_triplet['S'], energy_levels_triplet['2S']),
(x_positions_triplet['P'], energy_levels_triplet['4P'], x_positions_triplet['S'], energy_levels_triplet['2S']),
(x_positions_triplet['P'], energy_levels_triplet['4P'], x_positions_triplet['S'], energy_levels_triplet['3S']),
(x_positions_triplet['S'], energy_levels_triplet['3S'], x_positions_triplet['P'], energy_levels_triplet['2P']),
(x_positions_triplet['S'], energy_levels_triplet['4S'], x_positions_triplet['P'], energy_levels_triplet['2P']),
(x_positions_triplet['S'], energy_levels_triplet['4S'], x_positions_triplet['P'], energy_levels_triplet['3P']),
(x_positions_triplet['D'], energy_levels_triplet['3D'], x_positions_triplet['P'], energy_levels_triplet['2P']),
(x_positions_triplet['D'], energy_levels_triplet['4D'], x_positions_triplet['P'], energy_levels_triplet['3P']),
(x_positions_triplet['P'], energy_levels_triplet['4P'], x_positions_triplet['D'], energy_levels_triplet['3D']),
(x_positions_triplet['F'], energy_levels_triplet['4F'], x_positions_triplet['D'], energy_levels_triplet['3D']),
]
for x1, y1, x2, y2 in transitions:
ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
arrowprops=dict(arrowstyle="->", color='red', lw=1.5), zorder=0)
plt.tight_layout()
plt.show()