Global Prediction of COVID-19 Variant Emergence Using Dynamics-Informed Graph Neural Networks

In the context of infectious diseases, the term “variant” refers to a version of a virus with some changes in its genetic material, known as genome [12]. These changes happen through genetic mutations and can affect the virus’s characteristics, like how easily it spreads, how severe the illness it causes is, and whether it can evade the immune system or not. Variants can be categorized in multiple ways based on genetic differences, and the significance of these differences can vary. Some categorizations focus on a common ancestor, while others may specifically highlight mutations in key regions of the virus’s genome [12]. In our study, we rely on the categorizations provided by the CoVariants dataset [11].

Furthermore, the proportions of different variants can be analyzed over time. Some studies have observed that these proportions follow a straight line when plotted on a semi-logarithmic scale. They leverage this fact to understand the dynamics of variant prevalence in populations [13]. We also observe this in our data as shown in Figure 2. However, to the best of our knowledge, existing studies do not derive any dynamics capable of predicting delays in the emergence of variants between multiple regions/countries.

Many efforts were made to develop forecasting models capable of predicting the progression and future trends of the COVID-19 pandemic [14, 15, 16]. However, these forecasting approaches have primarily focused on predicting the spread of the virus rather than predicting delays in variant emergence. Predicting the timing of variant-specific outbreaks can significantly enhance public health preparedness and response strategies, allowing for timely interventions such as targeted vaccination campaigns and tailored public health measures [3].

Many Deep Learning (DL) techniques have been explored for this purpose, including GNNs which help capture the graph effects of the spread of infectious diseases over time across regions [17, 18, 19, 20, 21]. Other methods like PINNs  [22, 23, 24], agent-based models [25, 26], and a combination of spatio-temporal networks with location-aware features [27], including mobility data [28] have also been used.

Instead of only predicting virus spread, as others have done, we aim to predict when a new variant will emerge in a region given that it has emerged elsewhere. This involves deriving dynamics of variant spread and integrating these insights into a GNN. Notably, our strategy revolves around crafting relevant features rather than altering the loss function, providing a distinctive and effective approach.

First, we develop an understanding of the dynamics that play a role in the spread of a new variant and how we can utilize them to make informed predictions. We start by extracting each variant’s global growth rate from observed data. This value provides an indication of the dominance of the variants over other competing ones. Subsequently, we construct linear models based on the dynamics of infectious diseases and the derived growth rates to predict delays between regions, establishing this Dynamics-based Model as our baseline for Delay Prediction. Finally, we consider different approaches in which we can incorporate some disease dynamics using GNNs that are chosen for their ability to intuitively model spatial dynamics, with the goal of enhancing the prediction accuracy of our model.

Throughout the text, we use the terms “prevalence” and “prevalence ratio”, where the “prevalence”, $p$, is the proportion of cases of a given variant, and the “prevalence ratio” is defined as $r_{ij}=\frac{p_{i}}{p_{j}}$, where $p_{i}$ is the prevalence of the variant $i$, and $p_{j}$ is the prevalence of the most prevalent variant, $j$, that is not $i$. In our experiments, a variant is considered dominant if it reaches an $r_{ij}$ value of at least $\frac{1}{3}$. We define this value as the dominance threshold, $\theta$.

We first derive the dynamics of two competing variants in one region. While logistic regression has been widely used to fit these dynamics, for completeness, we present a derivation. We start by using a model of infectious diseases that generalizes a large class of models [29]:

where $\Delta I(t)$ is the new infections at time $t$, $\bar{\alpha}(\tau)$ are transmission parameters, $\bar{S}(t)$ is the fraction of susceptible population, and $\epsilon_{t}$ is some small number. The number of contacts at time $t$ is given by $\kappa(t)$. The approximation follows from Taylor expansion, and using the fact that $\Delta I(t)$ is a smooth function. The constant $b$ happens to be the mean serial interval – the average time between two successive infections. Now, given two variants $i$ and $j$, assuming full cross-immunity between all variants, i.e. $\bar{S}_{i}(t)=\bar{S}_{j}(t)=\bar{S}(t)$, the prevalence ratio is approximated as:

The above further assumes that the mean serial intervals of the two variants are close, i.e., $b_{i}=b_{j}=b$. Note that we assumed full cross-immunity which is true for a number of variants before the Omicron sub-variants [30]. If there is only a partial cross immunity, $\bar{S}_{i}(t)\neq\bar{S}_{j}(t)$. However, in a short interval, the new immunity (reduction is susceptibility) created by the variants is small. As a result, it can be shown that $\bar{S}_{i}(t)\approx\beta_{ij}^{\prime}\bar{S}_{j}(t)$ for some constant $\beta_{ij}^{\prime}$ which can be absorbed in the parameter $\beta_{ij}$ above. Solving the above recurrence relation, we get:

Taking the logarithm reduces the equation to:

Here, constant $b$ is the mean serial interval – the average time between two successive infections, and $\beta_{ij}$ is the ratio of transmission rates between the two competing variants. Parameter $S_{ij}$ can be referred to as the relative growth rate of variant $i$ over variant $j$ [29].

Now, consider a scenario in two regions $A$ and $B$ — $A$ has the variant $i$ emerging over the previously dominant variant $j$; but variant $j$ never reached region $B$ where variant $k$ is dominant. This makes it difficult to assess the potential impact of variant $i$ on region $B$ when we only know $S_{ij}$. To deal with this challenge, we note that the parameter $S{ij}$ should ideally be independent of the region (this may not hold due to the simplifying assumptions above, so we actually have a region-specific relative growth rate $S^{(p)}_{ij}$ for region $p$). Therefore, we can attempt to find the growth-rate of any variant $i$ with a fixed variant $0$, specifically, the original COVID-19 variant. For ease of notation, we represent the growth advantage of variant $i$ with respect variant $0$ as $S_{i}$, with $S_{0}=0$, and refer to these values as global growth rates. Ideally, $S_{ij}=S_{i}-S_{j}$, irrespective of the region where the $S_{ij}$ is estimated.

To estimate these the global growth rates for all the variants that exist by the time $t$ over all regions $p$, we can set up the following system of linear equations:

where $S^{(p)}_{ij}$ is the growth advantage of variant $i$ over variant $j$ in region $p$ and $\epsilon_{ijp}$ are error terms. The solution is given by minimizing the sum of squares of $\epsilon_{ijp}$.

We, then, define an objective function and solve for the $S_{i}$ value as

Finally, based on the global growth rates, we define global relative growth rates $S_{ij}$ as $S_{i}-S_{j}$.

These growth rates are used in our proposed methods for Delay Prediction.

Consider three variants $i$, $j$, and $k$ where $i$ is an emerging variant. We assume that either of the following holds – (i) Variants $j$ and $k$, respectively in regions $X$ and $Y$, are the variants that constitute most of the other infections; or (ii) All variants other than $i$ in the respective regions $X$ and $Y$ have similar enough global growth rates that can be grouped into variants $j$ and $k$. Suppose variant $i$ appears in region $X$ first, followed by region $Y$. Then we have:

For a new variant, initially, the number of infections increase because of the new infections coming from a different region (importations). After some point, when transmission starts to happen within region (community transmission), the imported infections become negligible in comparison, and this is the point after which above equations are valid.

Suppose we select $t_{0}$ as the time at which $p^{X}_{i}(t_{0})$ is large enough such that the community transmission dictates the dynamics in the region $Y$ rather than importations and so, both Equations 6 & 7 hold for $t\geq t_{0}$. Suppose that the threshold for prevalence in region $Y$ for community transmission to dominate is $\theta_{Y}$ and that for region $X$ is $\theta_{X}$. We assume that $\theta_{X}$ and $\theta_{Y}$ are independent of the variant. Then, by setting $t=t_{0}$ in equations 6 & 7, we get:

Therefore, both $C_{Y}^{\prime}$ & $C_{X}^{\prime}$ are variant independent. Let $t^{\theta}_{Y}$ and $t^{\theta}_{X}$ be the times at which variant $i$ reach target prevalence of $\theta$ in regions $Y$ and $X$, respectively. We want to find $\tau=t^{\theta}_{Y}-t^{\theta}_{X}$. From equations 6 and 7:

Also from equation 7:

Plugging this value into equation 9, we get:

Here $\theta$ is a given prevalence, $S_{ij}$, $S_{ik}$, and $S_{kj}$ are global relative growth rates estimated as shown in Section III-B. $C^{\prime}_{X}$ and $C^{\prime}_{Y}$ are unknowns. Therefore for each pair of regions $X$ and $Y$, to identify the delay in reaching a given prevalence $\theta$, we can build a linear model $Y=W^{T}Z$

We then use the calculated weights from our linear models in equation 12 and plug them into equation 11 to find the $\tau$ value between two regions. This is a pairwise linear model that estimates the delay between two regions rather than providing a delay from the current date. Algorithm 1 details how we compute the date at which a variant arrives in a region $Y$. We calculate the median of the outputs from all pairwise models from regions $X_{q}$ to region $Y$. We evaluate the performance of all the models $(X_{q},Y)$ until the current date and select the top three performing source regions among ${X_{1},X_{2},\dots}$. Notably, best performance here refers to the lowest error in predicting the delay. We identify three source regions that produced the lowest error in predicting the delays at $t-1$. Then, we take the median of the predictions produced by these selected models at time $t$ which represents the delay at time $t$ for the new variant to appear in region $Y$.

Key Idea: We observe that global growth rates $S_{i}$ play a crucial role in the dynamics. Furthermore, based on Equation 3, under some assumptions, the logarithm of the prevalence ratio $r_{ij}$ grows linearly with time. We hypothesize that using $\ln(r_{ij})$ and $S$ as features simplify the underlying patterns to be learned by an ML algorithm. An ML algorithm may also complement any violations of the assumptions made in the dynamics-based model.

We propose a GNN-based network for this problem as GNNs are adept at capturing spatial relationships, making them well-suited for problems where geographical proximity between regions, such as countries, plays a crucial role. Additionally, the problem induces a natural graph structure, where countries can be represented as nodes with edges denoting some relationship between the countries. Furthermore, we explore different techniques to incorporate disease dynamics into the GNN in an effort to capture the complex interactions and patterns associated with the spread of infectious diseases.

All training and validation of the models were performed retrospectively, as outlined in Algorithm 2. This means that the model is trained and validated at each time step using only the ”observed” data available up to that time step, denoted as $d$. For each variant, the retrospective algorithm iterates through all the biweekly data associated with that variant, training the model only on the weeks that have already passed (i.e., the observed data). Note that the pre-processing is also done retrospectively, accounting for any smoothing, calculation of the $S$ values, or interpolation in the case of the dynamics-based model. This ensures that the approach is applicable in a prospective setting, where no data from the future is available.

The train/validation split is also done temporally to ensure that validation is performed only on the most recently observed data. This aims to achieve a better fit to the latest data, enhancing the model’s predictive performance. An Early Stopper monitors the validation loss and halts the training when overfitting is detected, signified by a constant and substantial increase in the validation loss. The Early Stopper saves the best-performing weights, which are reused in the next iteration. This practice eliminates the need for the model to start training from scratch, leveraging knowledge gained from previous iterations.

We use MSE, or its adjusted version, as the loss function. Another crucial aspect is how to handle data related to variants that either do not appear in a given country or appear but fail to become dominant, i.e. rapidly diminish. In such cases, where regression targets would be undefined (potentially infinite), we address this challenge by creating a mask. This mask is employed to conceal nodes corresponding to these specific variants/country pairs, ensuring they do not contribute to our training loss. Essentially, our graph models are not trained on these nodes.

To validate our methodology, we’ve used a combination of multiple datasets which are detailed below.

Since we are the first to attempt solving this problem, we provide a benchmarking tool for the community. Figure 5 presents the comprehensive pipeline employed for retrospective training, validation, and testing of diverse models. Its design prioritizes user-friendliness, allowing seamless integration of new PyTorch models or direct utilization of scikit-learn models through a configuration file. Users only need to specify whether the model requires graph data, leaving the pipeline to handle the rest. Additionally, the data pre-processing is encapsulated in a separate function, providing users with the flexibility to make modifications.

As shown in Figure 5, the evaluation is also being done retrospectively. It means that for each date that the model is trained on, it is used to predict for the subsequent date representing a 1-timestep ahead forecast. This aligns with the biweekly nature of the data, resulting in a two-week ahead forecast. Furthermore, model performance is being evaluated per variant, meaning that for each given variant we retrospectively assess the trained models over time.

In the context of delay prediction, considering the biweekly data, the target is consistently a multiple of 14 days. Therefore, the output of our models is always rounded up to the nearest multiple of 14. Furthermore, evaluations exclude countries where a variant has already become dominant or never achieves dominance.

The initial pre-processing stage aligned the country names across diverse datasets, with an emphasis on inclusion based on data availability. Only countries present in all datasets were retained, ensuring a standardized set for subsequent analyses. For each country, the prevalence of each variant was calculated at each timestep. Variants with a prevalence of less than 5% were disregarded in our analysis. To calculate the prevalence ratios, the two variants with the highest prevalence were identified for each timestep. For each variant on that particular day, the prevalence ratio was calculated. Finally, variants that did not reach a prevalence ratio of 0.2 or those that were present in less than three countries were filtered out. These steps ensured that only variants with a significant presence and wide geographic distribution were included in the subsequent analysis. All subsequent pre-processing steps were conducted retrospectively. This implies that as time progresses, more data is used for computing the global growth rates ($S$) and for training.

To measure performance, we define Mean/Median Absolute Errors (MAE/MedAE) as follows:

where $i$ ranges from 1 to $N_{s}$ and $N_{s}$ represents the number of valid countries to predict for. The $t$, $v$ subscripts indicate that these errors are calculated for each variant at each timestep. Since this problem is retrospective, we also evaluate performance over time by finding the mean and median of the errors:

where $t$ ranges from 1 to $T$ and $T$ represents the total number of timesteps for which the variant $v$ circulates before reaching total dominance, i.e. it globally reached all its targets and begins to disappear. Finally, we conduct a mean across all the variants $v$ to get our four error metrics: MMAE, MMedAE, MedMAE, and MedMedAE. While all these metrics offer valuable insights, our main focus lies on MedMAE and MedMedAE, emphasizing median performance over time. This choice is made due to certain data inconsistencies, such as variants appearing in a single country for an extended period before spreading to other countries, resulting in large errors that may skew results.

We implemented a number of benchmark models that act as baselines for comparisons.

Mean Model – Delay Prediction: This model predicts delays based on average delays on the variants between the countries under consideration. Specifically, at time $t$, for a new variant $i$ on a target country $Y$, suppose $Q$ is the set of all countries where the variant has appeared. Then, the predicted time of arrival is given by

Here $\Delta_{j}(Y,q)$ is the delay between the appearance of variant $j$ from country $q$ to country $Y$. We take the mean of $\Delta_{j}(Y,q)$ over all prior variants $j$ that appeared in both countries $Y$ and $q$. Finally, we take the median of the delays over all countries in $Q$ where variant $i$ has already appeared.

Dynamics-based Model – Delay Prediction: This is the derived dynamics-based linear model discussed in the previous section. This model serves as a baseline mechanistic model, which, under the right assumptions, offers an explainable and efficient method for calculating delays.

Decision Tree: We use Decision Tree regression models to address Delay Prediction. This involves transforming the dataset into an $N\times 5$ matrix. Here, $N$ represents the total number of samples, where each sample corresponds to the 5 features of a given variant at a specific point in time in a specific country. Meaning $N=T\times V\times C$, where $T$ is the total number of timesteps, $V$ is the total number of variants, and $C$ is the total number of countries. These trees use the Gini impurity and mean squared error for their criterion functions. No maximum depth was set for either.

Multi Layer Perceptron: We employ a Multi-Layer Perceptron (MLP) regressor for Delay Prediction. They were configured with two hidden layers, consisting of 16 and 8 neurons, respectively. The selection of hyperparameters, including the number of hidden layers, activation function, and solver, was determined through comprehensive experimentation and optimization tailored to our specific problem. The training procedure closely followed the approach outlined in Section 3.3, with the only distinction being the structuring of our dataset into an $N\times 5$ matrix, as described for the decision tree. A batch size equivalent to the number of countries $C$ was used, where each batch represented the list of countries and their features at each timestep for each variant, ensuring that the gradient updates were performed collectively for each group of variants at a given timestep.

Gated Recurrent Unit Network We utilize a simple GRU network to encode the temporal information of variant transmissions and predict theirdelay. To train this model, first, we feed temporal features of the variant to the GRU architecture which embeds them for each country distinctly. Then, we flatten the embedded features and append them with $S$ and directly give them to an FC layer for predicting the time it takes for that variant to appear in a specific country. We employ GraphNorm layer here after the concatenation to accelerate training by smoothing graph aggregation distributions.

Graph Convolutional Network (GCN) We consider a simple GCN, with a basic architecture in an attempt to capture local relationships and features within the graph structure. It is comprised of two GCN layers, which output 32 and 16 channels respectively. Dropout is strategically employed between these two layers for regularization. The final graph embedding is flattened and fed into an FC layer which outputs the delay for each country.

Encoded Dynamics GCN Each country has a variable number of circulating variants, each with different growth rates $S$ (as derived earlier). These growth rates can be concatenated and represented as a feature vector $S_{qt}$, which correspond to the variants present in a region $q$ at time $t$. We can learn a latent vector representation $Z_{qt}$ of this data through the use of encoders. We implement this in two ways:

Using an Autoencoder (AE-GCN): We employ a simple autoencoder with two linear layers and ReLU activation in the encoder and decoder blocks. The model takes the feature vector $S_{qt}$ as input. We train the autoencoder using MSE loss to reconstruct the feature vector $S_{qt}$. After training, we take the intermediate latent encoding $Z_{qt}$ output from the encoder and append it to the feature matrix before constructing a spatio-temporal snapshot graph data for the GCN model.

Embedded Encoder (EE-GCN): Using an autoencoder implies that the model is guided towards learning an embedding using the decoder output, and these two are inherently separate models. The encoder weights are also unaffected during graph training. An alternative approach is to embed this into the larger GCN model. In this method, as before, the variable-sized $S_{qt}$ is fed to an encoder block (consisting of two linear layers with ReLU activation). Instead of using a decoder, we directly concatenate the resulting latent vector representation $Z_{qt}$ to the node features fed into the GCN. This means the encoder weights are updated in the graph training loop, allowing the model to learn a better latent representation.

In this section, we present our key results for Emergence Delay Prediction. We previously noted that prior research has concentrated on predicting the prevalence of a variant after its arrival in a specific region, with logistic regression demonstrating effectiveness in estimating its prevalence over time. However, as logistic regression can manage this task relatively easily, our evaluation focused solely on forecasting a variant’s prevalence before its arrival in a specific region. In other words, during evaluation, we excluded cases where a variant is present in a region but has not yet reached $\theta$.

Numerous models were tested and configured as benchmarks to test against the performance of our proposed method. Other implementations include a trivial Mean Model, the baseline Dynamics-based Model, a Decision Tree, a Multilayer Perceptron (MLP), a GRU, and a GCN without an adjusted loss function. In brief, the Mean Model predicts delays based on the average delays of all prior variants between all the countries in which the variant appeared and our intended country. The Decision Tree, MLP, and GRU regression models make use of $T$ extra features which consist of the average $\vec{r}_{c,i}$ of neighboring countries transforming the dataset into a matrix of size $N\times(2T+1)$. Here, $N=D\times V\times C$ is the total number of samples, $D$ is the total number of retrospective dates, $V$ is the total number of variants, and $C$ is the total number of countries. The full details on these models can be found in the Appendix.

To measure the performance of the models, we use the defined Mean/Median Absolute Errors (MAE/MedAE) in the Appendix and measure the performance over time by finding the mean and median of these errors as:

where $t$ ranges from 1 to $T_{D}$, $T_{D}$ being the total number of timesteps for which variant $v$ circulates before reaching total dominance, i.e. it globally reaches all its targets and begins to disappear. Finally, the mean across all the variants $v$ gets our two error metrics: MedMAE, and MedMedAE.

Table I presents the models’ performance metrics for Delay Prediction. The results indicate that our proposed FA-GCN model outperforms all other models in terms of MedMedAE and MedMAE, showcasing its effectiveness in capturing temporal dependencies for Delay Prediction. The dynamics-based model also has a significantly worse MedMAE than MedMedAE. This suggests that the dynamics-based model is susceptible to outlier predictions for some country pairs, which skews its results. Additionally, we observe that the PI-GCN does not surpass its GCN counterpart, highlighting the unnecessary use of dynamics as a regularization of the loss when employing the correct features.

Figure 6 provides a detailed representation of the results for the baseline dynamics-based model and the FA-GCN model. The heatmaps illustrate the number of countries eligible for prediction at each timestamp. The numbers indicate the errors over time for each variant throughout their existence. Notably, several variants in the dynamics-based model are denoted by $*$, signifying instances of model or evaluation failure. These failures stem from three primary factors. Firstly, the linear model requires a minimum of two common variants between two countries to build an effective model. Secondly, the linear model can encounter significant challenges when variant $j$ is identical in growth rate to variant $k$ (i.e., $S_{j}=S_{k}$). In such cases, the linear model fits a plane solely along a fixed axis and will thus always predict some constant $\frac{C}{S_{ij}}$, as seen in Equation 11. But as time passes and variant $j$ is no longer identical to variant $k$ for another variant, this model fails drastically since it would still fit on a fixed axis. These scenarios were more prevalent in the early stages of the pandemic when only a few variants were circulating, leading to their exclusion from the dynamics-based model’s training process. To ensure fair comparisons, all final errors across all models were calculated excluding these variants. Finally, there are cases where a variant appears in a country but fails to reach the threshold $\theta$. These instances were disregarded during the evaluation of all models as they were considered trivial.

The observations from Figure 6 give us two insights. Firstly, the primary source of error is associated with variants that persist for an extended duration, such as 21J.Delta or 23E.Omicron. Examining the color gradients, these variants appear to linger in only a few countries before spreading to others. In reality, variants typically do not endure for such prolonged periods without either being superseded by another variant or transmitting to additional countries [30, 38]. This suggests a potential issue in the data capturing these lingering variants. Secondly, despite this being a major source of error for FA-GCN, it exhibits overall robustness to outlier predictions, unlike the dynamics-based model as seen in Figure 6. This, in turn, severely affects the results of the PI-GCN model.

While the dynamics-based model does output good predictions, it fails often and dramatically in some cases. Figure 7 presents examples of 3D points in space utilized for fitting linear models. In Figure 7(a), the 3D points form a space where a plane can be fitted exceptionally well. However, as newer COVID variants emerge, the dynamics may evolve, causing some initial assumptions, such as full cross-immunity [30], to become invalid. This evolution makes it more challenging to perfectly fit a linear model, as depicted in Figure 7(c).

The pipeline depicted in Figure 5, was implemented in Python. The implementation was carried out on a machine equipped with a 32-core Intel(R) Xeon(R) Gold 5218 CPU running at 2.3 GHz and 64 GB of RAM. Focusing on the example of variant 23F.Omicron at a single timestep, which contains the majority of the data, and the heaviest model (FA-GCN), the average runtimes for different components were approximately 32 seconds for data pre-processing, 20 seconds for training, and 0.01 seconds for inference. The entire retrospective pipeline takes approximately 2 hours to complete. Notably, the most time-consuming aspect of the pipeline is the data pre-processing step, which could potentially be optimized further through memorization.

One limitation of our approach in addressing the problem of Delay Prediction lies in handling situations where a variant might not even become dominant in a given country. To tackle this issue, we define the following problem:

In future work, we plan to address this issue by adopting a two-step approach. First, we solve Dominance Prediction, which aims to predict whether a variant will become dominant in a given region. If the variant is predicted to become dominant, we can then proceed to solve Delay Prediction.