VizRobinson.R

#'Plot map in Robinson or other projections
#'
#'Transform a regular grid longitude-latitude data to a different projection and 
#'plot the map. The target projection must be a valid CRS string, preferrably be 
#'EPSG or ESRI code; check \link[sf]{st_crs} for more explanation. This function
#'is mainly tested for Robinson projection (ESRI:54030), but it can work with 
#'other projection types in theory.\cr
#'The map can be plotted by points or polygon. A legend can be plotted as either
#'a color bar or a discrete ggplot legend. Dots can be drawn on top of the data,
#'which can be used for significance test. A mask can be added to not plot the 
#'data specified. A number of options is provided to adjust aesthetics, like 
#'position, size, colors, etc. 
#'
#'@param data A numeric array with longitude and latitude dimensions. The grid
#'  should be regular grid. It can contain NA values. 
#'@param lon A numeric vector of longitude locations of the cell centers of the 
#'  grid of 'data'. Expected to be regularly spaced, within the range of either
#'  [-180, 180] or [0, 360]. 
#'@param lat A numeric vector of latitude locations of the cell centers of the 
#'  grid of 'data'. Expected to be regularly spaced, within the range [-90, 90]
#'  of ascending or descending order.
#'@param lon_dim A character string indicating the longitude dimension name in
#'  'data'. If it is NULL, the function tries to find the name in 
#'  \code{s2dv:::.KnownLonNames}. The default value is NULL.
#'@param lat_dim A character string indicating the latitude dimension name in
#'  'data'. If it is NULL, the function tries to find the name in 
#'  \code{s2dv:::.KnownLatNames}. The default value is NULL.
#'@param target_proj A character string indicating the target projection. It 
#'  should be a valid crs string. The default projection is Robinson 
#'  (ESRI:54030). Note that the character string may work differently depending
#'  on PROJ and GDAL module version.
#'@param legend A character string indicating the legend style. It can be 's2dv'
#'  (color bar by \code{ColorBarContinuous()}), 'ggplot2' (discrete legend by 
#'  ggplot2), or NULL (no legend),
#'@param style A character string indicating the plotting style. It can be 
#'  'point' or 'polygon'. The default value is 'point'. Note that 'polygon' may
#'  be time- and memory-consuming for global or high-resolution data.
#'@param dots An array with the same dimensions as 'data' of [0, 1] or logical
#'  indicating the grids to plot dots. The value 0 or FALSE is the point to be
#'  dotted.
#'@param mask An array with the same dimensions as 'data' of [0, 1] or logical
#'  indicating the grids to not plot data. The value 0 or FALSE is the point not
#'  to be plotted.
#'@param brks,cols,bar_limits,triangle_ends Usually only providing 'brks' is 
#'  enough to generate the desired color bar. These parameters allow to 
#'  define n breaks that define n - 1 intervals to classify each of the values 
#'  in 'data'. The corresponding grid cell of a given value in 'data' will be 
#'  colored in function of the interval it belongs to. These parameters are 
#'  sent to \code{ColorBarContinuous()} to generate the breaks and colours. 
#'  Additional colors for values beyond the limits of the colour bar are also 
#'  generated and applied to the plot if 'bar_limits' or 'brks' and 
#'  'triangle_ends' are properly provided to do so. See ?ColorBarContinuous for
#'  a full explanation.
#'@param col_inf,col_sup,colNA Colour identifiers to color the values that 
#'  excess the extremes of the color bar and to color NAs, respectively. 'colNA'
#'  takes attr(cols, 'na_color') if available by default, where cols is the 
#'  parameter 'cols' if provided or the vector of colors returned by 
#'  'color_fun'. 'col_inf' and 'col_sup' will take the value of 'colNA' if not 
#'  specified. See ?ColorBarContinuous for a full explanation.
#'@param color_fun,bar_extra_margin Set of 
#'  parameters to control the visual aspect of the drawn colour bar 
#'  (1/3). See ?ColorBarContinuous for a full explanation.
#'@param vertical A logical value indicating the direction of colorbar if 
#' parameter 'legend' is 's2dv'. The default value is TRUE.
#'@param toptitle A character string of the top title of the figure, scalable 
#'  with parameter 'title_size'.
#'@param caption A character string of the caption located at left-bottom of the
#'  plot.
#'@param units A character string of the data units, which is the title of the 
#'  legend.
#'@param crop_coastlines A named numeric vector [lonmin, lonmax, latmin, latmax] 
#'  indicating the region to plot coastlines. Note that the longitude range 
#'  cannot exceed 180 degrees.
#'@param point_size A number of the size of the data points if "style = 'point'".
#'  The default is 'auto' and the function tries to find the appropriate size.
#'@param title_size A number of the size of the top title. The default is 16.
#'@param dots_size A number of the size of the dots. The default is 0.5.
#'@param dots_shape A number indicating the dot shape recognized by parameter 
#'  'shape' in \code{geom_point()}.
#'@param coastlines_width A number indicating the width of the coastlines.
#'@param fileout A character string of the path to save the plot. If not 
#'  specified (default), a graphic device will pop up. The extension should be  
#'  accepted by \code{ggsave()}.
#'@param width A number of the plot width, in the units specified in parameter
#'  'size_units'. The default is 8.
#'@param height A number of the plot height, in the units specified in parameter 
#'  'size_units'. The default is 4.
#'@param size_units A character string of the units of the size of the device
#'  (file or window) to plot in. The default is 'in' (inches). See ?ggsave and
#'  ?Devices for details of the corresponding device.
#'@param res Resolution of the device (file or window) to plot in. The default 
#'  value is 300. See ?ggsave 'dpi' and ?Devices for details of the 
#'  corresponding device.
#'
#'@return A map plot with speficied projection, either in pop-up window or a 
#'  saved file.
#'
#'@examples
#'data <- array(rep(seq(-10, 10, length.out = 181), 360) + rnorm(360),
#'              dim = c(lat = 181, lon = 360))
#'dots <- data
#'dots[which(dots < 4 & dots > -4)] <- 0
#'dots[which(dots != 0)] <- 1
#'VizRobinson(data, lon = 0:359, lat = -90:90, dots = dots,
#'            brks = seq(-10, 10, length.out = 11),
#'            toptitle = 'synthetic example', vertical = F,
#'            caption = 'Robinson Projection',
#'            bar_extra_margin = c(0, 1, 0, 1), width = 8, height = 6)
#'VizRobinson(data, lon = 0:359, lat = -90:90, mask = dots, legend = 'ggplot2',
#'           target_proj = "+proj=moll", brks = seq(-10, 10, length.out = 11),
#'           color_fun = ClimPalette("purpleorange"), colNA = 'green',
#'           toptitle = 'synthetic example', caption = 'Mollweide Projection',
#'           width = 8, height = 6)
#'
#'@import sf ggplot2 rnaturalearth cowplot utils
#'@importFrom dplyr mutate group_by summarise
#'@importFrom ClimProjDiags Subset
#'@export
VizRobinson <- function(data, lon, lat, lon_dim = NULL, lat_dim = NULL,
                        target_proj = 54030, legend = 's2dv', style = 'point', 
                        dots = NULL, mask = NULL, brks = NULL, cols = NULL, bar_limits = NULL,
                        triangle_ends = NULL, col_inf = NULL, col_sup = NULL, colNA = NULL, 
                        color_fun = ClimPalette(), bar_extra_margin = rep(0, 4), vertical = TRUE,
                        toptitle = NULL, caption = NULL, units = NULL, crop_coastlines = NULL,
                        point_size = "auto", title_size = 16, dots_size = 0.5,
                        dots_shape = 47, coastlines_width = 0.3,
                        fileout = NULL, width = 8, height = 4, size_units = "in", 
                        res = 300) {

  # Sanity check
  # data
  data <- drop(data)
  if (length(dim(data)) != 2) {
    stop("Parameter 'data' must have two dimensions.")
  }
  dims <- dim(data)
  # lon, lon_dim
  if (is.null(lon_dim)) {
    lon_dim <- names(dims)[names(dims) %in% .KnownLonNames()]
    if (identical(lon_dim, character(0))) {
      stop("Cannot find known longitude name in data dimension. Please define parameter 'lon_dim'.")
    }
  }
  if (is.unsorted(lon)) {
    warning("Parameter 'lon' should be sorted to guarantee the correct result.")
  }
  # lat, lat_dim
  if (is.null(lat_dim)) {
    lat_dim <- names(dims)[names(dims) %in% .KnownLatNames()]
    if (identical(lat_dim, character(0))) {
      stop("Cannot find known latitude name in data dimension. Please define parameter 'lat_dim'.")
    }
  }
  if (!all(names(dims) %in% c(lat_dim, lon_dim))) {
    stop("Dimensions names in paramter 'data' should match 'lat_dim' and 'lon_dim.")
  }
  if (length(lon) != dims[lon_dim]) {
    stop("Length of parameter 'lon' should match longitude dimension in 'data'.")
  }
  if (length(lat) != dims[lat_dim]) {
    stop("Length of parameter 'lat' should match latitude dimension in 'data'.")
  }

  # Reorder data
  data <- aperm(data, match(names(dim(data)), c(lon_dim, lat_dim)))

  # Make lat always from 90 to -90
  sort_lat <- FALSE
  if (!is.unsorted(lat)) {
    lat <- rev(lat)
    data <- ClimProjDiags::Subset(data, along = lat_dim, indices = seq(length(lat), 1, -1))
    sort_lat <- TRUE
  }

  # original_proj: it can only be regular grid now
  original_proj <- st_crs(4326)
  # tartget_proj
  if (is.null(target_proj)) {
    stop("Parameter 'target_proj' cannot be NULL.")
  } else {
    target_proj_tmp <- st_crs(target_proj)
    if (is.na(target_proj_tmp)) {
      warning(paste0("Try ESRI code: ESRI:", target_proj))
      target_proj <- st_crs(paste0("ESRI:", target_proj))
    } else {
      target_proj <- target_proj_tmp
    }
  }

  # legend
  if (!is.null(legend) && (!legend %in% c('s2dv', 'ggplot2'))) {
    stop("Parameter 'legend' must be NULL, ggplot2 or s2dv.")
  }
  # style
  if (!style %in% c('point', 'polygon') || length(style) != 1) {
    stop("Parameter 'style' must be 'point' or 'polygon'.")
  }
  if (style == 'polygon') {
    # polygon is slow for global map (and may be wrong) Confirm if users want to proceed
    if ((abs(diff(range(lon))) > 350 & abs(diff(range(lat))) > 175) |
        (prod(dim(data)) >= (180 * 360))) {
      if (!isTRUE(utils::askYesNo("The region seems to be global and style 'polygon' is chosen. It may be time- and memory-consuming to plot the map. Are you sure that you want to continue?"))) {
       return(invisible())
      }
    }
  }
  # dots
  if (!is.null(dots)) {
    dots <- drop(dots)
    if (!is.array(dots) || any(!names(dim(dots)) %in% c(lon_dim, lat_dim))) {
      stop("Parameter 'dots' must have two dimensions named as longitude and latitude dimensions in 'data'.")
    } else {
      dots <- aperm(dots, match(names(dim(dots)), c(lon_dim, lat_dim)))
    }
    if (!identical(dim(dots), dim(data))) {
      stop("Parameter 'dots' must have the same dimensions as 'data'.")
    } else if (is.numeric(dots)) {
      if (all(dots %in% c(0, 1))) {
        dots <- array(as.logical(dots), dim = dim(dots))
      } else {
        stop("Parameter 'dots' must have only TRUE/FALSE or 0/1.")
      }
    } else if (is.logical(dots)) {
      if (!all(dots %in% c(T, F))) {
        stop("Parameter 'dots' must have only TRUE/FALSE or 0/1.")
      }
    } else {
      stop("Parameter 'dots' must be a logical or numerical array.")
    }
  }
  # mask
  if (!is.null(mask)) {
    mask <- drop(mask)
    if (!is.array(mask) || any(!names(dim(mask)) %in% c(lon_dim, lat_dim))) {
      stop("Parameter 'mask' must have two dimensions named as longitude and latitude dimensions in 'data'.")
    } else {
      mask <- aperm(mask, match(names(dim(mask)), c(lon_dim, lat_dim)))
    }
    if (!identical(dim(mask), dim(data))) {
      stop("Parameter 'mask' must have the same dimensions as 'data'.")
    } else if (is.numeric(mask)) {
      if (all(mask %in% c(0, 1))) {
        mask <- array(as.logical(mask), dim = dim(mask))
      } else {
        stop("Parameter 'mask' must have only TRUE/FALSE or 0/1.")
      }
    } else if (is.logical(mask)) {
      if (!all(mask %in% c(T, F))) {
        stop("Parameter 'mask' must have only TRUE/FALSE or 0/1.")
      }
    } else {
      stop("Parameter 'mask' must be a logical or numerical array.")
    }
  }

  # Color bar
  ## Check: brks, cols, bar_limits, color_fun, bar_extra_margin, units 
  ## Build: brks, cols, bar_limits, col_inf, col_sup
  var_limits <- c(min(data, na.rm = TRUE), max(data, na.rm = TRUE))
  colorbar <- ColorBarContinuous(brks = brks, cols = cols, vertical = vertical, subsampleg = NULL,
                       bar_limits = bar_limits, var_limits = var_limits, triangle_ends = triangle_ends, 
                       col_inf = col_inf, col_sup = col_sup, color_fun = color_fun, 
                       plot = FALSE, draw_ticks = TRUE, draw_separators = FALSE,
                       triangle_ends_scale = 1, extra_labels = NULL,
                       title = units, title_scale = 1, # units_scale
                       label_scale = 1, tick_scale = 1, #bar_tick_scale
                       extra_margin = bar_extra_margin, label_digits = 4)
  brks <- colorbar$brks
  cols <- colorbar$cols
  col_inf <- colorbar$col_inf
  col_sup <- colorbar$col_sup
  bar_limits <- c(head(brks, 1), tail(brks, 1))
  # colNA
  if (is.null(colNA)) {
    if ('na_color' %in% names(attributes(cols))) {
      colNA <- attr(cols, 'na_color')
      if (!.IsColor(colNA)) {
        stop("The 'na_color' provided as attribute of the colour vector must be a valid colour identifier.")
      }
    } else {
      colNA <- 'pink'
    }
  } else if (!.IsColor(colNA)) {
    stop("Parameter 'colNA' must be a valid colour identifier.")
  }
  # toptitle
  if (!is.null(toptitle) && !is.character(toptitle)) {
    stop("Parameter 'toptitle' must be a character string.")
  }
  # caption
  if (!is.null(caption) && !is.character(caption)) {
    stop("Parameter 'caption' must be a character string.")
  }
  # crop_coastlines
  if (!is.null(crop_coastlines)) {
    # if crop_coastlines doesn't have name, [lonmin, lonmax, latmin, latmax]
    if (is.null(names(crop_coastlines))) {
      names(crop_coastlines) <- c("lonmin", "lonmax", "latmin", "latmax")
    } else if (!identical(sort(names(crop_coastlines)), sort(c("latmax", "latmin", "lonmax", "lonmin")))) {
      stop("Parameter 'crop_coastlines' needs to have names 'latmax', 'latmin', 'lonmax', 'lonmin'.")
    }
  }

  # point_size
  if (point_size == 'auto') {
    # 360x181 with default setting, 0.05
    point_size <- round(0.05 * (360 * 181) / (length(lon) * length(lat)), 2)
  } else if (!is.numeric(point_size)) {
    stop("Parameter 'point_size' must be a number.")
  }
  #

#=================================================================

  # Adapt s2dv ColorBar parameters to ggplot plot
  # If legend is NULL, still tune with s2dv legend way
  if (is.null(legend) || legend == 's2dv') {
    # the colorbar triangle color. If it is NULL (no triangle plotted), use colNA
    col_inf_image <- ifelse(is.null(col_inf), colNA, col_inf)
    col_sup_image <- ifelse(is.null(col_sup), colNA, col_sup)
    cols_ggplot <- c(col_inf_image, cols, col_sup_image)

    # Add triangles to brks
    brks_ggplot <- brks
    if (max(data, na.rm = T) > tail(brks, 1)) {
      brks_ggplot <- c(brks_ggplot, max(data, na.rm = T))
    } else {
      brks_ggplot <- c(brks_ggplot, tail(brks, 1) + diff(tail(brks, 2)))
    }
    if (min(data, na.rm = T) < brks[1]) {
      brks_ggplot <- c(min(data, na.rm = T), brks_ggplot)
    } else {
      brks_ggplot <- c(brks[1] - diff(brks[1:2]), brks_ggplot)
    }

  } else {  # ggplot2 legend
    brks_ggplot <- brks
    cols_ggplot <- cols
  }

  # Build data dataframe
  lonlat_df <- data.frame(lon = rep(as.vector(lon), length(lat)),
                          lat = sort(rep(as.vector(lat), length(lon)), decreasing = TRUE))
  data_df <- lonlat_df %>%
    dplyr::mutate(dat = as.vector(data))

  lonlat_df_ori <- NULL
  # Remove the points where mask = FALSE
  if (!is.null(mask)) {
    # Save original lonlat_df to plot with expected region
    lonlat_df_ori <- st_as_sf(lonlat_df, coords = c("lon", "lat"), crs = original_proj)
    lonlat_df_ori <- st_transform(lonlat_df_ori, crs = target_proj)
    lonlat_df_ori <- as.data.frame(st_coordinates(lonlat_df_ori))
    names(lonlat_df_ori) <- c('long', 'lat')

    if (sort_lat) {
      mask <- ClimProjDiags::Subset(mask, along = lat_dim, indices = seq(length(lat), 1, -1))
    }
    mask_df <- data.frame(lon = rep(as.vector(lon), length(lat)),
                          lat = sort(rep(as.vector(lat), length(lon)), decreasing = TRUE),
                          mask = as.vector(mask))
    data_df <- data_df[mask_df$mask == TRUE, ]
    lonlat_df <- data_df[, 1:2]
  }

  #NOTE: if target_proj = "ESRI:54030", Nord3v2 has different behavior from hub and ws!!
  data_df <- st_as_sf(data_df, coords = c("lon", "lat"), crs = original_proj)
  data_df <- st_transform(data_df, crs = target_proj)
  data_df <- data_df %>%
    dplyr::mutate(long = st_coordinates(data_df)[, 1],
                  lat = st_coordinates(data_df)[, 2])

  # Re-project dots
  if (!is.null(dots)) {
    if (sort_lat) {
      dots <- ClimProjDiags::Subset(dots, along = lat_dim, indices = seq(length(lat), 1, -1))
    }
    dots_df <- data.frame(lon = rep(as.vector(lon), length(lat)),
                          lat = sort(rep(as.vector(lat), length(lon)), decreasing = TRUE),
                          dot = as.vector(dots))

    dots_df <- st_as_sf(dots_df, coords = c("lon", "lat"), crs = original_proj)
    dots_df <- st_transform(dots_df, crs = target_proj)
    dots_df <- dots_df %>%
      dplyr::mutate(long = st_coordinates(dots_df)[, 1],
                    lat = st_coordinates(dots_df)[, 2])
    dots_df <- subset(dots_df, dot == FALSE)
  }

  # coastlines
  coastlines <- rnaturalearth::ne_coastline(scale = "medium", returnclass = "sf")
  ## crop the coastlines to the desired range
  if (!is.null(crop_coastlines)) {
    suppressWarnings({
    coastlines <- st_crop(coastlines, 
                          xmin = as.numeric(crop_coastlines['lonmin']), 
                          xmax = as.numeric(crop_coastlines['lonmax']), 
                          ymin = as.numeric(crop_coastlines['latmin']), 
                          ymax = as.numeric(crop_coastlines['latmax']))
    })
  }
  coastlines <- st_transform(coastlines, crs = target_proj) 

  if (style == 'polygon') {
    # Calculate polygon points from regular lat/lon
    #NOTE: The original grid must be regular grid with same space
    d_lon <- abs(lon[2] - lon[1]) / 2
    d_lat <- abs(lat[2] - lat[1]) / 2
    lon_poly <- lat_poly <- rep(NA, 4 * dim(lonlat_df)[1])
    for (ii in 1:dim(lonlat_df)[1]) {
      lon_poly[(ii*4-3):(ii*4)] <- c(lonlat_df$lon[ii] - d_lon, lonlat_df$lon[ii] + d_lon,
                                     lonlat_df$lon[ii] + d_lon, lonlat_df$lon[ii] - d_lon)
      lat_poly[(ii*4-3):(ii*4)] <- c(lonlat_df$lat[ii] - d_lat, lonlat_df$lat[ii] - d_lat,
                                     lonlat_df$lat[ii] + d_lat, lonlat_df$lat[ii] + d_lat)
    }
    #  # To prevent out-of-global lon
    #  lon_poly[which(lon_poly >= 180)] <- 179.9
    #  lon_poly[which(lon_poly < -180)] <- -180
    # To prevent out-of-global lat
    lat_poly[which(lat_poly > 90)] <- 90
    lat_poly[which(lat_poly < -90)] <- -90

    lonlat_df <- data.frame(lon = lon_poly, lat = lat_poly)
    # Transfer lon/lat to projection
    proj_lonlat <- st_as_sf(lonlat_df, coords = c("lon", "lat"), crs = original_proj)
    #NOTE: if target_proj = "ESRI:54030", on Nord3v2, st_transform has lon and lat swapped!
    proj_lonlat <- st_transform(proj_lonlat, crs = target_proj)
    lonlat_df_proj <- st_coordinates(proj_lonlat)

    # Use id to create groups for each polygon
    ids <- factor(paste0("id_", 1:dim(data_df)[1]))
    values <- data.frame(id = ids,
                         value = data_df$dat)
    positions <- data.frame(id = rep(ids, each = 4),
                            x = lonlat_df_proj[, 1],
                            y = lonlat_df_proj[, 2])
    datapoly <- merge(values, positions, by = "id")
    datapoly <- st_as_sf(datapoly, coords = c("x", "y"), crs = target_proj)
    datapoly <- datapoly %>%
                dplyr::group_by(id) %>%
                dplyr::summarise() %>%  #NOTE: VERY SLOW if plot global
                dplyr::mutate(value = values[order(values$id), ]$value) %>%
                st_cast("POLYGON") %>%
                st_convex_hull()  # maintain outer polygen (no bowtie shape)
  }

  # Plots
  if (style == 'polygon') {
    res_p <- ggplot(data = data_df) +  #NOTE: must be data_df?
             geom_sf(data = datapoly,
                     aes(col = cut(value, breaks = brks_ggplot, include.lowest = T),
                         fill = cut(value, breaks = brks_ggplot, include.lowest = T)))
  } else if (style == 'point') {
    res_p <- ggplot(data = data_df) +
             geom_point(aes(x = long, y = lat,
                            col = cut(dat, breaks = brks_ggplot, include.lowest = T)),
                        #NOTE: These two lines make point size vary with lat
                        #size = point_size /  (data_df$lat / min(data_df$lat))) +
                        #size = (sort(rep(as.vector(lat), length(lon))) / max(lat)) * point_size) +
                        size = point_size)
  }

  if (is.null(lonlat_df_ori)) {
    coord_sf_lim <- c(range(data_df$long), range(data_df$lat))
  } else {
    coord_sf_lim <- c(range(lonlat_df_ori$long), range(lonlat_df_ori$lat))
  }
  res_p <- res_p +
           geom_sf(data = coastlines, colour ='black', size = coastlines_width) +
           # Remove background grid and lat/lon label; add white background
           theme_void() + theme(plot.background = element_rect(fill = 'white', colour = 'white')) +
           # crop the projection
           coord_sf(xlim = coord_sf_lim[1:2], ylim = coord_sf_lim[3:4],
                    expand = TRUE, datum = target_proj)

  if (!is.null(dots)) {
    res_p <- res_p + geom_point(data = dots_df, aes(x = long, y = lat),
                                shape = dots_shape, size = dots_size)
                                #NOTE: This line makes point size vary with lat
                                #size = dots_size / (dots_df$lat / min(dots_df$lat)))
  }

  if (identical(legend, 'ggplot2')) {
    if (style == 'polygon') {
    res_p <- res_p + scale_colour_manual(values = cols_ggplot,
                                         aesthetics = c("colour", "fill"),
                                         drop = FALSE, na.value = colNA) +
             guides(fill = guide_legend(title = units, override.aes = list(size = 1)),
                    color = "none")
   } else if (style == 'point') {
    res_p <- res_p + scale_colour_manual(values = cols_ggplot,
                                         drop = FALSE, na.value = colNA) +
             guides(colour = guide_legend(title = units, override.aes = list(size = 1)))
   }

  } else {  # s2dv or NULL
    if (style == 'polygon') {
      res_p <- res_p + scale_colour_manual(values = cols_ggplot,
                                           aesthetics = c("colour", "fill"),
                                           drop = FALSE, na.value = colNA)
    } else if (style == 'point') {
      res_p <- res_p + scale_colour_manual(values = cols_ggplot,
                                           drop = FALSE, na.value = colNA)
    }
    # Remove ggplot legend
    res_p <- res_p + theme(legend.position = "none", plot.margin = margin(0.5, 0, 0, 0, 'cm'))
  }

  if (!is.null(toptitle)) {
    res_p <- res_p + ggtitle(toptitle) +
      theme(plot.title = element_text(size = title_size, hjust = 0.5, vjust = 3))
  }
  if (!is.null(caption)) {
    res_p <- res_p + labs(caption = caption) +
      theme(plot.caption = element_text(hjust = 0, vjust = 1, margin = margin(0, 0, 0, 0, 'cm')))
  }

  # s2dv legend fun to put in cowplot::plot_grid
  if (identical(legend, 's2dv')) {
    fun_legend <- function() {
      if (vertical) {
        par(mar = c(7.1, 2.2, 7.1, 3.1), mgp = c(3, 1, 0))
      } else {
        par(mar = c(1.1, 1.2, 0.1, 1.1), mgp = c(3, 1, 0))
      }
      ColorBarContinuous(brks = brks, cols = cols, vertical = vertical, subsampleg = NULL,
               bar_limits = bar_limits, var_limits = var_limits, triangle_ends = triangle_ends,
               col_inf = col_inf, col_sup = col_sup, color_fun = color_fun,
               plot = TRUE, draw_ticks = TRUE, draw_separators = FALSE,
               triangle_ends_scale = 1, extra_labels = NULL,
               title = units, title_scale = 1, # units_scale
               label_scale = 1, tick_scale = 1, #bar_tick_scale
               extra_margin = bar_extra_margin, label_digits = 4)
    }
    if (vertical) {
      res_p <- cowplot::plot_grid(res_p, fun_legend, rel_widths = c(6, 1))
    } else {
      res_p <- cowplot::plot_grid(res_p, fun_legend, rel_heights = c(5, 1), ncol = 1)
    }
    res_p <- res_p +  theme(plot.background = element_rect(fill = "white", colour = "white"))
  }

  if (!is.null(fileout)) {
    ext <- regmatches(fileout, regexpr("[a-zA-Z0-9]*$", fileout))
    ggsave(fileout, res_p, width = width, height = height, dpi = res, units = size_units, 
           device = ext)
  } else { # pop-up window
    dev.new(units = size_units, res = res, width = width, height = height)
    res_p 
  } 

}