R/nls.R

In mikldk/staid: Start Aid for Fitting Certain Statistical Models

#' Prepare nls start aid
#' 
#' Names must be given such that the order is known 
#' when providing a point later in optimisation.
#' 
#' @importFrom stats nls
#' 
#' @export
prepare_nls <- function(formula, data, var_names) {
  stopifnot(is(formula, "formula"))
  stopifnot(is.data.frame(data))
  
  lhs <- as.expression(as.list(formula)[[2]])
  rhs <- as.expression(as.list(formula)[[3]])
  
  formula_names <- all.vars(rhs)
  var_names_detect <- setdiff(formula_names, colnames(data))
  
  if (!isTRUE(all.equal(sort(var_names_detect), 
                        sort(var_names)))) {
    warning("var_names not the same as the difference between names in formula and column names")
  }
  
  r_grad_expr <- vector("list", length(var_names))
  for (i in seq_along(var_names)) {
    r_grad_expr[[i]] <- as.expression(D(rhs, var_names[i]))
  }
  
  pars_to_list <- function(pars) {
    pars_lst <- as.list(pars)
    names(pars_lst) <- var_names
    pars_lst
  }
  
  phi <- function(pars) {
    pars_lst <- pars_to_list(pars)
    
    with(data, eval(rhs, pars_lst))
  }
  
  r <- function(pars) {
    pars_lst <- pars_to_list(pars)
    
    ypred <- phi(pars)
    y <- with(data, eval(lhs, pars_lst))
    
    ypred - y
  }
  
  r_grad <- function(pars) {
    pars_lst <- pars_to_list(pars)
    
    gvec <- vector("list", length(pars))
    
    for (i in seq_along(pars)) {
      # i <- 1
      
      gvec[[i]] <- with(data, eval(r_grad_expr[[i]], pars_lst))
    }
    
    return(gvec)
  }
  
  f <- function(pars) {
    rvec <- r(pars)
    sum(rvec^2)
  }
  
  g <- function(pars) {
    rvec <- r(pars)
    J <- do.call(cbind, r_grad(pars))
    
    t(J) %*% rvec
  }
  
  ret <- list(
    var_names = var_names,
    f = f, 
    g = g)
  
  return(ret)
}