Utils.R 78.4 KB
Newer Older

.LoadSampleData <- function(var, exp = NULL, obs = NULL, sdates, 
                            nmember = NULL, nmemberobs = NULL, 
                            nleadtime = NULL, leadtimemin = 1, 
                            leadtimemax = NULL, storefreq = 'monthly', 
                            sampleperiod = 1, lonmin = 0, lonmax = 360, 
                            latmin = -90, latmax = 90, output = 'areave', 
                            method = 'conservative', grid = NULL, 
                            maskmod = vector("list", 15), 
                            maskobs = vector("list", 15), 
                            configfile = NULL, suffixexp = NULL, 
                            suffixobs = NULL, varmin = NULL, varmax = NULL, 
                            silent = FALSE, nprocs = NULL) {
  ## This function loads and selects sample data stored in sampleMap and 
  ## sampleTimeSeries and is used in the examples instead of Load() so as
  ## to avoid nco and cdo system calls and computation time in the stage 
  ## of running examples in the CHECK process on CRAN.
  selected_start_dates <- match(sdates, c('19851101', '19901101', '19951101', 
                                          '20001101', '20051101'))
  start_dates_position <- 3
  lead_times_position <- 4

  if (output == 'lonlat') {
    sampleData <- s2dv::sampleMap
    if (is.null(leadtimemax)) {
      leadtimemax <- dim(sampleData$mod)[lead_times_position]
    }
    selected_lead_times <- leadtimemin:leadtimemax

    dataOut <- sampleData
    dataOut$mod <- sampleData$mod[, , selected_start_dates, selected_lead_times, , ]
    dataOut$obs <- sampleData$obs[, , selected_start_dates, selected_lead_times, , ]
aho's avatar
aho committed
  } else if (output == 'areave') {
    sampleData <- s2dv::sampleTimeSeries
    if (is.null(leadtimemax)) {
      leadtimemax <- dim(sampleData$mod)[lead_times_position]
    }
    selected_lead_times <- leadtimemin:leadtimemax

    dataOut <- sampleData
    dataOut$mod <- sampleData$mod[, , selected_start_dates, selected_lead_times]
    dataOut$obs <- sampleData$obs[, , selected_start_dates, selected_lead_times]
  }

  dims_out <- dim(sampleData$mod)
  dims_out[start_dates_position] <- length(selected_start_dates)
  dims_out[lead_times_position] <- length(selected_lead_times)
  dim(dataOut$mod) <- dims_out

  dims_out <- dim(sampleData$obs)
  dims_out[start_dates_position] <- length(selected_start_dates)
  dims_out[lead_times_position] <- length(selected_lead_times)
  dim(dataOut$obs) <- dims_out

  invisible(list(mod = dataOut$mod, obs = dataOut$obs, 
                 lat = dataOut$lat, lon = dataOut$lon))
}

.ConfigGetDatasetInfo <- function(matching_entries, table_name) {
  # This function obtains the information of a dataset and variable pair,
  # applying all the entries that match in the configuration file.
  if (table_name == 'experiments') {
    id <- 'EXP'
  } else {
    id <- 'OBS'
  }
aho's avatar
aho committed
  defaults <- c(paste0('$DEFAULT_', id, '_MAIN_PATH$'), 
                paste0('$DEFAULT_', id, '_FILE_PATH$'), 
                '$DEFAULT_NC_VAR_NAME$', '$DEFAULT_SUFFIX$', 
                '$DEFAULT_VAR_MIN$', '$DEFAULT_VAR_MAX$')
  info <- NULL

  for (entry in matching_entries) {
    if (is.null(info)) {
      info <- entry[-1:-2]
      info[which(info == '*')] <- defaults[which(info == '*')]
    } else {
      info[which(entry[-1:-2] != '*')] <- entry[-1:-2][which(entry[-1:-2] != '*')]
    }
  }

  info <- as.list(info)
  names(info) <- c('main_path', 'file_path', 'nc_var_name', 'suffix', 'var_min', 'var_max')  
  info
}

.ReplaceGlobExpressions <- function(path_with_globs, actual_path, 
                                    replace_values, tags_to_keep, 
                                    dataset_name, permissive) {
  # The goal of this function is to replace the shell globbing expressions in
  # a path pattern (that may contain shell globbing expressions and Load() 
  # tags) by the corresponding part of the real existing path.
  # What is done actually is to replace all the values of the tags in the 
  # actual path by the corresponding $TAG$
  #
  # It takes mainly two inputs. The path with expressions and tags, e.g.:
  #   /data/experiments/*/$EXP_NAME$/$VAR_NAME$/$VAR_NAME$_*$START_DATE$*.nc
  # and a complete known path to one of the matching files, e.g.:
  #   /data/experiments/ecearth/i00k/tos/tos_fc0-1_19901101_199011-199110.nc
  # and it returns the path pattern but without shell globbing expressions:
  #   /data/experiments/ecearth/$EXP_NAME$/$VAR_NAME$/$VAR_NAME$_fc0-1_$START_DATE$_199011-199110.nc
  #
  # To do that, it needs also as inputs the list of replace values (the 
  # association of each tag to their value).
  #
  # All the tags not present in the parameter tags_to_keep will be repalced.
  #
  # Not all cases can be resolved with the implemented algorithm. In an
  # unsolvable case a warning is given and one possible guess is returned.
  #
  # In some cases it is interesting to replace only the expressions in the
  # path to the file, but not the ones in the file name itself. To keep the
  # expressions in the file name, the parameter permissive can be set to 
  # TRUE. To replace all the expressions it can be set to FALSE.
  clean <- function(x) {
    if (nchar(x) > 0) {
      x <- gsub('\\\\', '', x)
      x <- gsub('\\^', '', x)
      x <- gsub('\\$', '', x)
aho's avatar
aho committed
      x <- unname(sapply(strsplit(x, '[', fixed = TRUE)[[1]], function(y) gsub('.*]', '.', y)))
      do.call(paste0, as.list(x))
    } else {
      x
    }
  }

  strReverse <- function(x) sapply(lapply(strsplit(x, NULL), rev), paste, collapse = "")

  if (permissive) {
    actual_path_chunks <- strsplit(actual_path, '/')[[1]]
    actual_path <- paste(actual_path_chunks[-length(actual_path_chunks)], collapse = '/')
    file_name <- tail(actual_path_chunks, 1)
    if (length(actual_path_chunks) > 1) {
      file_name <- paste0('/', file_name)
    }
    path_with_globs_chunks <- strsplit(path_with_globs, '/')[[1]]
    path_with_globs <- paste(path_with_globs_chunks[-length(path_with_globs_chunks)], 
                             collapse = '/')
    path_with_globs <- .ConfigReplaceVariablesInString(path_with_globs, replace_values)
    file_name_with_globs <- tail(path_with_globs_chunks, 1)
    if (length(path_with_globs_chunks) > 1) {
      file_name_with_globs <- paste0('/', file_name_with_globs)
    }
    right_known <- head(strsplit(file_name_with_globs, '*', fixed = TRUE)[[1]], 1)
    right_known_no_tags <- .ConfigReplaceVariablesInString(right_known, replace_values)
    path_with_globs_rx <- utils::glob2rx(paste0(path_with_globs, right_known_no_tags))
aho's avatar
aho committed
    match <- regexpr(gsub('$', '', path_with_globs_rx, fixed = TRUE), 
                     paste0(actual_path, file_name))
aho's avatar
aho committed
      stop("Incorrect parameters to replace glob expressions. ",
           "The path with expressions does not match the actual path.")
    }
    if (attr(match, 'match.length') - nchar(right_known_no_tags) < nchar(actual_path)) {
      path_with_globs <- paste0(path_with_globs, right_known_no_tags, '*')
      file_name_with_globs <- sub(right_known, '/*', file_name_with_globs)
    } 
  }
  path_with_globs_rx <- utils::glob2rx(path_with_globs)
aho's avatar
aho committed
  values_to_replace <- NULL
  tags_to_replace_starts <- NULL
  tags_to_replace_ends <- NULL
  give_warning <- FALSE
  for (tag in tags_to_keep) {
    matches <- gregexpr(paste0('$', tag, '$'), path_with_globs_rx, fixed = TRUE)[[1]]
    lengths <- attr(matches, 'match.length')
    if (!(length(matches) == 1 && matches[1] == -1)) {
aho's avatar
aho committed
      for (i in seq_along(matches)) {
aho's avatar
aho committed
          left <- 
            .ConfigReplaceVariablesInString(substr(path_with_globs_rx, 1, 
                                                   matches[i] - 1), replace_values)
          left_known <- 
            strReverse(head(strsplit(strReverse(left), 
                            strReverse('.*'), fixed = TRUE)[[1]], 1))
        }
        right <- NULL
        if ((matches[i] + lengths[i] - 1) < nchar(path_with_globs_rx)) {
aho's avatar
aho committed
          right <- 
            .ConfigReplaceVariablesInString(substr(path_with_globs_rx, 
                                                   matches[i] + lengths[i], 
                                                   nchar(path_with_globs_rx)),
                                                   replace_values)
          right_known <- head(strsplit(right, '.*', fixed = TRUE)[[1]], 1)
        }
        match_limits <- NULL
        if (!is.null(left)) {
          left_match <- regexpr(paste0(left, replace_values[[tag]], right_known), actual_path)
          match_len <- attr(left_match, 'match.length')
aho's avatar
aho committed
          left_match_limits <- 
            c(left_match + match_len - 1 - nchar(clean(right_known)) - 
                nchar(replace_values[[tag]]) + 1,
              left_match + match_len - 1 - nchar(clean(right_known)))
          if (!(left_match < 1)) {
            match_limits <- left_match_limits
          }
        }
        right_match <- NULL
        if (!is.null(right)) {
          right_match <- regexpr(paste0(left_known, replace_values[[tag]], right), actual_path)
          match_len <- attr(right_match, 'match.length')
aho's avatar
aho committed
          right_match_limits <- 
            c(right_match + nchar(clean(left_known)),
              right_match + nchar(clean(left_known)) + 
                nchar(replace_values[[tag]]) - 1)
          if (is.null(match_limits) && !(right_match < 1)) {
            match_limits <- right_match_limits
          }
        }
        if (!is.null(right_match) && !is.null(left_match)) {
          if (!identical(right_match_limits, left_match_limits)) {
            give_warning <- TRUE
          }
        }
        if (is.null(match_limits)) {
          stop("Too complex path pattern specified for ", dataset_name,
               ". Specify a simpler path pattern for this dataset.")
        }
        values_to_replace <- c(values_to_replace, tag)
        tags_to_replace_starts <- c(tags_to_replace_starts, match_limits[1])
        tags_to_replace_ends <- c(tags_to_replace_ends, match_limits[2])
      }
    }
  }

  if (length(tags_to_replace_starts) > 0) {
    reorder <- sort(tags_to_replace_starts, index.return = TRUE)
    tags_to_replace_starts <- reorder$x
    values_to_replace <- values_to_replace[reorder$ix]
    tags_to_replace_ends <- tags_to_replace_ends[reorder$ix]
    while (length(values_to_replace) > 0) {
      actual_path <- paste0(substr(actual_path, 1, head(tags_to_replace_starts, 1) - 1),
                           '$', head(values_to_replace, 1), '$',
aho's avatar
aho committed
                           substr(actual_path, head(tags_to_replace_ends, 1) + 1, 
                                  nchar(actual_path)))
      extra_chars <- nchar(head(values_to_replace, 1)) + 2 - 
                           (head(tags_to_replace_ends, 1) - 
                              head(tags_to_replace_starts, 1) + 1)
      values_to_replace <- values_to_replace[-1]
      tags_to_replace_starts <- tags_to_replace_starts[-1]
      tags_to_replace_ends <- tags_to_replace_ends[-1]
      tags_to_replace_starts <- tags_to_replace_starts + extra_chars
      tags_to_replace_ends <- tags_to_replace_ends + extra_chars
    }
  }

  if (give_warning) {
    .warning(paste0("Too complex path pattern specified for ", dataset_name, 
aho's avatar
aho committed
                    ". Double check carefully the '$Files' fetched for this dataset ",
                    "or specify a simpler path pattern."))
  }

  if (permissive) {
    paste0(actual_path, file_name_with_globs)
  } else {
    actual_path
  }
}

.FindTagValue <- function(path_with_globs_and_tag, actual_path, tag) {
  tag <- paste0('\\$', tag, '\\$')
  path_with_globs_and_tag <- paste0('^', path_with_globs_and_tag, '$')
  parts <- strsplit(path_with_globs_and_tag, '*', fixed = TRUE)[[1]]
  parts <- as.list(parts[grep(tag, parts)])
aho's avatar
aho committed
  longest_couples <- NULL
  pos_longest_couples <- NULL
aho's avatar
aho committed
  for (i in seq_along(parts)) {
    parts[[i]] <- strsplit(parts[[i]], tag)[[1]]
    if (length(parts[[i]]) == 1) {
      parts[[i]] <- c(parts[[i]], '')
    }
    len_parts <- sapply(parts[[i]], nchar)
    len_couples <- len_parts[-length(len_parts)] + len_parts[2:length(len_parts)]
    pos_longest_couples <- c(pos_longest_couples, which.max(len_couples))
    longest_couples <- c(longest_couples, max(len_couples))
  }
  chosen_part <- which.max(longest_couples)
aho's avatar
aho committed
  parts[[chosen_part]] <- 
    parts[[chosen_part]][pos_longest_couples[chosen_part]:(pos_longest_couples[chosen_part] + 1)]
  if (nchar(parts[[chosen_part]][1]) >= nchar(parts[[chosen_part]][2])) {
    if (nchar(parts[[chosen_part]][1]) > 0) {
      matches <- gregexpr(parts[[chosen_part]][1], actual_path)[[1]]
      if (length(matches) == 1) {
        match_left <- matches
aho's avatar
aho committed
        actual_path <- 
          substr(actual_path, match_left + attr(match_left, 'match.length'), nchar(actual_path))
      }
    }
    if (nchar(parts[[chosen_part]][2]) > 0) {
      matches <- gregexpr(parts[[chosen_part]][2], actual_path)[[1]]
      if (length(matches) == 1) {
        match_right <- matches
        found_value <- substr(actual_path, 0, match_right - 1)
      }
    }
  } else {
    if (nchar(parts[[chosen_part]][2]) > 0) {
      matches <- gregexpr(parts[[chosen_part]][2], actual_path)[[1]]
      if (length(matches) == 1) {
        match_right <- matches
        actual_path <- substr(actual_path, 0, match_right - 1)
      }
    }
    if (nchar(parts[[chosen_part]][1]) > 0) {
      matches <- gregexpr(parts[[chosen_part]][1], actual_path)[[1]]
      if (length(matches) == 1) {
        match_left <- matches
aho's avatar
aho committed
        found_value <- 
          substr(actual_path, match_left + attr(match_left, 'match.length'), 
                 nchar(actual_path))
      }
    }
  }
  found_value
}

.FilterUserGraphicArgs <- function(excludedArgs, ...) {
  # This function filter the extra graphical parameters passed by the user in
  # a plot function, excluding the ones that the plot function uses by default.
  # Each plot function has a different set of arguments that are not allowed to
  # be modified.
  args <- list(...)
  userArgs <- list()
  for (name in names(args)) {
      if ((name != "") & !is.element(name, excludedArgs)) {
          # If the argument has a name and it is not in the list of excluded
          # arguments, then it is added to the list that will be used
          userArgs[[name]] <- args[[name]]
      } else {
        .warning(paste0("the argument '", name, "' can not be 
        modified and the new value will be ignored"))
      }
  }
  userArgs
}

.SelectDevice <- function(fileout, width, height, units, res) {
  # This function is used in the plot functions to check the extension of the 
  # files where the graphics will be stored and select the right R device to 
  # save them.
  # If the vector of filenames ('fileout') has files with different 
  # extensions, then it will only accept the first one, changing all the rest 
  # of the filenames to use that extension.

  # We extract the extension of the filenames: '.png', '.pdf', ...
  ext <- regmatches(fileout, regexpr("\\.[a-zA-Z0-9]*$", fileout))

  if (length(ext) != 0) {
    # If there is an extension specified, select the correct device
    ## units of width and height set to accept inches
    if (ext[1] == ".png") {
      saveToFile <- function(fileout) {
        png(filename = fileout, width = width, height = height, res = res, units = units)
      }
    } else if (ext[1] == ".jpeg") {
      saveToFile <- function(fileout) {
        jpeg(filename = fileout, width = width, height = height, res = res, units = units)
      }
    } else if (ext[1] %in% c(".eps", ".ps")) {
      saveToFile <- function(fileout) {
        postscript(file = fileout, width = width, height = height)
      }
    } else if (ext[1] == ".pdf") {
      saveToFile <- function(fileout) {
        pdf(file = fileout, width = width, height = height)
      }
    } else if (ext[1] == ".svg") {
      saveToFile <- function(fileout) {
        svg(filename = fileout, width = width, height = height)
      }
    } else if (ext[1] == ".bmp") {
      saveToFile <- function(fileout) {
        bmp(filename = fileout, width = width, height = height, res = res, units = units)
      }
    } else if (ext[1] == ".tiff") {
      saveToFile <- function(fileout) {
        tiff(filename = fileout, width = width, height = height, res = res, units = units)
      }
    } else {
      .warning("file extension not supported, it will be used '.eps' by default.")
      ## In case there is only one filename
      fileout[1] <- sub("\\.[a-zA-Z0-9]*$", ".eps", fileout[1])
      ext[1] <- ".eps"
      saveToFile <- function(fileout) {
        postscript(file = fileout, width = width, height = height)
      }
    }
    # Change filenames when necessary
    if (any(ext != ext[1])) {
aho's avatar
aho committed
      .warning(paste0("some extensions of the filenames provided in 'fileout' ",
                      "are not ", ext[1],
                      ". The extensions are being converted to ", ext[1], "."))
      fileout <- sub("\\.[a-zA-Z0-9]*$", ext[1], fileout)
    }
  } else {
    # Default filenames when there is no specification
    .warning("there are no extensions specified in the filenames, default to '.eps'")
    fileout <- paste0(fileout, ".eps")
    saveToFile <- postscript
  }

  # return the correct function with the graphical device, and the correct 
  # filenames
  list(fun = saveToFile, files = fileout)
}

.message <- function(...) {
  # Function to use the 'message' R function with our custom settings
  # Default: new line at end of message, indent to 0, exdent to 3, 
  #  collapse to \n*
  args <- list(...)

  ## In case we need to specify message arguments
  if (!is.null(args[["appendLF"]])) {
    appendLF <- args[["appendLF"]]
  } else {
    ## Default value in message function
    appendLF <- TRUE
  } 
  if (!is.null(args[["domain"]])) {
    domain <- args[["domain"]]
  } else {
    ## Default value in message function
    domain <- NULL
  }
  args[["appendLF"]] <- NULL
  args[["domain"]] <- NULL

  ## To modify strwrap indent and exdent arguments
  if (!is.null(args[["indent"]])) {
    indent <- args[["indent"]]
  } else {
    indent <- 0
  }
  if (!is.null(args[["exdent"]])) {
    exdent <- args[["exdent"]]
  } else {
    exdent <- 3
  }
  args[["indent"]] <- NULL
  args[["exdent"]] <- NULL

  ## To modify paste collapse argument
  if (!is.null(args[["collapse"]])) {
    collapse <- args[["collapse"]]
  } else {
    collapse <- "\n*"
  }
  args[["collapse"]] <- NULL

  ## Message tag
  if (!is.null(args[["tag"]])) {
    tag <- args[["tag"]]
  } else {
    tag <- "* "
  }
  args[["tag"]] <- NULL

aho's avatar
aho committed
  tmp <- paste0(tag, 
                paste(strwrap(args, indent = indent, exdent = exdent), collapse = collapse))
  message(tmp, appendLF = appendLF, domain = domain)
}

.warning <- function(...) {
  # Function to use the 'warning' R function with our custom settings
  # Default: no call information, indent to 0, exdent to 3, 
  #  collapse to \n
  args <- list(...)

  ## In case we need to specify warning arguments
  if (!is.null(args[["call."]])) {
    call <- args[["call."]]
  } else {
    ## Default: don't show info about the call where the warning came up
    call <- FALSE
  }
  if (!is.null(args[["immediate."]])) {
    immediate <- args[["immediate."]]
  } else {
    ## Default value in warning function
    immediate <- FALSE
  }
  if (!is.null(args[["noBreaks."]])) {
    noBreaks <- args[["noBreaks."]]
  } else {
    ## Default value warning function
    noBreaks <- FALSE
  }
  if (!is.null(args[["domain"]])) {
    domain <- args[["domain"]]
  } else {
    ## Default value warning function
    domain <- NULL
  }
  args[["call."]] <- NULL
  args[["immediate."]] <- NULL
  args[["noBreaks."]] <- NULL
  args[["domain"]] <- NULL

  ## To modify strwrap indent and exdent arguments
  if (!is.null(args[["indent"]])) {
    indent <- args[["indent"]]
  } else {
    indent <- 0
  }
  if (!is.null(args[["exdent"]])) {
    exdent <- args[["exdent"]]
  } else {
    exdent <- 3
  }
  args[["indent"]] <- NULL
  args[["exdent"]] <- NULL

  ## To modify paste collapse argument
  if (!is.null(args[["collapse"]])) {
    collapse <- args[["collapse"]]
  } else {
    collapse <- "\n!"
  }
  args[["collapse"]] <- NULL

  ## Warning tag
  if (!is.null(args[["tag"]])) {
    tag <- args[["tag"]]
  } else {
    tag <- "! Warning: "
  }
  args[["tag"]] <- NULL

aho's avatar
aho committed
  tmp <- paste0(tag, 
                paste(strwrap(args, indent = indent, exdent = exdent), collapse = collapse))
  warning(tmp, call. = call, immediate. = immediate, 
          noBreaks. = noBreaks, domain = domain)
}

.IsColor <- function(x) {
  res <- try(col2rgb(x), silent = TRUE)
aho's avatar
aho committed
  return(!is(res, "try-error"))
}

# This function switches to a specified figure at position (row, col) in a layout.
# This overcomes the bug in par(mfg = ...). However the mode par(new = TRUE) is 
# activated, i.e., all drawn elements will be superimposed. Additionally, after 
# using this function, the automatical pointing to the next figure in the layout
# will be spoiled: once the last figure in the layout is drawn, the pointer won't 
# move to the first figure in the layout.
# Only figures with numbers other than 0 (when creating the layout) will be
# accessible.
# Inputs: either row and col, or n and mat
.SwitchToFigure <- function(row = NULL, col = NULL, n = NULL, mat = NULL) {
  if (!is.null(n) && !is.null(mat)) {
    if (!is.numeric(n) || length(n) != 1) {
      stop("Parameter 'n' must be a single numeric value.")
    }
    n <- round(n)
    if (!is.array(mat)) {
      stop("Parameter 'mat' must be an array.")
    }
    target <- which(mat == n, arr.ind = TRUE)[1, ]
    row <- target[1]
    col <- target[2]
  } else if (!is.null(row) && !is.null(col)) {
    if (!is.numeric(row) || length(row) != 1) {
      stop("Parameter 'row' must be a single numeric value.")
    }
    row <- round(row)
    if (!is.numeric(col) || length(col) != 1) {
      stop("Parameter 'col' must be a single numeric value.")
    }
    col <- round(col)
  } else {
    stop("Either 'row' and 'col' or 'n' and 'mat' must be provided.")
  }
  next_attempt <- c(row, col)
  par(mfg = next_attempt)
  i <- 1
  layout_size <- par('mfrow')
  layout_cells <- matrix(1:prod(layout_size), layout_size[1], layout_size[2], 
                         byrow = TRUE)
  while (any((par('mfg')[1:2] != c(row, col)))) {
    next_attempt <- which(layout_cells == i, arr.ind = TRUE)[1, ]
    par(mfg = next_attempt)
    i <- i + 1
    if (i > prod(layout_size)) {
      stop("Figure not accessible.")
    }
  }
  plot(0, type = 'n', axes = FALSE, ann = FALSE)
  par(mfg = next_attempt)
}

# Function to permute arrays of non-atomic elements (e.g. POSIXct)
.aperm2 <- function(x, new_order) {
  old_dims <- dim(x)
  attr_bk <- attributes(x)
  if ('dim' %in% names(attr_bk)) {
    attr_bk[['dim']] <- NULL
  }
  if (is.numeric(x)) {
    x <- aperm(x, new_order)
  } else {
aho's avatar
aho committed
    y <- array(seq_along(x), dim = dim(x))
    y <- aperm(y, new_order)
    x <- x[as.vector(y)]
  }
  dim(x) <- old_dims[new_order]
  attributes(x) <- c(attributes(x), attr_bk)
  x
}

# This function is a helper for the function .MergeArrays.
# It expects as inputs two named numeric vectors, and it extends them
# with dimensions of length 1 until an ordered common dimension
# format is reached.
# The first output is dims1 extended with 1s.
# The second output is dims2 extended with 1s.
# The third output is a merged dimension vector. If dimensions with
# the same name are found in the two inputs, and they have a different
# length, the maximum is taken.
.MergeArrayDims <- function(dims1, dims2) {
aho's avatar
aho committed
  new_dims1 <- NULL
  new_dims2 <- NULL
  while (length(dims1) > 0) {
    if (names(dims1)[1] %in% names(dims2)) {
      pos <- which(names(dims2) == names(dims1)[1])
      dims_to_add <- rep(1, pos - 1)
      if (length(dims_to_add) > 0) {
        names(dims_to_add) <- names(dims2[1:(pos - 1)])
      }
      new_dims1 <- c(new_dims1, dims_to_add, dims1[1])
      new_dims2 <- c(new_dims2, dims2[1:pos])
      dims1 <- dims1[-1]
aho's avatar
aho committed
      dims2 <- dims2[-(1:pos)]
    } else {
      new_dims1 <- c(new_dims1, dims1[1])
      new_dims2 <- c(new_dims2, 1)
      names(new_dims2)[length(new_dims2)] <- names(dims1)[1]
      dims1 <- dims1[-1]
    }
  }
  if (length(dims2) > 0) {
    dims_to_add <- rep(1, length(dims2))
    names(dims_to_add) <- names(dims2)
    new_dims1 <- c(new_dims1, dims_to_add)
    new_dims2 <- c(new_dims2, dims2)
  }
  list(new_dims1, new_dims2, pmax(new_dims1, new_dims2))
}

# This function takes two named arrays and merges them, filling with
# NA where needed.
# dim(array1)
#          'b'   'c'         'e'   'f'
#           1     3           7     9
# dim(array2)
#    'a'   'b'         'd'         'f'   'g'
#     2     3           5           9     11
# dim(.MergeArrays(array1, array2, 'b'))
#    'a'   'b'   'c'   'e'   'd'   'f'   'g'
#     2     4     3     7     5     9     11
.MergeArrays <- function(array1, array2, along) {
  if (!(is.null(array1) || is.null(array2))) {
    if (!(identical(names(dim(array1)), names(dim(array2))) &&
        identical(dim(array1)[-which(names(dim(array1)) == along)],
                  dim(array2)[-which(names(dim(array2)) == along)]))) {
      new_dims <- .MergeArrayDims(dim(array1), dim(array2))
      dim(array1) <- new_dims[[1]]
      dim(array2) <- new_dims[[2]]
aho's avatar
aho committed
      for (j in seq_along(dim(array1))) {
        if (names(dim(array1))[j] != along) {
          if (dim(array1)[j] != dim(array2)[j]) {
            if (which.max(c(dim(array1)[j], dim(array2)[j])) == 1) {
              na_array_dims <- dim(array2)
              na_array_dims[j] <- dim(array1)[j] - dim(array2)[j]
              na_array <- array(dim = na_array_dims)
              array2 <- abind(array2, na_array, along = j)
              names(dim(array2)) <- names(na_array_dims)
            } else {
              na_array_dims <- dim(array1)
              na_array_dims[j] <- dim(array2)[j] - dim(array1)[j]
              na_array <- array(dim = na_array_dims)
              array1 <- abind(array1, na_array, along = j)
              names(dim(array1)) <- names(na_array_dims)
            }
          }
        }
      }
    }
    if (!(along %in% names(dim(array2)))) {
      stop("The dimension specified in 'along' is not present in the ",
           "provided arrays.")
    }
    array1 <- abind(array1, array2, along = which(names(dim(array1)) == along))
    names(dim(array1)) <- names(dim(array2))
  } else if (is.null(array1)) {
    array1 <- array2
  }
  array1
}

# only can be used in Trend(). Needs generalization or be replaced by other function.
.reorder <- function(output, time_dim, dim_names) {
  # Add dim name back
  if (is.null(dim(output))) {
    dim(output) <- c(stats = length(output))
  } else {  #is an array
    if (length(dim(output)) == 1) {
      if (!is.null(names(dim(output)))) {
        dim(output) <- c(1, dim(output))
        names(dim(output))[1] <- time_dim
      } else {
        names(dim(output)) <- time_dim
      }
    } else {  # more than one dim
      if (names(dim(output))[1] != "") {
        dim(output) <- c(1, dim(output))
        names(dim(output))[1] <- time_dim
      } else {   #regular case
        names(dim(output))[1] <- time_dim
      }
    }
  }
  # reorder
  pos <- match(dim_names, names(dim(output)))
  output <- aperm(output, pos)
  names(dim(output)) <- dim_names
  names(dim(output))[names(dim(output)) == time_dim] <- 'stats'
  return(output)
}

# to be used in AMV.R, TPI.R, SPOD.R, GSAT.R and GMST.R
.Indices <- function(data, type, monini, indices_for_clim, 
                     fmonth_dim, sdate_dim, year_dim, month_dim, na.rm) {
    data <- Season(data = data, time_dim = fmonth_dim,
                   monini = monini, moninf = 1, monsup = 12,
                   method = mean, na.rm = na.rm)
aho's avatar
aho committed
    names(dim(data))[which(names(dim(data)) == fmonth_dim)] <- fyear_dim
    if (identical(indices_for_clim, FALSE)) { ## data is already anomalies
    } else { ## Different indices_for_clim for each forecast year (to use the same calendar years)
      n_fyears <- as.numeric(dim(data)[fyear_dim])
      n_sdates <- as.numeric(dim(data)[sdate_dim])
      
      if (is.null(indices_for_clim)) { ## climatology over the whole (common) period
        first_years_for_clim <- n_fyears : 1
        last_years_for_clim <- n_sdates : (n_sdates - n_fyears + 1)
      } else { ## indices_for_clim specified as a numeric vector
        first_years_for_clim <- seq(from = indices_for_clim[1], by = -1, length.out = n_fyears)
aho's avatar
aho committed
        last_years_for_clim <- 
          seq(from = indices_for_clim[length(indices_for_clim)], 
              by = -1, length.out = n_fyears)
      
      data <- s2dv::Reorder(data = data, order = c(fyear_dim, sdate_dim))
      anom <- array(data = NA, dim = dim(data))
      for (i in 1:n_fyears) {
aho's avatar
aho committed
        clim <- mean(data[i, first_years_for_clim[i]:last_years_for_clim[i]], na.rm = na.rm)
        anom[i, ] <- data[i, ] - clim
aho's avatar
aho committed
  } else if (type %in% c('obs', 'hist')) {
aho's avatar
aho committed
    data <- multiApply::Apply(data = data, target_dims = month_dim, 
                              fun = mean, na.rm = na.rm)$output1
    if (identical(indices_for_clim, FALSE)) { ## data is already anomalies
      clim <- 0
aho's avatar
aho committed
    } else if (is.null(indices_for_clim)) { 
      ## climatology over the whole period
      clim <- multiApply::Apply(data = data, target_dims = year_dim, fun = mean, na.rm = na.rm)$output1
aho's avatar
aho committed
    } else { 
      ## indices_for_clim specified as a numeric vector
      clim <- multiApply::Apply(data = ClimProjDiags::Subset(x = data, along = year_dim, indices = indices_for_clim),
                                target_dims = year_dim, fun = mean, na.rm = na.rm)$output1
aho's avatar
aho committed
  } else {
    stop('type must be dcpp, hist or obs')
  }

#TODO: Remove from s2dv when PlotLayout can get colorbar info from plotting function directly. 
#      The function is temporarily here because PlotLayout() needs to draw the colorbars of
#      PlotMostLikelyQuantileMap().
#Draws Color Bars for Categories
#A wrapper of s2dv::ColorBar to generate multiple color bars for different 
#categories, and each category has different color set.
GradientCatsColorBar <- function(nmap, brks = NULL, cols = NULL, vertical = TRUE, subsampleg = NULL,
                                 bar_limits, var_limits = NULL,
                                 triangle_ends = NULL, plot = TRUE,
                                 draw_separators = FALSE,
aho's avatar
aho committed
                                 bar_titles = NULL, title_scale = 1, 
                                 label_scale = 1, extra_margin = rep(0, 4),
  # bar_limits
  if (!is.numeric(bar_limits) || length(bar_limits) != 2) {
    stop("Parameter 'bar_limits' must be a numeric vector of length 2.")
  }

  # Check brks
  if (is.null(brks) || (is.numeric(brks) && length(brks) == 1)) {
    num_brks <- 5
    if (is.numeric(brks)) {
      num_brks <- brks
    }
    brks <- seq(from = bar_limits[1], to = bar_limits[2], length.out = num_brks)
  }
  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) >= nmap) {
      chosen_sets <- 1:nmap
      chosen_sets <- chosen_sets + floor((length(col_sets) - length(chosen_sets)) / 2)
    } else {
aho's avatar
aho committed
      chosen_sets <- array(seq_along(col_sets), nmap)
    }
    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) != nmap) {
      stop("Parameter 'cols' must be a list of the same length as the number of ",
           "maps in 'maps'.")
    }
  }
aho's avatar
aho committed
  for (i in seq_along(cols)) {
    if (length(cols[[i]]) != (length(brks) - 1)) {
      cols[[i]] <- grDevices::colorRampPalette(cols[[i]])(length(brks) - 1)
    }
  }

  # Check bar_titles
  if (is.null(bar_titles)) {
    if (nmap == 3) {
      bar_titles <- c("Below normal (%)", "Normal (%)", "Above normal (%)")
    } else if (nmap == 5) {
      bar_titles <- c("Low (%)", "Below normal (%)",
                         "Normal (%)", "Above normal (%)", "High (%)")
    } else {
      bar_titles <- paste0("Cat. ", 1:nmap, " (%)")
    }
  }

  if (plot) {
    for (k in 1:nmap) {
      s2dv::ColorBar(brks = brks, cols = cols[[k]], vertical = FALSE, subsampleg = subsampleg,
#                     bar_limits = bar_limits, var_limits = var_limits,
                     triangle_ends = triangle_ends, plot = TRUE,
                     draw_separators = draw_separators,
                     title = bar_titles[[k]], title_scale = title_scale,
                     label_scale = label_scale, extra_margin = extra_margin)
    }
  } else {
    #TODO: col_inf and col_sup
    return(list(brks = brks, cols = cols))
  }

}