Annotating slices of heatmaps with different color palettes in ComplexHeatmap in r

I want to split a heatmap into 3 slices and use different color palettes for each slice. For example, the matrix mat is split into 3 slices and I want the foo annotation for slice 1 to be assigned various shades of green based on foo value, slice 2 to shades of red, slice 3 to shade of yellow.

library(ComplexHeatmap)
set.seed(123)
mat = matrix(runif(100, 0, 10), 10)
rownames(mat) = paste0("R", 1:10)
colnames(mat) = paste0("C", 1:10)

library(circlize)
col_funa = colorRamp2(c(-3, 0, 3), c("pink", "white", "purple"))
col_funb = colorRamp2(c(-3, 0, 3), c("yellow", "white", "green"))

ha = HeatmapAnnotation(
foo = cbind(a = rnorm(10), b = rnorm(10)), 
foo2 = cbind(a = rnorm(10), b = rnorm(10)),
col = list(foo = col_funa, foo2 = col_funb)
)

Heatmap(mat, name = "mat", top_annotation = ha, 
    column_dend_height = unit(2, "cm"), 
    row_dend_width = unit(2, "cm"), 
    column_split = c(rep(1, 3),rep(2, 3),rep(3, 4)))

This creates the heatmap and then I tried to use decorate_annotation to adjust the color palette for foo annotation in different slices but the approach did not work.

col_fun1 = colorRamp2(c(-3, 0, 3), c("blue", "white", "red"))
col_fun2 = colorRamp2(c(-3, 0, 3), c("pink", "white", "purple"))
col_fun3 = colorRamp2(c(-3, 0, 3), c("yellow", "white", "green"))
col_funL <- list(col_fun1, col_fun2, col_fun3)

for(i in 1:3) {
    decorate_annotation("foo", slice = i, {
        grid.rect(gp = gpar(fill = col_funL[[i]], col = NA))
    })
}

I want to know how to use different color palettes for the foo annotation in each of the 3 slices of heatmap.

##################################################################### Followup 1: Zuguang's solutions worked well. However, when I adapt the code to my own data, which has thousands of rows and columns, I was notified that I should use layer_fun instead of cell_fun. I tried to adapt Zuguang's code as follows:

hta = Heatmap(rbind(value), rect_gp = gpar(type = "none"),
    layer_fun = function(j, i, x, y, w, h, fill) {
        if(pindex(rbind(foo), i, j) == "a") {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_a(pindex(rbind(value), i, j)), col = NA))
        } else if(pindex(rbind(foo), i, j) == "b") {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_b(pindex(rbind(value), i, j)), col = NA))
        } else {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_c(pindex(rbind(value), i, j)), col = NA))
        }
    }, show_heatmap_legend = FALSE)

But I was returned the following error:

Error in if (pindex(rbind(foo), i, j) == "a") { : 
  the condition has length > 1

Followup 2: I noted that using %v% not only vertically concatenated the annotation heatmap and the body heatmap, but also reordered the columns of the body heatmap. I understand this is necessary in general. However, what if I want to vertically align the annotation heatmap and the body heatmap without changing the column order of the body heatmap. I want to ask if it is possible to do so in ComplexHeatmap or through some other tricks.

Here the color mapping is actually two-dimensional where the first dimension is the discrete grouping of slice 1-3 and the second dimension corresponds to the continues numeric values. Such 2D color mapping is in my to do-list but has not been implemented yet. However, current I can provide the following two temporary solutions for this problem:

1. transform the original data into three distinct "blocks" so that we can define different colors for the three blocks of values separately. In the following example, I assume the data is from a uniform distribution so that making block is easier

set.seed(123)
m = matrix(rnorm(10*24), nrow = 10)
foo = sample(c("a", "b", "c"), 24, replace = TRUE)
value = runif(24)

We have three groups (a, b and c). To make values distinguishable of which groups they are in, we can add 1 to the values in "b" group and add 2 to the values in the "c" group so that values in (0, 1) belong to "a", (1, 2) belongs to "b" and (2, 3) belongs to "c".

value2 = value
value2[foo == "b"] = value2[foo == "b"] + 1
value2[foo == "c"] = value2[foo == "c"] + 2

And we can make a "complex" color mapping function to map green, red and yellow to the three intervals:

col_fun = colorRamp2(c(0, 0.9999, 1.00001, 1.99999, 2.00001, 3),
                     c("white", "green", "white", "red", "white", "yellow"))

And make the heatmap:

ht = Heatmap(m, column_split = foo,
    top_annotation = HeatmapAnnotation(foo = value2, col = list(foo = col_fun), show_legend = FALSE)
)

Note here we do not directly generate heatmap for this complex color mapping, but we can manually generate the three color mapping legend for the three groups.

col_fun_a = colorRamp2(c(0, 1), c("white", "green"))
col_fun_b = colorRamp2(c(0, 1), c("white", "red"))
col_fun_c = colorRamp2(c(0, 1), c("white", "yellow"))

lgd_list = list(
    Legend(title = "level a", col_fun = col_fun_a),
    Legend(title = "level b", col_fun = col_fun_b),
    Legend(title = "level c", col_fun = col_fun_c)
)

draw(ht, annotation_legend_list = lgd_list)

2. The second way is to treat the annotation as a heatmap, then you can use cel_fun to self-define the colors.

The following is the main heatmap. Here I explicitly set the name of this heatmap because the main heatmap is not the first one, and later we need to refer to it with its name.

ht = Heatmap(m, name = "main", column_split = foo)

Next we make the second heatmap which tends to be the top annotation of the first heatmap:

ht = Heatmap(rbind(value), rect_gp = gpar(type = "none"),
    cell_fun = function(j, i, x, y, w, h, fill) {
        if(foo[j] == "a") {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_a(value[j]), col = NA))
        } else if(foo[j] == "b") {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_b(value[j]), col = NA))
        } else {
            grid.rect(x, y, w, h, gp = gpar(fill = col_fun_c(value[j]), col = NA))
        }
    }, show_heatmap_legend = FALSE) %v% ht

So in cell_fun, we can test which group the current value is in (by testing the value of foo[j]), and choose a proper color mapping function.

Also note now we vertically concatenate the heatmaps, thus we use %v%.

And finally we draw the heatmap lists:

draw(ht, main_heatmap = "main", column_dend_side = "top", column_sub_title_side = "top",
    heatmap_legend_list = lgd_list)