R/jacobian.R

In mailund/dfdr: Automatic Differentiation of Simple Functions

jacobian <- function(f, y, x, derivs = NULL, num_functions = NULL) {
  y <- rlang::ensym(y)
  y <- rlang::as_string(y)
  x <- rlang::ensym(x)
  x <- rlang::as_string(x)

  stopifnot("the function is missing" = !is.null(f))
  stopifnot("the variable y is missing" = !is.null(y))
  stopifnot("the variable x is missing" = !is.null(x))

  if (!is.null(num_functions)) {
    stopifnot("the num_functions arguments has to be of type numeric" = is.numeric(num_functions))
  }

  # create function list
  # ============================================================================
  fl <- init_fct_list()
  if (!is.null(derivs)) {
    fd <- derivs@funs
    for (i in seq_along(fd)) {
      fl <- append_fct(fl, fd[[i]]@name, fd[[i]]@dfdx, fd[[i]]@name_deriv, fd[[i]]@keep)
    }
  }
  fct_list <- get_names(fl)

  const_fcts <- character()
  for (i in fct_list) {
    if (get_keep(fl, fct_list[i])) const_fcts <- c(const_fcts, fct_list[i])
  }

  # extract body and replacing
  # ============================================================================
  body <- body_of_fct(f)
  to_be_replaced <- NULL
  counter <- 1
  to_replace <- NULL
  to_replace_index <- NULL
  all_vars_left <- NULL
  all_vars_right <- NULL
  body_new <- as.list(body)
  counter_bn <- 1

  for (i in seq_along(1:length(body))) {
    in_if <- FALSE
    v <- Vars$new(const_fcts)
    ast <- v$find_vars(body[[i]])

    if (body[[i]][[1]] == as.name("if")) {
      in_if <- TRUE
    }
    if (body[[i]][[1]] == as.name("return")) next

    ls <- v$get_ls()
    all_vars_left <- c(all_vars_left, ls)
    rs <- v$get_rs()

    if (identical(rs, list())) next
    independet_at_rs <- sapply(rs, function(y) {
      tocheck <- rbfv(y)
      if (tocheck == x) {
        return(y)
      }
    })
    independet_at_rs[sapply(independet_at_rs, is.null)] <- NULL
    all_vars_right <- c(all_vars_right, independet_at_rs)

    # find variables which have to be replaced in the following lines
    if (!in_if) {
      for (j in seq_along(1:length(rs))) {
        tocheck_right <- rbfv(rs[[j]])
        tocheck_left <- rbfv(ls[[1]])
        if (check(tocheck_right, x, y) && tocheck_left != y) {
          to_be_replaced[[counter]] <- c(
            index = i, to_be_replaced = ls,
            replace_with = deparse(body[[i]][[3]])
          )
          counter <- counter + 1
          break
        }
      }
    }

    # replace directly
    index <- to_be_replaced[[counter - 1]][[1]]
    if (is.null(index)) next else index <- index + 1
    temp_body <- as.list(body_new)[index:length(body_new)]
    replace <- to_be_replaced[[counter - 1]][[2]]
    replace_with <- to_be_replaced[[counter - 1]][[3]]
    tmp_bd <- replace_all(temp_body, replace, replace_with)
    for (k in seq_along(tmp_bd)) {
      tb <- tmp_bd[[k]]
      bn <- body_new[[index + k - 1]]
      if (!identical(tb, bn)) {
        body_new[[index + k - 1]] <- tb
      }
    }

    # check if block
    v <- Vars$new(const_fcts)
    ast <- v$find_vars(body_new[[i]])
    ls <- v$get_ls()
    if (in_if) {
      for (j in seq_along(1:length(ls))) {
        tocheck <- rbfv(ls[[j]])
        if (tocheck != y) stop("Only the dependent variable (Input argument = y) is allowed at the left side")
      }
    }
  } # end loop over body

  body <- body_new

  # jac matrix creation
  # ============================================================================
  all_vars_left <- as.character(c(all_vars_left, lapply(formalArgs(f), str2lang)))
  jac_mat <- "jac_mat"
  while (jac_mat %in% all_vars_left) {
    jac_mat <- paste0(jac_mat, "1")
  }

  # find all independent variables for which 'd' has to be called
  # ============================================================================
  to_diff <- unique(all_vars_right)

  # call dfdr::d for each term at right hand side
  # ============================================================================
  body_new <- list()
  if (is.null(num_functions)) {
    body_new[[1]] <- str2lang(paste(
      noquote(jac_mat),
      "<- matrix(0, length(",
      x, "),",
      "length(",
      x, "))"
    ))
  } else {
    body_new[[1]] <- str2lang(paste(
      noquote(jac_mat),
      "<- matrix(0,", num_functions,
      ",length(", x, "))"
    ))
  }
  counter <- 2
  ret_found <- FALSE
  for (i in seq_along(1:length(body))) {
    in_if <- FALSE
    if (body[[i]][[1]] == as.name("if")) {
      in_if <- TRUE
    }
    if (body[[i]][[1]] == as.name("return")) {
      ret_found <- TRUE
      body_new[[counter]] <- bquote(return(.(str2lang(jac_mat))))
      counter <- counter + 1
      next
    }

    codeline <- NULL
    if (in_if) {
      u <- Unfold$new(diff, fl, to_diff, y, jac_mat)
      ast <- u$uf(body[[i]])
      codeline <- u$get_code(ast)
      body_new[[counter]] <- codeline
      counter <- counter + 1
    } else {
      codeline <- body[[i]]
      deriv <- diff(codeline[[2]], codeline[[3]], to_diff, y, fl, jac_mat)
      body_new[[counter]] <- codeline
      counter <- counter + 1
      if (!is.null(deriv)) {
        for (j in seq_along(1:length(deriv))) {
          body_new[[counter]] <- deriv[[j]]
          counter <- counter + 1
        }
      }
    }
  }
  if (ret_found == FALSE) {
    body_new[[counter]] <- bquote(return(.(str2lang(jac_mat))))
    counter <- counter + 1
  }

  body_new[sapply(body_new, is.null)] <- NULL
  body_new <- as.call(c(as.symbol("{"), body_new))
  body(f) <- body_new

  return(f)
}