Effective COI (eCOI) Demonstration

Background

The legacy moire sampler jointly estimates complexity of infection (COI) and within-host relatedness. Those quantities are not always separately identifiable from genotyping data alone, which can produce multimodal posteriors and poor mixing along a ridge in \((m, r)\) space.

The marginalized eCOI model (marginal_ecoi = TRUE) instead treats effective COI (\(e\)) as the headline quantity. Effective COI is the expected number of distinct alleles at a locus with infinite diversity (\(H_e = 1\)). With zero within-host relatedness, \(e\) equals the true COI; as relatedness increases, \(e\) shrinks toward 1:

\[ e = (m - 1)(1 - r) + 1 \]

where \(m\) is COI and \(r\) is the probability that an additional strain is genetically identical to an existing one.

The in-fit model uses a hybrid partial-collapse scheme:

Effective COI \(e\) is updated by Metropolis, with discrete COI integrated out under a zero-truncated Poisson (ZTPois) prior on \(m\) and uniform relatedness (Jacobian only).
Discrete COI \(m\) is imputed each sweep by a Gibbs draw \(m \mid e, p, G\), with \(r = r(m, e)\).
Allele frequencies \(p\) are updated against the explicit transmission likelihood at the current \((m, r)\), not the fully collapsed \(\int m\) objective.

This split avoids the allele-frequency bias that arises when \(p\) is updated against a likelihood that has integrated \(m\) out. Allele-frequency recovery is sensitive to initialization: starting \(p\) near the generative truth yields tight posterior tracking, while the default clustering-based initializer can leave \(p\) in a basin that drifts toward the collapsed-\(\int m\) attractor. Use prepare_initial_allele_frequencies() when you have external population-frequency estimates.

Population-level COI heterogeneity enters through a ZTPois rate \(\lambda\) (population_lam_coi), with a low-favoring Gamma hyperprior (population_coi_lam_shape, population_coi_lam_rate). There is no explicit prior on \(e\) itself.

Simulation and MCMC

set.seed(20260622)

n_loci <- 20L
n_allele <- 5L
n_samp <- 60L
max_coi <- 25L

true_lam <- 1.5   # ZTPois rate for discrete COI
mean_coi_true <- true_lam / (1 - exp(-true_lam))

ecoi_of <- function(m, r) (m - 1) * (1 - r) + 1
r_of <- function(m, e) 1 - (e - 1) / (m - 1)

rztpois <- function(n, lambda) {
  vapply(seq_len(n), function(i) {
    repeat {
      x <- rpois(1, lambda)
      if (x >= 1L) return(x)
    }
  }, integer(1))
}

e_true <- numeric(n_samp)
m_true <- integer(n_samp)
r_true <- numeric(n_samp)
for (i in seq_len(n_samp)) {
  m <- min(max(rztpois(1, true_lam), 2L), max_coi)
  r <- runif(1, 0.05, 0.95)
  e <- ecoi_of(m, r)
  e <- min(e, max_coi - 1e-3)
  m <- min(max(2L, round((e - 1) / (1 - r) + 1)), max_coi)
  r <- r_of(m, e)
  if (r <= 0) r <- 1e-3
  m_true[i] <- m
  r_true[i] <- r
  e_true[i] <- ecoi_of(m, r)
}

locus_freq_alphas <- replicate(n_loci, rep(1, n_allele), simplify = FALSE)

simulated_ecoi_data <- moire::simulate_data(
  num_samples = n_samp,
  epsilon_pos = 0,
  epsilon_neg = 0,
  missingness = 0,
  sample_cois = m_true,
  internal_relatedness = r_true,
  locus_freq_alphas = locus_freq_alphas
)

Because collapsing COI removes the implicit cap that discrete COI places on latent genotype size, we use informative genotyping-error priors so the model cannot explain away observed diversity as false positives or dropouts. The per-locus false-positive rate is additionally anchored near zero with eps_pos_locus_alpha / eps_pos_locus_beta.

mcmc_ecoi_results <- moire::run_mcmc(
  simulated_ecoi_data,
  is_missing = simulated_ecoi_data$is_missing,
  allow_relatedness = TRUE,
  marginal_ecoi = TRUE,
  num_populations = 1,
  population_coi_lam_shape = 0.1,
  population_coi_lam_rate = 10,
  eps_neg_alpha = 1,
  eps_neg_beta = 200,
  max_eps_neg = 0.2,
  eps_pos_alpha = 1,
  eps_pos_beta = 200,
  max_eps_pos = 0.05,
  eps_pos_locus_alpha = 1,
  eps_pos_locus_beta = 2000,
  pt_chains = 1,
  burnin = 1500,
  samples_per_chain = 1500,
  thin = 1,
  max_coi = max_coi,
  verbose = TRUE
)

After running the MCMC, we summarize per-sample effective COI, population COI rate \(\lambda\), and allele frequencies.

ecoi_summary <- moire::summarize_effective_coi(mcmc_ecoi_results)
he_summary <- moire::summarize_he(mcmc_ecoi_results)
allele_freq_summary <- moire::summarize_allele_freqs(mcmc_ecoi_results)

sample_data <- data.frame(
  ecoi_summary,
  true_ecoi = e_true,
  true_coi = m_true,
  true_relatedness = r_true * (m_true > 1)
)

chain <- mcmc_ecoi_results$chains[[1]]
pop_draws <- function(x, p = 1L) vapply(x, function(v) v[[p]], numeric(1))

pop_hyperparams <- data.frame(
  param = c("lambda", "mean_coi"),
  true = c(true_lam, mean_coi_true),
  post_mean = c(
    mean(chain$population_lam_coi),
    mean(chain$mean_coi)
  ),
  post_lower = c(
    quantile(chain$population_lam_coi, 0.025),
    quantile(chain$mean_coi, 0.025)
  ),
  post_upper = c(
    quantile(chain$population_lam_coi, 0.975),
    quantile(chain$mean_coi, 0.975)
  )
)

he_data <- data.frame(
  he_summary,
  true_he = sapply(
    moire::calculate_naive_allele_frequencies(simulated_ecoi_data$true_genotypes),
    function(x) moire::calculate_he(x)
  ),
  naive_he = sapply(
    moire::calculate_naive_allele_frequencies(simulated_ecoi_data$data),
    function(x) moire::calculate_he(x)
  )
)

allele_freq_data <- data.frame(
  allele_freq_summary,
  naive_allele_frequency = unlist(
    moire::calculate_naive_allele_frequencies(simulated_ecoi_data$data)
  ),
  true_allele_frequency = unlist(
    moire::calculate_naive_allele_frequencies(simulated_ecoi_data$true_genotypes)
  )
)

p_true <- unlist(
  moire::calculate_naive_allele_frequencies(simulated_ecoi_data$true_genotypes)
)
p_post <- allele_freq_data$post_allele_freqs_mean
p_l1 <- mean(abs(p_post - p_true))
p_in_ci <- mean(
  p_true >= allele_freq_data$post_allele_freqs_lower &
    p_true <= allele_freq_data$post_allele_freqs_upper
)

Estimating effective COI

The marginalized sampler draws \(e\) directly while integrating discrete COI out of the \(e\)-move likelihood. Here we compare posterior means and 95% credible intervals to the simulated truth.

ggplot(sample_data, aes(
  x = true_ecoi,
  y = post_effective_coi_mean,
  ymin = post_effective_coi_lower,
  ymax = post_effective_coi_upper
)) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", alpha = 0.5) +
  geom_errorbar() +
  geom_point() +
  xlab("True effective COI") +
  ylab("Estimated effective COI") +
  theme_classic(base_size = 12) +
  expand_limits(x = 1, y = 1)

We can also examine estimation error ordered by true eCOI.

ecoi_compare_to_truth <- sample_data |>
  dplyr::select(
    post_effective_coi_mean, true_ecoi, sample_id,
    post_effective_coi_lower, post_effective_coi_upper
  ) |>
  dplyr::mutate(
    sample_id = forcats::fct_reorder(sample_id, true_ecoi, min),
    ymin = post_effective_coi_lower - true_ecoi,
    ymax = post_effective_coi_upper - true_ecoi
  ) |>
  dplyr::arrange(true_ecoi)

ecoi_bins <- sample_data |>
  dplyr::group_by(true_ecoi) |>
  dplyr::summarize(n = dplyr::n()) |>
  dplyr::mutate(pos = cumsum(n))

ggplot(
  data = ecoi_compare_to_truth,
  aes(
    x = sample_id,
    y = post_effective_coi_mean - true_ecoi,
    ymin = ymin,
    ymax = ymax
  )
) +
  geom_errorbar() +
  geom_point() +
  geom_hline(yintercept = 0, linetype = "dashed", alpha = .5) +
  geom_vline(data = ecoi_bins, aes(xintercept = pos), linetype = "dotted", alpha = .5) +
  xlab("Sample") +
  ylab("Estimate - true eCOI") +
  theme_classic(base_size = 12) +
  theme(
    axis.text.x = element_blank(),
    axis.ticks.x = element_blank()
  )

Population COI rate

Discrete COI is modeled as \(m \sim \mathrm{ZTPois}(\lambda)\) with population rate \(\lambda\) given a low-favoring Gamma hyperprior. Posterior draws are returned in population_lam_coi; implied mean COI is mean_coi \(= \lambda / (1 - e^{-\lambda})\).

ggplot(pop_hyperparams, aes(x = param, y = post_mean, ymin = post_lower, ymax = post_upper)) +
  geom_pointrange() +
  geom_point(aes(y = true), color = "red", size = 3) +
  xlab(NULL) +
  ylab("Value") +
  theme_classic(base_size = 12)

Estimating allele frequencies

With hybrid partial collapse, \(p\) is updated at explicit \((m, r)\) each sweep after the Gibbs imputation of \(m\). This bundled run uses the default clustering-based initializer for \(p\) (no oracle truth). Recovery is reasonable but not as tight as when \(p\) starts at the simulated truth; see inst/scripts/diagnose_ecoi_p_recovery.R (ECOI_TRUE_P=1 for the latter).

Clustering-based \(p\) init: mean per-locus L1 error \(= 0.024\); 36% of true allele frequencies fall inside the 95% credible interval (vs. ~99% with true-\(p\) init on the same simulation).

ggplot(allele_freq_data) +
  geom_errorbar(aes(
    y = post_allele_freqs_mean,
    x = true_allele_frequency,
    ymax = post_allele_freqs_upper,
    ymin = post_allele_freqs_lower
  )) +
  geom_point(aes(y = post_allele_freqs_mean, x = true_allele_frequency)) +
  geom_point(aes(y = naive_allele_frequency, x = true_allele_frequency), color = "red", alpha = .3) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", alpha = .5) +
  ylab("Mean posterior frequency") +
  xlab("True frequency") +
  theme_classic(base_size = 12) +
  expand_limits(x = 0, y = 0) +
  ggtitle("Allele frequency estimates vs truth")

Estimating expected heterozygosity

ggplot(he_data, aes(x = true_he)) +
  geom_errorbar(aes(y = post_stat_mean, ymax = post_stat_upper, ymin = post_stat_lower)) +
  geom_point(aes(y = post_stat_mean), color = "black") +
  geom_point(aes(y = naive_he), color = "red") +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", alpha = .5) +
  xlab("True heterozygosity") +
  ylab("Mean posterior heterozygosity") +
  theme_classic()