FR Discrete compartmental models for epidemiology

[fritzo]

This issue tracks tasks related to discrete compartmental models.
See design doc for motivation.

Modeling

• Example: discrete SIR models with HMC inference #2410 Initial prototype using HMC + enumeration
• Start a contrib.epidemiology module #2437 Implement a framework in pyro.contrib.epidemiology
• Implement a simple SEIR model #2446 Implement an SEIR model @fritzo
• Add a factory for Poisson / NegativeBinomial / Binomial / BetaBinomial #2450 Implement overdispersed SIR and SEIR models #2451 Generalize to overdispersed noise via BetaBinomial @fritzo
• Add example of sparsely observed SIR model #2457 Handle observations accumulated over an interval
• Add SIR model with unknown time of initial infection #2460 Allow stochastic init state and start date estimation
• Implement a regional SIR model with approximate inference #2466 Add a regional model
• Implement a coalescent time distribution #2474 Add a CoalescentTimes distribution
• Eliminate need for CompartmentalModel.transition_bwd() #2514 Eliminate need for ComparmentalModel.transition_bwd()
• Add a time-heterogeneous SIR model #2517 Add a model with time-varying global parameters
• Add advanced modeling section to epidemiology intro tutorial #2558 Add coalescent likelihoods to contrib.epidemiology #2468 Incorporate phylogeny observations into the likelihood (see OverdispersedSEIRModel)
• Add advanced modeling section to epidemiology intro tutorial #2558, Add an SEIRD example with custom .compute_flows() #2559 Add example of branching flows add branching flow example to epidemiology models  #2552
• Add example of causal intervention (ideally interactive on nbviewer)
• Add example of cyclic flows add branching flow example to epidemiology models  #2552
• Support enumeration of discrete globals, e.g. to condition on multiple phylogenetic samples

Inference

• Use full mass matrix for globals in SIR example #2428 Use block diagonal mass matrix after Implement structured mass matrix in HMC #2425 is merged
• Fix recovery_time to known value in SIR models #2429 Fix global parameters that are easy to physically measure
• Add smoothing parameter to Discrete Cosine transform and reparam #2430 Try a frequency-smoothed DCT transform @fritzo
• add quintic spline quantization strategy to sir_hmc.py #2434 Try larger enumeration windows @martinjankowiak
• Use SMC to initialize compartmental models #2452 Use SMC to generically initialize to a feasible trajectory
• Implement arrow-head precision adaptation scheme #2465 Try structured adaptation matrices in HMC @fehiepsi
• add HaarTransform #2492 Add a HaarReparam and use it in contrib.epidemiology #2493 Try a wavelet transform @martinjankowiak
• Implement a SplitReparam and use it in contrib.epidemiology #2495 Implement a SplitReparam to split off low-frequency components of Haar reparams
• Add naive relaxed inference for CompartmentalModel #2510 Implements relaxed inference for compartmental models #2513 Implement continuous-relaxed inference for CompartmentalModel #2522 Try continuous relaxed inference
• Add overdispersed models to contrib.epidemiology #2498 Make all Binomial distributions overdispersed
• Support HaarReparam for non-compartmental variables #2523 Support Haar reparameterization for non-compartmental time series
• Implement SVI inference for CompartmentalModel #2529 Try SVI on the continuous relaxed model
• Try IWAE with DReG gradients, compare with ELBO (see FR: IWAE_DReG importance weighted VAE with doubly reparametrized gradients #2598)
• Try one-step-lookahead SMC
• Try a windowed modified DCT transform

Engineering

• Fix SIR predict() use of reparameterizers #2431 Fix prediction bug when using --dct
• Fix clamping bug in SIR example's quantize helpers #2433 Fix cubic spline weights @martinjankowiak
• Add a HaarReparam and use it in contrib.epidemiology #2493 Use .with_cache(1) in DiscreteCosineTransform (see Add a .with_cache() method to distributions.Transform objects pytorch/pytorch#36882 )
• Make SafeLog jit-script-compatible (see comment in hmm.py)

Tidying up for first release DONE

• Reorganize contrib.epidemiology models #2499 combine sir.py + seir.py -> models.py
• Add intro tutorial to contrib.epidemiology #2553 Add a notebook tutorial introducing the module purpose and usage
• Add advanced modeling section to epidemiology intro tutorial #2558 Add epidemiology tutorial with a regional SEIR model #2518 Add a notebook tutorial on regional models
• Tweak defaults for CompartmentalModel.fit_mcmc() #2563 Add pyro.contrib.epidemiology experiment sandbox#13 Set good default values for inference

[martinjankowiak]

great task list. i'm going to try to put together something so we can try larger enumeration windows

[fritzo]

Re: splines, I think we will need different spline weights for differently sized windows. Cubic splines require at least four points. Each point added beyond four gives one new degree of freedom on the coefficients to be solved (the other three degrees of freedom are used by the boundary conditions). One natural way to constrain these extra degrees of freedom in spline coefficients is to minimize the sum of squares of cubic coefficients, i.e. a simple notion of maximally smooth.