Uniformly Accurate Machine Learning Based Hydrodynamic Models for Kinetic Equations

TL;DR

This paper introduces a new machine learning framework for creating reliable, interpretable hydrodynamic models for kinetic equations that maintain accuracy across different flow regimes, addressing multiscale challenges.

Contribution

The paper presents a novel ML-based approach for constructing uniformly accurate reduced models for kinetic equations, incorporating generalized moments, physical constraints, and active learning.

Findings

Achieves uniform accuracy across a wide range of Knudsen numbers.

Successfully applies to BGK model and Maxwell molecules collisions.

Ensures model reliability without dependence on numerical discretization.

Abstract

A new framework is introduced for constructing interpretable and truly reliable reduced models for multiscale problems in situations without scale separation. Hydrodynamic approximation to the kinetic equation is used as an example to illustrate the main steps and issues involved. To this end, a set of generalized moments are constructed first to optimally represent the underlying velocity distribution. The well-known closure problem is then solved with the aim of best capturing the associated dynamics of the kinetic equation. The issue of physical constraints such as Galilean invariance is addressed and an active learning procedure is introduced to help ensure that the dataset used is representative enough. The reduced system takes the form of a conventional moment system and works regardless of the numerical discretization used. Numerical results are presented for the BGK (Bhatnagar--Gross--Krook) model and binary collision of Maxwell molecules. We demonstrate that the reduced model achieves a uniform accuracy in a wide range of Knudsen numbers spanning from the hydrodynamic limit to free molecular flow.

Tables

Table 1. Table 1. Relative error (in percentages) of the different machine learning models for the three tasks with the BGK model. The first one in each cell denotes the relative absolute error (RAE) and the second denotes the relative squared error (RSE). The numbers in the parentheses denote the standard deviation of the last one or two digits computed from three independent runs. The results in the fourth column are obtained using the models learned in Task Mix directly.

Model	Wave	Mix	MixInTransition
EncMLC (no explor)	1.27(13), 1.60(35)	1.68(10), 2.35(12)	1.82(11), 2.49(11)
EncMLC (explor)	0.85(10), 1.01(14)	1.25(5), 1.75(8)	1.55(4), 2.06(7)
GalEncMLC	0.97(22), 1.11(23)	1.28(6), 1.85(10)	1.53(4), 2.11(8)
EncE2EMLC	0.76(6), 0.92(7)	–	–
HermMLC	0.34(5), 0.43(6)	0.92(4), 1.45(10)	1.77(25), 3.01(27)
DirectConv	0.92(3), 1.16(6)	0.83(5), 1.13(9)	2.57(20), 3.61(26)

Table 2. Table 2. Relative error (in percentages) of the different machine learning models for the three tasks for the Maxwell model. See Table  1 for more detailed explanation.

Model	Wave	Mix	MixInTransition
EncMLC (no explor)	0.83(7), 1.19(13)	1.30(11), 1.96(16)	1.50(7), 2.28(10)
GalEncMLC	0.77(6), 1.09(10)	1.18(10), 1.77(15)	1.42(7), 2.12(13)
HermMLC	0.30(4), 0.50(10)	1.12(7), 1.74(11)	1.22(4), 1.89(5)
DirectConv	1.49(8), 2.22(14)	1.00(6), 1.47(9)	2.62(52), 3.64(74)

Equations

$${\partial }_{t}f+\mathbf{v}\cdot {\nabla }_{\mathbf{x}}f=\frac{1}{\epsilon }Q(f),f(\mathbf{x},\mathbf{v},0)=f_0(\mathbf{x},\mathbf{v}),$$

$$Q(f)(\mathbf{v})=f_M(\mathbf{v})-f(\mathbf{v}).$$

$$f_M(\mathbf{v})=\frac{\rho }{(2\pi T{)}^{\frac{D}{2}}}\exp \left(-\frac{|\mathbf{v}-\mathbf{u}{|}^2}{2T}\right),$$

$$\rho ={\int }_{{\mathbb{R}}^D}f\mathrm{d}\mathbf{v},\mathbf{u}=\frac{1}{\rho }{\int }_{{\mathbb{R}}^D}f\mathbf{v}\mathrm{d}\mathbf{v},T=\frac{1}{D\rho }{\int }_{{\mathbb{R}}^D}f|\mathbf{v}-\mathbf{u}{|}^2\mathrm{d}\mathbf{v}.$$

$$Q(f)(\mathbf{v})={\int }_{{\mathbb{R}}^2}{\int }_{{\mathbb{S}}^1}\frac{1}{2\pi }(f({\mathbf{v}}^{\prime })f({\mathbf{v}}_{\ast }^{\prime })-f(\mathbf{v})f({\mathbf{v}}_{\ast }))\mathrm{d}\sigma \mathrm{d}{\mathbf{v}}_{\ast }.$$

$$\{\begin{array}{rl}& {\mathbf{v}}^{\prime }=\frac{\mathbf{v}+{\mathbf{v}}_{\ast }}{2}+\frac{|\mathbf{v}-{\mathbf{v}}_{\ast }|}{2}\sigma ,\\ & {\mathbf{v}}_{\ast }^{\prime }=\frac{\mathbf{v}+{\mathbf{v}}_{\ast }}{2}-\frac{|\mathbf{v}-{\mathbf{v}}_{\ast }|}{2}\sigma ,\end{array}$$

$${\int }_{{\mathbb{R}}^D}Q(f)(\mathbf{v})\mathbf{\varphi }(\mathbf{v})\mathrm{d}\mathbf{v}=0,\mathbf{\varphi }(\mathbf{v})=(1,\mathbf{v},|\mathbf{v}{|}^2/2{)}^T.$$

$$\{\begin{array}{rl}& {\partial }_t\rho +{\nabla }_{\mathbf{x}}\cdot (\rho \mathbf{u})=0,\\ & {\partial }_t(\rho \mathbf{u})+{\nabla }_{\mathbf{x}}\cdot (\rho \mathbf{u}\otimes \mathbf{u}+p\mathbf{I})=0,\\ & {\partial }_{t}E+{\nabla }_{\mathbf{x}}\cdot ((E+p)\mathbf{u})=0,\end{array}$$

$$\mathbf{U}=\left(\begin{array}{c}\rho \\ \rho \mathbf{u}\\ E\end{array}\right),{\mathbf{F}}_{\text{Euler}}(\mathbf{U})=\left(\begin{array}{c}\rho {\mathbf{u}}^T\\ \rho \mathbf{u}\otimes \mathbf{u}+p\mathbf{I}\\ (E+p){\mathbf{u}}^T\end{array}\right),$$

$${\partial }_t\mathbf{U}+{\nabla }_{\mathbf{x}}\cdot {\mathbf{F}}_{\text{Euler}}(\mathbf{U})=0.$$

$${\partial }_t\mathbf{M}+{\nabla }_{\mathbf{x}}\cdot {\int }_{{\mathbb{R}}^D}f\mathbf{m}(\mathbf{v}){\mathbf{v}}^T\mathrm{d}\mathbf{v}={\int }_{{\mathbb{R}}^D}\frac{1}{\epsilon }Q(f)\mathbf{m}(\mathbf{v})\mathrm{d}\mathbf{v}.$$

$$\{\begin{array}{rl}& {\partial }_t\mathbf{U}+{\nabla }_{\mathbf{x}}\cdot \mathbf{F}(\mathbf{U},\mathbf{W})=0,\\ & {\partial }_t\mathbf{W}+{\nabla }_{\mathbf{x}}\cdot \mathbf{G}(\mathbf{U},\mathbf{W})=\frac{1}{\epsilon }\mathbf{R}(\mathbf{U},\mathbf{W}).\end{array}$$

$$\mathbf{U}(\mathbf{x},0)={\int }_{{\mathbb{R}}^D}{f}_0(\mathbf{x},\mathbf{v})\mathbf{\varphi }(\mathbf{v})\mathrm{d}\mathbf{v},\mathbf{W}(\mathbf{x},0)={\int }_{{\mathbb{R}}^D}{f}_0(\mathbf{x},\mathbf{v})\mathbf{w}(\mathbf{v})\mathrm{d}\mathbf{v}.$$

$$f(\mathbf{v})=\sum _{\mathbf{\alpha }\in {\mathbb{N}}^D}{f}_{\mathbf{\alpha }}{\mathcal{H}}_{T,\mathbf{\alpha }}\left(\frac{\mathbf{v}-\mathbf{u}}{\sqrt{T}}\right),$$

$${\mathcal{H}}_{T,\mathbf{\alpha }}(\mathbf{v})=\frac{1}{(2\pi {)}^{\frac{D}{2}}}\prod _{j=1}^D{T}^{-\frac{{\alpha }_j+1}{2}}H{e}_{{\alpha }_j}(v_j)\exp (-v_j^2/2),$$

$${\mathbf{W}}_{\text{Enc}}=\Psi (f)={\int }_{{\mathbb{R}}^D}f(\mathbf{v})\mathbf{w}(\mathbf{v})\mathrm{d}\mathbf{v},\Phi (\mathbf{U},{\mathbf{W}}_{\text{Enc}})(\mathbf{v})=\exp (h(\mathbf{v};\mathbf{U},{\mathbf{W}}_{\text{Enc}})).$$

$$\eta (f)={\int }_{{\mathbb{R}}^D}-f\ln f\mathrm{d}\mathbf{v},$$

$${\mathbb{E}}_{f\sim \mathcal{D}}[\Vert f-\Phi (\mathbf{U},\Psi (f)){\Vert }^2+{\lambda }_{\eta }(\eta (f)-h_{\eta }(\mathbf{U},\Psi (f)){)}^2].$$

$$\{\begin{array}{rl}& {\partial }_t\mathbf{U}+{\nabla }_{\mathbf{x}}\cdot [{\mathbf{F}}_0(\mathbf{U})+\tilde{\mathbf{F}}(\mathbf{U},\mathbf{W})]=0,\\ & {\partial }_t\mathbf{W}+{\nabla }_{\mathbf{x}}\cdot [{\mathbf{G}}_0(\mathbf{U})+\tilde{\mathbf{G}}(\mathbf{U},\mathbf{W})]=\frac{1}{\epsilon }\mathbf{R}(\mathbf{U},\mathbf{W}),\end{array}$$

$${\mathbf{G}}_0(\mathbf{U})={\int }_{{\mathbb{R}}^D}{f}_M(\mathbf{v};\mathbf{U})\mathbf{w}(\mathbf{v}){\mathbf{v}}^T\mathrm{d}\mathbf{v}.$$

$${\hat{\mathbf{W}}}_{j,n+1}-{\mathbf{W}}_{j,n}\approx \mathcal{S}[\tilde{\mathbf{G}},\mathbf{R}]({\mathbf{U}}_{j-1,n},{\mathbf{U}}_{j,n},{\mathbf{U}}_{j+1,n},{\mathbf{W}}_{j-1,n},{\mathbf{W}}_{j,n},{\mathbf{W}}_{j+1,n};\Delta t,\Delta x,\epsilon )$$

$${\mathbb{E}}_{f\sim \mathcal{D}}\Vert {\mathbf{W}}_{j,n+1}-{\mathbf{W}}_{j,n}-\mathcal{S}({\mathbf{U}}_{j-1,n},{\mathbf{U}}_{j,n},{\mathbf{U}}_{j+1,n},{\mathbf{W}}_{j-1,n},{\mathbf{W}}_{j,n},{\mathbf{W}}_{j+1,n}){\Vert }^2.$$

$${\mathbf{W}}_{\text{Gal}}={\Psi }_{\text{Gal}}(f)={\int }_{{\mathbb{R}}^D}f(\mathbf{v})\mathbf{w}\left(\frac{\mathbf{v}-\mathbf{u}}{\sqrt{T}}\right)\mathrm{d}\mathbf{v}.$$

$${\partial }_t{\int }_{{\mathbb{R}}^D}f\mathbf{w}\left(\frac{\mathbf{v}-{\mathbf{u}}_j}{\sqrt{T_j}}\right)\mathrm{d}\mathbf{v}+{\nabla }_{\mathbf{x}}\cdot {\int }_{{\mathbb{R}}^D}f\mathbf{w}\left(\frac{\mathbf{v}-{\mathbf{u}}_j}{\sqrt{T_j}}\right){\mathbf{v}}^T\mathrm{d}\mathbf{v}={\int }_{{\mathbb{R}}^D}\frac{1}{\epsilon }Q(f)\mathbf{w}\left(\frac{\mathbf{v}-{\mathbf{u}}_j}{\sqrt{T_j}}\right)\mathrm{d}\mathbf{v}.$$

$${\partial }_t{\mathbf{W}}_{\text{Gal}}+{\nabla }_{\mathbf{x}}\cdot {\mathbf{G}}_{\text{Gal}}(\mathbf{U},{\mathbf{W}}_{\text{Gal}};{\mathbf{U}}_j)=\frac{1}{\epsilon }{\mathbf{R}}_{\text{Gal}}(\mathbf{U},{\mathbf{W}}_{\text{Gal}}).$$

$${\mathbb{E}}_{f\sim \mathcal{D}}\Vert {\int }_{\mathbb{R}}\frac{1}{2}f|\mathbf{v}{|}^2{\mathbf{v}}^T\mathrm{d}\mathbf{v}-(E+p){\mathbf{u}}^T-{\tilde{\mathbf{F}}}_3(\mathbf{U},{\mathbf{W}}_{\text{Enc}}){\Vert }^2.$$

$${\mathbb{E}}_{f\sim \mathcal{D}}\Vert {\int }_{\mathbb{R}}f\mathbf{w}(\mathbf{v}){\mathbf{v}}^T\mathrm{d}\mathbf{v}-{\int }_{\mathbb{R}}{f}_M(\mathbf{v};\mathbf{U})\mathbf{w}(\mathbf{v}){\mathbf{v}}^T\mathrm{d}\mathbf{v}-\tilde{\mathbf{G}}(\mathbf{U},{\mathbf{W}}_{\text{Enc}}){\Vert }^2,$$

$${\mathbb{E}}_{f\sim \mathcal{D}}\Vert {\int }_{\mathbb{R}}Q(f)\mathbf{w}(\mathbf{v})\mathrm{d}\mathbf{v}-\mathbf{R}(\mathbf{U},{\mathbf{W}}_{\text{Enc}}){\Vert }^2.$$

$${\partial }_t\mathbf{U}+{\nabla }_{\mathbf{x}}\cdot ({\mathbf{F}}_{\text{Euler}}(\mathbf{U})+\tilde{\mathbf{F}}[\mathbf{U}(\cdot );\epsilon ])=0.$$

$$\{\begin{array}{l}\rho (x,0)=a_{\rho }\sin (2k_{\rho }\pi x/L+{\psi }_{\rho })+b_{\rho },\\ u(x,0)=0,\\ T(x,0)=a_T\sin (2k_T\pi x/L+{\psi }_T)+b_T.\end{array}$$

$$f_{wave}=\frac{{\alpha }_1{f}_M(\mathbf{v};{\mathbf{U}}_1)+{\alpha }_2{f}_M(\mathbf{v};{\mathbf{U}}_2)}{{\alpha }_1+{\alpha }_2+10^{-6}},$$

Full text

∗To whom correspondence should be addressed. E-mail: jiequnh@princeton.edu

In scientific modeling, we often encounter the following dilemma: We are interested in modeling some macro-scale phenomenon but we only have a reliable model at some micro scale and the micro-scale model is too detailed and too expensive for practical use. The idea of multiscale modeling is to develop models or algorithms that can efficiently capture the macro-scale behavior of the system, but use only the micro-scale model (see [1, 2] for a review). This program has been quite successful for situations in which there is a clear gap between the active scales in the macro- and micro-scale processes. However, for problems that do not exhibit scale separation, we still lack effective tools that can be used to discover the relevant variables and find accurate approximations for the dynamics of these relevant variables.

Machine learning, the second topic that we are concerned with here, is equally broad and important. We are particularly interested in developing physical models using machine learning. There are significant differences between this and traditional machine learning tasks such as the ones in computer vision and data analytics. One is that in the current setting, instead of being given the data, the data is generated using the micro-scale model. The good news is that in principle we can generate an unlimited amount of data. The bad news is that often times the process of generating data is expensive. Therefore it is an important issue to find a data set that is as small as possible and yet representative enough. The second aspect is that we have to be concerned with physical constraints such as symmetry, invariance, and the laws of nature. We have to make a choice between enforcing these constraints explicitly, or enforcing them approximately as a byproduct of accurately modeling the underlying physical process. The third is the interpretability issue. Machine learning models are often black-box in nature. But as a physical model that we can rely on, just as the Schrödinger equation in quantum mechanics or Navier–Stokes equation in fluid mechanics, it cannot just be a black-box completely. Some of these issues have been studied in [3, 4, 5] in the context of modeling inter-atomic force fields.

The work presented in [3, 4, 5] is limited to the context of learning some functions. This paper examines the aforementioned issues systematically in the context of developing new dynamical models for multiscale problems. As a concrete example, we will study the problem where the micro-scale model is the kinetic equation, and the macro-scale model is the hydrodynamic model. Besides being a very representative problem in multiscale modeling, this example is also of great practical interest in its own right. Kinetic equations have found applications in modeling rarefied gases, plasmas, micro-flow, semiconductor devices, radiative transfer, complex fluids, and so on. From a computational point of view, kinetic equations are costly to solve due to the dimensionality of the problem and the multi-dimensional integral in the collision term. As a result, there has been a long history of interest in reducing kinetic models to hydrodynamic models, going back at least to the work of Grad [6]. Our work is in the same spirit.

A crucial dimensionless number that influences the behavior of the solutions of the kinetic equations is the Knudsen number, defined as the ratio between the mean free path of a particle and the typical macroscopic length scale of interest. When the Knudsen number is small, the velocity distribution stays close to local Maxwellians. This is the so-called hydrodynamic regime. In this regime, it is possible to derive hydrodynamic models for some selected macroscopic fluid variables (typically the mass, momentum, and energy densities) such as the Euler equations or the Navier–Stokes–Fourier equations [7, 8]. These hydrodynamic models are not only much less costly to solve, but also much easier to use to explain experimental observations. They are often in the form of conservation laws with some conserved quantities and fluxes. However, when the Knudsen number is no longer small, the hydrodynamic models break down. Direct simulation Monte Carlo (DSMC) method  [9, 10, 11] becomes a more preferred choice. DSMC works well in the collision-less limit when the Knudsen number is large, but the associated computational cost becomes prohibitive when approaching the hydrodynamic regime due to the high collision rate. The significant difference between these two types of approaches creates a problem when modeling transitional flows in which the Knudsen number varies drastically.

Significant effort has been devoted to developing generalized hydrodynamic models as a replacement of the kinetic equation in different regimes, with limited success. These generalized models can be put into two categories. The first are direct extensions of the Navier–Stokes–Fourier equations. In these models, no new variables are introduced, but new derivative terms are added. One example is the Burnett equation [12]. The second are the moment equations. In this case additional moments are introduced and their dynamic equations are derived through the process of “moment closure” using the kinetic equation. The most well-known example is Grad’s 13-moment equations [6]. Much effort has gone into making sure that such moment equations are mathematically well-posed and respect the most important physical constraints, see [8, 13, 14]. In both approaches, one has to introduce some uncontrolled truncation in order to obtain a closed system of partial differential equations (PDE).

To develop machine learning-based models, one can also proceed along these two separate directions. One can stay with the original set of hydrodynamic variables (mass, momentum and energy densities), but introduce corrections to the fluxes as functionals of the local (in space) instantaneous hydrodynamic variables. This approach is relatively straightforward to implement through convolutional neural networks, and it does give reasonably satisfactory results. However, the accuracy of such a model is limited beforehand by the choice of the hydrodynamic variables. Moreover, the resulting model depends heavily on the specific spatial discretization used to find the data and train the model, and it is difficult to obtain a continuous PDE-like model using this approach. In other words, the result is more like a specific numerical algorithm rather than a new reliable physical model. In addition, some of the issues that we mentioned above, which are important for more general multiscale modeling problems, do not arise. This means that as an illustrative example for a general class of problems, it has a limited value. For these reasons, the majority of our efforts is devoted to the second category of methods: the moment equations. Specifically, we revisit the conventional Grad-type moments and explore the possibility of learning new generalized moments, not necessarily polynomials, to represent the velocity distribution function. We also explore the possibility of learning accurately the dynamical model, instead of resorting to ad hoc closure assumptions. In addition, we develop an active learning procedure in order to ensure that the data set used to generate these models are representative enough.

Below we use the Boltzmann equation to explain the technical details of our approach. The algorithms are implemented for the BGK (Bhatnagar–Gross–Krook) model [15] and the Boltzmann equation modeling binary collision of Maxwell molecules. We work with a wide range of Knudsen numbers ranging from $10^{-3}$ to 10 and different initial conditions. It is easy to see that the methodology developed here is applicable to a much wider range of models, some of which are under current investigation.

There has been quite some activities recently on applying machine learning, especially deep learning to study dynamical systems, including representing the physical quantities described by PDEs with physics-informed neural networks [16, 17], uncovering the underlying hidden PDE models with convolutional kernels [18, 19] or sparse identification of nonlinear dynamical systems method [20, 21], predicting reduced dynamics with memory effects using long short-term memory recurrent neural networks [22, 23], calibrating the Reynolds stress in the Reynolds averaged Navier–Stokes equation [24, 25], and so on. Machine learning techniques have also been leveraged to find proper coordinates for dynamical systems in the spirit of dimensionality reduction [26, 27, 28, 29]. Despite the widespread interests in machine learning of dynamical systems, the examples presented in the literature are mainly ODE models or PDE models in which the ground truth is known. This paper stands in contrast to previous works with the aim of learning, from the original micro-scale model, new reduced PDEs at the macro scale accurately.

1. Preliminaries

Let $D$ be the spatial dimension and $f=f(\bm{x},\bm{v},t)\colon\mathbb{R}^{D}\times\mathbb{R}^{D}\times\mathbb{R}\rightarrow\mathbb{R}$ be the one-particle phase space probability density. Consider the following initial value problem of the Boltzmann equation in the absence of external forces

[TABLE]

where $\varepsilon$ is the dimensionless Knudsen number and $Q$ is the collision operator. Two examples of the collision operator are considered in this paper. The first example is the BGK model

[TABLE]

Here $f_{M}$ is the local Maxwellian distribution function, sometimes also called the local equilibrium,

[TABLE]

where $\rho$, $\bm{u}$ and $T$ are the density, bulk velocity and temperature fields. They are related to the moments of $f$ through

[TABLE]

In above the dependence on the location and time has been dropped for clarity of the notation. The second example considered is the model for binary collision of Maxwell molecules in 2-D

[TABLE]

Here $(\bm{v},\bm{v}_{*})$ and $(\bm{v}^{\prime},\bm{v}_{*}^{\prime})$ are the pairs of pre-collision and post-collision velocities, related by

[TABLE]

with $\sigma\in\mathbb{S}^{1}$. In general $Q$ satisfies the conditions

[TABLE]

These ensure that mass, momentum and energy are conserved during the evolution.

It can be shown that when $\varepsilon\ll 1$, $f$ stays close to the local Maxwellian distributions [30, 31, 32, 33, 34]. Taking the moments $\langle\cdot,\bm{\phi}\rangle:=\int\cdot\bm{\phi}(\bm{v})\mathop{}\!\mathrm{d}\bm{v}$ on both sides of the Boltzmann equation (1) with $\bm{\phi}$ in (6) and replacing $f$ by $f_{M}$, one obtains the compressible Euler equations:

[TABLE]

where $p=\rho T$ is the pressure and $E=\frac{1}{2}\rho|\bm{u}|^{2}+\frac{D}{2}\rho T$ is the total energy. Let

[TABLE]

we can rewrite the Euler equations in a succinct conservation form

[TABLE]

For larger values of $\varepsilon$ one would like to use the moment method to obtain similar systems of PDEs that serve as an approximation to the Boltzmann equation. To this end, one starts with a linear space $\mathbb{M}$ of functions of $\bm{v}$ (usually chosen to be polynomials) that include the components of $\bm{\phi}$ in (6). Denote by $\bm{m}$, a vector whose components form a basis of $\mathbb{M}$. Then the moments $\bm{M}$, as functions of $\bm{x}$ and $t$, are defined as $\int_{\mathbb{R}^{D}}f\bm{m}(\bm{v})\mathop{}\!\mathrm{d}\bm{v}$. From the Boltzmann equation, one obtains

[TABLE]

The challenge is to approximate the two integral terms in (10) as functions of $\bm{M}$ in order to obtain a closed system. This is the well-known “moment closure problem” and it is at this stage that various uncontrolled approximations are introduced. In any case, once this is done one obtains a closed system of the form

[TABLE]

Here the moments $\bm{M}$ are decomposed as $\bm{M}=(\bm{U},\bm{W})^{\operatorname{T}}$, where $\bm{U}\in\mathbb{R}^{D+2}$ denotes the conserved quantities including mass, momentum and energy densities, and $\bm{W}\in\mathbb{R}^{M}$ are the extra new moments. The basis functions $\bm{m}(\bm{v})$ are decomposed accordingly as $\bm{m}(\bm{v})=(\bm{\phi}(\bm{v}),\bm{w}(\bm{v}))^{\operatorname{T}}$. The initial condition of (11) is determined by the one of the Boltzmann equation (1):

[TABLE]

Similarly, the term $1/\varepsilon$ in (11) is inherited directly from (1).

2. Machine Learning-Based Moment System

We are interested in approximating a family of kinetic problems, in which the Knudsen number may span from the hydrodynamic regime ($\varepsilon\ll 1$) to the free molecular regime ($\varepsilon\sim 10$), and the initial conditions are sampled from a wide distribution of profiles. Fig. 1 presents a schematic diagram of the framework for the machine learning-based moment method. Below we first describe the method for finding generalized moments and addressing the moment closure problem. Then we discuss how to build the data set $\mathcal{D}$ incrementally to achieve efficient data exploration.

2.1. Finding Generalized Moments

For this discussion, we can neglect the dependence of $f$ on $(\bm{x},t)$ and view $f$ as a function of the velocity $\bm{v}$ only. We consider two different ways of finding additional moments $\bm{W}$. One is using the conventional Grad-type moments and the other is based on the autoencoder. The conventional moment system, such as Grad’s moment system, starts with the Hermite expansion of $f$

[TABLE]

where $\bm{\alpha}=(\alpha_{1},\dots,\alpha_{D})$ is a $D$-dimensional multi-index. Here the basis functions are defined as

[TABLE]

and $He_{k}(\cdot)$ is the $k$-th order Hermite polynomial. Using the orthogonality of Hermite polynomials, one can easily express $f_{\bm{\alpha}}$ as certain moment of $f$. If $f$ is at the local equilibrium, the only non-zero coefficient is $f_{\bm{0}}=\rho$. In general, one can truncate the Hermite expansion at an order $L$ and select those coefficients $f_{\bm{\alpha}}$ with $0<|\bm{\alpha}|\leq L$ to construct the additional moments $\bm{W}_{\text{Herm}}$.

The aforementioned Grad-type moments are based on the rationale that the truncated Hermite series is an effective approximation to $f$ near the local equilibrium. In a broader sense, we can use an autoencoder [35, 36] to seek the generalized moments $\bm{W}_{\text{Enc}}$ without relying on such domain knowledge. More concretely, we wish to find an encoder $\Psi$ that maps $f$ to generalized moments $\bm{W}_{\text{Enc}}\in\mathbb{R}^{M}$ and a decoder $\Phi$ that recovers the original $f$ from $(\bm{U},\bm{W}_{\text{Enc}})$:

[TABLE]

The exponential form is chosen in order to ensure positivity of the reconstructed density. As an auxiliary goal, we would also like to predict a macroscopic analog of the entropy $\eta$

[TABLE]

with all the available moments. Accordingly, the objective function to be minimized reads

[TABLE]

In (14), the unknown functions to optimize are $\bm{w},h,h_{\eta}$ and the trial space (also known as the hypothesis space in machine learning) for all three functions can be any machine learning models. In this work we choose them to be multilayer feedforward neural networks. In practice the $f$’s are always discretized into finite dimensional vectors and the quantities in (14) are actually squared $l^{2}$ norms of the associated vectors. Once the optimal functions $\bm{w},h,h_{\eta}$ are trained, $\bm{W}_{\text{Enc}}=\Psi(f)$, as a general alternative to $\bm{W}_{\text{Herm}}$, provides a new set of the generalized moments of the system.

2.2. Learning Moment Closure

Recall the dynamic equation (11) for the moment system, the goal of moment closure is to find suitable approximations of $\bm{F},\bm{G},\bm{R}$ as functions of $(\bm{U},\bm{W})$. We first rewrite (11) into

[TABLE]

where $\bm{F}_{0}(\bm{U}),\bm{G}_{0}(\bm{U})$ are the fluxes of the corresponding moments $\bm{U},\bm{W}$ under the local Maxwellian distribution, that is, $\bm{F}_{0}(\bm{U})\equiv\bm{F}_{\text{Euler}}(\bm{U})$ and

[TABLE]

Here we denote by $f_{M}(\cdot;\bm{U})$ the Maxwellian distribution defined by the macroscopic variables $\bm{U}$. Our experience has been that separating the terms $\bm{F}_{0}(\bm{U})$ and $\bm{G}_{0}(\bm{U})$ out from $\bm{F}$ and $\bm{G}$ serves to reduce the variance during training. In practice, we can calculate $\bm{w}(\cdot)$ on a few points in advance and use the Gauss-Legendre quadrature rule to approximate $\bm{G}_{0}(\bm{U})$ efficiently. Now our goal becomes finding $\tilde{\bm{F}},\tilde{\bm{G}},\bm{R}$. Note that the first component of $\tilde{\bm{F}}$, corresponding to the correction of the density flux, is always zero since the mass continuity equation is exact.

Take learning $\tilde{\bm{G}},\bm{R}$ for example. At this point, we need a more concrete definition for the dynamics encoded in (15). The easiest way of doing this is to introduce certain numerical scheme for (15). We should emphasize that this scheme only serves to define the dynamics for the PDE. The machine learning model obtained is independent of the details of this scheme so long as it defines the same dynamics (in the limit as the grid size goes to 0). We can think of a numerical scheme $\mathcal{S}$ as an operator, which takes the flux function, the stencil, and numerical discretization parameters as input, and outputs the increment of the function values at the targeted point in time. For instance, in the case when $D=1$, a simple example of $\mathcal{S}$ may take the form

[TABLE]

where $j$ and $n$ denote the spatial and temporal indices respectively. In this case the tuple $\mathcal{X}=(\bm{U}_{j-1,n},\bm{U}_{j,n},\bm{U}_{j+1,n},\bm{W}_{j-1,n},\bm{W}_{j,n},\bm{W}_{j+1,n},\bm{W}_{j,n+1})$ is an example of data needed. The loss function can be chosen as:

[TABLE]

Care should be exercised when choosing the numerical scheme, since the solutions of these systems may contain shocks. The task for learning $\tilde{\bm{F}}$ can be formulated similarly based on the dynamics of $\bm{U}$. The details of the numerical scheme we use can be found in SI Appendix, Sec. B.

2.3. Data Exploration

The quality of the proposed moment method depends on the quality of the data set $\mathcal{D}$ that we use to train the model. It consists of several solutions of the original Boltzmann equation (1) under different initial conditions. Unlike most conventional machine learning problems that rely on fixed given data sets, here the construction of the data set is completely our own choice, and is an important part of the algorithm. In general our objective is to achieve greater accuracy with fewer training data by choosing the training data wisely. In this sense it is close to that of active learning [37]. To achieve this, an interactive algorithm is required between the augmentation of the data set and the learning process.

In this work we adopt the following strategy. One starts with a relatively small data set and uses it to learn the models. Then a new batch of solutions are generated for both the original kinetic model (1) and the moment system (15). The error in the macroscopic variables $\bm{U}$ is calculated as an indicator and the ones with large errors are added to the data set for the next round of learning. These two steps are repeated until convergence is achieved, which indicates that the phase space has been sufficiently explored. The whole scheme works as a closed loop and forms a self-learning process.

One key question is how to initialize the new batch of solutions. In principle we would like to initialize them so that at the end of the active learning process, the configurations that occur in practice have been explored sufficiently. Unfortunately at the moment, there are no precise mathematical principles that we can use to quantify this. This is certainly one issue that we will continue to investigate in the future. More details of the exploration procedure used can be found in SI Appendix, Sec. E.

2.4. Symmetries and Galilean Invariant Moments

An important issue in building machine learning models is how to handle the symmetries in the system. The handling of translational, rotational and even permutational symmetries has been discussed in depth in the literature already [24, 38, 39, 40]. Besides these static symmetries, Boltzmann equation also possesses an important dynamic symmetry, the Galilean invariance. Specifically, for every $\bm{u}^{\prime}\in\mathbb{R}^{D}$, define $f^{\prime}$ by $f^{\prime}(\bm{x},\bm{u},t)=f(\bm{x}-t\bm{u}^{\prime},\bm{v}-\bm{u}^{\prime},t)$. If $f$ is a solution of the Boltzmann equation, then so is $f^{\prime}$. It is desirable that the reduced moment system also inherits such an invariance. Here we present a viable approach that achieves this goal.

The idea is to define the generalized moments $\bm{W}_{\text{Gal}}$ (the subscript Gal signifies Galilean invariance) properly such that they are Galilean invariant. Given the velocity $\bm{u}$ and temperature $T$ of $f$, we modify the encoder to be

[TABLE]

Note that the encoder $\Psi_{\text{Gal}}$ now depends nonlinearly on the first and second moments of $f$ (through $\bm{u}$ and $T$) and is invariant with respect to the choice of the Galilean reference frame. It is straightforward to see that the Grad-type moments $\bm{W}_{\text{Herm}}$ is a special case of (16) and thus are Galilean invariant. Modeling the dynamics of $\bm{W}_{\text{Gal}}$ becomes more subtle due to the spatial dependence in $\bm{u},T$. Below for simplicity we will work with a discretized form of the dynamic model.

Suppose we want to model the dynamics of $\bm{W}_{\text{Gal},j}$ at the spatial grid point indexed by $j$. Integrating the Boltzmann equation against the generalized basis at this grid point gives

[TABLE]

The collision term on the right-hand side evaluated at the grid point $j$ can still be approximated reasonably well by a function of $(\bm{U}_{j},\bm{W}_{\text{Gal},j})$ only since there is no spatial interaction involved. However, after the discretization, the flux term above is going to depend not only on $(\bm{U},\bm{W}_{\text{Gal}})$ but also on the basis quantities $(\bm{u}_{j},T_{j})$ chosen in (17). This motivates us to consider the following approximate moment equation

[TABLE]

Note that (18) is only meant to be used to model the dynamics of $\bm{W}_{\text{Gal},j}$. To model the dynamics of $\bm{W}_{\text{Gal},j^{\prime}}$ at another grid point $j^{\prime}$, a different basis information $\bm{U}_{j^{\prime}}$ is provided. Given the moment equation (18), the loss functions for $\bm{G}_{\text{Gal}},\bm{R}_{\text{Gal}}$ can be defined in the same way as in Sec. Learning Moment Closure. More discussion on the dynamics of Galilean invariant moment systems can be found in SI Appendix, Sec. D.

One should also note that while preserving invariances is an important issue, it is not absolutely necessary to preserve such invariances exactly. If we make the trial space too restrictive in order to preserve invariances, it may become difficult to find a model with satisfactory accuracy. On the other hand, if the reduced dynamics is sufficiently accurate, all the invariances of the original system should be satisfied approximately with similar accuracy.

2.5. An End-To-End Learning Procedure

In the previous sections we have introduced a two-step procedure to learn separately the moments $\bm{W}_{\text{Enc}}$ (or $\bm{W}_{\text{Gal}}$) and the dynamics of the moments. Here we also present an alternative end-to-end learning procedure.

For simplicity, we go back to the framework of Sec. Finding Generalized Moments and Sec. Learning Moment Closure (i.e. using $\bm{W}_{\text{Enc}}$ without specifying explicitly the Galilean invariance) in the 1-D case. The loss function for the autoencoder component is still defined as in  (14). The functions in the moment equations are learned using the following loss functions, guided by (10) and its approximation (15). For $\tilde{\bm{F}}$, only the third component $\tilde{\bm{F}}_{3}$, the energy flux, needs to be considered because the mass and momentum continuity equations in (7) are exact in the 1-D case. The corresponding loss function is given by

[TABLE]

Similarly, the loss function for learning $\tilde{\bm{G}}$ and $\bm{R}$ can be written as

[TABLE]

Linearly combining eqs. (14) and (19) to (21) defines a loss function that allows us to learn the moment system in a single optimization step. The benefits of such a learning strategy are twofold. On one hand, any single snapshot $f_{n}$ ($n$ denotes the temporal index when solving the kinetic equation) is sufficient for evaluating the loss function in the end-to-end approach while two consecutive discretized snapshots $f_{n},f_{n+1}$ are required in the loss function described in Sec. Learning Moment Closure and Sec. Symmetries and Galilean Invariant Moments. The resulting paradigm becomes more like unsupervised learning rather than supervised learning since solving the Boltzmann equation becomes unnecessary. On the other hand, accuracy can potentially be improved since all the parameters are optimized jointly.

2.6. An Alternative Direct Machine Learning Strategy

A more straightforward machine learning approach is to stay at the level of the original hydrodynamic variables $\bm{U}$ and directly learn a correction term for the Euler equations to approximate the dynamics of the Boltzmann equation:

[TABLE]

Here $\tilde{\bm{F}}[\bm{U}(\cdot);\varepsilon]$ is a functional of $\bm{U}(\cdot)$ that depends on $\bm{U}$, where $\bm{U}$ is viewed as a function on the whole space. It is quite straightforward to design machine learning models under this framework. For instance, assume that $D=1$. One can discretize $\bm{U}$ using $N_{x}$ grid points and train a network to approximate $\tilde{\bm{F}}$ in which the input is an $N_{x}\times 4$ matrix representing $\bm{U}$ and $\varepsilon$, and the output are the values of $\tilde{\bm{F}}$ at the $N_{x}$ grid points. In practice, to guarantee translational invariance and approximate locality, we choose a $1$-D convolutional neural network with small convolutional kernels. This approach is conceptually simple and easy to implement. Ideas like this have been used in various problems, including the turbulence models [24, 41, 42].

There are several problems with this. The first is that the approach does not offer any room for improving accuracy since no information can be used other than the instantaneous hydrodynamic variables. The second is that this procedure requires a specific discretization to begin with. The model obtained is tied with this discretization. We would like to have machine learning models that are more like PDEs. The third is that the models obtained is harder to interpret. For instance the role that the Knudsen number plays in the convolutional network is not as explicit as in the Boltzmann equation or the moment equation. In any case, it is our views that the moment system is a more appealing approach.

3. Numerical Results

In this section we report results for the methods introduced above for the BGK model (2) with $D=1$ and the Boltzmann equation for Maxwell molecules (5) with $D=2$. At the moment, there are no reliable conventional moment systems for the full Boltzmann equation for Maxwell molecules. In contrast, the methodology introduced here does not make use of the specific form of $Q(f)$ and can be readily applied to this or even more complicated kinetic models.

We consider the 1-D interval $[-0.5,0.5]$ in the physical domain with periodic boundary condition and the time interval $[0,0.1]$. Note that for the 2-D Maxwell model, we assume the distribution function is constant in the $y$ direction of the physical domain so that it is still sufficient to compute the solution on the 1-D spatial interval. Three different tasks are considered, termed Task Wave, Task Mix, and Task MixInTransition respectively. In the first two tasks the Knudsen number $\varepsilon$ is constant across the whole domain. This constant value is sampled from a log-uniform distribution on $[-3,1]$ respect to base 10, i.e., $\varepsilon$ takes values from $10^{-3}$ to 10. For data exploration, the initial conditions used in Task Wave consist of a few combination of waves randomly sampled from a probability distribution. For Task Mix the initial conditions contain a mixture of waves and shocks, both randomly sampled from their respective probability distributions. More details can be found in SI Appendix, Sec. A. In Task MixInTransition, the initial conditions are the same as the ones in Task Mix but the Knudsen number is no longer constant in the spatial domain. Instead it varies from $10^{-3}$ to 10. This is a toy model for transitional flows. We do not train any new models in Task MixInTransition but instead adopt the model learned from Task Mix to check its transferability.

We generate the training data with Implicit-Explicit (IMEX) scheme [43] for the 1-D BGK model or fast spectral method [44, 45, 46] for the 2-D Maxwell model. The spatial and temporal domains are both discretized into 100 grid points. The velocity domain is discretized into 60 grid points in the 1-D BGK model and 48 grid points in each dimension in the 2-D Maxwell model. In all the numerical experiments, Task Wave uses 100 paths as the training data set whereas Task Mix use 200 paths. All the results reported are based on 100 testing initial profiles sampled from the same distribution of the corresponding tasks. To evaluate the accuracy on the testing data, we consider two types of error measures for the macro quantities $\bm{U}$, the relative absolute error (RAE) and the relative squared error (RSE). More details of the tasks and data are provided in SI Appendix, Sec. A.

We name different moment systems using a combination of the definition of $\bm{W}$ and closure model. The moments considered in this paper, $\bm{W}_{\text{Herm}}$, $\bm{W}_{\text{Enc}}$, $\bm{W}_{\text{Gal}}$ are abbreviated as Herm, Enc, GalEnc, respectively. The machine learning-based closure and its end-to-end version is called MLC and E2EMLC for short. When $\bm{W}_{\text{Enc}}$ is used, the learning of the closure takes the form introduced in Sec. Learning Moment Closure and when $\bm{W}_{\text{Herm}}$ or $\bm{W}_{\text{Gal}}$ is used, it takes the form introduced in Sec. Symmetries and Galilean Invariant Moments. The model described in Sec. An Alternative Direct Machine Learning Strategy is called DirectConv. In all the numerical experiments presented in the main text, we always set $\bm{W}_{\text{Enc}}$, $\bm{W}_{\text{Gal}}$ to be 6-dimensional for the 1-D BGK model and 9-dimensional for the 2-D Maxwell model. When Hermite polynomials are used for the moments, due to the numerical instability of evaluating moments of high-order polynomials, we limit the order to be no larger than 5 for each variable. The resulting $\bm{W}_{\text{Herm}}$ is 3-dimensional for the 1-D BGK model and 6-dimensional for the 2-D Maxwell model, dropping the ones that are identically [math] due to symmetry. Results of EncMLC are always augmented by data exploration, unless specified. Other models do not use data exploration for ease of comparison. Table 1 and Table 2 report the relative error of the different models on the testing data for the three tasks. Fig. 2 shows an example of the profiles of the mass, momentum, and energy densities at $t=0,0.05,0.1$ for the same initial condition in Task Mix for the 2-D Maxwell model, obtained by solving the kinetic equation, the Euler equations, and HermMLC.

The benefits brought by data exploration is clearly shown through the comparison between the first two rows in Table 1. EncMLC with data exploration has a similar accuracy compared to GalEncMLC. However, when the two models are trained on the data sets of the same size and distribution without data exploration, GalEncMLC performs better than EncMLC, for both the 1-D BGK model and 2-D Maxwell model. The superiority of GalEncMLC on data efficiency is not surprising since it better captures the intrinsic features of the original dynamical system.

While the generalized moments and moment equations are learned differently in Task Wave and Task Mix under the associated data distribution, in Task MixInTransition (the third column in Table 1 and Table 2), we use the same models learned from Task Mix directly without any further training. The fact that the relative error is similar to that in Task Mix indicates that the machine learning-based moment system has satisfactory transferability. On the contrary, while the model learned in Task Wave is accurate enough for that particular task, it is not sufficiently transferrable in general.

EncE2EMLC has the best accuracy in Task Wave for the BGK model among all the models based on $\bm{W}_{\text{Enc}}$. Note that all the networks in EncE2EMLC have the same structure as in EncMLC, it seems that this improvement comes mainly from the end-to-end training process. However, when shocks are present, this model performs badly. For example, it may produce unphysical solutions. This should be due to the lack of any enforcement of the entropy condition, either explicitly through the entropy function or implicitly through supervision from the dynamics of the kinetic equation. On the other hand, the existence of shocks is a special feature of the physical problem considered here. This issue disappears in most other physical systems and EncE2EMLC should become a very attractive approach for those systems.

The good performance of HermMLC suggests that the proposed machine learning-based closure is applicable to different types of moments. The remarkable accuracy of HermMLC in Task Wave might be a result of the close proximity between $f$ and the local equilibrium in this task. The fact that GalEncMLC achieved a similar accuracy in Task Mix and Task MixInTransition suggests that the autoencoder is an effective tool for finding generalized moments in general situations.

The accuracy of DirectConv is quite good for both Task Wave and Task Mix. However, the model obtained is tied to the specific discretization algorithm used and it is unclear how it can be used for other discretization schemes. The machine learning-based moment systems do not have this problem since they behave more like conventional PDEs. Fig. 3 illustrates the solutions of GalEncMLC under the same initial condition as in Task Wave for the 2-D Maxwell model but different spatial discretization.

In Fig. 4 we display the log-log scatter plots of the relative error versus the Knudsen number $\varepsilon$ for both Task Wave and Task Mix for the 2-D Maxwell model. One can see that the accuracy of the machine learning-based moment system is almost uniform across the whole regime, with the same computational cost. This stands in striking contrast to the conventional hydrodynamic models or DSMC method. As for the computational cost, to solve for one path for the 2-D Maxwell model on a Macbook Pro with a 2.9GHz Intel Core i5 processor, it takes about 5 minutes to run the fast spectral method for the original Boltzmann equation. In contrast, the runtime of the machine learning-based moment method (GalEncMLC) is only half a second.

We refer the interested readers to SI Appendix, Sec. F for more numerical results, including the shape of the generalized basis functions in the 1-D BGK model, the accuracy under different number of moments, the solution on the test data under different spatial and temporal discretization, the growth of the relative error as a function of time in the three tasks, and more sample profiles in the three tasks. A brief introduction of the hyperbolic regularized Grad’s moment system [13, 14], a conventional moment method for the BGK model, and related results are provided in SI Appendix, Sec. C.

4. Discussion and Conclusion

This paper presents a new framework for multiscale modeling using machine learning in the absence of scale separation. We have put our emphasis on learning physical models, not just a particular algorithm. We have studied some of the main issues involved, including the importance of obeying physical constraints, actively learning, end-to-end models, etc. Our experience suggests that it is often advantageous to respect physical constraints, but there is no need to sacrifice a lot of accuracy just to enforce them, since if we can model the dynamics accurately, the physical constraints are also satisfied with similar accuracy. Even though we still lack a proper mathematical framework to serve as guidelines, active learning is very important in order to ensure the validity of the model under different physical conditions. Regarding the end-to-end model, even though it did not perform satisfactorily for problem with shocks, we feel that it may very well be the most attractive approach in the more general cases since it seems most promising to derive uniform error bounds in this case.

There are a few important issues that need to be addressed when applying the methodology presented here to other problems. The first is the construction of new relevant variables. Here our requirement is encoded in the autoencoder: the reduced variables need to be able to capture enough information so that the one-particle phase space distribution function can be accurately reproduced. This needs to be generalized when considering other models, for example when the microscopic model is molecular dynamics. The second is the starting point for performing “closure”. This component is also problem-specific. Despite the fact that these important components need to be worked out for each specific problem, we do feel that general procedure here is applicable for a wide variety of multiscale problems.

Going back to the specific example we studied, the BGK model or the Boltzmann equation for Maxwell molecules, we presented an interpretable generalized moment system that works well over a wide range of Knudsen numbers. One can think of these models just like conventional PDEs, except that some of the terms in the fluxes and forcing are stored as subroutines. This is not very different from the conventional Euler equations for complex gases where the equations of state are stored as look-up tables or sometimes subroutines. Regarding the three ingredients involved in learning the reduced models, namely, labeling the data, learning from the data, and exploring the data (as shown in Fig. 1), labeling the data is straightforward in this case: we just need to solve the kinetic equation for some short period of time under different initial conditions. Data exploration is carried out using Monte Carlo sampling from some prescribed initial velocity distributions, and the picking of these initial velocity distributions is still somewhat ad hoc. The learning problem is the part that we have studied most carefully. We have explored and compared several different versions of machine learning models. We are now ready to attack other problems such as kinetic models in plasma physics and kinetic models for complex fluids, building on this experience.

Acknowledgement

The work presented here is supported in part by a gift to Princeton University from iFlytek and the Office of Naval Research grant N00014-13-1-0338. We also thank the reviewers for the valuable comments which helped to improve the quality of the work and the paper.

Supplementary Information

A. Tasks and Data

We consider 1-D interval $[-0.5,0.5]$ in the physical domain with periodic boundary condition. We put down a uniform grid of $100$ grid points. Note that for the 2-D Maxwell model, we assume that the spatial domain is homogeneous along the $y-$axis so that $f$ and all the macro variables only depend on the $x$ coordinate. The time interval considered is $[0,0.1]$ with time step size 0.001. In the 1-D BGK model, the velocity domain is truncated to $[-10,10]$ and discretized using 60 nodes according to the Gauss-Legendre quadrature rule. In the 2-D Maxwell model, the velocity domain is truncated to $[-7.7,7.7]\times[-7.7,7.7]$ and discretized using 48 equidistant nodes in each dimension. In Task Wave and Task Mix the Knudsen number $\varepsilon$ is sampled from a log-uniform distribution on $[-3,1]$ respect base 10, i.e., $\varepsilon$ takes values from $10^{-3}$ to 10, constant across the domain. We consider two types of initial conditions in all the three tasks.

For Task Wave the initial condition $f_{wave}$ is sampled from a mixture of two local Maxwellian distributions. Two macroscopic functions $\bm{U}_{1}$, $\bm{U}_{2}$ are sampled from sine waves

[TABLE]

Here we assume $a_{\rho},\psi_{\rho},b_{\rho}$ are random variables sampled from the uniform distributions on $[0.2,0.3]$, $[0,2\pi]$, $[0.5,0.7]$, respectively. $k_{\rho}$ is a random integer sampled uniformly from the set $\{1,2,3,4\}$. $a_{T},\psi_{T},b_{T},k_{T}$ in $T(\cdot,0)$ are independent and identically distributed random variables as their counterparts in the function $\rho(\cdot,0)$. Finally, two local Maxwellian distributions are randomly mixed through

[TABLE]

in which $\alpha_{1},\alpha_{2}$ are two random variables sampled from the uniform distribution on $[0,1]$.

For Task Mix the initial condition $f_{mix}$ is sampled from a random superposition of two functions, $f_{wave}$ as defined above and $f_{shock}$. $f_{shock}$ is also point-wise local Maxwellian, except that the macroscopic functions $\bm{U}$ are made up from some Riemann problems. Consider the Riemann problem in which $\rho_{L},T_{L}$ are two independent random variables sampled from the uniform distribution on $[1,2]$, $\rho_{R},T_{R}$ are two independent random variables sampled from the uniform distributions on $[0.55,0.9]$, and $u_{L},u_{R}$ are 0. The initial condition $f_{shock}$ then has the form

[TABLE]

where $x_{1},x_{2}$ are two random variables sampled from the uniform distributions on $[-0.3,-0.1]$ and $[0.1,0.3]$. Finally, we linearly combine two initial conditions to obtain

[TABLE]

where $\alpha$ is a random variable sampled from the uniform distribution on $[0.2,0.6]$.

For Task MixInTransition the initial condition is the same as in Task Mix. The values of $\varepsilon$ vary from $10^{-3}$ to 10 in the domain, similar to the one used in [47],

[TABLE]

where $x_{0}$ is sampled from the uniform distribution on $[-0.2,0.2]$. This is where the Knudsen number is at its maximum. The sample profiles of mass and energy densities, and Knudsen number in Task MixInTransition are shown in Fig. 5.

We compute two types of error, the relative absolute error (RAE) and the relative squared error (RSE), to measure the accuracy of different models. By a slight abuse of notation, we consider $N$ independent profiles of conserved quantities, $\bm{U}^{(i)}$, $\hat{\bm{U}}^{(i)}$, $i=1,\dots,N$, computed from the kinetic equation and the machine learning-based model respectively. Assume that each profile is discretized using $N_{x}$ grid points indexed by $j$. We define

[TABLE]

There are other ways to measure the accuracy, but our experience suggests that the overall behavior is quite independent of the accuracy measure that we use.

B. Numerical Scheme

Here we introduce in detail the numerical scheme $\mathcal{S}$ used when learning the moment closure. This enters into the concrete form of the loss function for the dynamics of $\bm{U}$ and $\bm{W}$. Recall the dynamic equation

[TABLE]

where $\bm{G}(\bm{U})=\bm{G}_{0}(\bm{U})+\tilde{\bm{G}}(\bm{U},\bm{W})$, as explained in Sec. Learning Moment Closure. Applying the classical finite volume discretization to (23), we get

[TABLE]

where $\hat{\bm{U}}_{j,n+1}$ and $\hat{\bm{W}}_{j,n+1}$ denote the one-step solution of the machine learning-based model at the spatial grid point $j$ and time step $n+1$. This is a conservative scheme with first order accuracy in both time and space.

For $(\bm{F}_{\text{Euler}})_{j+1/2,n}$, any classical numerical scheme for solving the Euler equations can be applied since there is no parameter to optimize. In our implementation, we choose the $1$-D HLLC Riemann solver [48] with entropy fix, implemented in the open source package Clawpack [49, 50, 51].

Regarding the fluxes $\tilde{\bm{F}}_{j+1/2,n}$ and $\bm{G}_{j+1/2,n}$, we choose a generalized Lax-Friedrichs scheme based on a three-point stencil,

[TABLE]

where the constants $A_{\bm{U}},A_{\bm{W}}$ are two diagonal matrices denoting the numerical viscosity coefficients. We optimize these constants during training such that a suitable strength of the numerical viscosity can be found for the machine-learned fluxes.

Combining (24) and (25), we see that the one-step output of the numerical scheme $\mathcal{S}$ is continuous respect to all the parameters. Hence we can use stochastic gradient descent to optimize them.

C. Hyperbolic Moment Method

In this section, we briefly introduce a well-acknowledged moment method in the literature suitable for solving the BGK model, the hyperbolic regularized Grad’s moment system (see for example [13, 14]). As introduced in Sec. Finding Generalized Moments, the starting point is the Hermite expansion of $f$ (12). The considered moments are the expansion coefficients $f_{\bm{\alpha}}$ truncated at a predefined order $L$. Plugging the expansion (12) into the Boltzmann equation and noting that the BGK collision operator can also be expressed using the basis function, one can collect all the coefficients for each basis function and deduce a system of equations. The system is closed by setting $f_{\bm{\alpha}}=0$ for all the $|\bm{\alpha}|>L$. Finally a regularization term is added to the equation at the highest order to ensure that the moment system is hyperbolic. We refer the readers to [52, 53, 54, 55] for more details of this approach and its numerical implementation.

The machine learning-based approach shares a lot in common with this conventional approach: both attempt to solve the kinetic equation accurately; the issues of Galilean invariance are similar. However, there are two important differences. First, while the approach mentioned above can be viewed as a spectral method for the $\bm{v}$ component of the variables, the machine learning-based generalized moments presented has more flexibility in representing the data with adaptive basis functions. Second, while the approach mentioned above is only suitable for a few relaxation types of collision operator (the Maxwell model is not included), the machine learning-based closure presented is equally applicable to other more realistic collision models.

In order to get some ideas about quantitative comparison, we implemented the algorithm for solving the hyperbolic regularized moment equation according to [55] and tested it on all the three tasks considered in the paper for the 1-D BGK model. We use $200$ grid points in the spatial domain and $9$ moments in total. The relative RAE and RSE in percentages are 2.36(3), 3.07(4) (Task Wave), 2.13(2), 3.25(17) (Task Mix), 2.11(7), 3.39(8) (Task MixInTransition). This is slightly worse than the results obtained from the machine learning-based models. It should be pointed out that it is difficult to compare these results directly with the proposed machine learning-based moment system since there are a lot of factors that contribute to the accuracy. Nevertheless, we feel that for more challenging problems, the advantage of the machine learning-based approach will be more striking, not alone its versatility to other collision models.

D. Galilean Invariant Dynamics

The first step of GalEncMLC, finding generalized moments $\bm{W}_{\text{Gal}}$, naturally obeys Galilean invariance. It is more subtle for the dynamics of $\bm{W}_{\text{Gal}}$ to obey Galilean invariance since the associated PDE has additional convection terms involving $\nabla_{\bm{x}}\bm{u},\nabla_{\bm{x}}T$ compared to (11). The simplest approximate solution to this problem is to introduce some spatial dependence into the PDE models.

As discussed in Sec. Symmetries and Galilean Invariant Moments, in consideration of the local dynamics of $\bm{W}_{\text{Gal},j}$ at the grid point $j$, the exact flux function should be

[TABLE]

We need to approximate this quantity in the form of $\bm{G}_{\text{Gal}}(\bm{U},\bm{W}_{\text{Gal}};\bm{U}_{j})$ as proposed in (18), a function of the macroscopic variables $\bm{U},\bm{W}_{\text{Gal}}$, with the basis information $\bm{U}_{j}$ provided as well. Meanwhile, we would like to have an alternative expression similar to the flux terms in  (15) in order to reduce the variance during training. Note that in this setting $\bm{G}_{0}(\bm{U})$ is much more expensive to compute than its counterpart in EncMLC since the Gauss-Legendre quadrature rule now requires evaluating the generalized basis at different sets of the grid points for each single batch of data, due to the nonlinear dependence on $\bm{u}$ and $T$. Instead we consider another decomposition of the flux function

[TABLE]

Here the approximation $\bm{u}\approx\bm{u}_{j}$ in the second term is made on the grounds that the flux above is used only locally around the grid point $j$. If we also assume $\bm{u}\approx\bm{u}_{j}$ in the first term above, it will be equivalent to $\int_{\mathbb{R}^{D}}f(\bm{v}+\bm{u})\bm{w}(\frac{\bm{v}}{\sqrt{T}})\bm{v}^{\operatorname{T}}\mathop{}\!\mathrm{d}\bm{v}$. Such an expression is consistent with our consideration of the Galilean invariance because it only depends on the shape of the phase density rather than the choice of the Galilean reference frame. The above decomposition finally motivates us to write

[TABLE]

where $\tilde{\bm{G}}_{\text{Gal}}$ is the neural network to be optimized. This ensures that the dynamics of $\bm{W}_{\text{Gal}}$ also obeys Galilean invariance approximately.

E. Learning of Neural Networks

EncMLC, GalEncMLC, and HermMLC

We use the same architecture for the neural networks used in EncMLC, GalEncMLC, and HermMLC. The input is always the concatenation of all the variables listed. For the autoencoder, the basis function $\bm{w}$ of the encoder in (13) is represented by a fully-connected neural network with two hidden layers and $3M$ hidden nodes in each layer (recall $M$ denotes the dimension of the generalized moments $\bm{W}_{\text{Enc}}$ or $\bm{W}_{\text{Gal}}$). We use the same technique as in Batch Normalization [56] to normalize the output within each data batch to ensure zero mean and unit variance. It is observed that this operation improves the stability of training. The function $h$ in the decoder is represented by another neural network with two hidden layers, whose widths are $2M$ and $M$, respectively. The activation function is chosen to be the softplus function. $\lambda_{\eta}$ in (14) is chosen to be 0.01. The Adam optimizer [57] is used with learning rate 0.001 and batch size 100 for training the autoencoder. Usually the autoencoder is trained for $60$-$120$ epochs, depending on the type of initial condition and the size of the data set.

For the moment closure in (15), $\tilde{\bm{G}}$ and $\bm{R}$ are modeled by two neural networks both with $3$ hidden layers and $64$ hidden nodes in each layer. For $\tilde{\bm{F}}$ in the 1-D BGK model, only the component corresponding to energy is non-zero and it is modeled by a neural network with two hidden layers and $32$ hidden nodes per layer. For $\tilde{\bm{F}}$ in the 2-D Maxwell model, the components corresponding to momentum and energy should be both approximated. The former is modeled by a neural network with two hidden layers and $8$ hidden nodes per layer and the latter is modeled by a neural network with two hidden layers and $32$ hidden nodes per layer. All the neural networks use the residual connection [58] in the hidden layers. Again softplus is used as the activation function. When building the loss function, we have two separate terms: one for $\tilde{\bm{F}}$ and the other for $\tilde{\bm{G}}$ and $\bm{R}$, as described in Sec. Learning Moment Closure. We use a weighted sum of the two terms as our loss function, with equal weights $100$. An Adam optimizer is employed to train the networks. The learning rate starts from $0.01$ and decays exponentially to $0.001$. Training takes $20$ epochs with batch size $256$.

EncE2EMLC

For EncE2EMLC, we use the exact same architecture for all the networks as in EncMLC and GalEncMLC in both the autoencoder and the moment closure. The single loss function is a linear combination of  (14)(19)(20)(21) with weights 0.01/$SS_{1}$, 0.01/$SS_{2}$, 0.01/$SS_{3}$, 0.01/$SS_{4}$, respectively. $\lambda_{\eta}$ is set to [math] in (14). Here $SS_{1},\dots,SS_{4}$ denote the total sum of squares in the associated regression problem

[TABLE]

Hence the goal of the optimization is to minimize the four relative losses with equal weights 0.01. Noting that $SS_{3},SS_{4}$ actually depend on the encoder, we use the statistics within the batch during training to approximate them. An Adam optimizer is used for 90 epochs with batch size 100. The learning rate is constant during each 30-epoch periods, deceasing from 0.005 to 0.001, and then to 0.0005.

DirectConv

As described in Sec. An Alternative Direct Machine Learning Strategy, we use a $1$-D convolutional neural network to represent the non-zero components of functional $\tilde{\bm{F}}$, corresponding to the momentum and energy fluxes. The input of the network contains $4$ channels of length $N_{x}$. The architecture of the network can be expressed

[TABLE]

where $Conv$ represents $1$-D convolution with kernel size $4$ and periodic padding followed by a softplus activation, $Pooling$ is done by max pooling, $Res$ means a residual connection in the layer, and $Deconv$ means a $1$-D deconvolution operation with kernel size $4$ and stride $2$. The network is trained by an Adam optimizer with batch size $50$ and learning rate exponentially decreasing from $0.001$ to $0.0002$. Training is run for $5000$ epochs.

Data Exploration

For Task Wave, an autoencoder is initially trained for $120$ epochs with a data set containing $50$ paths, sampled from the distribution of the initial profiles. Then, $5$ loops of exploration is done as follows. In each loop, $100$ new paths are evaluated by both the original kinetic model and the moment system, in which $10$ paths with the largest errors are added to the data set, then the autoencoder is retrained for another $20$ epochs on the new data set. Finally the data set contains $100$ paths, the same as in the case without data exploration. For Task Mix the autoencoder is initially trained for $90$ epochs on $100$ paths, before $5$ loops of data exploration. In each loop, $20$ paths with the largest errors from $200$ randomly sampled paths are added to the data set, and then the autoencoder is retrained for $15$ epochs. The final data set contains $200$ paths, the same as in the case without data exploration. The autoencoders, fluxes, and production terms used in Task MixInTransition are the same as Task Mix. No additional training was used.

It is worth mentioning that in the current implementation of data exploration, the cost of generating truthful micro-scale data is not directly reduced because it is needed in evaluating the prediction error of the new data in the exploration stage. There are various ways to fix this problem, for instance, by using the variance from the predictions of an ensemble of networks optimized independently as an indicator of the error in the exploration, as was done in [5]. The ideal scenario is that given a fixed budget for generating the data, the exploration procedure should provide training samples of the highest quality so that the best testing performance can be achieved. This is left for future work.

F. Additional Results

We train GalEncMLC for Task Mix for the 1-D BGK model using 3 or 9 generalized moments, the resulted RAE and RSE are 2.00(12), 3.03(30) and 1.27(5), 1.92(18) respectively. We train GalEncMLC for Task Mix for the 2-D Maxwell model using 6 generalized moments, the resulted RAE and RSE are 1.39(10), 2.27(16). Comparing these results with Table 1 and Table 2, we see that choosing 6 additional generalized moments in the 1-D BGK model and 9 additional generalized moments in the 2-D Maxwell model is a suitable trade-off between accuracy and efficiency given the sizes of the networks and the data sets used in the current experiments.

Fig. 6 plots all six generalized moment functions obtained in GalEncMLC for Task Mix for the 1-D BGK model. As we can see they are all well-behaved functions.

As already discussed in the main text, the machine learning-based moment system behaves like conventional PDEs and are adaptive to different spatial and temporal discretization. Fig. 7 and Fig. 8 provide more evidence of this point.

Fig. 9 shows the log-log scatter plots of the relative error versus the Knudsen number $\varepsilon$ for both Task Wave and Task Mix for the 1-D BGK model.

Fig. 10 shows the growth of the relative error for the solutions of the Euler equations and machine learning-based moment system on 200 new paths in Task Wave, Task Mix, and Task MixInTransition for the 1-D BGK and 2-D Maxwell model.

Fig. 11–13 illustrate three sample profiles of mass, momentum, and energy densities in Task Wave, Task Mix, and Task MixInTransition for the 1-D BGK model, obtained from the kinetic equation, the Euler equations, and the machine learning based moment systems.

Fig. 14–16 illustrate three sample profiles of mass, momentum, and energy densities in Task Wave, Task Mix, and Task MixInTransition for the 2-D Maxwell model, obtained from the kinetic equation, the Euler equations, and the machine learning based moment systems.