# Sequential Deep Learning Operator Network (S-DeepONet) for Time-Dependent
Loads

Jaewan Park1,∗∗ Shashank Kushwaha1,∗∗ Junyan He1,∗∗111∗∗The three authors
contributed equally to this work Seid Koric1,2<EMAIL_ADDRESS>Diab
Abueidda2,3 Iwona Jasiuk1 1 Department of Mechanical Science and Engineering,
University of Illinois at Urbana-Champaign, Urbana, IL, USA
2 National Center for Supercomputing Applications, University of Illinois at
Urbana-Champaign, Urbana, IL, USA
3 Civil and Urban Engineering Department, New York University Abu Dhabi, Abu
Dhabi, UAE

###### Abstract

Deep Operator Network (DeepONet), a recently introduced deep learning operator
network, approximates linear and nonlinear solution operators by taking
parametric functions (infinite-dimensional objects) as inputs and mapping them
to solution functions in contrast to classical neural networks (NNs) that need
re-training for every new set of parametric inputs. In this work, we have
extended the classical formulation of DeepONets by introducing recurrent
neural networks (RNNs) in its branch in so-called sequential DeepONets
(S-DeepONets) thus allowing accurate solution predictions in the entire domain
for parametric and time-dependent loading histories. We have demonstrated this
novel formulation’s generality and exceptional accuracy with thermal and
mechanical random loading histories applied to highly nonlinear thermal
solidification and plastic deformation use cases. We show that once S-DeepONet
is properly trained, it can accurately predict the final solutions in the
entire domain and is several orders of magnitude more computationally
efficient than the finite element method for arbitrary loading histories
without additional training.

###### keywords:

Machine/Deep Learning , Recurrent Neural Netowrk (RNN) , Deep Operator Network
(DeepONet) , Heat Transfer , Plastic Deformation

††journal: Engineering Applications of Artificial Intelligence

*[inlinelist,1]label=(),

##  1 Introduction

Recent technological advances in high-performance computing hardware and
machine learning (ML) methods have given rise to a wide range of applications
in fields like autonomous driving, image and speech recognition,
bioinformatics, medical diagnosis, document categorization, and others.
Physics-based modeling has shown much interest in applications of Deep
Learning with Artificial Neural Networks (ANN), a branch of machine learning
motivated by the brain’s biological structure and operation. Without the need
for expensive computing power or modeling software, a well-trained surrogate
deep learning model can almost immediately produce (infer) outcomes that are
comparable with traditional modeling techniques. Many data-driven surrogate
deep learning models have been devised and trained lately to quickly solve
problems in additive manufacturing [1, 2], topologically optimized materials
and structures [3, 4, 5], random heterogeneous materials [6], composite
materials [7], nonlinear material responses [3, 8, 9], as well as a variety of
other applications. Besides data-driven models, collocation point-based
physics-informed neural networks (PINN) were created by Raissi et al. [10] and
Abueidda et al. [11] capable of solving partial differential equations
governing deformation and stress generation in solids or other physics without
the aid of finite elements or other conventional numerical techniques, other
than for validation. Similar to this, Nguyen-Thanh et al. [12], Samaniego et
al. [13], Abueidda et al. [14, 15], Fuhg et al. [16] and He et al. [17, 18]
devised a deep energy method (DEM), which makes use of potential energy to
resolve nonlinear material responses. Besides forward problems, Haghighat et
al. [19] used PINNs for inverse problems in solid mechanics. Cai et al. [20]
even used a combination of measured data and a physics-informed deep-learning
method to obtain a solution for an ill-posed thermal fluid flow that was
previously thought to be unsolvable.

Nevertheless, most of these methods require retraining or transfer learning if
input parameters like loads, boundary conditions, and material properties, or
geometry change. The same is true of traditional numerical methods such as
finite elements (FE) in that each new input parameter value calls for a new
independent simulation. To address this problem, the universal approximation
theorem for operators [21] inspired Lu et al.’s Deep Operator Network [22],
often known as DeepONet, an innovative operator learning architecture. For a
few linear and nonlinear PDEs in that landmark work, DeepONet successfully
mapped between unknown parametric functions and solution spaces in addition to
learning explicit operators like integrals. This gave rise to a powerful new
method for solving stochastic and parametric PDEs. In the so-called physics-
informed DeepONet, Wang et al. [23] improved the DeepONet formulation by
including information from the governing PDE. They found improved prediction
accuracy and data handling efficiency but at the expense of a higher computing
cost for training. Recently, DeepONet has been used in heat conduction with a
spatially variable heat source by Koric and Abueidda [24], fracture mechanics
by Goswami et al. [25], and multiscale modeling using elastic and hyperelastic
materials by Yin et al. [26], and elastic-plastic stress field prediction on
topologically optimized geometries [27].

Many nonlinear thermal, mechanical, and multiphysics analyses have time-
dependent loads often coupled with highly nonlinear thermo-mechanical
properties, including phase transformation and/or materially nonlinear path-
dependent constitutive models, such as in plastic deformation. Frequently,
only the final stress or temperature field solution, obtained through
incremental nonlinear finite element analysis, is of interest for analysis,
designs, and optimizations. Such problems can be particularly computationally
challenging and time-consuming for sensitivity analysis, uncertainty
quantification, topological optimization, and similar iterative procedures
where thousands, or even millions of forward evaluations, are needed to be
done by classical nonlinear analysis to achieve statistical convergence. Koric
et al. [28] used DeepONet for full-field stress prediction in plastic
deformation with time-dependent loads and variable material properties using
classical forward fully connected neural network (FNN) architecture in both
branches of DeepONet. While the predictions were reasonably accurate in that
work, the question was if advanced sequential learning methods, such as
recurrent neural networks of the LSTM (Long short-term memory) and GRU (Gated
recurrent unit) type, would provide universally and significantly more
accurate full-field predictions by encoding random loading histories from
various applications in the branch network of DeepONet. In this novel work, we
have addressed that question.

This paper is organized as follows: Section 2 introduces three neural networks
architectures and provides detail on the data generation method. Section 3
presents and discusses the performance of the three models. Section 4
summarizes the outcomes and highlights future works.

##  2 Methods

###  2.1 Neural network models

This work explored three different NN architectures to predict the final
temperature or von Mises stress distribution. All three NNs were implemented
in the DeepXDE framework [29] with a TensorFlow backend [30].

####  2.1.1 FNN DeepONet

A DeepONet model with FNN in both the branch and trunk networks is used as the
performance baseline. In an infinite-dimension functional input space $M$,
$m\in M$ represents an input function with history-dependent load magnitudes
defined on $n$ control points (or input sensors) and refers to an unknown
temperature or stress field solution in the functional solution space $S$. We
consider that for every input $m\in M$ there is a solution $s=s(m)\in S$ for
the temperature or stress field distribution from the Section 2.2.1 and
Section 2.2.2, which are also subject to respective boundary conditions (BCs).
Consequently, the mapping solution operator of a DeepONet $G:M\to S$ can be
defined as:

$G(m)=s(m).$ (1)

For a collection of $N$ points $\bm{X}$ on a domain, each denoted by its
coordinates $(x_{i},y_{i})$, the DeepONet considers both the load magnitudes
$m$ in its branch and positions $\bm{X}$ in its trunk and predicts the
solution operator $\hat{G}(m)(\bm{X})$ by combining intermediate encoded
outputs $b_{i}$ (from branch) and $t_{i}$ (from trunk) in a dot product
enhanced by bias $\beta$, as shown in a schematic of the FNN-based DeepONet
model in Fig. 1.

Figure 1: Schematic of the FNN-based DeepONet used in this work. $m$, $x$,
$y$, $b_{i}$, $t_{i}$, $HD$, $\beta$ and $\hat{G}$ denote the load magnitude,
X coordinate, Y coordinate, branch output, trunk output, hidden dimensions,
the bias vector and the approximated solution operator.

In a larger sense, it is possible to think of $\hat{G}(m)(\bm{X})$ as a
function of $\bm{X}$ conditioning on input $m$, and DeepONet is more general
and capable than other neural networks. In both example cases, the time-
dependent input load magnitude curves were discretized into $n=101$ time
steps, and were fed into the branch FNN, which has seven layers (including
input and output) with the following number of neurons:
$[101,100,100,100,100,100,HD]$. The trunk FNN network received the nodal
coordinates of all nodes in the mesh, and has the following number of neurons:
$[2,100,100,100,100,100,HD]$. The output hidden dimension ($HD$) of both
networks was set to 100, and the dot product operation combines the
intermediate outputs of the two networks to produce the final output solution
field. This network was applied to solve both problems introduced in Section
2.2. The total number of trainable parameters in the network is 111500. A
larger network with similar number of parameters as in Section 2.1.2 was
tested but showed similar performance as this smaller network, so the smaller
network was used in the result sections. The model was trained for 351000
epochs with a batch size of 64. The Adam optimizer [31] was used and the
scaled mean squared error (MSE) was used as the loss function.

####  2.1.2 S-DeepONet

The FNN-based DeepONet introduced in Section 2.1.1 uses a FNN to encode the
time series signal of the input magnitude. However, time series signals are
typically treated by RNNs [32, 9] due to their ability to capture causal
relationship in the input signal, which significantly affects the final
outcome for path-dependent problems like plasticity. However, in a FNN, the
causality of the input data is not retained. The ability of RNNs to capture
causality prompts us to apply RNNs in the branch to encode the time series
signal. RNNs are referred to as recurrent because they carry out the same
action for each element of a sequence, with the result depending on the
results of earlier time-step calculations. RNN predicts the corresponding
matching outputs for a given input sequence at each time step while minimizing
the loss function between predictions and actual targets. Both encoder and
decoder networks in Fig. 2 show an RNN cell symbolically being unfolded in
time for an entire sequence.

Figure 2: Schematic of the proposed S-DeepONet. The black RNN blocks returns a
sequence (2D outputs), while the green RNN block does not return a sequence
(1D output). $m$, $x$, $y$, $b_{i}$, $t_{i}$, $HD$, $\beta$ and $\hat{G}$
denote the load magnitude, X coordinate, Y coordinate, branch output, trunk
output, hidden dimensions, the bias vector and the approximated solution
operator.

There is a hidden state at every time step, also known as the ”memory” of the
recurrent network, since it captures information about what happened in all
the previous time steps. It is calculated from the previous hidden state and
the input at the current step. Very deep computational graphs that
continuously (recurrently) apply the same operation at each time step of a
lengthy sequence are a common component of recurrent neural networks. It is
difficult to train long sequence data with a simple RNN, which may result in
vanishing or exploding gradient difficulties during backpropagation. Long
short-term memory (LSTM) [33] and gated recurrent unit (GRU) [34] are
solutions to the vanishing gradient issues of a simple RNN. In the LSTM and
GRU advanced recurrent neural networks, hidden state (memory) cells are
intended to dynamically ”forget” some outdated and pointless information via
particular gated units that regulate the information flow inside a memory
cell, preventing the multiplication of a lengthy series of numbers during
temporal backpropagation.

In this work, a branch network based on two different RNN layers, the LSTM and
GRU, is proposed. Since the branch network of this novel DeepONet accepts
sequence data, we named this architecture the S-DeepONet for short. An
encoder-decoder structure is used in the branch, which was shown in many
previous sequence-to-sequence prediction tasks to significantly improve the
prediction accuracy [35, 34, 36]. In the context of DeepONet, the sequence-to-
sequence mapping maps the loading history to a latent representation (the
branch output), which is then merged with the latent representation of the
domain geometry (the trunk output) via a dot product [29] to produce the final
output prediction. Essentially, the proposed S-DeepONet aims to combine the
powerful temporal encoding capability of the RNN structure in the branch with
the spatial encoding capability of the DeepONet architecture. A schematic of
the proposed S-DeepONet is shown in Fig. 2. The hidden dimension ($HD$) for
the S-DeepONet is 101. The RNN blocks shown in Fig. 2 can be LSTM cells or GRU
cells, leading to two different S-DeepONet models. The number of trainable
parameters in the LSTM- and GRU-based S-DeepONet are 1038294 and 990358,
respectively. Identical loss function and optimizer as mentioned in Section
2.1.1 were used to train the proposed two S-DeepONet models for a total of
351000 epochs.

###  2.2 Data generation

In this work, we demonstrate the application of the proposed S-DeepONets in
transient heat transfer and structural deformation problems.

####  2.2.1 Transient heat transfer

In the first example, a nonlinear, transient heat transfer problem is studied,
which is representative of a solidifying shell of low-carbon steel in a
continuous caster [9]. Fig. 3 provides a schematic to illustrate the process.

Figure 3: 3D schematic of molten steel solidifying in a caster mold.

To simplify the simulation, a 2D slice cross-section of the caster is modeled,
which has a length of 30 mm and a width of 0.1 mm. As the domain moves down
the mold in a Lagrangian frame of reference, it is subject to a prescribed
time-dependent heat flux extracting heat on one end, and all other surfaces
are insulated. A schematic of the 2D problem domain is shown in Fig. 4.

Figure 4: Schematic of the 2D domain used in the heat transfer problem.

The governing equation and initial and boundary conditions of the transient
heat transfer are given by:

$\displaystyle\rho H(T)\frac{\partial T}{\partial
t}=\nabla\cdot\left[k(T)\nabla T\right],$ (2) $\displaystyle-k(T)\nabla
T=\bm{q}(t),\;\;\forall\bm{x}\in\partial\Omega_{q},$ $\displaystyle
T(x,0)=T_{0},$

where $t$ is time, $T$ is temperature, $\rho$ is mass density, and $\bm{q}(t)$
is the time-dependent heat flux. $T_{0}$ = 1540 oC is the uniform initial
temperature. $H(T)$ and $k(T)$ denote the temperature-dependent specific
enthalpy and isotropic thermal conductivity, respectively. It is highlighted
that the specific enthalpy used in this work includes the latent heat effect
during phase transformations, such as in solidification and transition from
$\delta$-ferrite to austenite, and brings significant nonlinearity to the
system [37]. The material is a low carbon steel with a density of 7400 kg/m3;
other relevant material properties are shown in Table 1 and Table 2.

Table 1: Temperature-dependent material properties Temperature [oC] |  Conductivity [W/(m$\cdot$K)] |  Specific heat [J/(kg$\cdot$K)]
---|---|---
800 | 28.934 | 695.44
1480.04 | 34.188 | 695.44
1519.73 | 39.000 | 824.61
Table 2: Latent heat properties | Latent heat [J] |  Liquidus temperature [oC] |  Solidus temperature [oC]
---|---|---|---
Value | 245100.0 | 1524.59 | 1475.43

The problem domain was discretized into 300 four-node bi-linear heat transfer
elements (DC2D4) with an element size of 0.1 mm, and the thermal equation Eq.
(2) is solved for a total of 17s that the slice spends in the caster traveling
down the mold by using implicit time integration in Abaqus/Standard [37]. The
temperature distribution at the end of the load step was extracted as the
ground truth for NN training. To define a complex, time-dependent boundary
heat flux, the sampling approach by Abueidda et al. [9] was used, where the
time-dependent boundary heat flux is defined by six control points. The first
and last control points (cp) correspond to $t=0$ and $t=17s$, respectively.
The time values ($t_{cp}$) for the four remaining control points were randomly
sampled from a uniform distribution in range $(0,17)s$. Based on experimental
measurements, the flux value $q_{cp}$ generally has a decaying profile and can
be approximated as:

$q_{cp}=A(t_{cp}+1)^{-B}+C,$ (3)

where $A\in[3,8]$, $B\in[0.3,0.7]$, and $C\in[-0.5,0.5]$ were randomly chosen
variables from their respective ranges. $C$ can be considered as a random
noise added to the flux magnitude to emulate additional fluctuations and
nonlinearities observed in practice in the actual flux profile due to changes
in contact and interfacial heat transfer between mold and steel [38, 39].
After all control points and the flux values are defined, a radial basis
interpolation with Gaussian function is used to interpolate the flux. A total
of 4000 cases were generated with distinct heat flux histories, and a typical
example of the time-dependent heat flux is shown in Fig. 5a.

(a)

(b)

Figure 5: Typical time dependent load magnitudes: LABEL:sub@heat_flux Boundary
heat flux in the transient heat transfer problem. LABEL:sub@disp Applied
displacement in the plastic deformation problem.

####  2.2.2 History-dependent plastic deformation

In the second example, plastic deformation of a dog bone specimen under a
time-dependent loading history is studied, which follows from the recent work
by Koric et al. [28]. The dog bone specimen has a length of 110 mm and a width
(at the grip section) of 30 mm, with a gauge region width of 20 mm. A
schematic of the domain along with the mesh used in FE simulations are shown
in Fig. 6.

Figure 6: Schematic of the dogbone specimen and the mesh used in FE
simulations. The applied displacement is along the global X axis.

A total of 4756 linear plane stress elements (four-node quadrilaterals and
three-node triangles) were used to mesh the specimen with a plane-stress
thickness of 1 mm. In the absence of any body and inertial forces, the
equilibrium equations and boundary conditions can be stated in terms of the
Cauchy stress $\bm{\sigma}$ as:

$\displaystyle\nabla\cdot\bm{\sigma}=\bm{0},\;\;\forall\bm{X}\in\Omega,$ (4)
$\displaystyle\bm{u}=\bar{\bm{u}},\;\;\forall\bm{X}\in\partial\Omega_{u},$
$\displaystyle\bm{\sigma}\cdot\bm{n}=\bar{\bm{t}},\;\;\forall\bm{X}\in\partial\Omega_{t},$

where $\bm{n}$, $\bar{\bm{u}}$, and $\bar{\bm{t}}$ denote the outward boundary
normal, prescribed displacement, and prescribed traction, respectively. Under
small deformation assumption, total strain is given by:

$\bm{\epsilon}=\frac{1}{2}(\nabla\bm{u}+\nabla\bm{u}^{T}),$ (5)

The small-strain formulation of plasticity was used, so the total strain is
decomposed additively into its elastic and plastic parts:

$\bm{\epsilon}=\bm{\epsilon}^{e}+\bm{\epsilon}^{p}.$ (6)

For linear elastic and isotropic material under plane stress condition, the
constitutive equation is:

$\begin{bmatrix}\sigma_{11}\\\ \sigma_{22}\\\
\sigma_{12}\end{bmatrix}=\begin{bmatrix}\frac{E}{1-\nu^{2}}&\frac{\nu
E}{1-\nu^{2}}&0\\\ \frac{\nu E}{1-\nu^{2}}&\frac{E}{1-\nu^{2}}&0\\\
0&0&\frac{E}{2(1+\nu)}\end{bmatrix}\begin{bmatrix}\epsilon_{11}\\\
\epsilon_{22}\\\ \epsilon_{12}\end{bmatrix},$ (7)

where $E$ and $\nu$ are the Young’s modulus and Poisson’s ratio. In this work,
$J_{2}$ plasticity with linear isotropic hardening was used:

$\sigma_{y}(\bar{\epsilon}_{p})=\sigma_{y0}+H\bar{\epsilon}_{p},$ (8)

where $\sigma_{y}$, $\bar{\epsilon}_{p}$, $\sigma_{y0}$, and $H$ denote the
flow stress, equivalent plastic strain, initial yield stress, and the
hardening modulus, respectively. The material properties of the elastic-
plastic material response are presented in Table 3. The material model was
integrated implicitly in Abaqus/Standard [37].

Table 3: Material properties of the elastic-plastic material model Property |  $E$ [MPa] |  $\nu$ [/] |  $\sigma_{y0}$ [MPa] | H [MPa]
---|---|---|---|---
Value |  2.09$\times 10^{5}$ | 0.3 | 235 | 800

The specimen was fixed on the left side and a prescribed, time-dependent
displacement was applied on the right edge. Six control points were used to
define the loading path. Besides the two end points at $t=0$ and $t=1s$, four
other control points were randomly sampled from the range $[0.1,0.9]s$. The
applied displacement is 0 at $t=0$. The displacement magnitude at each control
point was randomly selected such that the nominal axial strain magnitude is
below 5%. Radial basis interpolation was used to interpolate the applied
displacement at arbitrary time points. A typical example of the applied
displacement is shown in Fig. 5b. A total of 4000 cases were generated, and
the von Mises stress was stored as the ground truth labels in the NN training,
and is defined as:

$\bar{\sigma}=\sqrt{\sigma_{11}^{2}+\sigma_{22}^{2}+\sigma_{11}\sigma_{22}+3\sigma_{12}^{2}}.$
(9)

##  3 Results and discussion

FE simulations of the transient heat transfer and plastic deformation were
conducted with eight high-end AMD EPYC 7763 Milan CPU cores. All NN training
and inference were conducted using a single Nvidia A100 GPU card on Delta, an
HPC cluster hosted at the National Center for Supercomputing Applications
(NCSA). Of the 4000 data points for each problem, 3900 were used in training
the NNs, and 100 were used in testing. To evaluate the model performance in
the test set, the relative $L_{2}$ error was used, which is given by:

${\rm{Mean\;L_{2}\;error}}=\frac{1}{N_{T}}\sum_{i=1}^{N_{T}}\frac{|f_{FE}-f_{Pred}|_{2}}{|f_{FE}|_{2}},$
(10)

where $N_{T}$, $f_{FE}$, and $f_{Pred}$ denote the total number of test cases,
the FE-simulated field value, and the NN-predicted field value, respectively.

###  3.1 Transient heat transfer results

For the heat transfer problem, training of the FNN-, LSTM- and GRU-based
DeepONet took 5486s, 20315s, and 18653s, respectively. The inference time for
the three NN models compared to FE simulation time is shown in Table 4.

Table 4: Computational cost of the plastic deformation problem | FE simulation time [s] | Inference time [s] | Speed up compared to FE (X)
---|---|---|---
FE simulation | 60 | / | /
FNN | / |  5.13 $\times 10^{-3}$ |  6.4$\times 10^{4}$
LSTM | / |  2.24 $\times 10^{-2}$ |  2.7$\times 10^{3}$
GRU | / |  2.02 $\times 10^{-2}$ |  3.0$\times 10^{3}$

Key statistics of the relative $L_{2}$ error on the test set for the three
models are shown in Table 5. In an average sense, across all testing samples,
both LSTM- and GRU-based DeepONet were able to lower the $L_{2}$ error of the
FNN-DeepONet by half. To provide a better perspective of the error
distribution over the 100 test cases, histograms of the error distribution are
depicted in Fig. 7. To show the spatial distribution of the prediction error,
selected cases are displayed in Fig. 8 with the final temperature distribution
of the solidifying slice for the testing cases with the lowest, median, and
largest $L_{2}$ errors, also dubbed as best, 50th percentile, and worst case.
Line plots of the temperature at the center line of the domain as predicted by
different DeepONet models are depicted in Fig. 9. It is observable that in the
worst case, FNN in the branch fails to predict the temperature distribution
correctly, particularly at the solidification front and the mushy zone, i.e.,
between Solidus and Liquidus temperatures. At the same time, LSTM and GRU in
the branch provided an almost indistinguishable distribution from the target
(FE simulation) temperature results for this case. Considering that the
thickness of solidifying shell at the mold exit is calculated from the
position of the solidifying front, in this case, the result inferred by
classical DeepONet (with FNN in branch) will be of significantly less value
for design, optimization, and online controls of this critical steel making
process.

Table 5: Error statistics of three DeepONet models on the solidification heat transfer problem | Mean | Std. deviation | Minimum | Maximum
---|---|---|---|---
FNN | 0.0016 | 0.0025 | 0.0005 | 0.0201
LSTM | 0.0005 | 0.0003 | 0.0002 | 0.0028
GRU | 0.0006 | 0.0004 | 0.0003 | 0.0031

(a)

(b)

(c)

Figure 7: Histograms showing how the relative $L_{2}$ error distributes over
the 100 test cases for the transient heat transfer case. (a)  (b) (a) Best
case
---
(c)  (d) (b) 50th percentile
(e)  (f) (c) Worst case
Figure 8: Comparison between FNN-, GRU- and LSTM-based DeepONets for the heat
transfer example. Rows 1-3 correspond to the best, 50th percentile, and worst
cases, respectively. The heat fluxes are shown in the first column.

(a) Best case

(b) 50th percentile

(c) Worst case

Figure 9: Line plots for the temperature distribution across the center line
for different DeepONet models.

###  3.2 History-dependent plastic deformation results

For the plastic deformation problem, training of the FNN-, LSTM- and GRU-based
DeepONet took 7805s, 20406s, and 18455s, respectively. The inference time for
the three NN models compared to the FE simulation time is shown in Table 6.

Table 6: Computational cost of the plastic deformation problem | FE simulation time [s] | Inference time [s] | Speed up compared to FE (X)
---|---|---|---
FE simulation | 21 | / | /
FNN | / |  1.23 $\times 10^{-3}$ |  1.7$\times 10^{4}$
LSTM | / |  2.28 $\times 10^{-2}$ |  9.2$\times 10^{2}$
GRU | / |  2.06 $\times 10^{-2}$ |  1.0$\times 10^{3}$

Key statistics of the relative $L_{2}$ error on the test set for the three
models are shown in Table 7. To provide a better perspective of the error
distribution over the 100 test cases, histograms of the error distribution are
depicted in Fig. 10. To show the spatial distribution of the DeepONet
prediction error, selected cases are displayed in Fig. 11. Line plots of the
stress field along a downward diagonal line across the dogbone domain domain
as predicted by different DeepONet models are depicted in Fig. 12.

Table 7: Error statistics of three DeepONet models on the plastic deformation problem | Mean | Std. deviation | Minimum | Maximum
---|---|---|---|---
FNN | 0.0288 | 0.0712 | 0.0019 | 0.4449
LSTM | 0.0104 | 0.0237 | 0.0012 | 0.1281
GRU | 0.0091 | 0.0200 | 0.0010 | 0.1325

(a)

(b)

(c)

Figure 10: Histograms showing how the relative $L_{2}$ error distributes over
the 100 test cases for the plastic deformation of the dog bone specimen. (a)
(b) (a) Best case
---
(c)  (d) (b) 50th percentile
(e)  (f) (c) Worst case
Figure 11: Comparison between FNN- and RNN-based DeepONet for the plastic
deformation example. Rows 1-3 correspond to the best, 50th percentile, and
worst cases, respectively. The normalized applied displacements are shown in
the first column.

(a) Best case

(b) 50th percentile

(c) Worst case

Figure 12: Line plots for the stress field across a downward diagonal line for
different DeepONet models.

The results in Fig. 11 and Fig. 12 indicate that the RNN-based S-DeepONets
were able to predict the stress field significantly better than the regular
FNN-based DeepONet in the worst-case scenario. In addition to the superior
worst-case performance, on average, the S-DeepONets are twice as accurate
compared to the FNN-based DeepONet, as seen in Table 7. These findings prove
that the proposed S-DeepONet architecture effectively encodes complex, time-
dependent loading histories and performs better than the traditional FNN-based
DeepONet. The spatial distribution of the predicted stress field closely
resembles the FE ground truth, which indicates that the proposed S-DeepONets
combines the temporal encoding capability of the RNN structure with the
spatial encoding capability of the DeepONet architecture to improve prediction
accuracy. However, it is worth noting that the training time of the RNN-based
DeepONets were significantly longer than the FNN-based counterpart due to its
more complex architecture and its recurrent nature. However, once the models
were trained, the inference speed was several orders of magnitude faster than
the FE simulation, which makes the S-DeepONet an attractive candidate where
repeated function evaluation with different input loading histories are
needed.

##  4 Conclusions and future work

After introducing the classical data-driven DeepONet framework, which
leverages the fully connected neural networks (FNN) in the branch and trunk,
we introduced the S-DeepONet, which we have formulated and devised with
advanced, recurrent neural networks of LSTM and GRU types in the branch. We
have then focused on learning full-field temperature and stress solutions in
the two highly nonlinear practical applications of thermal and mechanical
types with complex random loading histories. In both cases, S-DeepONet
provided significantly more accurate predictions than classical DeepONet,
proving that sequential learning methods in the branch of DeepONet can
universally and effectively encode loading histories no matter the underlying
physics of the problem. S-DeepONet was able to half the average error among
all testing samples compared to DeepONet. The difference was even more
profound for the worst testing sample, where S-DeepONet was able to correctly
predict the temperature distribution in the solidification case and even
recover stress concentration details in the plastic deformation of a dog bone
specimen. We have shown that S-DeepONet can be adequately trained on the
current high-end GPUs within a few hours. Moreover, once the S-DeepONet is
properly trained, it can instantly infer accurate full-field results and
several orders of magnitude faster than classical nonlinear numerical methods.
Unlike classical DeepONet, which can provide instant full-field evaluations
for preliminary design and optimizations, the increased accuracy of S-DeepONet
may be robust enough to be used in full high-fidelity modeling, optimization,
and design scenarios whenever a high number of forward evaluations with
parametric histories are needed in many complex nonlinear and non-equilibrium
applications and processes in engineering and science. Furthermore, the
trained S-DeepONet weights and biases can be transported to laptops and even
edge computing devices and used for inference without the use of GPUs or even
modeling tools for instant predictions in many online control scenarios. As
the capabilities and memory of GPU hardware rapidly increase, in our future
work, we will modify S-DeepONets for learning full solution fields from three-
dimensional modeling domains.

## Replication of results

The data and source code supporting this study’s findings can be made
available to the corresponding author upon reasonable request.

## Conflict of interest

The authors declare that they have no conflict of interest.

## Acknowledgements

The authors would like to thank the National Center for Supercomputing
Applications (NCSA) at the University of Illinois, and particularly its
Research Computing Directorate, the Industry Program, and the Center for
Artificial Intelligence Innovation (CAII) for their support and hardware
resources. This research is a part of the Delta research computing project,
which is supported by the National Science Foundation (award OCI 2005572) and
the State of Illinois, as well as the Illinois Computes program supported by
the University of Illinois Urbana-Champaign and the University of Illinois
System. Finally, the authors would like to thank Professors George
Karniadakis, Lu Lu, and the Crunch team at Brown, whose original work with
DeepONets inspired this research.

## CRediT author contributions

Jaewan Park: Methodology, Formal analysis, Investigation, Writing - Original
Draft. Shashank Kushwaha: Methodology, Formal analysis, Investigation, Writing
- Original Draft. Junyan He: Methodology, Formal analysis, Investigation,
Writing - Original Draft. Seid Koric: Conceptualization, Methodology,
Supervision, Resources, Writing - Original Draft, Funding Acquisition. Diab
Abueidda: Supervision, Writing - Review & Editing. Iwona Jasiuk: Supervision,
Writing - Review & Editing.

## References

  * Perumal et al. [2023] V Perumal, D Abueidda, S Koric, and A Kontsos.  Temporal convolutional networks for data-driven thermal modeling of directed energy deposition.  _Journal of Manufacturing Processes_ , 85:405–416, 2023\.
  * Sadeghpour and Nonn [2022] Ebrahim Sadeghpour and Aida Nonn.  Data-driven models for structure-property prediction in additively manufactured steels.  _Computational Materials Science_ , 215:111782, 2022.
  * Mozaffar et al. [2019] M Mozaffar, R Bostanabad, W Chen, K Ehmann, Jian Cao, and MA Bessa.  Deep learning predicts path-dependent plasticity.  _Proceedings of the National Academy of Sciences_ , 116(52):26414–26420, 2019.
  * Kollmann et al. [2020] Hunter T Kollmann, Diab W Abueidda, Seid Koric, Erman Guleryuz, and Nahil A Sobh.  Deep learning for topology optimization of 2d metamaterials.  _Materials & Design_, 196:109098, 2020.
  * He et al. [2023a] Junyan He, Charul Chadha, Shashank Kushwaha, Seid Koric, Diab Abueidda, and Iwona Jasiuk.  Deep energy method in topology optimization applications.  _Acta Mechanica_ , 234(4):1365–1379, Apr 2023a.  ISSN 1619-6937.  doi: 10.1007/s00707-022-03449-3.  URL https://doi.org/10.1007/s00707-022-03449-3.
  * Hoq et al. [2023] Enjamamul Hoq, Osama Aljarrah, Jun Li, Jing Bi, Alfa Heryudono, and Wenzhen Huang.  Data-driven methods for stress field predictions in random heterogeneous materials.  _Engineering Applications of Artificial Intelligence_ , 123:106267, 2023.
  * Bleiziffer et al. [2021] Patrick Bleiziffer, Juergen Hofmann, Robert Zboray, Thorsten Wiege, and Roger Herger.  Predicting the fiber orientation in glass fiber reinforced polymers using the moment of inertia and convolutional neural networks.  _Engineering Applications of Artificial Intelligence_ , 104:104351, 2021.
  * Egli et al. [2021] Franziska S Egli, Richard C Straube, André Mielke, and Tim Ricken.  Surrogate modeling of a nonlinear, biphasic model of articular cartilage with artificial neural networks.  _PAMM_ , 21(1):e202100188, 2021.
  * Abueidda et al. [2021a] Diab W Abueidda, Seid Koric, Nahil A Sobh, and Huseyin Sehitoglu.  Deep learning for plasticity and thermo-viscoplasticity.  _International Journal of Plasticity_ , 136:102852, 2021a.
  * Raissi et al. [2019] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational physics_ , 378:686–707, 2019\.
  * Abueidda et al. [2021b] Diab W Abueidda, Qiyue Lu, and Seid Koric.  Meshless physics-informed deep learning method for three-dimensional solid mechanics.  _International Journal for Numerical Methods in Engineering_ , 122(23):7182–7201, 2021b.
  * Nguyen-Thanh et al. [2020] Vien Minh Nguyen-Thanh, Xiaoying Zhuang, and Timon Rabczuk.  A deep energy method for finite deformation hyperelasticity.  _European Journal of Mechanics-A/Solids_ , 80:103874, 2020\.
  * Samaniego et al. [2020] Esteban Samaniego, Cosmin Anitescu, Somdatta Goswami, Vien Minh Nguyen-Thanh, Hongwei Guo, Khader Hamdia, X Zhuang, and T Rabczuk.  An energy approach to the solution of partial differential equations in computational mechanics via machine learning: Concepts, implementation and applications.  _Computer Methods in Applied Mechanics and Engineering_ , 362:112790, 2020.
  * Abueidda et al. [2022] Diab W Abueidda, Seid Koric, Rashid Abu Al-Rub, Corey M Parrott, Kai A James, and Nahil A Sobh.  A deep learning energy method for hyperelasticity and viscoelasticity.  _European Journal of Mechanics-A/Solids_ , 95:104639, 2022\.
  * Abueidda et al. [2023] Diab W Abueidda, Seid Koric, Erman Guleryuz, and Nahil A Sobh.  Enhanced physics-informed neural networks for hyperelasticity.  _International Journal for Numerical Methods in Engineering_ , 124(7):1585–1601, 2023.
  * Fuhg and Bouklas [2022] Jan N Fuhg and Nikolaos Bouklas.  The mixed deep energy method for resolving concentration features in finite strain hyperelasticity.  _Journal of Computational Physics_ , 451:110839, 2022.
  * He et al. [2023b] Junyan He, Diab Abueidda, Rashid Abu Al-Rub, Seid Koric, and Iwona Jasiuk.  A deep learning energy-based method for classical elastoplasticity.  _International Journal of Plasticity_ , page 103531, 2023b.
  * He et al. [2023c] Junyan He, Diab Abueidda, Seid Koric, and Iwona Jasiuk.  On the use of graph neural networks and shape-function-based gradient computation in the deep energy method.  _International Journal for Numerical Methods in Engineering_ , 124(4):864–879, 2023c.
  * Haghighat et al. [2021] Ehsan Haghighat, Maziar Raissi, Adrian Moure, Hector Gomez, and Ruben Juanes.  A physics-informed deep learning framework for inversion and surrogate modeling in solid mechanics.  _Computer Methods in Applied Mechanics and Engineering_ , 379:113741, 2021.
  * Cai et al. [2021] Shengze Cai, Zhicheng Wang, Frederik Fuest, Young Jin Jeon, Callum Gray, and George Em Karniadakis.  Flow over an espresso cup: inferring 3-d velocity and pressure fields from tomographic background oriented schlieren via physics-informed neural networks.  _Journal of Fluid Mechanics_ , 915:A102, 2021.
  * Chen and Chen [1995] Tianping Chen and Hong Chen.  Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems.  _IEEE Transactions on Neural Networks_ , 6(4):911–917, 1995.
  * Lu et al. [2021a] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis.  Learning nonlinear operators via deeponet based on the universal approximation theorem of operators.  _Nature machine intelligence_ , 3(3):218–229, 2021a.
  * Wang et al. [2021] Sifan Wang, Hanwen Wang, and Paris Perdikaris.  Learning the solution operator of parametric partial differential equations with physics-informed deeponets.  _Science advances_ , 7(40):eabi8605, 2021.
  * Koric and Abueidda [2023] Seid Koric and Diab W Abueidda.  Data-driven and physics-informed deep learning operators for solution of heat conduction equation with parametric heat source.  _International Journal of Heat and Mass Transfer_ , 203:123809, 2023.
  * Goswami et al. [2022] Somdatta Goswami, Minglang Yin, Yue Yu, and George Em Karniadakis.  A physics-informed variational deeponet for predicting crack path in quasi-brittle materials.  _Computer Methods in Applied Mechanics and Engineering_ , 391:114587, 2022.
  * Yin et al. [2022] Minglang Yin, Enrui Zhang, Yue Yu, and George Em Karniadakis.  Interfacing finite elements with deep neural operators for fast multiscale modeling of mechanics problems.  _Computer Methods in Applied Mechanics and Engineering_ , 402:115027, 2022.
  * He et al. [2023d] Junyan He, Seid Koric, Shashank Kushwaha, Jaewan Park, Diab Abueidda, and Iwona Jasiuk.  Novel deeponet architecture to predict stresses in elastoplastic structures with variable complex geometries and loads.  _arXiv preprint arXiv:2306.03645_ , 2023d.
  * Koric et al. [2023] Seid Koric, Asha Viswantah, Diab W Abueidda, Nahil A Sobh, and Kamran Khan.  Deep learning operator network for plastic deformation with variable loads and material properties.  _Engineering with Computers_ , pages 1–13, 2023.
  * Lu et al. [2021b] Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis.  Deepxde: A deep learning library for solving differential equations.  _SIAM review_ , 63(1):208–228, 2021b.
  * Abadi et al. [2015] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015\.  URL https://www.tensorflow.org/.  Software available from tensorflow.org.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * He et al. [2023e] Junyan He, Shashank Kushwaha, Diab Abueidda, and Iwona Jasiuk.  Exploring the structure-property relations of thin-walled, 2d extruded lattices using neural networks.  _Computers & Structures_, 277:106940, 2023e.
  * Schmidhuber et al. [1997] Jürgen Schmidhuber, Sepp Hochreiter, et al.  Long short-term memory.  _Neural Comput_ , 9(8):1735–1780, 1997.
  * Cho et al. [2014] Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio.  Learning phrase representations using rnn encoder-decoder for statistical machine translation.  _arXiv preprint arXiv:1406.1078_ , 2014.
  * Malhotra et al. [2016] Pankaj Malhotra, Anusha Ramakrishnan, Gaurangi Anand, Lovekesh Vig, Puneet Agarwal, and Gautam Shroff.  Lstm-based encoder-decoder for multi-sensor anomaly detection.  _arXiv preprint arXiv:1607.00148_ , 2016.
  * Sutskever et al. [2014] Ilya Sutskever, Oriol Vinyals, and Quoc V Le.  Sequence to sequence learning with neural networks.  _Advances in neural information processing systems_ , 27, 2014.
  * SIMULIA [2020] SIMULIA.  Abaqus, 2020.
  * Zappulla et al. [2020] Matthew LS Zappulla, Seong-Mook Cho, Seid Koric, Hyoung-Jun Lee, Seon-Hyo Kim, and Brian G Thomas.  Multiphysics modeling of continuous casting of stainless steel.  _Journal of Materials Processing Technology_ , 278:116469, 2020.
  * Koric and Thomas [2006] Seid Koric and Brian G Thomas.  Efficient thermo-mechanical model for solidification processes.  _International journal for numerical methods in engineering_ , 66(12):1955–1989, 2006.


11institutetext:  Key Laboratory for Computational Astrophysics, National
Astronomical Observatories, Chinese Academy of Sciences, Beijing 100101, China
School of Astronomy and Space Science, University of Chinese Academy of
Sciences, Beijing 100049, China<EMAIL_ADDRESS>
Received 0000 xxxx 00; accepted 0000 xxxx 00

# PhotoNs-GPU:A GPU accelerated cosmological simulation code

Qiao Wang  Chen Meng

###### Abstract

We present a GPU-accelerated cosmological simulation code, PhotoNs-GPU, based
on algorithm of Particle Mesh Fast Multipole Method (PM-FMM), and focus on the
GPU utilization and optimization. A proper interpolated method for truncated
gravity is introduced to speed up the special functions in kernels. We verify
the GPU code in mixed precision and different levels of interpolated method on
GPU. A run with single precision is roughly two times faster that double
precision for current practical cosmological simulations. But it could induce
a unbiased small noise in power spectrum. Comparing with the CPU version of
PhotoNs and Gadget-2, the efficiency of new code is significantly improved.
Activated all the optimizations on the memory access, kernel functions and
concurrency management, the peak performance of our test runs achieves 48% of
the theoretical speed and the average performance approaches to $\sim 35\%$ on
GPU.

###### keywords:

methods: numerical – cosmology: theory - large-scale structure of universe

##  1 Introduction

High resolution N-body simulations are essential tools to understand the
formation and evolution of dark matter from sub-halos to large scale structure
of our Universe (Springel et al. 2008; Angulo et al. 2012; Wang et al. 2020).
During the past 50 years, the dramatic increase in the size of the simulation
is not only because of the rapid progress on the supercomputer hardware, but
also be owed to the rapid development of the N-body solvers (Ishiyama et al.
2012; Habib et al. 2016; Yu et al. 2017; Potter et al. 2017; Cheng et al.
2020; Ishiyama et al. 2020).

One of popular N-body method is the Fast multiple method (FMM), it has an
attracting feature with a time complexity O(N), ideal for performing extreme
large simulation (Greengard & Rokhlin 1987; Cheng et al. 1999; Dehnen 2002,
2014). Indeed some extremely large simulations have been performed with the
FMM N-body solver, for instance, Potter et al. (2017) completed a cosmological
simulation with two trillion particles on the Piz Daint, GPU supercomputer,
using the FMM code of PKDGRAV (Stadel 2001). Recently, a hybrid method of
Particle Mesh Fast Multipole Method (PM-FMM) is introduced to cosmological
simulations. Analogue to the TreePM method (Xu 1995; Bagla 2002; Dubinski et
al. 2004; Springel 2005; Ishiyama et al. 2009; Wang et al. 2018), it
calculates the short-range tree method by FMM. Wang (2021) presented the
details of the method, its advantage and accuracy, with the code PhotoNs-2.
Note, Springel et al. (2020) also provides an alternative implementation in
Gadget-4 code.

In the era of heterogeneous supercomputer, it is important to develop
efficient N-body solvers based on these heterogeneous systems, instead of
homogeneous ones (Makino & Taiji 1995; Makino 2004). In recent years, some
Branes-Hut Tree (Barnes & Hut 1986) and FMM N-body solvers have been
successfully developed on Graphics Processing Unit (GPU) platforms (Hamada et
al. 2009; Hamada & Nitadori 2010; Gaburov et al. 2010; Bédorf et al. 2014;
Gumerov & Duraiswami 2008; Yokota & Barba 2011; Yokota 2012).

In this paper, we develop a GPU-accelerated PM-FMM code based on PhotoNs-2,
especially on the NVidia architecture, which is a referred to as PhotoNs-GPU.

This paper is organized as follows. In section 2, we briefly introduce the
algorithm of PM-FMM method. The detail of GPU implementation is discussed in
section 3. Two test runs are presented on performance and accuracy of PhotoNs-
GPU in section 4. Finally, some potential issues are discussed in Section 5.

##  2 algorithm and Code

First, we begin with a brief review of the PM-FMM method developed by Wang
(2021), in particularly the relevant parts for GPU acceleration in following
sections. A Gaussian transition function is employed to split gravity into a
smoothed long-range and short-range part, $\phi=\phi_{long}+\phi_{short}$. The
$\phi_{long}$ is satisfied to Poisson’s equation which is solved by a
Particle-Mesh method based on the convolution of density field $\rho(x)$ with
Green function of long-range gravitational force on a regular mesh. Thus,
$\phi_{long}(x)=\mathcal{F}^{-1}[(\hat{\rho}_{k}/k^{2})\exp(-k^{2}/4r_{s}^{2})]$,
where $\hat{\rho}_{k}=\mathcal{F}[\rho(x)]$ is the density field in Fourier
space and the long-range potential is smoothed by a Gaussian function with a
split radius $r_{s}=1.2~{}\Delta_{g}$ ($\Delta_{g}\equiv$ BOXSIZE / NPM).

The short-range gravity is computed by a truncated FMM which dominates the
most amount of computation and is accelerated on GPU in this work. Similar to
a conventional FMM, all particles in a computing domain are organized into a
tree structure and the finest tree cells (or tree nodes) point to continuous
particle packs (or leaves). An Orthogonal Recursive Bisection (ORB) tree is
employed and the particles belonged to a cell are equally assigned into two
offspring cells till to leaves. The maximum particle number allowed in a leaf
is constrained by the parameter of ‘MAXLEAF’.

In FMM, gravity of a particle is not directly accumulated by tree nodes or
particles. Instead, the computation is based on particle packs or leaves.
First, the information of particle distribution in leaves is transferred into
the multipole of those leaves by operator P2M (Particle to Multipole). Thus
the multipole of lower nodes built up the higher one, M2M (Multipole to
Multipole). This process is called PASS-UP. The interaction between multipoles
is computed by operator M2L (Multipole to Local) and the local multipole is
passed down to the lower nodes by L2L (Local to Local). Finally, gravity of a
particle in local leaf is determined by the local multipole of gravitational
potential L2P (Local to Particle) and direct interaction from particles in
neighborhoods P2P (Particle to Particle).

All operators are straightforwardly derived from multipole expansion of
gravitational interaction and M2L is specified by equation

${\mathcal{L}}_{\mathbf{n}}({\mathbf{z}}_{B})=\sum_{|{\mathbf{m}}|=0}^{p-|{\mathbf{n}}|}{\mathcal{M}}_{\mathbf{m}}({\mathbf{z}}_{A}){\mathcal{D}}_{{\mathbf{n}}+{\mathbf{m}}}({\mathbf{z}}_{B}-{\mathbf{z}}_{A}),$

where
${\mathcal{D}}_{\mathbf{n}}\equiv\nabla^{\mathbf{n}}\psi(r)=f_{(n)}\bar{{\mathbf{r}}}^{\mathbf{n}}$
is a traceless operator. $\bar{{\mathbf{r}}}^{\mathbf{n}}$ is a displacement
tensor and the prefactors of

$f^{\rm inv}_{(n)}(r)=(-1)^{n}\frac{(2n-1)!!}{r^{2n+1}},$

for inverse-square law. But for truncated short-range gravity, it becomes
complicated. One can compute the prefactor $f_{(n)}$ of any order $n$ by
equation of

$\displaystyle(-1)^{n}r_{s}^{2n+1}f_{(n)}(x)=\frac{(2n-1)!!}{2^{{2n+1}}}\frac{\textrm{erfc}\left(x\right)}{x^{2n+1}}$
$\displaystyle+\sum_{q=1}^{n}\frac{2^{-q-n}(2n-1)!!}{(2n-2q+1)!!}\frac{e^{-x^{2}}}{\sqrt{\pi}x^{2q}},$

where ${x\equiv{r}/{2r_{s}}}$. We numerically truncate the function at cutoff
radius $\sim 6~{}\Delta_{g}$, following a traditional Gaussian splitting
approach. Meanwhile, the split function for P2P direction summation of
truncation gravity is computed by

${\mathbf{F}_{s}}({\mathbf{r}})=-\frac{\mathbf{r}}{r^{3}}T(r;r_{s}),$

where the truncation function reads

$T\equiv\textrm{erfc}\left(\frac{r}{2r_{s}}\right)+\frac{r}{r_{s}\sqrt{\pi}}\exp\left(-\frac{r^{2}}{4r_{s}^{2}}\right).$
(1)

We only summarize the relevant algorithms and formulas used for GPU
acceleration in this section, and we refer the readers for more technical
details to Wang (2021).

##  3 GPU implementation

In practise, the short-range gravity in PhotoNs-2 is estimated by walking the
multipole tree to determine whether operators M2L or P2P between two multipole
nodes/leaves need to be considered. We use two interaction lists to record the
interaction pair of M2L and P2P by one traversal, respectively. Meanwhile, we
identify a M2l or P2P pair as a task so that the interaction lists naturally
become task queues. Usually, the lists are long enough to guarantee the
concurrency on multi-core CPUs or GPUs. Two task lists are both dealt with on
CPU in PhotoNs-2. In the GPU version, we push the P2P task on GPU, because of
the dominant computation amount of P2P over others.

###  3.1 Memory layout

The information of each particle is stored in a predefined data structure. It
is suitable to communicate between computing domains, but does not match the
coalesced memory access mechanism of GPU when loading data. Therefore we
transfer an Array of Structure (AoS) of particles into Structure of Array
(SoA) of particle components, e.g., position, velocity, acceleration etc.
Since the particles are already reordered in the process of tree building,
such a transformation is natural and widely used in many N-body problems with
CUDA (Nyland et al. 2009). The transformation executes on host memory, then
the arrays of position is sent to GPU.

Figure 1: Memory layout and relation on CPU and GPU (shadow zone). The dark
gray box denotes a SM unit. Arrows indicates the flow of data.

Fig 1 is schematic diagram of data arrays in the host memory and GPUs. There
are three main data arrays for particle, leaf and tasks. Each particle
structure contains position, velocity and acceleration, etc. We abstract the
position information of x, y, z (Pos) into individual arrays (SoA) for further
optimal GPU L2 cache performance. The task list (Task) of P2P interaction is
an array of pair $<i,j>$ which indicates attractive force of particles in
$i$-leaf induced by $j$-leaf.

They are sent to the GPU global memory denoted by the gray shadow zone. One
Streaming Multiprocessor (SM) is shown with a dark gray box, which contains a
chip of level-1 data cache (L1D$) plus shared memory (SMEM) and a group of
sub-cores with registers and Float Point Units, etc. We allocate two main
arrays to store Particle and Leaf. Each leaf records starting index of first
particle and particle number in the leaf.

###  3.2 GPU kernel

Arrays of Task list, Leaf and Pos are transferred to the global memory of GPU
via PCIe port. The item with the same index in Task list is not continuous in
memory, which dramatically slows down the speed of data accesses into CUDA
cores. To fix this problem, we reorder Task list with respect to $i$ and
rename this list as TaskGroup. A CUDA library thrust is utilized for this
reordering. Meanwhile the index array of Task Group (GroupIdx) is produced to
record the starting points and length of tasks for every $i$-leaf in the task
list.

According to the scheduling mechanism of GPU, we deal with all interactions in
one task group with a thread block, which corresponds to a SM unit (Hamada et
al. 2009). Since every task group contains only one $i$-leaf and all particles
in $i$-leaf are continuously stored, each thread responds to one particle in
$i$-leaf and the particles in $j$-leaf are prefetched into Shared memory for
interaction calculation.

For instance, we assume that the gravity of leaf 8 is induced by leaf 9 and 10
in Fig. 1. As the 2nd column in ‘GroupIdx’, Thread block 1 is correspondent to
it. Four particles in leaf 8 is sent into the register of sub-cores and
particle position in leaf 9 and 10 are sent to shared memory of SM. Core
function returns the acceleration acc$\\_$i of leaf 8 (see implementation of
Core function in Alg. 2). After all calculation of gravity is done, array of
acceleration is transferred back to host memory.

Algorithm 1 P2P Kernel

__shared__ sh_pos_j

for $b\leftarrow$ blockIdx.x to $len$(GroupIdx) do

i$\\_$leaf = GroupIdx.il[b]

$t$ = threadIdx.x

if $t$ $\textless$ Leaf[i_leaf].np then

load pos[$t$ \+ Leaf[i$\\_$leaf].ip] to register pos_i

reset register acc_i with 0

end if

j_start = GroupIdx.jl[b]

j_end = GroupIdx.jl[b] + GroupIdx.jn[b]

for $j\leftarrow$ j_start to j_end do

j_leaf = TaskGroup.j[$j$]

__syncthreads__

if $t$ $\textless$ Leaf[j_leaf].np then

preload pos[$t$+Leaf[j_leaf].ip]

to SMEM sh_pos_j

end if

__syncthreads__

nj = Leaf[j_leaf].np

acc_i $+=$ Core(pos_i, sh_pos_j; nj)

$\triangleright$ goto Algorithm 2

end for

accumulate register acc_i to global memory acc

$b$ += gridDim.x

end for$\triangleright$ block level parallelism

###  3.3 Interpolation function

In previous section, we present the kernel and memory layout on GPU. In
function of Core, two special functions are necessary in this kind of
splitting method. Compared with the computation of pair-wise inverse-square
gravity, they seriously decrease the efficiency of P2P kernel. The Gaussian
shaped prefactor contains an exponential function (exp) and an error function
(erf). Unlike function rsqrt estimated by an ‘intrinsic’ implementation, exp
could consume about 5.5 times the operations of add/mul in fast math library
on a NVidia GPU. We catch the effective number of floating-point operations
(flops) by the tool of nvprof and find that erf needs to implicitly call
functions of rcp, exp and several extra operations. This estimation is also
consistent with the measurement of Arafa et al. (2019).

In order to speed up the evaluation of special functions, we combine those two
functions into a table for the interpolation algorithm. The truncation
function (Eq. 1) can be transformed to an integral form,

$T(x\equiv\frac{r}{2r_{s}})=\frac{4}{\sqrt{\pi}}\int_{x}^{\infty}u^{2}e^{-u^{2}}du,$

by using partial integration. Obviously, its Taylor’s series reads

$\displaystyle T(x)$ $\displaystyle=T_{i}$ $\displaystyle:T_{[0]}$ (2)
$\displaystyle+x_{i}^{2}E_{i}\epsilon$ $\displaystyle:T_{[1]}$ $\displaystyle-
x_{i}(x_{i}^{2}-1)E_{i}\epsilon^{2}$ $\displaystyle:T_{[2]}$
$\displaystyle+(\frac{2}{3}x_{i}^{4}-\frac{5}{3}x_{i}^{2}+\frac{1}{3})E_{i}\epsilon^{3}$
$\displaystyle:T_{[3]}$ $\displaystyle-
x_{i}(\frac{1}{3}x_{i}^{4}-\frac{3}{2}x_{i}^{2}+1)E_{i}\epsilon^{4}$
$\displaystyle:T_{[4]}$ $\displaystyle+\mathcal{O}(\epsilon^{5}),$

where $\epsilon=x-x_{i}$ , $E_{i}=-4{\exp}({-x_{i}^{2}})/{\sqrt{\pi}}$ and
$T_{i}=\textrm{erfc}\left(x_{i}\right)-x_{i}E_{i}/2$ at the $i$-th gird of
table in the range of [0,3]. In this case, we set grid number of $E_{i}$ and
$T_{i}$ to 512, considering the suitable capacity of shared memory (light and
heavy brown arrays in Fig. 1).

Figure 2: From top to bottom, The curves denote the maximum envelop of
relative error with respect to exact splitting function at the interpolation
level of $T_{[0]}$, $T_{[1]}$, $T_{[2]}$, $T_{[3]}$ and $T_{[4]}$,
respectively.

We measure the precision of interpolation method to link accuracy with
interpolation level. Fig. 2 shows that the maximum value of relative errors at
different levels. We show the absolute value so that the transition of
positive and negative value occurs at the position of the cusps of curves.
According to the measurement result, $T_{[2]}$ is adequate for the single-
precision calculation and $T_{[4]}$ is for double-precision. In the next
section, we run some practical simulations to observe the overall
interpolation effect.

Eq. 2 contains thirty nine floating point operations and $loads$ two float
number from shared memory. The floating point operations can be reduced to
fourteen by variable reuse and Fused Multiply-Add (FMA). Combining with other
optimizations, the amount of flops can be suppressed by $\sim 60\%$ in total.
The modified ‘Core’ is presented in Alg. 2.

Algorithm 2 Core function

function Core(pos_i, sh_pos_j; nj)

for $j\leftarrow$ 0 to nj do

pos$\\_$j = sh_pos_j[$j$] $\triangleright$ read from SMEM

dx = pos$\\_$j.x - pos$\\_$i.x

dy = pos$\\_$j.y - pos$\\_$i.y

dz = pos$\\_$j.z - pos$\\_$i.z

r2 = dx*dx + dy*dy + dz*dz

determine table index $d$ and $\epsilon$ in Eq. 2

load $T_{d}$

load $E_{d}$

idr = rsqrt(r2)

pref = $T_{[0]}(d,\epsilon)$

pref += $T_{[1]}(d,\epsilon)$

pref += $T_{[2]}(d,\epsilon)$ $\triangleright$ single precision

pref += $T_{[3]}(d,\epsilon)$

pref += $T_{[4]}(d,\epsilon)$ $\triangleright$ double precision

grav = GM*idr*idr*idr*pref

ax = grav*dx

ay = grav*dy

az = grav*dz

end for

return acc_i = (ax, ay, az)

end function$\triangleright$ thread level parallelism

##  4 test run

We modify the CPU-based PM-FMM code to a GPU-accelerated version that is
referred to as PhotoNs-GPU. For testing the precision and efficiency of our
code, we carry out two groups of cosmological simulations on a 1.5 TB memory
GPU server with thirty-two cores of Intel Xeon 5218 CPU (2.3 GHz) and two
NVidia Tesla 32GB V100S GPUs. The theoretic single precision (SP) performance
of the machine is about 32.7 TFlops and the PCIe bandwidth is about 10 GB/s
for each GPU.

###  4.1 Precision check

We test the accuracy of our code by comparing simulations run with Gadget-2
and PhotoNs-GPU. We employ $256^{3}$ particle in 1 $h^{-1}$Gpc simulation box.
The initial condition at $z_{i}=99$ is generated by 2LPTic (Crocce et al.
2006), which is evolved for four simulations by Gadget-2 and PhotoNs-GPU,
respectively.

We use 16 MPI processes to carry out the simulation to the present $z=0$ with
N${}_{\bf PM}=256$ and $100~{}h^{-1}$kpc softening length. Gadget-2 executes
single step by 59 seconds of wall clock time and about 11,000 seconds in
total. The power spectrum at $z=0$ is indicated with gray points in Fig. 3.
Based on the same parameters and settings, we run two different simulations to
check the influence of floating point precision on GPU. The calculation on GPU
with single-precision (SP) is indicated with Golden dashed curve and double-
precision (DP) is blue dash-dotted curve. In the same condition, PhotoNs-GPU
executes single step by $\sim 1.3$ seconds of wall clock time and about 1,000
seconds in total, using 16 MPI processes and two V100 GPUs.

Figure 3: The comparison of power spectrum at z=0. The gray points denote
Gadget-2 result, the heavy orange dashed line denotes SP run and the blue
dash-dotted line denotes DP run by PhotoNs-GPU. In the bottom panel, the
residuals of power spectrum show that the DP and SP runs are both consistent
with the fiducial one.

The forth simulation is carried out as a fiducial one, which adopt direct P2P
method to replace the Tree or FMM in order to exactly compute short-range
gravity. Meanwhile we increase time steps by roughly three times. The bottom
panel of Fig. 3 is the residual of power spectrum with respect to the fiducial
simulation. The contrast of power spectrum is smaller than 1$\%$. Fig 3 shows
that the SP case is well consistent with DP one, but more noisy at small scale
($k\sim 1$). Since PhotoNs-GPU employs a larger cutoff radius and more P2P
contribution for gravity solver, our accuracy is slightly better than Gadget-2
in power spectrum statistics at the same settings.

###  4.2 Planck’s cosmology

According to the recent observation of the best-fit ’plik’ cosmological
parameters from Planck 2018 (Planck Collaboration et al. 2020):
$\Omega_{\mathrm{\Lambda}}$=0.684, $\Omega_{\mathrm{c}}$= 0.265,
$\Omega_{\mathrm{b}}$ = 0.0494, $H_{0}$= 67.32 $km~{}s^{-1}$Mpc-1,
$\sigma_{8}$=0.812 and $n_{s}$ = 0.966, we employ $512^{3}$ particles in 100
$h^{-1}$Mpc simulation box. Since the test in the previous section indicates
that a mixed-precision run is also consistent with the fiducial results in DP,
we set all variable as SP to reduce the memory occupation and amount of
communications and to utilize more Floating points units. The execution time
is about five hours of wall clocks from redshift $z=99$ to $z=0$, using
sixteen MPI processes and two Tesla V100S GPUs.

Figure 4: The comparison of mass function (left panel) and power spectrum
(right panel) of test run with theoretical curves. The (red) points, (green)
squares and (blue) triangles denote the simulation results at $z=0,1,3$,
respectively.

The initial condition is constructed by a built-in IC module, following the
Zel’dovich approximation method(Springel 2015) to effectively match our domain
decomposition. The built-in module is based on the subroutine of 2LPTic and
the notations are following the appendix A of Crocce et al. (2006). The input
initial power spectrum is calculated by CAMB (Lewis et al. 2000).

Dark haloes mass functions at redshift $z=0,1,3$ are shown in the left panel
of Fig. 4 and power spectra in the right panel. The mass function is measured
by an on-the-fly Friend-of-Friend (FoF) halos finder and the theoretical
curves are produced by Reed et al. (2007). The non-linear power spectrum is
predicted by the HALOFIT (Takahashi et al. 2012; Mead et al. 2016). The
deviation of power spectrum of $z=3$ (triangles) occurs at scale of $k\sim 5$.
It could be due to the lack of sampling. Higher resolution test shows that
that deviation disappears and the simulation result is well consistent with
the predictions.

##  5 discussion

We present a GPU implementation of the PM-FMM algorithm by using an
interpolated improvement for truncated pairwise interactions, then we tune the
data structure to speed up GPU memory accesses and adjust the order of
instructions in the kernel. A SP level test shows that SP runs can catch
adequate accuracy on the statistics of large scale structure.

Briefly, after memory arrangement, interpolation and instruction fine tuning,
the efficiency of code is significantly improved. In the case of the $512^{3}$
simulations with Planck’s cosmology , the number of particle-particle
interactions per step averaged in the first ten steps is $\sim 1.74\times
10^{12}$. The mean wall-clock time per step is 7.62 seconds, in which the
kernel time is 5.5 seconds. Thus, the average performance is 11.65 TFlops and
the peak performance of the kernel is 16.134 TFlops. Here, we use the
operation count of 51 per interaction. On our 32.7 TFlops server, the measured
efficiency of averaged and peak performance reaches 35% and 48% on GPU,
respectively.

Although the optimizations and tests in this work are carried out on a small
server, PhotoNs-GPU is potentially powerful for massive parallelism on
supercomputers. Therefore we add a built-in IC generator and an on-the-fly
halo finder (Sun et al. 2020) into the code to save data storage. More on-the-
fly data processes will be included into the code in the future.

We only focus on the P2P operator of the algorithm on GPUs in this work. We
measure its performance and verify the feasibility for cosmological
simulations. But there still needs more considerations and improvements on the
other functional modules. For instance, M2L operator becomes more important
for small box simulations or in isolated boundary conditions. Similarly, FFT
in PM method could become dominant in extremely large simulations. Further
optimizations and considerations about them will be done in our future works.

###### Acknowledgements.

We acknowledge the support from National SKA Program of China (Grant No.
2020SKA0110401), National Natural Science Foundation of China (Grant No.
12033008) and K.C.Wong Education Foundation. WQ thanks the useful discussions
with J. Makino, M. Iwasawa, and L. Gao.

## References

  * Angulo et al. (2012) Angulo, R. E., Springel, V., White, S. D. M., et al. 2012, MNRAS, 426, 2046
  * Arafa et al. (2019) Arafa, Y., Badawy, A.-H., Chennupati, G., Santhi, N., & Eidenbenz, S. 2019, arXiv e-prints, arXiv:1905.08778
  * Bagla (2002) Bagla, J. S. 2002, Journal of Astrophysics and Astronomy, 23, 185
  * Barnes & Hut (1986) Barnes, J., & Hut, P. 1986, Nature, 324, 446
  * Bédorf et al. (2014) Bédorf, J., Gaburov, E., Fujii, M. S., et al. 2014, in Proceedings of the International Conference for High Performance Computing, 54
  * Cheng et al. (1999) Cheng, H., Greengard, L., & Rokhlin, V. 1999, Journal of Computational Physics, 155, 468
  * Cheng et al. (2020) Cheng, S., Yu, H.-R., Inman, D., et al. 2020, arXiv e-prints, arXiv:2003.03931
  * Crocce et al. (2006) Crocce, M., Pueblas, S., & Scoccimarro, R. 2006, MNRAS, 373, 369
  * Dehnen (2002) Dehnen, W. 2002, Journal of Computational Physics, 179, 27
  * Dehnen (2014) Dehnen, W. 2014, Computational Astrophysics and Cosmology, 1, 1
  * Dubinski et al. (2004) Dubinski, J., Kim, J., Park, C., & Humble, R. 2004, New Astronomy, 9, 111
  * Gaburov et al. (2010) Gaburov, E., Bédorf, J., & Zwart, S. P. 2010, Procedia Computer Science, 1, 1119, iCCS 2010
  * Greengard & Rokhlin (1987) Greengard, L., & Rokhlin, V. 1987, Journal of Computational Physics, 73, 325
  * Gumerov & Duraiswami (2008) Gumerov, N. A., & Duraiswami, R. 2008, Journal of Computational Physics, 227, 8290
  * Habib et al. (2016) Habib, S., Pope, A., Finkel, H., et al. 2016, New A, 42, 49
  * Hamada et al. (2009) Hamada, T., Narumi, T., Yokota, R., et al. 2009, in Proceedings of the Conference on High Performance Computing Networking, Storage and Analysis, SC ’09 (New York, NY, USA: Association for Computing Machinery)
  * Hamada & Nitadori (2010) Hamada, T., & Nitadori, K. 2010, 2010 ACM/IEEE International Conference for High Performance Computing, Networking, Storage and Analysis, 1
  * Ishiyama et al. (2009) Ishiyama, T., Fukushige, T., & Makino, J. 2009, PASJ, 61, 1319
  * Ishiyama et al. (2012) Ishiyama, T., Nitadori, K., & Makino, J. 2012, in Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis, SC ’12 (Los Alamitos, CA, USA: IEEE Computer Society Press), 5:1
  * Ishiyama et al. (2020) Ishiyama, T., Prada, F., Klypin, A. A., et al. 2020, arXiv e-prints, arXiv:2007.14720
  * Lewis et al. (2000) Lewis, A., Challinor, A., & Lasenby, A. 2000, ApJ, 538, 473
  * Makino (2004) Makino, J. 2004, Publications of the Astronomical Society of Japan, 56, 521
  * Makino & Taiji (1995) Makino, J., & Taiji, M. 1995, in Proceedings of the 1995 ACM/IEEE Conference on Supercomputing, Supercomputing ’95 (New York, NY, USA: Association for Computing Machinery), 63–es
  * Mead et al. (2016) Mead, A. J., Heymans, C., Lombriser, L., et al. 2016, MNRAS, 459, 1468
  * Nyland et al. (2009) Nyland, L., Harris, M., & Prins, J. 2009, in GPU Gems 3 (Addison–Wesley), 677–695
  * Planck Collaboration et al. (2020) Planck Collaboration, Aghanim, N., Akrami, Y., et al. 2020, A&A, 641, A6
  * Potter et al. (2017) Potter, D., Stadel, J., & Teyssier, R. 2017, Computational Astrophysics and Cosmology, 4, 2
  * Reed et al. (2007) Reed, D. S., Bower, R., Frenk, C. S., Jenkins, A., & Theuns, T. 2007, MNRAS, 374, 2
  * Springel (2005) Springel, V. 2005, MNRAS, 364, 1105
  * Springel (2015) Springel, V. 2015, N-GenIC: Cosmological structure initial conditions
  * Springel et al. (2020) Springel, V., Pakmor, R., Zier, O., & Reinecke, M. 2020, arXiv e-prints, arXiv:2010.03567
  * Springel et al. (2008) Springel, V., Wang, J., Vogelsberger, M., et al. 2008, MNRAS, 391, 1685
  * Stadel (2001) Stadel, J. G. 2001, Cosmological N-body simulations and their analysis, PhD thesis, UNIVERSITY OF WASHINGTON
  * Sun et al. (2020) Sun, S.-P., Liao, S.-H., Guo, Q., Wang, Q., & Gao, L. 2020, Research in Astronomy and Astrophysics, 20, 046
  * Takahashi et al. (2012) Takahashi, R., Sato, M., Nishimichi, T., Taruya, A., & Oguri, M. 2012, ApJ, 761, 152
  * Wang et al. (2020) Wang, J., Bose, S., Frenk, C. S., et al. 2020, Nature, 585, 39
  * Wang (2021) Wang, Q. 2021, Research in Astronomy and Astrophysics, 21, 003
  * Wang et al. (2018) Wang, Q., Cao, Z.-Y., Gao, L., et al. 2018, Research in Astronomy and Astrophysics, 18, 062
  * Xu (1995) Xu, G. 1995, ApJS, 98, 355
  * Yokota (2012) Yokota, R. 2012, arXiv e-prints, arXiv:1209.3516
  * Yokota & Barba (2011) Yokota, R., & Barba, L. A. 2011, arXiv:1110.2921
  * Yu et al. (2017) Yu, H.-R., Emberson, J. D., Inman, D., et al. 2017, Nature Astronomy, 1, 0143