diff --git a/DESCRIPTION b/DESCRIPTION
index 96265e0..41235de 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -20,6 +20,7 @@ Imports:
     fdrtool,
     ggplot2,
     ggrepel,
+    grDevices,
     grid,
     httr,
     limma,
diff --git a/NAMESPACE b/NAMESPACE
index 7b27d45..74ae4fc 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -54,7 +54,9 @@ importFrom(ComplexHeatmap,draw)
 importFrom(ComplexHeatmap,row_order)
 importFrom(ConsensusClusterPlus,ConsensusClusterPlus)
 importFrom(MSnbase,MSnSet)
+importFrom(RColorBrewer,brewer.pal)
 importFrom(RColorBrewer,brewer.pal.info)
+importFrom(RColorBrewer,colorRampPalette)
 importFrom(S4Vectors,"metadata<-")
 importFrom(S4Vectors,metadata)
 importFrom(SNFtool,SNF)
@@ -113,6 +115,7 @@ importFrom(ggplot2,scale_x_continuous)
 importFrom(ggplot2,theme)
 importFrom(ggplot2,theme_bw)
 importFrom(ggrepel,geom_text_repel)
+importFrom(grDevices,colorRampPalette)
 importFrom(grid,gpar)
 importFrom(grid,unit)
 importFrom(httr,GET)
diff --git a/R/FragPipeAnalystR-package.R b/R/FragPipeAnalystR-package.R
index 78f650d..0aad3b6 100644
--- a/R/FragPipeAnalystR-package.R
+++ b/R/FragPipeAnalystR-package.R
@@ -61,6 +61,7 @@
 #' @importFrom ggplot2 theme
 #' @importFrom ggplot2 theme_bw
 #' @importFrom ggrepel geom_text_repel
+#' @importFrom grDevices colorRampPalette
 #' @importFrom grid gpar
 #' @importFrom grid unit
 #' @importFrom httr GET
@@ -76,6 +77,7 @@
 #' @importFrom plotly layout
 #' @importFrom plotly plot_ly
 #' @importFrom purrr map_df
+#' @importFrom RColorBrewer brewer.pal
 #' @importFrom RColorBrewer brewer.pal.info
 #' @importFrom readr parse_factor
 #' @importFrom S4Vectors metadata
diff --git a/R/heatmap.R b/R/heatmap.R
index 281ed71..c87d754 100644
--- a/R/heatmap.R
+++ b/R/heatmap.R
@@ -41,7 +41,10 @@ plot_missval_heatmap <- function(se) {
 #' @export
 plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper = 1,
                                      pal = "PRGn", pal_rev = FALSE, indicate = NULL,
-                                     font_size = 12, plot = TRUE, exp = "LFQ", use="complete.obs",...) {
+                                     font_size = 12, exp = NULL, use="complete.obs",...) {
+  if (is.null(exp)) {
+    exp <- metadata(dep)$exp
+  }
   # Show error if inputs are not the required classes
   assertthat::assert_that(
     inherits(dep, "SummarizedExperiment"),
@@ -56,16 +59,14 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
     is.logical(pal_rev),
     length(pal_rev) == 1,
     is.numeric(font_size),
-    length(font_size) == 1,
-    is.logical(plot),
-    length(plot) == 1
+    length(font_size) == 1
   )
 
   # Check for valid lower and upper values
   if (!(lower >= -1 & upper >= -1 & lower <= 1 & upper <= 1)) {
     stop("'lower' and/or 'upper' arguments are not valid
          Run plot_pca() with 'lower' and 'upper' between -1 and 1",
-      call. = FALSE
+         call. = FALSE
     )
   }
 
@@ -75,10 +76,10 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
     filter(category != "qual")
   if (!pal %in% pals$rowname) {
     stop("'", pal, "' is not a valid color panel",
-      " (qualitative panels also not allowed)\n",
-      "Run plot_pca() with one of the following 'pal' options: ",
-      paste(pals$rowname, collapse = "', '"), "'",
-      call. = FALSE
+         " (qualitative panels also not allowed)\n",
+         "Run plot_pca() with one of the following 'pal' options: ",
+         paste(pals$rowname, collapse = "', '"), "'",
+         call. = FALSE
     )
   }
 
@@ -95,20 +96,20 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
     columns <- colnames(col_data)
     if (any(!indicate %in% columns)) {
       stop("'",
-        paste0(indicate, collapse = "' and/or '"),
-        "' column(s) is/are not present in ",
-        deparse(substitute(dep)),
-        ".\nValid columns are: '",
-        paste(columns, collapse = "', '"),
-        "'.",
-        call. = FALSE
+           paste0(indicate, collapse = "' and/or '"),
+           "' column(s) is/are not present in ",
+           deparse(substitute(dep)),
+           ".\nValid columns are: '",
+           paste(columns, collapse = "', '"),
+           "'.",
+           call. = FALSE
       )
     }
 
     # Get annotation
     anno <- colData(dep) %>%
       data.frame() %>%
-      select(indicate)
+      select(all_of(indicate))
 
     # Annotation color
     names <- colnames(anno)
@@ -118,17 +119,17 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
       var <- anno[[i]] %>%
         unique() %>%
         sort()
+      print(length(var))
       if (length(var) == 1) {
         cols <- c("black")
-      }
-      if (length(var) == 2) {
+      } else if (length(var) == 2) {
         cols <- c("orangered", "cornflowerblue")
-      }
-      if (length(var) < 7 & length(var) > 2) {
-        cols <- RColorBrewer::brewer.pal(length(var), "Pastel1")
-      }
-      if (length(var) >= 7) {
-        cols <- RColorBrewer::brewer.pal(length(var), "Set3")
+      } else if (length(var) < 7 & length(var) > 2) {
+        cols <- brewer.pal(length(var), "Pastel1")
+      } else if (length(var) <= 12) {
+        cols <- brewer.pal(length(var), "Set3")
+      } else {
+        cols <- colorRampPalette(brewer.pal(12, "Set3"))(length(var))
       }
       names(cols) <- var
       anno_col[[i]] <- cols
@@ -149,9 +150,9 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
     # Check for significant column
     if (!"significant" %in% colnames(rowData(dep, use.names = FALSE))) {
       stop("'significant' column is not present in '",
-        deparse(substitute(dep)),
-        "'\nRun add_rejections() to obtain the required column",
-        call. = FALSE
+           deparse(substitute(dep)),
+           "'\nRun add_rejections() to obtain the required column",
+           call. = FALSE
       )
     }
 
@@ -172,32 +173,28 @@ plot_correlation_heatmap <- function(dep, significant = FALSE, lower = -1, upper
 
   # Plot heatmap
   ht1 <- Heatmap(cor_mat,
-    col = circlize::colorRamp2(
-      seq(lower, upper, ((upper - lower) / 7)),
-      if (pal_rev) {
-        rev(RColorBrewer::brewer.pal(8, pal))
-      } else {
-        RColorBrewer::brewer.pal(8, pal)
-      }
-    ),
-    heatmap_legend_param = list(
-      color_bar = "continuous",
-      legend_direction = "horizontal",
-      legend_width = unit(5, "cm"),
-      title_position = "topcenter"
-    ),
-    name = "Pearson correlation",
-    column_names_gp = gpar(fontsize = font_size),
-    row_names_gp = gpar(fontsize = font_size),
-    top_annotation = ha1,
-    ...
+                 col = circlize::colorRamp2(
+                   seq(lower, upper, ((upper - lower) / 7)),
+                   if (pal_rev) {
+                     rev(RColorBrewer::brewer.pal(8, pal))
+                   } else {
+                     RColorBrewer::brewer.pal(8, pal)
+                   }
+                 ),
+                 heatmap_legend_param = list(
+                   color_bar = "continuous",
+                   legend_direction = "horizontal",
+                   legend_width = unit(5, "cm"),
+                   title_position = "topcenter"
+                 ),
+                 name = "Pearson correlation",
+                 column_names_gp = gpar(fontsize = font_size),
+                 row_names_gp = gpar(fontsize = font_size),
+                 top_annotation = ha1,
+                 ...
   )
-  if (plot) {
-    draw(ht1, heatmap_legend_side = "top")
-  } else {
-    df <- as.data.frame(cor_mat)
-    return(df)
-  }
+
+  return(ht1)
 }