In light of the millions of confirmed COVID-19 cases across the globe, understanding the transmission dynamics of infectious diseases through modeling, analysis, and simulation is essential for not only predicting disease spread but also informing effective policy interventions to combat it. After the outbreak of severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2) was first identified, governments were forced to interpret the evolving pandemic landscape to create effective policies, relying on data visualization tools and statistical models to understand disease spread both locally and globally.
To address this demand for accurate intervention simulation, Kerr et al. (2021b) as vaccination, social distancing, and contact tracing (Kucharski et al. (2020)), some ABMs, including Covasim, have been deployed by public health officials to inform policy decisions in over a dozen countries (Kerr et al. (2021a); Panovska-Griffiths et al. (2020); Quach et al. (2021)).
However, the computational cost of ABMs, at best, scales linearly with population size.
While new techniques such as dynamic rescaling are able to maintain a constant level of precision and computation time throughout the simulation, ABMs are challenged by a constant tradeoff between computational cost and the inaccurate discretization of the population.
Compartmental models replace this discretization with mathematical theory, describing the progression of a population through disease states using differential equations. The Susceptible-Infected-Recovered (SIR) model by Kermack and McKendrick (1927) is one of the earliest such mathematical formulations, illustrated in Figure 1.
The simple SIR model assumes that the system is closed, ignoring vital dynamics such as birth and death rates. The target population, N , is then subdivided into three compartments by disease state; S(t), I(t), and R(t) correspond to the number of susceptible, infected, and recovered individuals respectively at a given time. The flux of individuals between compartments can be described using the Law of Mass Action and two disease parameters: the transmission rate, β, and the recovery rate, α. Together with initial conditions S(0), I(0), and R(0), the model's dynamics are fixed, as N = S(0) + I(0) + R(0). The SIR model not only provides theoretical epidemiological insights but also provides a way for generating realistic outbreaks using a numerical integration technique such as Runge-Kutta 4/5 (Butcher (2016)).
Compared to ABMs, the versatility of compartmental models is limited, as developing a mathematical system that encapsulates the complexity of human social networks is challenging.
However, recent advancements have allowed compartmental models to simulate many behaviors.
For example, Yang et al. (2021) evaluated the effect of lockdown policy in Spain with a model that introduced compartments for exposed, asymptomatic, pre-symptomatic, mild COVID, and severe COVID. Bouchnita et al. (2021) used an SEIR model to predict that social distancing measures would be needed to shorten COVID-19 waves in Vietnam.
Similarly, this work uses a SIR-based model to improve vaccination programs, which help populations achieve a level of immunity that prevents infectious people from causing further outbreaks. While ABMs are able to evaluate the effects of various vaccination strategies (Sulis and Terna (2021)), they are unable to inform optimal policy decisions. In order to compute an optimal intervention strategy, simulations must be repetitively run, compounding the computational cost of a single simulation. Instead, our work exploits the simplicity of compartmental models to create a framework for efficient, accurate parameter estimation to drive further optimization using gradient descent.
A variety of approaches have been developed to estimate disease parameters from real-world data including non-parametric estimation Smirnova et al. (2019), particle swarm optimization Akman et al. (2018), Bayesian techniques Akman et al. (2016), and inverse methods. Statistical approaches such as maximum-likelihood and Poisson regression methods have also been applied Capaldi et al. (2012). Some of this work has shown that estimation precision increases with the amount of outbreak data used, and sensitivity analysis techniques such as Latin Hypercube Sampling have been able to rank the relative importance of model parameters.
We take the approach of simultaneously estimating disease parameters while approximating a compartmental system, and several works have also attempted to learn differential equations from data. Ling et al. (2016) modeled the Reynolds stress anisotropy tensor using deep neural networks, and E et al. (2017) solved parabolic partial differential equations (PDEs) using reinforcement learning. However, recent advancements have pioneered the use of physicsinformed models to estimate solutions. In fact, physics-informed neural networks (PINNs) were first benchmarked with the SIR model in Raissi et al. (2019b).
Our framework uses PINNs in conjunction with a meta-population model to estimate disease parameters which are used to construct an optimal vaccine distribution plan. The paper is sectioned as follows. In Section 2, we provide necessary background information and present the structure of our approach. Section 3 analyzes our approach through a series of computational experiments, and Section 4 addresses the limitations of our framework while outlining future research avenues.
In this section, we establish the motivation for using meta-population models and PINNs for parameter estimation. We also present an optimization technique for optimally distributing vaccines from these parameters.
A deep neural network consists of multiple connected neurons, each capable of producing a realvalued activation. In a feed-forward neural network such as Figure 2, neurons are organized in sets of L connected layers. Densely-connected layers, in particular, connect each pair of neurons in neighboring layers, x (l-1) j and x (l) i , with a weight, w (l) ij . A neuron's activation is then computed as a weighted sum of the previous layer's activations plus a bias term, b, and passed into an activation function σ(t). In general, the network is evaluated as follows
where x is the provided input and x (L) = y is the computed output. Treating the hidden layers as a black box, a neural network can be described as the vector-valued function f net ( x) = y, and by the Universal Approximation Theorem, f net can approximate any continuous function f : R n → R m provided that σ(t) is non-polynomial and the network has arbitrary width Hornik (1991). During supervised learning, the network trains on a labeled dataset of pairs ( x, y), using the back-propagation algorithm to update the network's weights and biases to minimize a loss function which captures the error between the predicted y pred and the true y for each sample. In the case of mean squared error (MSE),
, providing a method to estimate any vector-valued system.
In (2022), but by modeling infectious disease spread as an ODE system, such as the SIR model, PINNs can also be applied in the context of epidemiology.
For ODE problems, the network, as seen in Figure 3, consists of an input layer of one neuron corresponding to the input time t, which is normalized to the range [-1, 1] during pre-processing, with σ = tanh(t) being the preferred activation function.
The network is trained to approximate the ODE system, u(t), by minimizing a loss function consisting of both differential equation error and data prediction error. For the former, each equation in the system is rearranged and associated with an error E i , which is then summed across the system:
The data prediction error is computed in standard fashion by evaluating the Mean Square Error (MSE) between the predicted compartment values of the output layer and the data provided.
While Raissi et al. (2019a) add the data prediction error and differential equation error directly into the loss function, we linearly weight each term according to a new hyperparameter, λ.
The computed loss, L, is then minimized using stochastic gradient descent via the Adam opti- TensorFlow v2), we create trainable variables with appropriate bounds for each equation (disease) parameter and allow the optimizer to vary both network and equation parameters during training. After sufficient training, the model will not only approximate the ODE system but the parameters of the system will also be calibrated to the training data.
The pure SIR model assumes that modeled populations are homogenous, but real-world populations are better described as heterogeneous meta-populations with each individual exhibiting differing levels of activity and risk. Person-to-person contact rates vary with population density and personal characteristics including age, gender, and location, inhibiting the generalizability of simple compartmental models to larger populations Feng et al. (2015). We attempt to remedy this assumption by combining PINNs with the simplest meta-population model that captures the heterogeneity of sub-population contact rates and preferential mixing between groups in the meta-population Jacquez et al. (1988).
We begin by dividing the meta-population, N , into n sub-populations, N i , of individuals exhibiting similar activity levels. We then model each sub-population using SIR, partitioning N i into compartments based on disease progression: N i = S i + I i + R i . However, we replace the former force of infection βI with λ i , the per-capita force of infection for the i th population.
Formally, let
where a i is the average contact rate of the i th sub-population, β is the net probability of infection from an infectious contact, and c ij is the proportion of the i th sub-population's contacts that are with the j th sub-population. The term a i c ij I j N j can be interpreted as the average infectious contact rate of the i th population with the j th population, and by summing each force contribution across the sub-populations, we arrive at the net force of infection for the population of interest.
In the manner of Jacquez et al. (1988), the mixing matrix c ij is computed by reserving a portion of each sub-population's contacts for itself and then proportionally distributing the remaining contacts to the other sub-populations Jacquez et al. (1988). Formally, let
where δ ij is the Kronecker delta function and i is the proportion of contacts of the i th subpopulation reservered for itself. δ ij allows us to isolate the additional i contribution when i = j, and f j accounts for the proportional distribution of the remaining contacts by normalizing the general term (1j )a j N j .
We also account for the vital dynamics of the population with μ, merging both birth and death rates, and α, the recovery rate or the reciprocal of the mean infectious period. To avoid further complexity, we assume perfect vaccine efficacy and remove p i of each population where p i is the proportion of vaccinated individuals in the i th population. In total, this results in the following SIR-based system of 3n compartmental equations and 3n + 3 parameters.
While PINNs is able to use differential equation structure to fix a large number of equation parameters while approximating the ODE system, we attempt to limit the number of estimated parameters to expedite convergence. We introduce the substitution A i = a i β and rewrite λ i and f j as follows:
This allows us to eliminate β from the system. In the context of preferential mixing, β acts as a scaling term for a i and can be completely absorbed into a i → A i . In practice, β is already known from approximating the pure SIR model, in which case, a i can be directly estimated.
Further, the vaccination rate of each sub-population is typically known beforehand, eliminating n parameters of the form p i . Additional assumptions about μ can be made, depending on the target meta-population. During training, we remove fixed parameters from the optimizer's trainable variables, allowing us to only train parameters of interest. These optimizations allow us to reduce the number of parameters to 2n + 1.
Once the preferential mixing parameters are estimated, we calculate R v , the effective reproduction number of the meta-population or the average number of secondary infections per infectious person Diekmann et al. (2010). We first define the basic reproduction number for the i th sub-population
We then remove the proportion of vaccinated individuals to find R vi , the effective reproduction number for the sub-population.
To calculate R v for the entire meta-population, we consider the next-generation matrix K, computed by multiplying the diagonal matrix of sub-population effective reproduction numbers with the contact proportion matrix with elements c ij .
The effective reproduction number for the entire meta-population is the spectral radius, or largest eigenvalue, of this matrix Diekmann et al. (2010).
Fixing the remaining parameters, we can analyze R v as a function of the sub-population vaccination rates p i in the manner of Feng et al. (2015). We initialize p := p 0 to the estimated or known vaccination rates of the sub-populations. Then, using gradient descent, we update
This provides the ideal vaccination rates over time for each sub-population, p i (t), that minimizes the spread of the infection as fast as possible.
We propose PINNs as a method for directly informing infectious disease policy, particularly in the case of vaccine distribution (See Figure 5).
Using collected data from a target meta-population, we use PINNs to approximate a modified SIR model, producing additional disease parameters. The parameters are then used to calculate and minimize the effective meta-population reproduction number, R v , creating an optimal vaccine distribution strategy for the meta-population.
In this section, we substantiate the accuracy and versatility of PINNs parameter estimation on SIR-based models. We also explore the influence of model hyperparameters and the efficacy of the end-to-end framework in various mixing scenarios.
Using Runge-Kutta 4/5, we simulate SIR dynamics for the following initial conditions and disease parameters:
We then sample each curve at 15 random time steps to produce training data. We initialize each trainable disease parameter to 0 and then train a 1 × 20 × 20 × 20 × 3 PINNs architecture with a learning rate of 1.0 × 10 -3 and error weight of λ = 0.5 while constraining each parameter to the range [0, ∞). After ≈ 10 5 iterations, we arrive at the following approximation (See Figure 6).
Using the simulated curves and predefined parameter values, we precisely compute the PINNs estimation error by finding the MSE between each simulation compartment and the corresponding PINNs approximation using a dense time mesh. For each disease parameter, we simply compare the simulation value with the final value of the corresponding trainable parameter. TABLE 2: Training Dynamics with Varying λ
We find that λ = 0.5 achieves optimal convergence time, suggesting that an equal weighting of data prediction error and differential equation error yields the fastest approximations.
Having established PINNs' accuracy, we use a popular London boarding school dataset from Anonymous (1978) to evaluate PINNs' ability to perform parameter estimation and compartment approximations on real-world infectious disease data. The dataset records the size of the infected population in a school of N = 763 students over the course of 14 days. Due to the lack of data for the Susceptible and Recovered compartments, we only use the infected population to calculate data prediction error. Using λ = 0.5 as recommended by our previous results, we achieve the following estimation (See Figure 7) after 10 minutes of training on an NVIDIA RTX 2060 GPU.
Despite being provided with only infected data, PINNs is still able to use the mathematical structure of the data to infer the dynamics of S(t) and R(t) as well.
Analyzing the loss landscape during training (See Figure 8), we see that the disease parameters descend to one local minimum over the course of training, with a final approximation of β ≈ 0.0228 and α ≈ 0.454.
Suppose the London Boarding School dataset had an equal number of total and infected girls and boys * at t = 0. Fixing α = 0.5, μ = 0, 1 = 0.75, 2 = 0.5, p i = 0, a 1 = 1.5, and a 2 = 2.5, we can again generate and sample training data. We constrain all parameters to the range [0, ∞), except i which is further restricted to [0, 1]. To avoid numerical computation errors when the parameters are exactly 0, we pad each range by a small δ = 1.0 × 10 -9 such that
After ≈ 1.5 × 10 5 training iterations, we estimate the following dynamics (See Figure 9). We again compute the associated error for each approximated compartment and parameter using the simulated dynamics (See TABLE 3: Error Analysis of Preferential Mixing Estimation
While the MAPEs of R i and i are significantly higher, PINNs is still able to accurately estimate each compartment and parameter. As previously described, we employ gradient descent to calculate and minimize the effective reproduction number of the meta-population (See Figure 10).
* This distinction is arbitrary and can be replaced with any plausible division of the student population that results in different contact patterns and activity levels. † Note: The term "value" is only applicable for parameters, so N/A has been used for compartments. The partial derivative of R v with respect to each p i represents the need for vaccines within the sub-population. Formally, the proportion of vaccines that should be allocated to the i th subpopulation at time t is: Using this formulation, we construct a visual aid for policymakers to calculate and implement the best vaccine distribution proportion. At any given time, policymakers can use Figure 11 to look up the ideal proportions P i given the vaccination rate of one of the sub-populations, such as N 1 .
Further, for specific types of mixing, PINNs estimation can be greatly simplified.
1. If i = 1, the sub-population is isolated and can be separately approximated from the rest of the meta-population.
2. If i = 0 for 1 ≤ i ≤ n, all contacts are proportional and the i δ ij term can be omitted from c ij calculation.
3. If i = for 1 ≤ i ≤ n, the preferential mixing is homogeneous and all trainable variables corresponding to i can be replaced with one trainable variable . This also
4. If i = j for some i = j, the mixing is preferential and heterogeneous and no further simplifications can be made.
For example, assuming that the preferential mixing is homogeneous as described in Case 3, slightly transforms the contours of R 12).
Journal of Machine Learning for Modeling and Computing 4. DISCUSSION 4.1 Limitations Due to the diversity of PDE and ODE models describing infectious disease spread, PINNs can be applied to estimate a large number of simulation parameters, beyond those mentioned in this work. PINNs is also able to avoid local minima and converge to valid parameter combinations despite limited, noisy data. However, as the size of the differential equation system increases, in practice, PINNs may experience convergence issues. When approximating the preferential mixing model, PINNs sometimes fails to estimate i despite satisfying the stop loss -the loss value at which the model halts training. FIG. 13: Grid Search of i over PINNs Loss A grid search over i reveals a curve of i parameter sets that result in low values of L,
suggesting that an additional constraint or substitution can be made involving i that simplifies the estimation problem. Alternatively, the range of i could be further restricted from ∈ [0, 1] to a smaller interval (See Figure 13).
Despite achieving training times under 10 minutes on average, PINNs convergence times scale with system complexity. While additional compartments are able to account for multiple vaccination doses and other complex behavior, added differential equations and parameters also require additional training data and larger PINNs networks. PINNs are universal function approximators in theory Raissi et al. (2019a), but in practice, parameter estimation tasks are inherently limited to data availability.
This work establishes PINNs as a valuable method for informing infectious disease policy through parameter estimation. Combined with optimization approaches, PINNs is able to directly suggest optimal policy decisions, and while compartmental models are perhaps less versatile than ABMs, our coupled approach empirically requires less computational cost than traditional techniques. For example, with the vaccine distribution plan proposed, government officials are able to deploy the optimal vaccine distribution strategy that minimizes the spread of the infectious disease as fast as possible.
Reducing the number of parameters from 3n + 3 to 2n + 1 through substitution and assumption allowed PINNs to converge to a parameter solution, and hyperparameter experiments recommend a value of λ = 0.5 during training. Overall, PINNs is a powerful parameter estimation technique that can drive optimization, but PINNs frameworks should aim to first simplify mathematical formulations before tackling complex compartmental systems.
For vaccine distribution, we aim to re-parameterize the preferential mixing model to elucidate additional constraints between i , as conjectured from the previously discussed grid search.
Decreasing the number of parameters would facilitate faster parameter estimation and provide additional insight into the mathematical theory behind meta-population models.
By expanding PINNs to other compartmentalized models, parameter estimations for the latent period or vaccination-specific parameters could be computed. With additional compartments, we could also revise our assumption of perfect vaccine efficacy to account for re-infection. Faster PINNs parameter estimation could enable daily tracking of infectious disease parameters alongside raw reported data. With greater efficiency, PINNs could become the backbone for highly-available, cloud-based, machine-learning dashboards that help policymakers track and minimize infectious disease spread.
Volume x, Issue x, 2023
Journal of Machine Learning for Modeling and Computing