R/group_assign.R

In vjf2/SocGen: Analysis of Genetic and Social Behavior Data

#' Assign individuals to groups based on 2-dimensional location and group size 
#'
#' This function clusters individuals into groups based on their locations and some mean group size. 
#' @param data list of simulation data generated by swap or UD method
#' @param id individual identifier field
#' @param xcoords x coordinates
#' @param ycoords y coordinates
#' @param time time step variable
#' @param group_size average group size in results
#' @param group_vector vector of group sizes per time step
#' @param method grouping method, either kmeans or hierarchical clustering
#' @keywords group ids, cluster ids
#' @export
#' @examples
#' group_assign()

group_assign <- function(data = data,
                         id = "id",
                         xcoord = "x",
                         ycoord = "y",
                         time = NULL,
                         group_size = NULL,
                         group_vector = NULL,
                         method = c("hclust", "kmeans")) {
  lapply(data, function(w) {
    
    each_days_assoc <- lapply(seq_along(w), function(day_idx) {
      
      x1<- w[[day_idx]]
      
      if(!is.null(group_size)){
        
        num_clust <- ceiling(nrow(x1) / group_size)}
      
      else{num_clust <- group_vector[day_idx]}
      
      if (num_clust > 1 & method=="hclust") {
        xc <- hclust(dist(x1[, c(xcoord, ycoord)]))
        groups <- cutree(xc, k = num_clust)
      }
      
      if (num_clust > 1 & method=="kmeans") {
        xk <- kmeans(dist(x1[, c(xcoord, ycoord)]), num_clust, algorithm="MacQueen")
        groups <- as.numeric(xk$cluster)
      }
      if (num_clust == 1) {
        groups <- 1
      }
      
      # daygroup <- cbind(as.character(x1[,id]), paste0(time[day_idx],"_",groups))
      daygroup <- data.frame(id=as.character(x1[,id]), 
                             group=paste0(time[day_idx],"_",groups))
      rownames(daygroup) <- NULL
      return(daygroup)
    })
    
    eda <- do.call("rbind", each_days_assoc)
    # eda$observation_id <- paste0(names(group_vector), "_", eda$groups)
    rownames(eda) <- NULL
    return(eda)
  })
}