sepinn 1.0.0

Solving the Schrödinger Equation via Physics-Informed Machine Learning

Summary

SE-PINN is a physics-informed neural network in PyTorch that solves 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.

SE-PINN was developed at Vanderbilt University in collaboration with Alexander Ahrens and under the supervision of Prof. Ipek Oguz (https://engineering.vanderbilt.edu/bio/?pid=ipek-oguz).

---

Demonstration

Figure 1 and Figure 2 are visualizations 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 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: SE-PINN without Enforcement of Symmetry

Figure 2: SE-PINN with Enforcement of Symmetry

---

Usage

(1) Install

pip install sepinn

(2) Import

from sepinn.wrappedpinn import WrappedPINN

model = WrappedPINN(...)

model.train(...)

---

Documentation

A Jupyter notebook is available for reference in the docs folder as well as through Google Colab and nbviewer.

Google Colab (Interactive):

https://colab.research.google.com/github/Tiger-Du/SE-PINN/blob/main/docs/quantum_harmonic_oscillator.ipynb

nbviewer (Non-interactive):

https://nbviewer.org/github/Tiger-Du/SE-PINN/blob/main/docs/quantum_harmonic_oscillator.ipynb

---

Citation

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}
}

---

License

SE-PINN is published under the GPL-3.0 license.