How to best impute missing values of county-level time series data using R? I have a dataset consisting of mobility data at the county-level for the US for about one year. So the number of observations is >1m. Apart from the county code, the date, and the mobility index, the dataset also has a number of socio-economic variables at the county level, which do not vary temporally (i.e. there is one single value for all day-county combinations) and also do not have any missing values. The mobility data contain a lot of missing values clustered both temporally and spatially, meaning that some counties have more missing values than others and some days have more missing values than others.
I'm looking for an imputation method that takes into account temporal, spatial correlation as well as the association with some sociodemographic county-level variables such as income, education, and so forth. I'm thinking that multiple imputation would be the way to go (like Hmisc or mice) but how would I incorporate the spatial and temporal correlation? Is there an R package that does what I'm looking for?
 A: This sounds like a tough problem.
One of the options that would be appropriate if the missingness can be explained by the observed information and does not depend on the unobserved values (an unverifiable assumption) is to use a Bayesian hierarchical model (usually fit using MCMC methods) with weakly informative priors. Multiple imputation is usually nothing more than fitting a Bayesian model and then simulating the missing data from MCMC draws of the posterior, after all. In such a model you could have a county level effect that in part randomly varies across counties, but also get explained by county level covariates. One can then draw predictions from such a model (e.g. one per MCMC sample). The most difficult bit would be to appropriately model the time and spatial correlations - in some sense (if you want to allow temporal patterns to vary by location - e.g. in Minnesota icy roads might be a problem early on in winter, later in winter in New Jersey and in Florida not really), we are talking about a 3D correlation. Perhaps one could have a model for time patterns for each country, but partially pool the parameters of these models (e.g. if you just want to model seasonality these parameters could be the county-level coefficients for sine(day of year/period) and cosine(day of year/period) terms) across counties with the strongest pooling amongst neighboring counties? There must be a way of doing that (e.g. I'm guessing someone in the Stan community knows how to do this in Stan - Andrew Gelman recently commented on hierarchical imputation models), but I'm not aware of it. I'm also a little apprehensive about how long it would take to fit such a model, but with a large enough compute cluster within-chain parallelization (see also) might be a good option
The Amelia R package fits a quite simple Bayesian model (with a latent multivariate normal distribution in the background that gets used to construct categorical and ordinal variables, too) that often does a nice job of creating multiple imputations. One could probably give it categorical predictors (e.g. counties), variables that construct time patterns (e.g. sine(day/period), cosine(day/period) for different choice of period) and spatial patterns, and then try to use its time series option. However, I have to admit that I've usually struggled to get it to work well with time series (I've got really sensible imputations for non-time series tabular data from it though, so definitely worth knowing, if you have not looked at it before).
Another thought is to use a Bayesian neural network as a very flexible modeling approach, I'm just note sure what exactly would fit your problem best. E.g. you could have an embedding layer for each county and each record within the county (this would deal decently well with counties/records for which you have at least some data, but less well with ones that you don't have any data for). You would either just throw all records from all time periods together (and just have time and geographic coordinates as features) in a simple tabular neural network (see e.g. what is implemented in the fastai library, with the PyTorch code to define the model located here). You might hope that such a model - as long as you cross-validate it properly - could be made to capture time and location dependencies appropriately. Alternatively, a time-series neural network like a LSTM might work okay for this problem (would be a good choice for capturing the time-series nature of the problem), or perhaps a graph-neural-network that could explicitly capture which counties neighbor each other. In any case, you'd presumably fit these with drop-out on the various layers (esp. the embedding layers and the fully connected layers). If you did that, then by leaving drop-out on during inference, you get variation in the predictions for the missing values and that would approximate the posterior of a proper Bayesian NN (however, you'd want to ensure that you predict all the missing records with the same random number seed to ensure correlations are as they should be within each imputed dataset - annoyingly this is not well supported, yet, by the big deep learning frameworks). See this paper (or the authors PhD thesis) and this for more details.
