R/dendrogramAdjustmentFunctions.R

In milescsmith/WGCNA: Weighted Correlation Network Analysis

## This file contains several functions which can be used to adjust the dendrogram
#   in ways which keep the dendrogram mathematically identical (ie, branch swapping,
#   branch reflection, etc).  The goal is to biologically optimize the dendrogram.

# ----------------------------------------------------------------------------- #

orderBranchesUsingHubGenes <- function(hierTOM, datExpr = NULL, colorh = NULL, type = "signed",
                                       adj = NULL, iter = NULL, useReflections = FALSE, allowNonoptimalSwaps = FALSE) {

  # First read in and format all of the variables
  hGenes <- hierTOM$labels
  if (is.null(adj)) {
    genes <- chooseOneHubInEachModule(datExpr, colorh, type = type)
    adj <- adjacency(datExpr[, genes], type = type, power = (2 - (type == "unsigned")))
    colnames(adj) <- rownames(adj) <- genes
  }
  genes <- rownames(adj)
  if (length(genes) != length(intersect(genes, hGenes))) {
    write("All genes in the adjacency must also be in the gene tree.  Check to make sure", "")
    write("that names(hierTOM$labels) is set to the proper gene or probe names and that", "")
    write("these correspond to the expression / adjacency gene names.", "")
    return(0)
  }
  genes <- hGenes[is.element(hGenes, genes)]
  adj <- adj[genes, genes]
  if (is.null(iter)) iter <- length(genes)^2
  iters <- (1:iter) / iter
  swapAnyway <- rep(0, length(iters)) # Quickly decreasing chance of random swap
  if (allowNonoptimalSwaps) swapAnyway <- ((1 - iters)^3) / 3 + 0.001

  # Iterate random swaps in the branch, only accepting the new result
  #  if it produces a higher correlation than the old result OR if the
  #  random variable says to swap (which gets less likely each iteration)
  changes <- NULL
  for (i in 1:iter) {
    swap <- 1
    if (useReflections) swap <- sample(0:1, 1)
    gInd <- sample(1:length(genes), 2)
    g <- genes[gInd]
    if (swap == 1) {
      hierTOMnew <- swapTwoBranches(hierTOM, g[1], g[2])
    } else {
      hierTOMnew <- reflectBranch(hierTOM, g[1], g[2], TRUE)
    }
    oldSum <- .offDiagonalMatrixSum(adj)
    oGenesNew <- hGenes[hierTOMnew$order]
    oGenesNew <- oGenesNew[oGenesNew %in% genes]
    adjNew <- adj[oGenesNew, oGenesNew]
    newSum <- .offDiagonalMatrixSum(adjNew)
    if ((newSum > oldSum) | ((sample(1:1000, 1) / 1000) < swapAnyway[i])) {
      hierTOM <- hierTOMnew
      changes <- rbind(changes, c(i, ifelse(swap == 1, "Swap", "Reflect"), g, oldSum, newSum))
      adj <- adjNew
    }
    write(paste("Interation", i, "of", iter), "")
    collectGarbage()
  }

  # Perform all of the suggested swappings on the input network.

  # Output the results
  colnames(changes) <- c("Iter.#", "Swap?", "Gene1", "Gene2", "OldScore", "NewScore")
  out <- list(geneTree = hierTOM, changeLog = changes)
  return(out)
}

# ----------------------------------------------------------------------------- #

selectBranch <- function(hierTOM, g1, g2) {
  ## This function selects of all genes in a branch given a gene in the
  ##  branch (g1) and a gene in a neighboring branch (g2), returning the
  ##  indices for genes in the branch in the hierTOM$labels vector

  # Convert genes to UNORDERED indices (if given, indices should be ordered)
  if (is.numeric(g1)) g1 <- hierTOM$order[g1]
  if (is.numeric(g2)) g2 <- hierTOM$order[g2]
  if (!is.numeric(g1)) g1 <- which(hierTOM$labels == g1)
  if (!is.numeric(g2)) g2 <- which(hierTOM$labels == g2)
  if ((length(g1) == 0) | (length(g2) == 0) | (max(c(g1, g2)) > length(hierTOM$labels))) {
    write("Input genes are not both legal indices", "")
    return(hierTOM)
  }

  # Now determine which branch is the correct one, and find the genes
  len <- length(hierTOM$height)
  tree1 <- which(hierTOM$merge == (-g1)) %% len
  continue <- length(which(hierTOM$merge == tree1)) > 0
  while (continue) {
    nextInd <- which(hierTOM$merge == tree1[length(tree1)]) %% len
    tree1 <- c(tree1, nextInd)
    continue <- length(which(hierTOM$merge == nextInd)) > 0
  }

  branchIndex <- which(hierTOM$height == .minTreeHeight(hierTOM, g1, g2))
  branch <- hierTOM$merge[branchIndex, ]
  b1 <- NULL
  if (is.element(branch[1], tree1)) {
    b1 <- .getBranchMembers(hierTOM, branch[1], b1)
  } else {
    b1 <- .getBranchMembers(hierTOM, branch[2], b1)
  }
  collectGarbage()
  return(b1)
}

# ----------------------------------------------------------------------------- #

reflectBranch <- function(hierTOM, g1, g2, both = FALSE) {
  ## This function reverses the ordering of all genes in a branch of the
  ##  clustering tree defined by the minimal branch possible that contains
  ##  both g1 and g2 (as either ORDERED index or gene names), or just by
  ##  the genes in g1

  b1 <- selectBranch(hierTOM, g1, g2)
  if (both) b1 <- c(b1, selectBranch(hierTOM, g2, g1))

  # Now reorder the hierTOM correctly
  ord <- hierTOM$order
  i1 <- which(ord %in% b1)
  b <- 1:(min(i1) - 1)
  if (b[length(b)] < b[1]) b <- NULL
  e <- (max(i1) + 1):length(ord)
  if ((max(i1) + 1) > length(ord)) e <- NULL
  ord <- ord[c(b, i1[order(i1, decreasing = T)], e)]
  hierTOM$order <- ord
  return(hierTOM)
}

# ----------------------------------------------------------------------------- #

swapTwoBranches <- function(hierTOM, g1, g2) {
  ## This function re-arranges two branches in a heirarchical clustering tree
  ##  at the nearest branch point of two given genes (or indices)

  # Convert genes to indices (ORDERED AS ON THE PLOT)
  if (is.numeric(g1)) g1 <- hierTOM$order[g1]
  if (is.numeric(g2)) g2 <- hierTOM$order[g2]
  if (!is.numeric(g1)) g1 <- which(hierTOM$labels == g1)
  if (!is.numeric(g2)) g2 <- which(hierTOM$labels == g2)
  if ((length(g1) == 0) | (length(g2) == 0) | (max(c(g1, g2)) > length(hierTOM$labels))) {
    write("Input genes are not both legal indices", "")
    return(hierTOM)
  }

  # Now determine the genes in each branch
  branchIndex <- which(hierTOM$height == .minTreeHeight(hierTOM, g1, g2))
  b1 <- b2 <- NULL
  b1 <- .getBranchMembers(hierTOM, hierTOM$merge[branchIndex, 1], b1)
  b2 <- .getBranchMembers(hierTOM, hierTOM$merge[branchIndex, 2], b2)

  # Now reorder the hierTOM correctly
  ord <- hierTOM$order
  i1 <- which(ord %in% b1)
  i2 <- which(ord %in% b2)
  if (min(i1) > min(i2)) {
    tmp <- i1
    i1 <- i2
    i2 <- tmp
    rm(tmp)
  }
  b <- 1:(min(i1) - 1)
  if (b[length(b)] < b[1]) b <- NULL
  e <- (max(i2) + 1):length(ord)
  if ((max(i2) + 1) > length(ord)) e <- NULL
  ord <- ord[c(b, i2, i1, e)]
  hierTOM$order <- ord
  return(hierTOM)
}

# ----------------------------------------------------------------------------- #

chooseOneHubInEachModule <- function(datExpr, colorh, numGenes = 100,
                                     omitColors = "grey", power = 2, type = "signed", ...) {
  ## This function returns the gene in each module with the highest connectivity, given
  #   a number of randomly selected genes to test.

  numGenes <- max(round(numGenes), 2)
  keep <- NULL
  isIndex <- FALSE
  modules <- names(table(colorh))
  numCols <- table(colorh)
  if (!(is.na(omitColors)[1])) modules <- modules[!is.element(modules, omitColors)]
  if (is.null(colnames(datExpr))) {
    colnames(datExpr) <- 1:dim(datExpr)[2]
    isIndex <- TRUE
  }

  for (m in modules) {
    num <- min(numGenes, numCols[m])
    inMod <- which(is.element(colorh, m))
    keep <- c(keep, sample(inMod, num))
  }
  colorh <- colorh[keep]
  datExpr <- datExpr[, keep]
  return(chooseTopHubInEachModule(datExpr, colorh, omitColors, power, type, ...))
}

# ----------------------------------------------------------------------------- #

chooseTopHubInEachModule <- function(datExpr, colorh, omitColors = "grey",
                                     power = 2, type = "signed", ...) {
  ## This function returns the gene in each module with the highest connectivity.

  isIndex <- FALSE
  modules <- names(table(colorh))
  if (!(is.na(omitColors)[1])) modules <- modules[!is.element(modules, omitColors)]
  if (is.null(colnames(datExpr))) {
    colnames(datExpr) <- 1:dim(datExpr)[2]
    isIndex <- TRUE
  }

  hubs <- rep(NA, length(modules))
  names(hubs) <- modules
  for (m in modules) {
    adj <- adjacency(datExpr[, colorh == m], power = power, type = type, ...)
    hub <- which.max(rowSums(adj))
    hubs[m] <- colnames(adj)[hub]
  }
  if (isIndex) {
    hubs <- as.numeric(hubs)
    names(hubs) <- modules
  }
  return(hubs)
}

#################################################################################
# Internal functions.............................................................

options(expressions = 50000) # Required for .getBranchMembers

.getBranchMembers <- function(hierTOM, ind, members) {
  # This is a recursive function that gets all the indices of members of
  #  a branch in an hClust tree.
  if (ind < 0) {
    return(c(members, -ind))
  }
  m1 <- hierTOM$merge[ind, 1]
  m2 <- hierTOM$merge[ind, 2]
  if (m1 > 0) {
    members <- .getBranchMembers(hierTOM, m1, members)
  } else {
    members <- c(members, -m1)
  }
  if (m2 > 0) {
    members <- .getBranchMembers(hierTOM, m2, members)
  } else {
    members <- c(members, -m2)
  }
  return(members)
}

# ----------------------------------------------------------------------------- #

.minTreeHeight <- function(hierTOM, l1, l2) {
  ## This function finds the minimum height at which two leafs
  ##  in a hierarchical clustering tree are connected.  l1 and
  ##  l2 are the UNORDERED indices for the two leafs.

  ## Return 2 (larger than 1, if l1 or l2 is negative). This represents
  ##  positions that are off the edge of the tree.
  if ((l1 < 0) | (l2 < 0)) {
    return(2)
  }

  ## Get the tree for l1
  len <- length(hierTOM$height)
  tree1 <- which(hierTOM$merge == (-l1)) %% len
  continue <- length(which(hierTOM$merge == tree1)) > 0
  while (continue) {
    nextInd <- which(hierTOM$merge == tree1[length(tree1)]) %% len
    tree1 <- c(tree1, nextInd)
    continue <- length(which(hierTOM$merge == nextInd)) > 0
  }

  ## Get the tree for l2
  tree2 <- which(hierTOM$merge == (-l2)) %% len
  continue <- length(which(hierTOM$merge == tree2)) > 0
  while (continue) {
    nextInd <- which(hierTOM$merge == tree2[length(tree2)]) %% len
    tree2 <- c(tree2, nextInd)
    continue <- length(which(hierTOM$merge == nextInd)) > 0
  }

  ## Now find the index where the two trees first agree
  minTreeLen <- min(c(length(tree1), length(tree2)))
  tree1 <- tree1[(length(tree1) - minTreeLen + 1):length(tree1)]
  tree2 <- tree2[(length(tree2) - minTreeLen + 1):length(tree2)]
  treeInd <- tree1[min(which(tree1 == tree2))]

  ## Now find and return the minimum tree height
  return(hierTOM$height[ifelse(treeInd == 0, len, treeInd)])
}

# ----------------------------------------------------------------------------- #

.offDiagonalMatrixSum <- function(adj) {
  len <- dim(adj)[1]
  output <- sum(diag(adj[1:(len - 1), 2:len]))
  return(output)
}