This is the official repository for Solving the Schrödinger Equation via Physics-Informed Machine Learning.
Here a physics-informed neural network is developed in PyTorch for solving the Schrödinger equation of quantum mechanics. The model is constrained to predict quantum-mechanical states that respect the mathematical-physical properties of symmetry, normality, and orthogonality — all via (1) a custom loss function and (2) a custom architectural layer. In addition, the model learns not through supervised learning but through reinforcement learning (RL) via feedback from the Schrödinger equation itself.
This research was in collaboration with Alexander Ahrens and under the supervision of Prof. Ipek Oguz (https://engineering.vanderbilt.edu/bio/ipek-oguz) at Vanderbilt University.
Figure 1 and Figure 2 are animations of the ground state (left) and the energy of the ground state (right) that are predicted by the model as it trains. The physical system of interest is the quantum harmonic oscillator, which is used to model diatomic molecules such as diatomic nitrogen, diatomic oxygen, and the hydrogen halides.
The enforcement of exact symmetry on the prediction of the ground state via a special architectural layer of the model — a "hub layer" — improves its convergence to the correct energy, as visualized in Figure 2.
| Figure 1: Without Architectural Enforcement of Exact Symmetry |
|---|
 |
| Figure 2: With Architectural Enforcement of Exact Symmetry |
|---|
 |
Installation from PyPI via pip:
pip install sepinn
Usage in Python:
from sepinn.wrappedpinn import WrappedPINN
model = WrappedPINN(...)
model.train(...)
A Jupyter notebook is available for reference in the docs folder as well as through Google Colab and nbviewer.
Google Colab (Interactive):
nbviewer (Non-interactive):
https://nbviewer.org/github/Tiger-Du/SE-PINN/blob/main/docs/quantum_harmonic_oscillator.ipynb
SE-PINN is citable via the BibTeX entry below.
@techreport{DuAhrensOguz2023,
author={Du, Tiger and Ahrens, Alexander and Oguz, Ipek},
institution={Vanderbilt University},
title={Solving the Schrodinger Equation via Physics-Informed Machine Learning},
year={2023}
}
SE-PINN is available under the GPL-3.0 license.