2.1 Spatial modeling in epidemiology‌‌

Epidemiological dynamics show spatial auto-correlation because of human mixing, as well as the dependence of some outbreaks on socioeconomic and demographic covariates which do not change erratically in space or time. The COVID-19 pandemic was observed and recorded with fine spatiotemporal granularity, and the data has been subjected to much spatiotemporal modeling. Many such studies found that the spread of the disease was mediated mostly by human mixing, rather than by socioeconomic factors; Huang et al. [13] found it to be so in Hubei province in China. Geng et al. [14] analyzed US data and found that spatial patterns’ length-scales ranged from the county-level to the nation; similar results were found by Schuler et al. [15] for Germany. McMahon et al. [16] analyzed data at the county-level and found that the correlation length-scales for spatial variability changed during the course of the pandemic; further, the correlation in epidemiological dynamics between urban centers were stronger than elsewhere. Indika et al. [17] analyzed data from the counties of Virginia and found that spatial auto-correlation of case-counts, as quantified by Moran test statistics, were impacted by, and linked to, executive orders at the state level. Thus, COVID-19 has presented us with much evidence of spatial auto-correlation in epidemiological dynamics.

The incorporation of spatial auto-correlation in the modeling and estimation of infection-rate fields is rare; however, it has been extensively used in disease mapping. In disease maps, one develops a field, called relative risk r_k, that is used to adjust an expected value of morbidity e_k to the locally observed case-counts in an areal unit (e.g., county) k, often via a Poisson or Negative Binomial link, i.e., . The expected value e_k is usually obtained from a regional average. The relative risk r_k is modeled using covariates of disease activity, e.g., socioeconomic conditions , where z_k are co-variate risk factors for areal unit k, are regression weights and captures auto-correlated random effects in space using a random field model. The simplest random field model is iCAR (intrinsic Conditional AutoRegressive [18]), a specific type of Gaussian Markov Random Field (GMRF). Thus

where W is the adjacency matrix of the areal units (i.e., if areal units i and j share a boundary). The object of estimation from data is . The precision matrix Q tends to be sparse. Another common model is the Besag-York-Mollie (BYM) model [19] which decomposes as and . An adaptation of the BYM is used in this paper. More details on spatial models is available in our preprint [4,6], where we also test the spatial model employed in the current paper on a smaller problem with three counties.

The method described in this paper is a spatial extension of a method in our previous paper [6] which was designed to provide short-term forecasts of an outbreak, conditional on daily case-counts data. One starts with a parameterized model of the time-varying (latent) infection-rate (f_inf) and convolves it with the cumulative distribution of the incubation period (F_inc) to obtain an integral equation for the for the case-counts (very similar to Eq. 2). f_inf and F_inc are fashioned as Gamma functions/distributions. This model was used to forecast COVID-19 case-counts in a number of countries and US states, each of them treated as a single areal unit. This method was then extended to multi-wave epidemics [4] and used to forecast outbreaks in California, New Mexico and Florida after multiple waves. These were high-dimensional estimation problems, as each wave had its own infection-rate parameters, and were solved using AMCMC. The spatial model for forecasting and detecting the “Fall 2020” wave on multiple areal units was developed in Refs. [8,9]. One of the innovation in Refs. [8,9] was the method used to model the spatial correlations in disease dynamics., driven by case-count data; it was found that, in NM, this was limited to nearest-neighbors among NM counties. Note that this will not hold true for other states with differently-sized areal units, population densities and mixing patterns. Refs. [8,9] also contain details on the GMRF model and how it was selected; the same GMRF is used in this paper. This infection-rate estimation was performed using AMCMC; due to its limited scalability, the estimation was limited to three adjoining counties in Ref. [8]. The estimation and the forecasts were very accurate and detection of the Fall 2020 wave was feasible within about a week’s worth of data. It was this study that motivated us to extend our work to all 33 NM counties by developing a scalable, if approximate, MFVI method that is described in this paper. It adapts the method in Refs. [8,9], via transforms, to be amenable for MVFI and unconstrained optimization.

The work done by Lawson and collaborators [10–12] is the closest to our own. Fundamentally we extend an older model of ours [4,6], meant for a single areal unit, to encompass multiple areal units. Each areal unit contains its own parametrized infection-rate representation, but are “stitched” together using a BYM model. The parameters are estimated by solving an inverse problem, conditioned on case-count data from areal units. Our model also includes a model for the incubation period. In contrast, Lawson and co-workers model case-counts directly; the clearest exposition of the model in in Ref. [10], and it has been used with COVID-19 data from South Carolina [20] and the UK [21]. Lawson and co-workers, much like us, have used the disagreement between calibrated model forecasts and data as signs of epidemiological anomalies and devised metrics such as the Surveillance Kullback-Liebler [22] (SKL) and Surveillance Conditional Predictive ordinate [23] (SCPO) to detect them.

2.2 Variational inference

Inverse problems for model calibration often require the approximation of intractable probability densities arising from Bayesian inference. A standard approach based on sampling is Markov Chain Monte Carlo (MCMC) [24] but suffers from scalability issues due to slow convergence rates for high-dimensional and/or multi-modal distributions [25,26]. Variational inference (VI) [27] provides an alternative to sampling techniques where approximate inference is recast as seeking a member of a family of approximating densities which minimizes a discrepancy measure such as KL-divergence. Originally developed for probabilistic graphical models [28] where some degree of analytical tractability is maintained, it has more recently been extended to many-parameter models, such as those seen in deep learning, through gradient-based iterative schemes adapted to a probabilistic setting [29–31]. These techniques are often termed Stochastic Variational Inference (SVI) and can exploit the automatic differentiation available in large ML models. They offer significantly improved scalability over MCMC while potentially sacrificing some approximation quality depending on the set of approximating distributions used. VI has seen successful applications to a number high-dimensional inverse problems in areas including medical classification [32,33] and segmentation [34,35], computer vision and image processing [36], natural language processing [37,38], and physics-based models [39,40].

In addition to this broad range of application spaces, VI has more recently been adopted in a growing number of epidemiological modeling challenges. Neural ODEs have been extended to a Bayesian setting where they can be calibrated with VI and applied to state space epidemiological models [41]. Model selection and dynamic causal modeling based on evolving real-world time series using VI have been applied to COVID-19 outbreaks [42,43] to provide online forecasting tools. Dynamical system inference for spatio-temporal modeling of infectious diseases using ODE and PDE formulations has been carried out at the state scale and combined with Bayesian neural networks [2]. Model calibration with VI has also been explored using alternatives to standard state-space models such as graph-coupled Hidden Markov Models [44]. Hassan et al. [45] developed a modified SEIR mathematical model to forecast COVID-19 trends. This study combined deterministic modeling with a Continuous-Time Markov Chain (CTMC) approach to calculate outbreak probability, emphasizing the critical need for coordinated public health measures. Kamrujjaman et al. [46] developed a modified SEIR mathematical model for COVID-19 that combines exposed and asymptomatic populations into a single compartment and includes the effects of panic, tension, and anxiety on mortality rates across susceptible, exposed, and infected classes. They identified disease transmission rates and panic-related mortality as the most sensitive parameters. In addition to predictive model calibration, generative probabilistic modeling using VI-based approximations of distributions has also been explored to generate mission information about disease spread [47].

3.1 Epidemiological model

The epidemiological model is defined by a spatio-temporally varying infection-rate model and an incubation model given by

(1)

where the infection-rate f_inf is a Gamma distribution with shape and scale parameters k^r and , respectively. Note that 1 ≤ r ≤ R indexes the spatial region. is the cumulative distribution of the incubation period for COVID-19, taken from Lauer et al. [49]; further details to the temporal model are in [4,6]. The parameter represents the start of the outbreak and will be inferred along with the infection-rate parameters. The number of people that turn symptomatic over the time interval [t_i−1, t_i] is given by

(2)

so that and i represents the time-dependence of the predictions. Here, N^r is the fourth and final region-dependent parameter and represents the total number of people infected during the entire epidemic wave in spatial region r normalized by the population of region r.

In the discussion above, f_inf was chosen to be a Gamma distribution, primarily for simplicity of parameterization. For more complex temporal behavior of the latent infection-rate, more involved parameterizations may be required, e.g., the parameterization for multi-wave outbreaks, as described in Ref. [4].

The noisy model predictions are defined as

(3)

Here , where , are the region-specific model parameters for and is an epidemiological / disease model that uses . In addition, is the vector of observed case-counts on day i for all R regions, is the vector of R case-count predictions for the same day using the model and is the “noise” observed case-counts on day i that cannot be explained by (mostly reporting errors). To account for spatial correlations and heteroscedastic noise seen in case-counts, the noise is assumed to be composed of two terms where the first is given by a Gaussian Markov Random Field (GMRF) model while the second represents temporally-varying, independent Gaussian noise. Letting and represent the data and parameters, respectively, the likelihood then takes the form

(4)

where is given by

(5)

The first term in Eq. (5) forms the precision matrix of a GMRF component of the noise where the strength of correlations induced by adjacent regions is governed by . The relative topology of regions is encoded by W, the county adjacency matrix, defined as

(6)

Here, where g_i is the number of regions adjacent to region i. Note that this formulation assumes that the R spatial regions are self-contained with negligible population mixing with the region around them. The derivation of the adjacency matrix W and its subsequent inclusion in is presented in detail in Refs. [8,9], and we reproduce a summary here. In Refs. [8,9], we plotted the population-normalized case-counts, summed over 90-day periods, for all the NM counties and noticed a spatial correlation between counties along the populated Rio Grande valley; furthermore, the spatial correlation persisted over a moving 90-day period. We examined the spatial correlation between a county and its neighborhood, where a neighborhood was defined as the set of counties that shared a border with the county in question, i.e., “one-hop” neighbors. A two-sided Moran’s I-statistic test revealed that the spatial correlation was supported by the evidence. We also considered a neighborhood of one- and two-hop neighbors and the same test found no evidence for it. This allowed us to design an adjacency matrix consisting of immediate neighbors only and formulate a simple precision matrix for the GMRF model for the correlated errors (defined as the difference between observed case-counts on a day i in a given county and model predictions ). We also add a heteroscedastic reporting error to the diagonal of the GMRF’s covariance matrix equal to . For the purposes of this paper, we have modeled NM in isolation, i.e., the adjacency matrix W does not contain terms that link counties along the NM border to the neighboring counties in adjoining states. This is an approximate boundary condition of no population mixing outside NM which is necessary to limit the scope of the problem, else W would grow unbounded. Note that the counties along the NM boundary are large, sparsely-populated, mountainous or desert regions with little mixing with the population outside, making feasible this approximation.

The second term captures prediction-dependent, uncorrelated noise with additive and multiplicative components governed by and , respectively. The relative contribution between the correlated GMRF noise and the uncorrelated noise is controlled by . Full details of the development of P can be found in Refs. [8,9]. Using Bayes’ rule, one can obtain an expression for the posterior density of , the parameters to be estimated from data, as:

where is defined in § 3.2 and . For low-dimensional problems, the density can be obtained by sampling directly. This was performed in Ref. [8,9] for three adjoining NM counties (Bernalillo, Santa Fe and Valencia) using AMCMC [7] as a test of the correctness of the formulation. Using data from the Summer wave, we developed PDFs of , and and forecast the case-counts for two weeks beyond September 15^th, 2020, in order to detect the arrival of the Fall 2020 wave (the forecasts beyond September 15 and the observed case-counts would disagree). We could detect the arrival of the Fall wave with a week’s worth of observations. We also tested the method using data gathered till August 15^th to check the detector’s susceptibility to false positives. However, due to the lack of scalability of AMCMC (in terms of parameters being estimated and parallel computing), we cannot use it to estimate the infection-rate field in all 33 counties of NM and therefore take recourse to a mean-field variational inference techniques that, though approximate, will scale to the 126 parameters that the NM-wide estimation will require.

3.2 Statistical inference

The set of parameters defining the likelihood in Eq. (4) is given by

(7)

where are the global noise parameters. Inference consists of forming the posterior distribution over uncertain parameters . As the posterior is intractable, we instead look to approximate it using VI. Hence, the following sections describe how VI is formulated carried out to approximate the posterior for the outbreak model as well as how the prior is defined to regularize the inverse problem.

3.2.1 Variational inference.

We will compare the Bayesian posterior sampled with AMCMC with posterior models obtained using MFVI which recasts approximate inference as an optimization problem. In particular, as the exact posterior is intractable, we consider a family of approximating densities and seek to find a density that minimizes the KL-divergence with respect to the posterior

(8)

This can be re-expressed as minimizing the objective function based on the evidence lower bound (ELBO) [30]

(9)

where the first term in Eq. (9) is the entropy of the surrogate posterior and the second, data-dependent term is an expectation with respect to the surrogate posterior that reflects both the expected data-fit and the prior. Here we take to be the set of mean-field Gaussian distributions, i.e.,

(10)

where , . We arrive at an optimization problem over 2d parameters where d is the number of parameters defining the epidemiological model . To carry out the above minimization problem, we aim to use a gradient-based iterative scheme as the expectation in Eq. (9) cannot be evaluated explicitly due to the nonlinearity of the forward model. Furthermore, is potentially a non-convex objective. Note that the gradient and expectation operators do not commute, i.e.,

so some care has to be taken to arrive at a Monte Carlo estimator for the gradient . Two widely used approaches are: (a) the score function estimator, described in § A.1 (in the Appendix), which forms the basis of black-box VI and requires only evaluations of the log-likelihood, and (b) the reparametrization approach which requires gradients of the log-likelihood. The score function estimator typically displays much larger variance as seen in Kucukelbir et al. [50] where two orders of magnitude more samples were needed to arrive at the same variance as a reparametrization estimator. A similar trend was confirmed for the outbreak problem (Fig 9) suggesting that the reparametrization approach would lead to superior scalability. Reparametrization proceeds by expressing as a differentiable transformation of a -independent random variable such that . This allows the gradient to be expressed as

(11)

where gradients of the entropy term in Eq. (11) are available analytically for the Gaussian surrogate posterior and the second term can now be approximated with Monte Carlo given a method to compute the required gradients. For many machine learning models, automatic differentiation can be exploited to calculate the gradient of the log-likelihood with respect to parameters . Here, the objective function involves the log of the likelihood (Eq. (4)) where derivatives of matrix inverses and determinants with respect to parameters are required to compute the gradient. Gradients such as these are not available using most automatic differentiation libraries. Instead, matrix calculus and quadrature were used to compute the derivatives of the log-likelihood with respect to model predictions and to approximate the derivatives of the model predictions with respect to parameters, respectively. For details, see § A (in the Appendix).

Note that some of the parameters comprising are required to satisfy constraints for the noise and epidemiological models to be well-defined. For example, noise parameters , , and as well as model parameters N^r, k^r and for should be positive while should satisfy . Sampling from the mean-field Gaussian (Eq. (10)) during the Monte Carlo estimation of the gradient (Eq. (11)) may result in violations of these constraints. To maintain the required properties without resorting to constrained optimization, we express a constrained parameter as an invertible, differentiable transformation of an unconstrained . Hence, the distribution governing is the push-forward density of through f_i, i.e., the components of the mean-field surrogate posterior (Eq. (10)) have modified probability densities

(12)

where . This results in mean-field approximation where some of the factors are Gaussian and others non-Gaussian. Each factor is still defined by a and parameter. The transformations are listed in Table 2.

The initial motivation for using Gaussian distributions can be seen in Fig 5 where the posterior approximations from AMCMC display negligible skew and kurtosis suggesting a variational distribution with vanishing higher-order moments is sufficient to capture the behavior of interest. Second, as the goal was to produce a method for efficient high-dimensional inference, scalability and analytical tractability had to be balanced with accuracy tradeoffs. Note that MFVI also has a tendency to underestimate the uncertainty / variance in the estimated parameters [51] and in § 4, we will check if this poses a limitation on the usefulness of the approximate solution, especially in the predictive skill of model and therefore the detection of the start of outbreaks.

3.2.3 Posterior predictive tests.

The mean-field variational inference (MFVI) described in § 3.2.1 results in a multivariate Gaussian posterior distribution (Eq. (12)) that then needs to be verified against data . To do so, we take samples , and using Eq. (3), generate predictions . For this paper, J = 100. The time-series Y_j result in a “fantail” of predictions for each areal unit r which should statistically reproduce , e.g., the inter-quartile range of the samples Y_j should bound 50% of the observations and the 5^th and 95^th percentiles should bound 90% of the individual data points in . This test is called a posterior predictive test (PPT) which we will use extensively in § 4. In addition, we define a score to summarize the agreement of Y_j with . Predictive distributions constituted out of samples are called “push-forward“ (PF) predictions.

Define to be the model prediction for region r for day i by combining Eq. (3) and Eq. (7). Let be the prediction corresponding to a sample drawn from the posterior. Let be the corresponding observation (from ). Let be the cumulative distribution function (CDF) for the model predictions arising from the posterior distribution . The Continuous Ranked Predictive Score (CRPS [52]) is defined as

(13)

The CRPS has units of case-counts and will be larger for areal units with larger total case-counts , and we will use the ratio to compare across areal units in § 4. In practice, the empirical CDF computed using J samples of Y_j is used to approximate . This score function has been used in judge the quality of the model to capture the spread in the data used for calibration. [6,53,54] We use the implementation in the R Statistical Software [55] (R version 4.3.2 (2023-10-31)) package verification [56], specifically the function crpsDecomposition(), to compute the CRPS.

4.1 Three-county inversion

In this section we investigate the effect of using an approximate MFVI inversion by comparing against AMCMC solution [8,9]. We also check the effect of estimating the infection-rate across multiple areal units vis-à-vis independently.

4.1.1 Joint versus independent calibrations.

First, we perform an infection-rate estimation for three adjoining counties – Bernalillo, Santa Fe and Valencia – using MFVI independently and jointly using the spatial (GMRF) model (Eq. (5)); see Fig 2 for their positions. In Fig 3, the results of a PPT for all three counties are displayed in both the joint and independently calibrated cases. The case-count data were smoothed with a 7-day running average. The median prediction is plotted with the red line and the dashed lines denote the 5^th and 95^th percentiles. We see that most of the observations (filled symbols) are within these bounds. We also see the 2-week-ahead forecast beyond September 15th and the unfilled symbols showing the observed data from that period. The forecasts and the data do not match for any of the three counties, indicating a change in the epidemiological dynamics. This change is due to the arrival of the Fall 2020 wave of COVID-19 in NM, and the figure shows that the wave arrived approximately simultaneously in all three counties (see Fig 1 (left) for a clearer picture of the three COVID-19 waves encountered in 2020 in NM). Observe that the jointly calibrated predictive distribution displays less uncertainty in the predictions (i.e., small ) than the independently calibrated version. This is likely due to a combination of two factors. The first is that incorporated spatial correlations between counties regularizes the calibration and results in more certainty about the true underlying model parameters. The second is that the uncorrelated MFVI approximation is known to underestimate uncertainty for highly correlated distributions. Note also that uncertainty is largest in the predictive distribution for Santa Fe consistent with the markedly noisier behavior of the case-counts for this county. The infection-rate profiles for these counties are in Fig 14 in the Appendix. There is not much difference between them indicating that the infection-rate parameter estimates in the two cases might be similar, whereas the estimate of the noise, which affects PPT results, vary between the two formulations due to the fact that the spatial noise model has to accommodate all the noise, across all the areal units, together. These findings echo what we observed when the same study was performed using AMCMC as the estimation procedure [8,9].

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. The counties of New Mexico.

The shaded ones are where we will present results. Abbreviations: B = Bernalillo; CI = Cibola; CU = Curry; DA = Doña Ana; RA = Rio Arriba SF = Santa Fe and V = Valencia. The shapefiles of the counties were downloaded from the US Census Bureau website [59].

https://doi.org/10.1371/journal.pone.0350090.g002

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Comparison of the predictive distribution for the MFVI inversions of Bernalillo (left), Santa Fe (middle), and Valencia (right) done jointly using the GMRF model and independently for each county.

The solid lines show median predictions and the dashed lines bound the 5^th–95^th quantile interval. Case-count data was smoothed with a 7-day running average. The filled symbols are the data used in the inversion. The unfilled symbols are the observations beyond September 15^th and are used to compare with the two-week forecasts.

https://doi.org/10.1371/journal.pone.0350090.g003

4.1.2 AMCMC versus MFVI estimation.

Next, we study the effect of using our approximate MFVI method versus the estimates computed using AMCMC [8,9]. In Fig 4 we plot the PPT results from the joint MFVI (top) and AMCMC (bottom) for the same three counties. Here we see that the MFVI estimates a larger – the 5^th and 95^th percentile bounds (dashed lines) are far wider for MFVI results vis-à-vis AMCMC results below. By dint of having wider bounds, the MFVI estimate is also better at bounding the data used to compute the infection-rate field for the three counties. Again the arrival of the Fall 2020 wave is clearly discerned in the figure. The infection-rate profiles for these counties are in Fig 15 in the Appendix and there is not much difference between them.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 4. Comparison of the predictive distribution for the joint inversion of Bernalillo (left), Santa Fe (middle), and Valencia (right) done jointly using MFVI and AMCMC.

The solid lines show median predictions and the dashed lines bound the 5^th–95^th quantile interval. Case-count data (filled and unfilled circles) are the same as in Fig 3. The AMCMC results are taken from Refs. [8,9].

https://doi.org/10.1371/journal.pone.0350090.g004

In Fig 5, we summarize the marginalized posteriors for the infection-rate parameters for Bernalillo, Santa Fe, and Valencia computed using MFVI and AMCMC, jointly and independently. assumes negative values as it is measured from June 10^th, 2020; the PDFs peak around −20 (for AMCMC results) and imply that the infections for the Summer wave started in late May, about 20 days before June 10^th. The MFVI results also agree approximately with this estimation. We see that apart from N_r, the MFVI and AMCMC posteriors do not match. The MFVI posteriors are extremely narrow, providing a spurious degree of certainty in the estimates; this arises from the form of the posterior distribution – independent Gaussians – that we postulate in Eq. 10. In addition, as is clear from the AMCMC results, the “true” posteriors are not Gaussian. The marginalized posteriors computed via AMCMC, for joint and independent estimation, do match (sometimes very well, as in the case of Valencia), but they are wide apart for the MFVI, showing the effect of the Gaussian approximation. Yet the infection-rate profiles in Fig 14 (in the Appendix) do not show much of a difference, nor do the PPT results in Fig 3, leading us to conjecture that the influence of some of these parameter on the infection-rate may be muted. This can also be surmised from the marginal distributions computed from AMCMC – they are quite wide.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 5. Comparison of the posterior over model parameters t₀, N, k, and for 3-county inference of Bernalillo (left), Santa Fe (middle), and Valencia (right) using both MCMC (orange) and MFVI (blue).

Both joint (solid lines) and independent calibration (dashed lines) are results are displayed. t₀ values are negative as it is measured from June 10^th, 2020, and the PDFs imply that infections for the Summer wave started in late May. ’VI indep.’ in the legend implies an estimate obtained for each county independently using MFVI.

https://doi.org/10.1371/journal.pone.0350090.g005

The difference in the posterior densities of the parameters (Fig 4, joint estimation only) deserves an explanation. They arise from the assumptions of independence inherent in MFVI. In reality, the posterior density is strongly correlated, and these can be computed using AMCMC (see Refs. [8,9]). The Gaussian assumption in MFVI also plays a role, but it is of lower consequence as the AMCMC posteriors are not excessively skewed, and a Gaussian is an acceptable approximation.

Finally, we summarize the predictive skill of the joint estimates computed using MFVI and AMCMC in Table 1 using CRPS. The CRPS, computed over data between June 1^st and September 15^th, 2020, summarizes the agreement of the PPTs with observations for each of the counties. We see that the CRPS (a measure of the error between predictions and data) is about 10 cases per day, for Bernalillo (where case-counts peaked at about 100 cases/day; see Fig 3) and 2.5 cases a day for Valencia and Santa Fe (which peaked at about 12 cases a day). What is remarkable is the difference in the CRPS as computed using MFVI and AMCMC – it is less than a case-count per day. This small change in the predictive skill leads us to believe that the scalable, but approximate, MFVI approach might be sufficiently accurate to allow us to detect the arrival of the Fall 2020 wave in an automated manner.

Download:

• PNG
  larger image
• TIFF
  original image

Table 1. Predictive skill of PPTs generated using infection-rate estimates computed using different procedures. All estimations are performed using the joint formulation, using the spatial model. The PPTs are scored using CRPS. They have units of “case-counts”.

https://doi.org/10.1371/journal.pone.0350090.t001

To summarize, from Figs 3–5 it is clear that:

• Joint estimation performed with a GMRF model leads to more certain forecasts compared to independent estimation of infection-rates in areal units (i.e., counties).
• Marginalized posterior obtained from MFVI shows spuriously low levels of uncertainty compared to AMCMC.
• Despite the difference in the posterior densities computed using MFVI and AMCMC, the forecasts are not very sensitive to them; this is because the differences in the infection-rate are compensated by the reporting errors’ parameters and . However, the forecasts performed using MFVI’s posterior densities are slightly more uncertain. Therefore anomaly detection with them will also be a little less accurate.

4.2 Joint inversion of all NM counties

Next, calibration of the full 33-county model using MFVI was carried out, followed by PPT runs which were then summarized using CRPS computed for each NM county. In Fig 6, we plot the CRPS normalized by the total number T of cases a function of T, for all counties. Data between June 1^st and September 15^th was used in the infection-rate estimation as well as the computation of CRPS is a measure of the “goodness of model fit” to data. We see that the CRPS, as a ratio of the total cases, decreases with increasing number of total cases, as the disease model fits larger outbreaks in counties like Bernalillo. Others, like Cibola, do not agree with model predictions, due to flaws in the data, as we will see later in Fig 12. The straight line fit to data has the form

Download:

• PNG
  larger image
• TIFF
  original image

Fig 6. CRPS/T plotted as a function of total number of cases T.

Both the axes are log-axes. We see, as expected, that the predictive skill of the disease model, post calibration, is better for counties with larger case-counts T where noise variance is low. The horizontal dashed lines are the first and third quartiles of .

https://doi.org/10.1371/journal.pone.0350090.g006

This equation shows that scales (approximately) as the fourth root of the total number of cases, a slow reduction indeed.

The horizontal lines in Fig 6 show the first and third quartiles of . The three counties studied in § 4.1, Bernalillo, Santa Fe and Valencia, are marked and fall in the lowest quartile, i.e., their data is good and the calibrated disease model is predictive. We will use the county of Doña Ana as another member of the “good” class, while Curry and Rio Arriba, which fall in the inter-quartile range of , will serve as exemplars of the “middling” class of calibration. Cibola, which falls in the last quartile, will be an exemplar of the “bad” class of calibration. We now examine the quality of the inversion in each of these counties. In Fig 7 we plot the infection-rate parameters, as estimated from the Summer wave case-count data, for 7 counties marked in Fig 6. We see that the standard deviations are spuriously tiny, in line with what was observed in Fig 5; thus MFVI consistently underestimates the uncertainty in the parameter estimates. Further, this spuriously low uncertainty is pervasive – counties with high CRPS such as Cibola and Rio Arriba show much the same estimation uncertainties as counties such as Bernalillo and Santa Fe with CRPSs a factor of three smaller. Thus the uncertainty in the parameters’ estimates do not seem credible and we will omit them from further discussion. However, despite the quality of the case-count data, the inversion completes stably and provides plausible results. However, as Fig 6 shows, the infection-rate may not be estimated very accurately in some counties and this might hamper the task of detecting the Fall 2020 wave reliably.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 7. Means (light gray) and standard deviation (dark gray) of various infection-rate parameters, for select counties marked in Fig 6.

Top left: The start time of the Summer wave t₀. Top right: The total size of the Summer outbreak N. Bottom left: k, the shape parameter of the Gamma profile of the infection-rate. Bottom right: θ, the scale parameter of the Gamma profile.

https://doi.org/10.1371/journal.pone.0350090.g007

In Fig 8 we compare some infection-rate parameters for Bernalillo, Santa Fe and Valencia, in the 3-county (as plotted in Fig 5) and 33-county inversions. MFVI was used for both the computations. We see that the parameters are not the same and there does not seem to be a clear trend, except that increases when we include counties (some with poor quality data) in the estimation. This implies that the PPTs for the “good” counties will likely be wider than what could be achieved with AMCMC and this might have repercussions regarding detecting the Fall 2020 wave.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 8. Select infection-rate and noise parameters, estimated jointly among three counties, compared with their counterparts from a 33-county inversion.

The noise parameter is markedly larger in the 33-county inversion.

https://doi.org/10.1371/journal.pone.0350090.g008

To summarize, the uncertainty in the estimated parameters, computed using MFVI, are not very credible. However, the MFVI inversion is stable, though the uncertainty in the PPTs will be larger because of an inflated , required to accommodate counties with very noisy data.

4.3 Algorithmic results

Finally, we investigate the numerical aspects of the reparametrized algorithm described in § 3.2.1. The convergence of MFVI is depicted in Fig 9 where the ELBO and the norm of its gradient are shown as a function of gradient descent iterations.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 9. (Top) Convergence of the ELBO for the 33-county inversion along with the norm of the ELBO gradient as a function of gradient descent iterations for the reparametrized gradient formulation of MFVI.

(Bottom) Convergence of the ELBO for a 1-county inverse problem along with the norm of the ELBO gradient for the black box formulation of MFVI. In both cases, samples were used for the MC estimators of the gradient.

https://doi.org/10.1371/journal.pone.0350090.g009

The top of Fig 9 shows the ELBO and gradient for the 33-county calibration using the reparametrization formulation while the bottom provides a comparison to using black box VI for 1-county calibration of Bernalillo. In both cases, samples were used for the reparametrization and score function MC estimators of the gradient. Note that even for a single county, black box VI shows significantly higher variance in the gradient leading to poor convergence in comparison to the much larger 33-county problem calibrated with reparametrization. The convergence of the ELBO in the top row is also quite smooth suggesting that significantly less samples could be used to obtain good estimates of the gradient with the reparametrization approach. Hence, it is clear that reparametrization is necessary to scale the calibration to the 33-county inversion despite the added complexity of obtaining gradients of the log-likelihood. Note that the ELBO shows a sudden drop-off in Fig 9. High-dimensional, nonlinear objective functions tend to be non-convex and such sharp drop offs are common. This is not seen as often with lower-dimensional problems which explains why we see it only in the 33-county inversion.

The top two rows of Fig 10 display convergence information from two counties, Bernalillo and Rio Arriba, taken from the 33-county inversion. Bernalillo and Rio Arriba were chosen as they have larger and smaller populations, respectively. The mean of the initial condition for MFVI is given by a MLE solution shown in red. Intermediate solutions are shown in blue along with the final solution in green. Observe that while similar to the MLE, MFVI subtly expands the shape of the wave to better cover the tail of the outbreak. This is potentially an effect of the tendency of the KL-divergence to increase the overlap of the surrogate and true posterior distributions at some expense of the mean prediction fitting the data less accurately. Comparing Figs 3–10, we can see that the coupled, 33-county inversion introduces a bias in the parameter estimates for Bernalillo. This is due to the multitude of NM counties that are sparsely populated and display significant noise in their daily case counts. Despite this effect, the outbreak detection performs well on the 33-county inversion data suggesting that the predictive uncertainties provided by the calibration remain informative.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 10. (Top) Convergence information corresponding to 33-county inversion using MFVI.

Results corresponind to Bernalillo (left) and Rio Arriba (right). (Top) The MFVI mean initial condition is the MLE solution shown in red. Intermediate solutions are shown in blue and the final mean solution in green. (Middle) Corresponding convergence of the model parameters. (Bottom) The ELBO objective (left) and its gradient (right) as a function of iteration for the full 33 county inversion.

https://doi.org/10.1371/journal.pone.0350090.g010

To summarize, reparametrization provided us with the scalability needed to solve the high-dimensional inverse problem conditioned on data from all 33 NM counties.

5.1 Temporal detection

The argument above is used to fashion an outbreak detector using PPT runs. The detector works as follows. We sample the posterior in the same manner as for PPT runs, and use Y to compute an “outlier boundary”; we define it as the percentile of the forecasts. Y is computed using data from June 1^st to an end date (usually August 15^th or September 15^th) where we test for a change in the epidemiological dynamics. The actual test consists of comparing a two-week-ahead forecast with the data that was observed during that period. Any day with case-counts above the “outlier boundary” is deemed an outlier. Three consecutive outlier days cause an “alarm”, corresponding to an anomalous change in the disease dynamics. Using this detector, when solving the problem with AMCMC [8,9], we found that we could detect the arrival of the Fall 2020 wave correctly, when tested using data up to September 15^th. In Fig 11 we repeat the same test, but the infection-rate is estimated using MFVI. The top row depicts the outbreak detector being applied beyond September 15^th, 2020. We see outliers and alarms for all three counties with a week of September 15^th, i.e., the Fall 2020 wave was easily detected within a week’s worth of data. The bottom row repeats the test, but applied to August 15^th. We see that Santa Fe and Valencia incur false positives, whereas Bernalillo does not. This implies that the approximations in MFVI may lead to erroneous detections for borderline cases (i.e., counties with small-count data and high variance noise) but areal units with large case-counts might be unaffected. The reason for the false positives is simple – the data is very noisy and the outbreak detector makes no attempt to reduce the variance in the noise. Comparing the correct detection on the top row of Fig 11 with the false positives in the bottom row, we see that outliers and alarms are plentiful after September 15^th and sporadic after August 15^th. A slightly more sophisticated detector that performed temporal averaging could eliminate such false positives. We will address this below, using a simple spatio-temporal method.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 11. The infection-rate detector implemented using the infection-rate from the joint MFVI estimation of three counties (corresponding to Fig 4 for Bernalillo (left column), Santa Fe (middle column), and Valencia (right column).

The symbols are the case-counts, the solid blue line the median forecast with the infection-rate, post-calibration and the solid red line is the alarm boundary. Outliers are circled and an alarm, corresponding to three consecutive outliers, is encased in a square. The vertical line is the point beyond which we forecast and thus test for the arrival of the Fall 2020 wave. Top row: Testing for arrival after September 15^th, 2020. Bottom row: Testing for arrival after August 15^th, 2020.

https://doi.org/10.1371/journal.pone.0350090.g011

In Fig 12, we demonstrate the outbreak detector on a “good” (Doña Ana), two “middling” (Curry and Rio Arriba) and one “bad” county (Cibola), as determined using CRPS in Fig 6. The plotting conventions are the same as in Fig 11. We see that we correctly detect the Fall 2020 wave when testing on September 15^th and do not incur any false positive on August 15^th. This performance is rather fortuitous for Rio Arriba and Cibola, which have unexplained case-count peaks in mid-July (Rio Arriba) and in late July (Cibola). It would be difficult to estimate an infection-rate profile for either of these counties independently (this is especially true for Cibola, where the reported case-counts show no wave-like structure); however, our spatial GMRF model regularizes the inversion and constructs an infection-rate stably. This also shows the necessity of smoothing in time guided by a process immune to flaws in the data – in our case, this is achieved via the incubation period distribution (F_inc) and the infection-rate profile f_inf (see Eq. 1). (This smoothing would be far more difficult with a purely data-driven approach such as a pCAR-ST, which would attempt to accommodate the July peaks in the two counties.) Forecasts using this infection-rate do not match the case-count data at all, as is clear in the plots for Cibola, but provides us with stable PPT results and a perhaps meaningless detection. However, numerical and algorithmic stability in the face of poor quality data encourage us to believe that the MFVI algorithm and outbreak detection can be automated.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 12. Performance of the outbreak detector for a “good” (Doña Ana), two “middling” (Curry and Rio Arriba) and one “bad” county (Cibola), as determined using CRPS in Fig 6.

Left column: Testing for Fall 2020 wave arrival on September 15^th, 2020. Right column: The same, but detection performed on August 15^th. We do not see any false positives in the right column.

https://doi.org/10.1371/journal.pone.0350090.g012

5.2 Detecting spatial patterns

The results above show that the MFVI estimates of the infection-rate have the ability to detect the Fall 2020 wave, especially for the counties in the first quartile (plotted in Fig 6), though our crude detector, which does not smooth / de-noise the data, might suffer from false positives due to high variance noise in observed case-count. However, the MFVI infection-rate field allows prediction in space and time, which allows spatio-temporal assimilation of data. We address this next.

Define “exceedance” for region r and day i. Here is the observed case-counts and is the alarm boundary (the 99^th percentile computed from Y). would denote an outlier. Let

where N_smooth is a time-period over which we will average out the temporal noise and detect spatial patterns. Fig 13 (top row) shows a map of computed over a two-week period after August (on the left) and a two-week period after September (on the right) when the Fall 2020 wave had arrived. We see from the colormap that the before the arrival in all counties except one; specifically, the false positives seen in Fig 11 for Santa Fe and Valencia are no longer visible. After the arrival, we see clear spatial structures on the right. Thus simply averaging the data in time provides a more robust detection, though with a loss of timeliness. However, we see two adjacent counties (in shades of yellow) with very high .

Download:

• PNG
  larger image
• TIFF
  original image

Fig 13. Top: Plot of mean exceedance for NM, where is averaged over August 15^th and August 31^st (left) before the arrival of the Fall 2020 wave and September 16^th and October 1^st (right), when the Fall 2020 wave had arrived.

Bottom left: Dendrogram from hierarchical clustering, with the cut at a height of 0.6, resulting in clustering of counties. Bottom right: Disease clusters from the dendrogram.

https://doi.org/10.1371/journal.pone.0350090.g013

Epidemiological activity spreads due to population mixing and is largely between adjoining NM counties because of their large spatial expanse. Therefore adjacent counties are expected to have similar epidemiological behavior. If the epidemiological activity is represented using a derived / inferred / estimated quantity, and adjoining areal units differ greatly, it is either a consequence of erroneous estimation or an artifact of noise in data, i.e., spatial auto-correlation might be a simple mechanism for removing erratic behavior in data. To this end, we subject the map in Fig 13 (top right) to clustering. The centroids of the counties serve as the spatial feature whereas serves as a measure of how strong the Fall 2020 wave is, in each county. The data was Z-scored and subjected to hierarchical clustering using the R Statistical Software [55] (R version 4.3.2 (2023-10-31)) package stats, specifically the function hclust(). The resulting dendrogram is in Fig 13 (bottom left), which is cut at a height of 0.6 (corresponding to a quantile of 0.15) to reveal 3 clusters of counties. These are plotted in Fig 13 (bottom right) and reveal the same clustering as was seen in Fig 13 (top right). However, the counties in yellow were eliminated – their level of , though high, were too different and thus violated the need for spatial auto-correlation. Two of the clusters surround the county of Bernalillo, the population center of NM, where COVID-19 was very active in the city of Albuquerque.

Temporal epidemiological anomaly detections and spatial clusterings can be forged into a rough-and-ready test for the credibility of an outbreak detection. Questions of detection credibility arise due to the quality of the data used in outbreak detection and the approximations inherent in MFVI, which lead to narrow posterior distributions. Diseases spread, and an areal unit which shows persistently high disease levels while surrounded by quiescent areal units, is difficult to explain. This unexplainable behavior may simply be an erroneous detection and thus could be ignored.

A Variational Inference

A.1 Score gradients of the ELBO.

We briefly review the score estimator, or black-box, approach to estimating the gradient of the ELBO. Recall that the ELBO is given by 9 which, for the sake of clarity, can be written in a more generic form

(14)

where encapsulates the dependence on the random vector . To derive an estimator for the gradient, we can carry out the following manipulations

(15)(16)(17)(18)(19)

Hence, the gradient can be expressed as an expectation with respect to where only the log of the surrogate posterior needs to be differentiated with respect to the variational parameters .

A.2 Reparametrization gradients of the ELBO.

The likelihood and log likelihood are given by

(20)(21)

Using the reparametrization trick, we can write the ELBO (9) and its gradient in the form

(22)(23)

where , with . Here, is a positive transformation of the unconstrained variable to ensure the variance is constrained to be positive. A Monte Carlo estimator of the gradient can then be written as

(24)(25)(26)

where the last line is given by the chain rule and the fact that is defined by an element-wise transformation of . Observe that

(27)(28)(29)(30)

so that it remains to compute the gradients of the log-likelihood and prior.

A.3 Gradients of the log likelihood.

As the log likelihood factors independently across the data , it suffices to compute the gradients and for a particular day i. The differentials of the these two terms are computed using matrix calculus [64] as

(31)(32)(33)

where the differential of is

(34)(35)

We have the following derivatives:

Case: are the unconstrained model variables for region r and is a particular variable.

(36)(37)(38)(39)(40)

Case:

(41)(42)(43)(44)

Case:

(45)(46)(47)(48)

Case:

(49)(50)(51)(52)

Case:

(53)(54)(55)(56)

The variable transformations are listed in Table 2.

Download:

• PNG
  larger image
• TIFF
  original image

Table 2. Variable transformations. is the softplus function and satisfies .

https://doi.org/10.1371/journal.pone.0350090.t002

A.4 Approximation of model predictions and gradients via quadrature..

The model predictions y_i, given by (2), involve a convolution integral that cannot be expressed in closed form. Hence, we approximate the predictions by integral quadrature

(57)(58)(59)

where w_j, are quadrature weights and points given by a method such as Gaussian quadrature. As the function does not depend on parameters , we can write it as for simplicity of notation. By the Leibniz integral rule, we can write the derivatives of the model predictions as

(60)(61)

This requires that the functions and are continuous in and in a region of the plane including , a condition that’s easily met by ensuring certain constraints are satisfied by the parameters for . These constraints are enforced via the variable transformations in (12). Hence, we can approximate (60) and (61) via quadrature in the form

(62)(63)

Hence, in this implementation of VI, gradients of the ELBO have a pseudo-analytic form where “outer” gradients of the log likelihood with respect to model predictions are exact and “inner” gradients of the model predictions with respect to parameters are approximated via quadrature. This allows for accurate gradient approximations that can be calculated efficiently leading to a scalable VI algorithm that can be applied to the high-dimensional inverse problem for the outbreak model.

B Infection-rate estimates

4.1 describes the estimation of the infection-rate field over Bernalillo, Santa Fe and Valencia using MFVI jointly (where the GMRF spatial model is used) and individually. Fig 3 plots the PPT runs, driven by a distribution of infection-rate fields. The corresponding fields are in Fig 14.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 14. Comparison of the infection rate curves that determine predictions in Fig 3.

These are for a 3-county inference of Bernalillo (left), Santa Fe (middle), and Valencia (right) done jointly using the GMRF model and independently for each county. The solid lines shows the median predictions and the dashed lines are the 5^th and 95^th percentiles.

https://doi.org/10.1371/journal.pone.0350090.g014

The same section describes the estimation of the infection-rate field over Bernalillo, Santa Fe and Valencia using AMCMC and the MFVI. Fig 4 plots the PPT runs, driven by a distribution of infection-rate fields. The corresponding fields are in Fig 15.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 15. Comparison of the infection rate curves that determine predictions in Fig 4.

These are for a 3-county inference of Bernalillo (left), Santa Fe (middle), and Valencia (right) done jointly using the MFVI and AMCMC (taken from Refs. [8,9]). The solid lines shows the median predictions and the dashed lines are the 5^th and 95^th percentiles.

https://doi.org/10.1371/journal.pone.0350090.g015