R/parse.R

In vital: Tidy Analysis Tools for Mortality, Fertility, Migration and Population Data

# This file borrowed from fabletools
# It would be better to export it from there if possible
# The only change is I've removed help file generation.

parse_specials <- function(call = NULL, specials = NULL){
  if(!is.null(call)){ # No model specified
    call <- enexpr(call)

    # Don't parse xreg_specials - leave them to the xregs
    nm <- setdiff(names(specials), attr(specials, "xreg_specials"))

    parsed <- traverse_call(!!call,
                            .f = function(.x, ...) {
                              merge_named_list(.x[[1]], .x[[2]])},
                            .g = function(.x){
                              map(as.list(get_expr(.x))[-1], expr)
                            },
                            .h = function(x){ # Base types
                              x <- get_expr(x)
                              if(!is_call(x) || !(call_name(x) %in% nm)){
                                list(list(x))
                              }
                              else{# Current call is a special function
                                set_names(list(list(x)), call_name(x))
                              }
                            },
                            base = function(.x){
                              .x <- get_expr(.x)
                              !is_call(.x) || call_name(.x) != "+"
                            }
    )
  } else {
    parsed <- list()
  }

  bare_xreg <- names_no_null(parsed) == ""
  if(any(bare_xreg)){
    parsed$xreg[[length(parsed$xreg) + 1]] <- expr((!!sym("xreg"))(!!!parsed[[which(bare_xreg)]]))
    parsed[[which(bare_xreg)]] <- NULL
  }

  # Add required_specials
  missing_specials <- setdiff(attr(specials, "required_specials"), names(parsed))
  parsed[missing_specials] <- map(missing_specials, function(x) list(call(x)))

  parsed
}

# Validate the user provided model
#
# Appropriately format the user's model for evaluation. Typically ran as one of the first steps
# in a model function.
# @param model A quosure for the user's model specification
# @param data A dataset used for automatic response selection
#
# @keywords internal
validate_formula <- function(model, data = NULL){
  # Clean inputs
  if(!is_quosure(model$formula)){
    model$formula <- new_quosure(model$formula)
  }

  if(quo_is_missing(model$formula)){
    model$formula <- guess_response(data)
  }
  else{
    if(possibly(compose(is.formula, eval_tidy), FALSE)(model$formula)){
      f_env <- get_env(model$formula)
      model$formula <- eval_tidy(model$formula)
      environment(model$formula) <- f_env

      # Add response if missing
      if(is.null(model_lhs(model))){
        model$formula <- new_formula(
          lhs = guess_response(data),
          rhs = model_rhs(model),
          env = get_env(model$formula)
        )
      }
    }
    else{
      model$formula <- get_expr(model$formula)
    }
  }

  invisible(model$formula)
}

# Parse the model specification for specials
#
# Using a list of defined special functions, the user's formula specification and data
# is parsed to extract important modelling components.
#
# @param model A model definition
#
# @keywords internal
parse_model <- function(model){
  # Parse model
  list2(
    model = model,
    !!!parse_model_lhs(model),
    specials = parse_model_rhs(model)
  )
}

# Parse the RHS of the model formula for specials
#
# @inheritParams parse_model
# @keywords internal
parse_model_rhs <- function(model){
  # if(length(model$specials) == 0){
  #   return(list(specials = NULL))
  # }
  rhs <- model_rhs(model)
  specials <- parse_specials(rhs, specials = model$specials)
  map(specials, function(.x){
      map(.x, eval_tidy, data = model$data, env = model$specials)
  })
}

# Parse the RHS of the model formula for transformations
#
# @inheritParams parse_model
# @keywords internal
parse_model_lhs <- function(model){
  model_lhs <- model_lhs(model)
  if(is_call(model_lhs) && call_name(model_lhs) == "vars"){
    model_lhs[[1]] <- sym("exprs")
    model_lhs <- eval(model_lhs)
  }
  else{
    model_lhs <- list(model_lhs)
  }
  is_resp <- function(x) is_call(x) && x[[1]] == sym("resp")

  # Traverse call removing all resp() usage
  # This is used to evaluate the response from the input data
  response_exprs <- lapply(model_lhs, traverse,
                           .f = function(x, y) {
                             if(is_resp(y)) x[[1]] else call2(call_name(y), !!!x)
                           },
                           .g = function(x) x[-1],
                           # .h = function(x) if(is_resp(x)) x[[length(x)]] else x,
                           base = function(x) is_syntactic_literal(x) || is_symbol(x)
  )

  # Traverse call to parse out AST for transformations
  traversed_lhs <- lapply(model_lhs, traverse,
     .f = function(x, y) {
       if(any(resp_pos <- map_lgl(x, function(x) any(names(x) %in% "response")))){
         if(sum(resp_pos) != 1) abort("The `resp()` function can only be used once per response variable. For multivariate modelling, use `vars()`.")
         names(x)[resp_pos] <- "response"
       }
       `attr<-`(x, "call", y[[1]])
     },
     .g = function(x) x[-1],
     .h = function(x) {
       if(is_resp(x)){
         if(length(x) > 2) abort("The response variable accepts only one input. For multivariate modelling, use `vars()`.")
         list(response = x)
       } else list(x)
     },
     base = function(x) is_syntactic_literal(x) || is_symbol(x) || is_resp(x)
  )

  # Reduce traversal down to the response
  # If the response is set via resp(), remove all usage of resp() from the traversal
  # If the response is not set via resp(), identify the response by the maximum length object until encountering ties
  traversed_lhs <- lapply(traversed_lhs, traverse,
    .f = function(x, y){
      # Capture parent expression of base case
      cl <- NULL
      if(length(x) == 0){
        # Multiple length `n` variables found and cannot disambiguate response
        # Start with most disaggregated result of computation as response.
        x <- if(is.null(attr(y, "call"))) list(y[[1]]) else syms(as_label(attr(y, "call")))
      }
      else{
        if(is.null(attr(y, "call"))){
          if(is_resp(x[[1]])){
            x[[1]] <- x[[1]][[2]]
          }
        }
        else{
          # Remove resp() from call
          cl <- attr(y,"call")
          if(any(names(y) == "response")){
            cl[[which(names(y) == "response")+1]] <- x[[1]][[length(x[[1]])]]
          }
        }
      }
      c(x[[1]], cl)
    },
    .g = function(x){
      if(all(names(x) != "response") && !is.null(attr(x, "call"))){
        # parent_len <- length(eval(attr(x, "call") %||% x[[1]], envir = model$data))
        len <- map_dbl(x, function(y) length(eval(attr(y, "call") %||% y[[1]], envir = model$data, enclos = model$specials)))
        if(sum(len == max(len)) == 1){
          names(x)[which.max(len)] <- "response"
        }
      }
      if("response" %in% names(x)) x["response"] else list()
    },
    base = function(x) !is.list(x)
  )

  # Obtain parsed out response variable
  responses <- map(traversed_lhs, function(x) x[[1]])
  responses <- map_chr(responses, as_label)

  # Obtain transformation expression applied to response variable
  transform_exprs <- lapply(traversed_lhs, function(x) x[[length(x)]])

  # Invert transformation applied to response variable
  inverse_exprs <- lapply(traversed_lhs, function(x){
    x <- rev(x)
    result <- x[[length(x)]]
    for (i in seq_len(length(x) - 1)){
      result <- undo_transformation(x[[i]], x[[i + 1]], result)
    }
    result
  })

  # Produce transformation class functions for bt() usage
  make_transforms <- function(exprs, responses){
    map2(exprs, responses, function(x, response){
      new_function(args = set_names(list(missing_arg()), response), x, env = model$env)
    })
  }
  transformations <- map2(
    make_transforms(transform_exprs, responses),
    make_transforms(inverse_exprs, responses),
    new_transformation
  )

  # Test combinations of transformations to see if they are valid
  comb_trans <- map_lgl(transformations, function(x)
    length(all.names(body(x))) > length(all.vars(body(x))) + 1)
  map2(transformations[comb_trans], responses[comb_trans], function(trans, resp){
    dt <- model$data
    environment(trans) <- new_environment(dt, get_env(trans))
    inv <- invert_transformation(trans)
    environment(inv) <- new_environment(dt, get_env(inv))
    y <- eval_tidy(rlang::parse_expr(resp), data = dt, env = model$env)
    valid <- all.equal(y, inv(trans(y)))
    if(!isTRUE(valid)){
      abort(
"Could not identify a valid back-transformation for this transformation.
Please specify a valid form of your transformation using `new_transformation()`.")
    }
  })

  list(
    expressions = response_exprs,
    response = syms(responses),
    transformation = transformations
  )
}