R/egenvar.R

In R4VN: Health Data Analysis and Publication-Ready Reporting

# ============================================================================
# R4VN egenvar(): extended variable generation
# ============================================================================

.r4vn_egen_group_key <- function(data, by_names = character()) {
  n <- nrow(data)
  if (!length(by_names)) return(rep.int("__ALL__", n))
  parts <- lapply(data[by_names], function(x) {
    z <- as.character(x)
    z[is.na(z)] <- "<NA>"
    paste0(nchar(z), ":", z)
  })
  do.call(paste, c(parts, sep = "\r"))
}

.r4vn_egen_by_apply <- function(x, key, FUN, simplify = TRUE) {
  idx <- split(seq_along(key), key, drop = TRUE)
  out <- vector("list", length(key))
  for (ii in idx) {
    val <- FUN(x[ii], ii)
    if (length(val) == 1L) val <- rep(val, length(ii))
    if (length(val) != length(ii)) {
      stop("An egenvar group function returned an invalid length.", call. = FALSE)
    }
    out[ii] <- as.list(val)
  }
  if (!isTRUE(simplify)) return(out)
  unlist(out, recursive = FALSE, use.names = FALSE)
}

.r4vn_egen_num <- function(x, arg = "variable") {
  if (is.logical(x)) x <- as.integer(x)
  if (!is.numeric(x)) stop("`", arg, "` must be numeric for this egenvar function.", call. = FALSE)
  x
}

.r4vn_egen_row_matrix <- function(exprs, data) {
  if (!length(exprs)) stop("A row function needs at least one variable.", call. = FALSE)
  nms <- unique(unlist(lapply(exprs, .r4vn_selector, data_names = names(data)), use.names = FALSE))
  if (!length(nms)) stop("No variables were selected for the row function.", call. = FALSE)
  vals <- lapply(nms, function(nm) .r4vn_egen_num(data[[nm]], nm))
  out <- do.call(cbind, vals)
  if (is.null(dim(out))) out <- matrix(out, ncol = 1L)
  colnames(out) <- nms
  out
}

.r4vn_egen_row_stat <- function(fun, exprs, data) {
  m <- .r4vn_egen_row_matrix(exprs, data)
  present <- rowSums(!is.na(m))
  ans <- switch(
    fun,
    rowmin = apply(m, 1L, function(z) if (all(is.na(z))) NA_real_ else min(z, na.rm = TRUE)),
    rowmax = apply(m, 1L, function(z) if (all(is.na(z))) NA_real_ else max(z, na.rm = TRUE)),
    rowmean = {
      z <- rowMeans(m, na.rm = TRUE); z[present == 0L] <- NA_real_; z
    },
    rowsum = {
      z <- rowSums(m, na.rm = TRUE); z[present == 0L] <- NA_real_; z
    },
    rowmedian = apply(m, 1L, function(z) if (all(is.na(z))) NA_real_ else stats::median(z, na.rm = TRUE)),
    rowsd = apply(m, 1L, function(z) {
      z <- z[!is.na(z)]
      if (length(z) < 2L) NA_real_ else stats::sd(z)
    }),
    rowmiss = rowSums(is.na(m)),
    rownonmiss = present,
    rowfirst = apply(m, 1L, function(z) {
      z <- z[!is.na(z)]; if (length(z)) z[1L] else NA_real_
    }),
    rowlast = apply(m, 1L, function(z) {
      z <- z[!is.na(z)]; if (length(z)) z[length(z)] else NA_real_
    }),
    stop("Unknown egenvar row function.", call. = FALSE)
  )
  unname(ans)
}

.r4vn_egen_group_stat <- function(fun, x, key, args = list()) {
  if (fun %in% c("mean", "sd", "min", "max", "median", "total", "count", "z", "pctile", "rank")) {
    if (missing(x) || is.null(x)) stop("`", fun, "()` needs a variable.", call. = FALSE)
  }
  if (fun %in% c("mean", "sd", "min", "max", "median", "total", "z", "pctile")) x <- .r4vn_egen_num(x)

  if (fun == "mean") return(.r4vn_egen_by_apply(x, key, function(z, ii) if (all(is.na(z))) NA_real_ else mean(z, na.rm = TRUE)))
  if (fun == "sd") return(.r4vn_egen_by_apply(x, key, function(z, ii) { z <- z[!is.na(z)]; if (length(z) < 2L) NA_real_ else stats::sd(z) }))
  if (fun == "min") return(.r4vn_egen_by_apply(x, key, function(z, ii) if (all(is.na(z))) NA_real_ else min(z, na.rm = TRUE)))
  if (fun == "max") return(.r4vn_egen_by_apply(x, key, function(z, ii) if (all(is.na(z))) NA_real_ else max(z, na.rm = TRUE)))
  if (fun == "median") return(.r4vn_egen_by_apply(x, key, function(z, ii) if (all(is.na(z))) NA_real_ else stats::median(z, na.rm = TRUE)))
  if (fun == "total") return(.r4vn_egen_by_apply(x, key, function(z, ii) sum(z, na.rm = TRUE)))
  if (fun == "count") return(.r4vn_egen_by_apply(x, key, function(z, ii) sum(!is.na(z))))
  if (fun == "n") return(.r4vn_egen_by_apply(seq_along(key), key, function(z, ii) length(ii)))
  if (fun == "seq") return(.r4vn_egen_by_apply(seq_along(key), key, function(z, ii) seq_along(ii)))
  if (fun == "z") return(.r4vn_egen_by_apply(x, key, function(z, ii) {
    mu <- if (all(is.na(z))) NA_real_ else mean(z, na.rm = TRUE)
    ss <- if (sum(!is.na(z)) < 2L) NA_real_ else stats::sd(z, na.rm = TRUE)
    if (!is.finite(ss) || ss == 0) rep(NA_real_, length(z)) else (z - mu) / ss
  }))
  if (fun == "pctile") {
    p <- args$p %||% args$prob %||% args$arg2 %||% 50
    if (!is.numeric(p) || length(p) != 1L || !is.finite(p)) stop("`p` in pctile() must be one finite number.", call. = FALSE)
    if (p > 1) p <- p / 100
    if (p < 0 || p > 1) stop("`p` in pctile() must be between 0 and 1, or between 0 and 100.", call. = FALSE)
    return(.r4vn_egen_by_apply(x, key, function(z, ii) if (all(is.na(z))) NA_real_ else unname(stats::quantile(z, probs = p, na.rm = TRUE, names = FALSE, type = 7))))
  }
  if (fun == "rank") {
    ties <- args$ties %||% args$arg2 %||% "average"
    valid <- c("average", "first", "last", "random", "max", "min")
    if (!is.character(ties) || length(ties) != 1L || !ties %in% valid) stop("`ties` must be one of: ", paste(valid, collapse = ", "), ".", call. = FALSE)
    return(.r4vn_egen_by_apply(x, key, function(z, ii) rank(z, na.last = "keep", ties.method = ties)))
  }
  stop("Unknown egenvar group function.", call. = FALSE)
}

.r4vn_egen_eval <- function(expr, data, env, key) {
  if (!is.call(expr)) return(.r4vn_eval(expr, data, env))
  head <- as.character(expr[[1L]])
  row_fun <- c("rowmin", "rowmax", "rowmean", "rowsum", "rowmedian", "rowsd", "rowmiss", "rownonmiss", "rowfirst", "rowlast")
  group_fun <- c("mean", "sd", "min", "max", "median", "total", "count", "n", "seq", "z", "pctile", "rank")

  if (head %in% row_fun) {
    return(.r4vn_egen_row_stat(head, as.list(expr)[-1L], data))
  }
  if (head %in% group_fun) {
    a <- as.list(expr)[-1L]
    if (head %in% c("n", "seq")) {
      if (length(a)) stop("`", head, "()` does not take a variable.", call. = FALSE)
      return(.r4vn_egen_group_stat(head, NULL, key))
    }
    if (!length(a)) stop("`", head, "()` needs a variable.", call. = FALSE)
    x <- .r4vn_eval(a[[1L]], data, env)
    extra <- list()
    if (length(a) > 1L) {
      nm <- names(a)[-1L]
      for (j in 2:length(a)) {
        val <- eval(a[[j]], envir = env)
        key_name <- if (!is.null(nm) && length(nm) >= j - 1L && nzchar(nm[j - 1L])) nm[j - 1L] else paste0("arg", j)
        extra[[key_name]] <- val
      }
    }
    return(.r4vn_egen_group_stat(head, x, key, extra))
  }
  .r4vn_eval(expr, data, env)
}

#' Extended variable generation
#'
#' Creates row-wise, group-wise, ranking, standardization, and identifier
#' variables using a compact syntax inspired by Stata's `egen`, while following
#' the active-data conventions of [genvar()].
#'
#' @param ... Named expressions in the form `new_variable = function(...)`.
#'   An explicit data frame may also be supplied first for compatibility with
#'   other R4VN data-editing commands.
#' @param data Optional explicit data frame. When omitted, the active data frame
#'   selected by [usedf()] is used.
#' @param by Optional grouping variables, for example `by = sex` or
#'   `by = vars(site, sex)`. Group summaries are repeated on every observation
#'   in the corresponding group.
#' @param label Optional variable label or one label per generated variable.
#' @param values Optional named value labels applied to generated variables.
#'
#' @details
#' Row-wise functions accept individual variables, variable ranges, `vars()`
#' selections, and wildcard selectors: `rowmin()`, `rowmax()`, `rowmean()`,
#' `rowsum()`, `rowmedian()`, `rowsd()`, `rowmiss()`, `rownonmiss()`,
#' `rowfirst()`, and `rowlast()`. Thus `rowmean(q1:q10)` is valid R4VN syntax.
#'
#' Group/overall functions are `mean()`, `sd()`, `min()`, `max()`, `median()`,
#' `total()`, `count()`, `n()`, `seq()`, `z()`, `pctile()`, and `rank()`.
#' Without `by`, they operate over the complete data set. With `by`, they
#' operate separately within groups. Missing values are ignored by summary
#' functions; `count()` counts non-missing values.
#'
#' Special identifier functions are `group()` and `tag()`. They use one or more
#' variables supplied inside the function and do not depend on the `by`
#' argument. `group(site, sex)` gives consecutive integer IDs for observed
#' combinations; `tag(id)` marks the first occurrence of each distinct value.
#'
#' @return The edited data frame invisibly. If active data are linked to a
#'   visible object through [usedf()], that object is updated as well.
#' @export
#'
#' @examples
#' d <- data.frame(
#'   id = c(1, 1, 2, 3),
#'   sex = c("F", "F", "M", "M"),
#'   q1 = c(2, 4, 3, NA), q2 = c(3, 5, 2, 4), q3 = c(4, NA, 1, 5),
#'   bmi = c(20, 22, 25, 27)
#' )
#' egenvar(d, min_score = rowmin(q1:q3),
#'         max_score = rowmax(q1:q3),
#'         mean_score = rowmean(q1:q3))
#'
#' egenvar(d, mean_bmi = mean(bmi), z_bmi = z(bmi), by = sex)
#' egenvar(d, p75_bmi = pctile(bmi, p = 75), rank_bmi = rank(bmi), by = sex)
#' egenvar(d, n_group = n(), sequence = seq(), by = sex)
#' egenvar(d, person_group = group(id, sex), first_id = tag(id))
egenvar <- function(..., data = NULL, by = NULL, label = NULL, values = NULL) {
  env <- parent.frame()
  exprs <- as.list(substitute(list(...)))[-1L]
  context <- .r4vn_edit_context(
    exprs = exprs,
    data_expr = substitute(data),
    data_supplied = !missing(data),
    env = env
  )
  d <- context$data
  exprs <- context$args
  if (!length(exprs)) stop("No new variables were specified.", call. = FALSE)
  variables <- names(exprs)
  if (is.null(variables) || any(!nzchar(variables))) stop("Every generated variable must use `new_name = expression`.", call. = FALSE)
  if (anyDuplicated(variables)) stop("Generated variable names must be unique.", call. = FALSE)

  by_names <- character()
  if (!missing(by) && !is.null(substitute(by)) && !identical(substitute(by), quote(NULL))) {
    by_names <- .r4vn_selector(substitute(by), names(d))
  }
  key <- .r4vn_egen_group_key(d, by_names)
  labels <- .r4vn_labels(label, variables)

  for (i in seq_along(exprs)) {
    expr <- exprs[[i]]
    head <- if (is.call(expr)) as.character(expr[[1L]]) else ""

    if (head %in% c("group", "tag")) {
      a <- as.list(expr)[-1L]
      if (!length(a)) stop("`", head, "()` needs at least one variable.", call. = FALSE)
      nms <- unique(unlist(lapply(a, .r4vn_selector, data_names = names(d)), use.names = FALSE))
      k <- .r4vn_egen_group_key(d, nms)
      if (head == "group") {
        value <- match(k, unique(k))
      } else {
        value <- as.integer(!duplicated(k))
      }
    } else {
      value <- .r4vn_egen_eval(expr, d, env, key)
    }

    if (length(value) == 1L && nrow(d) > 1L) value <- rep(value, nrow(d))
    if (length(value) != nrow(d)) stop("Generated variable `", variables[i], "` must contain one value per observation.", call. = FALSE)
    if (is.logical(value)) value <- as.integer(value)
    d[[variables[i]]] <- .r4vn_apply_lab(value, label = labels[[i]], values = values)
  }
  .r4vn_commit_context(context$target, d)
}