PINNs Part 3: Solving the Hydrogen Atom Schrodinger Equation
10/26/2025
As described in Part 2, an atom consists of a nucleus and electron(s). The Hydrogen atom has one electron and we will apply the Schrodinger equation to that electron. I will then use a PINN to solve the equation.
The Hydrogen atom Schrodinger equation
The Hamiltonian operator for Hydrogen is defined
$$\begin{eqnarray} \hat{H} = -\frac{\hbar^2}{2m}\nabla^2 + V(\vec{r}) , \end{eqnarray}$$
where $\nabla^2$ is the Laplace operator, $m$ is the mass of the electron, and $V(\vec{r})$ is the potential of the electron defined by
$$\begin{eqnarray} V(\vec{r}) = -\frac{e^2}{4 \pi \varepsilon_0 |\vec{r}|} , \end{eqnarray}$$
where $\vec{r}$ is the position of the electron in 3D cartesian coordinates, $e$ is the charge of the electron, and $\varepsilon_0$ is the permittivity of free space.
Unlike the infinite square well differential equation, the Hydrogen atom Schrodinger has a non-zero potential term. This complicates the equation and makes it significantly harder to solve analytically.
The full Schrodinger equation is
$$\begin{eqnarray} -\frac{\hbar^2}{2m} \left(\frac{\partial^2 \psi}{\partial x^2} + \frac{\partial^2 \psi}{\partial y^2} + \frac{\partial^2 \psi}{\partial z^2} \right) - \frac{e^2}{4 \pi \varepsilon_0 |\vec{r}|}\psi = E\psi . \end{eqnarray}$$
To solve this analytically, the equation must be expressed in spherical coorindates. Skipping some details, the Schrodinger equation becomes
$$\begin{eqnarray}-\frac{\hbar^2}{2m} \left[\frac{1}{r^2} \frac{\partial}{\partial r} \left(r^2 \frac{\partial \psi}{\partial r} \right) + \frac{1}{r^2 sin \theta} \frac{\partial}{\partial \theta} \left(sin \theta \frac{\partial \psi}{\partial \theta} \right) + \frac{1}{r^2 sin^2 \theta} \left(\frac{\partial^2 \psi}{\partial \phi^2} \right) \right] - \frac{e^2}{4 \pi \varepsilon_0 r}\psi = E \psi . \end{eqnarray}$$
A separation of variables can be applied to this equation to break it down into a radial component, $R(r)$, and an angular component, $Y(\theta, \phi)$. This lets us express the wavefunction as
$$\begin{eqnarray}\psi(r, \theta, \phi) = R(r) Y(\theta, \phi) , \end{eqnarray}$$
and gives two differential equations
$$\begin{eqnarray} l(l + 1) &=& \frac{1}{R} \frac{d}{dr} \left(r^2 \frac{dR}{dr} \right) - \frac{2mr^2}{\hbar^2}[V(r) - E ] , \\ -l(l + 1) &=& \frac{1}{Y} \left[\frac{1}{sin \theta} \frac{\partial}{\partial \theta} \left(sin \theta \frac{\partial Y}{\partial \theta} \right) + \frac{1}{sin^2 \theta} \frac{\partial^2 Y}{\partial \phi^2} \right] . \end{eqnarray}$$
Finding the analytical solutions to $R$ and $Y$ involves a ton of steps that I won't go over. If you're interested in the entire solution, I encourage you to check out chapter 4 of Introduction to Quantum Mechanics (Third Edition) by David J. Griffiths and Darrell F. Schroeter or Richard Behiel's YouTube series.
The analytical solution to $R$ is given by
$$\begin{eqnarray} R_{nl}(r) = \sqrt{\left(\frac{2}{na}\right)^3 \frac{(n - l - 1)!}{2n(n + l)!}}e^{-r/na} \left(\frac{2r}{na}\right)^l \left[L^{2l + 1}_{n - l - 1}\left(\frac{2r}{na}\right)\right] , \end{eqnarray}$$
where $a$ is the Bohr radius, and the associated Laguerre and Laguerre polynomials are defined as
$$\begin{eqnarray} L^p_q(x) &\equiv& (-1)^p \left(\frac{d}{dx}\right)^p L_{p + q}(x) , \\ L_q(x) &\equiv& \frac{e^x}{q!} \left(\frac{d}{dx}\right)^q (e^{-x}x^q) . \end{eqnarray}$$
The analytical solution to $Y$ is given by
$$\begin{eqnarray} Y^m_l(\theta, \phi) = \sqrt{\frac{(2l + 1)(l - m)!}{4 \pi (l + m)!}} e^{im \phi} P^m_l(cos \theta) , \end{eqnarray}$$
with the associated Legendre function and the Legendre polynomial defined as
$$\begin{eqnarray} P^m_l(x) &\equiv& (-1)^m (1 - x^2)^{m/2} \left(\frac{d}{dx}\right)^m P_l(x) , \\ P_l(x) &\equiv& \frac{1}{2^l l!} \left(\frac{d}{dx}\right)^l (x^2 - 1)^l , \end{eqnarray}$$
and $n$, $l$, $m$ are quantum numbers (constants) introduced when solving for $R$ and $Y$.
So, the full wavefunction for the Hydrogen atom is
$$\begin{eqnarray} \psi_{nlm} = \sqrt{\left(\frac{2}{na}\right)^3 \frac{(n - l - 1)!}{2n(n + l)!}}e^{-r/na} \left(\frac{2r}{na}\right)^l \left[L^{2l + 1}_{n - l - 1}\left(\frac{2r}{na}\right)\right] \sqrt{\frac{(2l + 1)(l - m)!}{4 \pi (l + m)!}} e^{im \phi} P^m_l(cos \theta) , \end{eqnarray}$$
with the energy eigenvalue (found when solving $R$),
$$\begin{eqnarray} E_n = -\left[\frac{m}{2 \hbar^2} \left(\frac{e^2}{4 \pi \varepsilon_0} \right)^2 \right] \frac{1}{n^2}. \end{eqnarray}$$
Easy!
Clearly, this analytical solution is a steep jump in complexity compared to the simple sinusoid solution encountered in the infinite well problem discussed in Part 2. The Hydrogen atom is the only atom in the periodic table that can be (or at least, has been) solved analytically. All atoms with more electrons do not have analytical solutions and must be approximated through other numerical methods for predictions. Getting a PINN to solve the full 3D Hydrogen atom is such an interesting task to me because if successful, the methods could be extended to predict the wavefunctions and energy levels of multi-electron atoms.
Units
Just like in Part 2, I will use Coulomb units in this problem. That means that I will set
$$\begin{eqnarray} \frac{\hbar^2}{2m} &=& \frac{1}{2}, \\ \frac{e^2}{4 \pi \varepsilon_0} &=& 1, \\ E_n &=& -\frac{1}{2n^2}. \end{eqnarray}$$
The Bohr radius, $a_0$, is a physical constant approximately defined as the most probable distance between the nucleus and electron in the Hydrogen atom. In Coulomb units, the Bohr radius conveniently becomes
$$\begin{eqnarray} a_0 &=& \frac{4 \pi \varepsilon_0 \hbar^2}{m e^2} \\ &=& \frac{\hbar^2}{mC} \\ &=& 1. \end{eqnarray}$$
So, the unit of length is equal to one Bohr radius.
It took me an unreasonably long time to get these units right, and I observed an endless amount of bizarre PINN predictions as a consequence of not having consistent and correct units in my problem setup.
Implementing a PINN to solve the 1D radial equation
In order to debug and check the validity of my units, I quickly implemented a PINN to predict the 1D radial equation solution. I didn't bother with scanning across $E$ values since I was just checking to see if I had a correct setup.
The $R$ differential equation is solved analytically with the substitution $u(r) = rR(r)$, making it
$$\begin{eqnarray} -\frac{\hbar^2}{2m} \frac{d^2u}{dr^2} + \left[V(r) + \frac{\hbar^2}{2m} \frac{l(l+1)}{r^2} \right]u = Eu . \end{eqnarray}$$
In Coulomb units, the potential becomes $V(r)=-\frac{e^2}{4 \pi \varepsilon_0}\frac{1}{r}=-\frac{C}{r}=-\frac{1}{r}$. Thus, the radial equation can be expressed as
$$\begin{eqnarray} -\frac{1}{2} \frac{d^2u}{dr^2} + \left[-\frac{1}{r} + \frac{1}{2} \frac{l(l+1)}{r^2} \right]u = -\frac{1}{2n^2}u . \end{eqnarray}$$
For the analytical solutions to the radial equation please see page 2 of this text or see chapter 4 of Griffiths.
This differential equation has a similar form to the 1D infinite square well differential equation. For simplicity's sake, I will set $l=0$. $f$ is defined by moving all terms to one side
$$\begin{eqnarray} f = -\frac{1}{2} \frac{d^2u}{du^2} - u \left( \frac{1}{r} + E \right).\end{eqnarray}$$
$L_f$ is defined in the same way as in Part 2 and I use the parametric scaling function to handle boundary conditions again.
Data generation simply requires the creation of values $r \in [r_{min}, r_{max}]$ where $r_{min}$ and $r_{max}$ are the defined minimum and maximum radii of the system, respectively, and a fixed step size of $\delta_r = 0.01$ is used.
With the correct system of units finally determined, the first three analytical $E_n$ values were set and the PINN was trained for 10,000 training steps on each $E_n$. A simple feedforward architecture with sine activation functions was used again. The PINN attained the following $L_f$ values and predicted the first three wavefunctions with decent accuracy.
The PINN's predicted solutions attain $L_f$ values on the order of $10^{-4}$ to $10^{-3}$ and MSE values of $10^{-5}$ to $10^{-1}. I believe better accuracy could be obtained with more training steps, but again, this was just a spot check to see if I had my units system defined correctly. My main focus is getting a PINN to solve the 3D equation.
TODO use 1D radial PINN as spot check on units + other setup details (maybe just fit first 3 solns w no scan and present them as being done just to check that pinn can find solns w these units correctly)