Uniformly Accurate Machine Learning Based Hydrodynamic Models for Kinetic Equations

Preliminaries

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

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

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

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

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

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

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

we can rewrite the Euler equations in a succinct conservation form

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

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

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.

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$

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

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 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$

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

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. 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

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

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:

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 . 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.

4. Symmetries and Galilean Invariant Moments

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

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

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.

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.

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.

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:

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.

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 $$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 for the 1-D BGK model or fast spectral method 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 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 , a conventional moment method for the BGK model, and related results are provided in SI Appendix, Sec. C.

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.

References

Supplementary Information

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 $$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$$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

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

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

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 $$,$\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

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

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 ,

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

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

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 with entropy fix, implemented in the open source package Clawpack .

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 ). 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 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 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

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

E. Learning of Neural Networks

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 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 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.

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 in (14). Here $SS_{1},\dots,SS_{4}$ denote the total sum of squares in the associated regression problem

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

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 . 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.