import numpy as np
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from scipy.special import erfc # Import the error function complement
# --- Universal Parameters and Grid ---
a0 = 1.0
x = np.linspace(-10, 10, 300)
z = np.linspace(-10, 10, 300)
X, Z = np.meshgrid(x, z)
# --- Figure Setup ---
# 6 rows (one for each MO) and 3 columns (AO_A, AO_B, MO)
fig, axes = plt.subplots(6, 3, figsize=(7.5, 14), gridspec_kw={'width_ratios': [1, 1, 1.4]})
plt.rcParams.update({'font.size': 10}) # Adjust font size for more rows
#fig.suptitle(r'$\mathbf{H}_2^+$ Molecular Orbitals from LCAO', fontsize=16, y=0.985)
# Define separate levels for 1s and 2p to ensure good contour detail in each case
levels_1s = np.linspace(-1.0, 1.0, 140)
# Max magnitude for 2p_sigma is around 0.35, for 2p_pi (using Z component) it's similar
levels_2p = np.linspace(-0.8, 0.8, 140)
# --- Helper Function to Plot a Row ---
def plot_mo_row(row_idx, ax_row, R, phiA, phiB, phi_MO, levels_AO, levels_MO, mo_label, op_label, AO_B_plot_sign, AO_limits, MO_limits):
"""Plots a single row of AO_A, AO_B, and MO."""
# Coordinates for Nuclei
rA_coords = (-R / 2, 0)
rB_coords = (R / 2, 0)
nuclei_coords = [rA_coords, rB_coords]
# --- Column 1: Atomic Orbital A (phiA) ---
ax_row[0].contourf(X, Z, -phiA, levels=levels_AO, cmap='bwr')
ax_row[0].scatter([0], [rA_coords[1]], color='k', s=30)
ax_row[0].set_aspect('equal')
ax_row[0].set_xlim(AO_limits[0])
ax_row[0].set_ylim(AO_limits[1])
ax_row[0].set_xticks([])
ax_row[0].set_yticks([])
for spine in ax_row[0].spines.values():
spine.set_visible(False)
if row_idx == 0:
ax_row[0].set_title(r'AO ($\phi_A$)', fontsize=14)
# --- Column 2: Atomic Orbital B (phiB) ---
ax_row[1].contourf(X, Z, AO_B_plot_sign * phiB, levels=levels_AO, cmap='bwr')
ax_row[1].scatter([0], [rB_coords[1]], color='k', s=30)
ax_row[1].set_aspect('equal')
ax_row[1].set_xlim(AO_limits[0])
ax_row[1].set_ylim(AO_limits[1])
ax_row[1].set_xticks([])
ax_row[1].set_yticks([])
for spine in ax_row[1].spines.values():
spine.set_visible(False)
if row_idx == 0:
ax_row[1].set_title(r'AO ($\phi_B$)', fontsize=14)
# --- Column 3: Molecular Orbital (phi_MO) ---
ax_row[2].contourf(X, Z, -phi_MO, levels=levels_MO, cmap='bwr')
ax_row[2].scatter([n[0] for n in nuclei_coords], [n[1] for n in nuclei_coords], color='k', s=30)
ax_row[2].set_aspect('equal')
ax_row[2].set_xlim(MO_limits[0])
ax_row[2].set_ylim(MO_limits[1])
ax_row[2].set_xticks([])
ax_row[2].set_yticks([])
for spine in ax_row[2].spines.values():
spine.set_visible(False)
# Add MO label to the right
ax_row[2].text(MO_limits[0][1] + 1.0, 0, mo_label, fontsize=16, ha='left', va='center') #, bbox=dict(boxstyle="round,pad=0.5", fc="white", alpha=0.6))
if row_idx == 0:
ax_row[2].set_title(r'MO ($\Phi_{\mathrm{MO}}$)', fontsize=14)
# Add operation signs to the figure
# Calculate the normalized y position for the center of the current row
# The figure has 6 rows.
row_center_y_norm = 0.85 - (row_idx) / 6.5
# + or - sign
fig.text(0.275, row_center_y_norm,
op_label, fontsize=30, ha='center', va='center')
# = sign
fig.text(0.56, row_center_y_norm,
r'$=$', fontsize=30, ha='center', va='center')
# --- 1. $\sigma_g(1s)$: Bonding MO (R=2.0) ---
R1 = 2.0
rA1 = np.sqrt((X + R1 / 2) ** 2 + Z**2)
rB1 = np.sqrt((X - R1 / 2) ** 2 + Z**2)
rC1 = np.sqrt(X ** 2 + Z**2)
phi1sA = np.exp(-rA1 / a0)
phi1sB = np.exp(-rB1 / a0)
phi1s = np.exp(-rC1 / a0)
phi_g1s = phi1sA + phi1sB
AO_limits_1s = (-5.5, 5.5), (-5, 5)
MO_limits_1s = (-7, 7), (-5, 5)
plot_mo_row(
0, axes[0], R1, phi1s, phi1s, phi_g1s,
levels_1s, levels_1s,
r'$\sigma_{g}(1s)$', r'$+$', -1,
AO_limits_1s, MO_limits_1s
)
# --- 2. $\sigma_u^\ast(1s)$: Antibonding MO (R=2.0) ---
phi_u1s = phi1sA - phi1sB
plot_mo_row(
1, axes[1], R1, phi1s, phi1s, phi_u1s,
levels_1s, levels_1s,
r'$\sigma_{u}(1s)$', r'$+$', 1,
AO_limits_1s, MO_limits_1s
)
# --- 3. $\sigma_g(2p)$: Bonding MO (R=3.0) ---
R2 = 3.0
rA2 = np.sqrt((X + R2/2)**2 + Z**2)
rB2 = np.sqrt((X - R2/2)**2 + Z**2)
# --- FADE WEIGHTING (The "Cheat") ---
rC2 = np.sqrt(X**2 + Z**2)
fade_radius = 8.5 # The radius where the MO should be fully faded
fade_sharpness = 0.8 # Controls the steepness of the fade (2.0 is smooth)
weight = 0.5 * erfc(fade_sharpness * (rC2 - fade_radius))
# $2p\sigma$ orbital (aligned along the internuclear x-axis): $\psi \propto x \cdot e^{-r/2}$
phi2p_sigA = (X + R2/2) * np.exp(-rA2 / (2*a0))
phi2p_sigB = (X - R2/2) * np.exp(-rB2 / (2*a0))
phi2p_sig = (X) * np.exp(-rC2 / (2*a0))
phi_g2p = phi2p_sigA + phi2p_sigB
AO_limits_2p = (-10, 10), (-10, 10)
MO_limits_2p = (-14, 14), (-10, 10)
levels_2pmo = np.linspace(-np.max(np.abs(phi_g2p)), np.max(np.abs(phi_g2p)), 140)
plot_mo_row(
3, axes[3], R2, phi2p_sig*weight, phi2p_sig*weight, phi_g2p*weight,
levels_2p, levels_2pmo,
r'$\sigma_{u}(2p)$', r'$+$', -1,
AO_limits_2p, MO_limits_2p
)
# --- 4. $\sigma_u^\ast(2p)$: Antibonding MO (R=3.0) ---
phi_u2p = phi2p_sigA - phi2p_sigB
levels_2pmo = np.linspace(-np.max(np.abs(phi_u2p)), np.max(np.abs(phi_u2p)), 140)
plot_mo_row(
2, axes[2], R2, phi2p_sig*weight, phi2p_sig*weight, phi_u2p*weight,
levels_2p, levels_2pmo,
r'$\sigma_{g}(2p)$', r'$+$', 1,
AO_limits_2p, MO_limits_2p
)
# --- 5. $\pi_u(2p)$: Bonding MO (R=3.0) ---
# For pi orbitals, we'll use the 2pz type orbitals, which have lobes along the z-axis (perpendicular to internuclear axis)
# $\psi_{2p\pi} \propto z \cdot e^{-r/2}$
phi2p_piA = Z * np.exp(-rA2 / (2*a0))
phi2p_piB = Z * np.exp(-rB2 / (2*a0))
phi2p_pi = Z * np.exp(-rC2 / (2*a0))
phi_u2p_pi = phi2p_piA + phi2p_piB # This is pi_u (bonding) because combination adds constructively above/below the axis
# Note the 'u' (ungerade) for bonding pi and 'g' (gerade) for antibonding pi, which is opposite to sigma.
# This is due to the symmetry of the atomic orbitals themselves.
levels_2pmo = np.linspace(-np.max(np.abs(phi_u2p_pi)), np.max(np.abs(phi_u2p_pi)), 140)
plot_mo_row(
5, axes[5], R2, phi2p_pi*weight, phi2p_pi*weight, phi_u2p_pi*weight,
levels_2p, levels_2pmo,
r'$\pi_{u}(2p)$', r'$+$', -1, # Original code used -phiA and -phiMO, let's keep that visual consistency
AO_limits_2p, MO_limits_2p
)
# --- 6. $\pi_g^\ast(2p)$: Antibonding MO (R=3.0) ---
phi_g2p_pi = phi2p_piA - phi2p_piB # This is pi_g_star (antibonding)
levels_2pmo = np.linspace(-np.max(np.abs(phi_g2p_pi)), np.max(np.abs(phi_g2p_pi)), 140)
plot_mo_row(
4, axes[4], R2, phi2p_pi*weight, phi2p_pi*weight, phi_g2p_pi*weight,
levels_2p, levels_2pmo,
r'$\pi_{g}(2p)$', r'$+$', 1,
AO_limits_2p, MO_limits_2p
)
# --- Final Adjustments ---
# Label the columns once at the top
#fig.text(0.12, 0.965, r'$\mathbf{\psi_A}$', fontsize=14, ha='center', va='center')
#fig.text(0.40, 0.965, r'$\mathbf{\psi_B}$', fontsize=14, ha='center', va='center')
#fig.text(0.74, 0.965, r'$\mathbf{\psi_{\mathrm{MO}} = \psi_A \pm \psi_B}$', fontsize=14, ha='center', va='center')
plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()
#plt.savefig('combined_molecular_orbitals_6rows.png')