PlotMultiMap.R

PlotMultiMap <- function(fun, var, ..., nrow = NULL, ncol = NULL,
                         toptitle = NULL, row_titles = NULL, col_titles = NULL,
                         bar_scale = 1, title_scale = 1, subtitle_scale = 1,
                         brks = NULL, cols = NULL, drawleg = 'S',
                         titles = NULL, subsampleg = NULL,
                         bar_limits = NULL, triangle_ends = NULL,
                         color_fun = clim.colors, draw_ticks = TRUE,
                         draw_separators = FALSE, triangle_ends_scale = 1,
                         bar_extra_labels = NULL, units = NULL, units_scale = 1,
                         bar_label_scale = 1, bar_tick_scale = 1,  
                         bar_extra_margin = rep(0, 4), bar_label_digits = 4,
                         fileout = NULL, width = NULL, height = NULL) {
  # If there is any filenames to store the graphics, process them
  # to select the right device 
  if (!is.null(fileout)) {
    deviceInfo <- .SelectDevice(fileout)
    saveToFile <- deviceInfo$fun
    fileout <- deviceInfo$files
  }

  is_single_na <- function(x) ifelse(length(x) > 1, FALSE, is.na(x))
  # Check var
  if (is.array(var) || (length(var) == 1 && anyNA(var))) {
    var <- list(var)
  } else if (is.list(var)) {
    if (!all(sapply(var, is.array) | sapply(var, is_single_na))) {
      stop("Parameter 'var' must be an array or a list of arrays (or NA values).")
    }
  } else {
    stop("Parameter 'var' must be an array or a list of arrays.")
  }

  # Check fun
  if (!is.function(fun)) {
    stop("Parameter 'fun' must be a function.")
  }

  # Check nrow
  if (!is.null(nrow)) {
    if (!is.numeric(nrow)) {
      stop("Parameter 'nrow' must be numeric or NULL.")
    }
    nrow <- round(nrow)
  }

  # Check ncol
  if (!is.null(ncol)) {
    if (!is.numeric(ncol)) {
      stop("Parameter 'ncol' must be numeric or NULL.")
    }
    ncol <- round(ncol)
  }

  # Check toptitle
  if (is.null(toptitle) || is.na(toptitle)) {
    toptitle <- ''
  }
  if (!is.character(toptitle)) {
    stop("Parameter 'toptitle' must be a character string.")
  }

  # Check row_titles
  if (!is.null(row_titles)) {
    if (!is.character(row_titles)) {
      stop("Parameter 'row_titles' must be a vector of character strings.")
    }
  }

  # Check col_titles
  if (!is.null(row_titles)) {
    if (!is.character(row_titles)) {
      stop("Parameter 'row_titles' must be a vector of character strings.")
    }
  }

  # Check drawleg
  if (is.character(drawleg)) {
    if (drawleg %in% c('up', 'u', 'U', 'top', 't', 'T', 'north', 'n', 'N')) {
      drawleg <- 'N'
    } else if (drawleg %in% c('down', 'd', 'D', 'bottom', 'b', 'B', 'south', 's', 'S')) {
      drawleg <- 'S'
    } else if (drawleg %in% c('right', 'r', 'R', 'east', 'e', 'E')) {
      drawleg <- 'E'
    } else if (drawleg %in% c('left', 'l', 'L', 'west', 'w', 'W')) {
      drawleg <- 'W'
    } else {
      stop("Parameter 'drawleg' must be either TRUE, FALSE or a valid identifier of a position (see ?PlotMultiMap).")
    }
  } else if (!is.logical(drawleg)) {
    stop("Parameter 'drawleg' must be either TRUE, FALSE or a valid identifier of a position (see ?PlotMultiMap).")
  }
  if (drawleg != FALSE && all(sapply(var, is_single_na)) && 
      (is.null(brks) || length(brks) < 2)) {
    stop("Either data arrays in 'var' or breaks in 'brks' must be provided if 'drawleg' is requested.")
  }

  # Check the rest of parameters (unless the user simply wants to build an empty layout)
  if (!(drawleg == FALSE)) {
    var_limits <- NULL
    if (!all(sapply(var, is_single_na))) {
      var_limits <- c(min(unlist(var), na.rm = TRUE), max(unlist(var), na.rm = TRUE))
      if ((any(is.infinite(var_limits)) || var_limits[1] == var_limits[2]))  {
        stop("Arrays in parameter 'var' must contain at least 2 different values.")
      }
    }
    colorbar <- ColorBar(brks, cols, FALSE, subsampleg, bar_limits,
                         var_limits,
                         triangle_ends, color_fun, plot = FALSE, draw_ticks, 
                         draw_separators, triangle_ends_scale, bar_extra_labels,
                         units, units_scale, bar_label_scale, bar_tick_scale,
                         bar_extra_margin, bar_label_digits)
  }
 
  # Check bar_scale
  if (!is.numeric(bar_scale)) {
    stop("Parameter 'bar_scale' must be numeric.")
  }

  # Check title_scale
  if (!is.numeric(title_scale)) {
    stop("Parameter 'title_scale' must be numeric.")
  }

  # Check subtitle_scale
  if (!is.numeric(subtitle_scale)) {
    stop("Parameter 'subtite_scale' must be numeric.")
  }

  # Check titles
  if (!all(sapply(titles, is.character))) {
    stop("Parameter 'titles' must be a vector of character strings.")
  }

  # Check width
  if (is.null(width)) {
    if (is.null(fileout)) {
      width <- 7
    } else {
      width <- 11
    }
  }
  if (!is.numeric(width)) {
    stop("Parameter 'width' must be numeric.")
  }

  # Check height
  if (is.null(height)) {
    if (is.null(fileout)) {
      height <- 7
    } else {
      height <- 8
    }
  }
  if (!is.numeric(height)) {
    stop("Parameter 'height' must be numeric.")
  }

  # Count the total number of maps and reorder each array of maps to have the lat and lon dimensions at the end.
  n_maps <- 0
  nlats <- NULL
  nlons <- NULL
  map_array_i <- 1
  for (map_array in var) {
    if (is_single_na(map_array)) {
      var[[map_array_i]] <- NA
      n_maps <- n_maps + 1
    } else {
      dimnames <- NULL
      if (!is.null(names(map_array))) {
        dimnames <- names(map_array)
      } else if (!is.null(attr(map_array, 'dimensions'))) {
        dimnames <- attr(map_array, 'dimensions')
      }
      if (!is.null(dimnames)) {
        if (!('lon' %in% dimnames || 'lat' %in% dimnames)) {
          stop("Parameter 'var' contains arrays without the dimensions 'lat' and/or 'lon'.")
        }
        lat_dim <- which(dimnames == 'lat')[1]
        lon_dim <- which(dimnames == 'lon')[1]
        map_array_nlats <- dim(map_array)[lat_dim]
        map_array_nlons <- dim(map_array)[lon_dim]
        n_maps <- n_maps + prod(dim(map_array)[-c(lat_dim, lon_dim)])
        var[[map_array_i]] <- aperm(var[[map_array_i]], c((1:length(dim(map_array)))[-c(lat_dim, lon_dim)], lat_dim, lon_dim))
      } else {
        .warning(paste0("Assuming the ", map_array_i, "th array provided in 'var' has 'lat' and 'lon' as last dimensions (right-most)."))
        dims <- tail(c(1, 1, dim(map_array)), 2)
        map_array_nlats <- dims[1]
        map_array_nlons <- dims[2]
        if (length(dim(var[[map_array_i]])) < 2) {
          dim(var[[map_array_i]]) <- dims
        }
        n_maps <- n_maps + prod(head(c(1, 1, dim(map_array)), length(dim(map_array))))
      }
      if (length(dim(var[[map_array_i]])) == 2) {
        dim(var[[map_array_i]]) <- c(1, dim(var[[map_array_i]]))
      }
      if (is.null(nlats)) {
        nlats <- map_array_nlats
        nlons <- map_array_nlons
      } else if (map_array_nlons != nlons || map_array_nlats != nlats) {
        stop(paste0("All the provided arrays in 'var' must have the same number of latitudes and longitudes (the ", map_array_i, "th array was the first unmatching array)."))
      }
    }
    map_array_i <- map_array_i + 1
  }
  if (is.null(nrow) && is.null(ncol)) {
    ncol <- ceiling(sqrt(n_maps))
    nrow <- ceiling(n_maps/ncol)
  } else if (is.null(ncol)) {
    ncol <- ceiling(n_maps/nrow)
  } else if (is.null(nrow)) {
    nrow <- ceiling(n_maps/ncol)
  } else if (nrow * ncol < n_maps) {
    stop("There are more maps in 'var' than cells defined by 'nrow' x 'ncol'.")
  }

  if (is.logical(drawleg) && drawleg) {
    if (nrow > ncol) {
      drawleg <- 'S'
    } else {
      drawleg <- 'E'
    }
  }
  vertical <- drawleg %in% c('E', 'W')

  # Open connection to graphical device
  if (!is.null(fileout)) {
    saveToFile(fileout, width = width, height = height)
  } else {
    dev.new(width = width, height = height)
  }

  # Take size of device and set up layout
  device_size <- par('din')
  cs <- char_size <- par('csi')
  title_cex <- 2.5 * title_scale
  title_margin <- 0.5 * title_cex
  subtitle_cex <- 1.5 * subtitle_scale
  subtitle_margin <- 0.5 * sqrt(nrow * ncol) * subtitle_cex
  mat_layout <- 1:(nrow * ncol)
  mat_layout <- matrix(mat_layout, nrow, ncol, byrow = TRUE)
  fsu <- figure_size_units <- 10  # unitless
  widths <- rep(fsu, ncol)
  heights <- rep(fsu, nrow)
  n_figures <- nrow * ncol
  if (length(row_titles) > 0) {
    mat_layout <- cbind(rep(0, dim(mat_layout)[1]), mat_layout)
    widths <- c(((subtitle_cex + subtitle_margin) * cs / device_size[2]) * ncol * fsu, widths)
  }
  if (length(col_titles) > 0) {
    mat_layout <- rbind(rep(0, dim(mat_layout)[2]), mat_layout)
    heights <- c(((subtitle_cex + subtitle_margin) * cs / device_size[2]) * nrow * fsu, heights)
  }
  if (drawleg != FALSE) {
    if (drawleg == 'N') {
      mat_layout <- rbind(rep(n_figures + 1, dim(mat_layout)[2]), mat_layout)
      heights <- c(round(bar_scale * 2 * nrow), heights)
    } else if (drawleg == 'S') {
      mat_layout <- rbind(mat_layout, rep(n_figures + 1, dim(mat_layout)[2]))
      heights <- c(heights, round(bar_scale * 2 * nrow))
    } else if (drawleg == 'W') {
      mat_layout <- cbind(rep(n_figures + 1, dim(mat_layout)[1]), mat_layout)
      widths <- c(round(bar_scale * 3 * ncol), widths)
    } else if (drawleg == 'E') {
      mat_layout <- cbind(mat_layout, rep(n_figures + 1, dim(mat_layout)[1]))
      widths <- c(widths, round(bar_scale * 3 * ncol))
    }
    n_figures <- n_figures + 1
  }
  if (toptitle != '') {
    mat_layout <- rbind(rep(0, dim(mat_layout)[2]), mat_layout)
    heights <-  c(((title_cex + title_margin) * cs / device_size[2]) * nrow * fsu, heights)
  }
  layout(mat_layout, widths, heights)
  # Draw titles
  if (toptitle != '' || length(col_titles) > 0 || length(row_titles) > 0) {
    plot(0, type = 'n', ann = FALSE, axes = FALSE, xaxs = 'i', yaxs = 'i', 
         xlim = c(0, 1), ylim = c(0, 1))
    width_lines <- par('fin')[1] / par('csi')
    plot_lines <- par('pin')[1] / par('csi')
    plot_range <- par('xaxp')[2] - par('xaxp')[1]
    plot_units_per_line <- plot_range / plot_lines
    if (toptitle != '') {
      title_x_center <- par('xaxp')[1] - par('mar')[2] * plot_units_per_line + 
                        ncol * width_lines * plot_units_per_line / 2
      title_y_center <- par('mar')[3] + title_margin / 2
      if (length(col_titles > 0)) {
        title_y_center <- title_y_center + subtitle_margin + subtitle_cex
      }
      mtext(toptitle, cex = title_cex, line = title_y_center, at = title_x_center)
    }
    if (length(col_titles) > 0) {
      t_x_center <- par('xaxp')[1] - par('mar')[2] * plot_units_per_line
      for (t in 1:ncol) {
        mtext(col_titles[t], cex = subtitle_cex,
              line = par('mar')[3] + subtitle_margin / 2,
              at = t_x_center + (t - 0.5) * width_lines * plot_units_per_line)
      }
    }
    height_lines <- par('fin')[2] / par('csi')
    plot_lines <- par('pin')[2] / par('csi')
    plot_range <- par('yaxp')[2] - par('yaxp')[1]
    plot_units_per_line <- plot_range / plot_lines
    if (length(row_titles) > 0) {
      t_y_center <- par('yaxp')[1] - par('mar')[1] * plot_units_per_line
      for (t in 1:nrow) {
        mtext(row_titles[t], cex = subtitle_cex,
              line = par('mar')[2] + subtitle_margin / 2, 
              at = t_y_center - (t - 1.5) * height_lines * plot_units_per_line, side = 2)
      }
    }
    .SwitchToFigure(n = 1, mat = mat_layout)
  }
  map_number <- 1
  # For each array provided in var
  lapply(var, function(x) {
    if (length(x) == 1 && is.na(x)[1]) {
      if (map_number < nrow * ncol) {
        .SwitchToFigure(n = map_number + 1, mat = mat_layout)
      }
      map_number <<- map_number + 1
    } else {
      # For each of the maps provided in that array
      apply(x, (1:length(dim(x)))[1:(length(dim(x)) - 2)], function(y) {
        # Plot the map
        fun(y, toptitle = titles[map_number], brks = colorbar$brks, 
                                             cols = colorbar$cols, 
               triangle_ends = colorbar$triangle_ends, drawleg = FALSE, ...)
        map_number <<- map_number + 1
      })
    }
  })

  # Create output containers

  # Draw the color bar
  if (drawleg != FALSE) {
    # Set as next figure the slot kept for the colorbar
    .SwitchToFigure(n = n_figures, mat = mat_layout)
    ColorBar(colorbar$brks, colorbar$cols, vertical, subsampleg, 
             bar_limits, var_limits = NULL,
             triangle_ends = colorbar$triangle_ends, color_fun, 
             plot = TRUE, draw_ticks, 
             draw_separators, triangle_ends_scale, bar_extra_labels,
             units, units_scale, bar_label_scale, bar_tick_scale, 
             bar_extra_margin, bar_label_digits)
    # Set as next figure the next free slot, if any
    .SwitchToFigure(n = (map_number - 1) %% (nrow * ncol) + 1, mat = mat_layout)
  }

  # Set the first figure as next figure if the user requested empty layout
  if (all(sapply(var, is_single_na))) {
    .SwitchToFigure(n = 1, mat = mat_layout)
  }

  # If the graphic was saved to file, close the connection with the device
  if (!is.null(fileout)) dev.off()

  invisible(list(brks = brks, cols = cols, triangle_ends = triangle_ends,
                 layout_matrix = mat_layout))
}