This repository contains a collection of physics-informed neural networks (PINNs) that are trained to solve problems in quantum chemistry, thermodynamics, and fluid dynamics by learning directly from physics rather than from labeled data. Most of the models can be trained in a few minutes using an M2 Mac. Pretrained models can be found in this Huggingface model.
- Getting Started
- What is a PINN?
- Heat Equation
- Burgers' Equation
- Schrödinger Equation
- Model Performance Summary
- Why do you need AI?
- What I Learned
- What's Next?
- Acknowledgements
- Python 3.12+
- PyTorch
- NumPy
- Matplotlib
git clone https://github.com/sr5434/pinns.git
cd pinns
pip install -r requirements.txt# Train the 3D heat equation model
cd heat_equation
python heat_equation_3D.py
# Generate visualizations
python heat_equation_visualizer_3D.pycd heat_equation
# Generate visualizations
curl -L "https://huggingface.co/sr5434/PINN-Collection/resolve/main/heat_equation_3d.pt?download=true" -o heat_equation_3d.pt
python heat_equation_visualizer_3D.py# Train the 1D Burgers' equation model
cd burgers_equation
python burgers_equation_1D.py
# Generate visualizations
python burgers_equation_visualization_1D.pycd burgers_equation
# Generate visualizations
curl -L "https://huggingface.co/sr5434/PINN-Collection/resolve/main/burgers_equation_1d.pt?download=true" -o burgers_equation_1d.pt
python burgers_equation_visualization_1d.pycd schrodingers_equation
python schrodingers_equation_1d.py
# Generate visualizations
python schrodingers_visualization_1d.pycd schrodingers_equation
# Generate visualizations
curl -L "https://huggingface.co/sr5434/PINN-Collection/resolve/main/schrodingers_equation_1d.pt?download=true" -o schrodingers_equation_1d.pt
python schrodingers_visualization_1d.pycd schrodingers_equation
python schrodingers_equation_hydrogen.py
# Generate visualizations
python schrodingers_visualization_hydrogen.pycd schrodingers_equation
# Generate visualizations
curl -L "https://huggingface.co/sr5434/PINN-Collection/resolve/main/schrodingers_equation_hydrogen.pt?download=true" -o schrodingers_equation_hydrogen.pt
python schrodingers_visualization_hydrogen.pycd schrodingers_equation
python schrodingers_equation_h2.py
# Generate visualizations
python schrodingers_visualization_h2.pycd schrodingers_equation
# Generate visualizations
huggingface-cli download sr5434/PINN-Collection --include "h2_models/*" --local-dir ./
python schrodingers_visualization_h2.py
python3 plot_h2_pes.py
python3 calculate_observables.py # No need to recompile LEVEL16 on Mac, as precompiled binaries are included in the repoTL;DR: PINNs are neural networks that learn to solve physics problems by learning from the underlying physical laws, rather than from labeled data.
PINNs are just neural networks that approximate functions described by Partial Differential Equations (PDEs). The main thing that is special about PINNs is not their architecture, but rather how they are trained. PINNs are unique because they are trained to satisfy different conditions, unlike most neural networks, which are trained to mimic labeled examples. These conditions are expressed through loss functions, which are detailed below.
The PDE loss function ensures that the model's solution is valid. Essentially, it plugs the model's solution back into the model and compares the left hand side to the right hand side, similar to how a student in Algebra might check their work after solving for a variable. All PINNs must be trained with the PDE loss.
This loss function checks that the model satisfies boundary conditions, which dictate model behavior at "boundaries". Examples of boundaries are faces of a cube or ends of a rod. Note that this loss is optional if the boundary is enforced mathematically in the model's code (all training scripts in this repository enforce boundary conditions at the model level). All models use Dirichlet boundary conditions, which mandate that the model's output at the boundary is equal to 0.
The initial conditions loss makes sure that the model's outputs at timestep 0 follow the initial conditions of the problem. The main purpose of this is to ensure the model "starts strong", as poor initial performance will only get worse over time. Like the boundary conditions, this loss is optional when the model is designed to always follow the initial conditions of the problem. Only the script for the 1D heat equation uses an initial conditions loss, with the other 2 scripts enforcing initial conditions at the architectural level.
The loss can be formalized as a weighted sum of the different losses described above:
The heat equation model places a weight of 2 on the initial conditions loss, while the other models use a weight of 1 for all losses. The weights can be adjusted to prioritize certain losses over others, depending on the problem at hand.
This implementation uses dimensionless quantities normalized to [0, 1] for numerical stability and generality. The solutions can be scaled to any physical units by applying appropriate transformations. This is standard practice in computational physics and ensures the neural network trains effectively.
heat_equation_3d_visualization.mp4
This repository contains code to train PINNs on the 1D, 2D, and 3D heat equations. It also contains code to generate visualizations from the 2D and 3D models (the 3D visualization is just a slice from the middle of a cube). The trained models predict how heat diffuses through a rod, a tile, and a cube respectively. This is the 3D heat equation (for 2D, remove the second derivative w.r.t. z, and for 1D, also remove the second derivative w.r.t. y):
Here,
- Electrical Engineering: Modeling heat flow in electronics
- Civil Engineering: Designing cooling systems that maximize energy efficiency in buildings
- Food Sciences: Simulating cooking or baking
burgers_equation_1d_visualization.mp4
The repository also contains a script to train a PINN on the 1-dimensional variant of Burgers' Equation, which predicts the instantaneous velocity of a particle in a fluid. Burgers' Equation is as follows:
Because we are solving the 1D variant of the equation, the gradient simplifies to a single derivative w.r.t.
- Public policy: Modeling traffic flow
- Acoustics: Describing sound waves
schrodinger_equation_1d.mp4
There is a script to train a model to predict the wavefunction of a quantum particle in a 1D infinite square well over time, following the Time-Dependent Schrödinger Equation:
Where
Due to the fact that this is such a small number, a value of 1 is used as a simplification to avoid an underflow.
The Hamiltonian operator returns the total energy of the quantum system(the sum of potential and kinetic energy) given the wavefunction
The squared magnitude of the quantum wavefunction can be used to estimate the probability density that when observed at a given time, a particle in quantum superposition with a certain energy level will collapse to a given location. Probability density is similar to probability, and the probability that a particle will be observed in the range
The model for Schrödinger Equation is our largest by far, with 4 hidden layers and a hidden size of 256(except for our last hidden layer, which returns a tensor with 128 channels). Also, the model takes in sinusoidal features generated based on the energy level of the particle, as this helps the model adjust to differences in oscillations between lower and higher levels. Due to the oscillatory nature of higher energy levels, the highest level our model supports is 3. The model was trained with a Cosine decay learning rate schedule that had a warm restart whenever the maximum energy level present in the data was increased(the maximum energy level was increased every 15,000 steps, and the warmup occurred slightly before this). The learning rate started at 0.001 and decayed to 0.0001. After 45,000 steps, the learning rate plateaued at the minimum. Like the other 2 models, the tanh activation function and Adam optimizer are used here. When enforcing the initial conditions, the model's raw output is scaled by
Here, L is the length of the infinite square well (which is 1 in our dimensionless system). Intuitively, this can be thought of as a metric to verify that the probabilities in the distribution generated by taking the squared magnitude of the wavefunction sum to 100%. This loss is implemented by performing Monte Carlo integration on the squared magnitude of the wavefunction at several points in space, multiplying by the width between points, and comparing that to 1.
It should also be noted that unlike the other 2 models, this model outputs 2 channels, representing the real and imaginary components of the complex wavefunction.
hydrogen_molecule_R_1.4.mp4
The repository also contains code to train and visualize a PINN that solves the Time-Independent Schrödinger Equation for the radial portion (more on this below) of Hydrogen's wavefunction. The model supports the 1s, 2s, and 2p orbitals. The Time-Independent equation was used because the probability density of hydrogen orbitals does not change over time. The Time-Independent Schrödinger Equation is as follows:
Where
Where the Laplacian operator in spherical coordinates is defined as:
And the Coulomb potential term,
This system solves the quantum eigenvalue problem for the hydrogen atom, meaning that the analytical values of E are not used in training. The exact values are present in the codebase for evaluation purposes, however. The Rayleigh quotient is used to estimate the energy level of the particle based on the model's current prediction for the wavefunction:
To improve the accuracy of the Rayleigh quotient, a loss based on the Virial Theorem is also used during training:
Where
This can be thought of as a continuous average of the operator over all space, weighted by the probability density of the particle being at each location.
Similar to the model for the TDSE, the magnitude loss is defined as follows:
To prevent the 2s and 1s orbitals from collapsing to the same state, an orthogonality loss is defined as follows:
This loss is only enforced between 1s and 2s, and is not enforced for any other orbital combinations.
The model architecture is the same as the one used for the Time-Dependent Schrödinger Equation, except that the input layer has been modified to accept 3 channels instead of 5. The model was trained on 600,000,000 samples using the Adam optimizer with the same cosine schedule as the Time-Dependent model, but without warm restarts. In each step, there were 3 sets of collocation points: general points, central points, and deterministic points. General points were sampled from a Gamma distribution and had a maximum radius of 30. Central points were uniformily sampled within a 3 unit radius to prioritize coverage where most of the action was happening. Deterministic points were sampled from a fixed grid. General and central points were only used for calculating the residual loss, and the deterministic points were used for all other losses. Trapezoidal integration was selected for all integrals because it is more accurate than Monte Carlo integration. Unlike the other models, this model uses spherical coordinates as inputs. The model was evaluated by comparing its results to the analytical solutions for the 1s, 2s, and 2p orbitals. Error is measured with mean absolute error between the predicted and analytical radial wavefunctions.
To use the model to estimate the wavefunction of a dihydrogen cation(
By default, the visualization script renders all of the Hydrogen orbitals supported by the model and the bonding and antibonding molecular orbitals of
hydrogen_molecule_R_1.4.mp4
The Schrödinger Equation for the hydrogen molecule is a more complex problem than the hydrogen atom, as it has 2 electrons instead of 1. Because of this, it is a many body problem and has no exact analytical solution. The Hamiltonian operator for the hydrogen molecule is defined as:
Where
The model also used a rayleigh consistency loss to penalize the model for sampling different energies when fed different sobol sets. No orthogonality loss was used because the model was only trained on the ground state, but virial loss was kept and an energy minimization loss was added to encourage the model to find the lowest energy solution. The energy minimization loss compared the model energy to the exponential moving average of step energies to calculate an advantage, and this advantage was clipped (sort of like in policy gradient reinforcement learning) to prevent the model from sacrificing physical accuracy for lower energy. The minimization term was also gated based on how well the model followed normalization, consistency, and virial losses. The model used an LCAO and Jastrow (three-body) ansatz to help the model learn the correct wavefunction. The log of
Where schrodingers_equation/assets folder. There is also data for an ablation with just energy and normalization losses, which as expected, performs very poorly.
- Physics Research: Modeling cold atom traps
- Chemistry: Understanding atomic structures and reactions
- Materials Science: Calculating properties of materials, such as specific heat capacity
| Model | Architecture | Training Samples | Max Error |
|---|---|---|---|
| Heat 1D | 1 layer, 50 hidden | 20M | <1% |
| Heat 2D | 1 layer, 50 hidden | 75M | <1.5% |
| Heat 3D | 1 layer, 50 hidden | 170M | <4% |
| Burgers' 1D | 2 layers, 100 hidden | 50M | <4.5% |
| Schrödinger 1D | 4 layers, 256/256/128 hidden | 400M | 1% to 6.5% depending on energy level1 |
| Schrödinger Hydrogen atom | 4 layers, 256/256/128 hidden | 600M | <$10^{-4}$ MAE for all cases |
| Schrödinger Hydrogen molecule | 4 layers, 256/256/128 hidden | 679.44M | <1 mHa for most distances, better than chemical accuracy for all distances |
It is true that analytical solutions to the heat equation, Burgers' Equation, and Schrödinger Equation are far more efficient than using a PINN. However, there are many unique attributes that make PINNs useful. For example, given the outputs of the model and all spatial/temporal inputs, it is possible to solve for the thermal diffusivity of an object, the viscosity of a fluid, or the energy level of a particle. Also, in more complex scenarios, analytical solutions may not exist, meaning PINNs are the only way to approximate the solution to a PDE. This is especially true in quantum mechanics, where even simple systems like a Helium atom are difficult to solve numerically or analytically. PINNs also have the advantage of being mesh-free, meaning they can make predictions at any point in space and time without needing to be retrained or interpolated.
I learned a lot about physics and multivariate calculus from doing this project. This project also helped me realize how simple natural concepts like heat diffusion (which require a couple thousand parameters to model) are compared to man-made constructs like language (which require billions or trillions of parameters to model effectively).
Working on this project brought back some nostalgia for a time when I was very passionate about physics, and it made me feel as if I was reconnecting with my past self.
- Scale up Burgers' Equation to 2D and 3D
- Implement more complex PDEs, such as the Navier-Stokes Equations
- Experiment with newer optimizers like Muon
- Enable inverse problems, where the model solves for physical constants given observations of a system
I want to thank Krivan Semlani for inspiring me to work on PINNs and encouraging me to keep up the work. I also want to thank Prakash Adhikari and Aryan Senthilkumar for helping me understand some of the physics behind Hydrogen's wavefunction. Finally, I want to thank Professor George C. McBane for suggesting that I use LEVEL16 to solve the nuclear Schrödinger equation for the hydrogen molecule.
Footnotes
-
The error for the Schrödinger Equation model varies based on the energy level of the particle. Lower energy levels tend to have lower error, while higher energy levels exhibit higher error due to their increased oscillatory behavior. ↩