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 PMF length of one delay, i.e. its maximum delay plus one
 *
 * @param i Index of the delay among all delays
 * @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_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 The length of the PMF of delay i
 *
 * @ingroup delay_handlers
 */
int get_delay_length(
  int i, array[] int delay_types_p, array[] int delay_types_id,
  array[] int delay_max, array[] int delay_np_pmf_groups
) {
  int id = delay_types_id[i];
  if (delay_types_p[i]) {
    return delay_max[id] + 1;
  }
  return delay_np_pmf_groups[id + 1] - delay_np_pmf_groups[id];
}
 
/**
 * 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 = rep_array(0, delay_types);
  for (i in 1:delay_types) {
    for (j in delay_types_groups[i]:(delay_types_groups[i + 1] - 1)) {
      ret[i] += get_delay_length(
        j, delay_types_p, delay_types_id, delay_max, delay_np_pmf_groups
      ) - 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.
 *   The cumulative sum is taken before any reversal, so setting both gives
 *   the reversed CDF, as used for right truncation.
 * @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
) {
  vector[len] pmf = rep_vector(0, len);
  int n = 0; // length of the PMF combined so far
  for (i in
         delay_types_groups[delay_id]:(delay_types_groups[delay_id + 1] - 1)) {
    int id = delay_types_id[i];
    int m = get_delay_length(
      i, delay_types_p, delay_types_id, delay_max, delay_np_pmf_groups
    );
    int new_n = n == 0 ? m : n + m - 1;
    // non-parametric PMFs are sliced directly, without a named copy
    if (delay_types_p[i]) {
      vector[m] delay_pmf = discretised_pmf(
        delay_params[delay_params_groups[id]:(delay_params_groups[id + 1] - 1)],
        m, delay_dist[id], 0
      );
      if (n == 0) {
        pmf[1:m] = delay_pmf;
      } else {
        pmf[1:new_n] = convolve_with_rev_pmf(
          pmf[1:n], reverse(delay_pmf), new_n
        );
      }
    } else {
      int start = delay_np_pmf_groups[id];
      if (n == 0) {
        pmf[1:m] = delay_np_pmf[start:(start + m - 1)];
      } else {
        pmf[1:new_n] = convolve_with_rev_pmf(
          pmf[1:n], reverse(delay_np_pmf[start:(start + m - 1)]), new_n
        );
      }
    }
    n = new_n;
  }
  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) {
        target += weight[d] *
          normal_lpdf(
            delay_params[s] | 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;
}