R/runes.R

In asshah4/archetypes: Formula Spells

# rune ---------------------------------------------------------------

#' Reading Runes and Terms
#'
#' @description
#'
#' `r lifecycle::badge('experimental')`
#'
#' @param x An object of the following types that can be coerced to a
#'   `rune` object. If it is an object that contains multiple terms,
#'   such as `formula`, the parameters are pluralized and should be contained
#'   via a list of formulas. See details for further explanation.
#'
#'   * `character`
#'
#'   * `formula`
#'
#'   * `lm`
#'
#'   * `glm`
#'
#' @param side states the side of the formula the variable belongs on:
#'
#'   * __left__: for variables that are intended to be dependent
#'
#'   * __right__: for variables that are intended to be independent
#'
#'   * __meta__: for variables that are intended to explain relationships
#'
#'   * __unknown__: for variables that have unknown or undetermined sides, such
#'   as unknown position between other variables (e.g. potential mediators,
#'   conditioning variables, etc)
#'
#' @param role Specific roles the variable plays within the formula. These are
#'   of particular importance, as they serve as special terms that can effect
#'   how a formula is interpreted. Please see the _Roles_ section below for
#'   further details. The options for roles are as below:
#'
#'   * __outcome__: outcome/dependent variable that serves as an individual
#'   variable in the \eqn{exposure -> outcome} relationship (DEFAULT for LHS variables)
#'
#'   * __predictor__: predictors of the outcomes (DEFAULT for RHS variables)

#'   * __exposure__: predictor variable that serves as a primary or key
#'   variable in the \eqn{exposure -> outcome} relationship
#'
#'   * __confounder__: predictor variable that is thought to be a confounder of
#'   the causal relationship in the \eqn{exposure <- confounder -> outcome}
#'   pathway, normally thought of as an adjustment or controlling variable
#'
#'   * __mediator__: predictor variable that is thought to be a causal
#'   intermediary in the \eqn{exposure -> mediator -> outcome} pathway
#'
#'   * __interaction__: predictor variable that is proposed as an interaction
#'   term with a exposure variable, and currently only supported if exposure
#'   variables are declared
#'
#'   * __unknown__: default role of a variable that has not yet been assigned a
#'   place, such as a potential intermediary object
#'
#' @param tier Grouping variable names for covariates or __confounders__ for
#'   modeling terms together
#'
#' @param label Display-quality label describing the variable
#'
#' @param description Option for further descriptions or definitions needed for
#'   the rune, potentially part of a data dictionary
#'
#' @param distribution If its associated with a data vector, describes the
#'   distribution pattern of the original rune
#'
#' @param class Class of the variable itself, either expected or measured, such
#'   as `character` or `numeric` or `factor`
#'
#' @param type Type of variable, either categorical (qualitative) or
#'   continuous (quantitative)
#'
#' @param subtype How the variable itself is more specifically subcategorized,
#'   e.g. ordinal, continuous, dichotomous, etc
#'
#' @param operation Modification of the term to be applied when
#'   combining with data
#'
#' @section Pluralized Arguments:
#'
#'   For the arguments that would be dispatched for objects that are plural,
#'   e.g. containing multiple terms such as a `formula` object, the input should
#'   be wrapped within a `list()`.
#'
#'   For example, for the __role__ argument, it would be written:
#'
#'   `role = list(X ~ "exposure", M ~ "mediator", C ~ "confounder")`
#'
#'   This applies for all others plural objects and arguments.
#'
#' @inheritSection spells Roles
#'
#' @name runes
#' @export
rx <- function(x = unspecified(), ...) {
  UseMethod("rx", object = x)
}

#' @rdname runes
#' @export
rx.character <- function(x,
                         side = character(),
                         role = character(),
                         tier = character(),
                         label = character(),
                         description = character(),
                         distribution = character(),
                         operation = character(),
                         type = character(),
                         subtype = character(),
                         ...) {

  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  # missing values
  if (length(side) == 0) {
    side <- "unknown"
  }
  if (length(role) == 0) {
    role <- "predictor"
  }
  if (length(tier) == 0) {
    tier <- NA
  }
  if (length(operation) == 0) {
    operation <- NA
  }
  if (length(label) == 0) {
    label <- NA
  }
  if (length(description) == 0) {
    description <- NA
  }
  if (length(distribution) == 0) {
    distribution <- NA
  }
  if (length(type) == 0) {
    type <- NA
  }
  if (length(subtype) == 0) {
    subtype <- NA
  }

  # Casting
  x <- vec_cast(x, character())
  side <- vec_cast(side, character())
  role <- vec_cast(role, character())
  tier <- vec_cast(tier, character())
  label <- vec_cast(label, character())
  description <- vec_cast(description, character())
  distribution <- vec_cast(distribution, character())
  type <- vec_cast(type, character())
  subtype <- vec_cast(subtype, character())
  operation <- vec_cast(operation, character())

  new_rune(
    runes = x,
    side = side,
    role = role,
    tier = tier,
    operation = operation,
    label = label,
    description = description,
    distribution = distribution,
    type = type,
    subtype = subtype
  )
}

#' @rdname runes
#' @export
rx.formula <- function(x,
                       role = list(),
                       tier = list(),
                       label = list(),
                       description = list(),
                       distribution = list(),
                       type = list(),
                       subtype = list(),
                       ...) {

  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  # validate
  validate_class(role, "list")
  validate_class(tier, "list")
  validate_class(label, "list")
  validate_class(description, "list")
  validate_class(distribution, "list")
  validate_class(type, "list")
  validate_class(subtype, "list")
  roles <- formula_to_named_list(role)
  tiers <- formula_to_named_list(tier)
  labels <- formula_to_named_list(label)
  descriptions <- formula_to_named_list(description)
  distributions <- formula_to_named_list(distribution)
  types <- formula_to_named_list(type)
  subtypes <- formula_to_named_list(subtype)

  # All terms are needed to build rx record
  left <- lhs(x)
  right <- rhs(x, tidy = TRUE)
  all <- c(left, right)
  n <- length(all)

  # Roles and operations need to be identified (on which terms they apply)
  right_ops <-
    x |>
    all.names() |>
    {
      \(.x) {
        # These will be named roles
        var_names <- character()
        var_roles <- character()
        for (i in seq_along(.x)) {
          if (.x[i] %in% template_shortcuts) {
            var_names <- append(var_names, .x[i + 1])
            var_roles <- append(var_roles, .x[i])
          }
        }

        names(var_roles) <- var_names
        var_roles |>
          as.list()
      }
    }()

  # Warn and validate for interaction (as needs exposure variable)
  if ("In" %in% right_ops & !("X" %in% right_ops)) {
    warning("As a specific interaction term was included, an exposure variable must be included as well or this cannot later be expanded to an appropriate formula.")
  }

  # check to see if it is a "role" or a data transformation
  which_ops <- right_ops %in% template_shortcuts
  role_ops <- right_ops[which_ops]
  data_ops <- right_ops[!which_ops]

  other <- right[!(right %in% names(role_ops))]
  other_ops <- rep("predictor", length(other))
  names(other_ops) <- other
  other_ops <- as.list(other_ops)

  left_ops <- rep("outcome", length(left))
  names(left_ops) <- left
  left_ops <- as.list(left_ops)

  role_ops <- c(role_ops, left_ops, other_ops)

  # Interaction term is already included by name
  for (i in seq_along(role_ops)) {
    if (role_ops[[i]] == "O") {
      role_ops[[i]] <- "outcome"
    }
    if (role_ops[[i]] == "X") {
      role_ops[[i]] <- "exposure"
    }

    if (role_ops[[i]] == "M") {
      role_ops[[i]] <- "mediator"
    }

    if (role_ops[[i]] == "C") {
      role_ops[[i]] <- "confounder"
    }

    if (role_ops[[i]] == "S") {
      role_ops[[i]] <- "strata"
    }

    if (role_ops[[i]] == "In") {
      role_ops[[i]] <- "interaction"
    }
  }

  # create runes
  rune_vector <- new_rune()

  for (i in 1:n) {
    # make parameters
    t <- all[i]

    # Sides and meta runes
    side <- if (t %in% names(role_ops[role_ops == "strata"])) {
      "meta"
    } else if (t %in% left) {
      "left"
    } else if (t %in% right) {
      "right"
    }

    # Data transforms
    op <- if (t %in% names(data_ops)) {
      data_ops[[t]]
    } else {
      NA
    }

    # Roles
    role <- if (t %in% names(role_ops)) {
      role_ops[[t]]
    } else {
      NA
    }

    # Tiers
    tier <-
      if (t %in% names(tiers) & t %in% names(role_ops[role_ops %in% c("exposure", "mediator", "strata", "outcome")])) {
        message(
          "The rune `",
          t,
          "` cannot be given a tier as it is not an ordinary predictor."
        )
      } else if (t %in% names(tiers)) {
        tiers[[t]]
      } else {
        NA
      }

    # Labels
    lab <- if (t %in% names(labels)) {
      labels[[t]]
    } else {
      NA
    }

    # place into rx list after casting appropriate classes
    rn <- rx.character(
      x = vec_cast(t, character()),
      side = vec_cast(side, character()),
      role = vec_cast(role, character()),
      tier = vec_cast(tier, character()),
      operation = vec_cast(op, character()),
      label = vec_cast(lab, character())
    )

    rune_vector <- append(rune_vector, rn)

  }

  # return as a record of runes
  rune_vector
}


#' @rdname runes
#' @export
rx.lm <- function(x,
                  role = list(),
                  tier = list(),
                  label = list(),
                  description = list(),
                  distribution = list(),
                  type = list(),
                  subtype = list(),
                  ...) {


  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  # obtain original formula
  f <- stats::formula(x)

  # generate runes
  rx.formula(
    f,
    role = role,
    tier = tier,
    label = label,
    description = description,
    distribution = distribution,
    type = type,
    subtype = subtype
  )
}

#' @rdname runes
#' @export
rx.glm <- rx.lm

#' @rdname runes
#' @export
rx.coxph <- rx.lm

#' @rdname runes
#' @export
rx.model_fit <- function(x,
                         role = list(),
                         tier = list(),
                         label = list(),
                         description = list(),
                         distribution = list(),
                         type = list(),
                         subtype = list(),
                         ...) {

  # Early break and validation
  if (validate_empty(x)) {
    return(new_rune())
  }

  # Get model fit and pass to appropriate rx dispatcher
  m <- x$fit
  validate_models(m)

  # Return
  rx(m)
}

#' @rdname runes
#' @export
rx.rune <- function(x, ...) {
  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  # Return the same
  x
}

#' @rdname runes
#' @export
rx.fmls <- function(x, ...) {
  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  get_runes(x)

}

#' @rdname runes
#' @export
rx.spell <- function(x, ...) {
  # Early Break if needed
  if (validate_empty(x)) {
    return(new_rune())
  }

  # Return to runes
  field(x, "runes")[[1]]
}

#' @rdname runes
#' @export
rx.default <- function(x = unspecified(), ...) {
  # Early break
  if (length(x) == 0) {
    return(new_rune())
  }

  stop("`rx()` is not defined for a `",
    class(x)[1],
    "` object.",
    call. = FALSE
  )
}


#' @rdname runes
#' @export
distill_rune <- rx

# Record definition ------------------------------------------------------------

#' record of formula rune
#' @keywords internal
#' @noRd
new_rune <- function(runes = character(),
                     side = character(),
                     role = character(),
                     tier = character(),
                     label = character(),
                     description = character(),
                     distribution = character(),
                     operation = character(),
                     type = character(),
                     subtype = character(),
                     order = integer()) {
  # Validation
  vec_assert(runes, ptype = character())
  vec_assert(side, ptype = character())
  vec_assert(role, ptype = character())
  vec_assert(tier, ptype = character())
  vec_assert(label, ptype = character())
  vec_assert(description, ptype = character())
  vec_assert(distribution, ptype = character())
  vec_assert(operation, ptype = character())
  vec_assert(type, ptype = character())
  vec_assert(subtype, ptype = character())
  vec_assert(order, ptype = integer())

  # Forced order
  if (length(runes) > 0) {
    order <- 0L
  }

  new_rcrd(
    list(
      "runes" = runes,
      "side" = side,
      "role" = role,
      "tier" = tier,
      "label" = label,
      "description" = description,
      "distribution" = distribution,
      "operation" = operation,
      "type" = type,
      "subtype" = subtype,
      "order" = order
    ),
    class = "rune"
  )
}

#' @keywords internal
#' @noRd
methods::setOldClass(c("rune", "rcrds_rcrd"))

# Output -----------------------------------------------------------------------

#' @export
format.rune <- function(x, ...) {
  tms <- vec_data(x)
  fmt <- character()

  if (vec_size(x) == 0) {
    fmt <- new_rune()
  } else if (has_cli() & vec_size(x) > 0) {
    for (i in 1:nrow(tms)) {
      if (tms$role[i] == "outcome") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_yellow(t))
      }

      if (tms$role[i] == "predictor") {
        t <- tms$runes[i]
        fmt <- append(fmt, t)
      }

      if (tms$role[i] == "exposure") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_magenta(t))
      }

      if (tms$role[i] == "mediator") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_cyan(t))
      }

      if (tms$role[i] == "confounder") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_blue(t))
      }

      if (tms$role[i] == "strata") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_br_white(t))
      }

      if (tms$role[i] == "interaction") {
        t <- tms$runes[i]
        fmt <- append(fmt, cli::col_silver(t))
      }

      if (tms$role[i] == "unknown") {
        t <- tms$runes[i]
        fmt <- append(fmt, t)
      }

    }
  } else {
    for (i in 1:nrow(tms)) {
      fmt <- append(fmt, tms$runes[i])
    }
  }

  # return
  fmt
}

#' @export
obj_print_data.rune <- function(x, ...) {
  if (vec_size(x) == 0) {
    new_rune()
  } else if (vec_size(x) > 1) {
    cat(format(x), sep = "\n")
  } else {
    cat(format(x))
  }
}

#' @export
vec_ptype_full.rune <- function(x, ...) {
  "rune"
}

#' @export
vec_ptype_abbr.rune <- function(x, ...) {
  "rx"
}

# Casting and coercion ---------------------------------------------------------

### rune() ###

#' @export
vec_ptype2.rune.rune <- function(x, y, ...) {
  x
}

#' @export
vec_cast.rune.rune <- function(x, to, ...) {
  x
}

### character() ###

#' @export
vec_ptype2.rune.character <- function(x, y, ...) {
  # `x` is rune
  # `y` is character
  y
}

#' @export
vec_ptype2.character.rune <- function(x, y, ...) {
  # `x` is character
  # `y` is rune
  x
}

#' @export
vec_cast.rune.character <- function(x, to, ...) {
  # order is flipped, such that `x` is character
  attributes(x) <- NULL
  x[[1]]
}

#' @export
vec_cast.character.rune <- function(x, to, ...) {
  # order is flipped, such that `x` is rune
  attributes(x) <- NULL
  x[[1]]
}

### list_of() ###

#' @export
vec_ptype2.rcrds_list_of.rune <- function(x, y, ...) {
  x
}

#' @export
vec_ptype2.rune.rcrds_list_of <- function(x, y, ...) {
  y
}

#' @export
vec_cast.rcrds_list_of.rune <- function(x, to, ...) {
  tl <- as.list(x) # convert to list
  lot <- new_list_of(tl, ptype = new_rune()) # make new list of
  lot # return list of rune
}

#' @export
vec_cast.rune.rcrds_list_of <- function(x, to, ...) {
  t <- append(x, new_rune()) # convert to a flattened record
  t # return record of rune
}

# Formulas

#' @export
#' @importFrom stats formula
formula.rune <- function(x, ...) {

  # If rune has interaction term and is appropriate complexity, this should be
  # expanded to show interaction with exposure variable
  rls <- roles(x)

  if ("interaction" %in% rls & "exposure" %in% rls) {
    In <- as.character(get_runes(x, field = "role", value = "interaction"))
    X <- as.character(get_runes(x, field = "role", value = "exposure"))
    right <- c(rhs(x), paste(rep(X, each = length(In)), In, sep = ":"))
    message("Recreating interaction term with the specified exposure(s).")
  } else {
    right <- rhs(x)
  }

  paste(lhs(x), collapse = " + ") |>
    paste(paste(right, collapse = " + "), sep = " ~ ") |>
    stats::as.formula()
}