R/permute.R

In edibble: Encapsulating Elements of Experimental Design

#' @export
order_trts.blocksdesign <- function(x, trts_table, units_table, unit, constrain, prov, model = NULL, ...) {
  if(!requireNamespace("blocksdesign")) abort("You need to install `blocksdesign` package.")
  names(trts_table) <- paste0("T", 1:ncol(trts_table))
  # drop any treatment factors that appear only once
  tcount <- sapply(trts_table, function(x) length(unique(x)))
  tlist <- lapply(trts_table[tcount > 1], as.factor)
  if(length(tlist)==0) return(rep(1, nrow(units_table)))
  # drop any unit factors that have completely distinct labels
  ucount <- sapply(units_table, function(x) length(unique(x)))
  udf <- data.frame(lapply(units_table[ucount < nrow(units_table)], as.factor))
  res <- blocksdesign::design(treatments = tlist, blocks = udf, treatments_model = model,
                              weighting = 0.5)
  trts_table$..trtid.. <- 1:nrow(trts_table)
  out <- res$Design
  out$..id.. <- 1:nrow(out)
  out <- merge(out, trts_table)
  out[order(out$..id..), "..trtid..", drop = TRUE]
}



permute_parent_one <- function(vid, udf, ntrts) {
  blocksizes <- as.data.frame(table(table(udf[[as.character(vid)]])))
  blocksizes$size <- as.numeric(as.character(blocksizes$Var1))

  for(isize in seq(nrow(blocksizes))) {
    if(blocksizes$size[isize] <= ntrts) {
      comb <- utils::combn(ntrts, blocksizes$size[isize])
      blocksizes$rows[isize] <- list(comb)
    } else {
      nrep <- floor(blocksizes$size[isize] / ntrts)
      nremain <- blocksizes$size[isize] %% ntrts
      comb <- utils::combn(ntrts, nremain)
      blocksizes$rows[isize] <- list(rbind(comb,
                                           matrix(rep(1:ntrts, nrep * ncol(comb)), ncol = ncol(comb))))
    }
    blocksizes$select[isize] <- list(sample(rep(sample(ncol(comb)), length.out = blocksizes$Freq[isize])))
  }
  blocksizes$wselect <- blocksizes$select
  gpar_tab <- as.data.frame(table(udf[[as.character(vid)]]))
  out <- vector("integer", length = nrow(udf))
  for(ianc in seq(nrow(gpar_tab))) {
    imatch <- which(blocksizes$size == gpar_tab$Freq[ianc])
    iselect <- blocksizes$wselect[imatch][[1]][1]
    blocksizes$wselect[imatch] <- list(blocksizes$wselect[imatch][[1]][-1])
    out[as.character(udf[[as.character(vid)]])==as.character(gpar_tab$Var1[ianc])] <- sample(blocksizes$rows[imatch][[1]][,iselect])
  }
  out
}


permute_parent_more_than_one <- function(xparents, udf, ntrts, nparents = NULL) {
  if(length(xparents)==1) {
    if(is_null(nparents) | length(nparents) == 0) {
      permute_parent_one_alg(xparents, udf, ntrts)
    } else {
      split_by_nested <- split(udf[as.character(c(xparents, nparents))], udf[as.character(nparents)])
      results_by_nested <- list()
      for(isplit in seq_along(split_by_nested)) {
        audf <- split_by_nested[[isplit]]
        audf$alloc <- permute_parent_more_than_one(xparents, audf[as.character(xparents)], ntrts)
        results_by_nested[[isplit]] <- audf
      }
      alloc_table <- do.call(rbind, results_by_nested)
      udf[["...order..."]] <- 1:nrow(udf)
      merged_alloc <- tibble::as_tibble(merge(udf, alloc_table))
      merged_alloc[order(merged_alloc[["...order..."]]), ]$alloc
    }
  } else if(is_null(nparents) | length(nparents) == 0) {
    # no nested parents
    vlevs <- lapply(udf[as.character(xparents)], unique)
    lvls <- lengths(vlevs[as.character(xparents)])
    oa <- latin_array(dim = lvls, ntrts)
    index <- lapply(xparents, function(id) match(udf[[as.character(id)]], vlevs[[as.character(id)]]))
    map_int(seq(nrow(udf)), function(i) do.call("[", c(list(oa), lapply(index, function(x) x[i]))))
  } else {
    split_by_nested <- split(udf[as.character(c(xparents, nparents))], udf[as.character(nparents)])
    results_by_nested <- list()
    for(isplit in seq_along(split_by_nested)) {
      audf <- split_by_nested[[isplit]]
      audf$alloc <- permute_parent_more_than_one(xparents, audf[as.character(xparents)], ntrts)
      results_by_nested[[isplit]] <- audf
    }
    alloc_table <- do.call(rbind, results_by_nested)
    udf[["...order..."]] <- 1:nrow(udf)
    merged_alloc <- tibble::as_tibble(merge(udf, alloc_table))
    merged_alloc[order(merged_alloc[["...order..."]]), ]$alloc
  }
}



permute_parent_one_alg <- function(vid, udf, ntrts) {
  udf$.id <- 1:nrow(udf)
  udf <- udf[order(udf[[as.character(vid)]]),]
  blocksizes <- table(udf[[as.character(vid)]])
  # if(min(blocksizes) > ntrts) {
  #   permute_parent_one(.edibble, vid, udf, ntrts)
  # } else {
  # AlgDesign::optBlock takes really long when at least one block size > ntrts
  # in this case subtract total trts
  blocksizes_adj <- blocksizes
  blocksizes_adj[blocksizes > ntrts] <- blocksizes_adj[blocksizes > ntrts] %% ntrts
  nB <- length(blocksizes_adj)
  nBlock <- sum(blocksizes_adj)
  # ALgDesign doesn't like it when the number of units is about the same as the number
  # of blocks
  if(nBlock < (ntrts - 1 + nB)) {
    return(sample(rep(sample(1:ntrts), length.out = nrow(udf))))
  }
  utils::capture.output({
    # prevent "No improvement over initial random design" print out from AlgDesign
    # where the design is balanced
    res <- tryCatch({
      AlgDesign::optBlock(~.,
                          withinData = data.frame(tindex = factor(1:ntrts)),
                          blocksizes = blocksizes_adj)
    }, error = function(x) {
      return(permute_parent_one(vid, udf, ntrts))
    })
  })
  if(is.integer(res)) return(res)
  reps <- ceiling(blocksizes / ntrts) - 1L
  out <- unname(unlist(lapply(seq_along(blocksizes), function(i) {
    blocksizes
    xv <- c(as.integer(res$Blocks[[i]]$tindex), rep(1:ntrts, reps[i]))
    if(length(xv)==1) xv else sample(xv)
  })))
  udf$.res <- out
  udf[order(udf$.id), ".res", drop = TRUE]
}

#' A custom ordering algorithm
#'
#' @param x A string specifying the class
#' @param ... Other arguments.
#'
#' @export
order_trts <- function(x, ...) {
  UseMethod("order_trts")
}

#' @export
order_trts.default <- function(x, ...) {
  abort(paste("The", cli::col_blue(x), "order is not implemented."))
}

#' @export
order_trts.dae <- function(x, prov, constrain, trts, ...) {
  # FIXME
  dat <- assign_trts(prov$design, order = "systematic", constrain = constrain, .record = FALSE) %>%
    serve_table(label_nested = everything()) %>%
    lapply(as.factor) %>%
    as.data.frame()
  out <- dae::designRandomize(allocated = dat[prov$trt_names],
                              recipient = dat[prov$unit_names],
                              nested.recipients = constrain)
  trtsv <- stats::setNames(1:nrow(trts), do.call(paste0, trts[prov$trt_names]))
  otrtsv <- do.call(paste0, dat[prov$trt_names])
  unname(trtsv[otrtsv])
}