utils.R

define.cellular.fractions <- function(df) {
  vars <- c("image_id", "tile_x", "tile_y", "region")
  vars <- vars[vars %in% colnames(df)]
  has.count <- "count" %in% colnames(df)
  has.area <- "area" %in% colnames(df)
  has.tissue.area <- "tissue_area" %in% colnames(df)
  ret <- ddply(df, .variables = vars,
               .fun = function(sdf) {
                 if(has.count) {
                   sdf$frac.cell.count <- sdf$count / sum(sdf$count)
                 }
                 if(has.area) { 
                   sdf$frac.cellular.area <- sdf$area / sum(sdf$area)
                   if(has.tissue.area) {
                     sdf$frac.regional.area <- sdf$area / sdf$tissue_area
                   }
                 }
                 return(sdf)
               })
  return(ret)  
}

widen.hovernet.results <- function(df, metrics = c("count", "area", "frac.cell.count", "frac.cellular.area", "frac.regional.area")) {
  # metrics <- c("count", "area", "frac.cell.count", "frac.cellular.area", "frac.regional.area")
  metrics <- metrics[metrics %in% colnames(df)]
  names(metrics) <- metrics
  lhs.vars <- c("image_id", "tile_x", "tile_y", "region")
  lhs.vars <- lhs.vars[lhs.vars %in% colnames(df)]
  form <- as.formula(paste0(paste0(lhs.vars, collapse = " + "), " ~ cell_type_label"))
  dfs <- llply(metrics,
               .fun = function(col) {
                 tmp <- dcast(df, form, value.var = col, fill = 0)
                 old.value.cols <- colnames(tmp)[!(colnames(tmp) %in% c(lhs.vars, "cell_type_label"))]
                 new.value.cols <- paste0(old.value.cols, "_", col)
                 setnames(tmp, old.value.cols, new.value.cols)
                 return(tmp)
               })
  ret <- Reduce(function(...) merge(..., all=T), dfs)
  count.cols <- colnames(ret)[grepl(colnames(ret), pattern="_count")]
  ret[,"total_cell_count"] <- unlist(apply(ret[, count.cols], 1, sum))
  area.cols <- colnames(ret)[grepl(colnames(ret), pattern="_area")]
  if(length(area.cols) > 0) {
    ret[,"total_cellular_area"] <- unlist(apply(ret[, area.cols], 1, sum))
  }
  for(met in metrics) {
    ret[, paste0("stroma_", met)] <- ret[, paste0("connec_", met)] + ret[, paste0("inflam_", met)]
  }
  cell.types <- c("neopla", "stroma", "necros", "inflam", "connec")
  combined <- as.vector(outer(cell.types, metrics, FUN = function(x,y) paste0(x,"_",y)))
  combined <- c("total_cell_count", "total_cellular_area", "region", combined)
  combined <- combined[combined %in% colnames(ret)]
  ret <- ret[do.call(order, ret[,combined]), ]
  ret <- ret[, c(lhs.vars, combined)]
  return(ret)
}

# Partition tiles into training and test set at the patient level. Try to balance
# the number of tiles within each region across test and training. To this 
# using binary quadratic programming. 
# NB: df is assumed to come from a single dataset.
# df is assumed to have the following columns:
# region
# image_id
# num.pt.tiles -- i.e., the # of tiles of that region assigned to that patient (i.e., image_id)
# num.dataset.tiles -- i.e., the # of tiles of that region across the entire dataset
partition.patients.into.training.and.test <- function(df, test.fraction = 0.5) {
  num.pts <- length(unique(df$image_id))
  n <- num.pts
  # Take half the patients for the test set
  test.size <- floor(n * test.fraction)
  test.size <- max(1, test.size)
  test.size <- min(n-1, test.size)
  # Define n_i^r: the # of tiles for region r in patient i
  n_ir <- acast(df[, c("region", "image_id", "num.pt.tiles")], image_id ~ region, fill = 0)
  # Try to optimize such that test.size patients are selected such
  # that approximately test.size / n of each region's tiles (across all patients)
  # are included in the test set. Call n^r: number of tiles we would like for region r
  # i.e., n^r = test.fraction * sum_i n_i^r
  # n_r <- 0.5 * as.vector(colSums(n_ir))
  # target.pt.frac <- test.fraction
  target.pt.frac <- test.size / n

  n_r <- floor(target.pt.frac * as.vector(colSums(n_ir)))
  print(c(test.size, target.pt.frac, n_r, colSums(n_ir) - n_r))
    
  # Formulate this as binary optimization problem where
  # x_i is a binary variable indicating whether patient i is included in the test set
  # We want to optimize
  # \sum_r ( \sum_i x_i n_i^r - n^r )^2
  # = \sum_r ( \sum_i,j x_i x_j n_i^r n_j^r - 2 n^r \sum_i x_i n_i^r + n^r^2 )
  # = \sum_r ( \sum_i,j>i 2 x_i x_j n_i^r n_j^r + \sum_i x_i x_i n_i^r^2 - 2 n^r \sum_i x_i n_i^r + n^r^2 )
  # \propto \sum_r ( \sum_i,j>i 2 x_i x_j n_i^r n_j^r + \sum_i x_i x_i n_i^r^2 - 2 n^r \sum_i x_i n_i^r )
  # Noting that for binary variables x_i = x_i x_i
  # = \sum_r ( \sum_i,j>i 2 x_i x_j n_i^r n_j^r + \sum_i x_i n_i^r^2 - 2 n^r \sum_i x_i n_i^r )
  # Collecting terms and re-arranging sums ..
  # = \sum_r [ \sum_i x_i n_i^r (n_i^r - 2 n^r ) + \sum_i,j>i 2 x_i x_j n_i^r n_j^r ]
  # And re-arranging r and i sums
  # = \sum_i x_i \sum_r n_i^r (n_i^r - 2 n^r ) + \sum_i,j>i 2 x_i x_j \sum_r n_i^r n_j^r ]
  # = \sum_i x_i c_i + \sum_i,j>i x_i x_j C_ij (*)
  # with
  # c_i = \sum_r n_i^r (n_i^r - 2 n^r)
  # C_ij = 2 \sum_r n_i^r n_j^r
  
  # Now note that the optimization (*) is in the same form as eqn (3) of
  # https://www.hindawi.com/journals/jam/2020/5974820/
  # Use the auxiliary variables
  # w_ij = x_i x_j
  # and the standard linearization defined there to transform this quadratic
  # problem into a linear one.
  # Namely, introduce the constraints of equations (4) - (7)
  # w_ij <= x_i             for all i < j
  # w_ij <= x_j             for all i < j
  # w_ij >= x_i + x_j - 1   for all i < j
  # w_ij >= 0               for all i < j
  
  # The optimization then becomes
  # = \sum_i x_i c_i + \sum_i,j>i w_ij C_ij
  
  # Finally, define our variables z_i s.t.
  # z_k = x_i     k = i; i <= n (n = num.pts)
  # z_k = w_ij    where the mapping between k and i,j is provided by
  # trans = which(upper.tri(C, diag = FALSE), arr.ind=T)
  # and the k = n + rij where rij is the row index of trans and i and j and the values of that row
  # i.e., i = trans[rij, 1]; j = trans[rij, 2]
  
  # Define constants c and C in optimization 
  c = rep(0,n)
  C = matrix(rep(0,n*n),nrow=n)
  trans = as.data.frame(which(upper.tri(C, diag = FALSE), arr.ind=T))
  
  # c_i = \sum_r n_i^r (n_i^r - 2 n^r)
  for(i in 1:n) {
    c[i] <- 0
    for(r in 1:ncol(n_ir)) {
      tmp <- n_ir[i,r] * (n_ir[i,r] - 2 * n_r[r])
      c[i] <- c[i] + tmp
    }
  }
  # C_ij = 2 \sum_r n_i^r n_j^r
  for(i in 1:(n-1)) {
    for(j in (i+1):n) {
      C[i,j] <- 0
      for(r in 1:ncol(n_ir)) {
        tmp <- 2 * n_ir[i,r] * n_ir[j,r]
        C[i,j] <- C[i,j] + tmp
      }
    }
  }
  
  # Define the objective function by its coefficients, one per variable
  n.x.vars <- n
  n.aux.vars <- nrow(trans) # auxiliary w_ij variables
  n.tot.vars <- n.x.vars + n.aux.vars
  f.obj <- rep(0, n.tot.vars)
  for(i in 1:n.x.vars) {
    f.obj[i] <- c[i]
  }
  for(aux.idx in 1:nrow(trans)) {
    i <- trans[aux.idx,1]
    j <- trans[aux.idx,2]
    f.obj[n.x.vars + aux.idx] <- C[i,j]
  }
  
  # Define the constraint matrix, constraint direction, and constraint right hand side, 
  # with the matrix having one row per constraint and one column per variable
  # There is only summation constraint over the original x_i variables
  # and 4 constraints for each of the auxiliary variables
  n.constraints <- 1 + 4 * n.aux.vars
  add.min.num.region.tiles <- FALSE
  if(add.min.num.region.tiles) {
    n.constraints <- n.constraints + 2 * length(n_r)
  }
  f.con <- matrix(rep(0,n.constraints * n.tot.vars), nrow=n.constraints)
  f.dir <- rep(0, n.constraints)
  f.rhs <- rep(0, n.constraints)
  
  nxt.constraint <- 1
  # Add the summation constraint over the original variables first
  for(i in 1:n.x.vars) {
    f.con[nxt.constraint,i] <- 1
  }
  f.dir[nxt.constraint] <- "=="
  f.rhs[nxt.constraint] <- test.size
  
  # Add a constraint that we have at least half of the _desired_ fraction of tiles of a given
  # region in the test set and at most twice.
  # \sum_i x_i * n_ir >= n_r  for all r
  if(add.min.num.region.tiles) {
    for(r in 1:ncol(n_ir)) {
      nxt.constraint <- nxt.constraint + 1
      f.con[nxt.constraint, 1 : n.x.vars ] <- n_ir[ 1 : n.x.vars, r]
      f.dir[nxt.constraint] <- ">="
      f.rhs[nxt.constraint] <- (1.0 / 2.0) * n_r[r]
      nxt.constraint <- nxt.constraint + 1
      f.con[nxt.constraint, 1 : n.x.vars ] <- n_ir[ 1 : n.x.vars, r]
      f.dir[nxt.constraint] <- "<="
      f.rhs[nxt.constraint] <- (2.0) * n_r[r]
    }
  }
  
  # Now add the four constraints for each of the auxiliary variables
  for(aux.idx in 1:nrow(trans)) {
    i <- trans[aux.idx,1]
    j <- trans[aux.idx,2]
    
    # w_ij <= x_i             for all i < j
    # -> w_ij - x_i <= 0
    nxt.constraint <- nxt.constraint + 1
    f.con[nxt.constraint, n.x.vars + aux.idx] <- 1
    f.con[nxt.constraint, i] <- - 1
    f.dir[nxt.constraint] <- "<="
    f.rhs[nxt.constraint] <- 0
    
    # w_ij <= x_j             for all i < j
    # -> w_ij - x_j <= 0
    nxt.constraint <- nxt.constraint + 1
    f.con[nxt.constraint, n.x.vars + aux.idx] <- 1
    f.con[nxt.constraint, j] <- - 1
    f.dir[nxt.constraint] <- "<="
    f.rhs[nxt.constraint] <- 0
    
    # w_ij >= x_i + x_j - 1   for all i < j
    # -> w_ij - x_i - x_j >= -1
    nxt.constraint <- nxt.constraint + 1
    f.con[nxt.constraint, n.x.vars + aux.idx] <- 1
    f.con[nxt.constraint, i] <- - 1
    f.con[nxt.constraint, j] <- - 1
    f.dir[nxt.constraint] <- ">="
    f.rhs[nxt.constraint] <- -1
    
    # w_ij >= 0               for all i < j
    nxt.constraint <- nxt.constraint + 1
    f.con[nxt.constraint, n.x.vars + aux.idx] <- 1
    f.dir[nxt.constraint] <- ">="
    f.rhs[nxt.constraint] <- 0
  }
  
  # All variables are binary
  binary.vec <- 1:n.tot.vars
  lp.res <- lp("max", f.obj, f.con, f.dir, f.rhs, binary.vec = binary.vec)
  trans$wij <-paste0("w_", trans[,1], ",", trans[,2])
  soln <- data.frame(sol = lp.res$solution)
  soln$varname <- ""
  soln[1:n.x.vars, "varname"] <- paste0("x_", 1:n.x.vars)
  soln[(n.x.vars+1):n.tot.vars, "varname"] <- trans$wij
  soln$patient <- ""
  soln[1:n.x.vars, "patient"] <- rownames(n_ir)
  
  test.df <- subset(df, image_id %in% subset(soln, sol == 1)$patient)
  train.df <- subset(df, image_id %in% subset(soln, sol == 0)$patient)
  
  ret <- list("soln" = soln, "test.df" = test.df, "train.df" = train.df)
  return(ret)
}


# Partition tiles into training and test set at the patient level. Try to balance
# the number of tiles within each region across test and training. To this 
# using binary quadratic programming and independently across each dataset.
# df is assumed to have the following columns:
# dataset
# region
# image_id
# num.pt.tiles -- i.e., the # of tiles of that region in that dataset assigned to that patient (i.e., image_id)
# num.dataset.tiles -- i.e., the # of tiles of that region across the entire dataset
partition.patients.into.training.and.test.across.groups <- function(df, groups = c("dataset", "Diagnosis"), test.fraction = 0.5) {
  test.train.dfs <- 
    dlply(df, .variables = groups,
          .fun = function(df.dataset) { 
            partition.patients.into.training.and.test(df.dataset, test.fraction = test.fraction)
          })
  test.df <-
    ldply(test.train.dfs, .fun = function(l) l[["test.df"]])
  colnames(test.df)[1] <- "dataset"
  
  train.df <-
    ldply(test.train.dfs, .fun = function(l) l[["train.df"]])
  colnames(train.df)[1] <- "dataset"
  
  return(list("train.df" = train.df, "test.df" = test.df))
}

p_load(viridis)
plot.metrics <- function(tbl, title) {
  
  metrics <- c("frac.cell.count")
  cell.types <- c("neopla", "necros", "inflam", "connec")
  combined <- as.vector(outer(cell.types, metrics, FUN = function(x,y) paste0(x,"_",y)))
  
  tbl <- tbl[do.call(order, tbl[,c("region", "dataset", combined)]), ]
  tbl$index <- 1:nrow(tbl)
  
  gs <- llply(combined, .fun = function(var) {
    g <- ggplot(data=tbl, aes_string(x="index",y=var)) + geom_point()
    g <- g + theme(axis.title=element_blank(), axis.text=element_blank(), axis.ticks=element_blank())
    g <- g + ggtitle(var)
    g
  })
  g.region <- ggplot(data=tbl, aes(x=index,colour=region,fill=region,y=1)) + geom_tile() + scale_fill_viridis(discrete = TRUE) + scale_colour_viridis(discrete = TRUE) + theme(legend.position="bottom") 
  g.region <- g.region + theme(axis.title=element_blank(), axis.text=element_blank(), axis.ticks=element_blank())
  g.dataset <- ggplot(data=tbl, aes(x=index,colour=dataset,fill=dataset,y=1)) + geom_tile() + theme(legend.position="bottom") + scale_fill_viridis(discrete = TRUE) + scale_colour_viridis(discrete = TRUE)
  g.dataset <- g.dataset + theme(axis.title=element_blank(), axis.text=element_blank(), axis.ticks=element_blank())
  g.tot <- do.call(grid.arrange, c(c(gs, list(g.region, g.dataset)), "ncol"=1))
  title <- ggdraw() + draw_label(title)
  g.final <- grid.arrange(title, g.tot, heights=c(0.1,1))
  return(g.final)
}

load.all.metadata <- function(metadata.files, sites = NULL) {
  nms <- names(metadata.files)
  names(nms) <- nms
  metadata <-
    llply(nms,
          .fun = function(nm) {
            metadata.file <- metadata.files[[nm]]
            df <- openxlsx::read.xlsx(metadata.file, sheet = 1, startRow = 3)
            if(!is.null(sites)) {
              df$Site <- sites[[nm]]
            }
            df
          })
  
  all.metadata <- ldply(metadata)
  colnames(all.metadata)[1] <- "dataset"
  
  all.metadata
}

define.folds <- function(tiles_data, fold_file, n.folds = 5, seed = 4321, allow.duplicates.across.patients = FALSE, stratify.col = "Diagnosis") {
  if(!file.exists(fold_file)) {
    cat("Defining folds and weights")
    un <- unique(tiles_data[, c("Diagnosis", "image_id", "Patient.ID", "Site")])
    if(allow.duplicates.across.patients) {
      un <- unique(tiles_data[, c("Diagnosis", "Patient.ID", "Site")])
    } else {
      # At this point, we assume that only one sample has been selected per patient.
    }
    stopifnot(!any(duplicated(un$Patient.ID)))
    folds <- 
      ddply(un, .variables = c("Site"),
            .fun = function(df) {
              set.seed(seed)
              df.folds <- createFolds(factor(df[, stratify.col]),k = n.folds)
              ldply(df.folds, .fun = function(fold) df[fold,])
            })
    colnames(folds)[1] <- "fold"
    # Don't merge folds with data -- i.e., folds will just have Diagnosis, Patient.ID, Site, and fold
    #cols <- c("minx", "tile_x", "maxx", "miny", "tile_y", "maxy", "tile", "image_id", "Diagnosis", "Patient.ID", "Site", "dataset")
    #cols <- cols[cols %in% colnames(tiles_data)]
    #tmp <- tiles_data[, cols]
    #folds <- merge(tmp, folds)
    
    write.table(file = fold_file, folds, sep=",", row.names=FALSE, col.names=TRUE, quote = FALSE)
  } else {
    cat("Fold and weight file already saved; retrieving it\n")
  }
  folds <- read.table(fold_file, sep=",", header=TRUE)
  return(folds)
}

# Define weights equally across class, across diagnoses within class, across site within class and diagnosis, 
# across patients within class, diagnosis, and site, across images within class, diagnosis, site, and patient,
# and across tiles within class, diagnosis, site, patient, and image
define.weights <- function(df, col.hierarchy = c("status", "Diagnosis", "Site", "Patient.ID", "image_id")) {
  
  # Define the number of tiles at the deepest level of the hierarchy,  e.g., within image_id
  ret <-
    ddply(df, .variables = col.hierarchy, 
          .fun = function(df.sub) { data.frame(num.tiles = nrow(df.sub))})
  df <- merge(ret, df)

  col.names <- paste0("num.", col.hierarchy)
  
  # Count the number of instances at each level of the hierarchy, e.g., the number of images
  # within each patient within each site within each diagnosis within each status
  for(indx in length(col.names):2) {
    col.name <- col.names[indx]
    ret <-
      ddply(unique(df[, col.hierarchy[1:indx], drop=FALSE]), .variables = col.hierarchy[1:(indx-1)], 
          .fun = function(df.sub) { 
            ret.df <- data.frame(foo = nrow(df.sub))
            colnames(ret.df)[1] <- col.name
            ret.df
          })
    df <- merge(ret, df)
  }
  
  df[col.names[1]] <- length(unique(df[, col.hierarchy[1]]))
  df$num.diagnoses <- length(unique(df$Diagnosis))
  df$total.weight <- 1 / df$num.tiles
  for(col.name in col.names) {
    df$total.weight <- df$total.weight * 1 / df[, col.name]
  }

  eps <- 0.001

  # Ensure total weight is the same for each class  
  tot <- ddply(df,
               .variables = col.hierarchy[1],
               .fun = function(df.sub) {
                 data.frame(tot = sum(df.sub$total.weight))
               })
  if(!similar.values(tot$tot, tot[1,"tot"])) { stop(paste0("Got dissimilar weights across ", col.hierarchy[1], "\n"))}
  
  # Now, ensure that total weight is the same for each subtree within a given level of the hierarchy
  # e.g., total weight should be the same for each patient within a site within a diagnosis within a class
  for(indx in 2:length(col.names)) {
    d_ply(df, .variables = col.hierarchy[1:(indx-1)],
        .fun = function(df.outer) {
          tot <- ddply(df.outer,
                       .variables = col.hierarchy[indx],
                       .fun = function(df.sub) {
                         data.frame(tot = sum(df.sub$total.weight))                         
                       })
          if(!similar.values(tot$tot, tot[1,"tot"])) { stop(paste0("Got dissimilar weights across ", col.hierarchy[indx], "\n")) }
        })
  }
  
  # Handle the bottom level of the hierarchy
  d_ply(df,
        .variables = col.hierarchy,
        .fun = function(df.sub) {
          if(!similar.values(df.sub$total.weight, df.sub[1,"total.weight"])) { stop(paste0("Got dissimilar weights across ", "tiles", "\n")) }
        })
  
  return(df)
}