nickbrazeau / polysimibd Goto Github PK

get_ver_realized_ibd <- function(swfsim, smpl1, smpl2) {
  # pull out ARG pieces
  # trees point left --  don't care about number of strains w/in, if any ibd, then loci is ibd
  arg <- polySimIBD::get_arg(swf = swfsim, host_index = c(smpl1, smpl2))
  conn <- purrr::map(arg, "c")

  # now look at loci by loci
  get_loci_pairwise_ibd <- function(conni, this_coi) {
    smpl1con <- conni[1:this_coi[1]]
    smpl2con <- conni[(this_coi[1]+1):(cumsum(this_coi)[2])]
    # get IBD connections between 1 and 2
    pwconn <- which(smpl2con %in% 0:(this_coi[1]-1) )
    # if any, IBD
    pwconn <- sum(pwconn)
    return(pwconn >= 1)
  }
  # iterate through
  lociibd <- purrr::map_dbl(.x = conn,
                             .f = get_loci_pairwise_ibd, this_coi = swfsim$coi[c(smpl1, smpl2)])
  # loci IBD
  return(sum(lociibd)/length(lociibd))
}

polySimIBD::get_realized_pairwise_ibd
function (swf, host_index = NULL) 
{
    assert_custom_class(swf, "swfsim")
    assert_pos_int(host_index)
    assert_length(host_index, 2)
    arg <- polySimIBD::get_arg(swf = swf, host_index = host_index)
    coi1 <- polySimIBD::get_realized_coi(swf = swf, host_index = host_index[[1]])
    coi1 <- max(coi1)
    coi2 <- polySimIBD::get_realized_coi(swf = swf, host_index = host_index[[2]])
    coi2 <- max(coi2)
    conn <- purrr::map(arg, "c")
    tm <- purrr::map(arg, "t")
    get_pairwise_ibd <- function(conni, tmi, this_coi) {
        smpl1con <- conni[1:this_coi[1]]
        smpl2con <- conni[(this_coi[1] + 1):(cumsum(this_coi)[2])]
        pwconn <- which(smpl2con %in% 0:(this_coi[1] - 1))
        locimatches <- rep(1, length(pwconn))
        if (length(pwconn) != 0) {
            for (i in 1:length(pwconn)) {
                haplotypeindex <- this_coi[1] + pwconn[i] - 1
                internalconn <- which(smpl2con %in% haplotypeindex)
                if (length(internalconn) != 0) {
                  for (i in 1:length(internalconn)) {
                    internalhaplotypeplace <- this_coi[1] + internalconn[i]
                    if (tmi[internalhaplotypeplace] < tmi[this_coi[1] + 
                      internalconn[i]]) {
                      locimatches[i] <- locimatches[i] + 1
                    }
                  }
                }
            }
        }
        return(sum(locimatches))
    }
    effCOI <- c(coi1, coi2)
    numerator <- purrr::map2_dbl(.x = conn, .y = tm, .f = get_pairwise_ibd, 
        this_coi = swf$coi[host_index])
    numerator <- ifelse(numerator > min(effCOI), min(effCOI), 
        numerator)
    pairwiseIBD <- sum(numerator)/(min(effCOI) * length(conn))
    return(pairwiseIBD)
}

# COI_equilibrium.R
#
# Author: Bob Verity
# Date: 2022-08-16
#
# Purpose:
# Derive the equilibrium COI distribution under the PolySimIBD model for
# infinite population size.
#
# ------------------------------------------------------------------

library(copula) # needed for Stirling numbers

# define model parameters
lambda <- 3
m <- 0.1
K_max <- 20

# calculate transition matrix
K_mat <- matrix(NA, K_max, K_max)
for (K_prev in 1:K_max) {
  
  p2 <- rep(0, K_max)
  for (i in 1:K_max) {
    N_clonal <- i
    k <- 1:min(N_clonal, K_prev)
    p1 <- exp(lfactorial(K_prev) - lfactorial(K_prev - k) - N_clonal*log(K_prev)) * Stirling2.all(N_clonal)[k]
    p2 <- p2 + dpois(i, lambda * (1 - m) / K_prev) * c(p1, rep(0, K_max - length(p1)))
  }
  p2 <- c(dpois(0, lambda * (1 - m) / K_prev), p2)
  
  p3 <- rep(NA, K_max + 1)
  lambda2 <- lambda * m + lambda * (1 - m) * (1 - 1 / K_prev)
  for (i in 0:K_max) {
    p3[i + 1] <- sum(p2[(0:i) + 1] * dpois(i:0, lambda2))
  }
  
  K_mat[K_prev,] <- p3[-1]
}
K_mat <- K_mat / (1 - dpois(0, lambda))

# calculate equilibrium solution by applying transition a large number of times
# (I'm having trouble getting the Eigenvalue method to work)
K_init <- rep(1/K_max, K_max)
for (i in 1:100) {
  K_init <- K_init %*% K_mat
}

# barplot of equilibrium distribution. Red dots show the raw Poisson(lambda)
# distribtion, i.e. what we would expect under pure super-infection from an
# infinitely large population
barplot(K_init, space = 0, names.arg = 1:K_max, col = grey(0.6))
points(1:K_max - 0.5, dpois(1:K_max, lambda) / (1 - exp(-lambda)), pch = 20, col = 2)