R:mgcv add colorbar to 2D heatmap of GAM

It would be easier (IMHO) to do this reliably within the ggplot2 ecosphere.

I'll show a canned approach using my {gratia} package but also checkout {mgcViz}. I'll also suggest a more generic solution using tools from {gratia} to extra information about your model's smooths and then plot them yourself using ggplot().

library('mgcv')
library('gratia')
library('ggplot2')
library('dplyr')

# load your snippet of data via df <- structure( .... )

# then fit your model (note you have y as response & in the tensor product
# I assume z is the response below and x and y are coordinates
m <- gam(z ~ te(x, y), data=df, method='REML')

# now visualize the mode using {gratia}
draw(m)

This produces:

{gratia}'s draw() methods can't plot everything yet, but where it doesn't work you should still be able to evaluate the data you need using tools in {gratia}, which you can then plot with ggplot() itself by hand.

To get values for your smooths, i.e. the data behind the plots that plot.gam() or draw() display, use gratia::smooth_estimates()

# dist controls what we do with covariate combinations too far
# from support of the data. 0.1 matches mgcv:::plot.gam behaviour
sm <- smooth_estimates(m, dist = 0.1)

yielding

r$> sm                                                                          
# A tibble: 10,000 × 7
   smooth  type   by      est    se     x     y
   <chr>   <chr>  <chr> <dbl> <dbl> <dbl> <dbl>
 1 te(x,y) Tensor NA     35.3 11.5      2  94.4
 2 te(x,y) Tensor NA     35.5 11.0      2  94.6
 3 te(x,y) Tensor NA     35.7 10.6      2  94.9
 4 te(x,y) Tensor NA     35.9 10.3      2  95.1
 5 te(x,y) Tensor NA     36.2  9.87     2  95.4
 6 te(x,y) Tensor NA     36.4  9.49     2  95.6
 7 te(x,y) Tensor NA     36.6  9.13     2  95.9
 8 te(x,y) Tensor NA     36.8  8.78     2  96.1
 9 te(x,y) Tensor NA     37.0  8.45     2  96.4
10 te(x,y) Tensor NA     37.2  8.13     2  96.6
# … with 9,990 more rows

In the output, x and y are a grid of values over the range of both covariates (the number of points in the grid in each covariate is controlled by n such that the grid for a 2d tensor product smooth is of size n by n). est is the estimated value of the smooth at the values of the covariates and se its standard error. For models with multiple smooths, the smooth variable uses the internal label that {mgcv} gives each smooth - these are the labels used in the output you get from calling summary() on your GAM.

We can add a confidence interval if needed using add_confint().

Now you can plot your smooth(s) by hand using ggplot(). At this point you have two options

1. if draw() can handle the type of smooth you want to plot, you can use the draw() method for that object and then build upon it, or
2. plot everything by hand.

Option 1

# evaluate just the smooth you want to plot
smooth_estimates(m, smooth = "te(x,y)", dist = 0.1) %>%
draw() +
  geom_point(data = df, alpha = 0.2) # add a point layer for original data

This pretty much gets you what draw() produced when given the model object itself. And you can add to it as if it were a ggplot object (which is not the case of the objects returned by gratia:::draw.gam(), which is wrapped by {patchwork} and needs other ways to interact with the plots).

Option 2

Here you are in full control

sm <- smooth_estimates(m, smooth = "te(x,y)", dist = 0.1)
ggplot(sm, aes(x = x, y = y)) +
  geom_raster(aes(fill = est)) +
  geom_point(data = df, alpha = 0.2) + # add a point layer for original data
  scale_fill_viridis_c(option = "plasma")

which produces

A diverging palette is likely better for this, along the lines of the one gratia:::draw.smooth_estimates uses

sm <- smooth_estimates(m, smooth = "te(x,y)", dist = 0.1)
ggplot(sm, aes(x = x, y = y)) +
  geom_raster(aes(fill = est)) +
  geom_contour(aes(z = est), colour = "black") +
  geom_point(data = df, alpha = 0.2) + # add a point layer for original data
  scale_fill_distiller(palette = "RdBu", type = "div") +
  expand_limits(fill = c(-1,1) * abs(max(sm[["est"]])))

which produces

Finally, if {gratia} can't handle your model, I'd appreciate you filing a bug report here so that I can work on supporting as many model types as possible. But do try {mgcViz} as well for an alternative approach to visualsing GAMs fitted using {mgcv}.

A base plot solution would be to use fields::image.plot directly. Unfortunately, it require data in a classic wide format, not the long format needed by ggplot.

We can facilitate plotting by grabbing the object returned by plot.gam(), and then do a little manipulation of the object to get what we need for image.plot()

Following on from @Anke's answer then, instead of plotting with plot.gam() then using image.plot() to add the legend, we proceed to use plot.gam() to get what we need to plot, but do everything in image.plot()

plt <- plot(df.gam)
plt <- plt[[1]] # plot.gam returns a list of n elements, one per plot

# extract the `$fit` variable - this is est from smooth_estimates
fit <- plt$fit
# reshape fit (which is a 1 column matrix) to have dimension 40x40
dim(fit) <- c(40,40)
# plot with image.plot
image.plot(x = plt$x, y = plt$y, z = fit, col = heat.colors(999, rev = TRUE))
contour(x = plt$x, y = plt$y, z = fit, add = TRUE)
box()

This produces:

You could also use the fields::plot.surface() function

l <- list(x = plt$x, y = plt$y, z = fit)
plot.surface(l, type = "C", col = heat.colors(999, rev = TRUE))
box()

This produces:

See ?fields::plot.surface for other arguments to modify the contour plot etc.

As shown, these all have the correct range on the colour bar. It would appear that @Anke's version the colour bar mapping is off in all of the plots, but mostly just a little bit so it wasn't as noticeable.

Following Gavin Simpson's answer and this thread (How to add colorbar with perspective plot in R), I think I've come up with a solution that uses plot.gam() (though I really love that {gratia} takes it into a ggplot universe and will definitely look more into that)

require(fields)
df.gam<-gam(y~te(x,z), data=df, method='REML')
sm <- as.data.frame(smooth_estimates(df.gam, dist = 0.1))
plot(df.gam, scheme=2, hcolors=heat.colors(999, rev =T), contour.col='black', rug=F,  main='', cex.lab=1.75, cex.axis=1.75)
image.plot(legend.only=T, zlim=range(sm$est), col=heat.colors(999, rev =T), legend.shrink = 0.5, axis.args = list(at =c(-10,-5,0,5, 10, 15, 20)))

I hope I understood correctly that gratia:smooth_estimates() actually pulls out the partial effects.

For my model with multiple terms (and multiple tensor products), this seems to work nicely by indexing the sections of the respective terms in sm. Except for one, where the colorbar and the heatmap aren't quite matching up. I can't provide the actual underlaying data, but add that plot for illustration in case anyone has any idea. I'm using the same approach as outlined above. In the colorbar, dark red is at 15-20, but in the heatmap the isolines just above 0 already correspond with the dark red (while 0 is dark yellow'ish in the colorbar).