Error-Bounded Learned Scientific Data Compression with Preservation of Derived Quantities

Abstract

Scientific applications continue to grow and produce extremely large amounts of data, which require efficient compression algorithms for long-term storage. Compression errors in scientific applications can have a deleterious impact on downstream processing. Thus, it is crucial to preserve all the “known” Quantities of Interest (QoI) during compression. To address this issue, most existing approaches guarantee the reconstruction error of the original data or primary data (PD), but cannot directly control the problem of preserving the QoI. In this work, we propose a physics-informed compression technique that is composed of two parts: (i) reduction of the PD with bounded errors and (ii) preservation of the QoI. In the first step, we combine tensor decompositions, autoencoders, product quantizers, and error-bounded lossy compressors to bound the reconstruction error at high levels of compression. In the second step, we use constraint satisfaction post-processing followed by quantization to preserve the QoI. To illustrate the challenges of reducing the reconstruction errors of the PD and QoI, we focus on simulation data generated by a large-scale fusion code, XGC, which can produce tens of petabytes in a single day. The results show that our approach can achieve a high compression amount while accurately preserving the QoI within scientifically acceptable bounds.

1. Introduction

As high-performance computing (HPC) enters the exascale era, scientific simulations are becoming capable of predicting physics phenomena with high accuracy, while generating extremely large volumes of data at high velocity. The increase in data volumes is rapidly outpacing the growth of storage and I/O capacities. For example, climate simulations [1] can generate 260 $\mathrm{TB}$ every 16 s when running large ensembles of high-fidelity 1 km × 1 km simulations of 15 years of climate. This necessitates the development of data compression algorithms for scientific applications that are reliable and can achieve a large reduction ratio. There have been successful approaches for traditional image and video compression [2], but they are not directly applicable to scientific data compression because of scientific post-processing requirements, which markedly differ from visual perception. Unlike image and video processing, where the dimensionality of data is two or three, scientific datasets are typically high-dimensional tensors and have underlying spatial and temporal structure.

Compressing scientific data is challenging because we need a significant reduction of the Primary Data (PD) with a guaranteed error bound, and we need to preserve the downstream Quantities Of Interest (QoI) as accurately as possible. The necessity of guaranteed error bounds of scientific data is due to the need for high accuracy after compression. For example, in turbulence simulations, the paths of turbulence can be estimated by properties such as latitude, longitude, pressure, and so on, which should be preserved to high fidelity during data reduction. This has to be achieved while ensuring that compression algorithms continue to cater to high-performance computing standards to reduce the time spent storing and retrieving the data. The mathematical relationships between the point-wise errors of PD and the errors of QoI are central to scientific data compression and create unique challenges for the reduction community. For this reason, we focus on simultaneously achieving lossy data compression while preserving the QoI to within specific error tolerances.

Data reduction for scientific applications has been extensively studied. Error-bounded lossy compression techniques are considered to be one of the most successful for these applications. Compared to lossless compressors [3,4] that achieve limited compression [5]), error-bounded lossy compressors can significantly reduce the data and, moreover, achieve high fidelity of the reconstructed data according to the user’s required data distortion bounds. The most well-known and state-of-the-art error-bounded lossy compressors in the literature are based on prediction (SZ [6,7,8]), block transform (ZFP [9]), and multilevel decomposition (MGARD [10,11,12]). SZ is designed based on data prediction models. It first performs a preprocessing step such as linearization of multi-dimensional data and predicts each data point based on its neighboring values using a Lorenzo predictor [13]. ZFP is based on orthogonal block transforms and embedded coding. It breaks the data into a set of blocks (4 × 4 × 4 for 3D data) and then applies an orthogonal block transformation to decorrelate the data. The generated coefficients are encoded to produce compressed data. Most error-bounded lossy compression algorithms such as SZ [6,7,8] and ZFP [9] just focus on preserving point-wise values and provide guaranteed error control on the primary data. However, this does not ensure QoI preservation. Derived quantities are critical to accurate post-analysis. Inconsistent moments and energy in XGC applications that simulate gyrokinetic particle motion, for example, cause totally different fusion simulation analysis. MGARD [10,11,12] is currently the leading error-controlled lossy compression technique with QoI preservation. MGARD can bound the errors of linear QoI by using a non-uniform multilevel subgrid algorithm and providing a mathematical error analysis. Despite this, MGARD cannot currently control the reconstruction errors of non-linear or complex QoI that are often necessary for scientific discovery.

In the machine learning (ML) community, there have been many approaches that successfully leverage deep learning for image compression [14,15,16]. The AutoEncoder (AE) technique is a standard deep learning model for data compression. It has an encoder–decoder architecture and is trained to minimize the point-wise error between the PD and the output of the decoder. While numerous AE variations have been developed, the application of ML techniques to scientific datasets is in its infancy. A fully connected AE has been studied for scientific data compression [17]. The network consists of seven fully connected layers: three layers for an encoder and decoder and a bottleneck layer. The model is, however, applied only on small-scale 1D scientific data. For a specific scientific application, a convolutional AE is developed to compress turbulent flow simulation data [18]. The model includes skip connections with residual block structure [19] to reduce the amount of information in both the encoder and the decoder. This enables the network to effectively learn the complex turbulent flow features. However, these approaches do not guarantee the error of reconstructed data.

There have been approaches that improve the efficiency of ML models by training under constraints. Although they are not studied for data reduction, these are closely related to QoI preservation, since QoI preservation can be treated as constraint satisfaction. Physics-informed neural networks [20] have recently received a lot of attention. Many physics simulations require powerful computational resources and are quite expensive to execute. Physics-informed neural networks are investigated to predict simulation results, and therefore can simplify and reduce the cost of entire simulations. Using the fact that many fields of science have closed-form constraints, these are enforced by incorporating the equations into the architecture of neural networks. Therefore, the approaches are limited to specific physics systems such as Reynolds-averaged Navier–Stokes equations [21,22] and advection equations [23,24]. The merits of the models have been demonstrated in solving a number of forward and inverse problems governed by partial differential equations. There are approaches for incorporating QoI using constrained optimization [25,26] that are closely related to our work. While augmented Lagrangians have been used to incorporate physical constraints into neural networks [27,28], some methods merely use soft penalties which violate QoI error tolerances [29,30]. These simulation approaches have not yet been used for compression, which is the focus of our present work.

In this paper, we propose a novel, hybrid, algorithmic framework to address these issues. We combine tensor decompositions, ML autoencoders, error-bounded lossy compressors, constrained optimization, and product quantization to achieve high compression ratios and fidelity of both the PD and the QoI. The approach is general and broadly applicable to any tensor dataset with moment-based constraint QoI preservation (and error-bounded reconstruction). We focus our attention on a large-scale fusion code—XGC [31], which produces petabytes of five-dimensional data in a day. The key data structures for the target applications are arranged in the form of tensors on a structured, irregular, or hybrid grid. Our methods allow suitable trade-offs between the accuracy and compression ratio for such applications, and comprise the following techniques:

• Transformation: We perform tensor decompositions as a preprocessing step to project the tensors within coherent segments into a more suitable space for reduction. We obtain core tensors (or coefficients) and bases. The tensor basis functions are similar to wavelets and/or local spatial frequencies, and are therefore good at expressing the correlations among the tensors.
• Autoencoder-based compression: AEs are trained on the transformed tensor coefficient space. Since tensor basis functions can take care of spatial correlation, AEs encode and decode tensor basis coefficients. We use AEs to obtain the encoded coefficients using standard loss measures and employ a quantization scheme to store the encoded coefficients.
• Error guarantees on PD: Error guarantees in the form of bounded reconstruction error are achieved by selectively applying error-bounded compressors on the residuals. (The residual is the difference between the reconstructed data of the AE and the PD). The residuals are processed with error-bounded lossy compressors to guarantee bounds on the overall PD error.
• QoI satisfaction: Preservation of QoI is addressed using constraint satisfaction on the reconstructed PD. After residual post-processing, we perform constrained optimization while minimizing the changes to the PD reconstruction. The parameters produced as a result of the QoI satisfaction step are quantized and stored with the encoded coefficients and residuals.

Our experimental results show that our framework can achieve high compression ratios when the sizes of the latent space and the quantization dictionary are controlled. Additionally, we show that selective residual post-processing prevents significant distortion of data even at high compression ratios. Furthermore, the constraint satisfaction step can directly manage QoI preservation with tolerable errors (up to two to three orders of magnitude in compression ratios). We compare our framework to MGARD, which is currently the leading error-bounded lossy compressor with QoI preservation, and show that our approach can achieve a significantly better fidelity of PD and QoI.

The rest of the paper is organized as follows. In Section 2, we discuss the X-point included Gyrokinetics Code (XGC) dataset in detail. In Section 3, our framework is presented. The experimental results are discussed in Section 4. Section 5 concludes the paper.

2. XGC Application

XGC [31,32] is one of the most important first principles fusion codes used to understand the physics in the edge of a fusion experiment, and is part of the DOE Exascale Computing Program. XGC uses a Particle-In-Cell (PIC) technique and evolves a massively large number of particles based on a five-dimensional gyrokinetic equation of motion, three in physical space and two in the velocity space. In XGC, particles are moved according to the physics equations in a time-stepping manner, and then regularly summarized to a discretized 3D space by planes and meshes for solving electric and magnetic field equations. One of the large datasets gathered in the mesh is a particle distribution function, known as a velocity histogram, called F-data (Figure 1).

XGC scientists typically reduce the trillions of particles to build F-data and examine them at every timestep. The size of F-data ranges from hundreds of GB to tens of TBs, depending on the simulation scale. To study the next-generation fusion devices, such as ITER and DEMO, data compression of F-data is necessary for multiple purposes, such as saving memory/disk spaces and analysis speed-up. However, XGC scientists are eager to ensure that various physics properties derived from the distribution, such as mass, canonical angular momentum, and energy, are conserved. Therefore, reducing compression error is needed for this XGC dataset.

More specifically, for F-data compression, XGC scientists are concerned with QoI preservation as well as compression errors. We investigate six physics quantities specified by XGC scientists. The QoI we consider are density n, parallel flow $u_{\Vert }$, perpendicular temperature $T_{\perp }$, parallel temperature $T_{\Vert }$, flux surface averaged density $n_0^{avg}$, and flux surface averaged temperature $T_0^{avg}$. In the XGC dataset, each mesh node x has a 2D velocity histogram $\mathit{f}$. The density n and parallel flow $u_{\Vert }$ are computed by integrating $\mathit{f}$ with volume weights and with the parallel velocity, respectively:

$$\begin{array}{ccc}\hfill n\left(x,t\right)& =& \int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert },\hfill \end{array}$$

(1)

$$\begin{array}{ccc}\hfill u_{\Vert }\left(x,t\right)& =& \frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)v_{\Vert }\left(\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)},\hfill \end{array}$$

(2)

where t represents the timestep, $\mathrm{vol}\left(x,\mathbf{v}\right)$ is the volume of the $x\mathrm{th}$ mesh node, and $v_{\Vert }$ is the parallel velocity. $T_{\perp }$ and $T_{\Vert }$ are the integrals of perpendicular and parallel kinetic energy, $m{v}^2$,

$$\begin{array}{ccc}\hfill T_{\perp }\left(x,t\right)& =& \frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{v}_{\perp }^2\left(\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)},\hfill \end{array}$$

(3)

$$\begin{array}{ccc}\hfill T_{\Vert }\left(x,t\right)& =& \frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathbf{v}\right)m{\left(v_{\Vert }\left(\mathit{v}\right)-u_{\Vert }\left(x,t\right)\right)}^{2}d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)},\hfill \end{array}$$

(4)

where m is the mass of atomic particles corresponding to $\mathit{f}$ and $v_{\perp }$ is the perpendicular velocity. The flux surface averaged density $n_0^{avg}$ and flux surface averaged temperature $T_0^{avg}$ are computed by averaging the density and temperature over a thin volume between toroidal magnetic flux surface contours. For simplicity, we omit the detailed equations for the flux surface averaged quantities [33]. Except for density, these QoI are non-linear on $\mathit{f}$ and are obtained via post-analysis of the data. In summary, the QoI described here result in four derived quantities (density, parallel flow, perpendicular, and parallel temperatures) which can be used in a constraint satisfaction procedure connecting the primary data ($\mathit{f}$) and the QoI. This is precisely the route taken in this paper.

3. The Compression Framework

A high-level description of our reduction/compression pipeline is described in Figure 2. The framework mainly comprises two pipelines—Primary Data (PD) reduction and Quantities of Interest (QoI) satisfaction. For PD reduction, we combine Tucker decomposition, an autoencoder (AE), and a product quantizer (PQ) for effective global compression. An error-bounded lossy compressor is used for residual post-processing after global learning to provide the guaranteed error bound on PD. For QoI satisfaction, we perform local constrained optimization on the reconstructed data from the PD reduction pipeline. This step enables us to directly control QoI preservation. We use another product quantizer to encode the Lagrange multipliers produced by this step. Each of these pipelines is now described in detail in the remainder of this section.

3.1. Reduction of the Primary Data

Basis functions can span the spectrum from simple principal/independent component analysis (PCA/ICA) to overcomplete dictionaries. Here, we use the Tucker decomposition [34] to obtain a simple factorization. Each tensor is thereby converted into a set of coefficients, and subsequently, we use an ML autoencoder (AE) with the mean squared error (MSE) measure on the converted Tucker coefficients. The output of the decoder is projected back into the original tensor space to produce a reconstruction. The residual of the reconstruction is then analyzed using an error-bounded lossy compressor to ensure that the PD are within prespecified error bounds. In our approach, we process residuals using MGARD. Note that to get an actual codec, we have to quantize the latent space coefficients in the AE. We use PQ to obtain a codec, with the result being a set of integers for each tensor. Varying the latent space cardinality of the AE yields different compression ratios.

3.1.1. Data Transformation

The Tucker decomposition [34] is widely used in data reduction by considering a subset of coefficients that carry most of the information in the dataset and eliminate lower-rank coefficients, which typically has no adverse effect on tensor reconstruction. The original Tucker decomposition can be viewed as a higher-order singular-value decomposition (HOSVD), and we do not consider more complex variants such as HOOI [35] in this paper. It decomposes a tensor T of order N into a core tensor with the same order while generating N orthonormal bases. If we consider a dataset as a tensor $T\in {\mathbb{R}}^{d_1\times d_2\times d_3}$ of order 3, its Tucker decomposition is $T=\mathcal{T}{\times }_1{U}^{\left(1\right)}{\times }_2{U}^{\left(2\right)}{\times }_3{U}^{\left(3\right)}$, where $\mathcal{T}\in {\mathbb{R}}^{d_1\times d_2\times d_3}$ is the core tensor which, as a lower rank approximation of T, gives a representation basis for it. Unitary or factor matrices contain basis vectors that help project T into the basis space $\mathcal{T}$.

3.1.2. The Autoencoder

An autoencoder is a neural network that is trained to reproduce the input at its output while going through a bottleneck. Internally, it has hidden layers and a final bottleneck hidden layer h that describes a code used to represent the input. The network may be viewed as consisting of two parts: an encoder function $h=f\left(x\right)$ and a decoder that produces a reconstruction $r=g\left(h\right)$. If an autoencoder succeeds in simply learning to set $g\left(f\left(x\right)\right)=x$ everywhere, then it is not especially useful. Instead, autoencoders are designed to be unable to learn to copy perfectly. Usually, they are restricted in ways that allow them to copy only approximately. The model is forced to prioritize which aspects of the input should be copied, and therefore it often learns useful properties of the data.

We obtain the encoded coefficients (the latent space) using standard loss measures in the AE. Consequently, the PD errors are minimized by the AE while learning a good codec. Error guarantees in the form of bounded reconstruction error (at a given level of reduction) on the PD will be achieved by applying error-bounded lossy compressors to the residuals. The loss function of the AE network compares the coefficients of the Tucker decomposition with the reconstruction from the bottleneck latent space:

$$\mathrm{Loss}=\mathrm{MSE}\left(X,\widehat{X}\right),$$

(5)

where X is the input Tucker coefficients and $\widehat{X}$ is the reconstructed coefficients.

We use an AE which constrains the number of nodes present in the hidden layers of the network, limiting the amount of information that can flow through the network. By penalizing the network according to the reconstruction error, our model can learn the most important attributes of the input data and how to best reconstruct the original input from an encoded state. Ideally, this encoding will learn and describe latent attributes of the input data. Therefore, the bottleneck layer contains the data compression.

3.1.3. Quantization of the Latent Space

Merely obtaining a latent space is not sufficient to obtain good compression. Entropy encoding of the latent space—also known as quantization—is required to efficiently store the bottleneck layer coefficients. Quantization is therefore an important step in our current pipeline. We use PQ [36] to convert the floating-point latent space coefficients into integers. PQ is a vector quantization [37] technique that breaks a high dimensional vector into numerous lower dimensional vectors and quantizes them independently using a simple clustering algorithm such as K-means. This effectively creates a Cartesian product of smaller subspaces to represent the original vector space (hence the term ‘product’ quantizer). Although very simple, the advantage of this approach is that the number of vectors that can be represented accurately by the codebook can be quite large with respect to the size of the codebook. The size of the quantized data depends on the number of bits that a PQ dictionary uses to represent codewords in the codebook (PQ bits) and the dimension of input vectors. Assume we have an L-dimensional vector consisting of floating-point numbers from an AE encoder, and let B be the PQ bit encoding. Then, each floating-point number can be represented by B bits, and therefore the size of the quantized vector is $L\times B$ bits.

3.1.4. Residual Post-Processing

Error guarantees are often required in scientific data compression and this is a significant departure from image and video compression. Here, we require that the error (${\ell }_{\infty }$ or an ${\ell }_2$ measure) does not exceed a prespecified bound. For standalone images or tensors, a straightforward approach that guarantees this is using truncated basis function representations based on MGARD [10,11,12]: here, the input domain is recursively subdivided with a fixed basis function used within each subdomain to represent the data. If the fixed basis is unable to meet the error requirement, it is abandoned and the subdomain further divided. In this way, error guarantees are met and a subdivision procedure hammered out per image. Depending on the methods of dividing grids, MGARD has two versions. One is MGARD uniform, which only focuses on guaranteeing PD error bounds and achieving the best possible compression without considering QoI preservation. It divides grids and assigns quantization bins uniformly. The other is MGARD non-uniform, which considers both guaranteed PD error bounds and the preservation of linear QoI. It uses an additional s parameter that assigns different weights and quantization bins depending on the QoI. Hence, it divides grids non-uniformly. MGARD non-uniform achieves a lower compression amount compared to MGARD uniform with the same error bound, since it uses precise quantization bins for some coefficients due to QoI preservation. In our approach, we apply MGARD uniform on residuals—the difference between the output of an AE and the PD—instead of applying MGARD on the PD. Since we have an additional step for QoI preservation, we choose MGARD uniform for the residual post-processor. Furthermore, we use the reconstructed residuals selectively based on an error level of each instance. We need to store not only the compressed data from an AE, but the compressed residuals as well, which affects the overall compression ratio. Therefore, we set error thresholds to reduce the number of residuals being used. If an error of an instance is larger than the threshold, we process the residual of the instance with MGARD so that the error does not exceed the threshold.

3.2. Preservation of the Derived Quantities

The QoI we seek to preserve in XGC are essentially moments and energy properties of the PD. Ideally, we can attempt to preserve the QoI using constraint satisfaction during AE training. While this is potentially an interesting avenue for exploration, in this work we restrict ourselves to performing constraint satisfaction on the reconstructed PD after residual post-processing. The different moment constraints are first set up below, and we then proceed with a standard Lagrangian constraint satisfaction approach to QoI preservation. Since the loss function we use (between the reconstructed PD and the PD obtained after constraint satisfaction) is convex and the constraints linear, we employ a standard dual maximization approach using Newton optimization (since constraint cardinality is $\mathcal{O}\left(1\right)$). Subsequently, after obtaining the Lagrange parameters for the QoI, we use PQ for vector quantization of these parameters with care taken to satisfy QoI error thresholds. In the experiments, we justify the Lagrange parameter quantization choices. These steps are explained below.

3.2.1. Constraint Satisfaction

QoI preservation is crucial to the XGC application, and thus, QoI errors need to be within tolerance limits after decompression. We first define a general QoI preservation problem and then apply it to the XGC application. Let ${\mathit{f}}^0$ be the decompressed output from any compressor that has a bounded error on PD and $\mathit{f}$ be the output that we are going to obtain after the constraint satisfaction step. We consider a general equality-constrained optimization problem [25,26],   

$$\begin{array}{cc}\hfill & \mathrm{minimize}d\left(\mathit{f},{\mathit{f}}^0\right)\hfill \\ \hfill & \mathrm{subject}\mathrm{to}h\left(\mathit{f}\right)=0,\hfill \end{array}$$

(6)

where $d\left(\mathit{f},{\mathit{f}}^0\right)$ is a distance measure between $\mathit{f}$ and ${\mathit{f}}^0$, and $h:{\mathbb{R}}^n\to {\mathbb{R}}^m$ computes the m number of QoI that need to be preserved. This approach enables us to obtain the updated output $\mathit{f}$ that satisfies QoI preservation while minimizing the distance between $\mathit{f}$ and ${\mathit{f}}^0$, so that PD error can be still bounded within tolerance. The Lagrangian is

$$\mathcal{L}\left(\mathit{f},\mathit{\lambda }\right)=d\left(\mathit{f},{\mathit{f}}^0\right)+{\mathit{\lambda }}^{T}h\left(\mathit{f}\right),$$

(7)

where $\mathit{\lambda }=\left[{\mathit{\lambda }}_1,\cdots ,{\mathit{\lambda }}_m\right]$ is a Lagrange multiplier vector.

For constraint satisfaction of the XGC application, we focus our attention on the four QoI—density n, parallel flow $u_{\Vert }$, perpendicular temperature $T_{\perp }$, parallel temperature $T_{\Vert }$ with flux surface averaged density $n_0^{avg}$ and flux surface averaged temperature $T_0^{avg}$ being averages over the space. Except for density, the QoI are non-linear, and so we make several assumptions and modifications to QoI computations in Equations (1)–(4) to simplify the problem. First, we rewrite $T_{\Vert }$ in Equation (4) to remove $u_{\Vert }$ in the integral. Since $\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)}=1$, we can use $\frac{f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)}{n\left(x,t\right)}$ as the probability density function. We will use expected value notation which considerably simplifies the development. For example,

$$\mathrm{E}\left(v_{\Vert }\right)=\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{v}_{\Vert }\left(\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)}.$$

(8)

Then, $T_{\Vert }$ can be written as follows:

$$\begin{array}{cc}\hfill T_{\Vert }(x,t)& =\frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{\left(v_{\Vert }\left(\mathit{v}\right)-u_{\Vert }\left(x,t\right)\right)}^{2}d{v}_{\perp }d{v}_{\Vert }}{n\left(x,t\right)}\hfill \\ & =\frac{1}{2}\mathrm{E}\left[m{\left(v_{\Vert }\left(\mathit{v}\right)-u_{\Vert }\left(x,t\right)\right)}^2\right]\hfill \\ & =\frac{m}{2}\left[\mathrm{E}\left(v_{\Vert }^2\left(\mathit{v}\right)\right)-u_{\Vert }\left(x\right)\mathrm{E}\left(v_{\Vert }\left(\mathit{v}\right)\right)+u_{\Vert }^2\left(x,t\right)\right]\hfill \\ & =\frac{m}{2}\left[\mathrm{E}\left(v_{\Vert }^2\left(\mathit{v}\right)\right)-u_{\Vert }^2\left(x,t\right)\right],\hfill \end{array}$$

(9)

where the expected value notation $\mathrm{E}(·)$ is deployed for the sake of convenience, and furthermore, $\mathrm{E}\left(v_{\Vert }\left(\mathit{v}\right)\right)=u_{\Vert }(x,t)$ following Equation (2). Therefore, we can eliminate $u_{\Vert }$ from the integral in the $T_{\Vert }$ computation. We get

$$R_{\Vert }\left(x,t\right)=T_{\Vert }\left(x,t\right)+\frac{m}{2}{u}_{\Vert }^2\left(x,t\right)=\frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{v}_{\Vert }^2\left(\mathit{v}\right)d{\mathit{v}}_{\perp }d{\mathit{v}}_{\Vert }}{n\left(x,t\right)}$$

(10)

where we have defined a new derived quantity $R_{\Vert }$:

$$R_{\Vert }\left(x,t\right)\equiv T_{\Vert }\left(x,t\right)+\frac{m}{2}{u}_{\Vert }^2\left(x,t\right).$$

(11)

Secondly, we use the true (correct) density $n^c\left(x,t\right)$ of the original PD to compute the non-linear QoI, so that we can treat the density as a constant for each node x, thereby making the QoI computation linear. We hence obtain the four independent constraints,

$$\begin{array}{ccc}\hfill n^c\left(x,t\right)& =& \int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert },\hfill \end{array}$$

(12)

$$\begin{array}{ccc}\hfill u_{\Vert }^c\left(x,t\right)& =& \frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)v_{\Vert }\left(\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n^c\left(x,t\right)},\hfill \end{array}$$

(13)

$$\begin{array}{ccc}\hfill T_{\perp }^c\left(x,t\right)& =& \frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{v}_{\perp }^2\left(\mathit{v}\right)d{v}_{\perp }d{v}_{\Vert }}{n^c\left(x,t\right)},\hfill \end{array}$$

(14)

$$\begin{array}{ccc}\hfill R_{\Vert }^c\left(x,t\right)& =& \frac{1}{2}\frac{\int f\left(x,\mathit{v},t\right)\mathrm{vol}\left(x,\mathit{v}\right)m{v}_{\Vert }^2\left(\mathit{v}\right)d{\mathit{v}}_{\perp }d{\mathit{v}}_{\Vert }}{n^c\left(x,t\right)},\hfill \end{array}$$

(15)

where the superscript c indicates the true (correct) value.

Now, we perform constrained optimization. We rewrite the constraints in discrete form since XGC dataset is discretized. We consider a 2D velocity histogram $\mathit{f}$ at a mesh node x from a single timestep for notational brevity. The discretized constraints are

$$\begin{array}{ccc}\hfill n^c& =& \sum _{i,j}{f}_{ij}{\mathrm{vol}}_{ij},\hfill \end{array}$$

(16)

$$\begin{array}{ccc}\hfill u_{\Vert }^c{n}^c& =& \sum _{i,j}{f}_{ij}{\mathrm{vol}}_{ij}{v}_j^{\Vert },\hfill \end{array}$$

(17)

$$\begin{array}{ccc}\hfill T_{\perp }^c{n}^c& =& \frac{1}{2}m\sum _{i,j}{f}_{ij}{\mathrm{vol}}_{ij}{\left(v_i^{\perp }\right)}^2,\hfill \end{array}$$

(18)

$$\begin{array}{ccc}\hfill R_{\Vert }^c{n}^c& =& \frac{1}{2}m\sum _{i,j}{f}_{ij}{\mathrm{vol}}_{ij}{\left(v_j^{\Vert }\right)}^2,\hfill \end{array}$$

(19)

where $v_j^{\Vert }$ and $v_j^{\perp }$ indicate parallel velocity $v_{\Vert }\left(\mathit{v}\right)$ and perpendicular velocity $v_{\perp }\left(\mathit{v}\right)$ respectively. For the objective function in Equation (7), we choose the Bregman divergence for the $xlogx$ entropy function [38,39] since all the velocity histogram values are positive and the range of histograms is fairly large $({10}^6$∼${10}^{17})$. Any suitable convex error measure can be used. Since the PD error is optimized in the AE and residual post-processing, we want to obtain the updated reconstruction with constraint satisfaction while minimizing the changes in the PD reconstruction. Let the reconstructed histogram of AETM or any compressor be ${\mathit{f}}^0$. We use Rectified Linear Units [40] (ReLU) with a small positive value to remove non-positive values, as a compressor might produce some non-positive values due to the large range of histograms.

$$f_{ij}^0=\left\{\begin{array}{c}{f}_{ij}^0+ϵ\mathrm{if}{f}_{ij}^0>0,\hfill \\ ϵ\mathrm{otherwise},\hfill \end{array}\right.$$

(20)

where $ϵ$ is a very small positive value compared to the histograms, so it does not affect overall compression accuracy. Let the updated reconstruction with constraint satisfaction be $\mathit{f}$. The Lagrangian for the XGC application is

$$\begin{array}{c}\hfill \mathcal{L}\left(\mathit{f},\mathit{\lambda }\right)=\sum _{i,j}\left(f_{ij}log\frac{f_{ij}}{f_{ij}^0}-f_{ij}+f_{ij}^0\right)+{\mathit{\lambda }}_1\left(\sum _{i,j}{f}_{ij}{\mathrm{V}}_{ij}^1-n^c\right)+{\mathit{\lambda }}_2\left(\sum _{i,j}{f}_{ij}{\mathrm{V}}_{ij}^2-u_{\Vert }^c{n}^c\right)\\ \hfill +{\mathit{\lambda }}_3\left(\sum _{i,j}{f}_{ij}{\mathrm{V}}_{ij}^3-T_{\perp }^c{n}^c\right)+{\mathit{\lambda }}_4\left(\sum _{i,j}{f}_{ij}{\mathrm{V}}_{ij}^4-R_{\Vert }^c{n}^c\right),\end{array}$$

(21)

where $\mathit{\lambda }=\left[{\mathit{\lambda }}_1,\cdots ,{\mathit{\lambda }}_4\right]$, ${\mathrm{V}}_{ij}^1={\mathrm{vol}}_{ij}$, ${\mathrm{V}}_{ij}^2={\mathrm{vol}}_{ij}{v}_j^{\Vert }$, ${\mathrm{V}}_{ij}^3=\frac{1}{2}m{\mathrm{vol}}_{ij}{\left(v_i^{\perp }\right)}^2$, and ${\mathrm{V}}_{ij}^4=\frac{1}{2}m{\mathrm{vol}}_{ij}{\left(v_j^{\Vert }\right)}^2$.

Since the objective function is convex, we set the gradient of $\mathcal{L}\left(\mathit{f},\mathit{\lambda }\right)$ with regard to $\mathit{f}$ to zero to find the minimum:

$$\frac{\partial \mathcal{L}\left(\mathit{f},\mathit{\lambda }\right)}{\partial f_{ij}}=log\frac{f_{ij}}{f_{ij}^0}+{\mathit{\lambda }}_1{\mathrm{V}}_{ij}^1+{\mathit{\lambda }}_2{\mathrm{V}}_{ij}^2+{\mathit{\lambda }}_3{\mathrm{V}}_{ij}^3+{\mathit{\lambda }}_4{\mathrm{V}}_{ij}^4=0,$$

(22)

from which we obtain

$$f_{ij}=f_{ij}^{0}exp\left(-K_{ij}\right),$$

(23)

where $K_{ij}={\mathit{\lambda }}_1{\mathrm{V}}_{ij}^1+{\mathit{\lambda }}_2{\mathrm{V}}_{ij}^2+{\mathit{\lambda }}_3{\mathrm{V}}_{ij}^3+{\mathit{\lambda }}_4{\mathrm{V}}_{ij}^4$. Using the fact that ${\mathit{f}}^T{\nabla }_{\mathit{f}}\mathcal{L}\left(\mathit{f},\mathit{\lambda }\right)=0$ at the minimum, we obtain the dual $q\left(\mathit{\lambda }\right)$ of Equation (21)

$$\begin{array}{cc}\hfill & q\left(\mathit{\lambda }\right)=\sum _{i,j}\left(-f_{ij}+f_{ij}^0\right)-{\mathit{\lambda }}_1{n}^c-{\mathit{\lambda }}_2{u}_{\Vert }^c{n}^c-{\mathit{\lambda }}_3{T}_{\perp }^c{n}^c-{\mathit{\lambda }}_4{R}_{\Vert }^c{n}^c\hfill \\ \hfill & =\sum _{i,j}\left(-f_{ij}^{0}exp\left(-K_{ij}\right)+f_{ij}^0\right)-{\mathit{\lambda }}_1{n}^c-{\mathit{\lambda }}_2{u}_{\Vert }^c{n}^c-{\mathit{\lambda }}_3{T}_{\perp }^c{n}^c-{\mathit{\lambda }}_4{R}_{\Vert }^c{n}^c.\hfill \end{array}$$

(24)

Finally, our objective function is the negative of the dual. We therefore seek to minimize

$$\underset{\mathit{\lambda }}{min}-q\left(\mathit{\lambda }\right).$$

(25)

To obtain the optimal $\mathit{\lambda }$, we use Newton’s method [41], which finds the minimum via a standard iterative procedure using the gradient $g\left(\mathit{\lambda }\right)$ and the Hessian $H\left(\mathit{\lambda }\right)$ of the objective function. We update $\mathit{\lambda }$ by

$${\mathit{\lambda }}^{k+1}={\mathit{\lambda }}^k-H^{-1}\left({\mathit{\lambda }}^{\mathit{k}}\right)g\left({\mathit{\lambda }}^{\mathit{k}}\right),$$

(26)

where the superscript k indicates the kth update. After finding the optimal $\mathit{\lambda }$, we obtain the updated reconstruction using Equation (23).

3.2.2. Quantization of the Lagrange Parameters

The constraint satisfaction step produces four Lagrange multipliers per $39\times 39$ histogram, and therefore we need to store four additional floating-point numbers per node. Since the parameters ${\mathrm{V}}^1,\dots ,{\mathrm{V}}^4$ can be obtained from the metadata to which all users have access and represent the common geometry across all timesteps in the XGC application, we do not consider those parameters as additional costs. To improve storage efficiency, we apply PQ to quantize the Lagrange multipliers. The number of sub-quantizers is set as equal to the number of Lagrange multipliers. To quantize the Lagrange parameters with $B_L$ bits, each sub-quantizer of PQ can use $2^{B_L}$ centroids to represent each Lagrange parameter, and thus the number of entries in the entire PQ dictionary is $4\times 2^{B_L}$. We store both the quantized Lagrange parameter indices and the PQ dictionary. The $B_L$ parameter can be set to achieve a suitable trade-off between the overall compression ratio and the QoI error.

Importantly, the change of PD error after the constraint satisfaction step needs to be considered. Suppose the maximum quantization error of Lagrange multipliers is ${ϵ}_L$. From Equation (23), $\frac{f_{ij}}{f_{ij}^0}$ is bounded:

$$\frac{f_{ij}}{f_{ij}^0}\le exp\left(-K_{ij}+\left|{ϵ}_L\left({\mathrm{V}}_{ij}^1+{\mathrm{V}}_{ij}^2+{\mathrm{V}}_{ij}^3+{\mathrm{V}}_{ij}^4\right)\right|\right).$$

(27)

Before QoI optimization, we have a reconstructed histogram ${\mathbf{f}}^0$ that has a PD error bound ${ϵ}_{f_0}$. Therefore, each entry of the histogram is also bounded:

$$\left|f_{ij}^0-f_{ij}^p\right|\le {ϵ}_{f_0},$$

(28)

where $f_{ij}^p$ indicates the original PD. Let $E_{ij}$ be $-K_{ij}+\left|{ϵ}_L\right({\mathrm{V}}_{ij}^1+{\mathrm{V}}_{ij}^2$$+{\mathrm{V}}_{ij}^3+{\mathrm{V}}_{ij}^4\left)\right|$. From Equations (27) and (28), the PD error bound is

$$\begin{array}{c}\hfill \left|f_{ij}-f_{ij}^p\right|\le \left|f_{ij}^{0}exp{E}_{ij}-f_{ij}^p\right|=\left|\left(f_{ij}^0-f_{ij}^p\right)exp{E}_{ij}+f_{ij}^{p}exp{E}_{ij}-f_{ij}^p\right|\\ \hfill \le {ϵ}_{f_0}exp{E}_{ij}+\left|f_{ij}^p\left(exp{E}_{ij}-1\right)\right|.\end{array}$$

(29)

Therefore, users can anticipate the increase in the error bound due to the constraint satisfaction step and choose the number of bits for Lagrange parameter quantization, with this step simultaneously impacting the PD and QoI errors.

4. Experimental Results

In this section, we present experimental results on data collected from XGC simulations. Compressing all the data generated from XGC simulations is unrealistic due to extremely large storage requirements. To give a sense of scale, a simulation model of the ITER-scale [42] problem typically contains trillions of particles, runs for thousands of timesteps, and can produce over 200 petabytes of data each day. Hence, we compress discretized data of XGC simulations. We first apply our algorithm to the 420th timestep of XGC data, which consists of eight cross-sections and 16,395 nodes per cross-section. Each node represents a $39\times 39$ 2D velocity histogram of electrons. Other XGC time point simulation data are deployed as test data, with results shown using the pretrained machine. All the experiments are performed on Summit nodes in the Oak Ridge National Laboratory. The Summit nodes contain POWER9 processors with 512 GB of DRAM and NVIDIA TESLA V100 GPUs.

4.1. Error Metrics

We use the following two inter-related and key metrics to assess the usefulness of our algorithms and software:

• NRMSE. Since the XGC dataset represents the velocity of electrons, the scale of the data is extremely large (∼10${}^{17}$). Therefore, instead of using the Root Mean Squared Error (RMSE), we measure the error using the Normalized Root Mean Square Error (NRMSE), defined as follows:

  $$\mathrm{NRMSE}\left(x,\widehat{x}\right)=\frac{\sqrt{{\sum }_{i=0}^N{\left(x_i-{\widehat{x}}_i\right)}^2/N}}{max\left(x\right)-min\left(x\right)},$$

  (30)

  where x is the primary or original data (PD), $\widehat{x}$ is the reconstructed data, and N is the number of nodes. We will ensure that error bounds for PD are met. Some of these may be tight error bounds and some of them may be probabilistic and vary within the PD.
• Amount of Data Reduction. The key feature of the data reduction algorithms is the amount of reduction achieved. We measure the storage sizes of the compressed data and residuals, and Lagrange multipliers as well as the size of the models, Tucker bases, and the dictionaries of product quantizers. Once the machine architecture is fixed (after training), the decoder size becomes a fixed compression cost in terms of the size, and it should be included in the overall size when computing the compression levels. Tucker bases are fixed compression costs, since the decoded coefficients need to be projected back into the original data space. The product quantizer’s dictionaries for both compressed data and Lagrange multipliers are also fixed compression costs which need to be included in all calculations. Considering all the costs that are needed for the reconstruction, we compute the compression ratio (CR) defined as shown below:

  $$\mathrm{CR}=\frac{\mathrm{Sizeof}\left(\mathrm{Primary}\mathrm{Data}\right)}{\mathrm{Sizeof}\left(\mathrm{Encoded}\mathrm{Data}\right)+\mathrm{Sizeof}\left(\mathrm{Model}\right)}.$$

  (31)
  The $\mathrm{Encoded}\mathrm{Data}$ in Equation (31) includes encoded coefficients of the AE, compressed residuals and quantized Lagrange multipliers. The $\mathrm{Model}$ encompasses the models of AE and MGARD, Tucker bases, and PQ dictionaries for the AE and Lagrange multipliers.

We discuss trade-offs of combinations of these two metrics for the XGC dataset so that XGC scientists can choose parameters based on their requirements.

4.2. Baseline

MGARD is the leading error-bounded lossy compression technique that considers both error bounds on PD and QoI preservation. We compare our framework to the two versions of MGARD, as described in Section 3.1.4. MGARD uniform only focuses on achieving the best compression within PD error bounds and does not consider QoI preservation. MGARD non-uniform tries to reduce the errors of QoI with guaranteed PD error bounds. We set the s parameter of MGARD non-uniform to be $-1$, as it is the best setting for the XGC dataset for QoI preservation. MGARD uniform is much better than MGARD non-uniform in terms of PD errors, whereas MGARD non-uniform performs better than MGARD uniform in terms of maintaining the QoI. Therefore, we use MGARD uniform for PD error comparison and MGARD non-uniform for QoI error comparison. The abbreviations and acronyms used are described in

Table 1. Table of acronyms.

Acronyms	Meaning
PD	Primary data (original data)
QoI	Quantities of interest
AE	Autoencoder
PQ	Product quantizer
AET	Combination of Tucker decomposition, AE, and PQ
AETM	AET followed by MGARD for residual post-processing
AETMC	AETM followed by constraint satisfaction for QoI preservation
MGARDU	MGARD uniform
MGARDN	MGARD non-uniform

.

4.3. Parameter Settings

We experimented with a reasonable cross-product of parameter choices in our data reduction pipeline. The configuration space of parameters is shown in

Table 2. Configuration space of our framework.

Reduction	Components		Parameters
PD	AET Compression	Bottleneck layer	5	⋯	300
		PQ bits	4	⋯	12
	Residual post-processing	PD threshold	${10}^{-5}$	⋯	${10}^{-2}$
		MGARD error bound	${10}^{-5}$	⋯	${10}^{-2}$
QoI	Lagrange Multipliers	PQ bits	6	⋯	10

. Our framework consists of mainly two parts—PD compression and QoI satisfaction. There are four parameters that affect overall PD compression and one parameter in the reduction of QoI which affects storage efficiency of the Lagrange multipliers. We report our experimental results with optimal combinations of parameters depending on the reduction amounts and error levels. The implementation details of the AE architecture and algorithm along with the computation time are described in Appendix A, Appendix B and Appendix C.

4.4. Results

4.4.1. Reduction of the Primary Quantities

The PD reduction aspect in our framework comprises global data compression with AET and residual post-processing with MGARD. In AET compression, we varied the size of AE bottleneck layers and the number of bits in the PQ dictionary (PQ bits) to represent codewords. Let N be the number of nodes, L be the bottleneck layer size, and B be PQ bits. Then, the compressed data size is $N\times L\times B$ bits, as discussed in Section 3.1.3. When we increase the number of PQ bits, codeword cardinality in the PQ dictionary is increased, and thus, the accuracy of quantization is improved. The compression ratio, however, is decreased due to the increased number of bits for quantization. In the XGC application, using four bits for PQ and a small size for the bottleneck layer (∼30) is sufficient, since we have an additional processing step with residuals.

Residual post-processing enables us to guarantee the PD error in our framework. We compress residuals with MGARD to bound the reconstruction error of residuals, and thus the overall PD error is bounded. In this step, any error-bounded lossy compressors can be used. We call the reduction of PD part in our framework AETM (AET followed by MGARD). To optimize the amount of compression achieved, we limit this step to only nodes where the errors are high, and therefore it is conditional post-processing. As discussed in Section 3.1.4, the residual of a node is processed if the NRMSE of the node is larger than a user-defined PD threshold. Therefore, we can achieve different compression amounts by changing the PD threshold and MGARD error bound parameters. Figure 3a shows that the AETM outperforms the two versions of MGARD in PD NRMSE. This is due to the efficient compression in AET followed by selective residual post-processing on a subset of the output.

4.4.2. Satisfaction of Derived Quantities

Figure 3 shows that AETM works better than MGARD uniform, but worse than MGARD non-uniform in QoI preservation. After we obtain the reconstructed output via AETM, we perform the reduction of QoI errors. We call our entire framework, including QoI satisfaction, AETMC (AETM followed by constraint satisfaction). To analyze particle simulation properties, it is crucial for XGC scientists to maintain low errors of QoI (less than 0.001 NRMSE) after decompression. As we discussed in Section 3.2.1, the QoI satisfaction process produces four Lagrange multipliers per node, and therefore we need to store four additional floating-point numbers per node. To improve storage efficiency, we use PQ to quantize the Lagrange multipliers. Hence, there is a trade-off between the amount of the reduction of the QoI errors and compression ratios as shown in Figure 4. Additionally, we need to consider the increase in PD errors after the constraint satisfaction step, as discussed in Section 3.2.1.

Table 3. Primary data error increase after constraint satisfaction. The numbers in the parentheses here and in Figure 4 are the quantization bits for Lagrange multipliers.

Framework	PD Error Increase $\left(\%\right)$
AETMC (6)	10
AETMC (8)	9.4
AETMC (10)	9.3

and Figure 5a show that the increase in PD errors is tolerable and the AETMC outperforms MGARD in PD NRMSE. Figure 5 and Figure 6 show the overall results and the performance on each QoI.

4.4.3. Framework with Test Data

In the previous experiment, we did not split training and test data, as the case considered was static in time. However, since XGC is a spatio-temporal simulation, we now consider one set of timestep data as a training set and other set of timestep data as the test set. We first train an AE model at timestep 400 and then use the pretrained AE model to test timesteps 500, 600, and 700. Here, only the AE aspect of our framework is related to training and testing with other parts of our framework such as residual processing and constraint satisfaction, as before. Figure 7 shows the comparison between the pretrained AE model and AE models separately trained at each timestep. The results show that the performance of the pretrained model is very close to the models trained at each timestep. Consequently, this indicates that it is possible to improve machine efficiency, since the cost can be amortized across time. We plan to explore this in the future.

5. Conclusions

We have presented a comprehensive framework and a set of algorithms for error-bounded scientific data compression, which are broadly applicable to tensor datasets with moment constraints. Tucker decompositions, product quantization, and error-bounded compression are combined with an autoencoder (AE) in the framework. An AE with the Tucker decomposition can learn efficient global representations, and the reconstruction errors are bounded by selective residual post-processing with an error-bounded compressor. This leads to high levels of compression with bounded errors. Furthermore, the constraint satisfaction step enables us to reduce the errors of the QoI, and therefore the QoI are preserved during compression. The application to XGC and comparison with MGARD has shown that we can achieve compression ratios in excess of 500 while satisfying QoI error bounds. We also show that a pretrained AE model shows good performance across nearby timestep data points in the XGC simulation. Therefore, the framework can compress multiple timestep data efficiently by omitting retraining the AE model for each timestep. Immediate future work will focus on improving machine efficiency with amortization across time, and end-to-end learning machines which integrate latent space entropy estimation and constraint satisfaction. Non-linear QoI that cannot be simplified need further investigation as well. Finally, we plan to develop methods for other scientific applications (such as climate) which require error-bounded compression.