-
Notifications
You must be signed in to change notification settings - Fork 6
/
hydrogen_wavefunction.py
212 lines (171 loc) · 7.54 KB
/
hydrogen_wavefunction.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
# --- --- --- --- --- --- --- --- ---
# Hydrogen Atom - Wavefunction and Electron Density Visualization
# Sebastian Mag | August 2023
# https://github.com/ssebastianmag/hydrogen-wavefunctions
# Modeling and visualization of hydrogen atom wavefunctions
# and electron probability density.
# --- --- --- --- --- --- --- --- ---
from scipy.constants import physical_constants
import matplotlib.pyplot as plt
import scipy.special as sp
import seaborn as sns
import numpy as np
def radial_function(n, l, r, a0):
""" Compute the normalized radial part of the wavefunction using
Laguerre polynomials and an exponential decay factor.
Args:
n (int): principal quantum number
l (int): azimuthal quantum number
r (numpy.ndarray): radial coordinate
a0 (float): scaled Bohr radius
Returns:
numpy.ndarray: wavefunction radial component
"""
laguerre = sp.genlaguerre(n - l - 1, 2 * l + 1)
p = 2 * r / (n * a0)
constant_factor = np.sqrt(
((2 / n * a0) ** 3 * (sp.factorial(n - l - 1))) /
(2 * n * (sp.factorial(n + l)))
)
return constant_factor * np.exp(-p / 2) * (p ** l) * laguerre(p)
def angular_function(m, l, theta, phi):
""" Compute the normalized angular part of the wavefunction using
Legendre polynomials and a phase-shifting exponential factor.
Args:
m (int): magnetic quantum number
l (int): azimuthal quantum number
theta (numpy.ndarray): polar angle
phi (int): azimuthal angle
Returns:
numpy.ndarray: wavefunction angular component
"""
legendre = sp.lpmv(m, l, np.cos(theta))
constant_factor = ((-1) ** m) * np.sqrt(
((2 * l + 1) * sp.factorial(l - np.abs(m))) /
(4 * np.pi * sp.factorial(l + np.abs(m)))
)
return constant_factor * legendre * np.real(np.exp(1.j * m * phi))
def compute_wavefunction(n, l, m, a0_scale_factor):
""" Compute the normalized wavefunction as a product
of its radial and angular components.
Args:
n (int): principal quantum number
l (int): azimuthal quantum number
m (int): magnetic quantum number
a0_scale_factor (float): Bohr radius scale factor
Returns:
numpy.ndarray: wavefunction
"""
# Scale Bohr radius for effective visualization
a0 = a0_scale_factor * physical_constants['Bohr radius'][0] * 1e+12
# z-x plane grid to represent electron spatial distribution
grid_extent = 480
grid_resolution = 680
z = x = np.linspace(-grid_extent, grid_extent, grid_resolution)
z, x = np.meshgrid(z, x)
# Use epsilon to avoid division by zero during angle calculations
eps = np.finfo(float).eps
# Ψnlm(r,θ,φ) = Rnl(r).Ylm(θ,φ)
psi = radial_function(
n, l, np.sqrt((x ** 2 + z ** 2)), a0
) * angular_function(
m, l, np.arctan(x / (z + eps)), 0
)
return psi
def compute_probability_density(psi):
""" Compute the probability density of a given wavefunction.
Args:
psi (numpy.ndarray): wavefunction
Returns:
numpy.ndarray: wavefunction probability density
"""
return np.abs(psi) ** 2
def plot_wf_probability_density(n, l, m, a0_scale_factor, dark_theme=False, colormap='rocket'):
""" Plot the probability density of the hydrogen
atom's wavefunction for a given quantum state (n,l,m).
Args:
n (int): principal quantum number, determines the energy level and size of the orbital
l (int): azimuthal quantum number, defines the shape of the orbital
m (int): magnetic quantum number, defines the orientation of the orbital
a0_scale_factor (float): Bohr radius scale factor
dark_theme (bool): If True, uses a dark background for the plot, defaults to False
colormap (str): Seaborn plot colormap, defaults to 'rocket'
"""
# Quantum numbers validation
if not isinstance(n, int) or n < 1:
raise ValueError('n should be an integer satisfying the condition: n >= 1')
if not isinstance(l, int) or not (0 <= l < n):
raise ValueError('l should be an integer satisfying the condition: 0 <= l < n')
if not isinstance(m, int) or not (-l <= m <= l):
raise ValueError('m should be an integer satisfying the condition: -l <= m <= l')
# Colormap validation
try:
sns.color_palette(colormap)
except ValueError:
raise ValueError(f'{colormap} is not a recognized Seaborn colormap.')
# Configure plot aesthetics using matplotlib rcParams settings
plt.rcParams['font.family'] = 'STIXGeneral'
plt.rcParams['mathtext.fontset'] = 'stix'
plt.rcParams['xtick.major.width'] = 4
plt.rcParams['ytick.major.width'] = 4
plt.rcParams['xtick.major.size'] = 15
plt.rcParams['ytick.major.size'] = 15
plt.rcParams['xtick.labelsize'] = 30
plt.rcParams['ytick.labelsize'] = 30
plt.rcParams['axes.linewidth'] = 4
fig, ax = plt.subplots(figsize=(16, 16.5))
plt.subplots_adjust(top=0.82)
plt.subplots_adjust(right=0.905)
plt.subplots_adjust(left=-0.1)
# Compute and visualize the wavefunction probability density
psi = compute_wavefunction(n, l, m, a0_scale_factor)
prob_density = compute_probability_density(psi)
# Here we transpose the array to align the calculated z-x plane with Matplotlib's y-x imshow display
im = ax.imshow(np.sqrt(prob_density).T, cmap=sns.color_palette(colormap, as_cmap=True))
cbar = plt.colorbar(im, fraction=0.046, pad=0.03)
cbar.set_ticks([])
# Apply dark theme parameters
if dark_theme:
theme = 'dt'
background_color = sorted(
sns.color_palette(colormap, n_colors=100),
key=lambda color: 0.2126 * color[0] + 0.7152 * color[1] + 0.0722 * color[2]
)[0]
plt.rcParams['text.color'] = '#dfdfdf'
title_color = '#dfdfdf'
fig.patch.set_facecolor(background_color)
cbar.outline.set_visible(False)
ax.tick_params(axis='x', colors='#c4c4c4')
ax.tick_params(axis='y', colors='#c4c4c4')
for spine in ax.spines.values():
spine.set_color('#c4c4c4')
else: # Apply light theme parameters
theme = 'lt'
plt.rcParams['text.color'] = '#000000'
title_color = '#000000'
ax.tick_params(axis='x', colors='#000000')
ax.tick_params(axis='y', colors='#000000')
ax.set_title('Hydrogen Atom - Wavefunction Electron Density', pad=130, fontsize=44, loc='left', color=title_color)
ax.text(0, 722, (
r'$|\psi_{n \ell m}(r, \theta, \varphi)|^{2} ='
r' |R_{n\ell}(r) Y_{\ell}^{m}(\theta, \varphi)|^2$'
), fontsize=36)
ax.text(30, 615, r'$({0}, {1}, {2})$'.format(n, l, m), color='#dfdfdf', fontsize=42)
ax.text(770, 140, 'Electron probability distribution', rotation='vertical', fontsize=40)
ax.text(705, 700, 'Higher\nprobability', fontsize=24)
ax.text(705, -60, 'Lower\nprobability', fontsize=24)
ax.text(775, 590, '+', fontsize=34)
ax.text(769, 82, '−', fontsize=34, rotation='vertical')
ax.invert_yaxis()
# Save and display the plot
plt.savefig(f'({n},{l},{m})[{theme}].png')
plt.show()
# - - - Example probability densities for various quantum states (n,l,m)
if __name__ == '__main__':
plot_wf_probability_density(2, 1, 1, 0.6, True)
plot_wf_probability_density(2, 1, 1, 0.6)
plot_wf_probability_density(3, 2, 1, 0.3, True)
plot_wf_probability_density(3, 2, 1, 0.3)
plot_wf_probability_density(4, 3, 0, 0.2, dark_theme=True, colormap='magma')
plot_wf_probability_density(4, 3, 0, 0.2, colormap='magma')
plot_wf_probability_density(4, 3, 1, 0.2, dark_theme=True, colormap='mako')