Inferring daily Covid-19 infections from both emergency service calls and deaths

Hi all,

I’ve built a model for inferring the daily number of Covid-19 infections in a geographical region from both coronavirus symptoms being reported through emergency service calls and coronavirus-related deaths. I hoped that fusing high-noise low-latency call data with low-noise high-latency death data would result in better estimates of the current number of daily infections. Unfortunately, I’m having difficulties fitting the model:

functions {
  real[] seeiittd(real time,
                  real[] state,
                  real[] params,
                  real[] real_data,
                  int[] integer_data) {

    // Unpack integer data values
    int T = integer_data[1];
    int n_beta_pieces = integer_data[2];
    int n_disease_states = integer_data[3];

    // Unpack real data values
    real beta_left_t[n_beta_pieces] = real_data[1:n_beta_pieces];
    real beta_right_t[n_beta_pieces] = real_data[n_beta_pieces+1:2*n_beta_pieces];
    real population = real_data[2*n_beta_pieces+1];

    // Unpack parameter values
    real beta_left[n_beta_pieces] = params[1:n_beta_pieces];
    real grad_beta[n_beta_pieces] = params[n_beta_pieces+1:2*n_beta_pieces];
    real nu = params[2*n_beta_pieces+1];
    real gamma = params[2*n_beta_pieces+2];
    real kappa = params[2*n_beta_pieces+3];
    real omega = params[2*n_beta_pieces+4];

    // Unpack state
    real S = state[1];
    real E1 = state[2];
    real E2 = state[3];
    real I1 = state[4];
    real I2 = state[5];
    real T1 = state[6];
    real T2 = state[7];
    real D = state[8];

    real infection_rate;

    for (i in 1:n_beta_pieces) {
      if(time >= beta_left_t[i] && time < beta_right_t[i]) {
        real beta = grad_beta[i] * (time - beta_left_t[i]) + beta_left[i];
        infection_rate = beta * (I1 + I2) * S / population;
      }
    }

    real nuE1 = nu * E1;
    real nuE2 = nu * E2;
    real gammaI1 = gamma * I1;
    real gammaI2 = gamma * I2;
    real kappaT1 = kappa * T1;
    real kappaT2 = kappa * T2;

    real dS_dt = -infection_rate;
    real dE1_dt = infection_rate - nuE1;
    real dE2_dt = nuE1 - nuE2;
    real dI1_dt = nuE2 - gammaI1;
    real dI2_dt = gammaI1 - gammaI2;
    real dT1_dt = gammaI2 * omega - kappaT1;
    real dT2_dt = kappaT1 - kappaT2;
    real dD_dt = kappaT2;

    return {dS_dt, dE1_dt, dE2_dt, dI1_dt, dI2_dt, dT1_dt, dT2_dt, dD_dt};
  }

  real[ , ] integrate_ode_explicit_trapezoidal(real[] initial_state, real initial_time, real[] times, real[] params, real[] real_data, int[] integer_data) {
    real h;
    vector[size(initial_state)] dstate_dt_initial_time;
    vector[size(initial_state)] dstate_dt_tidx;
    vector[size(initial_state)] k;
    real state_estimate[size(times),size(initial_state)];

    h = times[1] - initial_time;
    dstate_dt_initial_time = to_vector(seeiittd(initial_time, initial_state, params, real_data, integer_data));
    k = h*dstate_dt_initial_time;
    state_estimate[1,] = to_array_1d(to_vector(initial_state) + h*(dstate_dt_initial_time + to_vector(seeiittd(times[1], to_array_1d(to_vector(initial_state)+k), params, real_data, integer_data)))/2);

    for (tidx in 1:size(times)-1) {
      h = (times[tidx+1] - times[tidx]);
      dstate_dt_tidx = to_vector(seeiittd(times[tidx], state_estimate[tidx], params, real_data, integer_data));
      k = h*dstate_dt_tidx;
      state_estimate[tidx+1,] = to_array_1d(to_vector(state_estimate[tidx,]) + h*(dstate_dt_tidx + to_vector(seeiittd(times[tidx+1], to_array_1d(to_vector(state_estimate[tidx,])+k), params, real_data, integer_data)))/2);
    }

    return state_estimate;
  }
}
data {
  real initial_time;
  int<lower=1> n_beta_pieces;
  real<lower=0> beta_left_t[n_beta_pieces];
  real<lower=0> beta_right_t[n_beta_pieces];
  int<lower=1> n_rho_pieces;
  int<lower=0> rho_left_t[n_rho_pieces];
  int<lower=0> rho_right_t[n_rho_pieces];
  int<lower=1> T;
  real times[T];
  int<lower=1> n_disease_states;
  real<lower=0> population;
  int<lower=1> deaths_length;
  int<lower=1> deaths_starts[deaths_length];
  int<lower=1> deaths_stops[deaths_length];
  int<lower=0> deaths[deaths_length];
  int<lower=1> calls_111_length;
  int<lower=1> calls_111_start;
  int<lower=0> calls_111[calls_111_length];
  int real_data_length;
  real real_data[real_data_length];
  int integer_data_length;
  int integer_data[integer_data_length];
}
transformed data {
  real mu_dL = 4.00;
  real sigma_dL = 0.20;
  real mu_dI = 3.06;
  real sigma_dI = 0.21;
  real mu_dT = 16.00;
  real sigma_dT = 0.71;
  int max_lag = 13;
}
parameters {
  simplex[3] initial_state_raw;
  real<lower=0> beta_left[n_beta_pieces];
  real<lower=0> beta_right[n_beta_pieces];
  real<lower=0> dL;
  real<lower=0> dI;
  real<lower=0> dT;
  real<lower=0, upper=1> omega;
  real<lower=0> reciprocal_phi_deaths;
  real<lower=0> reciprocal_phi_calls_111;
  real<lower=0, upper=1> rho_phi;
  real<lower=0.1> rho_lambda;
  real<lower=0, upper=1> rho[n_rho_pieces];
  simplex[max_lag+1] lag_weights;
}
transformed parameters {
  real initial_state[n_disease_states];
  real grad_beta[n_beta_pieces];
  real nu;
  real gamma;
  real kappa;
  real phi_deaths;
  real phi_calls_111;
  real rho_alpha;
  real rho_beta;
  real state_estimate[T,n_disease_states];
  vector[T+1] S;
  vector[T+1] E1;
  vector[T+1] E2;
  vector[T+1] I1;
  vector[T+1] I2;
  vector[T+1] T1;
  vector[T+1] T2;
  vector[T+1] D;
  vector[T] daily_infections;
  vector[T] daily_deaths;
  vector[T] calls_111_lagged_daily_infections;
  vector[T] daily_calls_111;

  initial_state[1] = (population-5.0)*initial_state_raw[1] + 1.0;
  initial_state[2] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[3] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[4] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[5] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[6] = 0.0;
  initial_state[7] = 0.0;
  initial_state[8] = 0.0;
  grad_beta = to_array_1d((to_vector(beta_right) - to_vector(beta_left))./(to_vector(beta_right_t) - to_vector(beta_left_t)));
  nu = 2.0/dL;
  gamma = 2.0/dI;
  kappa = 2.0/dT;
  phi_deaths = 1.0 / reciprocal_phi_deaths;
  phi_calls_111 = 1.0 / reciprocal_phi_calls_111;
  rho_alpha = rho_lambda * rho_phi;
  rho_beta = rho_lambda * (1 - rho_phi);

  {
    real params[2*n_beta_pieces+4];
    params[1:n_beta_pieces] = beta_left;
    params[n_beta_pieces+1:2*n_beta_pieces] = grad_beta;
    params[2*n_beta_pieces+1] = nu;
    params[2*n_beta_pieces+2] = gamma;
    params[2*n_beta_pieces+3] = kappa;
    params[2*n_beta_pieces+4] = omega;

    state_estimate = integrate_ode_explicit_trapezoidal(initial_state, initial_time, times, params, real_data, integer_data);
  }

  S = append_row(initial_state[1], to_vector(state_estimate[, 1]));
  E1 = append_row(initial_state[2], to_vector(state_estimate[, 2]));
  E2 = append_row(initial_state[3], to_vector(state_estimate[, 3]));
  I1 = append_row(initial_state[4], to_vector(state_estimate[, 4]));
  I2 = append_row(initial_state[5], to_vector(state_estimate[, 5]));
  T1 = append_row(initial_state[6], to_vector(state_estimate[, 6]));
  T2 = append_row(initial_state[7], to_vector(state_estimate[, 7]));
  D = append_row(initial_state[8], to_vector(state_estimate[, 8]));

  daily_infections = S[:T] - S[2:] + machine_precision();
  daily_deaths = D[2:] - D[:T];

  calls_111_lagged_daily_infections = lag_weights[1]*daily_infections;

  for (i in 1:max_lag) {
    calls_111_lagged_daily_infections += lag_weights[i+1]*append_row(rep_vector(0.0, i), daily_infections[:T-i]);
  }

  daily_calls_111 = rep_vector(0.0, T);

  for (i in 1:n_rho_pieces) {
    daily_calls_111[rho_left_t[i]:rho_right_t[i]-1] = calls_111_lagged_daily_infections[rho_left_t[i]:rho_right_t[i]-1] * rho[i];
  }
}
model {
  initial_state_raw ~ dirichlet(rep_vector(0.1, 3));
  beta_left ~ std_normal();
  beta_right ~ std_normal();
  dL ~ normal(mu_dL, sigma_dL);
  dI ~ normal(mu_dI, sigma_dI);
  dT ~ normal(mu_dT, sigma_dT);
  omega ~ beta(100, 9803);
  reciprocal_phi_deaths ~ exponential(5);
  reciprocal_phi_calls_111 ~ exponential(5);
  rho_phi ~ beta(1, 1);
  rho_lambda ~ pareto(0.1, 1.5);
  rho ~ beta(rho_alpha, rho_beta);
  lag_weights ~ dirichlet(rep_vector(0.1, max_lag+1));

  for (i in 1:deaths_length) {
    target += 7.0*neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  for (i in 1:n_rho_pieces) {
    target += neg_binomial_2_lpmf(calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] | daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}
generated quantities {
  vector[T-1] growth_rate = (log(daily_infections[2:]) - log(daily_infections[:T-1]))*100;

  int pred_deaths[deaths_length];

  for (i in 1:deaths_length) {
    pred_deaths[i] = neg_binomial_2_rng(sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  int pred_calls_111[calls_111_length];

  for (i in 1:n_rho_pieces) {
    pred_calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] = neg_binomial_2_rng(daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}

I’ve been attempting to use the following script to fit the model to data from Liverpool (which is where I happen to live in the UK):

#!/bin/bash

for i in {1..6}
  do
    ./deaths_and_111_calls sample random seed=42 id=${i} data file=liverpool.data.json init=init${i}.json output file=samples${i}.csv &
  done

I’ve been using the init*.json files to initialise the chains in a region of parameter space that doesn’t result in errors being thrown. The initialisation files along with the model input data should all be attached to this post (with .txt extensions to allow uploading).

I’d be very grateful for suggestions about what to try to get the model working.

Thanks in advance,
Rob

init1.json.txt (297 Bytes) init2.json.txt (296 Bytes) init3.json.txt (298 Bytes) init4.json.txt (297 Bytes) init5.json.txt (297 Bytes) init6.json.txt (296 Bytes) liverpool.data.json.txt (3.0 KB)

5 Likes

I’m hoping that people like @charlesm93, @breckbaldwin and @jonah may be able to help you, @remoore: it’s really important that we get this model working well since we plan for it to be used (and not only inform some papers!).
Simon

1 Like

Hi Simon -

This might not be useful to you, but I have a COVID model that is more empirical in nature and might be easier to fit with your data. See paper:

https://osf.io/preprints/socarxiv/jp4wk

2 Likes

@saudiwin: thank you. We will certainly take a look. We do want to understand what is causing our current model to have issues though. I don’t like it when things don’t work for reasons we can’t explain!

Rob,
Is the issue that you are unhappy having to specify init values? I ran your code and data without inits as follows–note single chain, num_samples=10 and num_warmup=10. :

for i in {1..1} do ./deaths_and_111_calls sample num_samples=10 num_warmup=10 random seed=42 id=${i} data file=liverpool.data.json output file=samples${i}.csv
That fails badly (line number will be for my print-encrusted code below):

Rejecting initial value:
  Error evaluating the log probability at the initial value.
Exception: neg_binomial_2_lpmf: Location parameter is nan, but must be > 0! (in '../git/covid-19-uk/tmp/deaths_and_111_calls.stan', line 250, column 4 to column 115)

This stops the iteration in the first few data points every time I have run it. The actual code is:

target += 7.0*neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);

I added some print statements to your .stan program to see what the values are:

functions {
  real[] seeiittd(real time,
                  real[] state,
                  real[] params,
                  real[] real_data,
                  int[] integer_data) {

    // Unpack integer data values
    int T = integer_data[1];
    int n_beta_pieces = integer_data[2];
    int n_disease_states = integer_data[3];

    // Unpack real data values
    real beta_left_t[n_beta_pieces] = real_data[1:n_beta_pieces];
    real beta_right_t[n_beta_pieces] = real_data[n_beta_pieces+1:2*n_beta_pieces];
    real population = real_data[2*n_beta_pieces+1];

    // Unpack parameter values
    real beta_left[n_beta_pieces] = params[1:n_beta_pieces];
    real grad_beta[n_beta_pieces] = params[n_beta_pieces+1:2*n_beta_pieces];
    real nu = params[2*n_beta_pieces+1];
    real gamma = params[2*n_beta_pieces+2];
    real kappa = params[2*n_beta_pieces+3];
    real omega = params[2*n_beta_pieces+4];

    // Unpack state
    real S = state[1];
    real E1 = state[2];
    real E2 = state[3];
    real I1 = state[4];
    real I2 = state[5];
    real T1 = state[6];
    real T2 = state[7];
    real D = state[8];

    real infection_rate;

    for (i in 1:n_beta_pieces) {
      if(time >= beta_left_t[i] && time < beta_right_t[i]) {
        real beta = grad_beta[i] * (time - beta_left_t[i]) + beta_left[i];
        infection_rate = beta * (I1 + I2) * S / population;
      }
    }

    real nuE1 = nu * E1;
    real nuE2 = nu * E2;
    real gammaI1 = gamma * I1;
    real gammaI2 = gamma * I2;
    real kappaT1 = kappa * T1;
    real kappaT2 = kappa * T2;

    real dS_dt = -infection_rate;
    real dE1_dt = infection_rate - nuE1;
    real dE2_dt = nuE1 - nuE2;
    real dI1_dt = nuE2 - gammaI1;
    real dI2_dt = gammaI1 - gammaI2;
    real dT1_dt = gammaI2 * omega - kappaT1;
    real dT2_dt = kappaT1 - kappaT2;
    real dD_dt = kappaT2;

    return {dS_dt, dE1_dt, dE2_dt, dI1_dt, dI2_dt, dT1_dt, dT2_dt, dD_dt};
  }

  real[ , ] integrate_ode_explicit_trapezoidal(real[] initial_state, real initial_time, real[] times, real[] params, real[] real_data, int[] integer_data) {
    real h;
    vector[size(initial_state)] dstate_dt_initial_time;
    vector[size(initial_state)] dstate_dt_tidx;
    vector[size(initial_state)] k;
    real state_estimate[size(times),size(initial_state)];

    h = times[1] - initial_time;
    dstate_dt_initial_time = to_vector(seeiittd(initial_time, initial_state, params, real_data, integer_data));
    k = h*dstate_dt_initial_time;
    state_estimate[1,] = to_array_1d(to_vector(initial_state) + h*(dstate_dt_initial_time + to_vector(seeiittd(times[1], to_array_1d(to_vector(initial_state)+k), params, real_data, integer_data)))/2);

    for (tidx in 1:size(times)-1) {
      h = (times[tidx+1] - times[tidx]);
      dstate_dt_tidx = to_vector(seeiittd(times[tidx], state_estimate[tidx], params, real_data, integer_data));
      k = h*dstate_dt_tidx;
      state_estimate[tidx+1,] = to_array_1d(to_vector(state_estimate[tidx,]) + h*(dstate_dt_tidx + to_vector(seeiittd(times[tidx+1], to_array_1d(to_vector(state_estimate[tidx,])+k), params, real_data, integer_data)))/2);
    }

    return state_estimate;
  }
}
data {
  real initial_time;
  int<lower=1> n_beta_pieces;
  real<lower=0> beta_left_t[n_beta_pieces];
  real<lower=0> beta_right_t[n_beta_pieces];
  int<lower=1> n_rho_pieces;
  int<lower=0> rho_left_t[n_rho_pieces];
  int<lower=0> rho_right_t[n_rho_pieces];
  int<lower=1> T;
  real times[T];
  int<lower=1> n_disease_states;
  real<lower=0> population;
  int<lower=1> deaths_length;
  int<lower=1> deaths_starts[deaths_length];
  int<lower=1> deaths_stops[deaths_length];
  int<lower=0> deaths[deaths_length];
  int<lower=1> calls_111_length;
  int<lower=1> calls_111_start;
  int<lower=0> calls_111[calls_111_length];
  int real_data_length;
  real real_data[real_data_length];
  int integer_data_length;
  int integer_data[integer_data_length];
}
transformed data {
  real mu_dL = 4.00;
  real sigma_dL = 0.20;
  real mu_dI = 3.06;
  real sigma_dI = 0.21;
  real mu_dT = 16.00;
  real sigma_dT = 0.71;
  int max_lag = 13;
  print("deaths=",deaths);
  print("deaths_length=",deaths_length);
  
}
parameters {
  simplex[3] initial_state_raw;
  real<lower=0> beta_left[n_beta_pieces];
  real<lower=0> beta_right[n_beta_pieces];
  real<lower=0> dL;
  real<lower=0> dI;
  real<lower=0> dT;
  real<lower=0, upper=1> omega;
  real<lower=0> reciprocal_phi_deaths;
  real<lower=0> reciprocal_phi_calls_111;
  real<lower=0, upper=1> rho_phi;
  real<lower=0.1> rho_lambda;
  real<lower=0, upper=1> rho[n_rho_pieces];
  simplex[max_lag+1] lag_weights;
}
transformed parameters {
  real initial_state[n_disease_states];
  real grad_beta[n_beta_pieces];
  real nu;
  real gamma;
  real kappa;
  real phi_deaths;
  real phi_calls_111;
  real rho_alpha;
  real rho_beta;
  real state_estimate[T,n_disease_states];
  vector[T+1] S;
  vector[T+1] E1;
  vector[T+1] E2;
  vector[T+1] I1;
  vector[T+1] I2;
  vector[T+1] T1;
  vector[T+1] T2;
  vector[T+1] D;
  vector[T] daily_infections;
  vector[T] daily_deaths;
  vector[T] calls_111_lagged_daily_infections;
  vector[T] daily_calls_111;

  initial_state[1] = (population-5.0)*initial_state_raw[1] + 1.0;
  initial_state[2] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[3] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[4] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[5] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[6] = 0.0;
  initial_state[7] = 0.0;
  initial_state[8] = 0.0;
  grad_beta = to_array_1d((to_vector(beta_right) - to_vector(beta_left))./(to_vector(beta_right_t) - to_vector(beta_left_t)));
  nu = 2.0/dL;
  gamma = 2.0/dI;
  kappa = 2.0/dT;
  phi_deaths = 1.0 / reciprocal_phi_deaths;
  phi_calls_111 = 1.0 / reciprocal_phi_calls_111;
  rho_alpha = rho_lambda * rho_phi;
  rho_beta = rho_lambda * (1 - rho_phi);

  {
    real params[2*n_beta_pieces+4];
    params[1:n_beta_pieces] = beta_left;
    params[n_beta_pieces+1:2*n_beta_pieces] = grad_beta;
    params[2*n_beta_pieces+1] = nu;
    params[2*n_beta_pieces+2] = gamma;
    params[2*n_beta_pieces+3] = kappa;
    params[2*n_beta_pieces+4] = omega;

    state_estimate = integrate_ode_explicit_trapezoidal(initial_state, initial_time, times, params, real_data, integer_data);
  }

  S = append_row(initial_state[1], to_vector(state_estimate[, 1]));
  E1 = append_row(initial_state[2], to_vector(state_estimate[, 2]));
  E2 = append_row(initial_state[3], to_vector(state_estimate[, 3]));
  I1 = append_row(initial_state[4], to_vector(state_estimate[, 4]));
  I2 = append_row(initial_state[5], to_vector(state_estimate[, 5]));
  T1 = append_row(initial_state[6], to_vector(state_estimate[, 6]));
  T2 = append_row(initial_state[7], to_vector(state_estimate[, 7]));
  D = append_row(initial_state[8], to_vector(state_estimate[, 8]));

  daily_infections = S[:T] - S[2:] + machine_precision();
  daily_deaths = D[2:] - D[:T];

  calls_111_lagged_daily_infections = lag_weights[1]*daily_infections;

  for (i in 1:max_lag) {
    calls_111_lagged_daily_infections += lag_weights[i+1]*append_row(rep_vector(0.0, i), daily_infections[:T-i]);
  }

  daily_calls_111 = rep_vector(0.0, T);

  for (i in 1:n_rho_pieces) {
    daily_calls_111[rho_left_t[i]:rho_right_t[i]-1] = calls_111_lagged_daily_infections[rho_left_t[i]:rho_right_t[i]-1] * rho[i];
  }
}
model {
  print("Parameters from now ignored init*.json");
  print("initial_state_raw=", initial_state_raw);
  print("dL=", dL);
  print("dI=", dI);
  print("dT=", dT);
  print("omega=", omega);
  print("reciprocal_phi_deaths=", reciprocal_phi_deaths);
  print("reciprocal_phi_calls_111=", reciprocal_phi_calls_111);

  initial_state_raw ~ dirichlet(rep_vector(0.1, 3));
  beta_left ~ std_normal();
  beta_right ~ std_normal();
  dL ~ normal(mu_dL, sigma_dL);
  dI ~ normal(mu_dI, sigma_dI);
  dT ~ normal(mu_dT, sigma_dT);
  omega ~ beta(100, 9803);
  reciprocal_phi_deaths ~ exponential(5);
  reciprocal_phi_calls_111 ~ exponential(5);
  rho_phi ~ beta(1, 1);
  rho_lambda ~ pareto(0.1, 1.5);
  rho ~ beta(rho_alpha, rho_beta);
  lag_weights ~ dirichlet(rep_vector(0.1, max_lag+1));

  for (i in 1:deaths_length) {
    print("running loop, i=",i);
  }
  for (i in 1:deaths_length) {
    print("");
    print("deaths_length=", deaths_length);
    print("values going into neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths)");
    print("i=", i);
    print("deaths[i]=", deaths[i]); 
    print("sum=death_starts[i]:death_stops[i]=",sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]));
    print("death_starts[i]:death_stops[i]=",daily_deaths[deaths_starts[i]:deaths_stops[i]]);
    print("phi_deaths", phi_deaths);
    target += 7.0*neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  for (i in 1:n_rho_pieces) {
    target += neg_binomial_2_lpmf(calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] | daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}
generated quantities {
  vector[T-1] growth_rate = (log(daily_infections[2:]) - log(daily_infections[:T-1]))*100;

  int pred_deaths[deaths_length];

  for (i in 1:deaths_length) {
    pred_deaths[i] = neg_binomial_2_rng(sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  int pred_calls_111[calls_111_length];

  for (i in 1:n_rho_pieces) {
    pred_calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] = neg_binomial_2_rng(daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}

Example output:

Parameters from now ignored init*.json
initial_state_raw=[0.253532,0.201121,0.545347]
dL=1.04858
dI=0.178305
dT=6.01246
omega=0.858631
reciprocal_phi_deaths=0.903473
reciprocal_phi_calls_111=2.02008
running loop, i=1
running loop, i=2
running loop, i=3
.....
running loop, i=26
running loop, i=27

deaths_length=27
values going into neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths)
i=1
deaths[i]=0
sum=death_starts[i]:death_stops[i]=nan
death_starts[i]:death_stops[i]=[1.95461e+11,5.29288e+19,9.16438e+34,2.12192e+69,-2.20025e+142,1.96623e+299,nan]
phi_deaths1.10684

Rejecting initial value:
  Error evaluating the log probability at the initial value.
Exception: neg_binomial_2_lpmf: Location parameter is nan, but must be > 0! (in '../git/covid-19-uk/tmp/deaths_and_111_calls.stan', line 250, column 4 to column 115)

So sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]) needs to be positive.

I tried constraining daily_deaths to positive with vector<lower=0>[T] daily_deaths; but all that did was move the error to that line,.

Hopefully what I did will likely make it easier to debug what the issue is.

Sorry for not sorting it out.

Breck

1 Like

Hi Rob,

What are the difficulties you are encountering? Does the model not run? Did you attempt a fit, and if so, what warning messages did you get and what diagnostics did you do?

Based on answers to the above, we can start looking at trouble-shooting strategies.

Something to try, if you haven’t done this already, would be to set algorithm = fixed_param, run one iteration, and examine what kind of simulations generated quantities returns. This can help you debug the code or even the model.

You can also run Stan’s diagnose method, which compares the gradient computed by Stan (i.e with autodiff) and the gradient computed with finite differentiation. If there is a discrepancy between the two, this may be an issue.

This means you’re using Stan’s default inits, which may not be appropriate. Stan samples parameters on the unconstrained space from a uniform(-2, 2), which for some models leads to unreasonable parameter values and a numerically unstable evaluation of the log density (hence the nan). This is not uncommon for ODE based models.

Breck,

Thanks for attempting to fit the model.

As you discovered, the initial values get rejected if you use Stan’s default initialisation scheme with this model. I’ve been circumventing this issue by passing a randomly generated init*.json file to each instance of Stan, but I’m not sure this is the best solution.

Charles,

Thanks for commenting on the post.

I tried running the model with algorithm = fixed_param, as you suggested, and this made me realise that the priors on the beta_left and beta_right parameters were too wide. In my updated version of the model, these all have normal(0, 0.5) priors:

functions {
  real[] seeiittd(real time,
                  real[] state,
                  real[] params,
                  real[] real_data,
                  int[] integer_data) {

    // Unpack integer data values
    int T = integer_data[1];
    int n_beta_pieces = integer_data[2];
    int n_disease_states = integer_data[3];

    // Unpack real data values
    real beta_left_t[n_beta_pieces] = real_data[1:n_beta_pieces];
    real beta_right_t[n_beta_pieces] = real_data[n_beta_pieces+1:2*n_beta_pieces];
    real population = real_data[2*n_beta_pieces+1];

    // Unpack parameter values
    real beta_left[n_beta_pieces] = params[1:n_beta_pieces];
    real grad_beta[n_beta_pieces] = params[n_beta_pieces+1:2*n_beta_pieces];
    real nu = params[2*n_beta_pieces+1];
    real gamma = params[2*n_beta_pieces+2];
    real kappa = params[2*n_beta_pieces+3];
    real omega = params[2*n_beta_pieces+4];

    // Unpack state
    real S = state[1];
    real E1 = state[2];
    real E2 = state[3];
    real I1 = state[4];
    real I2 = state[5];
    real T1 = state[6];
    real T2 = state[7];
    real D = state[8];

    real infection_rate;

    for (i in 1:n_beta_pieces) {
      if(time >= beta_left_t[i] && time < beta_right_t[i]) {
        real beta = grad_beta[i] * (time - beta_left_t[i]) + beta_left[i];
        infection_rate = beta * (I1 + I2) * S / population;
      }
    }

    real nuE1 = nu * E1;
    real nuE2 = nu * E2;
    real gammaI1 = gamma * I1;
    real gammaI2 = gamma * I2;
    real kappaT1 = kappa * T1;
    real kappaT2 = kappa * T2;

    real dS_dt = -infection_rate;
    real dE1_dt = infection_rate - nuE1;
    real dE2_dt = nuE1 - nuE2;
    real dI1_dt = nuE2 - gammaI1;
    real dI2_dt = gammaI1 - gammaI2;
    real dT1_dt = gammaI2 * omega - kappaT1;
    real dT2_dt = kappaT1 - kappaT2;
    real dD_dt = kappaT2;

    return {dS_dt, dE1_dt, dE2_dt, dI1_dt, dI2_dt, dT1_dt, dT2_dt, dD_dt};
  }

  real[ , ] integrate_ode_explicit_trapezoidal(real[] initial_state, real initial_time, real[] times, real[] params, real[] real_data, int[] integer_data) {
    real h;
    vector[size(initial_state)] dstate_dt_initial_time;
    vector[size(initial_state)] dstate_dt_tidx;
    vector[size(initial_state)] k;
    real state_estimate[size(times),size(initial_state)];

    h = times[1] - initial_time;
    dstate_dt_initial_time = to_vector(seeiittd(initial_time, initial_state, params, real_data, integer_data));
    k = h*dstate_dt_initial_time;
    state_estimate[1,] = to_array_1d(to_vector(initial_state) + h*(dstate_dt_initial_time + to_vector(seeiittd(times[1], to_array_1d(to_vector(initial_state)+k), params, real_data, integer_data)))/2);

    for (tidx in 1:size(times)-1) {
      h = (times[tidx+1] - times[tidx]);
      dstate_dt_tidx = to_vector(seeiittd(times[tidx], state_estimate[tidx], params, real_data, integer_data));
      k = h*dstate_dt_tidx;
      state_estimate[tidx+1,] = to_array_1d(to_vector(state_estimate[tidx,]) + h*(dstate_dt_tidx + to_vector(seeiittd(times[tidx+1], to_array_1d(to_vector(state_estimate[tidx,])+k), params, real_data, integer_data)))/2);
    }

    return state_estimate;
  }
}
data {
  real initial_time;
  int<lower=1> n_beta_pieces;
  real<lower=0> beta_left_t[n_beta_pieces];
  real<lower=0> beta_right_t[n_beta_pieces];
  int<lower=1> n_rho_pieces;
  int<lower=0> rho_left_t[n_rho_pieces];
  int<lower=0> rho_right_t[n_rho_pieces];
  int<lower=1> T;
  real times[T];
  int<lower=1> n_disease_states;
  real<lower=0> population;
  int<lower=1> deaths_length;
  int<lower=1> deaths_starts[deaths_length];
  int<lower=1> deaths_stops[deaths_length];
  int<lower=0> deaths[deaths_length];
  int<lower=1> calls_111_length;
  int<lower=1> calls_111_start;
  int<lower=0> calls_111[calls_111_length];
  int real_data_length;
  real real_data[real_data_length];
  int integer_data_length;
  int integer_data[integer_data_length];
}
transformed data {
  real mu_dL = 4.00;
  real sigma_dL = 0.20;
  real mu_dI = 3.06;
  real sigma_dI = 0.21;
  real mu_dT = 16.00;
  real sigma_dT = 0.71;
  int max_lag = 13;
}
parameters {
  simplex[3] initial_state_raw;
  real<lower=0> beta_left[n_beta_pieces];
  real<lower=0> beta_right[n_beta_pieces];
  real<lower=0> dL;
  real<lower=0> dI;
  real<lower=0> dT;
  real<lower=0, upper=1> omega;
  real<lower=0> reciprocal_phi_deaths;
  real<lower=0> reciprocal_phi_calls_111;
  real<lower=0, upper=1> rho_phi;
  real<lower=0.1> rho_lambda;
  real<lower=0, upper=1> rho[n_rho_pieces];
  simplex[max_lag+1] lag_weights;
}
transformed parameters {
  real initial_state[n_disease_states];
  real grad_beta[n_beta_pieces];
  real nu;
  real gamma;
  real kappa;
  real phi_deaths;
  real phi_calls_111;
  real rho_alpha;
  real rho_beta;
  real state_estimate[T,n_disease_states];
  vector[T+1] S;
  vector[T+1] E1;
  vector[T+1] E2;
  vector[T+1] I1;
  vector[T+1] I2;
  vector[T+1] T1;
  vector[T+1] T2;
  vector[T+1] D;
  vector[T] daily_infections;
  vector[T] daily_deaths;
  vector[T] calls_111_lagged_daily_infections;
  vector[T] daily_calls_111;

  initial_state[1] = (population-5.0)*initial_state_raw[1] + 1.0;
  initial_state[2] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[3] = (population-5.0)*initial_state_raw[2]/2.0 + 1.0;
  initial_state[4] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[5] = (population-5.0)*initial_state_raw[3]/2.0 + 1.0;
  initial_state[6] = 0.0;
  initial_state[7] = 0.0;
  initial_state[8] = 0.0;
  grad_beta = to_array_1d((to_vector(beta_right) - to_vector(beta_left))./(to_vector(beta_right_t) - to_vector(beta_left_t)));
  nu = 2.0/dL;
  gamma = 2.0/dI;
  kappa = 2.0/dT;
  phi_deaths = 1.0 / reciprocal_phi_deaths;
  phi_calls_111 = 1.0 / reciprocal_phi_calls_111;
  rho_alpha = rho_lambda * rho_phi;
  rho_beta = rho_lambda * (1 - rho_phi);

  {
    real params[2*n_beta_pieces+4];
    params[1:n_beta_pieces] = beta_left;
    params[n_beta_pieces+1:2*n_beta_pieces] = grad_beta;
    params[2*n_beta_pieces+1] = nu;
    params[2*n_beta_pieces+2] = gamma;
    params[2*n_beta_pieces+3] = kappa;
    params[2*n_beta_pieces+4] = omega;

    state_estimate = integrate_ode_explicit_trapezoidal(initial_state, initial_time, times, params, real_data, integer_data);
  }

  S = append_row(initial_state[1], to_vector(state_estimate[, 1]));
  E1 = append_row(initial_state[2], to_vector(state_estimate[, 2]));
  E2 = append_row(initial_state[3], to_vector(state_estimate[, 3]));
  I1 = append_row(initial_state[4], to_vector(state_estimate[, 4]));
  I2 = append_row(initial_state[5], to_vector(state_estimate[, 5]));
  T1 = append_row(initial_state[6], to_vector(state_estimate[, 6]));
  T2 = append_row(initial_state[7], to_vector(state_estimate[, 7]));
  D = append_row(initial_state[8], to_vector(state_estimate[, 8]));

  daily_infections = S[:T] - S[2:] + machine_precision();
  daily_deaths = D[2:] - D[:T];

  calls_111_lagged_daily_infections = lag_weights[1]*daily_infections;

  for (i in 1:max_lag) {
    calls_111_lagged_daily_infections += lag_weights[i+1]*append_row(rep_vector(0.0, i), daily_infections[:T-i]);
  }

  daily_calls_111 = rep_vector(0.0, T);

  for (i in 1:n_rho_pieces) {
    daily_calls_111[rho_left_t[i]:rho_right_t[i]-1] = calls_111_lagged_daily_infections[rho_left_t[i]:rho_right_t[i]-1] * rho[i];
  }
}
model {
  initial_state_raw ~ dirichlet(rep_vector(0.1, 3));
  beta_left ~ normal(0, 0.5);
  beta_right ~ normal(0, 0.5);
  dL ~ normal(mu_dL, sigma_dL);
  dI ~ normal(mu_dI, sigma_dI);
  dT ~ normal(mu_dT, sigma_dT);
  omega ~ beta(100, 9803);
  reciprocal_phi_deaths ~ exponential(5);
  reciprocal_phi_calls_111 ~ exponential(5);
  rho_phi ~ beta(1, 1);
  rho_lambda ~ pareto(0.1, 1.5);
  rho ~ beta(rho_alpha, rho_beta);
  lag_weights ~ dirichlet(rep_vector(0.1, max_lag+1));

  for (i in 1:deaths_length) {
    target += 7.0*neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  for (i in 1:n_rho_pieces) {
    target += neg_binomial_2_lpmf(calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] | daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}
generated quantities {
  vector[T-1] growth_rate = (log(daily_infections[2:]) - log(daily_infections[:T-1]))*100;

  int pred_deaths[deaths_length];

  for (i in 1:deaths_length) {
    pred_deaths[i] = neg_binomial_2_rng(sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  int pred_calls_111[calls_111_length];

  for (i in 1:n_rho_pieces) {
    pred_calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] = neg_binomial_2_rng(daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}

I’ve also added custom initialisation for the beta_left and beta_right parameters, which are initialised between -2 and 0.5 in unconstrained space: init1.json.txt (1.2 KB) init2.json.txt (1.2 KB) init3.json.txt (1.2 KB) init4.json.txt (1.2 KB) init5.json.txt (1.2 KB) init6.json.txt (1.2 KB).

In answer to your questions,

sorry, I should have been more specific about the difficulties I’ve been encountering. The model runs, but I keep getting treedepth and divergence warnings when I run the diagnose executable:

../../cmdstan-2.23.0/bin/diagnose samples1.csv samples2.csv samples3.csv samples4.csv samples5.csv samples6.csv
Processing csv files: samples1.csv, samples2.csv, samples3.csv, samples4.csv, samples5.csv, samples6.csv

Checking sampler transitions treedepth.
8 of 6000 (0.13%) transitions hit the maximum treedepth limit of 10, or 2^10 leapfrog steps.
Trajectories that are prematurely terminated due to this limit will result in slow exploration.
For optimal performance, increase this limit.

Checking sampler transitions for divergences.
547 of 6000 (9.1%) transitions ended with a divergence.
These divergent transitions indicate that HMC is not fully able to explore the posterior distribution.
Try increasing adapt delta closer to 1.
If this doesn't remove all divergences, try to reparameterize the model.

Checking E-BFMI - sampler transitions HMC potential energy.
E-BFMI satisfactory for all transitions.

Effective sample size satisfactory.

Split R-hat values satisfactory all parameters.

Processing complete.

From what I’ve read, I don’t think I need to worry too much about the treedepth warnings but I understand that the divergent transitions can be a sign of incomplete exploration in regions of high curvature that result in biased estimates. Hopefully, that gives you enough to propose some troubleshooting strategies.

Cheers,
Rob

1 Like

That’s exactly right. So we need to work out where these high-curvature regions come from. I’m going to ask for a few more diagnostics. What we’re looking for, and which in my experience causes divergent transitions, are:

  1. strong or intricate correlations between parameters. If you plot the posterior samples for pairs of parameters, you can look for shapes such as thin ellipsis, funnels, etc.
  2. heavy tails or plateaus in the parameter space.

The pairs plots are useful here, because they’ll also tell you where the divergences and the exceeded tree depths occur. In Rstan, you can use e.g.

pairs(stanfit, pars = c(beta_left, beta_right))

There are two things to be careful about: (i) pairs plots over many parameters can be difficult to read, so you may have to use multiple plots; and (ii) ideally we’ll look at the parameter space over which HMC samples.

We can worry about (ii) later, unless the pairs plot give us an option to do that. @bgoodri and @jonah, is this doable?

HMC samples over an unconstrained space. So for example, when you specify real<lower=0> dL, Stan is running HMC over log dL. I’m not quite sure what transformation is used for simplices.

So, things to look at, for parameters declared in the parameters block,:

  • pairs plots
  • trace plots (to look for heavy tails)

If you want to look at parameters over the unconstrained space, great, but we can start with something simpler.

1 Like

It isn’t done, although there is the option to plot the logarithm of variables when all the draws are positive (irrespective of whether that was enforced by the parameterization).

Charles,

Thanks for the guidance.

I’ve written a little script that generates the diagnostic plots that you suggested I look at:

library(rio)
library(rstan)
library(bayesplot)
library(ggplot2)

color_scheme_set("darkgray")

data <- rio::import("liverpool.data.json")

fit1 <- read_stan_csv("samples1.csv")
fit2 <- read_stan_csv("samples2.csv")
fit3 <- read_stan_csv("samples3.csv")
fit4 <- read_stan_csv("samples4.csv")
fit5 <- read_stan_csv("samples5.csv")
fit6 <- read_stan_csv("samples6.csv")
fit <- sflist2stanfit(c(fit1, fit2, fit3, fit4, fit5, fit6))

posterior <- as.array(fit)

np <- nuts_params(fit)

parameters = c(lapply(1:3, function(i){sprintf("initial_state_raw[%s]", i)}),
               lapply(1:data$n_beta_pieces, function(i){sprintf("beta_left[%s]", i)}),
               lapply(1:data$n_beta_pieces, function(i){sprintf("beta_right[%s]", i)}),
               "dL",
               "dI",
               "dT",
               "omega",
               "reciprocal_phi_deaths",
               "reciprocal_phi_calls_111",
               "rho_phi",
               "rho_lambda",
               lapply(1:data$n_rho_pieces, function(i){sprintf("rho[%s]", i)}),
               lapply(1:14, function(i){sprintf("lag_weights[%s]", i)})
               )

for (i in 1:(length(parameters)-1)) {
  for (j in (i+1):length(parameters)){
    mcmc_scatter(posterior,
                 pars = c(toString(parameters[i]), toString(parameters[j])),
                 np = np,
                 size = 1)
    ggsave(paste0("pair_plots/", toString(parameters[i]), "_vs_", toString(parameters[j]), ".pdf"))
  }
}

for (i in 1:length(parameters)) {
  mcmc_trace(posterior, pars = toString(parameters[i]), np = np) + xlab("Post-warmup iteration")
  ggsave(paste0("trace_plots/", toString(parameters[i]), "_trace.pdf"))
}

Due to the number of model parameters, I’ve opted to generate a separate scatter plot for each of the 3741 unique pairs. The script is still running, but I thought it would be worth sharing some preliminary results.

The samples in all of the plots that feature initial_state_raw[1] form an unusual set of parallel ridges at constant values of initial_state_raw[1]. For some reason, it seems that the sample components in this dimension are being quantised. The following figure is typical of all of the plots featuring initial_state_raw[1]:

initial_state_raw[1]_vs_beta_right[19].pdf (291.2 KB)

The pair plots of the other two components of the initial_state_raw simplex don’t have these parallel ridges. The following figures are typical of this subset of plots:

initial_state_raw[2]_vs_beta_left[9].pdf (314.1 KB) initial_state_raw[3]_vs_dL.pdf (315.8 KB)

I suspect this simplex parameter is the cause of (at least some of) the divergences, but I’m keen to know what you think.

All of the plots of combinations of beta_left components, beta_right components, dL, dI, dT, omega, reciprocal_phi_deaths, reciprocal_phi_calls_111, rho_phi, and rho_lambda look like they would be fairly sensible Gaussian-like clouds in unconstrained space. The following figure is fairly typical of these plots:

beta_left[3]_vs_omega.pdf (341.5 KB)

The story is slightly different for plots involving components of rho. Elements at the start of the vector (for early time intervals) produce Gaussian-like clouds, but element towards the end have a concentration of points (at least in constrained space) near the upper constraint (of 1):

beta_left[5]_vs_rho[2].pdf (340.0 KB) beta_left[5]_vs_rho[7].pdf (342.2 KB)

My hunch is that these parameters aren’t causing issues, but are faithfully reflecting the increased uncertainty in rho at more recent times.

Finally, plots featuring components of lag_weights appear to have a reasonable concentration of samples near zero. I don’t know, but I suspect, this concentration at low values would be less apparent in unconstrained space. The following figure is fairly typical of this subset of plots:

beta_left[3]_vs_lag_weights[12].pdf (320.7 KB)

If it’s helpful, I can share the whole ~1.3GB collection of pairs plots with you when the script has finished.

I’m interested to know what you make of this.

Cheers,
Rob

2 Likes

That’s the limitation with the pairs plots: we can’t really examine all of them. That’s why we want to improve divergent transitions diagnostics.

I don’t see any clear patterns in the pairs plots you generated. Regarding the discretization for initial_state, I agree it’s odd. Note that a simplex also needs to be plotted on the unconstrained scale (it’ll be a logistic type transformation, and you’ll only worry about the first two elements).

If you think the simplex is causing the issue, you can try fixing it, and see how well the simplified model fits (i.e. do we remove the divergent transitions?). It’s usually good to have a simple reference we can fit, and then we can work out how the complications impact our inference.

In order to help, I need to learn more about the model, so I think it’d be good to talk. You can walk me through the different components of the model.

2 Likes

Yes, inspecting all of the pairs plots is too time-intensive. A tool that automates this process would be handy.

I thought you’d be interested to know that I’ve managed to get rid of the treedepth and divergence warnings by replacing the 3-simplex for the initial state with a couple of parameters constrained to take values between 0 and 1. I also had to deviate from Stan’s defaults and set adapt delta=0.9 and max_depth=12 to eliminate the last divergent transition and handful of treedepth warnings. The latest iteration of the Stan file is as follows:

functions {
  real[] seeiittd(real time,
                  real[] state,
                  real[] params,
                  real[] real_data,
                  int[] integer_data) {

    // Unpack integer data values
    int T = integer_data[1];
    int n_beta_pieces = integer_data[2];
    int n_disease_states = integer_data[3];

    // Unpack real data values
    real beta_left_t[n_beta_pieces] = real_data[1:n_beta_pieces];
    real beta_right_t[n_beta_pieces] = real_data[n_beta_pieces+1:2*n_beta_pieces];
    real population = real_data[2*n_beta_pieces+1];

    // Unpack parameter values
    real beta_left[n_beta_pieces] = params[1:n_beta_pieces];
    real grad_beta[n_beta_pieces] = params[n_beta_pieces+1:2*n_beta_pieces];
    real nu = params[2*n_beta_pieces+1];
    real gamma = params[2*n_beta_pieces+2];
    real kappa = params[2*n_beta_pieces+3];
    real omega = params[2*n_beta_pieces+4];

    // Unpack state
    real S = state[1];
    real E1 = state[2];
    real E2 = state[3];
    real I1 = state[4];
    real I2 = state[5];
    real T1 = state[6];
    real T2 = state[7];
    real D = state[8];

    real infection_rate;

    for (i in 1:n_beta_pieces) {
      if(time >= beta_left_t[i] && time < beta_right_t[i]) {
        real beta = grad_beta[i] * (time - beta_left_t[i]) + beta_left[i];
        infection_rate = beta * (I1 + I2) * S / population;
      }
    }

    real nuE1 = nu * E1;
    real nuE2 = nu * E2;
    real gammaI1 = gamma * I1;
    real gammaI2 = gamma * I2;
    real kappaT1 = kappa * T1;
    real kappaT2 = kappa * T2;

    real dS_dt = -infection_rate;
    real dE1_dt = infection_rate - nuE1;
    real dE2_dt = nuE1 - nuE2;
    real dI1_dt = nuE2 - gammaI1;
    real dI2_dt = gammaI1 - gammaI2;
    real dT1_dt = gammaI2 * omega - kappaT1;
    real dT2_dt = kappaT1 - kappaT2;
    real dD_dt = kappaT2;

    return {dS_dt, dE1_dt, dE2_dt, dI1_dt, dI2_dt, dT1_dt, dT2_dt, dD_dt};
  }

  real[ , ] integrate_ode_explicit_trapezoidal(real[] initial_state, real initial_time, real[] times, real[] params, real[] real_data, int[] integer_data) {
    real h;
    vector[size(initial_state)] dstate_dt_initial_time;
    vector[size(initial_state)] dstate_dt_tidx;
    vector[size(initial_state)] k;
    real state_estimate[size(times),size(initial_state)];

    h = times[1] - initial_time;
    dstate_dt_initial_time = to_vector(seeiittd(initial_time, initial_state, params, real_data, integer_data));
    k = h*dstate_dt_initial_time;
    state_estimate[1,] = to_array_1d(to_vector(initial_state) + h*(dstate_dt_initial_time + to_vector(seeiittd(times[1], to_array_1d(to_vector(initial_state)+k), params, real_data, integer_data)))/2);

    for (tidx in 1:size(times)-1) {
      h = (times[tidx+1] - times[tidx]);
      dstate_dt_tidx = to_vector(seeiittd(times[tidx], state_estimate[tidx], params, real_data, integer_data));
      k = h*dstate_dt_tidx;
      state_estimate[tidx+1,] = to_array_1d(to_vector(state_estimate[tidx,]) + h*(dstate_dt_tidx + to_vector(seeiittd(times[tidx+1], to_array_1d(to_vector(state_estimate[tidx,])+k), params, real_data, integer_data)))/2);
    }

    return state_estimate;
  }
}
data {
  real initial_time;
  int<lower=1> n_beta_pieces;
  real<lower=0> beta_left_t[n_beta_pieces];
  real<lower=0> beta_right_t[n_beta_pieces];
  int<lower=1> n_rho_pieces;
  int<lower=0> rho_left_t[n_rho_pieces];
  int<lower=0> rho_right_t[n_rho_pieces];
  int<lower=1> T;
  real times[T];
  int<lower=1> n_disease_states;
  real<lower=0> population;
  int<lower=1> deaths_length;
  int<lower=1> deaths_starts[deaths_length];
  int<lower=1> deaths_stops[deaths_length];
  int<lower=0> deaths[deaths_length];
  int<lower=1> calls_111_length;
  int<lower=1> calls_111_start;
  int<lower=0> calls_111[calls_111_length];
  int real_data_length;
  real real_data[real_data_length];
  int integer_data_length;
  int integer_data[integer_data_length];
}
transformed data {
  real mu_dL = 4.00;
  real sigma_dL = 0.20;
  real mu_dI = 3.06;
  real sigma_dI = 0.21;
  real mu_dT = 16.00;
  real sigma_dT = 0.71;
  int max_lag = 13;
}
parameters {
  real<lower=0, upper=1> initial_state_raw[2];
  real<lower=0> beta_left[n_beta_pieces];
  real<lower=0> beta_right[n_beta_pieces];
  real<lower=0> dL;
  real<lower=0> dI;
  real<lower=0> dT;
  real<lower=0, upper=1> omega;
  real<lower=0> reciprocal_phi_deaths;
  real<lower=0> reciprocal_phi_calls_111;
  real<lower=0, upper=1> rho_phi;
  real<lower=0.1> rho_lambda;
  real<lower=0, upper=1> rho[n_rho_pieces];
  simplex[max_lag+1] lag_weights;
}
transformed parameters {
  real initial_state[n_disease_states];
  real grad_beta[n_beta_pieces];
  real nu;
  real gamma;
  real kappa;
  real phi_deaths;
  real phi_calls_111;
  real rho_alpha;
  real rho_beta;
  real state_estimate[T,n_disease_states];
  vector[T+1] S;
  vector[T+1] E1;
  vector[T+1] E2;
  vector[T+1] I1;
  vector[T+1] I2;
  vector[T+1] T1;
  vector[T+1] T2;
  vector[T+1] D;
  vector[T] daily_infections;
  vector[T] daily_deaths;
  vector[T] calls_111_lagged_daily_infections;
  vector[T] daily_calls_111;

  initial_state[1] = (population-5.0)*initial_state_raw[1] + 1.0;
  initial_state[2] = (population-5.0)*(1.0-initial_state_raw[1])*initial_state_raw[2]/2.0 + 1.0;
  initial_state[3] = (population-5.0)*(1.0-initial_state_raw[1])*initial_state_raw[2]/2.0 + 1.0;
  initial_state[4] = (population-5.0)*(1.0-initial_state_raw[1])*(1.0-initial_state_raw[2])/2.0 + 1.0;
  initial_state[5] = (population-5.0)*(1.0-initial_state_raw[1])*(1.0-initial_state_raw[2])/2.0 + 1.0;
  initial_state[6] = 0.0;
  initial_state[7] = 0.0;
  initial_state[8] = 0.0;
  grad_beta = to_array_1d((to_vector(beta_right) - to_vector(beta_left))./(to_vector(beta_right_t) - to_vector(beta_left_t)));
  nu = 2.0/dL;
  gamma = 2.0/dI;
  kappa = 2.0/dT;
  phi_deaths = 1.0 / reciprocal_phi_deaths;
  phi_calls_111 = 1.0 / reciprocal_phi_calls_111;
  rho_alpha = rho_lambda * rho_phi;
  rho_beta = rho_lambda * (1 - rho_phi);

  {
    real params[2*n_beta_pieces+4];
    params[1:n_beta_pieces] = beta_left;
    params[n_beta_pieces+1:2*n_beta_pieces] = grad_beta;
    params[2*n_beta_pieces+1] = nu;
    params[2*n_beta_pieces+2] = gamma;
    params[2*n_beta_pieces+3] = kappa;
    params[2*n_beta_pieces+4] = omega;

    state_estimate = integrate_ode_explicit_trapezoidal(initial_state, initial_time, times, params, real_data, integer_data);
  }

  S = append_row(initial_state[1], to_vector(state_estimate[, 1]));
  E1 = append_row(initial_state[2], to_vector(state_estimate[, 2]));
  E2 = append_row(initial_state[3], to_vector(state_estimate[, 3]));
  I1 = append_row(initial_state[4], to_vector(state_estimate[, 4]));
  I2 = append_row(initial_state[5], to_vector(state_estimate[, 5]));
  T1 = append_row(initial_state[6], to_vector(state_estimate[, 6]));
  T2 = append_row(initial_state[7], to_vector(state_estimate[, 7]));
  D = append_row(initial_state[8], to_vector(state_estimate[, 8]));

  daily_infections = S[:T] - S[2:] + machine_precision();
  daily_deaths = D[2:] - D[:T];

  calls_111_lagged_daily_infections = lag_weights[1]*daily_infections;

  for (i in 1:max_lag) {
    calls_111_lagged_daily_infections += lag_weights[i+1]*append_row(rep_vector(0.0, i), daily_infections[:T-i]);
  }

  daily_calls_111 = rep_vector(0.0, T);

  for (i in 1:n_rho_pieces) {
    daily_calls_111[rho_left_t[i]:rho_right_t[i]-1] = calls_111_lagged_daily_infections[rho_left_t[i]:rho_right_t[i]-1] * rho[i];
  }
}
model {
  initial_state_raw[1] ~ beta(5.0, 0.5);
  initial_state_raw[2] ~ beta(1.1, 1.1);
  beta_left ~ normal(0, 0.5);
  beta_right ~ normal(0, 0.5);
  dL ~ normal(mu_dL, sigma_dL);
  dI ~ normal(mu_dI, sigma_dI);
  dT ~ normal(mu_dT, sigma_dT);
  omega ~ beta(100, 9803);
  reciprocal_phi_deaths ~ exponential(5);
  reciprocal_phi_calls_111 ~ exponential(5);
  rho_phi ~ beta(1, 1);
  rho_lambda ~ pareto(0.1, 1.5);
  rho ~ beta(rho_alpha, rho_beta);
  lag_weights ~ dirichlet(rep_vector(0.1, max_lag+1));

  for (i in 1:deaths_length) {
    target += 7.0*neg_binomial_2_lpmf(deaths[i] | sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  for (i in 1:n_rho_pieces) {
    target += neg_binomial_2_lpmf(calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] | daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}
generated quantities {
  vector[T-1] growth_rate = (log(daily_infections[2:]) - log(daily_infections[:T-1]))*100;

  int pred_deaths[deaths_length];

  for (i in 1:deaths_length) {
    pred_deaths[i] = neg_binomial_2_rng(sum(daily_deaths[deaths_starts[i]:deaths_stops[i]]), phi_deaths);
  }

  int pred_calls_111[calls_111_length];

  for (i in 1:n_rho_pieces) {
    pred_calls_111[rho_left_t[i]-calls_111_start+1:rho_right_t[i]-calls_111_start] = neg_binomial_2_rng(daily_calls_111[rho_left_t[i]:rho_right_t[i]-1], phi_calls_111);
  }
}

Even though I’m now getting an error-free fit, I’m keen to take you up on this offer. It would be great to get your thoughts on the model.

2 Likes

I know it’s quite generic, but check out Divergent transitions - a primer for some general strategies.

I think the part most immediately relevant is to simplify the model to find a) the most complex model you can reliably fit (preferably without altering tree depth or adapt_delta) and b) the least complex model that has issues. Similarly if the issues manifest also on a reduced dataset it is easier to debug with it. This can simplify debugging a lot.

1 Like

Thanks - I’ll make sure to give this a read!