Introducing Machine Learning in Teaching Quantum Mechanics

Abstract

In this article, we describe an approach to teaching introductory quantum mechanics and machine learning techniques. This approach combines several key concepts from both fields. Specifically, it demonstrates solving the Schrödinger equation using the discrete-variable representation (DVR) technique, as well as the architecture and training of neural network models. To illustrate this approach, a Python-based Jupyter notebook is developed. This notebook can be used for self-learning or for learning with an instructor. Furthermore, it can serve as a toolbox for demonstrating individual concepts in quantum mechanics and machine learning and for conducting small research projects in these areas.

1. Introduction

Machine learning is currently a rapidly developing domain of computer science with applications in various fields, including the natural sciences, engineering, social sciences, and digital humanities, as well as business and economics. Unlike the previous approaches to using computers for prediction and analysis, such as expert systems, machine learning algorithms are designed to learn about the phenomena under study directly from large datasets. Machine learning is broadly divided into categories such as supervised learning, unsupervised learning, and reinforcement learning, depending on the objective, which could be prediction or the extraction of properties from data. Traditional machine learning relies on algorithms like decision trees, random forests, support vector machines (SVMs), and kernel methods [1,2,3], which reached their peak in the 2000s.

Currently, the dominant paradigm in machine learning relies on neural networks [4,5,6], which are loosely inspired by the architecture of biological brains. It began with the conception of a perceptron [7,8,9,10] in the middle of the XXth century and its subsequent study [11,12,13,14]. Later on, in order to efficiently train neural networks, the back-propagation technique [15,16] was developed. However, because of conceptual issues, as well as limited computer power, spectacular popularity for neural networks had to wait until the 2010s when easy access to GPU-assisted computation became common. This started an era of development of neural networks known under the name of deep learning [17,18]. The neural network paradigm has found diverse uses: in image recognition, convolutional neural networks [19] have proven to be substantially more successful than the previously used machine learning approaches; long short-term memory or LSTM networks [20] have been widely applied to speech recognition, machine translation, and time series forecasting (in opposition to the traditional ARIMA models). More recently, transformers [21] have replaced LSTMs in these application, as well as becoming integral to so-called Large Language Models (LLMs), which are used to generate stochastic text.

Machine learning has also found specific applications in physics and quantum chemistry [22,23]. Statistical physics has utilised unsupervised and supervised learning with restricted Boltzmann machines and neural networks; particle physics and astrophysics have employed classification and regression machine learning models; condensed matter physics [24] has seen its use in simulations of strongly correlated systems and the phases of matter quantum states and to accelerate Monte Carlo simulations; atomic and molecular physics has used feedforward neural networks to predict the state distributions of atom–diatom reactions [25]; quantum chemistry builds the potential energy surfaces in chemical reactions using machine learning as an alternative to ab initio calculations [26]; and quantum measurement and metrology have leveraged neural networks and SVMs [27]. Neural networks have been used to learn and represent the wave functions of quantum many-body systems, particularly their ground states and dynamics, and to discover their compact representations [28].

Neural networks have also been used to facilitate solving ordinary [29] and partial differential equations [30,31]. They have been used to solve physics-related differential equations and the Schrödinger equation in particular [32,33,34,35,36], which is a thriving field, with many solvers proposed. Conventional solvers suffer from the fact that the dimensionality of the Hilbert space of states being a function of the number of electrons means that the computational cost scales exponentially, making finding the solutions to atomic, molecular, and condensed matter physics problems computationally difficult. In contrast, neural network solvers benefit from slower—polynomial—growth in computational cost.

In the context of quantum science education, the Python programming language and Jupyter notebooks [37] have been used as an aid for students and other learners. Jupyter notebooks are an implementation of computational notebooks which combine code, Markdown-formatted text, HTML elements, images, and other visualisations into compact, stand-alone documents. Their interactive nature allows for free experimentation, which provides opportunities to easily answer questions that could not be answered in a static text. For example, the quantum computing frameworks Qskit [38] and PennyLane [39] utilise Jupyter notebooks as a fundamental part of their demonstrations and tutorials.

In this paper, we will focus on describing a pedagogical ML model that learns about the ground states of quantum systems. Given the preponderance of neural-network-based Schrödinger equation solvers, we do not claim to have constructed a state-of-the-art, innovative solver—rather, we aim to provide a toy model example which can be used to acquaint students (and other learners) with neural networks in the context of quantum sciences in a way that is not completely artificial and in fact is being developed in the research. In Section 2, we introduce the concepts from quantum mechanics. We describe the Schrödinger equation, as well as three types of quantum mechanical potentials: the harmonic potential, the Morse potential, and the polynomial potential. We describe the wave functions associated with each of them when they allow analytic solutions and discuss the numerical method for solving the Schrödinger equation, called discrete-variable representation (DVR), when they do not. In Section 3, we describe the architecture of neural networks, in particular fully connected neural networks. A discussion of the loss functions and stochastic gradient descent is also included. We continue discussing neural networks in Section 4, which is dedicated to the process of training a machine learning model. In particular, we focus on the concepts of underfitting and overfitting and the process of splitting the dataset into training and test sets to avoid these issues, as well as choosing the hyperparameters of a model. In Section 5, we present the ML model trained on the solutions to the Schrödinger equation and the Python-based Jupyter notebook that realises it. Finally, in Section 6, we provide the conclusions. Appendix A contains a short summary of the method of integration by quadrature, while Appendix B provides the derivation of plane-wave DVR and orthogonal polynomial DVR methods for numerically solving the Schrödinger equation.

2. Quantum Mechanical Systems

In undergraduate quantum mechanics, students first encounter particles confined within a one-dimensional potential $V\left(x\right)$. While in the classical case one is interested in finding the position and momentum of a particle and how they evolve in time, in quantum mechanics, the object of interest is the less intuitive concept of a quantum state.

Quantum mechanical problems are often presented in terms of the Schrödinger equation in the position representation, in which the unknowns are the wave functions $\psi \left(x\right)$ and the associated spectrum of energies E. If the kinetic energy operator is denoted by $\widehat{T}$ and the potential $V\left(x\right)$ depends only on the position operator $\widehat{x}$, then the Schrödinger equation has the form

$$(\widehat{T}+V\left(\widehat{x}\right))\psi \left(x\right)=E\psi \left(x\right),$$

(1)

where $\widehat{H}=\widehat{T}+V\left(\widehat{x}\right)$ is known as the Hamiltonian operator.

If one only has a wave function $\psi$, which we assume to be normalised

$${\int }_{-\infty }^{+\infty }\mathrm{d}x{\left|\psi \left(x\right)\right|}^2=1,$$

(2)

then one can compute the average energy of that state using the formula

$$\langle E\rangle ={\int }_{-\infty }^{+\infty }\mathrm{d}x\psi {\left(x\right)}^{*}\widehat{H}\psi \left(x\right).$$

(3)

If the wave function $\psi$ is a solution to the Schrödinger Equation (1), then the average energy $\langle E\rangle$ returns simply the energy E.

2.1. The Harmonic Oscillator

The harmonic oscillator is a prototypical quantum system, described by a quadratic potential of the form

$$V\left(x\right)={\displaystyle \frac{1}{2}}m{\omega }^2{x}^2,$$

(4)

where m is the mass of the particle, and $\omega$ is the angular frequency. The harmonic oscillator has the benefit of being analytically solvable—its wave functions, which correspond to the energies $E_n=ħ\omega \left(n+{\displaystyle \frac{1}{2}}\right)$ for $n=0,1,2,\dots$, are given by

$${\psi }_n\left(x\right)={\left({\displaystyle \frac{m\omega }{\pi ħ}}\right)}^{{\displaystyle \frac{1}{4}}}{\displaystyle \frac{1}{\sqrt{2^{n}n!}}}{H}_n\left(\sqrt{{\displaystyle \frac{m\omega }{ħ}}}x\right)e^{-{\displaystyle \frac{m\omega }{2ħ}}{x}^2},$$

(5)

where the Hermite polynomials $H_n\left(y\right)$ are defined by the equation

$$H_n\left(y\right)={(-1)}^n{e}^{y^2}{\displaystyle \frac{{\mathrm{d}}^n}{\mathrm{d}{y}^n}}{e}^{-y^2}.$$

(6)

The harmonic potential together with the harmonic oscillator’s ground-state wave function and its associated energy are illustrated in Figure 1.

2.2. The Morse Potential

While small oscillations in a bound particle can often be approximated using the harmonic potential, in many situations, such as the interaction of two isolated atoms with each other (two-particle collisions or diatomic molecules, for example), the potential of interaction is better represented conceptually using the Morse potential [40,41,42]. It is described as

$$V\left(x\right)=D{\left(1-e^{-a(x-x_e)}\right)}^2-D,$$

(7)

where D, a are the positive-value parameters of the potential, and $x_e$ is the position of the global minimum of the potential well. The Schrödinger equation for the potential is exactly solvable, giving the energies

$$E_n=ħ\omega \left(n+{\displaystyle \frac{1}{2}}\right)-ħ\omega x_e{\left(n+{\displaystyle \frac{1}{2}}\right)}^2-D,$$

(8)

where $n=0,1,2,\dots$ and

$$\omega =\sqrt{{\displaystyle \frac{2D{a}^2}{m}}},x_e={\displaystyle \frac{ħ\omega }{4D}}.$$

(9)

The wave functions for the energies are

$${\psi }_n\left(x\right)=\sqrt{{\displaystyle \frac{a(l-2n-1)n!}{\Gamma (l-n)}}}{e}^{-{\displaystyle \frac{z}{2}}}{z}^{{\displaystyle \frac{l-2n-1}{2}}}{L}_n^{l-2n-1}\left(z\right),$$

(10)

where $l=x_e^{-1}$ and $z=l{e}^{-a(x-r_e)}$ and the associated Laguerre polynomials are defined by the formula

$$L_n^{\left(\alpha \right)}\left(x\right)={\displaystyle \frac{x^{-\alpha }{e}^x}{n!}}{\displaystyle \frac{{\mathrm{d}}^n}{\mathrm{d}{x}^n}}\left(e^{-x}{x}^{n+\alpha }\right).$$

(11)

The Morse potential together with a wave function and its associated energy are illustrated in Figure 2.

2.3. Polynomial Potentials

Not all quantum mechanical systems admit analytic solutions. One can approximate various physical systems, with a potential well, using a polynomial potential of the form

$$V\left(x\right)=\sum _{i=0}^n{\alpha }_i{x}^i,$$

(12)

where ${\alpha }_i$ are parameters with ${\alpha }_n>0$. Unfortunately, this system does not generally have an analytic solution, and therefore one has to use numerical methods to obtain the wave functions and the energy spectrum. Figure 3 gives an example of the polynomial potential of Equation (12) with $n=5$ and its ground-state wave function.

2.4. Numerical Solutions to the Schrödinger Equation

Equation (1) can be solved numerically. Many different numerical methods exist [43], each with its own advantages and drawbacks. To solve differential equations like the Schrödinger equation, one can use propagation methods, e.g., Runge–Kuta [44,45] or Numerov’s [46], or variational methods, like finite-element methods, spectral methods, or pseudo-spectral methods. (Pseudo-)spectral methods have the advantage of allowing a Hamiltonian to be represented as a finite-dimensional matrix, reducing solving a differential equation into an eigenvalue problem. As such, they approximate the low-energy spectrum automatically and have good error properties compared to finite-elements methods.

We will use a particular pseudo-spectral method—the discrete-variable method—when we solve the Schrödinger equation numerically. In particular, this has the advantage that the coefficients of the eigenvectors of the Hamiltonian are simply wave functions in the position representation evaluated at grid points in space and therefore have an immediate and straightforward interpretation (see Equation (27)).

One of the numerical methods used for this purpose is recasting the Schrödinger equation as an eigenvalue problem. Let us suppose that a basis ${\left\{{\phi }_i\left(x\right)\right\}}_{i=1}^{\infty }$ of square-integrable functions $\phi \in L^2\left(\mathbb{R}\right)$ exists that spans the entire $L^2\left(\mathbb{R}\right)$ and is orthonormal and complete. Then, any solution $\psi$ can be decomposed into this basis

$$\psi \left(x\right)=\sum _{i=0}^{\infty }{c}_i{\phi }_i\left(x\right),$$

(13)

where the coefficients ${\left\{c_i\right\}}_{i=0}^{\infty }$ constitute a square-summable series, i.e., ${\sum }_{i=0}^{\infty }{\left|c_i\right|}^2<+\infty$. Using the orthonormality of the basis, Schrödinger’s Equation (1) can be put into the form

$$\sum _{j=0}^{\infty }\left(T_{ij}+V_{ij}+E{\delta }_{i,j}\right)c_j=0,$$

(14)

where $T_{ij}={\int }_{-\infty }^{+\infty }\mathrm{d}x{\phi }_i^{*}\left(x\right)\widehat{T}{\phi }_j\left(x\right)$, $V_{ij}={\int }_{-\infty }^{+\infty }\mathrm{d}x{\phi }_i^{*}\left(x\right)V\left(x\right){\phi }_j\left(x\right)$. This is an infinite-dimensional eigenvalue problem for the eigenvector $c={(c_1{c}_2\dots )}^T$ and the eigenvalue E.

Equation (14) is not treatable using numerical methods given its infinite dimension. However, one can replace this problem with an approximate problem

$$\sum _{j=0}^{N-1}\left(T_{ij}+V_{ij}+E{\delta }_{i,j}\right)c_j=0,$$

(15)

with a wave function

$$\psi \left(x\right)=\sum _{i=0}^{N-1}{c}_i{\phi }_i\left(x\right),$$

(16)

where $N\in {\mathbb{Z}}_{>1}$ determines where the infinite basis is truncated to a finite basis.

• The choice of the truncation dimension N is important for minimising the error, which comes from approximation. The coefficients $c_i$ vanish quickly enough such that after some term N, the error ${\sum }_{i=1}^{\infty }|c_i{|}^2-{\sum }_{i=1}^N{\left|c_i\right|}^2\ll 1$ can be neglected. However, the choice of the correct truncation depends on the energy, and therefore the choice of N is important.
• The coefficients $T_{ij}$ and $V_{ij}$ have to be integrated. Numerical integration of the coefficients can be performed numerically or, in some cases, analytically.

The method called the discrete-variable method (DVR) [47,48,49,50,51,52,53] allows approximate values of the coefficients $T_{ij}$ and $V_{ij}$ to be obtained while minimising the errors. The coefficients of the kinetic energy operator are known analytically, while the potential energy operator has a diagonal form. The details of the method are provided in Appendix B. In particular, there are many different implementations [54,55,56], depending on the choice of the basis functions ${\phi }_i\left(x\right)$. Here, we will use the Fourier plane-wave basis [57], which is described below.

Let us consider a domain $[x_0,x_0+L]$ with a collection of N equally spaced points $x_i=x_0+i\Delta x$, where the length of the interval L and the spacing $\Delta x$ are related by $L=\Delta xN$. The collection of points ${\left\{x_i\right\}}_{i=0}^{N-1}$ forms a discrete grid which is used to calculate the coefficients of the Hamiltonian (note that the endpoint $x_N=x_0+L$ does not belong to this set). Moreover, let us assume that N is an even integer. The basis functions in which the wave functions are expanded are given by the sums of the plane waves in the form

$${\phi }_i\left(x\right)={\displaystyle \frac{1}{N}}\sum _{l=-N/2}^{N/2-1}{e}^{i{k}_l(x-x_i)},$$

(17)

where the wave numbers $k_l$ are given by

$$k_l={\displaystyle \frac{2\pi }{L}}l,-{\displaystyle \frac{N}{2}}\le l\le {\displaystyle \frac{N}{2}}-1.$$

(18)

The basis functions ${\phi }_i\left(x\right)$, illustrated in Figure 4, on the grid are proportional to the Kronecker function

$${\phi }_i\left(x_j\right)={\displaystyle \frac{1}{\sqrt{\Delta x}}}{\delta }_{i,j},$$

(19)

and are non-zero only at one of the grid points. The basis functions are linear combinations of plane waves. Plane waves are characterised by wave numbers. Because of the truncation, one can only represent a finite range of wave numbers. The lowest wave number that can be represented is

$$k_{\min }=-{\displaystyle \frac{\pi N}{L}},$$

(20)

while the highest is

$$k_{\max }={\displaystyle \frac{\pi }{L}}(N-2).$$

(21)

Because the wavelength and the wave number are related to each other

$$\lambda ={\displaystyle \frac{2\pi }{\left|k\right|}},$$

(22)

the constraints for the wave number imply that the minimal de Broglie wavelength that can be represented is

$${\lambda }_{\min }={\displaystyle \frac{2L}{N}}=2\Delta x.$$

(23)

This also means that the kinetic energy $T={\displaystyle \frac{{ħ}^2{k}^2}{2m}}$ is bounded

$$0\le T\le {\displaystyle \frac{1}{2m}}{\left({\displaystyle \frac{\pi ħN}{L}}\right)}^2,$$

(24)

where m is the mass of the particle. This gives a clear bound on the energies of the states that can be described using DVR, and this is in agreement with our heuristic description of truncation of the basis, which we discussed above. In order to faithfully capture the states of higher energy, the number of grid points has to be augmented (or alternatively, the number of basis functions) N.

The form of the kinetic energy operator coefficients $T_{ij}$ on the plane-wave basis (which holds for N being an even integer) is as follows:

$$T_{ij}={\displaystyle \frac{{ħ}^2}{m}}{\displaystyle \frac{{\pi }^2}{N^2{(\Delta x)}^2}}{(-1)}^{j-i}\left\{\begin{array}{cc}{\displaystyle \frac{N^2+2}{6}}\hfill & i=j,\hfill \\ {\left[sin\left({\displaystyle \frac{\pi (j-i)}{N}}\right)\right]}^{-2}\hfill & i\ne j,\hfill \end{array}\right.$$

(25)

while the potential energy operator is diagonal, and its coefficients are obtained by just evaluating the potential on the grid ${\left\{x_i\right\}}_{i=0}^{N-1}$

$$V_{ij}=V\left(x_i\right){\delta }_{i.j}.$$

(26)

With the coefficients of the kinetic energy operator and the potential energy operator defined, one can solve Equation (15) as an eigenvalue problem. In particular, one can use the function scipy.linalg.eigh from the Python package scipy [58] to perform this operation on a computer.

The coefficients of the wave function (16) are actually related to the values of the wave function on the grid points

$$c_i=\sqrt{\Delta x}\psi \left(x_i\right).$$

(27)

Therefore, the wave function is an envelope along those coefficients, as illustrated in Figure 5. It is worthwhile noting that the discrete counterpart of the normalisation (2) implies

$$\Delta x\sum _{i=0}^{N-1}|\psi \left(x_i\right){|}^2=\sum _{i,j=0}^{N-1}{\left|c_i\right|}^2=1.$$

(28)

3. Neural Network Models

In traditional computational sciences, one makes use of expert systems and dedicated algorithms (which come from statistics or the relevant sciences) in order to solve a given problem. These traditional methods have been augmented by the advent of machine learning techniques, where instead of encoding the dedicated algorithms directly, one lets the algorithm learn from the data, provided that a sufficient amount of relevant data exists. Often, that amount is extremely large.

Machine learning can respond to a variety of problems. Unsupervised learning concerns situations where only feature data is available, in the absence of target data, and the goal is to extract certain structures from the dataset itself. Common tasks in unsupervised learning are clustering (using clustering algorithms like k-means, hierarchical clustering, or DBSCAN) and dimensionality reduction (using, e.g., a principal component analysis, t-SNE, or autoencoders). Reinforcement learning is a paradigm of machine learning where a model is taught by offering rewards or punishments depending on the outcome, and it is used in the context of agent control, e.g., in robotics or game playing programmes. Meanwhile, unsupervised learning and reinforcement learning (and generative machine learning) are fields in their own right, and we are interested in supervised machine learning in particular. This is a subfield of machine learning which given pairs $(x_i,y_i)$ of (n-dimensional) input data and (m-dimensional) target data attempts to construct a function ${\widehat{y}}_i=f(x_i;\theta )$ such that the difference between the actual target data $y_i$ and the prediction ${\widehat{y}}_i$ is minimised. The parameters $\theta$ describe the model, and their exact values are determined from the input and target data. Supervised learning is used mostly in prediction tasks, like classification and regression. It has been very successfully applied to the tasks of image recognition and classification, machine translation, machine transcription, and natural language processing and understanding, among many others.

One could ask how one would choose the function $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$. The space of possible functions of this type is enormous, and the choice is between the possibility of encoding the relations behind the empirical data sufficiently and having as small a number of free parameters as possible. Various machine learning algorithms, like linear regression, logistic regression, and kernel support vector machines, provide different forms of functions f.

A particularly fruitful choice of architecture is a deep feedforward neural network or a multi-layer perceptron. For this particular choice of machine learning algorithm, the function $\widehat{y}=f(x;\theta )$ is a composition of multiple simpler functions, which constitute layers

$$\widehat{y}=f^{\left(k\right)}\circ \dots f^{\left(2\right)}\circ f^{\left(1\right)}\left(x_i\right),$$

(29)

each one with its own, smaller set of parameters ${\theta }^{\left(i\right)}$, $i=1,\dots ,k$. The fact that one can have multiple layers provided the name deep learning to neural networks.

The basic building block of a feedforward neural network is a neuron [5,6], which is a function ${\mathbb{R}}^n\to \mathbb{R}$ (c.f. Figure 6). Neurons can be grouped into collections to create fully connected layers, i.e., functions ${\mathbb{R}}^n\to {\mathbb{R}}^m$

$$y_i=g({\left(Wx\right)}_i+b_i),$$

(30)

where $x\in {\mathbb{R}}^n$ is the input, $y\in {\mathbb{R}}^m$ is the output, W is the mapping $({\mathbb{R}}^n\to {\mathbb{R}}^m)$, $b\in {\mathbb{R}}^m$ are the parameters of the neural network, and g is a nonlinear function ${\mathbb{R}}^m\to {\mathbb{R}}^m$. One uses a nonlinear function g that is applied element-wise to allow the neural network to produce more complicated functions than just linear ones, which are too restrictive in terms of the possible outcomes. Popular choices for the nonlinear function g are the hyperbolic tangent tanh, the softmax function, and the rectified linear unit ReLU.

Until now, we have described the architecture of the neural network depending on its weights $(W,b)$, which can give us the target values as a function of the input values. However, starting from a random choice of weights would lead to a model that would not correctly predict the data, as the space of parameters would be very large. Therefore, there are two issues to resolve:

• How can one evaluate the correctness of the predictions of a neural network?
• Given the correctness, how one can improve the weights iteratively so that they give better results?

The first question is answered using the so-called loss function. Loss functions are non-negative definite scalar functions

$$l=l(\widehat{y},y),$$

(31)

which depend on the target data y and the prediction $\widehat{y}$. They are chosen in such a way that they obtain the minimum $l(y,y)$ when the prediction is equal to the target data $\widehat{y}=y$, and any difference contributes positively. Loss functions have a probabilistic interpretation in terms of maximising the log-likelihood; however, usually, it is sufficient to choose the one that is most suitable for a given problem. In the context of classification problems, one tends to use the logarithm of sigmoid or softmax functions, while in the context of regression problems, the mean squared error is used: 

$$\mathrm{MSE}(\widehat{y},y)={\displaystyle \frac{1}{m}}\sum _{i=1}^m{({\widehat{y}}_i-y_i)}^2.$$

(32)

Having resolved the first issue, we can now focus on a way to iteratively improve the model’s weights so that it minimises the loss function. The method used for this purpose is called gradient descent [59]. One can imagine the multi-dimensional surface given by the loss function as a function of the weights of the neural network

$$(W,b)\mapsto l(\widehat{y}(x;W,b),y).$$

(33)

This surface will possibly have multiple minima—a global minimum and a collection of local minima. Ideally, one desires an algorithm which provides the location of the global minimum and therefore the best possible weights given the input and target data. However, this is not feasible, and a local method is used where the weights along the direction of the gradient of the loss function are updated:

$$\begin{array}{cc}\hfill W^{\prime }& =W-ϵ{\nabla }_{W}l(\widehat{y}(x;W,b),y),\hfill \\ \hfill b^{\prime }& =b-ϵ{\nabla }_{b}l(\widehat{y}(x;W,b),y),\hfill \end{array}$$

(34)

where $ϵ$ is called the learning rate. The simplest choice for the learning rate is a constant number. However, in deep learning, various adaptive algorithms have been developed that adjust the learning rate dynamically as a function of the gradients, like AdaGrad [60], RMSprop, and Adam [61].

In order to effectively perform gradient descent, one needs to quickly compute the gradients of the loss function. This is conducted with the help of an efficient back-propagation technique which computes the gradients using the chain rule by keeping track of the graph of elementary mathematical operations (e.g., addition, multiplication, exponentiation, etc.) that build up the loss function.

Given that the input and target data contains multiple (multi-dimensional) points, one can make a choice on when to update the weights. If the update happens after adding up the gradients for all of the points, then the gradient descent is deterministic. However, by bundling the data points into batches, computing the gradient for each batch, and updating the weights, one can stochastically approximate the deterministic gradient. This method is known under the name of stochastic gradient descent. Especially for high-dimensional problems, stochastic gradient descent reduces the computational complexity and therefore the demand on the computer hardware, at the price of a lower convergence rate.

4. Machine Learning Training and Prediction

The process that uses stochastic gradient descent in order to find the model weights that best represent the data is known under the name of training of the machine learning model (which we illustrate schematically in Figure 7). If one trains a model on the whole set of pairs of input data and target data, one can expect (under certain assumptions) that the loss function will either monotonically tend towards a local minimum or start to oscillate around it. Therefore, if the loss function on the data which are used for the training is the only measure of quality of the model, then spending more and more time on training iterations would give a better model. In reality, however, this is not the case—and this is related to the concept of the generalisability of the model.

Let us imagine that the initial data used for training comes from some larger probabilistic distribution and as such constitutes just a sample taken from this distribution. If the sample does not faithfully reproduce all of the properties of the probabilistic distribution but rather exhibits some random aberrations (e.g., the initial data consists of more points at a distance $3\sigma$ from the mean than expected, or it is skewed to one side of the mean), then the machine learning model will also learn the aberrations from the sample dataset. Once one evaluates the model on other data generated from the initial probability distribution, the model will not perform well because it learned the random properties of the training data, which were only present in the sample. This phenomenon is known under the name of overfitting—as opposed to underfitting, which happens when the model is not trained for long enough and does not learn the properties of the training data sufficiently.

Therefore, in order to prevent overfitting, one puts aside part of the initial data and avoids training the machine learning model on it. This part of the dataset that is put aside is called the test data. If the loss function evaluated on the training data is approximately equal to the loss on the test data after the process of training, one knows that both overfitting and underfitting have been prevented.

The situation is more complicated when one wants to test the performance of multiple different models or a model that has various adjustable hyperparameters (like the number of layers of the neural network, the training time, etc.). In that case, using the results of evaluating the models on the test data to choose the hyperparameters would not indicate under-/overfitting correctly. To aid with this, a validation procedure is developed. One proceeds with dividing the initial data into three distinct parts which have different purposes:

• Training data—the largest part of the initial data, which is used to train the machine learning models;
• Validation data—the part of the initial data which is used to determine the hyperparameters of the model;
• Test data—the part of the initial data (∼20%) which is used to determine the quality of the model and check whether it is under- or overfitted.

This does leave a particular issue, however—if the initial data is limited (e.g., it is observational data related to markets, climate, weather, etc.) and cannot be generated easily (as is the case in our application to quantum mechanics in Section 5), then setting aside both the training and validation data significantly cuts into the amount of data points that can be used to train the machine learning model. The technique of cross-validation has been developed to avoid this problem.

In the case of neural networks, outside of the quantities which determine the architecture (like the number of layers or the amount of neurons), the number of epochs is the most obvious hyperparameter. Training a neural network on the entirety of the training data once over using stochastic gradient descent constitutes one epoch. Since the training can be repeated indefinitely, reducing (under certain conditions) the training loss each time, the number of epochs equivalently measures the time spent training the model. One can set an early stopping condition on the training process that is applied when the validation loss ceases to improve. Figure 8 illustrates this procedure.

5. Neural Networks for Solving the Schrödinger Equation

In order to illustrate the above concepts, we developed a Python $3.8$-based Jupyter notebook in which a fully connected neural network was trained to predict the wave functions of a quantum mechanical system given its defining potential.

The initial data constitutes the pairs $(V,\psi )$, $V\in {\mathbb{R}}^N,\psi \in {\mathbb{R}}^N$ of potentials and wave functions evaluated on N equidistant grid points. These wave functions are the solutions to the Schrödinger equation with the potentials V obtained using the plane-wave DVR method, implemented using numpy [62] version 1.24.3 and scipy version 1.10.1 Python packages. We have chosen to take $N=200$ as the default value.

The training data and the validation data is composed of polynomial potentials, with the coefficients of the polynomials randomly chosen from uniform distributions, which ensures that the neural network can learn a variety of different potentials. The degree of the polynomial potential as well as the shape of the uniform probability distributions can be adjusted.

In addition to the training and validation data, we also use the test data to evaluate the performance of the final model. The test data contains polynomial potentials drawn from the same distribution as the training and validation data. Moreover, we also use (separately) data points generated from the harmonic oscillator and Morse potentials. Since these potentials admit analytical solutions, we use the analytical solutions instead of the numerical solutions provided by the DVR method. While the harmonic oscillator potential is a subtype of polynomial potential, the Morse potential is not (as it involves an exponential function) and cannot belong to the training data. These polynomials are added to the test data because it is instructive to learn how our neural network generalises to different, previously unseen potentials.

The Python package sklearn [63] version 1.3.0 is used to split the initial data containing (by default) 5000 polynomial potentials and wave functions $66\%:33\%$ between the training data and the validation data.

The machine learning model is supposed to take a potential evaluated on N points $V\in {\mathbb{R}}^N$ as an input and give a wave function evaluated on the same points $\widehat{\psi }\in {\mathbb{R}}^N$ as an output

$$\widehat{\psi }=\mathrm{MLP}\left(V\right).$$

(35)

The fully connected neural network that we use is composed of five layers—four of them are hidden layers, and the last one is the output layer. The hidden layers are composed using ReLU functions and have the output sizes 256, 256, 128, and 128, while the output layer does not have a nonlinear component, and its output size is equal to the number of grid points N, which is also the size of the input of the first layer. Each layer is separated by a dropout layer, which is used only during training—dropout selectively sets the output of individual neurons randomly to zero, which has been shown to improve the training by introducing noise into the process. At the level of implementation, the neural network uses the tensorflow version 2.12.1 [64] Python package, in particular its high-level keras interface.

The loss function which is used to evaluate the performance of the model is the mean squared error between the initial wave functions $\psi$ obtained either analytically or numerically through DVR and the predicted wave functions $\widehat{\psi }$, which are the output of the neural network

$$l(\widehat{\psi },\psi )=\mathrm{MSE}(\widehat{\psi },\psi ).$$

(36)

One trains the neural network with an early stopping condition which relies on the validation loss not improving for a given amount of epochs (by default 20), as well as with the initial learning rate of $0.0005$ adjusted by the Adam algorithm. We summarise all of this information in

Table 1. Table summarising the datasets, neural network architecture, loss function, and training process used.

Step of Machine Learning Process	Detailed Summary
Training and validation datasets	Polynomial–wave function pairs evaluated on grid points; polynomial degree: 5; number of grid points: 200; training data size: 3350; validation data size: 1650
Test datasets	Function–wave function pairs evaluated on grid points; polynomial subset size: 10; harmonic oscillator subset size: 10; Morse subset size: 10
Neural network	Five layers; input size: 200; output size: 200; hidden layer output sizes: 256, 256, 128, 128. activation functions: ReLU, ReLU, ReLU, ReLu, linear; dropout: after hidden layers
Loss function	Mean square error
Training	Learning rate: 0.0005; early stopping; gradient descent algorithm: Adam; batch size: 64

. Afterwards, the loss of the data with polynomial potentials can be examined, as can the graphs of the original and predicted wave functions be inspected visually (c.f. Figure 9). The wave functions are really close, as one would expect. We also examine the energies of the original and predicted wave functions, using the discrete analogue of Equation (3):

$$\begin{array}{cc}\hfill \langle E\rangle & \approx \Delta x\sum _{i,j=0}^{N-1}\psi {\left(x_i\right)}^{*}(T_{ij}+V_{ij})\psi \left(x_j\right)=\sum _{i,j=0}^{N-1}{c}_i^{*}(T_{ij}+V_{ij})c_j.\hfill \end{array}$$

(37)

In addition to the polynomial potentials, we also examine the performance of the neural network on harmonic and Morse potentials. As expected, the model performs worse on the Morse potential compared to polynomial potentials—the losses are approximately an order of magnitude higher. This is also evident from the wave functions; examples are shown in Figure 10 and Figure 11. This is because the Morse potential (7) is described by an analytic function and therefore admits an expansion in terms of an infinite series. Because the neural network was trained only on polynomials of the fifth degree, it does not generalise sufficiently well. One can improve the score for the Morse potential by using higher-order polynomials in the training data—however, it is surprising that the neural network is capable of relatively close predictions even when trained on fifth-degree polynomials. Overall, we conclude that the machine learning model trained on relatively low-degree polynomials can capture the wave functions of the Morse potential (which are given by exponential functions). We summarise these results in

Table 2. A table of the losses for the perceptron evaluated on various datasets using the procedure depicted in Figure 7. Note that the losses of the training and validation data for the polynomial potentials are given for 400 epochs (see Figure 8).

	Training Data	Validation Data	Test Data
	Polynomial Potentials	Polynomial Potentials	Polynomial Potentials	Harmonic Oscillator	Morse Potential
Loss (10−6)	2.60	6.07	3.49	4.05	35.7

.

6. Conclusions

In this paper, we described a Python-based Jupyter notebook developed with the purpose of pedagogically introducing concepts from both quantum mechanics and machine learning. We introduced various types of Schrödinger equations (depending on the potentials) and an efficient way to numerically obtain solutions to them using the DVR method. This constituted the initial dataset which we used to train a fully connected neural network, which predicted the wave function given a potential as an input. The developed Jupyter notebook is employed during the classes within the Erasmus Mundus Quarmen Master programme. During the programme, students learn the basics of machine learning by first implementing a single-layer perceptron and then multi-layer ones in Python. Afterwards, they are acquainted with the keras interface from the tensorflow package. Students also learn how to implement numerical solvers which can be used to study the quantum problems described by the Schrödinger equation. The notebook presented with this paper is used to combine the ideas which were previously learned by the participants from the programme. We hope that this type of introduction will be beneficial to students, who can later tackle fully fledged and realistic problems involving applying machine learning to quantum physics and quantum chemistry.