Inference with SEIR model and NPE

[francescopinotti92]

Description

I'm testing the ability of neural SBI methods to fit epidemic data in a simple setting where an epidemic is growing exponentially starting from a single infected case. I am using a SEIR model with transmission rate $\beta$, recovery rate $\mu$ and rate of infectiousness onset $\sigma$. "True" incidence in week $t$ is calculated from simulations as: $J_t = \sigma \int_{7t}^{7t+7}E(u)du$. Observed cases in the same week are assumed to be Poisson-distributed with mean $\rho \cdot J_t$, where $\rho$ is an under-reporting factor. The last parameter, $T_0$, denotes the time elapsed between the introduction of the first case into the population and the beginning of the observation window; this means I do not observe any data between times $-T_0$ and $0$.

In my experiments, case data cover 23 weeks. Also, I assume that $\mu$, $\sigma$, $\rho$ and population size $N$ are known, which leaves only two parameters to be inferred, $\beta$ and $T_0$. The transmission rate is actually calculated as $\beta = R_0 \mu$, where $R_0$ is my parameter of interest (instead of $\beta$). I chose uniform priors for both $R_0$ and $T_0$. More precisely, $R_0 \sim \text{Uniform}(1,5)$ and $T_0 \sim \text{Uniform}(0,50)$.

In an even simpler scenario, I further assume that $T_0$ is known so that $R_0$ is the only unknown parameter.

I've been playing around with sbi for a few weeks now but unfortunately I can't seem to get good posteriors in this example. More precisely, posteriors obtained with, say, NPE, are typically wider than the exact posteriors (which I obtained via standard MCMC methods and have the correct coverage). The results are actually satisfactory for relatively large $R_0$ values, e.g. 3, but they degrade considerably at lower values.

I believe that this is quite a simple problem that could be tackled by NPE, hence I wonder if I'm missing any fundamental concept. Any advice would be really welcome. Thanks for making this impressive toolkit available!

Code

import numpy as np
import torch
import sbi
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from sbi.inference import NPE
from sbi.utils import BoxUniform
from sbi.utils.user_input_checks import (
    check_sbi_inputs,
    process_prior,
    process_simulator,
)
from sbi.neural_nets import posterior_nn

#=== actual system of ODEs
def SEIR_ODE( y, t, N, beta, sigma, mu ):

    # C is tracking cumulative incidence 
    S, E, I, R, C = y

    dS_dt  = -beta * S * I / N
    dE_dt  =  beta * S * I / N - sigma * E
    dI_dt  =  sigma * E - mu * I
    dR_dt  =  mu * I
    dC_dt =  sigma * E

    return( [ dS_dt, dE_dt, dI_dt, dR_dt, dC_dt ] )

#=== simulator with 2 free parameters, R_0 and T_0
def SBI_simulator( theta ):

  #== fixed parameters
  N = 10628972.
  sigma = 1./11.4
  mu = 1./7.0
  rho = 0.5

  R0, T0 = theta
  beta = R0 * mu

  ts = [ 0., T0 ] + [ T0 + 7. * i for i in range( 1, 24 ) ]
  sol = odeint( SEIR_ODE, [ N - 1.0, 1.0, 0.0, 0.0, 0.0 ], ts, args = ( N, beta, sigma, mu ) )

  C   = sol[:,-1]
  inc = np.diff( C )
  inc_obs  = np.random.poisson( rho * inc )

  return inc_obs[1:]

#=== simulator with R_0 as sole free parameter
def SBI_simulator_R0( theta ):

    R0 = theta[0]
    global T0 # T0 is set outside the simulator

    return SBI_simulator( [ R0, T0 ] )


#=== simulation and fitting

T0 = 30. # set T_0

prior = BoxUniform( low  = torch.tensor( [ 1.] ),
                    high = torch.tensor( [ 5.] ) )

prior, num_parameters, prior_returns_numpy = process_prior( prior )
simulator = process_simulator( SBI_simulator_R0, prior, prior_returns_numpy )

theta, x = simulate_for_sbi( simulator, prior, num_simulations = 200000, num_workers = 4 )

neural_posterior = posterior_nn( model="maf" ) # options for density estimation
inferer_NPE = NPE( prior = prior, density_estimator = neural_posterior ) # inferer object
density_estimator_NPE = inferer_NPE.append_simulations( theta, x ).train( training_batch_size = 512, stop_after_epochs = 50  ) # density estimator
posterior_NPE = inferer_NPE.build_posterior( density_estimator_NPE ) 

# simulate some observed data
data_true = SBI_simulator_R0( torch.tensor([1.6]) ) 
samples_NPE = posterior_NPE.set_default_x( data_true ).sample( ( 5000, ) )

When setting $R_0=1.6$ I get plots like the following:

The true posterior, not shown here, is much tighter around the ground-truth value. I also tried using NSF instead of MAF but results do not improve, unfortunately.

[janfb]

Hi @francescopinotti92

thanks for sharing your application and filing the issue.

What is the dimensionality of the data here? It's the diff of C plus Poisson noise, so a noisy time-series?
The problem could be that the NPE density estimator has a hard time learning the posterior if the input dimension is too large, say >50 or so. Thus, it probably makes sense to either use a set of summary statistics here, e.g., the first five or so moments and auto-correlations etc, or use an embedding network to embed the time-series.

If you are not familiar with embedding networks, please have a look at our tutorial: https://sbi-dev.github.io/sbi/dev/tutorials/04_embedding_networks/
As a first step, we also have pre-configured embedding nets that could be a good first start, e.g., just a vanilla MLP here: https://github.com/sbi-dev/sbi/blob/main/sbi/neural_nets/embedding_nets/fully_connected.py
As a next step it then probably makes more sense to use a net that can learn temporal dependencies, e.g., an RNN, or a transformer (not implemented in sbi).

Does this help?

Cheers,
Jan

[francescopinotti92]

Thanks for your rapid answer!

Yes, my data is a noisy time series obtained by applying Poisson noise to the output of the ODE model (which is deterministic). C is just an intermediate variable that stores cumulative cases over time; hence np.diff(C) allows me to calculate weekly cases. The time series has length 23 in my examples.

I will try your suggestions and add the results here.

[francescopinotti92]

Hello,

I obtained major improvements on a simpler instance (SIR model instead of SEIR) using mixed density networks. These are the steps that I believed helped me:

1. Reducing the upper prior bound on R0 from 5 to 3. My guess is that larger values of R0 yield a faster dynamics and hence simulations displaying the classic "bell" shape; however I am only interested in data resembling exponential growth. Perhaps limiting the variability of simulated data helped here.
2. I used a feed-forward neural network with 5 hidden layers.
3. I set stop_after_epochs = 200 to increase chances of finding a good solution.

My code:

def SIR_ODE( y, t, N, beta, mu ):

    S, I, R, C = y

    dS_dt  = -beta * S * I / N
    dI_dt  =  beta * S * I / N - mu * I
    dR_dt  =  mu * I
    dC_dt  =  beta * S * I / N

    return( [ dS_dt, dI_dt, dR_dt, dC_dt ] )

def SBI_simulator( theta ):

  #== fixed parameters
  N = 10628972.
  mu = 1./20.0
  rho = 0.5

  R0, T0 = theta
  beta = R0 * mu

  ts = [ 0., T0 ] + [ T0 + 7. * i for i in range( 1, 15 ) ]
  sol = odeint( SIR_ODE, [ N - 1.0, 1.0, 0.0, 0.0 ], ts, args = ( N, beta, mu ) )

  C   = sol[:,-1]
  inc = np.diff( C )
  inc_obs  = np.random.poisson( rho * inc )

  return inc_obs[1:]

def SBI_simulator_R0( theta ):

    R0 = theta[0]
    global T0 # T0 is set outside the simulator

    return SBI_simulator( [ R0, T0 ] )

T0 = 30.
prior = BoxUniform( low  = torch.tensor( [ 0.5] ),
                    high = torch.tensor( [ 3.] ) )

prior, num_parameters, prior_returns_numpy = process_prior( prior )
simulator = process_simulator( SBI_simulator_R0, prior, prior_returns_numpy )

embedding_net = FCEmbedding(
input_dim=14,
num_hiddens=20,
num_layers=5,
output_dim=4,
)

neural_posterior = posterior_nn( model="mdn", num_components=10, hidden_features=40, z_score_x = 'independent', embedding_net = embedding_net ) # options for density estimation
inferer_NPE = NPE( prior = prior, density_estimator = neural_posterior ) # inferer object
density_estimator_NPE = inferer_NPE.append_simulations( theta, x ).train( training_batch_size = 128, stop_after_epochs = 200, validation_fraction=0.1  ) # density estimator
posterior_NPE = inferer_NPE.build_posterior( density_estimator_NPE ) 

The approximate posteriors are much closer to the true ones over a large range of R0 values (I'm still struggling close to R_0=1). I'm still unsure whether this will work with SEIR dynamics and with 2 parameters to infer at once but it's definitely an improvement. I guess I should spend more time tinkering with the prior and the embedding network. Are there any further checks I can make to diagnose problems with inference?

Thanks for the guidance!

[janfb]

Thanks for the update, that's great!

In general, I would suggest the following checks:

1. prior predictive check: if you are dealing with "real" observed data x_o, check whether x_o lies within the distribution of simulated data (with parameters sampled from prior)
2. after the training: have a look at the validation loss curve to see whether is has converged properly, e.g., a nice plateaux, no sudden bumps, etc
3. posterior calibration checks: run sbc or other calibration checks to see whether your posterior is well-calibrated
4. posterior predictive check: run simulation with samples from the posterior and check whether they reproduce the x_o the posterior was conditioned on. There should be some noise around it, e.g., due to simulation noise, but it should be close and definitively much more constrained than the prior predictive.

We have tutorials for all these steps on the website / in GitHub under tutorials/

I hope this helps! Feel free to report your results here.

I am moving this issue to GH discussion as it is more about a specific use case than about an issue with the toolbox.

Cheers,
Jan