## Supplementary Data for Kazer et al., 
## "An Integrated Single-Cell Analysis of Multicellular Immune Dynamics during Hyper-Acute HIV-1 Infection."
##
## Generated on 01/17/20 by SWK

## Required packages for code
require(WGCNA)
require(flashClust)
require(Hmisc)
require(dplyr)

## Required packages for saving results
require(openxlsx)
require(ggplot2)
require(cowplot)

#######################################################
## Functions for Module Discovery Adapted from WGCNA ##
#######################################################

# Choosing the appropriate power for generating the adjacency matrix.
FindPower <- function(datExpr){
  #choose soft-threshold power
  powers = c(c(1:10), seq(from = 12, to=20, by=2))
  sft=pickSoftThreshold(datExpr,dataIsExpr = TRUE,powerVector = powers,corFnc = cor,
                        corOptions = list(use = 'p', method = "pearson"),networkType = "signed")
  
  # Plot the results
  par(mfrow = c(1,2));
  cex1 = 0.9;
  
  # Scale-free topology fit index as a function of the soft-thresholding power
  plot(sft$fitIndices[,1], -sign(sft$fitIndices[,3])*sft$fitIndices[,2],xlab="Soft Threshold (power)",ylab="Scale Free Topology Model Fit, signed R^2",type="n", main = paste("Scale independence"));
  text(sft$fitIndices[,1], -sign(sft$fitIndices[,3])*sft$fitIndices[,2],labels=powers,cex=cex1,col="red");
  
  # Red line corresponds to using an R^2 cut-off
  abline(h=0.80,col="red")
  
  # Mean connectivity as a function of the soft-thresholding power
  plot(sft$fitIndices[,1], sft$fitIndices[,5],xlab="Soft Threshold (power)",ylab="Mean Connectivity", type="n",main = paste("Mean connectivity"))
  text(sft$fitIndices[,1], sft$fitIndices[,5], labels=powers, cex=cex1,col="red")
}

# Generating the adjacency matrix and performing clustering
ClusterTOM <- function(datExpr, softPower){
  #dev.off()
  #Calclute the adjacency matrix
  adj= adjacency(datExpr,type = "signed", power = softPower);
  
  #Turn adjacency matrix into a topological overlap matrix (TOM) to minimize the effects of noise and spurious associations.
  TOM=TOMsimilarityFromExpr(datExpr,networkType = "signed", TOMType = "signed", power = softPower, corType="bicor");
  
  colnames(TOM) = rownames(TOM) = colnames(datExpr)
  dissTOM=1-TOM
  
  #Hierarchical clustering of the genes based on the TOM dissimilarity measure
  geneTree = flashClust(as.dist(dissTOM),method="complete");
  
  #Plot the resulting clustering tree (dendrogram)
  plot(geneTree, xlab="", sub="",cex=0.3);
  
  return(list(dissTOM = dissTOM, geneTree = geneTree)) #returns list with dissimilarity TOM, and the clustered gene tree.
}

# Cut the resulting clustering dendrogram using the "tree" method for cutreeDynamic. Minimum module size can be specified.
CutTOMTree <- function(datExpr, dissTOM, geneTree, minModuleSize = 10){
  #dev.off()
  # Module identification using dynamic tree cut, you can also choose the hybrid method
  dynamicMods = cutreeDynamic(dendro = geneTree,  method="tree", minClusterSize = minModuleSize);
  #dynamicMods = cutreeDynamic(dendro = geneTree, distM = dissTOM, method="hybrid", deepSplit = 2, pamRespectsDendro = FALSE, minClusterSize = minModuleSize);
  
  #Get the module labels and the size of each module. Lable 0 is reserved for unassigned genes
  print(table(dynamicMods))
  
  #Plot the module assignment under the dendrogram; note: The grey color is reserved for unassigned genes
  dynamicColors = labels2colors(dynamicMods)
  table(dynamicColors)
  plotDendroAndColors(geneTree, dynamicColors, "Dynamic Tree Cut", dendroLabels = FALSE, hang = 0.03, addGuide = TRUE, guideHang = 0.05, main = "Gene dendrogram and module colors")
  
  #Set the diagonal of the dissimilarity to NA 
  diag(dissTOM) = NA;
  
  #extract modules
  module_colors= setdiff(unique(dynamicColors), "grey")
  modules = lapply(module_colors, function(x){colnames(datExpr)[which(dynamicColors==x)]})
  names(modules) = module_colors
  return(list(dyanmicColors = dynamicColors, modules = modules)) #returns list with module colors, and the modules themselves
}

# Merge modules with low dissimilarity. Cutoff for dissimilarity merge can be specified
MergeSimilarModules <- function(datExpr, dynamicColors, geneTree, MEDissThres = 0.5){
  #cacluate eigengenes
  MEList = moduleEigengenes(datExpr, colors=dynamicColors)
  MEs = MEList$eigengenes
  
  # Calculate dissimilarity of module eigengenes
  MEDiss = 1-cor(MEs);
  
  # Cluster module eigengenes
  METree = hclust(as.dist(MEDiss), method = "average");
  
  # Plot the result
  #sizeGrWindow(7, 6)
  plot(METree, main = "Clustering of module eigengenes",
       xlab = "", sub = "")
  abline(h = MEDissThres, lwd=2, col="red")
  
  # Call an automatic merging function
  merge = mergeCloseModules(datExpr, dynamicColors, cutHeight = MEDissThres, verbose = 3)
  # The merged module colors
  mergedColors = merge$colors;
  # Eigengenes of the new merged modules:
  mergedMEs = merge$newMEs;
  #plot showing how merged modules exist
  plotDendroAndColors(geneTree, cbind(dynamicColors, mergedColors),
                      c("Dynamic Tree Cut", "Merged dynamic"),
                      dendroLabels = FALSE, hang = 0.03,
                      addGuide = TRUE, guideHang = 0.05)
  
  #extract merged modules
  merged_module_colors= setdiff(unique(mergedColors), "grey")
  merged_modules = lapply(merged_module_colors, function(x){colnames(datExpr)[which(mergedColors==x)]})
  names(merged_modules) = merged_module_colors
  
  return(list(mergedColors = mergedColors, merged_modules = merged_modules)) #returns list with merged colors, and the merged modules themselves
}

############################################################################################################
## Functions for assaying module significance against background and temporal variability in module score ##
############################################################################################################

## Test to determine if the genes within the module are truly the least dissimilar compared to randomly generated modules of the same size.
TestModuleSignificance <- function(mod, dissTOM, expr.data, n_perm = 10000, pval = 0.05, n.bin = 10){
  #vectorize the actual distribution of (dis)similarities, and remove zeros!
  true.diss = as.vector(dissTOM[mod,mod])
  true.diss = true.diss[-which(true.diss == 0)]
  
  #size of module for permutations
  mod.size = length(mod)
  
  #bin all genes by expression
  expr.avg = rowMeans(expr.data)
  expr.avg = expr.avg[order(expr.avg)]
  expr.avg.cut = as.numeric(x = cut2(x = expr.avg, m=round(length(expr.avg)/n.bin)))
  names(expr.avg.cut) = names(expr.avg)
  
  #create a table of binnings of all genes and of our module genes
  all.bin.table = table(expr.avg.cut)
  mod.bin.table = table(expr.avg.cut[mod])
  
  #randomly generate module with same expression binning structure and run t.test and record results
  test.results = data.frame(statistic = rep(NA, n_perm), pval = rep(NA, n_perm)) #container for results
  
  for (i in 1:n_perm){ #by permutation
    random.mod = list() #create an empty list we will fill with gene names (each element will be gene names by bin)
    
    #for each bin of the mod bin table, randomly select that number of genes from the full set with that expression
    for (j in names(mod.bin.table)){ 
      bin.genes = sample(names(expr.avg.cut)[which(expr.avg.cut == as.numeric(j))], mod.bin.table[j], replace = FALSE)
      random.mod[[as.numeric(j)]] = bin.genes #stick those genes into the random.mod list
    }
    #unlist and vectorize the distribution of (dis)similarities (remember to remove zeros)
    random.mod = unlist(random.mod)
    random.diss = as.vector(dissTOM[random.mod,random.mod])
    random.diss = random.diss[-which(random.diss == 0)]
    
    #perform a one-sided wilcox.test and record the statistic and p-value.
    #Note, IMPORTANT: here we perform the test asking if the true diss is LESS THAN the random diss, as we are trying to minimize dissimilarity
    test = wilcox.test(x = true.diss, y = random.diss, alternative = "less")
    test.results[i,] = c(test$statistic, test$p.value)
  }
  
  #correct for multiple hypothesis testing, and then report the proportion of bad tests
  test.results$FDR = p.adjust(test.results$pval, method = "fdr")
  num.failed.tests = sum(test.results$FDR > pval)
  print(paste(paste(num.failed.tests, n_perm, sep="/"), "permutations failed the Mann-Whitney test.", sep=" "))
  
  #is the percentage of failed tests less than or equal to the p-val specified?
  return(num.failed.tests/n_perm <= pval) #returns a vector of booleans indicating if each module was significant based on the specific p.val
}

## Test for variation in module score as a function of time. Compares many samplings of scores between pre-infection and each time point.
# Here, our time points are labeled as specified in tps; pre = pre-infection, 2001 = 0 Weeks, 2002 = 1 Week, ..., 2024 = 6 Months, 2024 = 1 Year
# meta.data here is the meta.data data.frame that exists within a Seurat object
TestModuleTemporalVariation <- function(tps = c("pre","2001","2002","2003","2004","2005","2024","2048"), #which time points to run the function on
                                        meta.data, sample.size = 50, ntest = 1000, name.of.feature = "M"){ #name.of.feature refers to prefix for columns with module scores
  
  #apply to run over one module at a time
  mod.min.pvals <- apply(meta.data[,grep(name.of.feature, colnames(meta.data), value=TRUE)], 2, function(mod){
    #get the scores for this cluster separated by time points in a list
    scores.by.tp = lapply(tps, function(time) mod[meta.data$TimePoint == time])
    
    #run the wilcox test ntest times with sample.size between each timepoint and pre-infection, report the average p-val
    wilcox.pval.by.tp = mapply(scores = scores.by.tp[-1], tp = tps[-1], SIMPLIFY=FALSE, function(scores,tp){
      test.pval = replicate(ntest, expr={
        wilcox.test(x=sample(scores.by.tp[[1]], sample.size),y=sample(scores, sample.size))$p.value})
      return(mean(test.pval))
    })
    #return the min p-value for all comparison within that cluster, FDR corrected
    min.pval = min(p.adjust(unlist(wilcox.pval.by.tp)))
    return(min.pval)
  })
  return(mod.min.pvals) #returns the vector of minimum average p-value across all time point comparisons for each module
}

## Test for variation in module score as a function of time across all participants. Null hypothesis = Participant only, Alternative hypothesis = Participant + TimePoing
# Here, we use an anova to test two linear regressions of the hypotheses noted above.
# In the manuscript, we ran this using a binned timepoint scheme. Participants must share time points for this test to make sense.
TestModuleTemporalVariation2 <- function(meta.data, sample.size = 25, ntest = 100, 
                                         score.feature.prefix = "Cluster", condition.feature = "TimePoint",
                                         batch.feature = "Patient", hist.bins = 10){
  
  #add group column (i.e. combination of Batch & TimePoint)
  meta.data$group = paste(meta.data[,batch.feature], meta.data[,condition.feature], sep = ".")
  
  #determine which groups don't have enough cells
  kept.groups = names(table(meta.data$group))[which(table(meta.data$group) >= sample.size)]
  excluded.groups = names(table(meta.data$group))[which(table(meta.data$group) < sample.size)]
  print(paste0("Excluded these groups: ", paste(excluded.groups, collapse = ", ")))
  
  #apply over each cluster scores at a time
  cluster.names = grep(score.feature.prefix, colnames(meta.data), value=TRUE)
  clust.min.pvals <- lapply(cluster.names, function(clust){
    
    #make data.frame of cluster score, patient, timepoint, and patient.timepoint
    score.df = data.frame(Score = meta.data[,clust], 
                          Batch = meta.data[,batch.feature], 
                          Condition = meta.data[,condition.feature],
                          Batch.Condition = paste(meta.data[,batch.feature], meta.data[,condition.feature], sep = "."))
    
    #set up the formula for the linear models below
    f0 = as.formula(paste("Score", "Batch", sep = " ~ 1 + ")) # null = participant (batch) only
    f1 = as.formula(paste("Score", paste("Condition", "Batch", sep = " + "), sep = " ~ 1 + ")) # alternative = particpant (batch) + time (condition)
    
    #run the anova comparing batch and batch+condition models, all time points, drawing same # of cells from every Patient.TimePoint
    test.pval = replicate(ntest, expr={ #run the following ntest times
      
      #sample cell names from kept.groups
      sampled.cells = unlist(lapply(kept.groups, function(group){
        sample(x = rownames(score.df)[which(score.df$Batch.Condition == group)], size = sample.size)
      }))
      
      #subset the score.df for sampled cells
      sampled.score.df = score.df[sampled.cells,]
      
      #Build null (only batch) and alternative model (batch + )
      m0 = lm(formula = f0, data = sampled.score.df)
      m1 = lm(formula = f1, data = sampled.score.df)
      
      #run anova and report the p-value
      test = anova(m0, m1, test = "F")
      return(test$`Pr(>F)`[2])
    })
    #histogram of results
    p = ggplot(data = data.frame(log.p = -log10(test.pval)), aes(log.p)) + geom_histogram(bins = hist.bins) + labs(title = clust)
    mean.pval = mean(test.pval)
    results = list(p, mean.pval)
    names(results) = c("hist", "mean.p")
    return(results)
  })
  #plot histograms in one plot
  print(cowplot::plot_grid(plotlist = lapply(clust.min.pvals, "[[", 1)))
  return(unlist(lapply(clust.min.pvals, "[[", 2)))
}


##########################################################################################################################################

###########################################
## Example Use Code from a Seurat Object ##
###########################################

## This code starts from a Seurat object of CD4 T cells (S.CD4) from one individual with 8 time points, annotated. Scaling, normalization, 
## variable gene selection, and PCA has already been calculated.

# Get the expression matrix for the top genes contributing to significant structure in PC-space
CD4.nPCs = 7 #chosen based on PCElbowPlot and inspection of the genes in each PC
S.CD4 = ProjectPCA(S.CD4, genes.print = 4, do.print = FALSE)
CD4.top.genes = unique(PCTopGenes(S.CD4, pc.use=1:CD4.nPCs, num.genes=50, do.balanced=TRUE, use.full=TRUE))
CD4.mat.top.genes = as.matrix(S.CD4@data[CD4.top.genes,])

# Look at the structure of the Scale Free Topology and Mean Connectivity as a function of Soft Threshold (Power)
disableWGCNAThreads() #run on a single core
FindPower(datExpr = t(CD4.mat.top.genes))

# choose softPower from above, here we choose softPower = 5. We recommend playing with the softPower choice
# to see how it affects the number and size of modules that are generated by the functions below.

# Caclulate the TOM, cluster the genes, and cut the tree based on desired module parameters
CD4.ClusterTOM = ClusterTOM(datExpr = t(CD4.mat.top.genes), softPower = 5)
CD4.ColorsModules = CutTOMTree(datExpr = t(CD4.mat.top.genes), dissTOM = CD4.ClusterTOM$dissTOM,
                               geneTree = CD4.ClusterTOM$geneTree, minModuleSize = 10)

# Merge modules with disimilarity < 0.5
CD4.MergedModules = MergeSimilarModules(datExpr = t(CD4.mat.top.genes), 
                                        dynamicColors = CD4.ColorsModules$dyanmicColors,
                                        geneTree = CD4.ClusterTOM$geneTree, MEDissThres = 0.5)

# Print the modules generated by WGCNA
CD4.MergedModules$merged_modules

# Test for module significance against randomly generated modules
CD4.Modules.isSig = sapply(CD4.MergedModules$merged_modules, function(module){
  TestModuleSignificance(mod = module, dissTOM = CD4.ClusterTOM$dissTOM, expr.data = CD4.mat.top.genes,
                         n_perm = 10000, pval = 0.05, n.bin = 10)
})
CD4.Sig.Modules = CD4.MergedModules$merged_modules[CD4.Modules.isSig] # keep those modules that are significant

# Score each cell for each module using Seurat's built in AddModuleScore
S.CD4 = AddModuleScore(S.CD4, genes.list = CD4.Sig.Modules, ctrl.size=5, enrich.name = "M")

#Visualize scores over time by box & whisker. "TimePoint" refers to a column in S.CD4 that annotates time point (S.CD4@meta.data$TimePoint)
CD4.gene.clusters.names = grep("Cluster", colnames(S.CD4@meta.data), value = TRUE)
CD4.geneclusts.time = lapply(CD4.gene.clusters.names, function(c_name){
  p = ggplot(S.CD4@meta.data, aes_string(x="TimePoint", y=c_name, fill="TimePoint")) +
    geom_boxplot() + scale_fill_manual(values = Timepoint.palette) #User defined palette for time points
  return(p)
})
plot_grid(plotlist = CD4.geneclusts.time)

# Look at how many cells are present at each time point to choose the appropriate sample size and time points 
table(S.102.CD4@meta.data$TimePoint)

# Determine the smallest pval for wilcox.test for every timepoint compared to reference (pre-infection) for each module.
# tps is a userdefined list of time points, whose values are present in S.CD4@meta.data $TimePoint. 
# The first time point in tps is used as the reference time point to which all others are compared
CD4.clust.wilcox.tests.minp = TestModuleTemporalVariation(meta.data = S.CD4@meta.data, sample.size = 150, ntest=1000, tps = tps)
CD4.clust.wilcox.tests.minp #print the results

# keep any modules with a minimum pval < 0.05
CD4.Sig.Var.Modules = CD4.Sig.Modules[CD4.clust.wilcox.tests.minp <= 0.05]

# remove the original module scores, and add scores for only those modules that were variant in time
S.CD4@meta.data = S.CD4@meta.data[,-grep("M", colnames(S.CD4@meta.data), ignore.case = FALSE)]
S.CD4 = AddModuleScore(S.CD4, genes.list = CD4.Sig.Var.Modules, ctrl.size=5, enrich.name = "M")
CD4.gene.clusters.names.final = grep("M", colnames(S.102.CD4@meta.data), ignore.case = FALSE, value = TRUE)

# rename the modules with the appropraite naming scheme (e.g. M1.CD4)
names(CD4.Sig.Var.Modules) = CD4.gene.clusters.names.final
#write.xlsx(CD4.Sig.Var.Modules, "CD4 Final Modules.xlsx") #save to file

# plot box & whiskers of final modules
CD4.geneclusts.time.final = lapply(CD4.gene.clusters.names.final, function(c_name){
  p = ggplot(S.102.CD4@meta.data, aes_string(x="TimePoint", y=c_name, fill="TimePoint")) +
    geom_boxplot() + scale_fill_manual(values = Timepoint.palette) + theme(legend.position = "none")
  return(p)
})
plot_grid(plotlist = CD4.geneclusts.time.final)
#ggsave("CD4 Final Module Scores Over Time.pdf", height = 8, width=12) #save to file