# Data-driven, structure-preserving approximations to entropy-based moment
closures for kinetic equations††thanks: This manuscript has been authored, in
part, by UT-Battelle, LLC, under Contract No. DE-AC0500OR22725 with the U.S.
Department of Energy. The United States Government retains and the publisher,
by accepting the article for publication, acknowledges that the United States
Government retains a non-exclusive, paid-up, irrevocable, world-wide license
to publish or reproduce the published form of this manuscript, or allow others
to do so, for the United States Government purposes. The Department of Energy
will provide public access to these results of federally sponsored research in
accordance with the DOE Public Access Plan (http://energy.gov/downloads/doe-
public-access-plan).

William A. Porteous Department of Mathematics, University of Texas at Austin,
Austin, TX 78712 USA, (afpwilliam@gmail.com).  M. Paul Laiu Multiscale Methods
and Dynamics Group, Computer Science and Mathematics Division, Oak Ridge
National Laboratory, Oak Ridge, TN 37831 USA, (laiump@ornl.gov).  Cory D.
Hauck Multiscale Methods and Dynamics Group, Computer Science and Mathematics
Division, Oak Ridge National Laboratory, Oak Ridge, TN 37831 USA,
(hauckc@ornl.gov).

###### Abstract

We present a data-driven approach to construct entropy-based closures for the
moment system from kinetic equations. The proposed closure learns the entropy
function by fitting the map between the moments and the entropy of the moment
system, and thus does not depend on the space-time discretization of the
moment system and specific problem configurations such as initial and boundary
conditions. With convex and $C^{2}$ approximations, this data-driven closure
inherits several structural properties from entropy-based closures, such as
entropy dissipation, hyperbolicity, and H-Theorem. We construct convex
approximations to the Maxwell-Boltzmann entropy using convex splines and
neural networks, test them on the plane source benchmark problem for linear
transport in slab geometry, and compare the results to the standard,
optimization-based MN closures. Numerical results indicate that these data-
driven closures provide accurate solutions in much less computation time than
the MN closures.

##  1 Introduction

Kinetic equations describe the movement of particles and their interactions
with each other and with the surrounding environment. They are used to model
dilute particle systems in non-equilibrium regimes, with applications that
includes neutral gases [13], plasmas [25], radiation transport [35, 44], and
charge transport in materials [37]. The solution of a kinetic equation is a
density defined over position-velocity phase space. Weighted integrals of this
density with respect to velocity recover important physical quantities such as
mass, momentum, and energy densities.

In many applications, kinetic equations are too expensive to simulate
directly, and reduced models are needed. Moment methods are a class of reduced
models which approximate the evolution of a finite number of velocity
averages, or moments, of the kinetic distribution. The form of such models
depends on the closure which the information from the kinetic density that is
lost in the moment approach. The structure and behavior of such models depends
heavily on the choice of closure.

Entropy-based closures [34] approximate the unknown kinetic density by the
solution of a convex minimization (using the convention of convex entropy,
rather than concave) problem that is constrained by the known moments which
are input from the model. These models inherit many of the structural features
of kinetic equations, including hyperbolicity, entropy dissipation, and an
H-Theorem which relates the uniquely characterizes equilibrium states as
points where the entropy dissipation vanishes.

Despite their elegant structure, the practical implementation of entropy-based
moment closures presents three major challenges. First, in some degenerate
cases, the entropy minimization problem (see (9) below) may not have a
solution [24, 29, 11, 47]. Roughly speaking a finite infimum exists, but it is
never attained. Second, even in cases that the defining optimization problem
is well-posed, a complicated and expensive numerical procedure is required to
solve it via the convex dual (see (11) below) [23, 8, 7, 18, 1, 3, 2]. Third,
a necessary condition for the entropy optimization problem to have a solution
is that the moments in the constraint set be realizable; that is, they are
moments of a non-negative kinetic distribution. However, numerical methods
which preserve convex invariant domains (such as the realizable set) can be
challenging to construct. Limiters based on the strategy proposed in [48] have
been used in [39, 14, 4]. Another viable limiting strategy can be found in
[20].

An alternative to limiting in a numerical method is to regularize the
optimization problem. Such a strategy was introduced in [6]. The main idea is
to replace the moment constraints in the optimization problem by a
penalization term. The resulting moment system maintains hyperbolicity and
entropy dissipation properties of the original entropy-based closure, but
extends the realizable set to the entire space. Recently, a rigorous
convergence analysis has been initiated [5] to understand the errors made in
the regularization.

The relaxed minimization problem that defines the regularized closure is still
numerically expensive to solve. To address this problem, we take advantage of
the entropy structure to construct a data-driven approximation to the entropy-
based closure. The development of data-driven closures has been an active area
of research lately. Constructive models include the classical work in [30], as
well as various approximations of entropy-based closures [46, 42, 38]. Learned
models (i.e., those which rely on optimization of a loss function) can be
found in [10, 27, 21, 28, 45]. Of these learned models the approach presented
here is most closely related to the work in [28, 45]. The approach in [45],
which is applied to the Boltzmann equation of gas-dynamics, accepts the
entropy-based closure as sufficient and then seeks to approximate the map
between the moments and the solution of the dual problem. Thus the method does
not require training from a full kinetic simulation, but instead relies only
on the optimization problem. The consequences of this difference are not
entirely clear: on one hand approximating the entropy-based closure is likely
more robust, reproducible, flexible, and computationally efficient. However,
training with simulation data may provide important prior information that
leads to higher fidelity simulations, if that information can be transferred
to similar problem setups. The method in [28] uses kinetic simulation data to
train the model, and is therefore not an approximation of the entropy-based
closure. However, it does enforce hyperbolicity, at least for linear kinetic
models in slab geometries, like those studied in the numerical results section
of this paper.

In the current paper, we propose a model that learns the map between the
moments and the entropy of the moment system, rather than the moments and the
dual variables. Even though the dual variables are the derivative of the
entropy with respect to the moments, approximating them directly does not
enforce the convex structure required to ensure entropy dissipation and
hyperbolicity in the moment system. However, arguments established in [6] show
that a convex approximation to the entropy will preserve such properties. With
this fact in mind, we simulate the entropy-based systems M1 and M2 models
which are used to approximate radiative transfer. In our numerical results, we
consider a simple linear kinetic in slab geometry, like the one use in [28].
However, the approach is generalizable to arbitrary moment order and dimension
and to nonlinear problems. The only requirement is that the entropy-based
closure is well-posed and the approximation of the entropy is convex.

We construct convex approximations of the entropy using convex splines and
neural networks. In both cases, we rely on a normalization of the moment space
in order to reduce the dimensions of the problem and make the domain bounded.
We show that for the Maxwell-Boltzmann entropy, this normalization preserves
convexity and $C^{2}$ regularity of the entropy. The spline approximation is
fast and provably convex everywhere in the domain. However, they are used only
in the M1 case, which has two moment but reduces to a one-dimensional fitting
problem after the normalization. For M2, the domain for the fitting problem is
two-dimensional. Although convex splines exist is this setting, we are not
aware of a constructive approach that yields a $C^{2}$ function. For the
neural network approximation, convexity can be enforced at training points and
checked at test points. For the problems considered here, we have found that
the neural network fit is convex on a densely sampled set of test points.
However, this approach is not viable in higher dimensions and networks which
guarantee convexity [9] will be necessary.

The rest of paper is organized as follows. In Section 2, we briefly introduce
kinetic equations, entropy-based moments methods, and approximate entropy
closures. In Section 3, we describe the data-driven closures, using the
Maxwell-Boltzmann entropy and applied to a simple linear kinetic equation. In
Section 4, we present implementations details and numerical results. Summary
and conclusions are given in Section 5.

##  2 Preliminaries

###  2.1 Kinetic equation

For neutral particle systems, the governing kinetic equation takes the form

$\partial_{t}f(t,x,v)+v\cdot\nabla_{x}f=\mathcal{C}(f(t,x,\cdot))(v)\>,$ (1)

where the evolving kinetic density function $f\colon[0,\infty)\times X\times
V\to\mathbb{R}_{+}$ describes the particle density at time
$t\in\mathbb{R}_{+}:=[0,\infty)$ at position $x\in X\subseteq\mathbb{R}^{d}$
traveling at velocity $v\in V\subseteq\mathbb{R}^{d}$. The operator
$v\cdot\nabla_{x}(\cdot):=\sum_{i=1}^{d}v^{[i]}\partial_{x^{[i]}}(\cdot)$
models the advection of particles, where $v^{[i]}$ and $x^{[i]}$ respectively
denotes the velocity and position component in dimension $i$. The collision
operator $\mathcal{C}$ is an integral operator in $v$ at each $(t,x)$, which
introduces the interactions between particles as well as the interaction of
particle and the background medium. In order to be well-posed, (1) must be
equipped with appropriate initial and boundary conditions. Solutions to (1)
are often restricted to values in a set $B\subseteq\mathbb{R}_{+}$, which
reflects physical bounds on $f$ such as positivity.

The collision operator is endowed with certain structural features [34] (see
also [6]):

  1. (i)

Collision invariants: There exist function $e\colon V\to\mathbb{R}$ such that
$\langle e\mathcal{C}(g)\rangle=0$ for all
$g\in\operatorname{Dom}(\mathcal{C})$, where
$\langle\cdot\rangle:=\int_{V}\cdot\,dv$. As a result

$\partial_{t}\langle ef\rangle+\nabla_{x}\cdot\langle vef\rangle=0\>,$ (2)

where $\nabla_{x}\cdot\langle
g\rangle:=\sum_{i=1}^{d}\partial_{x^{[i]}}\langle g^{[i]}\rangle$. We denote
the set of collision invariants associated to $\mathcal{C}$ as $\mathbb{E}$.

  2. (ii)

Entropy dissipation: For an open set $D\subseteq\mathbb{R}$, there exists a
twice differentiable, strictly convex function $\eta\colon D\to\mathbb{R}$,
referred to as the kinetic entropy density, such that

$\langle\eta^{\prime}(g)\mathcal{C}(g)\rangle\leq 0\quad\text{for all $g$ in
Dom($\mathcal{C}$) such that Range($g$)$\subseteq D$}\>.$ (3)

Integrating (1) in $v$ against $\eta^{\prime}(f)$ then leads to the entropy
dissipation law

$\partial_{t}\langle\eta(f)\rangle+\nabla_{x}\cdot\langle
v\eta(f)\rangle=\langle\eta^{\prime}(f)\mathcal{C}(f)\rangle\leq 0\>.$ (4)

  3. (iii)

H-theorem: For any $g\in\operatorname{Dom}(\mathcal{C})$ such that
Range($g$)$\subseteq D$, the following statements are equivalent:

$\text{(a)}\,\,\langle\eta^{\prime}(g)\mathcal{C}(g)\rangle=0\>;\quad\text{(b)}\,\,\mathcal{C}(g)=0\>;\quad\text{(c)}\,\,\eta^{\prime}(g)\in\mathbb{E}\>.$
(5)

###  2.2 Entropy-based moment methods

Moment methods approximate the evolution of a finite set of integral averages,
i.e., moments, of $f$ given by

$\mathbf{u}_{f}(t,x)=\langle\mathbf{m}(\cdot)f(t,x,\cdot)\rangle,$ (6)

where the vector-valued function $\mathbf{m}(v)=[m_{0}(v),\dots,m_{n}(v)]^{T}$
collects a set of test functions in $v$. In most cases, these test functions
form the basis of the space $\mathbb{P}_{N}(V)$ consisting of polynomials on
$V$ up to a prescribed degree $N$. The moments $\mathbf{u}_{f}$ satisfy the
equation

$\partial_{t}\mathbf{u}_{f}+\nabla_{x}\cdot\langle
v\mathbf{m}f\rangle=\langle\mathbf{m}\mathcal{C}(f)\rangle,$ (7)

which is not closed. A kinetic closure is one in which an approximation for
$f$ is used to close the system. Let
$\mathbb{F}_{\mathbf{m}}=\\{g\in\text{Dom}(\mathcal{C})\colon\text{Range}(g)\subseteq
D~{}\text{and}~{}|\langle\mathbf{m}g\rangle|<\infty\\}$ and let
$\mathcal{R}_{\mathbf{m}}=\\{\mathbf{w}:\mathbf{w}=\langle\mathbf{m}g\rangle,g\in\mathbb{F}_{\mathbf{m}}\\}$.
A kinetic closure is prescribed via an ansatz
$F\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{F}_{\mathbf{m}}$ where
$F_{\mathbf{w}}$ is understood approximation of a kinetic distribution $g$
with moments $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$., i.e.,
$F_{\langle\mathbf{m}g(x,t,\cdot\rangle)}(v)\approx g(x,v,t)$. The resulting
(closed) moment system for $\mathbf{u}(x,t)\approx\mathbf{u}_{f}(x,t)$ take
the form

$\partial_{t}\mathbf{u}+\nabla_{x}\cdot\langle
v\mathbf{m}F_{\mathbf{u}}\rangle=\langle\mathbf{m}\mathcal{C}(F_{\mathbf{u}})\rangle.$
(8)

An entropy-based closure [34, 16] is one which specifies $F_{\mathbf{w}}$ as
the solution to the optimization problem

$\operatorname*{minimize}_{g\in\mathbb{F}_{\mathbf{m}}}\mathcal{H}(g):=\langle\eta(g)\rangle\qquad\text{subject
to}\quad\langle\mathbf{m}g\rangle=\mathbf{w}\>.$ (9)

When a solution for (9) exists, it is unique and takes the form [24, 29]
$G_{\hat{\boldsymbol{\alpha}}(\mathbf{w})}$ where

$G_{\boldsymbol{\alpha}}:=[\eta^{*}]^{\prime}(\boldsymbol{\alpha}\cdot\mathbf{m}),\quad\boldsymbol{\alpha}\in\mathbb{R}^{n+1};$
(10)

$\eta^{*}$ is the Legendre dual of $\eta$; and
$\hat{\boldsymbol{\alpha}}\colon\mathbb{R}^{n}\to\mathbb{R}^{n}$ is the map
from $\mathbf{w}$ to the solution of the dual problem to (9), i.e.,

$\hat{\boldsymbol{\alpha}}(\mathbf{w})=\operatorname*{argmax}_{\boldsymbol{\alpha}\in\mathbb{R}^{n}}\\{\boldsymbol{\alpha}\cdot\mathbf{w}-\langle\eta^{*}(\boldsymbol{\alpha}\cdot\mathbf{m})\rangle\\}\>.$
(11)

It follows from the equality constraints in (9) (or equivalently from the
first-order optimality condition for (11)) that
$\langle\mathbf{m}G_{\hat{\boldsymbol{\alpha}}{(\mathbf{w})}}\rangle=\mathbf{w}\>,$
i.e., the ansatz preserves moment vectors. Thus the inverse of
$\hat{\boldsymbol{\alpha}}$, denoted as
$\hat{\mathbf{w}}\colon\mathbb{R}^{n}\to\mathbb{R}^{n}$, is given by

$\hat{\mathbf{w}}(\boldsymbol{\alpha}):=\langle\mathbf{m}G_{\boldsymbol{\alpha}}\rangle\>.$
(12)

With this property, (8) becomes the moment equation

$\partial_{t}\mathbf{u}+\nabla_{x}\cdot\langle
v\mathbf{m}G_{\hat{\boldsymbol{\alpha}}(\mathbf{u})}\rangle=\langle\mathbf{m}\mathcal{C}(G_{\hat{\boldsymbol{\alpha}}(\mathbf{u})})\rangle\>.$
(13)

The system (13) is a symmetric hyperbolic balance law when expressed in terms
of the dual variable
$\boldsymbol{\beta}(x,t)=\hat{\boldsymbol{\alpha}}(\mathbf{u}(x,t))$ and that
solutions to (13) dissipate the strictly convex entropy

$h(\mathbf{u}):=\mathcal{H}(G_{\hat{\boldsymbol{\alpha}}(\mathbf{u})}),$ (14)

These results are easy to show once the known relationship
$h^{\prime}(\mathbf{w})=\hat{\boldsymbol{\alpha}}(\mathbf{w})$ is established;
see [34] for details.

###  2.3 Approximate entropy-based closures

Let $h_{\textup{a}}$ be a strictly convex $C^{2}$ approximation of $h$ and let
$h^{*}_{\textup{a}}$ be the Legendre dual of $h$, so that
$[h^{*}_{\textup{a}}]^{\prime}=[h^{\prime}_{\textup{a}}]^{-1}$. For any
$\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$, define
$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w})=h_{\textup{a}}^{\prime}(\mathbf{w})$.
The approximate entropy-based closure for $\mathbf{u}_{\textup{a}}$ is

$\partial_{t}\mathbf{u}_{\textup{a}}+\nabla_{x}\cdot\langle
v\mathbf{m}G_{\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{u}_{\textup{a}})}\rangle=\langle\mathbf{m}\mathcal{C}(G_{\hat{\boldsymbol{\alpha}}(\mathbf{u}_{\textup{a}})})\rangle\>.$
(15)

The following result comes directly from [6], where $h_{\textup{a}}$ is
defined to be $h_{\gamma}$, which is parameterized by a regularization
parameter $\gamma>0$ such that $\lim_{\gamma\to 0^{+}}h_{\gamma}=h$. (See [6]
for the explicit definition of $h_{\gamma}$.) Since it forms the basis of the
data-driven closure strategy, we present it again here, along with a complete
proof.

######  Theorem 2.1 ([6]).

The moment system (15) is symmetric hyperbolic in the variable
$\boldsymbol{\beta}(t,x):=\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{u}_{\textup{a}}(t,x))$.
The function $h_{\textup{a}}$ acts as an entropy of (15); that is,
$h_{{\textup{a}}}(\mathbf{u}_{{\textup{a}}})$ is dissipated by solutions of
(15). Moreover, the following statements are equivalent:

${\rm(a)}\,\,\boldsymbol{\beta}\cdot\langle\mathbf{m}\mathcal{C}(G_{\boldsymbol{\beta}})\rangle=0\>;\quad{\rm(b)}\,\,\langle\mathbf{m}\mathcal{C}(G_{\boldsymbol{\beta}})\rangle=0\>;\quad{\rm(c)}\,\,\boldsymbol{\beta}\cdot\mathbf{m}\in\mathbb{E}\>.$
(16)

###### Proof.

By definition,
$\boldsymbol{\beta}=h^{\prime}_{\textup{a}}(\mathbf{u}_{\textup{a}})$, which
implies that
$\mathbf{u}_{\textup{a}}=[h^{*}_{\textup{a}}]^{\prime}(\boldsymbol{\beta})$.
Thus in terms of $\boldsymbol{\beta}$, (15) becomes,

$\partial_{t}[h^{*}_{\textup{a}}]^{\prime}(\boldsymbol{\beta})+\nabla_{x}\cdot\langle
v\mathbf{m}G_{\boldsymbol{\beta}}\rangle=\langle\mathbf{m}\mathcal{C}(G_{\boldsymbol{\beta}})\rangle\>.$
(17)

The chain rules, along with the definition of $G$ in (10), then gives

$\partial_{t}[h^{*}_{\textup{a}}]^{\prime}(\boldsymbol{\beta})=[h^{*}_{\textup{a}}]^{\prime\prime}(\boldsymbol{\beta})\partial_{t}\boldsymbol{\beta}\quad\mbox{and}\quad\nabla_{x}\cdot\left\langle
v\mathbf{m}G_{\boldsymbol{\beta}}\right\rangle=\left\langle
v\mathbf{m}\,[\eta^{*}]^{\prime}(\boldsymbol{\beta}\cdot\mathbf{m})\right\rangle\cdot\nabla_{x}\boldsymbol{\beta}$
(18)

The symmetric hyperbolic form then follows since
$[h^{*}_{\textup{a}}]^{\prime\prime}(\boldsymbol{\beta})$ is symmetric and
positive definite and $\left\langle
v\mathbf{m}\mathbf{m}^{T}\,[\eta^{*}]^{\prime}(\boldsymbol{\beta}\cdot\mathbf{m})\right\rangle$
is symmetric.

To show entropy dissipation, we multiply (15) by $\boldsymbol{\beta}$:

$\partial_{t}h_{\textup{a}}(\mathbf{u}_{\textup{a}})+\boldsymbol{\beta}\cdot\nabla_{x}\cdot\langle
v\mathbf{m}G_{\boldsymbol{\beta}}\rangle=\langle(\boldsymbol{\beta}\cdot\mathbf{m})\mathcal{C}(G_{\boldsymbol{\beta}})\rangle$
(19)

and observe that, upon integration by parts and the reverse chain rule,

$\displaystyle\boldsymbol{\beta}\cdot\nabla_{x}\cdot\langle
v\mathbf{m}G_{\boldsymbol{\beta}}\rangle$
$\displaystyle=\nabla_{x}\cdot\langle
v(\boldsymbol{\beta}\cdot\mathbf{m})G_{\boldsymbol{\beta}}\rangle-\langle
v\mathbf{m}G_{\boldsymbol{\beta}}\rangle\cdot\nabla_{x}\boldsymbol{\beta}$
(20) $\displaystyle=\nabla_{x}\cdot\langle
v(\boldsymbol{\beta}\cdot\mathbf{m})G_{\boldsymbol{\beta}}\rangle-\nabla_{x}\cdot\langle
v\eta^{*}(\boldsymbol{\beta}\cdot\mathbf{m})\rangle$

Hence there is a locally conservative entropy flux for smooth solutions.
Meanwhile, setting
$[\eta^{*}]^{\prime}(g):=\boldsymbol{\beta}\cdot\mathbf{m}$, we recognize the
source term in (19) as

$\langle(\boldsymbol{\beta}\cdot\mathbf{m})\mathcal{C}(G_{\boldsymbol{\beta}})\rangle=[\eta^{*}]^{\prime}(g)\mathcal{C}(g)\leq
0$ (21)

This shows that the source terms dissipates the entropy. Moreover, the
conditions in (16) are inherited directly from the kinetic structure
formulated in (5). ∎

##  3 Data-driven approximations for the Maxwell-Boltzmann entropy

In this paper, we focus on the construction of a data-driven approximate
entropy to the Maxwell-Boltzmann entropy. The Maxwell-Boltzmann entropy is
commonly used to close the moment system due to its physical relevance.
However, its application is often limited by the high computational cost
incurred in the optimization procedure for calculating the ansatz. The
Maxwell-Boltzmann kinetic entropy density is given by

$\eta_{\mbox{\tiny{\sc MB}}}(z):=z\log
z-z\>,\quad\text{with}\quad\eta_{\mbox{\tiny{\sc
MB}}}^{*}(y)=[\eta_{\mbox{\tiny{\sc MB}}}^{*}]^{\prime}(y)=e^{y}\>.$ (22)

The domain of $\eta_{\mbox{\tiny{\sc MB}}}$ is $D=\mathbb{R}_{+}$, thus the
Maxwell-Boltzmann entropy leads to a nonnegative ansatz. Specifically,
plugging $\eta=\eta_{\mbox{\tiny{\sc MB}}}$ into (9)–(11) gives the ansatz
$G_{\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w})}=e^{\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w})\cdot\mathbf{m}}$, where

$\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w})=\operatorname*{argmax}_{\boldsymbol{\alpha}\in\mathbb{R}^{n}}\\{\boldsymbol{\alpha}\cdot\mathbf{w}-\langle
e^{\boldsymbol{\alpha}\cdot\mathbf{m}}\rangle\\}\>.$ (23)

The Maxwell-Boltzmann entropy $h_{\mbox{\tiny{\sc MB}}}$ is defined as
$h_{\mbox{\tiny{\sc MB}}}:=\langle\eta_{\mbox{\tiny{\sc
MB}}}^{*}(G_{\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w})})\rangle$, and the associated realizable set

$\mathcal{R}_{\mathbf{m},\mbox{\tiny{\sc
MB}}}:=\\{\mathbf{w}\in\mathbb{R}^{n}\colon\mathbf{w}=\langle\mathbf{m}g\rangle,\,g\geq
0,\,\langle|\mathbf{m}g|\rangle<\infty\\}\>.$ (24)

It is known [8, 33, 17] that the optimization procedure for evaluating
$\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc MB}}}(\mathbf{w})$ at given moment
vector $\mathbf{w}$ is usually the most computationally intensive part in
solving the moment system with Maxwell-Boltzmann moment closure. We aim to
construct a twice differentiable, convex approximate entropy
$h_{\textup{a}}\colon\mathbb{R}^{n}\to\mathbb{R}$ to the Maxwell-Boltzmann
entropy $h_{\mbox{\tiny{\sc MB}}}$, and approximate the optimal multiplier
$\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc MB}}}$ by

$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w}):=h_{\textup{a}}^{\prime}(\mathbf{w})\>.$
(25)

With a twice differentiable and convex $h_{\textup{a}}$, the closed moment
system inherits important features of the kinetic equation, including
hyperbolicity, collisional invariants, entropy dissipation, and H-theorem , as
shown in Section 2.3. Further, since the proposed closure is based on an
approximate entropy $h_{\textup{a}}$, it can be constructed using data from
the optimization problem (23), which does not require full kinetic
simulations. Thus, the resulting closure is independent of the space-time
discretization of moment models, initial and boundary conditions, and
collision operators. The closure only depends on the moment order $N$.

We propose to construct $h_{\textup{a}}$ based on a given dataset
$\\{(\mathbf{w}^{(i)},h_{\mbox{\tiny{\sc
MB}}}(\mathbf{w}^{(i)}),\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w}^{(i)}))\\}_{i\in\mathcal{I}}$, where $\\{h_{\mbox{\tiny{\sc
MB}}}(\mathbf{w}^{(i)})\\}_{i\in\mathcal{I}}$ and
$\\{\hat{\boldsymbol{\alpha}}_{\mbox{\tiny{\sc
MB}}}(\mathbf{w}^{(i)})\\}_{i\in\mathcal{I}}$ are respectively the function
values and the gradients of the Maxwell-Boltzmann entropy, $h_{\mbox{\tiny{\sc
MB}}}$, evaluated at the moments $\\{\mathbf{w}^{(i)}\\}_{i\in\mathcal{I}}$
sampled from the domain of $h_{\mbox{\tiny{\sc MB}}}$, i.e., the realizable
set $\mathcal{R}_{\mathbf{m},\mbox{\tiny{\sc MB}}}$. For the remainder of the
paper, we focus on Maxwell-Boltzmann entropy and will drop the MB subscript
for simplicity.

###  3.1 Dimension reduction via normalization

One key factor in the construction of data-driven closures is a good sampling
strategy of the dataset
$\\{(\mathbf{w}^{(i)},h(\mathbf{w}^{(i)}),\hat{\boldsymbol{\alpha}}(\mathbf{w}^{(i)}))\\}_{i\in\mathcal{I}}$.
Due to the nonlinear nature of the entropy function, the quality of data-
driven approximations depends on whether the sampled dataset provides a
reasonable coverage of the realizable set $\mathcal{R}_{\mathbf{m}}$. In the
context of Maxwell-Boltzmann entropy, the realizable set
$\mathcal{R}_{\mathbf{m}}$ is unbounded, e.g.,
$\mathbf{w}=[w_{0},0,\dots,0]^{T}$ is in $\mathcal{R}_{\mathbf{m}}$ for any
$w_{0}>0$. The unboundedness of $\mathcal{R}_{\mathbf{m}}$ makes sampling a
representative dataset difficult. In this section, we show that this issue can
be addressed by (i) sampling data on a “normalized" bounded realizable set
$\tilde{\mathcal{R}}_{\mathbf{m}}\subset\mathbb{R}^{n-1}$, (ii) constructing a
convex approximation of the Maxwell-Boltzmann entropy on
$\tilde{\mathcal{R}}_{\mathbf{m}}$, and (iii) extending the constructed
approximation to the full realizable set
$\mathcal{R}_{\mathbf{m}}\subset\mathbb{R}^{n}$.

From here on, we introduce the notation $w_{0}$ for the zero-th order moment
and $\tilde{\mathbf{w}}$ for the remainder moment vector, i.e.,
$\mathbf{w}=[w_{0},\tilde{\mathbf{w}}^{T}]^{T}$. With this notation, we
consider the normalized realizable set to be
$\tilde{\mathcal{R}}_{\mathbf{m}}:=\\{\tilde{\boldsymbol{\omega}}\in\mathbb{R}^{n-1}\colon\boldsymbol{\omega}=[1,\tilde{\boldsymbol{\omega}}^{T}]^{T}\in\mathcal{R}_{\mathbf{m}}\\}$,
where each vector in $\tilde{\mathcal{R}}_{\mathbf{m}}$ is a realizable moment
vector with the zero-th moment normalized to one. The set
$\tilde{\mathcal{R}}_{\mathbf{m}}$ is known to be bounded for reasonable
choices of $\mathbf{m}$, e.g. polynomial basis and spherical harmonic basis,
which makes it much easier to sample a representative dataset on
$\tilde{\mathcal{R}}_{\mathbf{m}}$. To this point, we aim to approximate
$\tilde{h}$, the restriction of Maxwell-Boltzmann entropy $h$ on
$\tilde{\mathcal{R}}_{\mathbf{m}}$, by a data-driven convex approximation
$\tilde{h}_{\textup{a}}\in C^{2}(\tilde{\mathcal{R}}_{\mathbf{m}})$ with data
collected on $\tilde{\mathcal{R}}_{\mathbf{m}}$. The function
$\tilde{h}_{\textup{a}}$ is then extended from
$\tilde{\mathcal{R}}_{\mathbf{m}}$ to the full realizable set
$\mathcal{R}_{\mathbf{m}}$ to approximate $h$.

We first show in the following lemma that, for each
$\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$, the Maxwell-Boltzmann entropy
$h(\mathbf{w})$ can be calculated from $w_{0}$ and
$\tilde{h}(\tilde{\mathbf{w}}/w_{0})$, where $\tilde{h}$ is the restriction of
$h$ on $\tilde{\mathcal{R}}_{\mathbf{m}}$.

######  Lemma 3.1.

Let $h$ be the Maxwell-Boltzmann entropy, $\tilde{h}$ be the restriction of
$h$ on $\tilde{\mathcal{R}}_{\mathbf{m}}$, i.e.,

$\tilde{h}(\tilde{\boldsymbol{\omega}}):=h([1,\tilde{\boldsymbol{\omega}}^{T}]^{T})\quad\text{for
all }\tilde{\boldsymbol{\omega}}\in\tilde{\mathcal{R}}_{\mathbf{m}},$ (26)

and $\tilde{\boldsymbol{\alpha}}$ be the gradient of $\tilde{h}$. Then, for
any moment vector
$\mathbf{w}=[w_{0},\tilde{\mathbf{w}}^{T}]^{T}\in\mathcal{R}_{\mathbf{m}}$, it
follows that

$h(\mathbf{w})=w_{0}\,\tilde{h}(\tilde{\mathbf{w}}/w_{0})+w_{0}\log w_{0}\>,$
(27)

and

$\hat{\boldsymbol{\alpha}}(\mathbf{w}):=h^{\prime}(\mathbf{w})=\left[\begin{array}[]{c}\tilde{h}(\tilde{\mathbf{w}}/w_{0})+\frac{1}{w_{0}}\tilde{\mathbf{w}}^{T}\tilde{\boldsymbol{\alpha}}(\tilde{\mathbf{w}}/w_{0})+\log
w_{0}+1\\\
\tilde{\boldsymbol{\alpha}}(\tilde{\mathbf{w}}/w_{0})\end{array}\right]\>.$
(28)

###### Proof.

It follows from the definition of $h$ in (14) and strong duality of the
primal-dual pair (9)-(11) that

$h(\mathbf{w})=\hat{\boldsymbol{\alpha}}(\mathbf{w})\cdot\mathbf{w}-\langle
e^{\hat{\boldsymbol{\alpha}}(\mathbf{w})\cdot\mathbf{m}}\rangle\quad\mbox{and}\quad\tilde{h}(\tilde{\mathbf{w}}/w_{0})={h}({\mathbf{w}}/w_{0})=\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{w}/w_{0}-\langle
e^{\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{m}}\rangle\>,$ (29)

From (12), we have

$\mathbf{w}/w_{0}=\hat{\mathbf{w}}(\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0}))=\langle\mathbf{m}\,e^{\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{m}}\rangle\>,$
(30)

which implies

$\hat{\mathbf{w}}(\hat{\boldsymbol{\alpha}}(\mathbf{w}))=\mathbf{w}=w_{0}\langle\mathbf{m}\,e^{\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{m}}\rangle=\langle\mathbf{m}\,e^{\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{m}+\log
w_{0}}\rangle\>.$ (31)

Thus
$\hat{\boldsymbol{\alpha}}(\mathbf{w})=\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})+[\log
w_{0},0,\dots,0]^{T}$. Plugging this into (29) leads to

$h(\mathbf{w})=\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{w}+w_{0}\log
w_{0}-\langle
e^{\hat{\boldsymbol{\alpha}}(\mathbf{w}/w_{0})\cdot\mathbf{m}+\log
w_{0}}\rangle=w_{0}\,\tilde{h}(\tilde{\mathbf{w}}/w_{0})+w_{0}\log w_{0}\>.$
(32)

Then (28) can be verified by taking the gradient of $h$. ∎

The result in Lemma 3.1 allows us to construct an approximation
$\tilde{h}_{\textup{a}}$ to $\tilde{h}$ on the bounded set
$\tilde{\mathcal{R}}_{\mathbf{m}}$, and then approximate $h$ on
$\mathcal{R}_{\mathbf{m}}$ by the extension
$h_{\textup{a}}\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{R}$ defined as

$h_{\textup{a}}(\mathbf{w})=w_{0}\,\tilde{h}_{\textup{a}}(\tilde{\mathbf{w}}/w_{0})+w_{0}\log
w_{0}\>,\quad\forall\mathbf{w}\in\mathcal{R}_{\mathbf{m}}\>.$ (33)

It is then left for us to prove that this $h_{\textup{a}}$ is twice
differentiable and convex with a properly chosen $\tilde{h}_{\textup{a}}$,
which is shown in the following theorem.

######  Theorem 3.2.

Suppose that
$\tilde{h}_{\textup{a}}\colon\tilde{\mathcal{R}}_{\mathbf{m}}\to\mathbb{R}$ is
a (strictly) convex and twice differentiable function, then the reconstructed
function $h_{\textup{a}}$ defined in (33) is also (strictly) convex and twice
differentiable.

###### Proof.

First, the twice differentiability of $h_{\textup{a}}$ directly follows from
(33) and the twice differentiability of $\tilde{h_{\textup{a}}}$. To show that
$h_{\textup{a}}$ is convex when $\tilde{h}_{\textup{a}}$ is convex, we prove
that $H(\mathbf{w})$, the Hessian of $h_{\textup{a}}$, is positive
semidefinite for all $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$. Let $H$ take the
form

${H}(\mathbf{w})=\left[\begin{array}[]{cc}a(\mathbf{w})&\mathbf{b}^{T}(\mathbf{w})\\\
\mathbf{b}(\mathbf{w})&M(\mathbf{w})\end{array}\right]\>,\quad\text{with}\quad
a\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{R}\>,\quad\mathbf{b}\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{R}^{n-1}\>,\quad
M=\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{R}^{(n-1)\times(n-1)}\>.$ (34)

Then direct calculation of the second derivatives leads to

$a(\mathbf{w})=\frac{1}{w_{0}^{3}}\tilde{\mathbf{w}}^{T}\tilde{H}(\tilde{\boldsymbol{\omega}})\tilde{\mathbf{w}}+\frac{1}{w_{0}}\>,\quad\mathbf{b}(\mathbf{w})=-\frac{1}{w_{0}^{2}}\tilde{H}(\tilde{\boldsymbol{\omega}})\tilde{\mathbf{w}}\>,\quad\mbox{and}\quad
M(\mathbf{w})=\frac{1}{w_{0}}\tilde{H}(\tilde{\boldsymbol{\omega}})\>,$ (35)

where $\mathbf{w}=:[w_{0},\tilde{\mathbf{w}}^{T}]^{T}$,
$\tilde{\boldsymbol{\omega}}:=\tilde{\mathbf{w}}/w_{0}$, and $\tilde{H}$ is
the Hessian of $\tilde{h}_{\textup{a}}$. The terms in (35) are well-defined
since $w_{0}>0$ for any $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$. We now show
that $H$ is positive semi-definite by verifying that
$\mathbf{p}^{T}H(\mathbf{w})\mathbf{p}\geq 0$ for all
$\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$, $\mathbf{p}\in\mathbb{R}^{n}$. Let
$\mathbf{p}=:[p_{0},\tilde{\mathbf{p}}]^{T}$, then

$\displaystyle\mathbf{p}^{T}H(\mathbf{w})\mathbf{p}$
$\displaystyle=p_{0}^{2}\,a(\mathbf{w})+2\,p_{0}\,\tilde{\mathbf{p}}^{T}\mathbf{b}(\mathbf{w})+\tilde{\mathbf{p}}^{T}M(\mathbf{w})\tilde{\mathbf{p}}$
(36)
$\displaystyle=\frac{p_{0}^{2}}{w_{0}^{3}}\tilde{\mathbf{w}}^{T}\tilde{H}(\tilde{\boldsymbol{\omega}})\tilde{\mathbf{w}}+\frac{p_{0}^{2}}{w_{0}}-2\,\frac{p_{0}}{w_{0}^{2}}\,\tilde{\mathbf{p}}^{T}\tilde{H}(\tilde{\boldsymbol{\omega}})\tilde{\mathbf{w}}+\frac{1}{w_{0}}\tilde{\mathbf{p}}^{T}\tilde{H}(\tilde{\boldsymbol{\omega}})\tilde{\mathbf{p}}$
$\displaystyle=\frac{p_{0}^{2}}{w_{0}}+\frac{1}{w_{0}}\left(\frac{p_{0}}{w_{0}}\tilde{\mathbf{w}}-\tilde{\mathbf{p}}\right)^{T}\tilde{H}(\tilde{\boldsymbol{\omega}})\left(\frac{p_{0}}{w_{0}}\tilde{\mathbf{w}}-\tilde{\mathbf{p}}\right)\geq
0\>,$

where the inequality follows from the positive semi-definiteness of
$\tilde{H}$ (convexity of $\tilde{h}_{\textup{a}}$).

Finally, we prove the strict convexity of $h_{\textup{a}}$ provided that
$\tilde{h}_{\textup{a}}$ is strictly convex. Suppose $\tilde{h}_{\textup{a}}$
is strictly convex, then its Hessian $\tilde{H}$ is positive definite. Thus
the inequality in (36) becomes an equality only when $p_{0}=0$ and
$\frac{p_{0}}{w_{0}}\tilde{\mathbf{w}}-\tilde{\mathbf{p}}=\mathbf{0}$, which
implies that, for all $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$,
$\mathbf{p}^{T}H(\mathbf{w})\mathbf{p}=0$ if and only if
$\mathbf{p}=\mathbf{0}$. Therefore, $H$ is positive definite and
$h_{\textup{a}}$ is strictly convex. ∎

To this point, we have shown that a twice differentiable, convex approximation
$h_{\textup{a}}$ to the Maxwell-Boltzmann entropy $h$ can be extended from a
twice differentiable, convex function $\tilde{h}_{\textup{a}}$ that
approximates $\tilde{h}$, the restriction of $h$ on
$\tilde{\mathcal{R}}_{\mathbf{m}}$. Since $\tilde{\mathcal{R}}_{\mathbf{m}}$
is a bounded set, it is more tractable to generate a dataset that provides a
good coverage of $\tilde{\mathcal{R}}_{\mathbf{m}}$. The proposed procedure is
summarized as follows.

1 Sample a normalized dataset
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$
with $\tilde{\boldsymbol{\omega}}^{(i)}\in\tilde{\mathcal{R}}_{\mathbf{m}}$;

2 Construct a twice differentiable, convex approximation
$\tilde{h}_{\textup{a}}$ from
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$;

3 Extend
$\tilde{h}_{\textup{a}}\colon\tilde{\mathcal{R}}_{\mathbf{m}}\to\mathbb{R}$ to
$h_{\textup{a}}\colon\mathcal{R}_{\mathbf{m}}\to\mathbb{R}$ by taking
$h_{\textup{a}}(\mathbf{w})=w_{0}\,\tilde{h}_{\textup{a}}(\tilde{\mathbf{w}}/w_{0})+w_{0}\log
w_{0}$ for all $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$ as given in (33). This
leads to an approximate multiplier function
$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w}):=h_{\textup{a}}^{\prime}(\mathbf{w})=\left[\begin{array}[]{c}\tilde{h}_{\textup{a}}(\tilde{\mathbf{w}}/w_{0})+\frac{1}{w_{0}}\tilde{\mathbf{w}}^{T}\tilde{\boldsymbol{\alpha}}_{\textup{a}}(\tilde{\mathbf{w}}/w_{0})+\log
w_{0}+1\\\
\tilde{\boldsymbol{\alpha}}_{\textup{a}}(\tilde{\mathbf{w}}/w_{0})\end{array}\right]\>,$
(37) where $w_{0}$, $\tilde{h}_{\textup{a}}$, and its gradient
$\tilde{\boldsymbol{\alpha}}_{\textup{a}}$ are needed to compute the gradient
of $h_{\textup{a}}$;

Approximation procedure 1 Steps for constructing $h_{\textup{a}}$ from data

A flow chart that describes the process of mapping any
$\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$ to the approximate multiplier
$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w})$ using the constructed
approximation $\tilde{h}_{\textup{a}}$ is included in Figure 1. We next
discuss two approaches to construct twice differentiable and convex
$\tilde{h}_{\textup{a}}$ from data.

Figure 1: Flow chart: The process of mapping a realizable moment $\mathbf{w}$
to the approximate multiplier
$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w})$ from the data-driven
approximate entropy function $h_{\textup{a}}$ is illustrated. The data-driven
approximation $\tilde{h}_{\textup{a}}$ on the normalized realizable set
$\tilde{\mathcal{R}}_{\mathbf{m}}$ is evaluated at the normalized moment
$\tilde{\boldsymbol{\omega}}$. The result is then extended to the full
realizable set $\mathcal{R}_{\mathbf{m}}$ using the formulas given in (33) and
(37).

###  3.2 Data-driven approximations

In this subsection, we present two approaches to construct twice
differentiable and convex approximations to the Maxwell-Boltzmann entropy from
the normalized dataset
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$.
Section 3.2.1 summarizes the construction of one-dimensional rational cubic
spline interpolants proposed in [19] that maintain the convexity of the data.
Section 3.2.2 describes the procedure for building an approximate entropy
using neural networks, which works in arbitrary dimensions.

From here on, we focus on a simplified kinetic model for neutral particle
systems with uniform particle traveling speed in a purely scattering medium
with slab geometry, which takes the form

$\partial_{t}f(t,x,\mu)+\mu\partial_{x}f=\sigma_{\textup{s}}(\frac{\langle
f\rangle}{2}-f)\>,$ (38)

where the position $x\in\mathbb{R}$, angle $\mu\in[-1,1]$, and the scattering
crosssection $\sigma_{\textup{s}}>0$ is a constant in $t$, $x$, and $\mu$.
With the basis function $m_{i}(\mu)$ in $\mathbf{m}$ chosen to be the $i$-th
order Legendre polynomials, $i=0,\dots,N$, the moment system for (38) with
approximate entropy $h_{\textup{a}}$ can be derived following (15), which
gives

$\partial_{t}\mathbf{u}_{\textup{a}}+\partial_{x}\langle\mu\mathbf{m}G_{\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{u}_{\textup{a}})}\rangle=-\sigma_{\textup{s}}R\mathbf{u}_{\textup{a}}\>,$
(39)

where
$G_{\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w})}=e^{\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w})\cdot\mathbf{m}}$
and
$\hat{\boldsymbol{\alpha}}_{\textup{a}}(\mathbf{w}):=h_{\textup{a}}^{\prime}(\mathbf{w})$
for all $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$, and the diagonal matrix
$R=\operatorname{diag}(0,1,\dots,1)\in\mathbb{R}^{n\times n}$. This
simplification allows for constructing the approximate entropy
$h_{\textup{a}}$ using the one-dimensional shape-preserving splines in Section
3.2.1 when the moment order $N=1$. On the other hand, the neural network
approach considered in Section 3.2.2 can approximate the entropy function in
multi-dimensions, thus is not limited to the slab geometry case and can work
on $N>1$.

####  3.2.1 Shape-preserving splines

When moment order $N=1$, the moment $\mathbf{u}(t,x)$ in (39) is a vector in
$\mathbb{R}^{2}$ for each $(t,x)$, i.e., $n=2$. With $\mathbf{m}$ the Legendre
polynomial basis, the realizable set becomes
$\mathcal{R}_{\mathbf{m}}=\\{\mathbf{w}\in\mathbb{R}^{2}\colon\,|w_{1}|<w_{0}\\}$,
and the normalized realizable set is simply
$\tilde{\mathcal{R}}_{\mathbf{m}}=[-1,1]$. The problem then reduces to:
finding a twice differentiable, convex function
$\tilde{h}_{\textup{a}}\colon[-1,1]\to\mathbb{R}$ that approximates
$\tilde{h}$ given the position, function value, and derivative data
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$,
where $\tilde{\boldsymbol{\omega}}^{(i)}\in[-1,1]$ for all $i\in\mathcal{I}$.
The shape-preserving spline interpolation method proposed in [19] constructs a
one-dimensional cubic spline interpolant
$\tilde{h}_{\textsf{Sp}}(\tilde{\boldsymbol{\omega}})$ that is twice
differentiable and preserves the monotonicity and convexity of the provided
data. Specifically, for
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$
sampled from a convex function $\tilde{h}$, the cubic spline
$\tilde{h}_{\textsf{Sp}}(\tilde{\boldsymbol{\omega}})$ is constructed such
that, for all $i\in\mathcal{I}$,

$\tilde{h}_{\textsf{Sp}}(\tilde{\boldsymbol{\omega}}^{(i)})=\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\quad\tilde{h}_{\textsf{Sp}}^{\prime}(\tilde{\boldsymbol{\omega}}^{(0)})=\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(0)}),\quad\tilde{h}_{\textsf{Sp}}^{\prime}(\tilde{\boldsymbol{\omega}}^{(|\mathcal{I}|)})=\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(|\mathcal{I}|)}),\quad\mbox{and}\quad\tilde{h}_{\textsf{Sp}}^{\prime\prime}(\tilde{\boldsymbol{\omega}}^{(i)}_{-})=\tilde{h}_{\textsf{Sp}}^{\prime\prime}(\tilde{\boldsymbol{\omega}}^{(i)}_{+}),$
(40)

where the first equality is the interpolation condition, the second and third
equalities enforce exact derivative at the two end-points, and the last
equality guarantees that $\tilde{h}_{\textsf{Sp}}\in C^{2}([-1,1])$ with
$\tilde{\boldsymbol{\omega}}_{-}$ and $\tilde{\boldsymbol{\omega}}_{+}$
denoting the left and right limit at $\tilde{\boldsymbol{\omega}}$,
respectively. To impose convexity, the spline shape parameters are chosen such
that $\tilde{h}_{\textsf{Sp}}^{\prime\prime}(\tilde{\boldsymbol{\omega}})\geq
0$ between the interpolation points. The spline coefficients are then obtained
by solving the equations in (40) via a Gauss-Seidel approach. See [19] for
details. Once the spline approximation $\tilde{h}_{\textsf{Sp}}$ is
constructed, it is straightforward to compute
$\tilde{\boldsymbol{\alpha}}_{\textsf{Sp}}$ by taking the derivative of
$\tilde{h}_{\textsf{Sp}}$. The approximate multiplier
$\hat{\boldsymbol{\alpha}}_{\textsf{Sp}}$ can then be reconstructed from
$\tilde{h}_{\textsf{Sp}}$ and $\tilde{\boldsymbol{\alpha}}_{\textsf{Sp}}$ via
(37).

While this shape-preserving spline approach gives stable, accurate, and twice
differentiable convex approximations to convex functions, its application is
limited to one-dimensional functions. To the authors’ knowledge, there is no
existing twice differentiable, convexity-preserving spline approximations in
dimensions higher than one.

####  3.2.2 Neural networks

To handle the more general case that $\tilde{\mathcal{R}}_{\mathbf{m}}$ is in
a multi-dimensional space, we use neural networks to approximate the
restricted Maxwell-Boltzmann entropy. We denote the neural network
approximation as $\tilde{h}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}})$. For
simplicity, we construct
$\tilde{h}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}})$ via training standard
feed-forward, fully-connected neural networks, i.e., multilayer perceptrons.
Specifically, we consider neural networks of the form

$\tilde{h}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}};\boldsymbol{\theta})=(L^{D+1}\circ
L^{D}\circ\dots\circ L^{1}\circ L^{0})(\tilde{\boldsymbol{\omega}})\>\text{
with
}L^{\ell}(\mathbf{z})=\sigma(\boldsymbol{\phi}^{\ell}\cdot\mathbf{z}+\mathbf{b}^{\ell})\>,\quad\ell=0,\dots,D+1\>,$
(41)

where $\sigma$ denotes the activation function of the network,
$\boldsymbol{\phi}^{\ell}$ and $\mathbf{b}^{\ell}$ denote the weights and
biases at the $\ell$-th layer $L^{\ell}$, and $\boldsymbol{\theta}$ denotes
the collection of all parameters
$\\{(\boldsymbol{\phi}^{\ell},\mathbf{b}^{\ell})\\}_{\ell=0}^{D}$. The number
of network parameters depends on the size of the network, e.g., a network of
depth $D$ and width $W$ has weights

$\boldsymbol{\phi}^{0}\in\mathbb{R}^{(n-1)\times
W},\quad\boldsymbol{\phi}^{\ell}\in\mathbb{R}^{W\times
W},\,\ell=1,\dots,D\,,\quad\mbox{and}\quad\boldsymbol{\phi}^{D+1}\in\mathbb{R}^{W\times
1},$ (42)

and biases $\mathbf{b}^{\ell}\in\mathbb{R}^{W}$, $\ell=0,\dots,D$ and
$\mathbf{b}^{D+1}\in\mathbb{R}$, which sum up to $DW^{2}+(n+D+1)W+1$
parameters. In order to guarantee that the network approximation
$\tilde{h}_{\textsf{NN}}$ is twice differentiable, we choose to use the smooth
softplus activation function, i.e.,
$\sigma(\mathbf{z})=\log(1+\exp(\mathbf{z}))$. In general, the network
function $\tilde{h}_{\textsf{NN}}$ in (41) is not guaranteed to be convex with
respect to $\tilde{\boldsymbol{\omega}}$. In this paper, we verify the
convexity of the constructed network approximations by confirming that the
Hessian of $\tilde{h}_{\textsf{NN}}$ is positive semidefinite on a dense test
dataset. To enforce convexity of $\tilde{h}_{\textsf{NN}}$, one can either
penalize negative eigenvalues of the Hessian of $\tilde{h}_{\textsf{NN}}$ in
the training process, or use special network architectures, e.g., [9], that
guarantee convexity. The former imposes convexity as a soft constraint on the
training dataset and is feasible only when the eigenvalues of Hessian can be
obtained efficiently. The latter gives networks that are convex on the entire
input domain, while the special architectures and limitations of the
parameters may complicate the training process.

###  3.3 Data acquisition strategy

The spline and neural network approximations discussed in Section 3.2 both
rely on some normalized dataset
$\\{(\tilde{\boldsymbol{\omega}}^{(i)},\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)}),\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)}))\\}_{i\in\mathcal{I}}$.
A naive approach to construct the normalized data set is to first sample
$\\{\tilde{\boldsymbol{\omega}}^{(i)}\\}_{i\in\mathcal{I}}$ from
$\tilde{\mathcal{R}}_{\mathbf{m}}$, evaluate $h$ and
$\hat{\boldsymbol{\alpha}}$ at
$\mathbf{w}=[1,\tilde{\boldsymbol{\omega}}^{T}]^{T}$ for each
$\tilde{\boldsymbol{\omega}}\in\\{\tilde{\boldsymbol{\omega}}^{(i)}\\}_{i\in\mathcal{I}}$,
and then compute
$\\{\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)})\\}_{i\in\mathcal{I}}$ and
$\\{\tilde{\boldsymbol{\alpha}}(\tilde{\boldsymbol{\omega}}^{(i)})\\}_{i\in\mathcal{I}}$
via (26) and (28). This approach is computationally expensive since the
evaluation of $h$ and $\hat{\boldsymbol{\alpha}}$, i.e., the solution of the
optimization problem (23), is required at each sampled moment. Further, the
accuracy of the dataset is limited by the tolerance in the optimization
procedure.

To address this issue, we propose to first sample the normalized multipliers
$\\{\tilde{\boldsymbol{\alpha}}^{(i)}\\}_{i\in\mathcal{I}}$ and then compute
the associated normalized moments
$\\{\tilde{\boldsymbol{\omega}}^{(i)}\\}_{i\in\mathcal{I}}$ and restricted
entropy function values
$\\{\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)})\\}_{i\in\mathcal{I}}$ from
$\\{\tilde{\boldsymbol{\alpha}}^{(i)}\\}_{i\in\mathcal{I}}$. This approach
avoids solving the optimization problem (23) in the data acquisition
procedure, and the accuracy of the data is only restricted by the quadrature
rule used to perform numerical integrations. Specifically, the proposed
sampling strategy is as follows.

1 Sample the normalized multipliers
$\\{\tilde{\boldsymbol{\alpha}}^{(i)}\\}_{i\in\mathcal{I}}$ from
$\mathbb{R}^{n-1}$;

2 For every $i\in\mathcal{I}$, set
$\boldsymbol{\alpha}^{(i)}=[\alpha_{0}^{(i)},(\tilde{\boldsymbol{\alpha}}^{(i)})^{T}]$
with $\alpha_{0}^{(i)}$ calculated such that $w_{0}^{(i)}:=\langle
m_{0}G_{\boldsymbol{\alpha}^{(i)}}\rangle=1$;

3 For every $i\in\mathcal{I}$, compute the moment and entropy function value
associated to $\boldsymbol{\alpha}^{(i)}$ by taking
$\mathbf{w}^{(i)}=\hat{\mathbf{w}}(\boldsymbol{\alpha}^{(i)})=\langle\mathbf{m}G_{\boldsymbol{\alpha}^{(i)}}\rangle\quad\mbox{and}\quad
h(\mathbf{w}^{(i)})=\langle\eta(G_{\hat{\boldsymbol{\alpha}}(\mathbf{w}^{(i)})})\rangle=\langle\eta(G_{\boldsymbol{\alpha}^{(i)}})\rangle\>.$
(43) These relations follow from (9)–(14) for a general entropy $\eta$;

4 Denote $\mathbf{w}^{(i)}=[w_{0}^{(i)},(\tilde{\mathbf{w}}^{(i)})^{T}]^{T}$.
Since $w_{0}^{(i)}=1$, the normalized moment
$\tilde{\boldsymbol{\omega}}^{(i)}=\tilde{\mathbf{w}}^{(i)}$ and the
restricted entropy function value
$\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)})=h(\mathbf{w}^{(i)})$;

Sampling strategy 2 Steps for efficient sampling data from
$\tilde{\mathcal{R}}_{\mathbf{m}}$

##  4 Numerical results

Numerical results for the proposed data-driven entropy-based closures are
presented in this section. Section 4.1 gives the implementation details,
including the sampling of datasets and the training setting for the neural
networks. Section 4.2 summarizes the constructed closures and their
approximation accuracy. Section 4.3 compares the computation time between the
neural network closures and the standard optimization approach for Maxwell-
Boltzmann entropy-based closures. Finally, the data-driven closures are
applied to solve a plane source benchmark problem, and the results are
reported in Section 4.4.

###  4.1 Implementation details

####  4.1.1 Sampled data

We follow the approach outlined in Section 3.3 to sample the data needed in
the construction of data-driven closures considered in this section. With the
proposed approach, we only need to determine a strategy for sampling the
normalized multipliers
$\\{\tilde{\boldsymbol{\alpha}}^{(i)}\\}_{i\in\mathcal{I}}$, and the remaining
parts of the data can be computed accordingly. In the tests considered in this
section, we sample each entry of $\tilde{\boldsymbol{\alpha}}$ on a uniform
grid of size $s$ in a prescribed interval
$[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]$, e.g., when
$\tilde{\boldsymbol{\alpha}}\in\mathbb{R}^{n-1}$, then each sample
$\tilde{\boldsymbol{\alpha}}^{(i)}$ is a point on the uniform tensor mesh of
size $s^{n-1}$ in the hypercube
$[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]^{n-1}$. With
$\tilde{\alpha}_{\min}<0$ sufficiently small and $\tilde{\alpha}_{\max}>0$
sufficiently large, this sampling strategy allocates more data points near the
boundary of the normalized realizable set $\tilde{\mathcal{R}}_{\mathbf{m}}$
for the Maxwell-Boltzmann entropy $h_{\mbox{\tiny{\sc MB}}}$, which
significantly improves the quality of data-driven closures since
$h_{\mbox{\tiny{\sc MB}}}$ is known to vary drastically near the boundary of
$\tilde{\mathcal{R}}_{\mathbf{m}}$.

We choose $[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]=[-65,65]$ when
training all data-driven closures for the case $N=1$. This choice covers
normalized moments in $[-1+0.015,1-0.015]$, whereas the normalized realizable
set $\tilde{\mathcal{R}}_{\mathbf{m}}=[-1,1]$. For the spline closures, the
data size $s$ is equivalent to the size of the spline approximation, i.e., the
number of interpolation points. For the network closures, we choose the data
size $s=10,000$ for training networks with various depths and widths. As for
training neural network closures for the case $N=2$, we choose
$[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]=[-10,10]$ with $100$ grid
points in the $\tilde{\alpha}_{1}$ direction and $50$ grid points in the
$\tilde{\alpha}_{2}$ direction.

In Section 4.2, we report the training and test errors of data-driven closures
constructed with the training datasets described above. The test datasets are
sampled using the same approach, but with a denser uniform grid for each entry
of $\tilde{\boldsymbol{\alpha}}$. In addition, to test the accuracy of the
reconstruction steps in (33) and (37), we allow the zero-th moment $w_{0}$ to
vary in the test datasets while $w_{0}$ in the training datasets is always
normalized to one. When $N=1$, the spline and network approximations are both
tested on a dataset with 160 $w_{0}$ values evenly distributed on interval
$[10^{-8},8]$. For each $w_{0}$ value, $\tilde{\boldsymbol{\alpha}}$ is
sampled on a uniform grid of 52,000 points with
$[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]=[-65,65]$. When $N=2$, the test
dataset for network approximations also has 160 evenly distributed $w_{0}$ in
$[10^{-8},8]$. For each $w_{0}$, we choose
$[\tilde{\alpha}_{\min},\tilde{\alpha}_{\max}]=[-10,10]$ and sample
$\tilde{\boldsymbol{\alpha}}$ on a uniform tensor grid with 200 grid points in
both the $\tilde{\alpha}_{1}$ and $\tilde{\alpha}_{2}$ directions.

####  4.1.2 Neural network training

The neural network approximations are implemented as Keras models with
tensorflow 2.0 in Python.

##### Network initialization

In the neural network training procedure, we initialize the network parameters
using a strategy considered in [32], which aims to normalize the variance of
each network layer output with respect to a standard normal random variable
input to one. This strategy is motivated from the common Kaiming
initialization [26] for neural networks with the Rectified Linear Unit (ReLU)
activation function. We adopt the initialization strategy and apply it on the
softplus activation function, $\sigma(\mathbf{z})=\log(1+\exp(\mathbf{z}))$,
which is used here to guarantee the regularity of network approximations.
Specifically, for a network of depth $D$ and width $W$, we initialize all
biases to be zero, i.e., $\mathbf{b}^{\ell}=0$, $\ell=0,\dots,D+1$, and draw
the initial weights $\boldsymbol{\theta}^{\ell}$ from
$\mathcal{N}(0,s^{2}_{\ell})$, the zero-mean normal distribution with variance
$s^{2}_{\ell}$, where the variance for each layer is given by

$s^{2}_{0}=\frac{1}{(n-1)\sigma^{\prime}(0)^{2}(1+\sigma(0)^{2})},\quad\mbox{and}\quad
s^{2}_{\ell}=\frac{1}{W\sigma^{\prime}(0)^{2}(1+\sigma(0)^{2})},\quad\ell=1,\dots,D+1\>.$
(44)

##### Training and validation data

We follow the standard practice and split the sampled data in Section 4.1.1
into the training and validation datasets. In each test, we uniformly draw 10%
of the sampled data to be the validation data, i.e., a 90–10 split. From here
on, we denote the index sets of the training and validation datasets as
$\mathcal{I}^{{\textup{train}}}$ and $\mathcal{I}^{{\textup{val}}}$,
respectively.

##### Network output — multiplier and symmetry

The network closure aims to provide an approximate multiplier
$\hat{\boldsymbol{\alpha}}_{\textsf{NN}}(\mathbf{w})$ at a given realizable
moment $\mathbf{w}\in\mathcal{R}_{\mathbf{m}}$. It follows from (37) that both
the network function value $\tilde{h}_{\textsf{NN}}$ and the network gradient
$\tilde{\boldsymbol{\alpha}}_{\textsf{NN}}$ are needed to compute
$\hat{\boldsymbol{\alpha}}_{\textsf{NN}}(\mathbf{w})$. To obtain the gradient
$\tilde{\boldsymbol{\alpha}}_{\textsf{NN}}$ from the fully-connected network
$\tilde{h}_{\textsf{NN}}$ discussed in Section 3.2.2, we append at the end of
the fully-connected network a custom layer that performs automatic
differentiation by calling GradientTape.gradient() in tensorflow. With the
choice of Legendre polynomials as moment basis $\mathbf{m}$, the entropy
function $h$ is known to be symmetric with respect to the moments of odd
orders, e.g., $h(\mathbf{w})=h(\mathbf{w}^{*})$ if $w_{i}^{*}=w_{i}$ for even
$i$ and $w_{i}^{*}=-w_{i}$ for odd $i$. To impose this symmetry in the neural
network approximations, we define

${\tilde{h}_{\textsf{NN}}}^{\text{sym}}(\tilde{\boldsymbol{\omega}}):=\frac{1}{2}({\tilde{h}_{\textsf{NN}}}(\tilde{\boldsymbol{\omega}})+{\tilde{h}_{\textsf{NN}}}(\tilde{\boldsymbol{\omega}}^{*}))\quad\text{for
all }\tilde{\boldsymbol{\omega}}\in\tilde{\mathcal{R}}_{\mathbf{m}}\>,$ (45)

where $\tilde{\omega}_{i}^{*}=\tilde{\omega}_{i}$ for even $i$ and
$\tilde{\omega}_{i}^{*}=-\tilde{\omega}_{i}$ for odd $i$. This symmetric
network approximation is used in the plane source tests discussed in Section
4.4.

##### Loss function

The training procedure aims to find the optimal network parameter
$\boldsymbol{\theta}$ by minimizing a loss function. The loss function
considered here takes the form

$L(\boldsymbol{\theta};\mathcal{I})=E_{h}^{2}(\boldsymbol{\theta};\mathcal{I})+\lambda(\mathcal{I})E_{\mathbf{w}}^{2}(\boldsymbol{\theta};\mathcal{I})\>,$
(46)

where
$\lambda(\mathcal{I})=(\sum_{i\in\mathcal{I}}{\|\mathbf{w}^{(i)}\|^{2}})^{-1}$
with $\mathbf{w}^{(i)}=[1,(\tilde{\boldsymbol{\omega}}^{(i)})^{T}]^{T}$ and
the errors are defined as

$E_{h}^{2}(\boldsymbol{\theta};\mathcal{I}):=\sum_{i\in\mathcal{I}}|\tilde{h}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta})-\tilde{h}(\tilde{\boldsymbol{\omega}}^{(i)})|^{2},\quad
E_{\mathbf{w}}^{2}(\boldsymbol{\theta};\mathcal{I}):=\sum_{i\in\mathcal{I}}\lVert{\mathbf{w}}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta})-\mathbf{w}^{(i)}\rVert^{2}.$
(47)

Here
${\mathbf{w}}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta})$
is the approximate moment in $\mathcal{R}_{\mathbf{m}}$ given by the neural
network closure, i.e.,
${\mathbf{w}}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta}):=\langle\exp(\mathbf{m}\cdot\hat{\boldsymbol{\alpha}}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta})\rangle$
with $\hat{\boldsymbol{\alpha}}_{\textsf{NN}}$ the approximate multiplier that
can be computed using $\tilde{h}_{\textsf{NN}}$ and
$\tilde{\boldsymbol{\alpha}}_{\textsf{NN}}$ via (37). We choose to include the
moment reconstruction error $E_{\mathbf{w}}$ in the loss function rather than
the multiplier error

$E_{\boldsymbol{\alpha}}^{2}(\boldsymbol{\theta};\mathcal{I}):=\sum_{i\in\mathcal{I}}\|\hat{\boldsymbol{\alpha}}_{\textsf{NN}}(\tilde{\boldsymbol{\omega}}^{(i)};\boldsymbol{\theta})-\hat{\boldsymbol{\alpha}}^{(i)}\|^{2},$
(48)

because, in the moment system (15), the closures are used to evaluate the flux
term and the collision term, and both of which are integrated quantities over
the velocity/angular space. Thus we conjecture that the moment reconstruction
error $E_{\mathbf{w}}$ would be a more reliable measure for the closure
accuracy than the multiplier error $E_{\boldsymbol{\alpha}}$. We will verify
this conjecture in the plane source tests reported in Section 4.4.

##### Hyperparameters and optimizer

The networks were scheduled to be trained over 15,000 epochs with a batch size
50. The training process terminates early if one of the following two criteria
is satisfied: (i) the validation error
$E_{\mathbf{w}}^{2}(\boldsymbol{\theta},\mathcal{I}^{{\textup{val}}})$ is less
than tolerance $10^{-8}$, and (ii) there is no improvement greater than
$10^{-9}$ on
$E_{\mathbf{w}}^{2}(\boldsymbol{\theta},\mathcal{I}^{{\textup{val}}})$ in the
past 1,500 epochs. To minimize the loss function (46), the Adam optimizer
[31], a first-order stochastic gradient descent algorithm, is used with an
adaptive learning rate $r=10^{-3+(m/5000)}$, where $m$ is the counter for the
current epoch.

###  4.2 Constructed data-driven closures

Here we list several data-driven closures constructed using the training
datasets specified in Section 4.1.1. The neural network closures are trained
with the specifics detailed in Section 4.1.2 and the spline closures are built
with the procedure described in Section 3.2.1, where a Gauss-Seidel approach
is used to solve for the spline coefficients up to a $10^{-12}$ relative
tolerance. From here on, we denote the neural network closures of various
sizes as $\texttt{NN}_{N}^{D\times W}$, where $D$ and $W$ are respectively the
depth and width of the network, and $N$ is the moment order. The spline
closures are denoted as $\texttt{S}^{P}$, where $P$ is the number of
interpolation points in the spline. Tables 1 and 2 report the training and
test errors of various spline and neural network closures for moment order
$N=1$. In both tables, the sizes are the number of parameters in the closures.
For neural network closures, the formula is given in Section 3.2.2. For spline
closures, the sizes are identical to the number of interpolation points. Here
the reported training and test errors for the entropy function values,
moments, and multipliers, are the relative root-mean-square errors (RMSEs),
e.g., the training and test errors in function value are defined respectively
as

$\textstyle
e_{h}^{{\textup{train}}}=(\sum_{i\in\mathcal{I}^{{\textup{train}}}}|h_{\textup{a}}(\tilde{\boldsymbol{\omega}}^{(i)})|^{2})^{-\frac{1}{2}}E_{h}(\boldsymbol{\theta},\mathcal{I}^{{\textup{train}}})\>,\quad\mbox{and}\quad
e_{h}^{{\textup{test}}}=(\sum_{i\in\mathcal{I}^{{\textup{test}}}}|h_{\textup{a}}(\tilde{\boldsymbol{\omega}}^{(i)})|^{2})^{-\frac{1}{2}}E_{h}(\boldsymbol{\theta},\mathcal{I}^{{\textup{test}}})\>,$
(49)

where $E_{h}$ is the RMSE as defined in the loss function (46). The test and
training errors for the moments and multipliers are all defined analogously.
Here the test datasets for $N=1$ and $2$ are also specified in Section 4.1.1.

Table 1 compares the spline and neural network closures that have testing
moment errors ($e_{\mathbf{w}}^{{\textup{test}}}$) around $10^{-2}$,
$10^{-3}$, and $10^{-4}$. Table 2 compares the accuracy of the spline and
neural network closures that are of similar sizes (around 100, 500, and 1000).
The results reported in these two tables suggest that neural network
approximations generally need many more parameters to achieve similar accuracy
as the spline approximations. When the two approximations are of similar size,
the test errors of the spline approximation are at least two orders of
magnitude lower than the ones of the network approximation, and the different
increases as the number of parameter grows. We note that, for spline
approximations, $e_{h}^{{\textup{train}}}$ is zero by construction while
$e_{\boldsymbol{\alpha}}^{{\textup{train}}}$ is nonzero since the spline
interpolation requires derivative to be exact only at the end points.

Closure | $\texttt{NN}^{1\times 15}_{1}$ | $\texttt{NN}^{3\times 15}_{1}$ | $\texttt{NN}^{5\times 30}_{1}$ | $\texttt{S}^{30}$ | $\texttt{S}^{60}$ | $\texttt{S}^{130}$
---|---|---|---|---|---|---
Size | 286 | 766 | 4741 | 30 | 60 | 130
$e_{h}^{{\textup{train}}}$ | 6.28e-3 | 5.11e-4 | 1.14e-4 |  0e0 |  0e0 |  0e0
$e_{\mathbf{w}}^{{\textup{train}}}$ | 7.31e-3 | 7.70e-4 | 1.93e-4 | 5.61e-3 | 4.52e-4 | 1.91e-5
$e_{\hat{\boldsymbol{\alpha}}}^{{\textup{train}}}$ | 1.00e-1 | 1.62e-2 | 4.36e-3 | 1.03e-3 | 7.02e-5 | 3.29e-6
$e_{h}^{{\textup{test}}}$ | 3.03e-3 | 2.46e-4 | 5.49e-5 | 4.30e-3 | 2.28e-4 | 9.48e-6
$e_{\mathbf{w}}^{{\textup{test}}}$ | 7.31e-3 | 7.69e-4 | 1.92e-4 | 1.06e-2 | 9.36e-4 | 7.86e-5
$e_{\hat{\boldsymbol{\alpha}}}^{{\textup{test}}}$ | 1.01e-1 | 1.64e-2 | 4.35e-3 | 1.11e-3 | 9.69e-5 | 8.67e-6
Table 1: Training and test errors in terms of entropy function values ($h$), moments ($\mathbf{w}$), and multipliers ($\hat{\boldsymbol{\alpha}}$) for neural network and spline approximations in the case $N=1$ with testing moment errors ($e_{\mathbf{w}}^{{\textup{test}}}$) around $10^{-2}$, $10^{-3}$, and $10^{-4}$. Closure | $\texttt{NN}^{2\times 15}_{1}$ | $\texttt{NN}^{0\times 45}_{1}$ | $\texttt{NN}^{4\times 15}_{1}$ | $\texttt{S}^{100}$ | $\texttt{S}^{500}$ | $\texttt{S}^{1000}$
---|---|---|---|---|---|---
Size | 136 | 526 | 1006 | 100 | 500 | 1000
$e_{h}^{{\textup{train}}}$ | 8.48e-2 | 1.76e-3 | 3.77e-4 |  0e0 |  0e0 |  0e0
$e_{\mathbf{w}}^{{\textup{train}}}$ | 7.14e-2 | 2.43e-3 | 5.87e-4 | 5.47e-5 | 8.78e-8 | 5.50e-9
$e_{\hat{\boldsymbol{\alpha}}}^{{\textup{train}}}$ | 4.71e-1 | 4.07e-2 | 1.18e-2 | 9.40e-6 | 1.53e-8 | 9.61e-10
$e_{h}^{{\textup{test}}}$ | 4.12e-2 | 8.48e-4 | 1.82e-4 | 2.75e-5 | 4.42e-8 | 2.79e-9
$e_{\mathbf{w}}^{{\textup{test}}}$ | 7.14e-2 | 2.42e-3 | 5.88e-4 | 1.80e-4 | 1.32e-6 | 1.64e-7
$e_{\hat{\boldsymbol{\alpha}}}^{{\textup{test}}}$ | 4.77e-1 | 4.10e-2 | 1.20e-2 | 1.96e-5 | 1.47e-7 | 1.85e-8
Table 2: Training and test errors in terms of entropy function values ($h$),
moments ($\mathbf{w}$), and multipliers ($\hat{\boldsymbol{\alpha}}$) for
neural network and spline approximations of size around 100, 500, and 1000 in
the case $N=1$.

To further explore the relation between the test errors and numbers of
parameters for the spline and neural network approximations, we have trained
and tested more approximations of different sizes and reported their test
errors in Figure 2. It can be observed from Figure 2 that (i) the spline
approximation is indeed third-order accurate; (ii) the networks generally
require substantially more parameters than the splines to reach the same level
of accuracy; and (iii) the network depth is more important than its width in
this test, i.e., for networks of similar size, deep but narrow networks
outperform wide but shallow ones. We also note that $e_{\boldsymbol{\alpha}}$
is often orders of magnitude higher than $e_{h}$ and $e_{\mathbf{w}}$. This is
because $\hat{\boldsymbol{\alpha}}$ becomes unbounded at the realizable
boundary, which makes the approximation difficult.

(a) Entropy function error $e_{h}^{{\textup{test}}}$ (b) Moment error
$e_{\mathbf{w}}^{{\textup{test}}}$ (c) Multiplier error
$e_{\boldsymbol{\alpha}}^{{\textup{test}}}$

Figure 2: Relative test errors for the constructed data-driven closures for
$N=1$. Neural network closures of various depths with width $W=$ 15, 30, and
45 are denoted as blue crosses, circles, and triangles. Spline closures of
various sizes are denoted as orange circles.

Table 3 shows three neural network approximations for the case $N=2$ that have
relative moment reconstruction (test) errors around $10^{-2}$, $10^{-3}$, and
$10^{-4}$. Here the spline closures are not included since the convex spline
approximation in Section 3.2.1 is restricted to one dimension and the
normalized realizable set $\tilde{\mathcal{R}}_{\mathbf{m}}$ is in
$\mathbb{R}^{2}$ when $N=2$. The results in Table 3 indicate that more
parameters are needed in the network approximations for $N=2$ to achieve
similar levels of accuracy compared to the network approximations for $N=1$,
which is as expected since the dimension of the domain of approximation
functions grows as $N$ increases.

Closures | $\texttt{NN}^{1\times 15}_{2}$ | $\texttt{NN}^{3\times 30}_{2}$ | $\texttt{NN}^{4\times 45}_{2}$
---|---|---|---
Size | 301 | 2911 | 8461
$e_{h}^{{\textup{train}}}$ | 1.89e-2 | 4.73e-4 | 3.24e-4
$e_{\mathbf{w}}^{{\textup{train}}}$ | 1.74e-2 | 8.13e-4 | 5.67e-4
$e_{\boldsymbol{\alpha}}^{{\textup{train}}}$ | 3.98e-1 | 2.25e-1 | 2.17e-1
$e_{h}^{{\textup{test}}}$ | 9.02e-3 | 2.34e-4 | 1.55e-4
$e_{\mathbf{w}}^{{\textup{test}}}$ | 1.76e-2 | 8.08e-4 | 5.63e-4
$e_{\boldsymbol{\alpha}}^{{\textup{test}}}$ | 4.05e-1 | 2.31e-1 | 2.21e-1
Table 3: Training and test errors in terms of entropy function values ($h$),
moments ($\mathbf{w}$), and multipliers ($\hat{\boldsymbol{\alpha}}$) for
neural network approximations in the case $N=2$ with testing moment errors
($e_{\mathbf{w}}^{{\textup{test}}}$) around $10^{-2}$, $10^{-3}$, and
$10^{-4}$.

Figure 3 reports the test errors of more neural network approximations of
various width and depth for $N=2$. The results are similar to the ones in
Figure 2 for the $N=1$ case, and we also observe that, to improve the network
accuracy, it is preferable to increase the network depth than the width.

(a) Entropy function error $e_{h}^{{\textup{test}}}$ (b) Moment error
$e_{\mathbf{w}}^{{\textup{test}}}$ (c) Multiplier error
$e_{\boldsymbol{\alpha}}^{{\textup{test}}}$

Figure 3: Relative test errors for the constructed neural network closures for
$N=2$. Closures of various depths with width $W=$ 15, 30, and 45 are denoted
as blue crosses, circles, and triangles.

The data-driven closures considered in this section are all twice
differentiable by construction. The spline approximations are guaranteed to be
convex. On the other hand, the convexity of the network approximations is
verified by checking the positive semidefiniteness of the Hessian matrix at
each point of the excessively refined test datasets given in Section 4.1.1.
For both the $N=1$ and $N=2$ cases, we do not observe any violation of the
convexity condition on the test datasets in the network approximations
considered here.

###  4.3 Computation time comparison

In this section, we compare the computation time needed for the network and
optimization approaches to map given moments $\mathbf{w}$ to the associated
multipliers $\hat{\boldsymbol{\alpha}}(\mathbf{w})$. The network approach maps
$\mathbf{w}$ to $\hat{\boldsymbol{\alpha}}(\mathbf{w})$ via a direct
evaluation, whereas the optimization approach generates
$\hat{\boldsymbol{\alpha}}(\mathbf{w})$ by solving (23). The optimization
algorithm compared here is a Newton-based solver with backtracking line search
considered in [8]. We set the optimization tolerance to be $10^{-8}$
throughout the tests in this section. It is known [8, 7, 6] that, for moments
close to the realizable boundary, the Hessian in (23) becomes ill-conditioned,
which leads to more iteration counts in the optimization. Here we will confirm
this result in the cases when the moment order $N=1$ and $2$, and verify that
the cost of evaluating neural network approximations remains roughly constant.

We first compare the network and optimization approaches in the $N=1$ case on
a set of moments $\\{\mathbf{w}^{(j)}\\}_{j=1}^{200}$ with
$\mathbf{w}^{(j)}:=[1,w_{1}^{(j)}]$ and $w_{1}^{(j)}$ evenly distributed in
$[0,1]$. Here we consider two variants of the optimization approach,
$\texttt{opt}_{\texttt{Q}}$ and $\texttt{opt}_{\texttt{A}}$, where
$\texttt{opt}_{\texttt{Q}}$ performs $\langle\cdot\rangle$, the angular
integration in (23), using a 30-point Gauss-Legendre quadrature, and
$\texttt{opt}_{\texttt{A}}$ uses integration based on an analytic formula,
which is available for the special case $N=1$. The network evaluation is
performed using Keras function model.predict(). Figure 4(a) shows the
iteration counts needed for the optimization algorithm to converge and Figure
4(b) shows the computation time for both the optimization and network
evaluation. In these results, the computation is restricted to one CPU core,
and the network is evaluated at one sample of moments at a time, i.e.,
model.predict() is called 200 times to evaluate the network at
$\\{\mathbf{w}^{(j)}\\}_{j=1}^{200}$, with a single point $\mathbf{w}^{(j)}$
evaluated each time. Figure 4(c) reports the median of network evaluation time
as well as the first and third quartile values over 20 trials. The results in
Figure 4 confirm that the iteration counts for both
$\texttt{opt}_{\texttt{Q}}$ and $\texttt{opt}_{\texttt{A}}$ grow as the
moments gets closer to the realizable boundary and the increased iteration
counts are reflected in the computation time. It can also be observed from the
results that the single-point neural network evaluation is significantly
slower than both variants of the optimization approach and that evaluating a
smaller network ($\texttt{NN}^{0\times 15}_{1}$) is only slightly faster than
evaluating a larger one ($\texttt{NN}^{4\times 45}_{1}$) at a single point. We
suspect that the overhead incurred when calling the Keras function
model.predict() may dominant the computation time in network evaluations. In
order to minimize the overhead, we next use model.predict() to perform batch
evaluation and compare the result to the optimization time in the next test.

(a) Iteration count

(b) Compuation time

(c) Compuation time (networks)

Figure 4: Iteration counts and computation time comparison between two
variants of the optimization approach and the network approximations for
multiplier calculation with $N=1$. Here the calculations are performed on
moments at various distance to the realizable boundary. Neural networks are
evaluated point-wisely, which results in longer evaluation time. Figure 4(c)
shows a zoomed-in version of the medium network computation time for the small
($\texttt{NN}^{0\times 15}_{1}$, green) and large ($\texttt{NN}^{4\times
45}_{1}$, red) networks over 20 trials, where the curves in lighter color
denote the upper and lower bounds of the interquartile range (IQR) of the
measured time.

Table 4 reports the total time needed for the optimization and network
approaches to compute multipliers for 100, 200, 500, 1,000, and 10,000
moments. These results are median values of 20 trails, in which the first and
third quartiles are both within $8\%$ deviation from the median. Here the
networks are evaluated by calling model.predict() once to evaluate all
moments, i.e., batch evaluation. This reduces the overhead for calling
model.predict() in network evaluation. The results in Table 4 show that the
batch network evaluation is about 6–60x faster than the optimization approach,
and the advantage grows as the number of moments increases. This confirms the
conjecture that the overhead in calling the Keras function model.predict() is
indeed the dominant cost in network evaluation, especially when the number of
moments to be evaluated is small. We note that the results reported in Table 4
are closer to realistic than the ones in Figure 4, since batch evaluation is
often allowed when solving moment systems with closures, e.g., at a given time
step, one can collect moments from all spatial points and perform a batch
evaluation to obtain all multipliers needed to evolve the time step.

$\\#$ of moments | $\texttt{opt}_{\texttt{Q}}$ | $\texttt{opt}_{\texttt{A}}$ | $\texttt{NN}^{0\times 15}_{1}$ | $\texttt{NN}^{4\times 45}_{1}$
---|---|---|---|---
100 | 1.76e-1 | 1.54e-1 | 2.46e-2 | 2.51e-2
200 | 3.56e-1 | 3.08e-1 | 2.74e-2 | 2.93e-2
500 | 9.13e-1 | 7.98e-1 | 3.21e-2 | 3.94e-2
1000 | 1.89e0 | 1.71e0 | 4.26e-2 | 5.19e-2
10000 | 1.99e1 | 1.87e1 | 2.24e-1 | 3.11e-1
Table 4: Comparison between the computation time for the optimization approach
with quadrature $(\texttt{opt}_{\texttt{Q}})$ and analytic
$(\texttt{opt}_{\texttt{A}})$ integrals and the batch evaluation time for the
small ($\texttt{NN}^{0\times 15}_{1}$) and large ($\texttt{NN}^{4\times
45}_{1}$) networks in the $N=1$ case. The time (sec) for evaluating 100, 200,
500, 1000, and 10000 moments are reported.

Next, we repeat the comparison between the optimization and neural network
approaches in the case $N=2$. Here we test both approaches with moments moving
towards the boundary of the normalized realizable set $\tilde{\mathcal{R}}$ in
four different directions, as shown in Figure 5(a). The tested moments are 200
evenly distributed points between the origin and the boundary point in each
direction. The iteration count for $\texttt{opt}_{\texttt{Q}}$ is reported in
Figure 5(b) along each direction. Although analytic integration formula exists
for the angular integrations in (23) when $N=2$, we do not include the
$\texttt{opt}_{\texttt{A}}$ into the comparison here since there is no
substantial difference between $\texttt{opt}_{\texttt{Q}}$ and
$\texttt{opt}_{\texttt{A}}$ when $N=1$. The $\texttt{opt}_{\texttt{Q}}$
optimization time and the network single-point evaluation time are compared in
Figure 5(c), where we observe higher optimization time and nearly constant
network evaluation time for moments near the realizable boundary. However, as
in the $N=1$ case, the single evaluation time for the networks is still higher
than the optimization time. Table 5 reports the total time for computing the
multiplier $\hat{\boldsymbol{\alpha}}$ from various numbers of given moments,
averaged over the four test directions. As in the results in Table 4, the
network evaluations here are performed in batches, i.e., model.predict() is
only called once to evaluate all moments. These results are median values of
20 trails, in which the first and third quartiles are both within $9\%$
deviation of the median. We observe that, for both the small
($\texttt{NN}^{0\times 15}_{2}$) and large ($\texttt{NN}^{4\times 45}_{2}$)
networks, the batch evaluation time is 7–62x faster than the optimization,
which is slightly higher than the speedup observed in Table 4 for the $N=1$
case. The results again confirm that the overhead of calling model.predict()
dominants the evaluation cost when evaluating smaller number of moments. We
also expect the speedup of using the neural network closures to become more
significant as the moment order $N$ goes higher, in which case the
optimization time increases drastically while the network evaluation time is
expected to remain in the same order of magnitude.

(a) Realizable set $\tilde{\mathcal{R}}_{\mathbf{m}}$

(b) Iteration count

(c) Computation time

Figure 5: Iteration counts and computation time comparison between the optimization approach ($\texttt{opt}_{\texttt{Q}}$) and two neural network approximations ($\texttt{NN}^{0\times 15}_{2}$ and $\texttt{NN}^{4\times 45}_{2}$) for multiplier calculation with $N=2$. The calculations are performed along four directions towards the realizable boundary as shown in Figure 5(a). Neural networks are evaluated point-wisely, which results in longer evaluation time. $\\#$ of moments | $\texttt{opt}_{\texttt{Q}}$ | $\texttt{NN}^{0\times 15}_{2}$ | $\texttt{NN}^{4\times 45}_{2}$
---|---|---|---
100 | 2.18e-1 | 2.84e-2 | 3.08e-2
200 | 4.25e-1 | 3.01e-2 | 3.46e-2
500 | 1.07e0 | 3.58e-2 | 4.50e-2
1000 | 2.15e0 | 4.68e-2 | 6.11e-2
10000 | 2.20e1 | 2.32e-1 | 3.55e-1
Table 5: Comparison between the computation time for the optimization approach
with quadrature $(\texttt{opt}_{\texttt{Q}})$ and analytic
$(\texttt{opt}_{\texttt{A}})$ integrals and the batch evaluation time for the
small ($\texttt{NN}^{0\times 15}_{2}$) and large ($\texttt{NN}^{4\times
45}_{2}$) networks in the $N=2$ case. The time (sec) for evaluating 100, 200,
500, 1000, and 10000 moments are reported.

###  4.4 Plane source problem

In this section, the data-driven closures are applied to solve a benchmark
problem: the plane source problem. This problem models the propagation of
neutral particles in a purely scattering medium from initial time $t_{0}=0$ to
final time $t_{\textup{final}}$. The kinetic model is given in (38), and we
consider $t_{\textup{final}}=1$ with a constant scattering crosssection
$\sigma_{\text{s}}=1$. The initial particle distribution is given by

$f(t_{0}=0,x,\mu)=0.5\delta(x)+f_{\textup{floor}},$ (50)

where $\delta$ is the Dirac-delta function and $f_{\textup{floor}}=10^{-8}$ is
the floor value for the distribution for safeguarding. To minimize the effect
of the boundaries, the computation domain in $(x,\mu)$ is chosen to be
$[x_{L},x_{R}]\times[-1,1]$, where $x_{L}=-(t_{\textup{final}}+0.1)$ and
$x_{R}=t_{\textup{final}}+0.1$ with boundary conditions

$f(t,x_{L},\mu)=f_{\textup{floor}}\quad\mbox{and}\quad
f(t,x_{R},\mu)=f_{\textup{floor}}.$ (51)

Here we solve the corresponding moment equation (39) with moment order $N=1$
and $2$ using Maxwell-Boltzmann entropy-based closure (MN), the PN closure
[43, 12, 36], and two data-driven closures — the spline and neural network
closures. The plane source problem is known to be difficult for the closures
since the discontinuous initial distribution leads to moments close to
realizable boundary during the evolution. We use it to demonstrate that the
data-driven closures can serve as efficient approximations to the expensive,
optimization-based MN closure.

The numerical scheme used to solve the moment system is the realizability-
preserving kinetic scheme developed in [8] based on earlier work in [15, 22,
40, 41]. This scheme discretizes (39) using a second-order finite volume
method in space and second-order strong-stability-preserving Runge-Kutta (SSP-
RK2) method in time. In the following tests, we solve (39) on a uniform
spatial grid of $100$ cells, and the boundary condition is implemented using
two ghost cells on each side of the boundary. The time step is given by
$\Delta t=0.95(\frac{2}{2+\theta})\Delta x$, where $\theta$ is a parameter in
the minmod limiter used in the numerical scheme. We use $\theta=2$ throughout
the tests.

##### Closure comparison for $N=1$.

Figure 6 shows the solutions to the plane source problem at time given by the
tested closures with moment order $N=1$. The M1 closure is implemented using
both the $\texttt{opt}_{\texttt{Q}}$ and $\texttt{opt}_{\texttt{A}}$
approaches discussed in Section 4.3, where the a 10-point Gauss-Legendre
quadrature rule is used in $\texttt{opt}_{\texttt{Q}}$. We show the solutions
from these two approaches in the same plot as they are indistinguishable. The
solutions of spline closures are shown in Figure 6(c), where we use the
coarser ($\texttt{S}^{30}$) and finer ($\texttt{S}^{130}$) spline
approximations reported in Table 1. The results of the network closures are
plotted in Figure 6(d), where the smaller ($\texttt{NN}^{1\times 15}_{1}$) and
larger ($\texttt{NN}^{5\times 30}_{1}$) networks in Table 1 are tested. To
preserve the symmetry of the solution, the symmetric network approximation
given in (45) is used here. From these results, we confirm that, with
sufficiently accurate approximations, the data-driven closures can provide
quality solutions similar to the ones from the more expensive M1 closures.

(a) M1, $\texttt{opt}_{\texttt{Q}}$/$\texttt{opt}_{\texttt{A}}$

(c) Spline closures

(b) P1

(d) Symmetric networks closures

Figure 6: Solutions to the plane source problem at $t_{\text{final}}=1$ using
various moment closures with $N=1$. Figures 6(c) and 6(d) show that the
solutions from the larger spline ($\texttt{S}^{130}$) and network
($\texttt{sym-NN}^{5\times 30}_{1}$) are nearly identical to the reference M1
solution.

Table 6 reports the computation time for the tested closures and the relative
error of the solutions from data-driven closures at the final time
$t_{\textup{final}}=1$. Here we define the error as

$e_{\mathbf{u}}:=\frac{\|\mathbf{u}_{\textup{a}}(t_{\textup{final}},\cdot)-\mathbf{u}_{\text{ref}}\|_{L^{2}([x_{L},x_{R}])}}{\|\mathbf{u}_{\text{ref}}\|_{L^{2}([x_{L},x_{R}])}},\quad\text{with
norm
}\|\mathbf{w}\|_{L^{2}([x_{L},x_{R}])}:=\big{(}\int_{x_{L}}^{x_{R}}\|\mathbf{w}\|_{2}^{2}\,dx\big{)}^{\frac{1}{2}}\>,$
(52)

where the reference solution $\mathbf{u}_{\text{ref}}$ is chosen to be the MN
($\texttt{opt}_{\texttt{Q}}$) solution at final time, since the data-driven
closures aim to approximate the MN closure. The computation time for the P1
closure serves as a baseline in the following comparison, since the cost of
computing the P1 closure is minimal. We observe from Table 6 that, for both
the spline and network closures, increasing the number of parameters improves
the solution accuracy with minimal effect in the computation time. When using
the finer spline ($\texttt{S}^{130}$) and larger network
($\texttt{NN}^{5\times 30}_{1}$) closures, the errors are of order $10^{-3}$.
The fast evaluation of the spline closures results in a 183x speedup comparing
to the M1 closure. Despite the slower evaluation of the symmetric network
closures due to the larger number of parameters and the overhead for calling
the model.predict() Keras function, the symmetric network closures is still
2.6x faster than the M1 closure.

Closure | M1 $\texttt{opt}_{\texttt{Q}}$ | M1 $\texttt{opt}_{\texttt{A}}$ | P1 | $\texttt{sym-NN}^{1\times 15}_{1}$ | $\texttt{sym-NN}^{5\times 30}_{1}$ | $\texttt{S}^{30}$ | $\texttt{S}^{130}$
---|---|---|---|---|---|---|---
Time (sec) | 1.66e1 | 1.62e1 | 2.15e-2 | 6.09e0 | 6.24e0 | 8.71e-2 | 8.86e-2
$e_{\mathbf{u}}$ | — | — | — | 6.33e-2 | 2.85e-3 | 1.07e-1 | 1.49e-3
Table 6: Computation time (sec) and relative errors in the plane source test
using various closures with $N=1$. The error $e_{\mathbf{u}}$ (see (52) for
definition) measures the relative $L^{2}$ error in the final time solutions
from data-driven closures with respect to the one from the M1
$\texttt{opt}_{\texttt{Q}}$ closure.

##### Closure comparison for $N=2$.

We then test the closures on the plane source problem for moment order $N=2$.
The solutions for the tested closures are plotted in Figure 7 and the
computation time and relative errors are reported in Table 7. Here we only
compare the M2 ($\texttt{opt}_{\texttt{Q}}$), P2, and the symmetric network
closures, since the spline approximation in Section 3.2.1 is limited to one
dimension. Figure 7 confirms that the solutions from symmetric network
closures successfully capture the qualitative behavior of the M2 solution.
Using the larger network ($\texttt{sym-NN}^{4\times 45}_{2}$), the network
closure solution is nearly identical to the reference M2 solution.

(a) M2, $\texttt{opt}_{\texttt{Q}}$ (b) P2

(c) Symmetric network closures

Figure 7: Solutions to the plane source problem at $t_{\text{final}}=1$ using
various moment closures with $N=2$. Figure 7(c) shows that the solution from
the larger network ($\texttt{sym-NN}^{4\times 45}_{2}$) is nearly identical to
the reference M2 solution.

The computation time in Table 7 shows a 3.2x speedup when replacing the M2
closure by the network closure, which is slightly larger than the 2.6x speedup
observed in Table 6 for the case $N=1$. We expect the speedup for using the
network closure to increase as the moment order $N$ goes higher, because the
optimization problem (23) becomes much harder to solve for large $N$ while the
network evaluation time does not appear to have a strong dependency on the
value of $N$.

Closure | M2 $\texttt{opt}_{\texttt{Q}}$ | P2 | $\texttt{sym-NN}^{1\times 15}_{2}$ | $\texttt{sym-NN}^{4\times 45}_{2}$
---|---|---|---|---
Time (sec) | 1.90e1 | 2.06e-2 | 5.78e0 | 6.01e0
$e_{\mathbf{u}}$ | — | — | 9.66e-2 | 5.49e-3
Table 7: Computation time (sec) and relative errors in the plane source test
using various closures with $N=2$. The error $e_{\mathbf{u}}$ (see (52) for
definition) measures the relative $L^{2}$ error in the final time solutions
from data-driven closures with respect to the one from the M2
$\texttt{opt}_{\texttt{Q}}$ closure.

##### Closure comparison on a refined spatial grid for $N=2$.

The slight differences between the computation time for the small and large
networks in Tables 6 and 7 suggest that the dominant cost in network
evaluations is the overhead for calling Keras function model.predict(), as
discussed in Section 4.3. To minimize the effect of the overhead, we repeat
the comparison test of the closures but on a refined spatial grid with the
number of spatial cells increased from 100 to 1000. In this test, each call of
model.predict() evaluates 1000, rather than 100, samples of moments, thus it
could give a more accurate comparison between the optimization cost and the
actual network evaluation cost. To avoid numerical difficulties, we also
replace the Dirac-delta function initial condition in (50) with a smoother
isotropic initial condition

$f(t_{0}=0,x,\mu)=0.5\chi_{[-\frac{1}{2},\frac{1}{2}]}(x)\cos^{2}(\pi
x)+f_{\textup{floor}},$ (53)

where $\chi_{[a,b]}(x)=1$ if $x\in[a,b]$ and $\chi_{[a,b]}(x)=0$ otherwise.
Table 8 shows the relative error and computation time of the compared closures
using the refined spatial grid with 1000 cells and the smooth initial
condition (53). In Table 8, a 9.7x speedup for using the large network
$\texttt{sym-NN}^{4\times 45}_{2}$ over the M2 closure is reported, which is
higher than the 3.2x speedup observed in Table 7. This result suggests that
the advantage in computation time of the network closures can be further
improved via an efficient implementation of the network evaluation function
with minimal overhead.

Closure | M2 $\texttt{opt}_{\texttt{Q}}$ | P2 | $\texttt{sym-NN}^{1\times 15}_{2}$ | $\texttt{sym-NN}^{4\times 45}_{2}$
---|---|---|---|---
Time (sec) | 1.06e3 | 2.96e-1 | 9.03e1 | 1.09e2
$e_{\mathbf{u}}$ | — | — | 6.43e-2 | 1.92e-3
Table 8: Computation time (sec) and relative errors in the test with 1000
spatial cells and smooth initial condition (53) using various closures with
$N=2$. The error $e_{\mathbf{u}}$ (see (52) for definition) measures the
relative $L^{2}$ error in solutions at final time $t_{\textup{final}}=1$ from
data-driven closures with respect to the one from the M2
$\texttt{opt}_{\texttt{Q}}$ closure.

##### Correlation between network moment error
$e_{\mathbf{w}}^{{\textup{test}}}$ and solution error $e_{\mathbf{u}}$.

In Figure 8, we explore the correlation between the test moment error
$e_{\mathbf{w}}^{{\textup{test}}}$ of the network approximations and the
relative $L^{2}$ error $e_{\mathbf{u}}$ in the plane source solution from the
network closures when $N=2$. The data shown in Figure 8 are collected from the
small ($\texttt{sym-NN}^{1\times 15}_{2}$) and large $(\texttt{sym-
NN}^{4\times 45}_{2})$ networks considered in Table 7 as well as several other
networks of different sizes. These results show that the test moment error
$e_{\mathbf{w}}^{{\textup{test}}}$ is indeed a good proxy of the error
$e_{\mathbf{u}}$ in the final solution, which justifies the use of the moment
reconstruction error in the loss function (46) when training the neural
network approximations.

Figure 8: Test moment errors $e_{\mathbf{w}}^{{\textup{test}}}$ versus $L^{2}$
error of the plane source solution $e_{\mathbf{u}}$ from the various symmetric
neural network closures for $N=2$. Closures of various depths with width $W=$
15, 30, and 45 are denoted as blue crosses, circles, and triangles.

##  5 Summary and discussion

We have developed data-driven closures that approximate the map between the
moments and the entropy of the moment system for kinetic equations. The
approximations are constructed to be convex and twice differentiable to ensure
entropy dissipation and hyperbolicity in the moment system. We have shown that
convexity of the entropy approximation is preserved in the extension from a
normalized realizable set to the full realizable set of moments, which allows
for the strategy of building a convex entropy approximation on the bounded
normalized realizable set and then extending the approximation to the
unbounded full realizable set. We have focused on the Maxwell-Boltzmann
entropy and considered two approaches — rational cubic splines and neural
networks — to approximate the Maxwell-Boltzmann entropy. The convexity-
preserving cubic spline approach gives fast, accurate, and data-efficient
approximations that are convex by construction, but this approach is limited
to the case that the normalized realizable set is an interval, which generally
implies the restriction of moment order $N=1$ in the simplified slab geometry.
On the other hand, the neural network approach does not have the restrictions
on the moment order or geometry, but it requires a substantial amount of data
in the training process to build an accurate approximation with convexity
enforced on the data points. We have tested the data-driven closures on the
plane source benchmark problem in slab geometry and compared them to the
standard, optimization-based MN closures. Numerical results indicate that the
data-driven closures are able to provide solutions that capture the
qualitative behavior of the MN solutions while requiring much less computation
time.

We plan to extend this work in various directions: (i) adopting neural network
architectures that guarantee convexity of the network approximation, e.g., [9]
and investigating tools or implementations that allow for faster network
evaluation; (ii) combining the data-driven closures with the entropy
regularization scheme developed in [6] to address the issue that multipliers
goes unbounded as moments approaching realizable boundary; (iii) investigating
the extension of proposed data-driven closures to more complex problems that
are in higher dimensions with higher moment order and potentially from other
applications, such as charge transport or plasma simulations; (iv) exploring
other data-driven entropy approximations, such as high-dimensional splines,
that can potentially be more data-efficient with faster evaluation; and (v)
developing strategies to circumvent the degeneracy issue that the entropy
minimization problem has no solution for some realizable moments (see [6] and
references therein for discussion), which is not addressed in either the
regularization scheme in [6] or the data-driven closures considered in this
paper.

## References

  * [1] Rafail V Abramov.  An improved algorithm for the multidimensional moment-constrained maximum entropy problem.  Journal of Computational Physics, 226(1):621–644, 2007.
  * [2] Rafail V Abramov.  The multidimensional moment-constrained maximum entropy problem: A bfgs algorithm with constraint scaling.  Journal of Computational Physics, 228(1):96–108, 2009.
  * [3] Rafail V Abramov et al.  The multidimensional maximum entropy moment problem: A review of numerical methods.  Communications in Mathematical Sciences, 8(2):377–392, 2010.
  * [4] Graham Alldredge and Florian Schneider.  A realizability-preserving discontinuous galerkin scheme for entropy-based moment closures for linear kinetic equations in one space dimension.  Journal of Computational Physics, 295:665–684, 2015.
  * [5] Graham W Alldredge, Martin Frank, and Jan Giesselmann.  On the convergence of the regularized entropy-based moment method for kinetic equations.  arXiv preprint arXiv:2105.10274, 2021.
  * [6] Graham W Alldredge, Martin Frank, and Cory D Hauck.  A regularized entropy-based moment method for kinetic equations.  SIAM Journal on Applied Mathematics, 79(5):1627–1653, 2019.
  * [7] Graham W Alldredge, Cory D Hauck, Dianne P O’Leary, and André L Tits.  Adaptive change of basis in entropy-based moment closures for linear kinetic equations.  Journal of Computational Physics, 258:489–508, 2014.
  * [8] Graham W Alldredge, Cory D Hauck, and Andre L Tits.  High-order entropy-based closures for linear transport in slab geometry ii: A computational study of the optimization problem.  SIAM Journal on Scientific Computing, 34(4):B361–B391, 2012.
  * [9] Brandon Amos, Lei Xu, and J Zico Kolter.  Input convex neural networks.  In International Conference on Machine Learning, pages 146–155. PMLR, 2017.
  * [10] Leo Bois, Emmanuel Franck, Laurent Navoret, and Vincent Vigon.  A neural network closure for the euler-poisson system based on kinetic simulations.  arXiv preprint arXiv:2011.06242, 2020.
  * [11] Russel E Caflisch and C David Levermore.  Equilibrium for radiation in a homogeneous plasma.  The Physics of fluids, 29(3):748–752, 1986.
  * [12] K. Case and P. Zweifel.  Linear Transport Theory.  Addison-Wesley, Reading, MA, 1967.
  * [13] Carlo Cercignani.  The Boltzmann equation.  In The Boltzmann Equation and Its Applications, pages 40–103. Springer, 1988.
  * [14] Ran Chu, Eirik Endeve, Cory D Hauck, and Anthony Mezzacappa.  Realizability-preserving dg-imex method for the two-moment model of fermion transport.  Journal of Computational Physics, 389:62–93, 2019.
  * [15] S. M. Deshpande.  Kinetic theory based new upwind methods for inviscid compressible flows.  In American Institute of Aeronautics and Astronautics, New York, 1986.  Paper 86-0275.
  * [16] Bruno Dubroca and Jean-Luc Feugeas.  Etude théorique et numérique d’une hiérarchie de modèles aux moments pour le transfert radiatif.  Comptes Rendus de l’Académie des Sciences-Series I-Mathematics, 329(10):915–920, 1999.
  * [17] C. K. Garrett and C. D. Hauck.  A comparison of moment closures for linear kinetic transport equations: the line source benchmark.  Transport Theor. Stat., 42:203 – 235, 2013.
  * [18] C Kristopher Garrett, Cory Hauck, and Judith Hill.  Optimization and large scale computation of an entropy-based moment closure.  Journal of Computational Physics, 302:573–590, 2015.
  * [19] J Gregory.  Shape preserving spline interpolation.  COMP. AIDED DESIGN., 18(1):53–57, 1986.
  * [20] Jean-Luc Guermond, Bojan Popov, and Ignacio Tomas.  Invariant domain preserving discretization-independent schemes and convex limiting for hyperbolic systems.  Computer Methods in Applied Mechanics and Engineering, 347:143–175, 2019.
  * [21] Jiequn Han, Chao Ma, Zheng Ma, and E Weinan.  Uniformly accurate machine learning-based hydrodynamic models for kinetic equations.  Proceedings of the National Academy of Sciences, 116(44):21983–21991, 2019.
  * [22] A. Harten, P. D. Lax, and Van Leer.  On upstream differencing and Godunov-type schemes for hyperbolic conservation laws.  SIAM Rev., 25:35–61, 1983.
  * [23] Cory D Hauck.  High-order entropy-based closures for linear transport in slab geometry.  Communications in Mathematical Sciences, 9(1):187–205, 2011.
  * [24] Cory D Hauck, C David Levermore, and André L Tits.  Convex duality and entropy-based moment closures: Characterizing degenerate densities.  SIAM Journal on Control and Optimization, 47(4):1977–2015, 2008\.
  * [25] Richard D Hazeltine.  The framework of plasma physics.  CRC Press, 2018.
  * [26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In Proceedings of the IEEE international conference on computer vision, pages 1026–1034, 2015.
  * [27] Juntao Huang, Yingda Cheng, Andrew J Christlieb, and Luke F Roberts.  Machine learning moment closure models for the radiative transfer equation i: directly learning a gradient based closure.  arXiv preprint arXiv:2105.05690, 2021.
  * [28] Juntao Huang, Yingda Cheng, Andrew J Christlieb, Luke F Roberts, and Wen-An Yong.  Machine learning moment closure models for the radiative transfer equation ii: enforcing global hyperbolicity in gradient based closures.  arXiv preprint arXiv:2105.14410, 2021.
  * [29] Michael Junk.  Maximum entropy for reduced moment problems.  Mathematical Models and Methods in Applied Sciences, 10(07):1001–1025, 2000.
  * [30] David S Kershaw.  Flux limiting natures own way–a new method for numerical solution of the transport equation.  Technical report, Lawrence Livermore National Lab., CA (United States), 1976.
  * [31] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [32] Siddharth Krishna Kumar.  On weight initialization in deep neural networks.  arXiv preprint arXiv:1704.08863, 2017.
  * [33] M. P. Laiu, C. D. Hauck, R. G. McClarren, D. P. O’Leary, and A. L. Tits.  Positive filtered PN moment closures for linear transport equations.  SIAM J. Numer. Anal., 54(6):3214–3238, 2016.
  * [34] C. D. Levermore.  Moment closure hierarchies for kinetic theories.  Journal of Statistical Physics, 83:1021–1065, 1996.
  * [35] E. E. Lewis and W. F. Miller.  Computational methods of neutron transport.  Wiley (New York), 1993.
  * [36] E. E. Lewis and Jr. W. F. Miller.  Computational Methods in Neutron Transport.  John Wiley and Sons, New York, 1984.
  * [37] Peter A Markowich, Christian A Ringhofer, and Christian Schmeiser.  Semiconductor equations.  Springer Science & Business Media, 2012.
  * [38] James McDonald and Manuel Torrilhon.  Affordable robust moment closures for cfd based on the maximum-entropy hierarchy.  Journal of Computational Physics, 251:500–523, 2013.
  * [39] Edgar Olbrant, Cory D Hauck, and Martin Frank.  A realizability-preserving discontinuous galerkin method for the M1 model of radiative transfer.  Journal of Computational Physics, 231(17):5612–5639, 2012.
  * [40] B. Perthame.  Boltzmann type schemes for gas dynamics and the entropy property.  SIAM J. on Numer. Anal., 27(6):1405–1421, 1990.
  * [41] B. Perthame.  Second-order Boltzmann schemes for compressible euler equations in one and two space dimensions.  SIAM J. Numer. Anal., 29(1):1–19, 1992.
  * [42] Teddy Pichard, Graham W Alldredge, Stéphane Brull, Bruno Dubroca, and M Frank.  An Approximation of the $M_{2}$ Closure: Application to Radiotherapy Dose Simulation.  Journal of Scientific Computing, 71(1):71–108, 2017.
  * [43] G. C. Pomraning.  Radiation Hydrodynamics.  Pergamon Press, New York, 1973.
  * [44] G. C. Pomraning.  The equations of radiation hydrodynamics.  Oxford: Pergamon, 1973.
  * [45] Mohsen Sadr, Manuel Torrilhon, and M Hossein Gorji.  Gaussian process regression for maximum entropy distribution.  Journal of Computational Physics, 418:109644, 2020.
  * [46] Joachim AR Sarr and Clinton PT Groth.  A second-order maximum-entropy inspired interpolative closure for radiative heat transfer in gray participating media.  Journal of Quantitative Spectroscopy and Radiative Transfer, 255:107238, 2020.
  * [47] Jacques Schneider.  Entropic approximation in kinetic theory.  ESAIM: Mathematical Modelling and Numerical Analysis, 38(3):541–561, 2004.
  * [48] X. Zhang and C.-W. Shu.  On maximum-principle-satisfying high order schemes for scalar conservation laws.  Journal of Computational Physics, 229:3091–3120, 2010.