Forecasting infections with spatio-temporal graph neural networks: a case study of the Dutch SARS-CoV-2 spread

Abstract

The spread of an epidemic over a population is influenced by a multitude of factors having both spatial and temporal nature, which are hard to completely capture using first principle methods. This paper concerns regional forecasting of SARS-Cov-2 infections 1 week ahead using machine learning. We especially focus on the Dutch case study for which we develop a municipality-level COVID-19 dataset. We propose to use a novel spatiotemporal graph neural network architecture to perform the predictions. The developed model captures the spread of infectious diseases within municipalities over time using Gated Recurrent Units and the spatial interactions between municipalities using GATv2 layers. To the best of our knowledge, this model is the first to incorporate sewage data, the stringency index, and commuting information into GNN-based infection prediction. In experiments on the developed real-world dataset, we demonstrate that the model outperforms simple baselines and purely spatial or temporal models for the COVID-19 wild type, alpha, and delta variants. More specifically, we obtain an average R2 of 0.795 for forecasting infections and of 0.899 for predicting the associated trend of these variants.