Replacing Physical Engines with Graph Neural Networks: A Neural Surrogate Model for N-Body Electrostatic Simulation

This project builds a neural surrogate model based on Graph Neural Networks (GNNs) to predict the motion states of N-body charged particles under Coulomb repulsion and gravity. It achieves real-time prediction 50 times faster than traditional physical simulation with high short-term accuracy. The project is open-sourced on GitHub, and its core idea aligns with applications by institutions like DeepMind and NVIDIA in climate modeling and drug discovery.

The computational complexity of the traditional N-body problem is O(N²), requiring repeated solving of physical equations, which is costly. The core of neural surrogate models is to train neural networks to learn physical laws, replacing expensive equation solving with fast prediction. Their advantages include high speed, differentiability (facilitating optimization/control), and generalization ability; the challenge is the accumulation of errors in long-term simulation.

This project's technical architecture consists of four parts:

1. Physical Simulator: Implemented based on NumPy, including softened Coulomb repulsion, gravity, Euler integration, and elastic boundaries, used for generating training data and baseline comparison.
2. Dataset: 500 independent simulations (500 steps each) generate 250,000 samples, stored in HDF5 format.
3. GNN Architecture: A message-passing network implemented with PyTorch Geometric. Particles are nodes (features: vx, vy, charge, mass), edges are particle pairs (features: dx, dy, distance), predicting Δvx, Δvy, Δx, Δy at the next moment.
4. Training Details: MSE loss (normalized changes), Adam optimizer (lr=0.0001), gradient clipping, feature normalization, and training tracked with Weights and Biases.

Training Performance: Initial loss was 0.483, dropping to 0.234 after 100 epochs. Error Analysis: MSE was approximately 0.0002 at step 0, ~4.0 at step 100, and ~13.0 at step 200. Key Findings: Short-term predictions are accurate, errors accumulate with steps, inference speed is 50 times faster than physical simulation, and the error pattern is consistent with DeepMind's 2020 research.

Application scenarios include:

1. Real-time Simulation and Games: Enabling real-time rendering of large-scale particle interactions (sandbox games, special effects).
2. Parameter Optimization: Accelerating the search for different parameter configurations (e.g., optimal initial charge distribution).
3. AI System Integration: Differentiability supports end-to-end systems like reinforcement learning and control optimization.
4. Educational Research: Serving as a learning case for GNNs and scientific machine learning.

Current Limitations: Increased error in long-term simulation (>100 steps), fixed number of particles, simplified physical model (no collision/friction). Improvement Directions: Try Graph Transformer/Equivariant GNN; multi-scale modeling; uncertainty quantification; add energy conservation constraints; transfer learning to new systems.

Project Value:

• Education: Clearly demonstrates the application of GNNs in scientific computing, serving as excellent material for getting started with scientific machine learning.
• Research: Validates the feasibility and challenges of neural surrogate models in physical simulation.
• Engineering: Provides an extensible code framework.
• Inspiration: Shows the potential of AI to change the paradigm of traditional scientific computing.