PlotMostLikelyQuantileMap.R

PlotMostLikelyQuantileMap <- function(probs, lon, lat, cat_dim = 'bin',
                                      brks = NULL, cols = NULL, bar_titles = NULL,
                                      legend_scale = 1, col_unknown_cat = 'white',
                                      mask = NULL, col_mask = 'grey',
                                      ...) {
  # Check probs
  error <- FALSE
  if (!is.numeric(probs)) {
    error <- TRUE
  }
  if (is.null(dim(probs))) {
    error <- TRUE
  }
  if (length(dim(probs)) != 3) {
    error <- TRUE
  }
  if (error) {
    stop("Parameter 'probs' must be a numeric array with 3 dimensions.")
  }
  dimnames <- names(dim(probs))

  # Check cat_dim
  if (is.character(cat_dim)) {
    if (is.null(dimnames)) {
      stop("Specified a dimension name in 'cat_dim' but no dimension names provided ",
           "in 'probs'.")
    }
    cat_dim <- which(dimnames == cat_dim)[1]
    if (length(cat_dim) != 1) {
      stop("Dimension 'cat_dim' not found in 'prob'.")
    }
  } else if (!is.numeric(cat_dim)) {
    stop("Parameter 'cat_dim' must be either a numeric value or a ",
         "dimension name.")
  }
  if (length(cat_dim) != 1) {
    stop("Parameter 'cat_dim' must be of length 1.")
  }
  cat_dim <- round(cat_dim)

  # Work out lon_dim and lat_dim
  lon_dim <- NULL
  if (!is.null(dimnames)) {
    lon_dim <- which(dimnames %in% c('lon', 'longitude'))[1]
  }
  if (length(lon_dim) < 1) {
    lon_dim <- (1:3)[-cat_dim][1]
  }
  lon_dim <- round(lon_dim)

  lat_dim <- NULL
  if (!is.null(dimnames)) {
    lat_dim <- which(dimnames %in% c('lat', 'latitude'))[1]
  }
  if (length(lat_dim) < 1) {
    lat_dim <- (1:3)[-cat_dim][2]
  }
  lat_dim <- round(lat_dim)

  # Check lon
  if (!is.numeric(lon)) {
    stop("Parameter 'lon' must be a numeric vector.")
  }
  if (length(lon) != dim(probs)[lon_dim]) {
    stop("Parameter 'lon' does not match the longitude dimension in 'probs'.")
  }

  # Check lat
  if (!is.numeric(lat)) {
    stop("Parameter 'lat' must be a numeric vector.")
  }
  if (length(lat) != dim(probs)[lat_dim]) {
    stop("Parameter 'lat' does not match the longitude dimension in 'probs'.")
  }

  # Check brks
  if (is.null(brks)) {
    brks <- ceiling(1 / dim(probs)[cat_dim] * 10 * 1.1) * 10
    brks <- seq(from = brks, to = 100, length.out = 5)
  }
  if (!is.numeric(brks)) {
    stop("Parameter 'brks' must be a numeric vector.")
  }

  # Check cols
  col_sets <- list(c("#A1D99B", "#74C476", "#41AB5D", "#238B45"),
                   c("#6BAED6FF", "#4292C6FF", "#2171B5FF", "#08519CFF"),
                   c("#FFEDA0FF", "#FED976FF", "#FEB24CFF", "#FD8D3CFF"),
                   c("#FC4E2AFF", "#E31A1CFF", "#BD0026FF", "#800026FF"),
                   c("#FCC5C0", "#FA9FB5", "#F768A1", "#DD3497"))
  if (is.null(cols)) {
    if (length(col_sets) >= dim(probs)[cat_dim]) {
      chosen_sets <- 1:(dim(probs)[cat_dim])
      chosen_sets <- chosen_sets + floor((length(col_sets) - length(chosen_sets)) / 2)
    } else {
      chosen_sets <- array(1:length(col_sets), dim(probs)[cat_dim])
    }
    cols <- col_sets[chosen_sets]
  } else {
    if (!is.list(cols)) {
      stop("Parameter 'cols' must be a list of character vectors.")
    }
    if (!all(sapply(cols, is.character))) {
      stop("Parameter 'cols' must be a list of character vectors.")
    }
    if (length(cols) != dim(probs)[cat_dim]) {
      stop("Parameter 'cols' must be a list of the same length as the number of ",
           "categories in 'probs'.")
    }
  }
  for (i in 1:length(cols)) {
    if (length(cols[[i]]) != (length(brks) - 1)) {
      cols[[i]] <- colorRampPalette(cols[[i]])(length(brks) - 1)
    }
  }

  # Check bar_titles
  if (is.null(bar_titles)) {
    if (!is.null(names(cols))) {
      bar_titles <- names(cols)
    } else {
      if (length(cols) == 3) {
        bar_titles <- list("Below normal (%)", "Normal (%)", "Above normal (%)")
      } else if (length(cols) == 5) {
        bar_titles <- list("Low (%)", "Below normal (%)", "Normal (%)", "Above normal (%)", "High (%)")
      } else {
        bar_titles <- paste0("Cat. ", 1:length(cols), " (%)")
      }
    }
  } else {
    if (!is.character(bar_titles)) {
      stop("Parameter 'bar_titles' must be a character vector.")
    }
    if (length(bar_titles) != length(cols)) {
      stop("Parameter 'bar_titles' must be of the same length as the number of ",
           "categories in 'probs'.")
    }
  }

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

  # Check col_unknown_cat

  # Check col_mask
  if (!is.character(col_mask)) {
    stop("Parameter 'col_mask' must be a character string.")
  }

 if (!is.null(mask)){ 
 # Check mask
  if (!is.numeric(mask)) {
    stop("Parameter 'mask' must be numeric.")
  }
  if (length(dim(mask)) != 2) {
    stop("Parameter 'mask' must have two dimensions.")
  }
  if ((dim(mask)[1] != dim(probs)[lat_dim]) || 
      (dim(mask)[2] != dim(probs)[lon_dim])) {
    stop("Parameter 'mask' must have dimensions c(lat, lon).")
  }
 }

  #----------------------
  # Identify the most likely category
  #----------------------
  MLCat <- function(x) {	
    a <- which.max(x) + max(x)
    if (length(a)) {
      return(a)
    } else {
      return(NA)
    }
  }
  mlcat <- apply(probs, c(lat_dim, lon_dim), MLCat)
  ncat <- dim(probs)[cat_dim]
  nlat <- length(lat)
  nlon <- length(lon)

  #----------------------
  # Set latitudes from minimum to maximum
  #----------------------
  if (lat[1] > lat[nlat]){
    lat <- lat[nlat:1]
    indices <- list(nlat:1, TRUE)
    mlcat <- do.call("[", c(list(x = mlcat), indices))
    if (!is.null(mask)){
      mask <- mask[nlat:1, ]
    }
  }

  #----------------------
  # Set layout and parameters
  #----------------------
  plot.new()
  par(font.main = 1)
  layout(matrix(c(rep(ncat+1,ncat),1:ncat),2,ncat,byrow=T), height = c(6, 1.5))


 #----------------------
  # Add colorbars 
  #----------------------
  for (k in 1:ncat){
    ColorBar(brks = brks, cols = cols[[k]], vertical = FALSE, 
             draw_separators = TRUE, extra_margin = c(2, 0, 2, 0), 
             label_scale = legend_scale * 1.5)
    mtext(bar_titles[[k]], 3, line =1, cex = 1.5)
  }


  #----------------------
  # Set colors and breaks and then PlotEquiMap
  #----------------------
  tcols <- c(col_unknown_cat, cols[[1]])
  for (k in 2:ncat) {
    tcols <- append(tcols, c(col_unknown_cat, cols[[k]]))
  }
  tbrks <- c(1, brks / 100 + rep(1:ncat, each=length(brks)))
  PlotEquiMap(var = mlcat, lon = lon, lat = lat, 
              brks = tbrks, cols = tcols, drawleg = FALSE, 
              filled.continents = FALSE, ...)
  
  #----------------------
  # Add overplot on top
  #----------------------
  if (!is.null(mask)) {
    cols_mask <- sapply(seq(from = 0, to = 1, length.out = 10), 
                        function(x) adjustcolor(col_mask, alpha.f = x))
    image(lon, lat, t(mask), axes = FALSE, col = cols_mask, 
          breaks = seq(from = 0, to = 1, by = 0.1), 
          xlab='', ylab='', add = TRUE)
    if (!exists('coast_color')) {
      coast_color <- 'black'
    }
    if (min(lon) < 0) {
      map('world', interior = FALSE, add = TRUE, lwd = 1, col = coast_color) # Low resolution world map (lon -180 to 180).
    } else {
      map('world2', interior = FALSE, add = TRUE, lwd = 1, col = coast_color) # Low resolution world map (lon 0 to 360).
    }
  }

}