1、Graph Neural Network for Large-scale SimulationsCollaboratorsStanford:Tailin Wu,Sophia Kivelson,Yinan Zhang,Jacqueline Yau,Jure LeskovecSLAC(Accelerator Laboratory):Jason Chou,Frederico FiuzaDeepMind Inc.Saudi AramcoUCLA:Paulo Alves1June 25,2022Rex YingOutlineSimulations in Science and Engineering,a

2、nd CharacteristicsGraphics and Physics SimulationReservoir SimulationScientific Simulation:Laser-plasma Interaction Simulation2OutlineSimulations in Science and Engineering,and CharacteristicsGraphics and Physics SimulationReservoir SimulationScientific Simulation:Laser-plasma Interaction Simulation

3、3Simulations in Science and Engineering1.Particle-particle interactions:Galaxy formationWater simulation4We construct graphs with particles as nodes and interactions as edgesSimulations in Science and EngineeringWeather predictionAerodynamics2.PDEs(on grid or mesh)Reservoir simulation5Graph based on

4、 grid/mesh structureSimulations in Science and Engineering3.Particle-in-Cell(involves both grid and particles)FusionLaser-plasma particle acceleration6Cosmic-ray accelerationGraph combining grid/mesh with particles!Simulations in Science and EngineeringCharacteristics:Large scale in size:at the fore

5、front of HPCNevertheless,even those large compute with long-time simulation may only do reasonably small systems in practiceE.g.for a reasonable 3D laser-plasma interaction system,it has 100B grid vertices,1T particles,over 100k time stepsLargest simulations(1/year):10-1 of that scale,most studies:O

6、(n2)Physics solution:also discretize the field and simulate the field evolution,so the update becomes local =Particle-in-Cell(PIC)method36How to simulate the interactions between ions,electrons and field?(1)Interpolate the field at the particle location:37How to simulate the interactions between ion

7、s,electrons and field?(2)Advance the particle velocity and position:38How to simulate the interactions between ions,electrons and field?(3)Deposit the current from particle to the field:39How to simulate the interactions between ions,electrons and field?(4)Update the field based on the neighboring f

8、ields:40Re-implement the above PIC using a heterogeneous GNN,test its performanceSwitch each components one-by-one to a learnable heterogeneous GNN,test its performance41Learning the dynamics with GNNLearning the dynamics with GNN(1)We re-implement the ground-truth solver in hetero-GNN format We see

9、 that even if we are at machine precision(10-11)for single step prediction,the error will grow exponentially and to discernible scale(10-1)in 200 steps 42Learning the dynamics with GNN(1)We re-implement the ground-truth solver in hetero-GNN format Electron rollout at 50,100,200 steps:Electron ground

10、-truth at 50,100,200 steps:We see that even if the RMSE of electron positions at 700 step is 10-1,the distribution is quite similarX-axis:positionY-axis:momentumColor:density43Learning the dynamics with GNN(1)We re-implement the ground-truth solver in hetero-GNN format Ion rollout at 50,100,200 step

11、s:Ion ground-truth at 50,100,200 steps:44Learning the dynamics with GNN(1)We re-implement the ground-truth solver in hetero-GNN format Field rollout at 50,100,200 steps:Field ground-truth at 50,100,200 steps:45X-axis:positionY-axis:amplitudeColor:Field valuesLearning the dynamics with GNN(1)We re-im

12、plement the ground-truth solver in hetero-GNN format 46Energy spectrum(solid:rollout,dashed:ground-truth):Learning the dynamics with GNN(2)We replace each ground-truth step by learnable module one at a time(2.1)Replacing the first interpolation part by learnable MLP:47Learning the dynamics with GNN(

13、2.1)Replacing the first interpolation part by learnable MLP:Electron rollout at 50,100,200 steps:Electron ground-truth at 50,100,200 steps:48Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Ion rollout at 50,100,200 steps:Ion ground-truth at 50,100,200 steps

14、:49Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Field rollout at 50,100,200 steps:Field ground-truth at 50,100,200 steps:50Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Energy spectrum(solid:rollout,dashed:gro

15、und-truth):51Learning the dynamics with GNN(2)We replace each ground-truth step by learnable module one at a time(2.2)Replacing all the components by learnable MLPs:52Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Electron rollout at 50,100,200 steps:Elect

16、ron ground-truth at 50,100,200 steps:53Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Ion rollout at 50,100,200 steps:Ion ground-truth at 50,100,200 steps:54Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Field ro

17、llout at 50,100,200 steps:Field ground-truth at 50,100,200 steps:55Learning the dynamics with GNN(2.1)Replacing the first interpolation part by learnable MLP:Energy spectrum(solid:rollout,dashed:ground-truth):56OutlineSimulations in science and engineering,and characteristicsGraphics Simulation and

18、Reservoir SimulationProblem SettingIntroduction of Graph Neural Simulator(GNS)Results Physics Simulation:Laser-plasma Interaction SimulationProblem SettingLearning the simulationLearning to speedup57How to speed up?58non-thermal particles thermal particles(Gaussian distribution in velocity)How to sp

19、eed up?Only small percentage of the particles(non-thermal)are accelerated but can carry 10-50%of system energy,while the majority are thermal,whose velocity distribution follows a Gaussian distribution.We think a good way to speed up is to simulate non-thermal particles with particle description,and

20、 thermal particles with a reduced description.We still maintain the field grid,and only compress the thermal particles.One major challenge:how to promote thermal particles in reduced description to non-thermal particles59How to speed up?How to represent the thermal particles with reduced description

21、?Fluid descriptionAt each field cell,use the mean velocity and higher-order moments to represent the thermal particles within the cellUse CNN/GNN to evolve the fluid,while use the GNN to evolve those non-thermal particles60Implementation in progress,anticipated to achieve 100 fold speedupSummaryLarg

22、e-scale simulations are important across many science and engineering domains.GNNs provides a natural way to learn the evolution of such systems.Our work on using GNS/hetero-GNS on graphics simulation,reservoir simulation and laser-plasma interaction show their capability in learning complex dynamics and generalization.61