R/stata_functions.R

In Steensson/valhalla: A helper for Sebastian

#' Function to mimic the _n function in Stata
#'
#' @param df Your dataframe
#' @param idvar A string containing the name of the group variable
#' @param counter_name The name of the counter variable. Defaults to "n".
#'
#' @return Your dataframe but with at counter with the name according to counter_name
#' @export
#'
#' @examples df <- data.frame(id = c("A", "A", "B", "B", "B", "C", "C"), val = rnorm(7))
#' df <- stata_n(df, "id")
stata_n <- function(df, idvar, counter_name = "n") {

    df <- cbind(df, v1 = 1)
    df$counter_name <- ave(df$v1, df[[idvar]], FUN = seq_along)
    colnames(df)[ncol(df)] <- counter_name
    df$v1 <- NULL

    return(df)
}

#' Function to mimic the _N function in Stata
#'
#' @param df Your dataframe
#' @param idvar A string containing the name of the group variable
#' @param counter_name The name of the counter variable. Defaults to "N".
#'
#' @return Your dataframe but with at counter with the name according to counter_name
#' @export
#'
#' @examples df <- data.frame(id = c("A", "A", "B", "B", "B", "C", "C"), val = rnorm(7))
#' df <- stata_N(df, "id")
stata_N <- function(df, idvar, counter_name = "N") {

    df <- cbind(df, v1 = 1)
    df$counter_name <- ave(df$v1, df[[idvar]], FUN = sum)
    colnames(df)[ncol(df)] <- counter_name
    df$v1 <- NULL

    return(df)
}