How to use multiple colour scales in ggplot2 in R

Given this matrix where on the x-axis are the latent values (1,2,3,4) and on the y-axis three levels for each drug (RAS, BB and AA). Each box of the matrix contains values between 0 and 1, I would like the boxes colored with an increasing colour gradient based on the value shown, with legends different for each drug.

Currently, for each row of each drug, the boxes are colored in the same color (RAS-ivory, BB-darkred, MRA-darkblue ). I would like to create three legends where each legend has a continuous scale of colors (from light to dark) and color the matrix according to its value.

#Dataframe
df <-data.frame(category=rep(c(0,1,2),12), 
                state=c(rep(1,3),rep(2,3),rep(3,3),rep(4,3)), 
                value=runif(36), 
                drug = c(rep("RAS",12),rep("BB",12),rep("AA",12)),
               cat_drug = paste(rep(c(0,1,2),12),sep="_",c(rep("RAS",12),rep("BB",12), rep("AA",12))))

#Different color for each drug
color_scale <- c("RAS" = "ivory", "BB" = "darkred", "AA" = "darkblue")

ggplot(Psi, aes(y = cat_drug, x = state, fill = drug, label = round(value, 3))) +
  geom_raster() +
  geom_text(color = "black") +
  scale_fill_manual(values = color_scale, name = "Drug") +
  guides(fill = guide_colourbar(title = "Drug")) +
  labs(y = "Level of adherence", x = "Latent state") +
  theme_bw() +
  theme(
    axis.text.x = element_text(size = 9, angle = 0, vjust = 0.3),
    axis.text.y = element_text(size = 9),
    plot.title = element_text(size = 11),
    legend.position = "right"
  )

One option to have multiple scales and legends for the same aesthetic is to use the ggnewscale package. However, doing so requires some additional effort in that we have to use multiple geom layers, i.e. one for each of your drugs and of course multiple scales. To avoid copy & paste I split the data by drug, then use purrr::imap to loop over the splitted data and add the layers. Unfortunately I ran in some issue with geom_raster to I switched to geom_tile which also to set the height of the tiles. Finally note that I use scale_fill_distiller and some of the Brewer color palettes.

library(ggplot2)
library(ggnewscale)

df_split <- split(df, df$drug)
pal_drug <- c("RAS" = "Blues", "BB" = "Greens", "AA" = "Oranges")

set.seed(123)

ggplot(df, aes(y = cat_drug, x = state, label = round(value, 3))) +
  purrr::imap(df_split, function(x, y) {
    list(
      geom_tile(data = x, aes(fill = value), height = 1),
      scale_fill_distiller(
        palette = pal_drug[[y]], name = y,
        direction = 1, limits = c(0, 1)
      ),
      ggnewscale::new_scale_fill()
    )
  }) +
  geom_text(color = "black") +
  labs(y = "Level of adherence", x = "Latent state") +
  theme_bw() +
  theme(
    axis.text.x = element_text(size = 9, angle = 0, vjust = 0.3),
    axis.text.y = element_text(size = 9),
    plot.title = element_text(size = 11),
    legend.position = "right"
  )

The ggh4x package offers one way to add multiple color (or fill) scales

library(ggh4x)
#> Loading required package: ggplot2
library(scales)

ggplot(dat, aes(y = drug, x = state, label = round(value, 3))) +
  geom_raster(aes(rasf = value),
              data = ~subset(dat, drug == 'RAS')
              ) +
  geom_raster(aes(bbf = value),
              data = ~subset(dat, drug == 'BB')
              ) +
  geom_raster(aes(aaf = value),
              data = ~subset(dat, drug == 'AA')
              ) +
  scale_fill_multi(
    aesthetics = c('rasf','bbf','aaf'),
    name = list('ivory', 'darkred', 'darkblue'),
    colors = list(
      brewer_pal(palette = 'Greys')(4),
      brewer_pal(palette = 'PuRd')(4),
      brewer_pal(palette = 'Blues')(4)
    ),
    guide = guide_colorbar(barheight = unit(60, 'pt'),
                           title = "Drug")) +
  geom_text(color = "black")  +
  labs(y = "Level of adherence", x = "Latent state") +
  theme_bw() +
  facet_wrap(vars(category), ncol = 1, dir = 'v', strip.position = 'left') +
  theme(
    axis.text.x = element_text(size = 9, angle = 0, vjust = 0.3),
    axis.text.y = element_text(size = 9),
    plot.title = element_text(size = 11),
    legend.position = "right",
    strip.placement = "outside"
  )
#> Warning in geom_raster(aes(rasf = value), data = ~subset(dat, drug == "RAS")):
#> Ignoring unknown aesthetics: rasf
#> Warning in geom_raster(aes(bbf = value), data = ~subset(dat, drug == "BB")):
#> Ignoring unknown aesthetics: bbf
#> Warning in geom_raster(aes(aaf = value), data = ~subset(dat, drug == "AA")):
#> Ignoring unknown aesthetics: aaf

<!-- -->

^{Created on 2023-04-19 with reprex v2.0.2}