R/score_mate.R

In danagibbon/MultifacitedChoice:

#' Score Mate
#'
#' Score the mating of a single female and male. \cr
#' Apply this function to multiple matings using the `rank_mates` function.
#'
#' @param DB RSQLite database generated by `make_database`
#' @param female Sample ID of the female to compare to a male. The Sample ID
#' must match the ones found in the `DB`.
#' @param male Sample ID of the male to compare to a female. The Sample ID
#' must match the ones found in the `DB`.
#' @param type The type of mating that is advantageous.
#' 1. `all_alleles`: the user specifies for each locus what is advantageous.
#' 2. `assortive`: Assortive (same alleles) is advantageous for all loci
#' 3. `disassortive`: Disassortive (different alleles) is advantageous for all loci
#' @param bonus Numeric value to increase mating score if the state of alleles
#' specified in `weighted_alleles` are met. If you specify .15 as your bonus,
#' and 2/3 of your loci match the advantageous, then .10 will be used as the bonus.
#' The final score will be increased by 10%. Must specify `weighted_alleles`.
#' @param weighted_alleles vector listing allele IDs to calculate the bonus. Must
#' specify `bonus`.
#' @return dataframe
#' @export
#' @examples
#' \dontrun{
#' single <- score_mate(DB = DB, female = female,
#'                      male = male, type = type,
#'                      bonus=bonus, weighted_alleles=weighted_alleles)
#' }
#' @export
#' @import dplyr
#' @import tibble
#' @import tidyr
#' @import tidyverse

score_mate <- function(DB, female, male, type,
                       bonus=NULL, weighted_alleles=NULL){
  # df
  df <- dbReadTable(DB, "original_df")
  try <- step.0(df,
                female = female,
                male = male)
  # get allele advantage table
  allele_info <- dbReadTable(DB, "allele_info")
  # get no calls
  no.calls <- get_nocalls(try,
                          female = female,
                          male = male)
  # get called
  called <- get_called(try,
                       female = female,
                       male = male)
  # both alleles
  both <- get_both(called,
                   female = female,
                   male = male)
  # none
  none <- get_none(called,
                   female = female,
                   male = male)
  # het
  hets <- get_hets(called,
                   female = female,
                   male = male)

  # Counts
  all_vars <- nrow(called)

  # check both alleles (in common)
  if(type == "all_alleles"){
    two_alleles <- merge(both, allele_info, by.x = "allele", by.y = "site_id") %>%
      filter(advantage == "assortive") %>%
      nrow()
  } else {
    two_alleles <- nrow(both)
  }

  # check no alleles (in common)
  if(type == "all_alleles"){
    zero_alleles <- nrow(none)
    if(zero_alleles > 0){
      zero_alleles_score <- merge(none, allele_info, by.x = "allele", by.y = "site_id") %>%
        filter(advantage == "disassortive") %>%
        nrow()
    } else if(zero_alleles == 0){
      zero_alleles_score <- 0
    }
  } else {
    zero_alleles <- nrow(none)
    zero_alleles_score <- 0
  }

  #one_alleles
  # Het Scores
  if(type == "all_alleles"){
    het_hom.t <- hets %>%
      filter(.[[2]] == .[[3]] | .[[4]] == .[[5]]) %>%
      nrow()
    if(het_hom.t > 0){
      # get het/hom combo
      het_hom.t.t <- hets %>%
        filter(.[[2]] == .[[3]] | .[[4]] == .[[5]])
      # get het/homs w/ a matching allele
      het_hom.t.a <- het_hom.t.t %>%
        filter(.[[2]] == .[[4]] | .[[2]] == .[[5]] | .[[3]] == .[[4]] | .[[3]] == .[[5]])
      if(nrow(het_hom.t.a) > 0){
        het_hom.t.a <- merge(het_hom.t.a, allele_info, by.x = "allele", by.y = "site_id") %>%
          filter(advantage == "assortive")
        het_hom.t.a.a <- het_hom.t.a$allele
        het_hom.t.a <- het_hom.t.a %>%
          nrow()
      } else if(nrow(het_hom.t.a) == 0){
        het_hom.t.a <- 0
        het_hom.t.a.a <- NULL
      }
      # get het/homs w/o a matching allele
      het_hom.t.d <- het_hom.t.t %>%
        filter(.[[2]] != .[[4]] & .[[2]] != .[[5]] & .[[3]] != .[[4]] & .[[3]] != .[[5]])
      if(nrow(het_hom.t.d) > 0){
        het_hom.t.d <- merge(het_hom.t.d, allele_info, by.x = "allele", by.y = "site_id") %>%
          filter(advantage == "disassortive")
        het_hom.t.d.a <- het_hom.t.d$allele
        het_hom.t.d <- het_hom.t.d %>%
          nrow()
      } else if(nrow(het_hom.t.d) == 0){
        het_hom.t.d <- 0
        het_hom.t.d.a <- NULL
      }
    } else if(het_hom.t == 0){
      het_hom <- 0
    }
    het_hom <- het_hom.t.a + het_hom.t.d
    # number of alleles not met by above
    het_hom.a <- c(het_hom.t.d.a, het_hom.t.a.a)
    het_hom.a <- which(!het_hom.a %in% het_hom.t.t$allele) %>% length()
  } else {
    het_hom <- hets %>%
      filter(.[[2]] == .[[3]] | .[[4]] == .[[5]]) %>%
      nrow()
    het_hom.a <- 0
  }
  ## het - het
  het_het <- hets %>%
    filter(.[[2]] != .[[3]] & .[[4]] != .[[5]]) %>%
    nrow()
  het_het <- het_het + het_hom.a

  if(is.null(bonus)){
    bonus <- 0
  } else {
    bonus <- calculate_bonus(bonus = bonus,
                             alleles = weighted_alleles,
                             both = both, none = none)
    #if ( type == "disassortive" ){
    #bonus <- bonus*(-1)
    #}
  }
  score_df <- data.frame(all_vars, two_alleles,
                         het_hom, het_het, zero_alleles,
                         zero_alleles_score, bonus)
  if ( type == "disassortive" ){
    score_df <- score_df %>%
      add_column(score = (((zero_alleles + .5*het_het + .6*het_hom)/all_vars) +
                            bonus*(zero_alleles + .5*het_het + .6*het_hom)/all_vars))
  } else {
    score_df <- score_df %>%
      add_column(score = (((two_alleles + .5*het_het + .6*het_hom + zero_alleles_score)/all_vars) +
                            bonus*(two_alleles + .5*het_het + .6*het_hom + zero_alleles_score)/all_vars))
  }
  # output
  out <- data.frame("female"=female, "male"=male,
                    "type"=type, "score"=score_df$score,
                    "bonus"=bonus,
                    "alleles_used"=((nrow(called)/(nrow(try)))*100))
  return(out)
}

# Helper functions for ranking

## Process df
step.0 <- function(df, female, male){
  out <- df %>%
    select(-Raw.Reads, -On.Target.Reads, -X.On.Target, -X.GT, -IFI) %>%
    filter(Sample == female |
             Sample == male) %>%
    t() %>% data.frame() %>%
    rename(!!female := X1, !!male := X2) %>%
    add_column(allele = row.names(.)) %>%
    slice(-1) %>%
    select(3, everything())
  return(out)
}

## Get no-calls
get_nocalls <- function(df, female, male){
  out <- df %>%
    filter(!!as.symbol(female) == "0" |
             !!as.symbol(male) == "0")
  return(out)
}

## Get called
get_called <- function(df, female, male){
  out <- df %>%
    filter(!!as.symbol(female) != "0" &
             !!as.symbol(male) != "0")
  return(out)
}

## Get SNPs with both Alleles the same
get_both <- function(df, female, male){
  out <- df %>%
    filter(!!as.symbol(female) == !!as.symbol(male))
  return(out)
}

## Get Hets
get_hets <- function(df, female, male){
  out <- df %>%
    filter(!!as.symbol(female) != !!as.symbol(male)) %>%
    separate(!!as.symbol(female),
             into = c(paste0(female, "-1"), paste0(female, "-2")),
             sep = 1) %>%
    separate(!!as.symbol(male),
             into = c(paste0(male, "-1"), paste0(male, "-2")),
             sep = 1) %>%
    filter(.[[2]] == .[[4]] | .[[2]] == .[[5]] |
             .[[3]] == .[[4]] | .[[3]] == .[[5]])
  return(out)
}

## Get None
get_none <- function(df, female, male){
  out <- df %>%
    filter(!!as.symbol(female) != !!as.symbol(male)) %>%
    separate(!!as.symbol(female),
             into = c(paste0(female, "-1"), paste0(female, "-2")),
             sep = 1) %>%
    separate(!!as.symbol(male),
             into = c(paste0(male, "-1"), paste0(male, "-2")),
             sep = 1) %>%
    filter(.[[2]] != .[[4]] & .[[2]] != .[[5]] &
             .[[3]] != .[[4]] & .[[3]] != .[[5]])
  return(out)
}

## Calculate bonus
calculate_bonus <- function(bonus, alleles, both, none){
  # not dealing with hets yet
  assort <- weighted_alleles %>%
    filter(type == "assortive") %>%
    filter(allele %in% both$allele)
  disassort <- weighted_alleles %>%
    filter(type == "disassortive")%>%
    filter(allele %in% none$allele)
  # calculate bonus
  new_bonus <- ((nrow(assort) + nrow(disassort))/nrow(weighted_alleles))*bonus
  return(new_bonus)
}