SE-PINN

Webpage | GitHub | PyPI

SE-PINN is a physics-informed neural network in PyTorch that solves the Schrödinger equation of quantum mechanics.

Through (a) a custom loss function and (b) a custom neural-network layer, SE-PINN is constrained to predict quantum-mechanical states that satisfy the mathematical-physical properties of symmetry, normality, and orthogonality. In addition, the model learns not through supervised learning on labeled data but through unsupervised learning via feedback from the Schrödinger equation itself.

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

The design of SE-PINN is based on https://arxiv.org/abs/2203.00451 and https://arxiv.org/abs/1904.08991.

In addition, SE-PINN supports the following features: - L-BFGS optimization - A class that supports deterministic training, checkpointing of training, and visualization - A Python package on PyPI with CI/CD (pip install sepinn) - Documentation as a webpage and a Google Colab notebook

Example

Figure 1 and Figure 2 are both visualizations of the ground state (left) and the energy of the ground state (right) that are predicted by SE-PINN as it trains. The physical system of interest here is the quantum harmonic oscillator, which can be used to model diatomic molecules such as diatomic nitrogen, diatomic oxygen, and the hydrogen halides.

As visualized in Figure 2, constraining SE-PINN via a special architectural layer — a hub layer — to predict a ground state that satisfies exact symmetry improves its convergence to both the correct ground state and the correct energy.

| Figure 1: SE-PINN without Enforcement of Symmetry | | --- | | |

| Figure 2: SE-PINN with Enforcement of Symmetry | | --- | | |

Usage

1. Install SE-PINN.

pip install sepinn

2. Import SE-PINN.

```python 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/SE-PINN/SE-PINN/blob/main/docs/quantumharmonicoscillator.ipynb

nbviewer (Non-interactive):

https://nbviewer.org/github/SE-PINN/SE-PINN/blob/main/docs/quantumharmonicoscillator.ipynb

Citation

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