R/spells.R

In asshah4/archetypes: Formula Spells

# Formula Magic ----------------------------------------------------------------

#' Spells
#'
#' @description
#'
#' `r lifecycle::badge('stable')`
#'
#' This function defines a magical (modified) `formula` class that has been
#' vectorized. The `spell` serves as a set of instructions or a _script_ for the
#' formula and its runes. It expands upon the functionality of formulas,
#' allowing for additional descriptions and relationships to exist between the
#' runes.
#'
#' @param x Objects of the following types can be used as inputs
#'
#'   * `rune`
#'
#'   * `formula`
#'
#' @inheritParams runes
#'
#' @param pattern This is the expansion pattern used to decide how the
#'   covariates will incorporated into the formulas. The options are
#'   `c("direct", "sequential", "parallel")`. See the details for further
#'   explanation.
#'
#'   * __direct__: the covariates will all be included in each formula
#'
#'   * __sequential__: the covariates will be added sequentially, one by one, or
#'   by tiers, as indicated
#'
#'   * __parallel__: the covariates or tiers of covariates will be placed in
#'   parallel

#' @param ... Arguments to be passed to or from other methods
#'
#' @section Roles:
#'
#' Specific roles the variable plays within the formula. These are of particular
#' importance, as they serve as special runes that can effect how a formula is
#' interpreted. The specialized options for roles are as below:
#'
#' * __exposure__ or `X(...)`
#'
#' * __outcome__ or `O(...)` or placement of variable on LHS of formula
#'
#' * __confounder__ or `C(...)`
#'
#' * __mediator__ or `M(...)`
#'
#' * __strata__ or `S(...)`
#'
#' * __interaction__ or `In()`
#'
#' Formulas can be condensed by applying their specific role to individual runes
#' as a function/wrapper. For example, `y ~ X(x1) + x2 + x3`. This would signify
#' that `x1` has the specific role of an exposure.
#'
#' @inheritSection runes Pluralized Arguments
#'
#' @section Patterns:
#'
#' The expansion pattern allows for instructions on how the covariates should be
#' included in different formulas. Below, assuming that _x1_, _x2_, and _x3_ are
#' covariates...
#'
#' \deqn{y ~ x1 + x2 + x3}
#'
#' __Direct__:
#'
#' \deqn{y ~ x1 + x2 + x3}
#'
#' __Seqential__:
#'
#' \deqn{y ~ x1}
#' \deqn{y ~ x1 + x2}
#' \deqn{y ~ x1 + x2 + x3}
#'
#' __Parallel__:
#'
#' \deqn{y ~ x1}
#' \deqn{y ~ x2}
#' \deqn{y ~ x3}
#'
#' @return An object of class `spell`
#' @name spells
#' @export
sx <- function(x = unspecified(),
               role = list(),
               tier = list(),
               label = list(),
               pattern = character(),
               ...) {

  # Break early if nothing is given
  # If appropriate class, but empty, then also break early but warn/message
  if (class(x)[1] == "vctrs_unspecified") {
    return(new_spell())
  }
  validate_class(x, c("rune", "formula"))
  if (validate_empty(x)) {
    return(new_spell())
  }

  # Check pattern
  if (length(pattern) == 0) {
    pattern <- "direct"
  }
  if (!pattern %in% template_patterns) {
    stop("The pattern ",
         deparse(pattern),
         " is not yet supported.",
         call. = FALSE)
  }

  # runes list (nested for field length equivalence)
  # Updated attributes/components internally
  t <-
    distill_rune(x) |>
    set_roles(roles = formula_to_named_list(role)) |>
    set_tiers(tiers = formula_to_named_list(tier)) |>
    set_labels(labels = formula_to_named_list(label)) |>
    unique()

  # Look at composition of runes
  order <- decipher(t)

  # Formula
  f <- deparse1(stats::formula(t))

  # Return
  new_spell(
    formula = f,
    runes = t,
    pattern = pattern,
    order = order
  )
}

#' @rdname spells
#' @export
cast_spell <- sx

# Vector Creation --------------------------------------------------------------

#' Formula vector
#' @keywords internal
#' @noRd
new_spell <- function(runes = distill_rune(),
                      formula = character(),
                      pattern = character(),
                      order = integer()) {

  # Validation of types
  vec_assert(runes, ptype = distill_rune())
  vec_assert(formula, ptype = character())
  vec_assert(pattern, ptype = character())
  vec_assert(order, ptype = integer())

  # Bend runes into a list
  if (vec_size(runes) == 0) {
    runes <- distill_rune()
  } else {
    runes <- list(runes)
  }

  # Everything needs to be the same length
  new_rcrd(
    fields = list(
      "formula" = formula,
      "runes" = runes,
      "pattern" = pattern,
      "order" = order
    ),
    class = "spell"
  )
}

#' @keywords internal
#' @noRd
methods::setOldClass(c("spell", "vctrs_vctr"))

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

#' @export
format.spell <- function(x, ...) {

  # Character representation of formula
  if (vec_size(x) == 0) {
    return()
  } else {
    fmt <-
      sapply(x, FUN = function(.x) {
        field(.x, "formula") |>
          format()
      })
  }
  # Return
  fmt
}

#' @export
obj_print_data.spell <- function(x, ...) {

  # Colorful printing
  if (vec_size(x) == 0) {
    fmt <- new_spell()
  } else {
    fmt <-
      sapply(
        x,
        FUN = function(.x) {
          t <- field(.x, "runes")[[1]]
          f <- stats::formula(field(.x, "formula"))
          left <- match_runes(t, lhs(f))
          right <- match_runes(t, rhs(f))

          f <-
            paste(format(left), collapse = " + ") |>
            paste(paste(format(right), collapse = " + "), sep = " ~ ")

          f
        }
      )

  }

  if (length(fmt) > 1) {
    cat(format(fmt), sep = "\n")
  } else {
    cat(format(fmt))
  }
}

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

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

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

# Arithmetic
vec_arith.spell <- function(op, x, y, ...) {
  UseMethod("vec_arith.spell", y)
}

vec_arith.spell.default <- function(op, x, y, ...) {
  stop_incompatible_op(op, x, y)
}


### self

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

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

### characters

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

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

#' @export
vec_cast.character.spell <- function(x, to, ...) {
  format(x) # Returns a character class by default
}

### sx

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

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

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

### base formula

#' @export
formula.spell <- function(x, ...) {
  format(x) |>
    stats::as.formula()
}