Split LINCS Cell Painting Pilot 1 data into train / test

library(magrittr)
library(tidyverse)

1 Read

metadata <- 
  read_tsv("https://raw.githubusercontent.com/broadinstitute/lincs-cell-painting/master/metadata/moa/repurposing_info_external_moa_map_resolved.tsv") 

metadata %<>% select(broad_sample, moa)

2 Tidy

metadata %<>%
  group_by(broad_sample) %>%
  nest() %>%
  mutate(x = map(
    data,
    function(y)
      data.frame(moa = str_split(y$moa, "\\|")[[1]])
  )) %>% 
  select(-data) %>%
  unnest(x) %>%
  distinct() %>%
  na.omit()

3 Report count

Number of MOAs (n_moa) with the specified number of compounds (n_compound)

compound_count_per_moa <- 
  metadata %>% 
  group_by(moa) %>% 
  tally(name = "n_compound") %>% 
  ungroup() %>% 
  group_by(n_compound) %>% 
  tally(name = "n_moa")

compound_count_per_moa

ggplot(compound_count_per_moa, aes(n_compound, n_moa)) + 
  geom_col() +
  ggtitle("Num. of MOAs (n_moa) with the specified num. of compounds (n_compound)")

Number of compounds (n_compound) with the specified number of MOA labels (n_moa)

moa_count_per_compound <- 
  metadata %>% 
  group_by(broad_sample) %>% 
  tally(name = "n_moa") %>% 
  ungroup() %>%
  group_by(n_moa) %>%
  tally(name = "n_compound")

moa_count_per_compound

ggplot(moa_count_per_compound, aes(n_moa, n_compound)) + 
  geom_col() +
  ggtitle("Num. of compounds (n_compound) with the specified num. of MOA labels (n_moa)")

4 Assign

metadata %<>% 
  mutate(train = FALSE, 
         test = FALSE)

metadata %<>%
  mutate(marked = train | test)

metadata_init <- metadata

Get list of MOA sorted in ascending order of class size

moas <- 
  metadata %>% 
  group_by(moa) %>% 
  tally(name = "n_compound") %>%
  arrange(n_compound) %>% 
  extract2("moa")

Assign compounds to train / test ensuring that the assignments are consistent across MOAs

# loop over moa's
for (moa_i in moas) {
  
  moa_current <- metadata %>% filter(moa == moa_i)
  
  # number of compounds in this MOA
  n_compounds <- nrow(moa_current)
  
  # number of train, test to be picked
  
  if (n_compounds == 1) {
    # for singleton MOAs, assign to train
    
    n_test <- 0
    n_train <- 1
    
  } else {
    n_test <- ceiling(n_compounds * 0.2)
    n_train <- floor(n_compounds * 0.8)
  }
  
  # number that are already assigned as train or test
  n_test_marked <- sum(moa_current$test)
  n_train_marked <- sum(moa_current$train)
  
  # number that need to be assigned as train or test
  n_train_needed <- n_train - n_train_marked
  n_test_needed <- n_test - n_test_marked
  
  # split compounds into those that have already been previously marked vs not
  moa_current_notmarked <- moa_current %>% filter(!marked)
  moa_current_marked    <- moa_current %>% filter( marked)
  
  # number of compounds in this MOA that are not marked
  n_compounds_notmarked <- nrow(moa_current_notmarked)
  
  n_compounds_needed <- n_test_needed + n_train_needed
  
  # if you hit this, we have a problem and would need to randomly shuffle the moa list 
  # and start over
  stopifnot(n_compounds_needed == n_compounds_notmarked)
  
  n_test_needed <- max(n_test_needed, 0)
  
  n_train_needed <- max(n_train_needed, 0)
  
  train_flag <- c(rep(TRUE, n_train_needed), rep(FALSE, n_test_needed))
  
  train_flag <- sample(train_flag, n_compounds_needed)
  
  test_flag <- 
    !train_flag
  
  moa_current_notmarked$train <- train_flag
  
  moa_current_notmarked$test <- test_flag
  
  moa_current_notmarked$marked <- TRUE
  
  moa_current <-
    bind_rows(
      moa_current_notmarked,
      moa_current_marked
    )
  
  moa_other <- metadata %>% filter(moa != moa_i)
  
  moa_other_marked    <- metadata %>% filter((moa != moa_i) &  (marked))
  
  moa_other_notmarked <- metadata %>% filter((moa != moa_i) & (!marked))
  
  moa_other_notmarked %<>% 
    select(broad_sample, moa) %>% 
    left_join(
      moa_current_notmarked %>% select(-moa), 
      by = "broad_sample") %>%
    replace_na(list(train = FALSE, test = FALSE, marked = FALSE))
  
  metadata <- 
    bind_rows(
      moa_current,
      moa_other_marked,
      moa_other_notmarked
    )
  
  #message(sprintf("%s:%d", moa_i, nrow(metadata)))
}

metadata %>% 
  group_by(moa) %>%
  summarise(n_train = sum(train), n_test = sum(test)) %>% 
  mutate(n = n_train + n_test) %>%
  arrange(desc(n))

Make sure that a compounds is either in train or in test (we need to check this because a compound can occur in more than one MOA)

consistent <- function(x) all(x) || all(!x)

# should have zero rows if consistent
metadata %>% 
  group_by(broad_sample) %>%
  summarise(train = consistent(train),
            test = consistent(test)) %>%
  filter(!train | !test)

metadata %>% 
  group_by(broad_sample) %>%
  summarise(train = all(train),
            test = all(test)) %>%
  ungroup() %>%
  summarize(train = sum(train),
            test = sum(test)) %>%
  mutate(total = train + test)

metadata %>% 
  group_by(moa) %>%
  summarize(train = sum(train),
            test = sum(test)) %>%
  mutate(test_ratio = test/(train + test)) %>% 
  filter(test_ratio == 0 & train > 1) %>%
  arrange(desc(train))

metadata %>%
  ungroup() %>%
  distinct(broad_sample) %>%
  count()

metadata %>% write_csv("~/Desktop/lincs_split.csv")