Machine-Learned HASDM Model with Uncertainty Quantification

HASDM is an assimilative framework using the Jacchia-Bowman 2008 Empirical Thermospheric Density Model (JB2008) as a background density model. HASDM improves upon the density correction work of Marcos et al. (1998) and Nazerenko et al. (1998) to modify 13 global temperature correction coefficients with its DCA algorithm. HASDM uses observations of more than 70 carefully chosen calibration satellites to estimate local density values. The satellite orbits span an altitude range of 190-900 km although a majority are between 300 and 600 km (Bowman and Storz 2003). The HASDM algorithm uses a prediction filter that employs wavelet and Fourier analysis for the correction coefficients (Storz et al. 2005). Another highlight of HASDM’s novel framework is its segmented solution for the ballistic coefficient (SSB). This allows the ballistic coefficient estimate to deviate over the fitting period for the satellite trajectory estimation.

SET validates the HASDM output each week and archives the results. The archived values from 2000-2020 make up the SET HASDM density database, upon which this work is based. The database contains density predictions with 15^∘ longitude, 10^∘ latitude, and 25 km altitude increments spanning from 175 - 825 km. This results in 12,312 grid points for every three hours from the start of 2000 to the end of 2019. For further details on HASDM, the reader is referred to Storz et al. (2005), and for details on SET’s validation process and on the database, the reader is referred to Tobiska et al. (2021).

JB2008, and subsequently HASDM, use solar and geomagnetic indices/proxies to drive them. F_10.7 (referred to as F₁₀ in this work) is a legacy solar proxy used widely in thermospheric density models. It represents 10.7 cm solar radio flux which does not directly interact with the thermosphere, but it is a reliable proxy for solar EUV heating. The S₁₀ index is a measure of the integrated 26-34 nm solar EUV emission which is absorbed by atomic oxygen in the middle thermosphere. M₁₀ is a proxy that represents the Mg II core-to-wing ratio and is a surrogate for far ultraviolet photospheric 160-nm Schumann-Runge Continuum emissions that cause molecular oxygen dissociation in the lower thermosphere. The last solar index used by JB2008 is Y₁₀ which is a hybrid measure of solar coronal X-ray emissions and Lyman-alpha. S₁₀, M₁₀, and Y₁₀ have no relation to the 10.7 cm wavelength but are converted to solar flux units (sfu) through linear regression with F₁₀. The 81-day solar centered averages of these four solar drivers are used as inputs to JB2008 and are denoted by the ”81c” subscript.

For temperature equations corresponding to geomagnetic activity, JB2008 utilizes a_p and Dst. The 3-hour a_p is indicative of overall geomagnetic activity and is only measured at low to mid-latitudes. It has 28 discrete values which are capped at 400 2nT leading to further limitations. To combat this, JB2008 uses Dst during geomagnetic storms if the minimum Dst ¡ -75 nT, and it is the primary parameter for energy being deposited into the thermosphere during storms. The data (drivers and density) is split into training, validation, and test data using 60%, 20%, and 20%, respectively as shown in Figure 1. Table 1 shows the number of time steps in the SET HASDM density database across various space weather conditions. The cutoff values for F₁₀ and a_p are obtained from Licata et al. (2020b) where they were used to bin space weather driver forecasts.

In Table 1, there is clear under-representation of geomagnetic storms in this vast dataset. This can cause limitations in model development, because over 98% of the dataset corresponds to a_p $\leq$ 50. Hierarchical modeling could be used for data of this nature, but we proceed with the development of a single comprehensive model.

Principal Component Analysis is an eigendecomposition technique that determines uncorrelated linear combinations of the data that maximize variance (F.R.S. 1901; Hotelling 1933). PCA has been previously used to analyze atmospheric density models and accelerometer-derived density sets. Researchers applied PCA to CHAllenging Minisatellite Payload (CHAMP) and Gravity Recovery and Climate Experiment (GRACE) accelerometer-derived density data in order to analyze the dominant modes of variation and identify phenomona encountered by the satellites (Matsuo and Forbes 2010; Lei et al. 2012; Calabia and Jin 2016). PCA has also been leveraged for dimensionality reduction to create linear dynamic reduced order models (Mehta and Linares 2017; Mehta et al. 2018; Mehta and Linares 2020; Gondelach and Linares 2020). Licata et al. (2021b) recently applied PCA to the SET HASDM density database for investigation.

The SET HASDM dataset is a prime candidate for PCA with the goal of predictive modeling due to its high dimensionality. The three spatial dimensions are first flattened to make the dataset two-dimensional (time $\times$ space). Then a common logarithm (log₁₀) of the density values is taken in order to reduce the variance of the dataset. Next, we subtract the temporal mean for each cell to center the data and save it for later use. Finally, we perform PCA using the svds function in MATLAB to acquire the $U$, $\Sigma$, and $V$ matrices (shown in Equation 2). PCA decomposes the data and separates spatial and temporal variations as in Equation 1,

$$\begin{split}\mathbf{x}\left(\mathbf{s},t\right)=\mathbf{\bar{x}}\left(\mathbf{s}\right)+\mathbf{\widetilde{x}}\left(\mathbf{s},t\right)\;\;\;\textrm{and}\;\;\;\mathbf{\widetilde{x}}\left(\mathbf{s},t\right)=\sum^{r}_{i=1}\alpha_{i}\left(t\right)U_{i}\left(s\right)\end{split}$$

(1)

where $\mathbf{x}\in\mathbb{R}^{n}$ is the log-transformed model output state (HASDM density on its full 3D grid), $\mathbf{\bar{x}}$ is the mean, $\mathbf{\widetilde{x}}$ is the deviation about the mean, $\mathbf{s}$ represents the spatial dimension, $r$ is the choice of order truncation (here $r$=10), $\alpha_{i}(t)$ are temporal coefficients, and $U_{i}$ are orthogonal modes, or basis functions. Equation 2 shows how to derive these modes.

$$\begin{split}\mathbf{X}=\left[\begin{array}[]{ccccc}\vrule&\vrule&\vrule&&\vrule\\ \mathbf{\widetilde{x}_{1}}&\mathbf{\widetilde{x}_{2}}&\mathbf{\widetilde{x}_{3}}&\ldots&\mathbf{\widetilde{x}_{m}}\\ \vrule&\vrule&\vrule&&\vrule\end{array}\right]\;\;\;\textrm{and}\;\;\;\mathbf{X}=U\Sigma V^{T}\end{split}$$

(2)

In Equation 2, $m$ represents the ensemble size (two solar cycles with HASDM) and $\mathbf{X}$ represents the log-transformed, mean-centered data . The left unitary matrix, $U$, is made of orthogonal vectors that represent the modes of variation, and $\Sigma$ is a diagonal matrix consisting of the squares of the eigenvalues that correspond to the vectors in $U$. We encode the data ($\mathbf{X}$) into temporal coefficients ($\alpha_{i}$) by performing matrix multiplication between $\Sigma$ and $V^{T}$. To decode back to the density, the coefficients are multiplied by $U$ with the temporal mean added at each cell, and taking the antilogarithm ($10^{y}$) of the resulting value completes the back-transformation. Figure 2 shows the first 10 PCA coefficients generated using the methods described above. The first five are discussed in detail by Licata et al. (2021b).

In machine learning, developing an accurate model previously required a manual architecture selection process that did not guarantee optimal performance for a given application. Recent developments in hyperparameter tuning make this process much quicker and more thorough. In this work, we use KerasTuner which allows the user to define the ranges and choices of any hyperparameter and to choose a search algorithm (O’Malley et al. 2019). We train all models with 100 trials (or architectures) with the first 25 being randomly sampled from our hyperparameter space while the last 75 trials take advantage of the tuner’s Bayesian optimization scheme.

The tuner trained three identical models for each trial and returned the lowest validation loss after 100 training iterations, or epochs. The Bayesian optimization scheme chooses future architectures based on the validation losses resulting from previous trials. Once completed, the tuner returns the best 10 models which we can evaluate on all three subsets of our data to find the most accurate and generalized model.

A traditional approach to regression modeling with UQ is to develop Gaussian process (GP) models (Rasmussen 2004; Chandorkar et al. 2017). In the early stages of model development, we attempted this approach – training GP regression models on each of the PCA coefficients. However, we could only use 10 years of data for training, limited by computational resources, which resulted in higher predictive error. In addition, the ten resulting models were 6.83 GB each making subsequent work cumbersome. Therefore, the results only pertain to the feedforward neural networks with MC dropout. There are three methods we explore in developing the best HASDM-ML model. First, we create a ROM using PCA for dimensionality reduction and a nonlinear ANN for prediction. We then modify this technique using a custom loss function in an attempt to obtain a model with calibrated uncertainty estimates. We test two loss functions to achieve this: NLPD and CRPS; details of which are given in the following section.

ANNs have two key components: features (or inputs) and labels (or outputs). We try using three separate input sets for our regression models, the first two are explained in Table 2. The first set is the JB2008 input set, referred to as JB. Licata et al. (2021b) showed that the SET HASDM density database contained evidence of post-storm cooling mechanisms which cannot be captured solely by geomagnetic indices at epoch. Therefore, we introduce a second set that is similar to the first but with a time history of the geomagnetic indices. We refer to this as JB_H (Historical JB2008). Unlike the actual JB2008 inputs, all of our input sets contain sinusoidal transformations to the day of year (doy) and universal time (UT) inputs. The four resulting temporal inputs (shown in Equation 3) represent annual (t₁ and t₂) and diurnal (t₃ and t₄) variations. The use of cyclical functions of time as inputs has been previously demonstrated (Weimer et al. 2020; Licata et al. 2021c).

$$\begin{split}t_{1}=sin\left(\frac{2\pi doy}{365.25}\right),\;\;\;\;t_{2}=cos\left(\frac{2\pi doy}{365.25}\right),\;\;\;\;t_{3}=sin\left(\frac{2\pi UT}{24}\right),\;\;\;\;t_{4}=cos\left(\frac{2\pi UT}{24}\right).\end{split}$$

(3)

In the JB_H set, the geomagnetic indices are extensive in an effort to improve storm-time and post-storm modeling. The ”A” subscript refers to the daily average, and the numerical subscripts refer to the value of the index that many hours prior to the epoch. The combination of two numbers references the number of previous hours the index is being averaged over (e.g. a_p12-33 refers to the average a_p value between 12 and 33 hours prior to the prediction epoch). The authors found that using different time histories of a_p and Dst (shown in Table 2) resulted in more optimal performance. For completeness, the results will also be shown using an input set that adopts the same time history for Dst as the a_p time history in Table 2. This input set will be referred to as JB_H0.

As the geomagnetic inputs in the JB_H and JB_H0 sets are in a time series format, we tested the time-delay neural network (TDNN) approach. TDNNs use convolutional layers to absorb the information of the time history of a given input and make the model time invariant (Waibel et al. 1989). We ran tuners for the two historical input sets using the MSE loss function and did not see error reductions worthy of further investigation. Additionally, the architecture required for a TDNN adds significant complexity when implementing the NLPD and CRPS loss functions explained in Section 2.5.

Dropout is a regularization tool often used in machine learning (ML) to prevent the model from overfitting to the training data (Srivastava et al. 2014). In standard feed-forward neural networks, each layer sends outputs from all nodes to those in the subsequent layer where they are introduced to weights and biases. Deep neural networks can have millions of parameters and thus are prone to overfitting. This causes undesired performance when interpolating or extrapolating.

Dropout layers use Bernoulli distributions, one for each node, with probability P. This makes the model probabilistic since the distributions are sampled each time that a set of inputs are given to the model. If a sample is ”true”, the node’s output is nullified, and the output of the layer is scaled based on the number of nullified outputs. Dropout is believed to make each node independently sufficient and not reliant on the outputs of other nodes in the layer (Alake 2020).

In traditional use, dropout layers are only activated during training to uphold the deterministic nature of the model. However, measures can be taken in order for this feature to remain activated during prediction making the model probabilistic. When passing the same input set to the model a significant number of times (e.g. 1,000), there is a distribution of model outputs for each unique input. This process is referred to as Monte Carlo (MC) dropout. Essentially, every time the model is presented with a set of inputs, random nodes are dropped out providing a different functional representation of the model. Gal and Ghahramani (2016) show that MC dropout is a Bayesian approximation of a Gaussian process.

In this case, the model uses the input set to predict the ten (10) PCA coefficients. Through the bootstrapping method, we estimate the mean and variance for each of the coefficients (Efron 1979). As the variance is different for each output/coefficient, this is a heteroskedastic problem. Each unique input can be passed to the model k times and there will be k $\times$ 10 outputs. The mean and variance are computed from those outputs with respect to the repeated axis, k. The two loss functions used to improve uncertainty estimation (in addition to MSE) are NLPD and CRPS. NLPD is based on the logarithm of the probability density function (pdf) of the Gaussian distribution, and is shown in Equation 4 (Matheson and Winkler 1976; Camporeale and Care 2021).

$$\begin{split}NLPD(y,\mu,\sigma)=\frac{\left(y-\mu\right)^{2}}{2\sigma^{2}}+\frac{log(\sigma^{2})}{2}+\frac{log(2\pi)}{2}\end{split}$$

(4)

In Equation 4, y is the ground truth, $\mu$ is the mean prediction, and $\sigma$ is the predictive standard deviation, each being computed for all 10 outputs making it a heteroskedastic loss function. For clarity, the log used in the NLPD loss function is the natural logarithm. The second loss function for uncertainty estimation is CRPS which is shown analytically in Equation 5 (Gneiting et al. 2005). The main difference between NLPD and CRPS is that CRPS is also includes the cumulative distribution function (cdf) of the Gaussian distribution as opposed to only the pdf,

$$\begin{split}CRPS(y,\mu,\sigma)=\sigma\left[\frac{y-\mu}{\sigma}\textrm{ erf}\left(\frac{y-\mu}{\sqrt{2}\sigma}\right)+\sqrt{\frac{2}{\pi}}\textrm{ exp}\left(-\frac{\left(y-\mu\right)^{2}}{2\sigma^{2}}\right)-\frac{1}{\sqrt{\pi}}\right].\end{split}$$

(5)

As with Equation 4, $\mu$ and $\sigma$ are computed using bootstrap methods, and this loss function is also heteroskedastic. An important aspect of using the loss functions described in Equations 4 and 5 is the preparation of the training data. The data is traditionally in the following form. The n features are set up as the number of samples, with n_I denoting the number of inputs, resulting in the input shape dimension as (n $\times$ n_I). The n labels are the number of samples with n_o the number of outputs, resulting in the output shape dimension as (n $\times$ n_o). To implement these loss functions, we stack each input and output by the number of MC samples, k. This is a repeated axis, meaning all samples along this axis are identical. The resulting shapes of the features and labels are (n $\times$ k $\times$ n_I) and (n $\times$ k $\times$ n_o), respectively. This allows us to determine the mean and standard deviation for each unique input within the loss function.

To assess the conditions in which HASDM-ML can improve, the global mean absolute errors are computed as a function of space weather conditions. The results are shown in Table 5 and the number of samples in each bin can be found in Table 1.

The results from Table 5 show that HASDM-ML is robust to different F₁₀ and a_p ranges when a_p $\leq$ 50. The only conditions in which the mean absolute error exceeds 11% is when a_p ¿ 50, which only accounts for 1.80% of the samples. This shows that more samples are required for this specific condition in both the training and evaluation phases. The last row contains the errors only as a function of F₁₀ which shows that across all four solar activity levels, the error deviates by only 1.24%. The bottom-right cell shows that the error across all 20 years of available data is only 9.71%.