Towards the Hydrogen Atom

The Hydrogen Atom

Now that we’ve discussed the angular momentum operator, we can tackle one of the most important problems in quantum mechanics: the hydrogen atom. The hydrogen atom consists of a proton and an electron bound together by the Coulomb force. This system is particularly important because:

1. It is the simplest atomic system and can be solved analytically
2. More complex atoms can be understood as perturbations of the hydrogen atom
3. It demonstrates key quantum mechanical concepts like quantization and orbital angular momentum

Let’s approach this step by step, starting with the Schrödinger equation in spherical coordinates.

The Schrödinger Equation in Spherical Coordinates

For a spherically symmetric potential, the time-independent Schrödinger equation takes the form:

$$\begin{array}{rl}0& =\hat{H}\psi (r,\vartheta ,\phi )-E\psi (r,\vartheta ,\phi )\\ 0& =\frac{1}{r^2}\frac{\partial }{\partial r}\left(r^2\frac{\partial \psi }{\partial r}\right)+\frac{1}{r^2\sin \vartheta }\frac{\partial }{\partial \vartheta }(\sin \vartheta \frac{\partial \psi }{\partial \vartheta })\\ & +\frac{1}{r^2{\sin }^2\vartheta }\frac{{\partial }^2\psi }{\partial {\phi }^2}+\frac{2m}{{\hslash }^2}(E-V\left(r\right))\psi .\end{array}$$

This equation looks intimidating, but we can break it down into manageable pieces using a powerful mathematical technique called separation of variables.

Separation of Variables

The key insight is that the wavefunction can be written as a product of functions, each depending on only one coordinate:

$$\psi (r,\vartheta ,\phi )=R\left(r\right)\cdot \mathrm{\Theta }\left(\vartheta \right)\cdot \mathrm{\Phi }\left(\phi \right).$$

This separation ansatz reflects the spherical symmetry of the problem. Each function will describe a different aspect of the electron’s behavior: - $R(r)$ describes how the probability density varies with distance from the nucleus - $\mathrm{\Theta }(\theta )$ and $\mathrm{\Phi }(\varphi )$ together describe the angular distribution of the electron

Substituting this into the Schrödinger equation and rearranging terms:

$$\begin{array}{rl}0& =\frac{1}{r^2}\frac{\partial }{\partial r}\left(r^2\frac{\partial R}{\partial r}\right)\mathrm{\Theta }\mathrm{\Phi }+\frac{1}{r^2\sin \left(\vartheta \right)}\frac{\partial }{\partial \vartheta }(\sin \left(\vartheta \right)\frac{\partial \mathrm{\Theta }}{\partial \vartheta })R\mathrm{\Phi }\\ & +\frac{1}{r^2{\sin }^2\left(\vartheta \right)}\frac{{\partial }^2\mathrm{\Phi }}{\partial {\phi }^2}R\mathrm{\Theta }+\frac{2m}{{\hslash }^2}(E-V\left(r\right))R\mathrm{\Theta }\mathrm{\Phi },\end{array}$$

For mathematical convenience, we multiply by $r^2{\sin }^2\left(\vartheta \right)/\left[R\mathrm{\Theta }\mathrm{\Phi }\right]$:

$$\begin{array}{rl}0& =\frac{{\sin }^2\left(\vartheta \right)}{R}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R}{\mathrm{d}r}\right)+\frac{\sin \left(\vartheta \right)}{\mathrm{\Theta }}\frac{\mathrm{d}}{\mathrm{d}\vartheta }(\sin \left(\vartheta \right)\frac{\mathrm{d}\mathrm{\Theta }}{\partial \vartheta })\\ & +\frac{1}{\mathrm{\Phi }}\frac{{\mathrm{d}}^2\mathrm{\Phi }}{\mathrm{d}{\phi }^2}+\frac{2m}{{\hslash }^2}(E-V\left(r\right))r^2{\sin }^2\left(\vartheta \right).\end{array}$$

We can now separate the $\varphi$-dependent terms:

$$\begin{array}{rl}-\frac{1}{\mathrm{\Phi }}\frac{{\mathrm{d}}^2\mathrm{\Phi }}{\mathrm{d}{\phi }^2}& =\frac{{\sin }^2\left(\vartheta \right)}{R}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R}{\mathrm{d}r}\right)+\frac{\sin \left(\vartheta \right)}{\mathrm{\Theta }}\frac{\mathrm{d}}{\mathrm{d}\vartheta }(\sin \left(\vartheta \right)\frac{\mathrm{d}\mathrm{\Theta }}{\partial \vartheta })\\ & +\frac{2m}{{\hslash }^2}(E-V\left(r\right))r^2{\sin }^2\left(\vartheta \right)\end{array}$$

This equation has a remarkable property: the left side depends only on $\varphi$, while the right side depends only on $r$ and $\vartheta$. For this equality to hold for all values of the coordinates, both sides must equal a separation constant, which we’ll call $C_1$.

The Azimuthal Equation

The equation for $\mathrm{\Phi }(\phi )$ is:

$$\frac{{\mathrm{d}}^2\mathrm{\Phi }\left(\phi \right)}{\mathrm{d}{\phi }^2}=-C_1\cdot \mathrm{\Phi }\left(\phi \right)$$

This is a simple harmonic oscillator equation with solution:

$$\mathrm{\Phi }\left(\phi \right)=A\cdot {\mathrm{e}}^{\pm i\sqrt{C_1}\phi }.$$

Because the wavefunction must be single-valued when $\phi$ changes by $2\pi$:

$$\begin{array}{rl}\mathrm{\Phi }(\phi +n\cdot 2\pi )& =\mathrm{\Phi }\left(\phi \right)\\ ⟶{\mathrm{e}}^{\pm 2\pi in\sqrt{C_1}}& =1\\ ⟶\sqrt{C_1}& =m(m\in \mathbb{Z})\end{array}.$$

This gives us our first quantum number m, which represents the z-component of angular momentum. The normalized solution is:

$${\mathrm{\Phi }}_m\left(\phi \right)=\frac{1}{\sqrt{2\pi }}{\mathrm{e}}^{im\phi }$$

These functions are orthonormal:

$${\int }_{\phi =0}^{2\pi }{\mathrm{\Phi }}_m^{\ast }\left(\phi \right){\mathrm{\Phi }}_n\left(\phi \right)\mathrm{d}\phi ={\delta }_{mn}.$$

The Polar Angle Equation

For the $\vartheta$-dependent part, we get:

$$\begin{array}{rl}{C}_1& =\frac{{\sin }^2\left(\vartheta \right)}{R\left(r\right)}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R\left(r\right)}{\mathrm{d}r}\right)+\frac{\sin \left(\vartheta \right)}{\mathrm{\Theta }\left(\vartheta \right)}\frac{\mathrm{d}}{\mathrm{d}\vartheta }(\sin \left(\vartheta \right)\frac{\mathrm{d}\mathrm{\Theta }\left(\vartheta \right)}{\mathrm{d}\vartheta })+\frac{2m}{{\hslash }^2}(E-V\left(r\right))r^2{\sin }^2\left(\vartheta \right)\end{array},$$

Dividing by ${\sin }^2\left(\vartheta \right)$ and rearranging:

$$\begin{array}{rl}\frac{C_1}{{\sin }^2\left(\vartheta \right)}-\frac{1}{\sin \left(\vartheta \right)\mathrm{\Theta }\left(\vartheta \right)}\frac{\mathrm{d}}{\mathrm{d}\vartheta }(\sin \left(\vartheta \right)\frac{\mathrm{d}\mathrm{\Theta }\left(\vartheta \right)}{\mathrm{d}\vartheta })& =\frac{1}{R\left(r\right)}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R\left(r\right)}{\mathrm{d}r}\right)+\frac{2m}{{\hslash }^2}(E-V\left(r\right))r^2\end{array},$$

Again, the left side depends only on $\vartheta$ while the right side depends only on r, introducing another separation constant $C_2$. For the $\vartheta$ equation:

$$\begin{array}{rl}-C_2& =\frac{1}{\sin \left(\vartheta \right)\mathrm{\Theta }\left(\vartheta \right)}\frac{\mathrm{d}}{\mathrm{d}\vartheta }(\sin \left(\vartheta \right)\frac{\mathrm{d}\mathrm{\Theta }\left(\vartheta \right)}{\mathrm{d}\vartheta })-\frac{C_1}{{\sin }^2\left(\vartheta \right)}\end{array}.$$

When $m=0$, using $\xi =\cos (\vartheta )$, this becomes the Legendre equation:

$$\frac{\mathrm{d}}{\mathrm{d}\xi }\left[(1-{\xi }^2)\frac{\mathrm{d}\mathrm{\Theta }\left(\vartheta \right)}{\mathrm{d}\xi }\right]+C_2\mathrm{\Theta }=0.$$

The solutions are Legendre polynomials:

$$\mathrm{\Theta }=\sum _{k=0}^{\mathrm{\infty }}{a}_k{\xi }^k$$

With the recursion relation:

$$a_{k+2}=\frac{k(k+1)-C_2}{(k+2)(k+1)}{a}_k.$$

For finite solutions, we need $C_2=l(l+1)$ where $l$ is an integer $\ge 0$. This gives us our second quantum number $l$, which represents the total angular momentum.

For $m\ne 0$, we get associated Legendre polynomials:

$$P_l^m(\cos \left(\vartheta \right))=const.\cdot {(1-{\xi }^2)}^{\left|\frac{m}{2}\right|}\frac{{\mathrm{d}}^{\left|m\right|}}{\mathrm{d}{\xi }^{\left|m\right|}}{P}_l(\cos \left(\vartheta \right))$$

With the condition $|m|\le l$ and normalization:

$${\int }_{\vartheta =0}^{\pi }{\left|P_l^m(\cos \left(\vartheta \right))\right|}^2\sin \left(\vartheta \right)\mathrm{d}\vartheta =1.$$

Spherical Harmonics

The angular part of the solution combines $\mathrm{\Theta }$ and $\mathrm{\Phi }$ into spherical harmonics:

$$Y_l^m(\vartheta ,\phi )=P_l^m(\cos \left(\vartheta \right))\cdot \mathrm{\Phi }\cos \left(\phi \right)$$

These are normalized:

$${\int }_{\vartheta =0}^{\pi }{\int }_{\phi =0}^{2\pi }{\left|Y_l^m(\vartheta ,\phi )\right|}^2\sin \left(\vartheta \right)\mathrm{d}\vartheta \mathrm{d}\phi =1.$$

The first few spherical harmonics are:

$l$	$m$	$Y_l^m(\vartheta ,\phi )$
0	0	$\frac{1}{2\sqrt{\pi }}$
1	$\pm 1$	$\mp \sqrt{\frac{3}{8\pi }}\sin \left(\vartheta \right){\mathrm{e}}^{\pm i\phi }$
1	0	$\sqrt{\frac{3}{4\pi }}\cos \left(\vartheta \right)$
2	$\pm 2$	$\sqrt{\frac{15}{32\pi }}{\sin }^2\left(\vartheta \right){\mathrm{e}}^{\pm 2i\phi }$
2	$\pm 1$	$\mp \sqrt{\frac{15}{8\pi }}\cos \left(\vartheta \right)\sin \left(\vartheta \right){\mathrm{e}}^{\pm i\phi }$
2	0	$\sqrt{\frac{5}{16\pi }}(3{\cos }^2\left(\vartheta \right)-1)$

The graph below shows the spherical harmonics for $l=0,1,2,3$ and $m=0,1,2,3$.

#
def plot_spherical_harmonic(l, m, ax):
    # Create meshgrid for theta and phi
    phi = np.linspace(0, 2*np.pi, 100)
    theta = np.linspace(0, np.pi, 100)
    phi, theta = np.meshgrid(phi, theta)

    # Calculate the spherical harmonic
    Y = sph_harm(m, l, phi, theta)

    # Convert to cartesian coordinates
    r = np.abs(Y)
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)

    # Plot the surface
    surf = ax.plot_surface(x, y, z, cmap='seismic')
    ax.set_title(f'l={l}, m={m}', fontsize=6, pad=0)
    ax.set_axis_off()

# Create figure with subplots
fig = plt.figure(figsize=get_size(12, 16))

# Plot l=0
ax = fig.add_subplot(4, 4, 1, projection='3d')
plot_spherical_harmonic(0, 0, ax)

# Plot l=1
for m in range(2):
    ax = fig.add_subplot(4, 4, 5+m, projection='3d')
    plot_spherical_harmonic(1, m, ax)

# Plot l=2
for m in range(3):
    ax = fig.add_subplot(4, 4, 9+m, projection='3d')
    plot_spherical_harmonic(2, m, ax)

# Plot l=3
for m in range(4):
    ax = fig.add_subplot(4, 4, 13+m, projection='3d')
    plot_spherical_harmonic(3, m, ax)

plt.tight_layout()
plt.show()

Figure 1— Spherical harmonics for l=0,1,2,3 and m=0,1,2,3

Solving the Radial Part: $R\left(r\right)$

Having solved the angular parts using spherical harmonics, we now tackle the final piece - the radial function $R(r)$ that describes how the wavefunction varies with distance from the nucleus. From our previous separation of variables, we found that $C_2$ satisfies:

$$\begin{array}{rl}{C}_2& =l(l+1)\\ & =\frac{1}{R\left(r\right)}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R\left(r\right)}{\mathrm{d}r}\right)+\frac{2m}{{\hslash }^2}(E-V\left(r\right))r^2\end{array}.$$

Multiplying by $R(r)$ and dividing by $r^2$ gives us:

$$\begin{array}{rl}\frac{l(l+1)}{r^2}R\left(r\right)& =\frac{1}{r^2}\frac{\mathrm{d}}{\mathrm{d}r}\left(r^2\frac{\mathrm{d}R\left(r\right)}{\mathrm{d}r}\right)+\frac{2m}{{\hslash }^2}(E-V\left(r\right))R\left(r\right)\end{array},$$

Here, the quantum number $l$ determines the angular momentum $\left|\overrightarrow{L}\right|=\sqrt{l(l+1)}\cdot \hslash$. For the hydrogen atom, we use the Coulomb potential:

$$V\left(r\right)=-\frac{Z{e}^2}{4\pi {\epsilon }_{0}r}$$

Where $Z$ is the atomic number ($Z=1$ for hydrogen) and ${\epsilon }_0$ is the vacuum permittivity. We also use the reduced mass $\mu$ instead of the electron mass $m$ to account for the finite nuclear mass:

$$\frac{1}{\mu }=\frac{1}{m_e}+\frac{1}{m_n}$$

After some algebra, we get the radial equation:

$$\begin{array}{rl}0& =\frac{{\mathrm{d}}^{2}R\left(r\right)}{\mathrm{d}{r}^2}+\frac{2}{r}\frac{\mathrm{d}R\left(r\right)}{\mathrm{d}r}+(\frac{2\mu }{{\hslash }^2}(E+\frac{Z{e}^2}{4\pi {\epsilon }_{0}r})-\frac{l(l+1)}{r^2})R\left(r\right)\end{array}.$$

This is a second-order differential equation. For bound states ($E<0$), the physical nature of the problem suggests an exponential decay at large distances. We can understand this intuitively: as $r\to \mathrm{\infty }$, the potential energy approaches zero, so the Schrödinger equation approximately becomes:

$$\frac{d^{2}R}{d{r}^2}\approx \frac{-2\mu E}{{\hslash }^2}R$$

This has solutions of the form $e^{\pm \kappa r}$ where $\kappa =\sqrt{\frac{-2\mu E}{{\hslash }^2}}$. Since we require the wavefunction to be normalizable, we must choose the decaying solution. This motivates the ansatz:

$$R\left(r\right)=u\left(r\right)\cdot {\mathrm{e}}^{-\kappa r},$$

where $\kappa =\sqrt{\frac{-2\mu E}{{\hslash }^2}}$. The exponential factor ensures the wavefunction goes to zero at infinity, as required for bound states. We also introduce:

$$a=\frac{\mu Z{e}^2}{4\pi {\epsilon }_0{\hslash }^2}$$

This gives us a simpler equation for $u(r)$:

$$\begin{array}{rl}0& =\frac{{\mathrm{d}}^{2}u\left(r\right)}{\mathrm{d}{r}^2}+2(\frac{1}{r}-\kappa )\frac{\mathrm{d}u\left(r\right)}{\mathrm{d}r}+(2\frac{a-\kappa }{r}-\frac{l(l+1)}{r^2})u\left(r\right)\end{array}.$$

We can solve this using a power series:

$$u\left(r\right)=\sum _j{b}_j{r}^j.$$

Substituting this into our equation and comparing coefficients gives us a recursive formula:

$$b_j=2b_{j-1}\cdot \frac{\kappa \cdot j-a}{j(j+1)-l(l+1)}.$$

For the wavefunction to be normalizable, this series must terminate. If we let $j=n-1$ be the last non-zero term:

$$j<n.$$

This condition, combined with the recursive formula, gives us:

$$a=n\cdot \kappa$$

From this, we can find the energy levels:

$$\begin{array}{rl}E=-\frac{a^2}{2\mu n^2}& =-\frac{\mu e^4}{32{\pi }^2{\epsilon }_0^2{\hslash }^2}\frac{Z^2}{n^2}\\ & =-R{y}^{\ast }\frac{Z^2}{n^2}=E_n.\end{array}$$

This is the famous Bohr formula for the energy levels of hydrogen-like atoms! Here $R{y}^{\ast }$ is the Rydberg energy. The quantum number $n$ must be a positive integer, explaining why atomic energy levels are quantized.

From the denominator of our recursive formula, we also get:

$$l\le j\le n-1$$

This gives us an important constraint on the angular momentum quantum number:

$$l\le n-1.$$

The complete radial function $R_{n,l}(r)$ therefore depends on two quantum numbers:

• $n$ (principal quantum number)
• $l$ (angular momentum quantum number)

The first few normalized radial functions are shown in the table below ($N={\left(Z/\left(a_{0}n\right)\right)}^{3/2}$, $x=Zr/\left(a_{0}n\right)$, $a_0=4\pi {\epsilon }_0{\hslash }^2/\left(\mu e^2\right)$ is the Bohr radius):

$n$	$l$	$R_{n,l}\left(r\right)$
1	0	$2N{\mathrm{e}}^{-x}$
2	0	$2N{\mathrm{e}}^{-x/2}(1-\frac{x}{2})$
2	1	$\frac{1}{\sqrt{3}}2N{\mathrm{e}}^{-x/2}\frac{x}{2}$
3	0	$2N{\mathrm{e}}^{-x/3}(1-\frac{2x}{3}+\frac{2x^2}{27})$
3	1	$\frac{\sqrt{2}}{3}2N{\mathrm{e}}^{-x/3}\frac{x}{3}(2-\frac{x}{3})$
3	2	$\frac{\sqrt{2}}{3\sqrt{10}}2N{\mathrm{e}}^{-x/3}{\left(\frac{x}{3}\right)}^2$

def R_nl(n, l, r):
    x = r # Simplified version without physical constants
    if n == 1 and l == 0:
        return 2 * np.exp(-x)  # Already normalized
    elif n == 2:
        if l == 0:
            return (1/np.sqrt(2)) * (2-x) * np.exp(-x/2)
        elif l == 1:
            return (1/np.sqrt(6)) * x * np.exp(-x/2)
    elif n == 3:
        if l == 0:
            return (2/np.sqrt(27)) * (6-6*x+x**2) * np.exp(-x/3)
        elif l == 1:
            return (8/(27*np.sqrt(6))) * x * (4-x) * np.exp(-x/3)
        elif l == 2:
            return (4/(81*np.sqrt(30))) * x**2 * np.exp(-x/3)
    return 0

r = np.linspace(0, 40, 1000)

fig, ((ax1, ax1p), (ax2, ax2p), (ax3, ax3p)) = plt.subplots(3, 2, figsize=get_size(10,12))

# n = 1
ax1.plot(r, R_nl(1,0,r), label='1s (l=0)')
ax1.set_title('n = 1')
ax1.legend()
ax1p.plot(r, R_nl(1,0,r)**2, label='1s (l=0)')
ax1p.set_title('n = 1 probability density')
ax1p.set_ylim(0,2)
ax1p.set_xlim(0,20)

# n = 2
ax2.plot(r, R_nl(2,0,r), label='2s (l=0)')
ax2.plot(r, R_nl(2,1,r), label='2p (l=1)')
ax2.set_title('n = 2')
ax2.legend()
ax2p.plot(r, R_nl(2,0,r)**2, label='2s (l=0)')
ax2p.plot(r, R_nl(2,1,r)**2, label='2p (l=1)')
ax2p.set_title('n = 2 probability density')
ax2p.set_ylim(0,2)
ax2p.set_xlim(0,20)


# n = 3
ax3.plot(r, R_nl(3,0,r), label='3s (l=0)')
ax3.plot(r, R_nl(3,1,r), label='3p (l=1)')
ax3.plot(r, R_nl(3,2,r), label='3d (l=2)')
ax3.set_title('n = 3')
ax3.legend()
ax3p.plot(r, R_nl(3,0,r)**2, label='3s (l=0)')
ax3p.plot(r, R_nl(3,1,r)**2, label='3p (l=1)')
ax3p.plot(r, R_nl(3,2,r)**2, label='3d (l=2)')
ax3p.set_title('n = 3 probability density')
ax3p.set_ylim(-0.1,2)
ax3p.set_xlim(0,20)


ax3.set_xlabel(r'$r/a_0$')
ax3p.set_xlabel(r'$r/a_0$')
fig.supylabel(r'$R(r)$ / $|R(r)|^2$')

plt.tight_layout()
plt.show()

Figure 2— Radial wavefunctions R(r) for n=1,2,3

A key result is that the energy depends only on n, not on l or m. For each l value, there are 2l+1 degenerate states (different m values with the same energy). The total number of states for a given n is:

$$\sum _{l=0}^{n-1}(2l+1)=n^2$$

giving us $n^2$ distinct wavefunctions:

$${\psi }_{n,l,m}(r,\vartheta ,\phi )=R_{n,l}\left(r\right)\cdot Y_l^m(\vartheta ,\phi )$$.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.special import sph_harm
import random

def hydrogen_wavefunction(n, l, m, r, theta, phi):
    # Include radial part for n=1, l=0
    if n == 1 and l == 0:
        return (1/np.sqrt(np.pi)) * np.exp(-r) * sph_harm(m, l, phi, theta)
    return sph_harm(m, l, phi, theta)

# Create meshgrid
r = np.linspace(0, 10, 50)
theta = np.linspace(0, np.pi, 50)
phi = np.linspace(0, 2*np.pi, 50)

# Number of total points to generate and sample from
n_points_total = 200000
n_points_plot = 10000

# Create figure
fig = plt.figure(figsize=get_size(12, 8))

# Plot for different quantum states
for idx, (l, m) in enumerate([(0,0), (1,0), (1,1)]):
    ax = fig.add_subplot(2, 3, idx+1, projection='3d')

    # Generate more points than we'll plot
    r_points = np.random.uniform(0, 10, n_points_total)
    theta_points = np.random.uniform(0, np.pi, n_points_total)
    phi_points = np.random.uniform(0, 2*np.pi, n_points_total)

    # Calculate wavefunction
    psi = hydrogen_wavefunction(1, l, m, r_points, theta_points, phi_points)
    probability = np.abs(psi)**2

    # Sample points according to probability
    prob_normalized = probability / np.sum(probability)
    indices = np.random.choice(n_points_total, n_points_plot, p=prob_normalized)

    r_sampled = r_points[indices]
    theta_sampled = theta_points[indices]
    phi_sampled = phi_points[indices]
    probability_sampled = probability[indices]

    # Convert to Cartesian coordinates
    x = r_sampled * np.sin(theta_sampled) * np.cos(phi_sampled)
    y = r_sampled * np.sin(theta_sampled) * np.sin(phi_sampled)
    z = r_sampled * np.cos(theta_sampled)

    # Plot points with probability-based transparency
    scatter = ax.scatter(x, y, z, c='blue',
                        alpha=0.04, s=2, edgecolors='none')

    ax.set_title(f'l={l}, m={m}')
    ax.set_axis_off()

# Add second row with l=2
for idx, m in enumerate([0, 1, 2]):
    ax = fig.add_subplot(2, 3, idx+4, projection='3d')

    # Generate more points than we'll plot
    r_points = np.random.uniform(0, 10, n_points_total)
    theta_points = np.random.uniform(0, np.pi, n_points_total)
    phi_points = np.random.uniform(0, 2*np.pi, n_points_total)

    # Calculate wavefunction
    psi = hydrogen_wavefunction(1, 2, m, r_points, theta_points, phi_points)
    probability = np.abs(psi)**2

    # Sample points according to probability
    prob_normalized = probability / np.sum(probability)
    indices = np.random.choice(n_points_total, n_points_plot, p=prob_normalized)

    r_sampled = r_points[indices]
    theta_sampled = theta_points[indices]
    phi_sampled = phi_points[indices]
    probability_sampled = probability[indices]

    # Convert to Cartesian coordinates
    x = r_sampled * np.sin(theta_sampled) * np.cos(phi_sampled)
    y = r_sampled * np.sin(theta_sampled) * np.sin(phi_sampled)
    z = r_sampled * np.cos(theta_sampled)

    # Plot points with probability-based transparency
    scatter = ax.scatter(x, y, z, c='blue',
                        alpha=0.1,s=2,edgecolors='none')

    ax.set_title(f'l=2, m={m}')
    ax.set_axis_off()

plt.tight_layout()
plt.show()

Hydrogen wavefunctions for n=1, l=0,1,2 and m=0,1,2