MechGNN-Epi: Mechanistically Constrained Spatiotemporal Graph Learning for Regional Epidemic Forecasting

Mechanistic compartmental models such as SIR provide interpretability and enforce conservation of population, but their predictive performance can degrade in heterogeneous settings where transmission varies over time and across regions. Conversely, spatiotemporal graph neural networks (GNNs) can learn flexible spatial interactions from data, yet unconstrained predictors may yield epidemiologically inconsistent dynamics.

We propose MechGNN-Epi, a hybrid framework that couples a spatiotemporal graph encoder with a differentiable SIR update. Rather than directly predicting future infections, the model learns region- and time-specific transmission and recovery parameters from graph-based latent representations and then propagates compartments through a mechanistic solver. This design yields epidemiologically constrained trajectories and produces region- and time-indexed parameter proxies that can be inspected as diagnostic signals, while not being guaranteed as causally identifiable mechanistic parameters. We evaluate the pipeline on U.S. state/territory COVID-19 surveillance data using strict chronological splits and report variability across multiple random seeds, multi-horizon rollouts, and ablation/sensitivity studies against graph and model variants.