Machine Learning Speeds Up and Enhances Physics Calculations

October 1, 2020

Interpreting data from NIF’s cutting-edge high energy density (HED) science experiments relies on physics calculations that are so complex they can challenge Lawrence Livermore National Laboratory supercomputers, which stand among the best in the world.

Now a close collaboration between researchers from LLNL and France’s Alternative Energies and Atomic Energy Commission (CEA) and Defense Procurement Agency-Defense Innovation Agency (DGA-AID) has resulted in a novel way to incorporate machine learning and neural networks to significantly speed up those calculations. This opens a path to achieving higher-quality physics simulations without increasing computing costs and could help shape the next generation of LLNL supercomputers.

“We’re trying to understand what the computer platform of the future is going to look like at Livermore,” said LLNL design physicist Brian Spears. “We looked at how supercomputing platforms and equipment that we’re buying now are going to look different in 2028 and beyond once we equip them with artificial intelligence (AI) workloads.”

LLNL’s Sierra supercomputer.

A paper published as an Editor’s Pick in Physics of Plasmas detailed “the first proof positive demonstration that this is actually going to work,” said Spears, one of several co-authors from LLNL. The lead author was Gilles Kluth, an applied mathematician for ICF simulations for France’s Laser Mégajoule.

“The research explored incorporating AI and neural networks into one of the most complex and expensive computer physics simulation codes–the non-local thermodynamic equilibrium (NLTE) atomic physics package. 

NLTE Within HYDRA

HYDRA is one of the multi-physics radiation hydrodynamics codes used to calculate hydrodynamics, neutronics, magnetic fields, and other physical processes occurring in and around an imploding target during the nanoseconds it takes for an indirect-drive inertial confinement fusion (ICF) experiment on NIF, the world’s largest, most energetic laser. 

Used within HYDRA, the NLTE package computes the opacity, emissivity, ionization, and equation of state of the hot plasma when those physical properties depend on details of the surrounding radiation. This produces a large volume of data that can take a day or two of computation time even on LLNL’s advanced high-performance computing systems.

The researchers wanted to see if they could tap the power of machine learning to do those calculations once instead of every time.

“Often times, we’re probably calculating the exact same thing thousands of times, but we could just be using those data to train the neural net instead and saving ourselves a lot of computational time down the road,” said LLNL design physicist Kelli Humbird, another co-author.

Reducing Computer Time

They trained a deep neural network to emulate the expensive, NLTE atomic kinetics calculation that is typically performed for every NLTE zone on every time step in HYDRA.

About 95 percent of the HYDRA code runs as it normally would. “But for the one piece of the calculation where it says tell me what the state of this atom is and what its optical properties are, how much radiation should it absorb and how much should it emit, rather than calling a code to do that, it calls our AI tool,” Spears said.

The result produced a tenfold reduction in overall computer processing time, “with significant room for further speed-up,” the researchers said. “This method for accelerating physics calculations also offers a straightforward path to including higher fidelity physics without changing the computational cost: the neural network models used in this work can be trained on data produced with more detailed atomic models.”

While the NLTE package comprises only a fraction of HYDRA, it’s one of the most computationally expensive parts of the calculation, easily taking up about 20 to 25 percent of the total computational costs. And the most accurate models take up about 90 percent of computational costs.

“Because it’s arguably the meatiest part of the whole code, speed-ups there translate into big wins overall,” added co-author and design physicist Luc Peterson.

Co-author Mehul Patel, a computational physicist with LLNL’s Weapons and Complex Integration (WCI) Directorate, explained the neural network is trained using the known results of previous atomic physics calculations.

“It’s devising its own connections so that when you give it an input without an output, it can predict what the output is,” said Patel. “And then you can compare that to doing the actual calculation to see how well it did.”

The Lassen supercomputer at LLNL.

DJINN Algorithm

The groundwork for the research started several years ago with Humbird’s work developing a Deep Jointly Informed Neural Networks (DJINN) algorithm as part of her PhD work at the Lab under the Livermore Graduate Research Scholar Program. Her advisor was Peterson.

“She authored the seminal paper that allowed this work to go forward,” Spears said.

DJINN automatically choses the neural network architecture that is appropriate for the data being computed without requiring the user to manually tune the settings. Spears compared DJINN to a derailleur that changes a bicycle chain to the appropriate gear.

“A lot of the challenge associated with building neural networks is figuring out how to tune all of the parameters,” Spears said. “They are very massive models that can be as complicated or as simple as you want them to be.”

Humbird said the stack of software makes it easy for anyone with a little bit of experience running a physics code to train the models.

“The real punch line of this story was that we trained these neural network models to essentially replace a physics calculation within a big multi-physics code,” she said. “At the end of the day, the user of this neural network doesn’t even necessarily need to know that they’re using a neural network.”

Teaming with CEA DGA-AID

Thousands of miles away from Livermore, the research took another step thanks to the CEA’s Gilles Kluth. “I was wondering how I could use the huge progress of deep learning to improve ICF simulations,” Kluth said.

Then he read a paper by Peterson published in Physics of Plasmas describing the use of machine learning to find optimal designs for ICF experiments (see “Machine Learning Points Toward New Laser Target Designs”).

Kluth reached out to Spears, the group leader, then arranged for funding from the CEA DGA-AID to work with the LLNL team starting in the late summer of 2018. He and Humbird started discussing how her work on DJINN could fit in with the interaction of plasma and radiative fields.

“Gilles wanted to learn machine learning,” Humbird said. “He came up with this idea of working with opacity data, something we had thought about. It was really a close collaboration with him to actually get this project built to the point where we could publish a satisfying paper to show this wasn’t just a crazy idea.”

Physicists can also use more data to run higher-quality simulation models. Before, they had to scale back data because they would require too much computer time.

“We typically use what we might call moderate fidelity atomic physics,” Spears said. “So we’ve turned one day to a half day with moderate fidelity. But if you’re willing to do all the capital investment up front, we can use extremely high fidelity physics models and run the simulation in a half day.”

Cognitive Simulation

The effort is part of the Lab’s broader Cognitive Simulation Director’s Initiative, led by Spears and funded through the Laboratory Directed Research and Development (LDRD) program.

Recently, LLNL scientists published related research on surrogate models supported by neural networks, combining simulations with HED experiments using AI to improve predictions (see “Deep Learning Surrogate Models Could Hasten New Scientific Discoveries”).

LLNL’s El Capitan is projected to be the world’s most powerful supercomputer when it is fully deployed in 2023. Illustration courtesy of Hewlett Packard Enterprise.

The neural net technology has also played a role in Lab-wide research to develop simulations on how COVID-19 spreads. That team effort, which is being used to by military and government decision-makers, recently won an award from the Weapons and Complex Integration Directorate that recognizes outstanding contributions and achievements in support of the Lab’s national security missions.

The surrogate models and looping AI into the HYDRA workflow described in Kluth’s paper are pieces of a bigger, long-term vision that could produce faster simulations that can be upgraded before, during, and after simulations.

“The interesting story is that we have this large director’s initiative in cognitive simulation and have made some targeted forays into particularly strategic projects,” Spears said. “This is a good example of one that has paid really nice dividends and it’s seeding ideas across the laboratory, especially for NIF and beyond.”

Both Humbird and Kluth are now working on separate projects to build on the research in tangential directions.

“What he showed in that one paper is a very promising start to what can really revolutionize how we run our big codes in the future,” said Humbird. “There’s a lot of potential to make these codes that have already been getting faster — just blazingly fast — and enable us to do better physics. We’re just getting started on this.”

Joining Kluth, Humbird, Spears, Patel, and Peterson on the paper, “Deep Learning for NLTE Spectral Opacities,” were LLNL’s Howard Scott,  Joe Koning, Marty Marinak, Laurent Divol, and Chris Young.

More information:

“The Internship that Launched a Machine-learning Target Revolution,” NIF & Photon Science News, October 1, 2020

“LLNL and HPE to Partner with AMD on El Capitan, Projected as World’s Fastest Supercomputer,” LLNL News Release, March 5, 2020

—Benny Evangelista and Charlie Osolin

Follow us on Twitter: @lasers_llnl