homing-gd-functions.R

packages = c("tidyverse", "deSolve", "DescTools", "here") 
installed = packages %in% installed.packages()[, "Package"]
if (any(!installed)) {
  install.packages(packages[!installed])
}
lapply(packages, library, character.only = TRUE)

source(here("deSolve-simulations/helper-functions.R"))

## ASSUMPTIONS ###
#   - No flux at the walls (absorbing boundaries)
#   - Constant density
#   - 50/50 ratio of males and females; ie same genotype frequencies in males and females
#   - Parents are sampled at probabilities equal to their genotype frequencies
#   - Genotype survival rates occur exactly at their viability probabilities
#   - Conversion occurs exactly at c


pg_hat_gene_drive = function(s,h,c){
  numerator = c - h*s*(1 + c)
  denominator = s*(1 - 2*h)
  return(numerator/denominator)
}

solve_for_h_gene_drive = function(s, c, pg_hat){
  numerator = pg_hat*s - c
  denominator = s*(2*pg_hat - c- 1)
  return(numerator/denominator)
}

ptilda_gene_drive = function(s,h,c){
  numer.term1 = (2/3)*(3*h + c*h - (c/s) - 1)
  denom = 2*h - 1
  sqrt.term.inside = ((4/9)*(c/s - c*h - 3*h + 1)^2) - 2*(2*h - 1)*(h + c*h - c/s)
  sqrt.term = sqrt(sqrt.term.inside)
  
  ptilda = (numer.term1 - sqrt.term) / denom
  
  return(ptilda)
}

x_bubble_gene_drive = function(s,h,c,p, D){
  # Arguments:
  #   s: selection coefficient
  #   h: dominance parameter
  #   c: conversion rate
  #   p: frequency; make sure it's between 0 and ptilda
  #   D: diffusion constant
  
  # Returns:
  #   the x that corresponds to this p in the critical bubble
  
  factor = sqrt(s/(h*s + c*(h*s - 1)))*sqrt(D/s)
  
  # alpha
  factor.alpha = 1/(3*sqrt(2*h - 1)*sqrt(s)*sqrt(h*s + c*(h*s - 1)))
  sqrt.term.alpha = sqrt(2 - 3*h -(c*(6*h - 5)*(h*s - 1))/s + (2*(c^2)*(h*s-1)^2)/(s^2))
  numer.alpha = s*sqrt(2) - 3*h*s*sqrt(2) - c*sqrt(2)*(h*s - 1) + s*sqrt.term.alpha
  alpha = factor.alpha*(numer.alpha)
  
  # beta 
  factor.beta = sqrt(s)/(3*sqrt(2)*sqrt(c*h*s + h*s - c))
  big.sqrt.term.beta = sqrt(18*h*(p - 1)^2 + 3*p*(4 - 3*p) - (6*c*(2*p - 3)*(h*s - 1))/s)
  beta = factor.beta*(-3*p*sqrt(2*h - 1) + big.sqrt.term.beta)
  
  x = factor*(log(1 - alpha) - log(1 + alpha) + log(1+beta) - log(1 - beta))
  
  return(x)
}

critical_bubble_vector_gene_drive = function(s, h, c, D, N = 100000, center_0 = F){
  # Arguments:
  #   s: selection coefficient
  #   h: dominance parameter
  #   c: conversion rate
  #   p: frequency; make sure it's between 0 and ptilda
  #   D: diffusion constant
  #   N: the grid length between 0 and 1; each step is 1/N
  #   center_0: if T, keep bubble centered at 0. If F, center at x = 0.5
  # Returns:
  #   List of x and p in the critical bubble (x is shifted to be centered at 0.5) and the plot
  
  dx = 1 / N
  p_dx = dx
  
  middle_index = ceiling(N/2)
  
  ptilda = ptilda_gene_drive(s, h, c) # value at the middle index
  
  if (N %% 2){
    # odd number of slices, so add ptilda at the middle
    pgrid = seq(p_dx, ptilda-p_dx, length.out = middle_index - 1)
    xgrid = x_bubble_gene_drive(s, h, c, pgrid, D)
    
    x_full = c(sort(-xgrid), 0,xgrid)
    p_full = c(pgrid, ptilda,pgrid)
  } else {
    pgrid = seq(p_dx, ptilda-p_dx, length.out = middle_index)
    xgrid = x_bubble_gene_drive(s, h, c, pgrid, D)
    
    x_full = c(sort(-xgrid), xgrid)
    p_full = c(pgrid, pgrid)
  }
  
  if (!center_0){
    data = tibble(x = x_full + 0.5, y = p_full) %>% arrange(x)  
  } else {
    data = tibble(x = x_full, y = p_full) %>% arrange(x)
  }
  
  x_full = data$x; p_full = data$y
  
  pghat = pg_hat_gene_drive(s,h,c)
  
  plot = ggplot(data, aes(x = x_full, y = p_full)) + geom_line(color = "red") +
    ylim(0,1) + xlab("x") + ylab("p") + geom_hline(yintercept = pghat, color = "black", linetype = "dashed")
  
  return(list(x = x_full, p = p_full, plot = plot, data = tibble(x = x_full, p = p_full)))
}

gene_drive_reaction = function(p, s,h,c,pg_hat, cubic_approximation = F){
  # The reaction equation for homing gene drive
  # Arguments:
  #   p: drive germline frequency
  #   s: fitness cost
  #   h: dominance
  #   c: conversion rate
  #   pg_hat: threshold frequency
  #   cubic_approximation: whether to use the cubic approximation
  # Returns:
  #   the expected change in frequency based on the reaction equation
  
  if (cubic_approximation){
    reaction = p*(p - 1)*(s*(h + p - 2*h*p) + c*(h*s - 1))
  } else {
    reaction = (p*(p - 1)*(s*(h + p - 2*h*p) + c*(h*s - 1)))/(1 + 2*p*h*s*(p - 1) - s*p^2)
  }
  return(reaction)
}


critical_bubble_auc_gene_drive = function(s, h, c, D, N = 100000){
  # Get critical width through numerical integration
  results = critical_bubble_vector_gene_drive(s = s, h = h, c = c, D=D, N=N)
  width = AUC(results$x, results$p) # find the AUC
  
  return(width)
}

############ 1D simulations ##################################

pde_function_gene_drive_1D = function(time, state, parms, N, D) {
  # Arguments:
  #     time: vector of timesteps 
  #     state: frequencies
  #     parms: list with s, h, and c, and cubic_approximation
  #     N: grid length between 0 and 1
  #     D: diffusion constant
  # Returns:
  #   dp/dt = D * d^2 p/ dx^2 + f(p)
  
  dx = 1/N
  with (as.list(parms), {
    # prevents floating point error from creating frequencies outside of 0,1
    state = pmax(state, 0)
    state = pmin(state, 1)
    
    P = state
    
    FluxP = -D * diff(c(P[1], P, P[N]))/dx # 0 flux at the walls
        
    # cubic reaction
    if (cubic_approximation){
      reaction = P*(P - 1)*(s*(h + P - 2*h*P) + c*(h*s - 1))
    } else {
      reaction = (P*(P - 1)*(s*(h + P - 2*h*P) + c*(h*s - 1)))/(1 + 2*P*h*s*(P - 1) - s*P^2)
    }
    
    ## Rate of change = -Flux gradient + Biology
    dP   = -diff(FluxP)/dx + reaction
    
    return (list(c(dP))) 
  })
}


simulate_homing_gene_drive_rde_1D = function(s, h, c, D, 
                          release_width=NULL, 
                          bubble = F,
                          p0 = 1,
                          factor = 1,
                          N = 50000, max_time = 100, 
                          cubic_approximation = T,
                          plot = T,
                          add_bubble_to_plot = F,
                          sleep_time = 0.4, return_out = F,
                          tol = 0){
  # Arguments:
  #   s: selection coefficient
  #   h: dominance parameter
  #   c: conversion rate
  #   D: diffusion constant
  #   release_width: square width or NULL if releasing at the bubble
  #   bubble: if T, start at the critical bubble (or a factor above/below). If F, start at a square release.
  #   p0: frequency in the square, if bubble = F,
  #   factor: applied to each point on the critical bubble. if above 1, then you're above the
  #           critical bubble. if below 1, then you're below the critical bubble.
  #   N: the grid length between 0 and 1; each step is 1/N
  #   max_time: number of ticks to simulate 
  #   cubic_approximation: whether to use the cubic reaction (more accurate) (T) or not (F)
  #   plot: whether to plot the deSolve simulation
  #   sleep_time: the time to wait between frames if plot=T
  #   plot: whether to plot the deSolve simulation
  #   add_bubble_to_plot: whether to add a static dashed line showing the critical bubble to all frames
  #   sleep_time: the time to wait between frames if plot=T
  #   return_out: whether to return the out ode.1D object
  #   tol: tolerance used in comparisons
  #   
  # Returns:
  #   if return_out: list with
  #       spread: T if the final overall frequency > the initial overall frequency (and more than 1e-5 apart)
  #       out: deSolve object; matrix with each row representing a timestep and columns for freq at each slice
  #   else just returns spread
  
  pars = c(s = s, h = h, c = c, cubic_approximation = cubic_approximation)
  dx = 1/N
  
  # Set up the release square 
  x_starts = seq(0, 1-dx, by = dx)
  x_ends = x_starts + dx
  x_midpoints = (x_starts + x_ends)/2
  
  yini = rep(0, N)
  
  if (!bubble){
    
    # Release from a square
    # Find out how much each grid slice overlaps with the release interval and set 
    # drive introduction frequencies accordingly
    prop_overlaps = get_overlap_proportions_1D(N = N, release_width = release_width)
    prop_overlaps[near(prop_overlaps,0)] = 0 # correct floating point errors
    prop_overlaps[near(prop_overlaps,1)] = 1
    inds_overlap = which(prop_overlaps > 0)
    yini[inds_overlap] = p0*prop_overlaps[inds_overlap]
  } else {
    res = critical_bubble_vector_gene_drive(s = s, h = h, c = c , D = D, N = N)
    actual_y = res$p
    actual_x = res$x 
    
    # Loop through the grid slice midpoint x values and find the res$x that's closest - use this frequency.
    for (i in 1:N){
      x_here = x_midpoints[i]
      if (x_here < min(actual_x) || x_here > max(actual_x)){
        yini[i] = 0
      } else {
        # find the closest value to this in the actual_x dataset
        ind = which.min(abs(actual_x - x_here))
        yini[i] = actual_y[ind]
      }
    }
    
    if (factor != 1){
      yini = yini*factor # increase or decrease below the critical bubble
    }
    
  }
  
  pd.0 = mean(yini)
  times = seq(0, max_time, by = 1)
  
  out = ode.1D(y = yini, times = times, func = pde_function_gene_drive_1D, parms = pars,
               N = N, D = D, nspec = 1)
  
  if (plot){
    
    if (add_bubble_to_plot){
      # add a dashed line showing the critical bubble
      res = critical_bubble_vector_gene_drive(s = s, h = h, c = c, D = D, N = N)
      actual_y = res$p
      actual_x = res$x
    } 
    
    for (i in 1:nrow(out)){
      time = i - 1
      freqs = out[i, 2:(N+1)]
      overall_freq = mean(freqs)
      plot(x = x_midpoints, y  = freqs, ylab = "p", ylim = c(0,1), xlab = "x",
           main = paste0("t = ",time,"\n","pd = ", round(overall_freq,5)), type = "l", col = "black", lwd = 2)
      
      if (add_bubble_to_plot){
        lines(x = actual_x, y = actual_y, lty = 2, lwd = 1.5, col = "red")  
      }
      
      
      Sys.sleep(sleep_time)
    }
  }
  
  pd.last = mean(out[nrow(out), 2:(N+1)])
  
  spread = (pd.last > (pd.0 + tol)) 
  
  if (return_out){
    return(list(spread = spread, out = out))
  }
  
  return(spread)
}

############ 2D simulations ##################################

pde_function_gene_drive_2D = function(time, state, parms, N, D) {
  # Arguments:
  #     time: vector of timesteps 
  #     state: a vector of length N*N, giving the drive frequency in each grid cell; this is converted into a N x N matrix
  #     parms: list with s, h, and c, and cubic_approximation
  #     N: grid length between 0 and 1 (for the x direction and the y direction; N x N cells total)
  #     D: diffusion constant
  # Returns:
  #   dp/dt = D * d^2 p/ dx^2 + f(p)
  
  dx = 1/N
  dy = 1/N
  NN = N*N

  with (as.list(parms), {
    
    # prevents floating point error from creating frequencies outside of 0,1
    state = pmax(state, 0)
    state = pmin(state, 1)
    
    P = matrix(nrow = N, ncol = N,state) 
    
    if (cubic_approximation){
      reaction = P*(P - 1)*(s*(h + P - 2*h*P) + c*(h*s - 1))
    } else {
      reaction = (P*(P - 1)*(s*(h + P - 2*h*P) + c*(h*s - 1)))/(1 + 2*P*h*s*(P - 1) - s*P^2)
    }
    
    zero = rep(0, N)
    
    ## 1. Fluxes in x-direction; zero flux at boundaries
    Flux_in_x = -D * rbind(zero,(P[2:N,] - P[1:(N-1),]), zero)/dx
    
    ## 2. Fluxes in y-direction; zero flux at boundaries
    Flux_in_y = -D * cbind(zero,(P[,2:N] - P[,1:(N-1)]), zero)/dx
    
    ## Add flux gradient to rate of change
    # dP/dt = f(p) - D d^2 p / dx^2 - D d^2 p/dy^2
    dP = reaction - (Flux_in_x[2:(N+1),] - Flux_in_x[1:N,])/dx - (Flux_in_y[,2:(N+1)] - Flux_in_y[,1:N])/dx
    
    return(list(c(as.vector(dP)))) # asVector will go by column
  })
}

simulate_homing_gene_drive_rde_2D = function(s, h, c, D, 
                                  release_diameter, 
                                  p0 = 1,
                                  N = 223, 
                                  max_time = 100,
                                  cubic_approximation = T,
                                  plot = T,
                                  sleep_time = 0.4, return_out = F,
                                  tol = 0){
  
  # Arguments:
  #     s: selection coefficient
  #     h: dominance parameter
  #     c: conversion rate
  #     D: diffusion constant
  #     release_diameter: diameter of the release shape
  #     p0: drive frequency in the circle (assuming homozygotes are released)
  #     N: the grid length in the x direction and in the y direction (N x N total grid cells)
  #     max_time: the number of ticks to simulate 
  #     cubic_approximation: whether to use Barton's cubic approximation in the reaction
  #     plot: whether to plot the deSolve simulation
  #     sleep_time: the time to wait between frames if plot=T
  #     return_out: whether to return the out ode.2D object
  #     tol: tolerance used in comparisons
  #
  # Returns:
  #   if return_out: list with
  #       spread: T if the final overall frequency > the initial overall frequency (and more than 1e-5 apart)
  #       out: deSolve object; matrix with each row representing a timestep and columns for freq at each slice
  #   else just returns spread
  
  pars = c(s = s, h = h, c = c, cubic_approximation = cubic_approximation)
  
  # Get overlap matrix for circular release
  prop_overlaps = get_overlap_proportions_2D(N = N, release_diameter = release_diameter)
  
  # Correct floating point errors
  prop_overlaps[near(prop_overlaps, 0)] = 0
  prop_overlaps[near(prop_overlaps, 1)] = 1
  
  yini = matrix(0, nrow = N, ncol = N)
  yini[prop_overlaps > 0] = p0 * prop_overlaps[prop_overlaps > 0]
  
  pd.0 = mean(yini)
  state = as.vector(yini)
  
  times = seq(0, max_time, by = 1)
  
  out = ode.2D(y = state, time = times, func = pde_function_gene_drive_2D,
               parms = pars, dimens = c(N,N), N = N, D=D, nspec = 1, method = rkMethod("rk45ck")) 
  
  if (plot){
    for (i in 1:nrow(out)){
      time = i - 1
      freqs = matrix(nrow = N, ncol = N, data = out[i, -1])
      
      overall_freq = mean(freqs)
      
      # for the contour plot, ensure values fall between 0 and 1 (could be slightly above/below due to floating point error)
      freqs.no.float.error = pmax(freqs, 0)
      freqs.no.float.error = pmin(freqs.no.float.error, 1)
      
      filled.contour(freqs.no.float.error, zlim = c(0,1), main =  paste0("t = ",time,"\n","pd = ", round(overall_freq,5)))
      
      Sys.sleep(sleep_time)
    }
  }
  
  pd.last = mean(matrix(nrow = N, ncol = N, data = out[nrow(out), -1]))
  
  spread = (pd.last > (pd.0 + tol)) 
  
  if (return_out){
    return(list(spread = spread, out = out))
  }
  
  return(spread)
  
}