PlotMultiMap.R

PlotMultiMap <- function(fun, var, ..., nrow = NULL, ncol = NULL, 
                         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,
                         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, 
                         bar_scale = 1) {
  # Check var
  if (is.array(var)) {
    var <- list(var)
  } else if (is.list(var)) {
    if (!all(sapply(var, is.numeric) || !all(sapply(var, is.array)))) {
      stop("Parameter 'var' must be an array or a list of arrays.")
    }
  } else {
    stop("Parameter 'var' must be an array or a list of arrays.")
  }

  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 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).")
  }
 
  # Check bar_scale
  if (!is.numeric(bar_scale)) {
    stop("Parameter 'bar_scale' must be numeric.")
  }

  # Check the rest of parameters
  colorbar <- ColorBar(brks, cols, FALSE, subsampleg, bar_limits,
                       c(min(unlist(var), na.rm = TRUE), max(unlist(var), na.rm = TRUE)),
                       triangle_ends, color_fun, plot = FALSE, draw_ticks, 
                       draw_separators, triangle_ends_scale, extra_labels,
                       units, units_scale, bar_label_scale, bar_tick_scale,
                       bar_extra_margin, bar_label_digits)

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

  # 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) {
    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')

  mat_layout <- 1:(nrow * ncol)
  mat_layout <- matrix(mat_layout, nrow, ncol, byrow = TRUE)
  widths <- rep(1, ncol)
  heights <- rep(1, nrow)
  if (drawleg == 'N') {
    mat_layout <- rbind(rep(nrow * ncol + 1, ncol), mat_layout)
    heights <- c(round(bar_scale * 2 * nrow), heights * 10)
  } else if (drawleg == 'S') {
    mat_layout <- rbind(mat_layout, rep(nrow * ncol + 1, ncol))
    heights <- c(heights * 10, round(bar_scale * 2 * nrow))
  } else if (drawleg == 'W') {
    mat_layout <- cbind(rep(nrow * ncol + 1, nrow), mat_layout)
    widths <- c(round(bar_scale * 3 * ncol), widths * 10)
  } else if (drawleg == 'E') {
    mat_layout <- cbind(mat_layout, rep(nrow * ncol + 1, nrow))
    widths <- c(widths * 10, round(bar_scale * 3 * ncol))
  }
  layout(mat_layout, widths, heights)
  map_number <- 1
  lapply(var, function(x) {
    apply(x, (1:length(dim(x)))[1:(length(dim(x)) - 2)], function(y) {
      fun(y, toptitle = titles[map_number], brks = colorbar$brks, 
                                            cols = colorbar$cols, 
             triangle_ends = colorbar$triangle_ends, drawleg = FALSE, ...)
      map_number <<- map_number + 1
    })
  })

  if (drawleg != FALSE) {
    # Set as next figure the slot kept for the colorbar
    next_figure <- which(mat_layout == (nrow * ncol + 1), arr.ind = TRUE)[1, ]
    .SwitchToFigure(next_figure[1], next_figure[2])
    colorbar <- 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, 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
    #next_figure <- which(mat_layout == (((map_number - 1) %% (nrow * ncol)) + 1), 
    #                     arr.ind = TRUE)[1, ]
    #.SwitchToFigure(next_figure[1], next_figure[2])
  }
}