helper-functions.R

## This script includes helper functions for simulating
## the underdominance system, homing gene drive system,
## TADE modification system, or TADE suppression system

########## 1D RDE helpers ####################################################

get_overlap_proportions_1D = function(N, release_width) {
  # Arguments:
  #    N: the grid length between 0 and 1; each step is 1/N (match to NN distance)
  #    release_width: the release width
  # Returns:
  #   the fraction each grid slice overlaps with (0.5 +/- release_width/2)
  
  dx = 1 / N
  x_starts = seq(0, 1-dx, by = dx)
  x_ends = seq(dx, 1, by = dx)
  
  # release interval centered at 0.5
  release_start = 0.5 - release_width / 2
  release_end   = 0.5 + release_width / 2
  
  overlaps = numeric(N)
  
  for (i in 1:N) {
    slice_start = x_starts[i]
    slice_end   = x_ends[i]
    
    # Find the overlap interval
    overlap_start = max(slice_start, release_start)
    overlap_end   = min(slice_end, release_end)
    
    # If there's overlap, compute the length
    if (overlap_start < overlap_end) {
      overlaps[i] = (overlap_end - overlap_start) / dx  # proportion of the slice overlapped
    } else {
      overlaps[i] = 0
    }
  }
  
  return(overlaps)
}

tade_ode_extract_info_1D = function(out, row, N, suppression = F){
  # USED FOR SUMMARY STATISTICS
  # Arguments:
  #   out: deSolve output with (t+1) rows and (1 + 6*N) columns
  #   row: the row number (or time point + 1) to look at
  #   N: the number of grid cells
  #   suppression: whether it's TADE suppression (T) or modification (F)
  # Returns:
  #   list with the frequency at the middle, overall frequency, and area under the drive allele curve
  
  freqs = matrix((out[row, -1]), nrow = N, ncol = 6)
  
  if (!suppression){
    drive_frequencies = apply(freqs, 1, function(row) row[2]*0.5 + row[3] + 0.5 * row[4] + row[5] + row[6])
    middle_ind = ceiling(N/2)
    middle_freq = drive_frequencies[middle_ind]
    overall_freq = mean(drive_frequencies)
    auc = sum(drive_frequencies*(1/N))
    return(list(middle_freq = middle_freq, overall_freq = overall_freq, auc = auc))
  } else {
    
    total_densities = rowSums(freqs)
    drive_densities = apply(freqs, 1, function(row) row[2] + row[3]*2 + row[4] + 2*row[5] + 2*row[6])
    drive_frequencies = ifelse(near(total_densities, 0, tol = 0.01) | (total_densities < 0), 1, drive_densities/(2*total_densities))
    middle_ind = ceiling(N/2)
    middle_freq = drive_frequencies[middle_ind]
    
    overall_density = sum(total_densities)
    overall_drive_density = sum(drive_densities)
    overall_freq = overall_drive_density/(2*overall_density)
    
    # for each slice, multiply pd by 1/N
    auc = sum(drive_frequencies*(1/N))
    last_density = overall_density
    
    if (near(overall_density, 0)){
      last_density = 0
      overall_freq = 1
      middle_freq = 1
    }
    
    return(list(middle_freq = middle_freq, overall_freq = overall_freq, auc = auc, last_density = last_density))
  }
}

########## 2D RDE helpers ####################################################

get_overlap_proportions_2D = function(N, release_diameter, num_subcells = 20) {
  # Arguments:
  #   N: number of grid cells per axis (total grid is NxN)
  #   release_diameter: diameter of circular release (centered at 0.5, 0.5)
  # Returns:
  #   NxN matrix of proportions of each cell overlapped by the circular release
  #   rows are the y position and columns are the x position
  
  dx = 1 / N
  dy = dx  # Square grid
  
  radius = release_diameter / 2
  center_x = 0.5
  center_y = 0.5
  
  # Cell centers
  x_centers = seq(dx / 2, 1 - dx / 2, length.out = N)
  y_centers = x_centers
  
  overlap_matrix = matrix(0, nrow = N, ncol = N)
  
  # Approximate the % of the cell that overlaps with the circle
  # by creating subcells and counting the number that are within the circle
  sub_dx = dx / num_subcells

  for (i in 1:N) {
    x_start = x_centers[i] - dx/2
    for (j in 1:N) {
      y_start = y_centers[j] - dy / 2
      
      # Create subcell center points
      sub_x = seq(x_start + sub_dx / 2, x_start + dx - sub_dx / 2, length.out = num_subcells)
      sub_y = seq(y_start + sub_dx / 2, y_start + dy - sub_dx / 2, length.out = num_subcells)
      grid = expand.grid(x = sub_x, y = sub_y)
      
      # Check how many fall inside the circular release region
      distances = sqrt((grid$x - center_x)^2 + (grid$y - center_y)^2)
      overlap_matrix[j, i] = mean(distances <= radius) # rows (vertical) show y position and columns (horizontal) show x position
    }
  }
  
  return(overlap_matrix)
}


tade_ode_cell_stats_2D = function(out, row, N, suppression = F){
  # Finds genotype frequencies (or densities) and other stats in each grid cell in 2D
  #
  # Arguments:
  #   out: deSolve out.2D object with (max_time + 1) rows and (1 + 6*N*N) columns
  #   row: the row to examine (timestep + 1)
  #   out: deSolve out.2D object with (max_time +1) rows and (1 + 6*N*N) columns
  #   row: the row number (or time point + 1) to look at
  #   N: the length of the grid in the xdirection and ydirection (N x N grid cells total)
  #   suppression: whether it's TADE suppression (T) or modification (F)
  #
  # Returns:
  #   If (!suppression) aka modification:
  #       List with p1, p2, p3, p5, p6, p9 matrices showing genotype frequencies in each grid cell and
  #       pd, providing the overall drive allele frequency in each cell
  #   else aka suppression:
  #       List with N1, N2, N3, N5, N6, N9 matrices showing genotype densities in each grid cell;
  #       Ntot, providing the overall densities in each cell;
  #       Ndrive, providing the overall number of drive alleles in each cell; and
  #       pd, providing the drive frequency in each cell (note: if density = 0, pd = 1 as suppression is assumed to have occurred there)
  
  # This is used to find, in each grid cell, the genotype frequencies A
  NN = N*N
  state = out[row, -1]
  if (!suppression){
    p1_states = matrix(nrow = N, ncol = N, data = state[1:NN])
    p2_states = matrix(nrow = N, ncol = N, data = state[(NN + 1):(2 * NN)])
    p3_states = matrix(nrow = N, ncol = N, data = state[(2 * NN + 1):(3 * NN)])
    p5_states = matrix(nrow = N, ncol = N, data = state[(3 * NN + 1):(4 * NN)])
    p6_states = matrix(nrow = N, ncol = N, data = state[(4 * NN + 1):(5 * NN)])
    p9_states = matrix(nrow = N, ncol = N, data = state[(5 * NN + 1):(6 * NN)])
    
    p_drive = p2_states*0.5 + p3_states + p5_states*0.5 + p6_states + p9_states
    
    return(list(p1 = p1_states, p2 = p2_states, p3 = p3_states, p5 = p5_states,
                p6 = p6_states, p9 = p9_states, pd = p_drive))
  } else {
    N1_states = matrix(nrow = N, ncol = N, data = state[1:NN])
    N2_states = matrix(nrow = N, ncol = N, data = state[(NN + 1):(2 * NN)])
    N3_states = matrix(nrow = N, ncol = N, data = state[(2 * NN + 1):(3 * NN)])
    N5_states = matrix(nrow = N, ncol = N, data = state[(3 * NN + 1):(4 * NN)])
    N6_states = matrix(nrow = N, ncol = N, data = state[(4 * NN + 1):(5 * NN)])
    N9_states = matrix(nrow = N, ncol = N, data = state[(5 * NN + 1):(6 * NN)])
    
    Ntot = N1_states + N2_states + N3_states + N5_states + N6_states + N9_states
    Ndrive = N2_states + 2*N3_states + N5_states + 2*N6_states + 2*N9_states
    
    p_drive = ifelse(Ntot > 0, Ndrive/(2*Ntot), 1) # if Ntot = 0, drive reached 100% freq and suppressed there
    
    return(list(N1 = N1_states, N2 = N2_states, N3 = N3_states,
                N5 = N5_states, N6 = N6_states, N9 = N9_states,
                Ntot = Ntot,
                Ndrive = Ndrive, 
                pd = p_drive))
  }
}


tade_ode_extract_info_2D = function(out, row, N, suppression = F){
  # USED FOR SUMMARY STATISTICS
  #
  # Arguments:
  #   out: deSolve out.2D object with (max_time +1) rows and (1 + 6*N*N) columns
  #   row: the row number (or time point + 1) to look at
  #   N: the length of the grid in the xdirection and ydirection (N x N grid cells total)
  #   suppression: whether it's TADE suppression (T) or modification (F)
  #
  # Returns:
  #   If (!suppression) aka modification:
  #       List with drive frequency at the middle of the landscape and the overall drive frequency
  #   else (if suppression):
  #       List with drive frequency at the middle of the landscape, the overall drive frequency, and overall population size
  
  cell_stats = tade_ode_cell_stats_2D(out = out, row = row, N = N, suppression = suppression)
  middle_ind = ceiling(N/2)
  middle_freq = cell_stats$pd[middle_ind, middle_ind]
  
  if (!suppression){
    overall_freq = mean(cell_stats$pd) # density assumed to be constant
    return(list(middle_freq = middle_freq, overall_freq = overall_freq))
    
  } else {
    overall_density = sum(cell_stats$Ntot)
    overall_num_genomes = 2*overall_density
    overall_num_drive_alleles = sum(cell_stats$Ndrive)
    overall_freq = overall_num_drive_alleles/overall_num_genomes
  
    return(list(middle_freq = middle_freq, overall_freq = overall_freq, overall_density = overall_density))
    
  }
  
}