R/lav_optim_nlminb_constr.R

In yrosseel/lavaan: Latent Variable Analysis

# constrained optimization
# - references: * Nocedal & Wright (2006) Chapter 17
#               * Optimization with constraints by Madsen, Nielsen & Tingleff
#               * original papers: Powell, 1969 and Rockafeller, 1974
# - using 'nlminb' for the unconstrained subproblem
# - convergence scheme is based on the auglag function in the alabama package
nlminb.constr <- function(start, objective, gradient = NULL, hessian = NULL,
                          ..., scale = 1, control = list(),
                          lower = -Inf, upper = Inf,
                          ceq = NULL, ceq.jac = NULL,
                          cin = NULL, cin.jac = NULL,
                          control.outer = list()) {
  # we need a gradient
  stopifnot(!is.null(gradient))

  # if no 'ceq' or 'cin' function, we create a dummy one
  if (is.null(ceq)) {
    ceq <- function(x, ...) {
      return(numeric(0))
    }
  }
  if (is.null(cin)) {
    cin <- function(x, ...) {
      return(numeric(0))
    }
  }

  # if no user-supplied jacobian functions, create them
  if (is.null(ceq.jac)) {
    if (is.null(ceq)) {
      ceq.jac <- function(x, ...) {
        matrix(0, nrow = 0L, ncol = length(x))
      }
    } else {
      ceq.jac <- function(x, ...) {
        numDeriv::jacobian(func = ceq, x = x, ...)
      }
    }
  }
  if (is.null(cin.jac)) {
    if (is.null(cin)) {
      cin.jac <- function(x, ...) {
        matrix(0, nrow = 0L, ncol = length(x))
      }
    } else {
      cin.jac <- function(x, ...) {
        numDeriv::jacobian(func = cin, x = x, ...)
      }
    }
  }

  # how many ceq and cin constraints?
  nceq <- length(ceq(start))
  ncin <- length(cin(start))
  ncon <- nceq + ncin
  ceq.idx <- cin.idx <- integer(0)
  if (nceq > 0L) ceq.idx <- 1:nceq
  if (ncin > 0L) cin.idx <- nceq + 1:ncin
  cin.flag <- rep(FALSE, length(ncon))
  if (ncin > 0L) cin.flag[cin.idx] <- TRUE

  # control outer default values
  control.outer.default <- list(
    mu0 = 100,
    lambda0 = 10,
    tol = 1e-06, # changed this in 0.4-12
    itmax = 100L,
    verbose = FALSE
  )
  control.outer <- modifyList(control.outer.default, control.outer)


  # construct augmented lagrangian function
  auglag <- function(x, ...) {
    # apply constraints
    ceq0 <- ceq(x, ...)
    cin0 <- cin(x, ...)
    con0 <- c(ceq0, cin0)
    # 'release' inactive constraints
    if (ncin > 0L) {
      slack <- lambda / mu
      inactive.idx <- which(cin.flag & con0 > slack)
      con0[inactive.idx] <- slack[inactive.idx]
    }
    objective(x, ...) - sum(lambda * con0) + (mu / 2) * sum(con0 * con0)
  }

  fgrad <- function(x, ...) {
    # apply constraints
    ceq0 <- ceq(x, ...)
    cin0 <- cin(x, ...)
    con0 <- c(ceq0, cin0)
    # jacobian
    JAC <- rbind(ceq.jac(x, ...), cin.jac(x, ...))
    lambda.JAC <- lambda * JAC

    # handle inactive constraints
    if (ncin > 0L) {
      slack <- lambda / mu
      inactive.idx <- which(cin.flag & con0 > slack)
      if (length(inactive.idx) > 0L) {
        JAC <- JAC[-inactive.idx, , drop = FALSE]
        lambda.JAC <- lambda.JAC[-inactive.idx, , drop = FALSE]
        con0 <- con0[-inactive.idx]
      }
    }

    if (nrow(JAC) > 0L) {
      (gradient(x, ...) - colSums(lambda.JAC) +
        mu * as.numeric(t(JAC) %*% con0))
    } else {
      gradient(x, ...)
    }
  }


  # initialization
  ceq0 <- ceq(start, ...)
  cin0 <- cin(start, ...)
  con0 <- c(ceq0, cin0)
  lambda <- rep(control.outer$lambda0, length(con0))
  mu <- control.outer$mu0
  inactive.idx <- integer(0)
  if (ncin > 0L) {
    slack <- lambda / mu
    inactive.idx <- which(cin.flag & con0 > slack)
    con0[inactive.idx] <- slack[inactive.idx]
  }
  K <- max(abs(con0))
  if (control.outer$verbose) {
    cat("init cin0 values: ", cin0, "\n")
    cat("init ceq0 values: ", ceq0, "\n")
    cat("init slack values: ", lambda / mu, "\n")
    cat("init inactive idx: ", inactive.idx, "\n")
    cat("init con0 values: ", con0, "\n")
    cat("K = max con0: ", K, "\n")
  }

  r <- obj <- objective(start, ...)
  feval <- 0L
  geval <- 0L
  niter <- 0L
  ilack <- 0L
  Kprev <- K
  mu0 <- control.outer$mu0 / Kprev
  if (is.infinite(mu0)) mu0 <- 1.0
  mu <- mu0

  K <- Inf
  x.par <- start
  for (i in 1:control.outer$itmax) {
    x.old <- x.par
    r.old <- r
    ############################################################
    if (control.outer$verbose) {
      cat("\nStarting inner optimization [", i, "]:\n")
      cat("lambda: ", lambda, "\n")
      cat("mu: ", mu, "\n")
    }
    optim.out <- nlminb(
      start = x.par, objective = auglag,
      gradient = fgrad, control = control,
      lower = lower, upper = upper,
      scale = scale, ...
    )
    ############################################################
    x.par <- optim.out$par
    r <- optim.out$objective
    feval <- feval + optim.out$evaluations[1]
    geval <- geval + optim.out$evaluations[2]
    niter <- niter + optim.out$iterations

    # check constraints
    ceq0 <- ceq(x.par, ...)
    cin0 <- cin(x.par, ...)
    con0 <- c(ceq0, cin0)
    if (ncin > 0L) {
      slack <- lambda / mu
      inactive.idx <- which(cin.flag & con0 > slack)
      con0[inactive.idx] <- slack[inactive.idx]
    }
    K <- max(abs(con0))
    if (control.outer$verbose) {
      cat("cin0 values: ", cin0, "\n")
      cat("ceq0 values: ", ceq0, "\n")
      cat("active threshold: ", lambda / mu, "\n")
      cat("inactive idx: ", inactive.idx, "\n")
      cat("con0 values: ", con0, "\n")
      cat("K = max con0: ", K, " Kprev = ", Kprev, "\n")
    }

    # update K or mu (see Powell, 1969)
    if (K <= Kprev / 4) {
      lambda <- lambda - (mu * con0)
      Kprev <- K
    } else {
      mu <- 10 * mu
    }

    # check convergence
    pconv <- max(abs(x.par - x.old))
    if (pconv < control.outer$tol) {
      ilack <- ilack + 1L
    } else {
      ilack <- 0L
    }

    if ((is.finite(r) && is.finite(r.old) &&
      abs(r - r.old) < control.outer$tol && K < control.outer$tol) |
      ilack >= 3) {
      break
    }
  }

  # output
  a <- list()

  if (i == control.outer$itmax) {
    a$convergence <- 10L
    a$message <- "nlminb.constr ran out of iterations and did not converge"
  } else if (K > control.outer$tol) {
    a$convergence <- 11L
    a$message <- "Convergence due to lack of progress in parameter updates"
  } else {
    a$convergence <- 0L
    a$message <- "converged"
  }
  a$par <- optim.out$par
  a$outer.iterations <- i
  a$lambda <- lambda
  a$mu <- mu
  # a$value <- objective(a$start, ...)
  # a$cin <- cin(a$start, ...)
  # a$ceq <- ceq(a$start, ...)
  a$evaluations <- c(feval, geval)
  a$iterations <- niter
  # a$kkt1 <- max(abs(a$fgrad)) <= 0.01 * (1 + abs(a$value))
  # a$kkt2 <- any(eigen(a$hessian)$value * control.optim$objectivescale> 0)

  # jacobian of ceq and 'active' cin
  ceq0 <- ceq(a$par, ...)
  cin0 <- cin(a$par, ...)
  con0 <- c(ceq0, cin0)
  JAC <- rbind(ceq.jac(a$par, ...), cin.jac(a$par, ...))
  inactive.idx <- integer(0L)
  cin.idx <- which(cin.flag)
  # ceq.idx <- which(!cin.flag)
  if (ncin > 0L) {
    # FIXME: slack value not too strict??
    slack <- 1e-05
    # cat("DEBUG:\n"); print(con0)
    inactive.idx <- which(cin.flag & con0 > slack)
    # if(length(inactive.idx) > 0L) {
    #    JAC        <-        JAC[-inactive.idx,,drop=FALSE]
    # }
  }
  attr(JAC, "inactive.idx") <- inactive.idx
  attr(JAC, "cin.idx") <- cin.idx
  attr(JAC, "ceq.idx") <- ceq.idx
  a$con.jac <- JAC

  a
}