Hyperbolic Machine Learning Moment Closures for the BGK Equations
Andrew J. Christlieb Mingchang Ding Juntao Huang Nicholas A. Krupansky
We introduce a hyperbolic closure for the Grad moment expansion of the
Bhatnagar-Gross-Krook's (BGK) kinetic model using a neural network (NN) trained
on BGK's moment data. This closure is motivated by the exact closure for the
free streaming limit that we derived in our paper on closures in transport
\cite{Huang2022-RTE1}. The exact closure relates the gradient of the highest
moment to the gradient of four lower moments. As with our past work, the model
presented here learns the gradient of the highest moment in terms of the
coefficients of gradients for all lower ones. By necessity, this means that the
resulting hyperbolic system is not conservative in the highest moment. For
stability, the output layers of the NN are designed to enforce hyperbolicity
and Galilean invariance. This ensures the model can be run outside of the
training window of the NN. Unlike our previous work on radiation transport that
dealt with linear models, the BGK model's nonlinearity demanded advanced
training tools. These comprised an optimal learning rate discovery, one cycle
training, batch normalization in each neural layer, and the use of the
\texttt{AdamW} optimizer. To address the non-conservative structure of the
hyperbolic model, we adopt the FORCE numerical method to achieve robust
solutions. This results in a comprehensive computing model combining learned
closures with methods for solving hyperbolic models. The proposed model can
capture accurate moment solutions across a broad spectrum of Knudsen numbers.
Our paper details the multi-scale model construction and is run on a range of
test problems.