delays.stan

Go to the documentation of this file.

/**
 * Delay Distribution Functions
 *
 * This group of functions handles the creation, manipulation, and application
 * of delay distributions in the model. Delay distributions represent the time
 * between events (e.g., infection to symptom onset, symptom onset to reporting).
 *
 * Per-delay values (PMFs, parameters) are packed into ragged vectors: one
 * vector holds the concatenated values for every delay, and a companion
 * `*_groups` array of lookup indices marks where each delay's segment
 * starts and ends within that vector. See the Stan User's Guide section on
 * ragged data structures:
 * https://mc-stan.org/docs/stan-users-guide/sparse-ragged.html#ragged-data-structs.section
 *
 */
 
/**
 * Get the maximum delay for each delay type
 *
 * @param delay_types Number of delay types
 * @param delay_types_p Array indicating whether each delay is parametric (1) or non-parametric (0)
 * @param delay_types_id Array mapping delay types to their respective IDs
 * @param delay_types_groups Array of lookup indices for the ragged groups
 *   of delay types
 * @param delay_max Array of maximum delays for parametric distributions
 * @param delay_np_pmf_groups Array of lookup indices for the ragged
 *   non-parametric PMF vector
 * @return An array of maximum delays for each delay type
 *
 * @ingroup delay_handlers
 */
array[] int get_delay_type_max(
  int delay_types, array[] int delay_types_p, array[] int delay_types_id,
  array[] int delay_types_groups, array[] int delay_max,
  array[] int delay_np_pmf_groups
) {
  array[delay_types] int ret;
  for (i in 1:delay_types) {
    ret[i] = 0;
    for (j in delay_types_groups[i]:(delay_types_groups[i + 1] - 1)) {
      if (delay_types_p[j]) { // parametric
        ret[i] += delay_max[delay_types_id[j]];
      } else { // nonparametric
        ret[i] += delay_np_pmf_groups[delay_types_id[j] + 1] -
          delay_np_pmf_groups[delay_types_id[j]] - 1;
      }
    }
  }
  return ret;
}
 
/**
 * Get the reversed probability mass function for a delay
 *
 * @param delay_id Identifier for the specific delay distribution
 * @param len Length of the output PMF
 * @param delay_types_p Array indicating whether each delay is parametric (1) or non-parametric (0)
 * @param delay_types_id Array mapping delay types to their respective IDs
 * @param delay_types_groups Array of lookup indices for the ragged groups
 *   of delay types
 * @param delay_max Array of maximum delays for parametric distributions
 * @param delay_np_pmf Ragged vector of probability mass functions for
 *   non-parametric delays
 * @param delay_np_pmf_groups Array of lookup indices for the ragged
 *   non-parametric PMF vector
 * @param delay_params Ragged vector of parameters for parametric delay
 *   distributions
 * @param delay_params_groups Array of lookup indices for the ragged
 *   parameter vector
 * @param delay_dist Array of distribution types using primarycensored
 *   convention (1: lognormal, 2: gamma, 3: weibull, 4: exponential)
 * @param left_truncate Left truncation point (0 for no truncation)
 * @param reverse_pmf Whether to reverse the PMF (1) or not (0)
 * @param cumulative Whether to return cumulative (1) or daily (0) values
 * @return A vector containing the (reversed) PMF of length len
 *
 * @ingroup delay_handlers
 */
vector get_delay_rev_pmf(
  int delay_id, int len, array[] int delay_types_p, array[] int delay_types_id,
  array[] int delay_types_groups, array[] int delay_max,
  vector delay_np_pmf, array[] int delay_np_pmf_groups,
  vector delay_params, array[] int delay_params_groups, array[] int delay_dist,
  int left_truncate, int reverse_pmf, int cumulative
) {
  // loop over delays
  vector[len] pmf = rep_vector(0, len);
  int current_len = 1;
  int new_len;
  for (i in
         delay_types_groups[delay_id]:(delay_types_groups[delay_id + 1] - 1)) {
    if (delay_types_p[i]) { // parametric
      int start = delay_params_groups[delay_types_id[i]];
      int end = delay_params_groups[delay_types_id[i] + 1] - 1;
      vector[delay_max[delay_types_id[i]] + 1] new_variable_pmf =
        discretised_pmf(
          delay_params[start:end],
          delay_max[delay_types_id[i]] + 1,
          delay_dist[delay_types_id[i]],
          0
      );
      new_len = current_len + delay_max[delay_types_id[i]];
      if (current_len == 1) { // first delay
        pmf[1:new_len] = new_variable_pmf;
      } else { // subsequent delay to be convolved
        pmf[1:new_len] = convolve_with_rev_pmf(
          pmf[1:current_len], reverse(new_variable_pmf), new_len
        );
      }
    } else { // nonparametric
      int start = delay_np_pmf_groups[delay_types_id[i]];
      int end = delay_np_pmf_groups[delay_types_id[i] + 1] - 1;
      new_len = current_len + end - start;
      if (current_len == 1) { // first delay
        pmf[1:new_len] = delay_np_pmf[start:end];
      } else { // subsequent delay to be convolved
        pmf[1:new_len] = convolve_with_rev_pmf(
          pmf[1:current_len], reverse(delay_np_pmf[start:end]), new_len
        );
      }
    }
    current_len = new_len;
  }
  if (left_truncate) {
    pmf = append_row(
      rep_vector(0, left_truncate),
      pmf[(left_truncate + 1):len] / sum(pmf[(left_truncate + 1):len])
    );
  }
  if (cumulative) {
    pmf = cumulative_sum(pmf);
  }
  if (reverse_pmf) {
    pmf = reverse(pmf);
  }
  return pmf;
}
 
/**
 * Update log density for delay distribution priors
 *
 * @param delay_params Ragged vector of parameters for parametric delay
 *   distributions
 * @param delay_params_mean Ragged vector of prior means for delay parameters
 * @param delay_params_sd Ragged vector of prior standard deviations for
 *   delay parameters
 * @param delay_params_groups Array of lookup indices for the ragged
 *   parameter vectors
 * @param delay_dist Array of distribution types using primarycensored
 *   convention (1: lognormal, 2: gamma, 3: weibull, 4: exponential)
 * @param weight Array of weights for each delay distribution in the log density
 *
 * @ingroup delay_handlers
 */
void delays_lp(vector delay_params,
               vector delay_params_mean, vector delay_params_sd,
               array[] int delay_params_groups,
               array[] int delay_dist, array[] int weight) {
  int n_delays = num_elements(delay_params_groups) - 1;
  if (n_delays == 0) {
    return;
  }
  for (d in 1:n_delays) {
    int start = delay_params_groups[d];
    int end = delay_params_groups[d + 1] - 1;
    for (s in start:end) {
      if (delay_params_sd[s] > 0) {
        if (weight[d] > 1) {
          target += weight[d] *
            normal_lpdf(
              delay_params[s] | delay_params_mean[s], delay_params_sd[s]
            );
        } else {
          delay_params[s] ~ normal(delay_params_mean[s], delay_params_sd[s]);
        }
      }
    }
  }
}
 
/**
 * Update log prior density for estimated nonparametric delays
 *
 * Applies Gamma(alpha, 1) priors to the raw vector that backs the
 * estimated nonparametric PMFs. Once normalised within each ragged
 * segment, this induces a Dirichlet(alpha) prior on the segment.
 * See https://mc-stan.org/docs/stan-users-guide/simplexes.html for
 * the gamma-normalisation construction.
 *
 * @param delay_np_est_raw Raw gamma-distributed values backing each
 *   estimated nonparametric PMF segment.
 * @param delay_np_est_alpha Dirichlet concentration parameters,
 *   matched element-wise to delay_np_est_raw.
 *
 * @ingroup delay_handlers
 */
void delays_np_lp(
  vector delay_np_est_raw, vector delay_np_est_alpha
) {
  if (num_elements(delay_np_est_raw) == 0) return;
  delay_np_est_raw ~ gamma(delay_np_est_alpha, 1);
}
 
/**
 * Combine fixed and estimated nonparametric delay PMFs
 *
 * Returns a copy of the fixed PMF vector with the estimated entries
 * overwritten by per-segment normalisation of the raw gamma values.
 * Fixed entries and structural zeros are left unchanged.
 *
 * @param delay_np_pmf Fixed PMF vector (concatenated ragged array).
 * @param delay_n_np_est Number of estimated nonparametric delays.
 * @param delay_np_est_groups Ragged-array boundaries into
 *   delay_np_est_raw (length delay_n_np_est + 1).
 * @param delay_np_est_pos For each estimated element, its position
 *   within delay_np_pmf.
 * @param delay_np_est_raw Raw gamma values backing each segment.
 * @return A vector of length num_elements(delay_np_pmf) containing
 *   the combined PMF.
 *
 * @ingroup delay_handlers
 */
vector combine_np_pmf(
  vector delay_np_pmf, int delay_n_np_est,
  array[] int delay_np_est_groups, array[] int delay_np_est_pos,
  vector delay_np_est_raw
) {
  vector[num_elements(delay_np_pmf)] ret = delay_np_pmf;
  for (i in 1:delay_n_np_est) {
    int es = delay_np_est_groups[i];
    int ee = delay_np_est_groups[i + 1] - 1;
    real seg_sum = sum(delay_np_est_raw[es:ee]);
    for (j in es:ee) {
      ret[delay_np_est_pos[j]] = delay_np_est_raw[j] / seg_sum;
    }
  }
  return ret;
}
 
/**
 * Generate random samples from a normal distribution with lower bounds
 *
 * @param mu Vector of means
 * @param sigma Vector of standard deviations
 * @param lb Vector of lower bounds
 * @return A vector of random samples from the truncated normal distribution
 *
 * @ingroup delay_handlers
 */
vector normal_lb_rng(vector mu, vector sigma, vector lb) {
  int len = num_elements(mu);
  vector[len] ret;
  for (i in 1:len) {
    real p = normal_cdf(lb[i] | mu[i], sigma[i]);  // cdf for bounds
    real u = uniform_rng(p, 1);
    ret[i] = (sigma[i] * inv_Phi(u)) + mu[i];  // inverse cdf for value
  }
  return ret;
}