Physics-Informed Neural Network for simulating two interacting quantum particles in an isolated 2D system by embedding the time-dependent Schrodinger equation directly into the network's loss function.
Traditional numerical solvers for quantum mechanics (split-step Fourier, Crank-Nicolson, exact diagonalization) compute the wavefunction at every grid point at every time step. They work, but the cost scales with grid resolution and time horizon, and each new initial condition requires a fresh solve from scratch.
This project takes a different approach: a fully connected feedforward neural network learns the time-evolution operator itself. The time-dependent Schrodinger equation is not solved externally and fed to the network as labels alone — it is embedded directly into the loss function as a physics residual (a Physics-Informed Neural Network, or PINN). The network is simultaneously penalized for:
- Data mismatch — deviation from numerically generated ground-truth wavefunction snapshots.
- PDE violation — residual of the Schrodinger equation evaluated on the network's own predictions.
- Boundary condition violation — amplitude leaking through the reflective walls.
This composite loss forces the network to learn physically consistent dynamics rather than merely interpolating training frames. Once trained, the network produces the next wavefunction state in a single forward pass — no iterative time-stepping required.
- Two quantum particles (A and B) confined to a square 2D domain with perfectly reflective walls (Dirichlet boundary conditions: wavefunction amplitude is zero at the boundary).
- The particles interact through a density-dependent mean-field potential: each particle's evolution depends on the probability density of the other, creating coupled, nonlinear dynamics.
- Probability is normalized to 1 inside the bounded region at every time step — a hard physical constraint enforced during both data generation and training.
A standard supervised network trained only on MSE can fit the training distribution but has no reason to respect conservation laws, boundary conditions, or the PDE outside the training samples. The physics residual term acts as a regularizer grounded in first principles: it penalizes outputs that violate the Schrodinger equation regardless of whether a matching label exists, improving generalization and physical plausibility.
The trained network achieves 99.9% agreement with benchmark analytical/numerical solutions, measured as the complement of the mean percentage error between predicted and ground-truth wavefunction snapshots across the test horizon.
The simulation produces probability density maps for both particles evolving over time. Below is an animated visualization of the two confined particles:
Probability densities |psi_A|^2 and |psi_B|^2 evolving inside the 2D box with reflective walls. Generated by the prototype simulator (Simulation/ProtoType/Simulator.py).
The Results/ directory contains the full training record:
losses.png— Total composite loss vs. training epoch, showing convergence behavior.error.png— Percentage error vs. training epoch, tracking the gap between predicted and true wavefunctions.plots/— Actual-vs-predicted wavefunction snapshots saved every 30 epochs (all four channels: real and imaginary parts of both particles), allowing visual inspection of how the network's predictions improve over training.
Training data is generated by numerically integrating the coupled Schrodinger equations using forward Euler on a discrete spatial grid:
d psi_A / dt = (i hbar / 2m) * nabla^2 psi_A - i * alpha * |psi_B|^2 * psi_A
d psi_B / dt = (i hbar / 2m) * nabla^2 psi_B - i * alpha * |psi_A|^2 * psi_B
where alpha controls the interaction strength (positive = repulsion, negative = attraction).
Data generator (Simulation/Generate Data/Sim_torch.py):
- GPU-accelerated via PyTorch with CUDA support.
- 64x64 spatial grid, dt = 0.0005, 1000 time steps per sample.
- Initial conditions: Gaussian wave packets with randomized positions, momenta (
kx,ky), and interaction strengths (alphain [-5, 5]). - Generates 100 independent simulations saved in HDF5 format.
- Laplacian computed via 5-point finite-difference stencil with Dirichlet (zero-padding) boundaries.
- Wavefunctions are re-normalized at every time step to enforce probability conservation.
Prototype simulator (Simulation/ProtoType/Simulator.py):
- NumPy-based CPU implementation at higher resolution (128x128 grid).
- Same physics, used for visualization and GIF generation.
The PINN is a fully connected feedforward network (Model/pinn.py):
Input: [4, 32, 32] (real + imaginary parts of psi_A and psi_B)
|
Flatten --> 4096
|
Linear(4096, 1024) + GELU
|
Linear(1024, 1024) + GELU
|
Linear(1024, 4096)
|
Unflatten --> [4, 32, 32]
Output: predicted next-timestep wavefunction
The network takes a single wavefunction snapshot (4 channels: Re(psi_A), Im(psi_A), Re(psi_B), Im(psi_B)) and predicts the next time step.
The total loss is a weighted sum of three terms:
L_total = 10000 * L_MSE + 0.0001 * L_PDE + L_BC
| Term | Role | Computation |
|---|---|---|
| L_MSE | Data fidelity | Mean squared error between predicted and true next-frame wavefunctions |
| L_PDE | Schrodinger residual | Finite-difference Laplacian + finite-difference time derivative plugged into the Schrodinger equation; the residual's magnitude is penalized |
| L_BC | Boundary conditions | Sum of squared amplitudes at all four walls — drives edge values to zero (reflective walls) |
The heavy weighting on MSE (10,000x) anchors the network to the ground truth, while the PDE residual provides a physics-consistent gradient signal even in regions where the MSE gradient is small. The boundary loss explicitly enforces the hard-wall constraint.
| Parameter | Value |
|---|---|
| Optimizer | Adam |
| Learning rate | 1.2e-5 |
| Epochs | 600 |
| Batch size | 5 |
| Training samples | 5 simulations, 10 time steps each |
| Hardware | CUDA GPU |
Actual-vs-predicted comparison plots are saved every 30 epochs to Results/plots/ for visual monitoring.
Quantum-Simulator/
├── Model/
│ └── pinn.py # PINN architecture, composite loss, training loop,
│ # inline data generation, visualization utilities
├── Results/
│ ├── error.png # Percentage error over training
│ ├── losses.png # Composite loss over training
│ └── plots/ # Actual-vs-predicted snapshots every 30 epochs
│ ├── 29actualvpred.png
│ ├── 59actualvpred.png
│ ├── ...
│ └── 599actualvpred.png
├── Simulation/
│ ├── Generate Data/
│ │ └── Sim_torch.py # GPU-accelerated batch data generator (PyTorch + CUDA)
│ └── ProtoType/
│ └── Simulator.py # CPU prototype simulator (NumPy) + GIF animation
├── quantum_particles.gif # Animated simulation visualization
├── .gitattributes
└── README.md
Model/pinn.py is the central file. It contains the network definition, all three loss functions (MSE, Schrodinger residual, reflective BC), a self-contained data generator for quick experiments, and the training loop. It can be run standalone.
Simulation/Generate Data/Sim_torch.py is the production data generator — it produces larger, more varied datasets (100 samples with randomized initial conditions and interaction strengths) at higher spatial resolution (64x64) and saves them in HDF5 format for offline training.
Simulation/ProtoType/Simulator.py is the original NumPy prototype. It runs at 128x128 resolution and includes the animation code that produces quantum_particles.gif.
- Python 3.8+
- PyTorch (with CUDA recommended for GPU acceleration)
- NumPy
- h5py
- matplotlib
- tqdm
pip install torch numpy h5py matplotlib tqdmFor GPU support, install PyTorch with the appropriate CUDA version from pytorch.org.
cd Simulation/Generate\ Data/
python Sim_torch.pyThis produces quantum_wavefunctions_complex.h5 containing 100 simulated wavefunction evolutions.
cd Model/
python pinn.pyThe script generates a small inline dataset (5 samples), trains for 600 epochs, and saves loss curves (losses.png, error.png) and comparison plots to Results/plots/.
cd Simulation/ProtoType/
python Simulator.pyProduces quantum_particles.gif — the animated probability density visualization.
- Forward Euler integration. Both the data generator and the PDE residual use first-order Euler time stepping. This is simple and stable at the chosen dt, but higher-order integrators (RK4, split-step spectral) would improve accuracy for longer time horizons or stiffer interaction regimes.
- Fixed grid resolution. The PINN is trained on a 32x32 grid (the data generator produces 64x64, but
pinn.pyuses 32x32 internally). Resolution is not a learned parameter — the network would need retraining for a different grid. - Mean-field interaction only. The density-dependent coupling (
alpha * |psi_other|^2 * psi) is a mean-field approximation, not a full two-body Coulomb integral. This is a modeling simplification, not a bug — it captures the qualitative physics of interacting confined particles while keeping the computational cost tractable. - Small training set. The default configuration in
pinn.pytrains on 5 simulations of 10 time steps each (45 input-output pairs). The standalone data generator (Sim_torch.py) produces 100 varied simulations — using that larger dataset would likely improve generalization. - No rollout evaluation. The current training evaluates single-step predictions only. Multi-step autoregressive rollout (feeding predictions back as inputs) is the harder test and is not yet benchmarked.
If you use this work, please cite the preprint:
@article{goel2025quantum,
title = {Neural Simulation of Quantum Interactions in a Confined System},
author = {Goel, Arpit},
year = {2025},
note = {ResearchGate preprint 394354557},
url = {https://www.researchgate.net/publication/394354557_Neural_Simulation_of_Quantum_Interactions_in_a_Confined_System}
}This repository does not currently include a license file. All rights are reserved by the author by default. If you intend to use, modify, or distribute this code, please contact the author for permission.
Author: Arpit Goel