R/expand_sch.R

In tessaleejohnson/corclus: Simulate Data from the Correlated Cluster Multiple Membership Model

#---------------------------------
# External Dependencies:
# glue
#
# Internal Dependencies:
#
#---------------------------------


#' expand_sch
#'
#' This function expands school-level information generated by
#' \code{\link{gen_u_mmrem}}, repeating each school's info a number of times
#' equal to the number of students attending each school. That is, if there
#' are \code{.n_stu} per school and \code{.n_sch}, the final output should have
#' \code{.n_stu * .n_sch} rows.
#'
#' @inheritParams corclus_params
#'
#' @return This function returns a dataframe with a number of rows equal to
#' \code{.n_stu * .n_sch}. The columns of the input data (\code{.sch_dat}) are
#' retained in the output.
#'
#' @export
#'
#' @examples \dontrun{
#'
#' # generate data using the gen_u_mmrem function
#' dat <- gen_u_mmrem(.n_sch = 10)
#'
#' # expand the data to repeat school information for each student
#' # attending a given school
#' expand_sch(dat, .n_stu = 10, .n_sch = 10)
#'
#' }
expand_sch <-
  function(
    .sch_dat,
    .n_stu,
    .n_sch
  ) {

    ##--setup--##

    # check if .sch_dat is a dataframe with .n_sch rows
    if (!is.data.frame(.sch_dat)) {
      stop(".sch_dat must have data.frame class.")
    }

    if (NROW(.sch_dat) != .n_sch) {
      stop(
        glue::glue('
          .sch_dat must have {.n_sch} rows but it has {NROW(.sch_dat)}.
          '))
    }


    ##--expand school info across students--##

    # repeat each row of .sch_dat a number of times equal to .n_stu and combine
    # expanded school data with student ID (then remove the weird rownames)
    rep_seq <- rep(1:.n_sch, each = .n_stu)
    sch_exp <- .sch_dat[rep_seq, ]
    rownames(sch_exp) <- NULL


    ##--output--##
    sch_exp

  }