R/partitions.R

In skgallagher/TBornotTB: Analyze Tuberculosis Patients in Maryland

#' Generate list of lists of partitions
#'
#' @param n maximum number of people in a cluster
#' @param verbose logical whether to print loops
#' @return a list of lists, the first list has entries named "nx" where x is the number of individuals in the cluster.  In the second level of lists, the names are "gy" where y is the number of non-empty generations.
generate_part_list <- function(n, verbose = FALSE){
  out_list <- vector(mode = "list", length = n)
  names(out_list) <- paste0("n", 1:n)
  for(ii in 1:n){
    if(verbose) print(paste("Cluster of size", ii))
    out_list[[ii]] <- generate_parts(ii, verbose)
  }
  return(out_list)
}


#' Generate all unique partitions* of size n
#'
#' @param n number of people in cluster
#' @param verbose if TRUE print the loop number
#' @return a list of length n, where each list entry is named "g{y}" where y is the number of generations (between 1 and n).   Each entry is a matrix with number of rows y where each column sums to n, each entry in the matrix is strictly positive.
generate_parts <- function(n, verbose = FALSE){
  sub_list <- vector(mode = "list", length = n)
  names(sub_list) <- paste0("g", 1:n)
  sub_list[["g1"]] <- matrix(n, ncol = 1)
  if(n == 1) return(sub_list)

  for(ii in 2:n){
    if(verbose) print(paste("Loop", ii))
   # browser()
    parts <- partitions::restrictedparts(n = n, m = ii,
                                         include.zero = FALSE)
  #  print(ii)
   # if(ii == (n-1)) browser()
    permuted_parts <- t(do.call('rbind',
                                plyr::alply(.data = parts, .margins = 2,
                                            .fun = function(x){
                                              if(length(unique(x)) == 1) return(matrix(x, ncol = ii))
                                              RcppAlgos::permuteGeneral(v = sort(unique(x)),
                                                                        m = ii, freqs = table(x))
                                            })))


    sub_list[[ii]] <- permuted_parts
  }
  return(sub_list)
}

#' Generate unique permutations
#'
#' @param function d vector
#' @return unique permutations
#' @details Straight from https://stackoverflow.com/questions/5671149/permute-all-unique-enumerations-of-a-vector-in-r.  Thanks!
unique_perm <- function(d) {
  if(length(d) == 1) return(d)
  dat <- factor(d)
  N <- length(dat)
  n <- tabulate(dat)
  ng <- length(n)
  if(ng==1) return(d)
  a <- N-c(0,cumsum(n))[-(ng+1)]
  foo <- lapply(1:ng, function(i) matrix(combn(a[i],n[i]),nrow=n[i]))
  out <- matrix(NA, nrow=N, ncol=prod(sapply(foo, ncol)))
  xxx <- c(0,cumsum(sapply(foo, nrow)))
  xxx <- cbind(xxx[-length(xxx)]+1, xxx[-1])
  miss <- matrix(1:N,ncol=1)
  for(i in seq_len(length(foo)-1)) {
    l1 <- foo[[i]]
    nn <- ncol(miss)
    miss <- matrix(rep(miss, ncol(l1)), nrow=nrow(miss))
    k <- (rep(0:(ncol(miss)-1), each=nrow(l1)))*nrow(miss) +
      l1[,rep(1:ncol(l1), each=nn)]
    out[xxx[i,1]:xxx[i,2],] <- matrix(miss[k], ncol=ncol(miss))
    miss <- matrix(miss[-k], ncol=ncol(miss))
  }
  k <- length(foo)
  out[xxx[k,1]:xxx[k,2],] <- miss
  out <- out[rank(as.numeric(dat), ties="first"),]
  foo <- cbind(as.vector(out), as.vector(col(out)))
  out[foo] <- d
  if(is.null(dim(out))) out <- matrix(out, ncol = 2)
  out
}

#' Get the weights of g
#'
#' @param part_list output from \code{generate_part_list}
#' @return a list of length of n, the length of part list.  each entry ii will be a vector of length ii, which denotes the times of having a generation size of jj given cluster size ii
get_weight_list <- function(part_list){
  n <- length(part_list)
  g_weights_list <- vector(mode = "list",
                           length = n)
  for(ii in 1:n){
    sub_list <- part_list[[ii]]
    g_weights_list[[ii]] <- sapply(sub_list, ncol)
  }
  return(g_weights_list)

}