Iterative lsmr solver does not work

[ncooder]

I noticed that (at least) for the electricity example described here:
https://cran.r-project.org/web/packages/stR/vignettes/stRvignette.html

The iterative lsmr solver returns an error:

Error in optim(par = log(initP), fn = f, method = ifelse(length(initP) >  (compare.R#107): function cannot be evaluated at initial parameters

I checked in detail and for this test with dummy data:

library(Matrix)
library(foreach)
library(stR)

create_dummy_data <- function(n_samples = 100, n_features = 5, n_folds = 3) {
  set.seed(42)

  trainData <- list()
  trainC <- list()
  for (i in seq_len(n_folds)) {
    y <- seq(0, 100, length.out = n_samples)
    nan_indices <- sample(seq_len(n_samples), 5, replace = FALSE)
    y[nan_indices] <- NA_real_
    trainData[[i]] <- y

    C <- matrix(rnorm(n_samples * n_features), nrow = n_samples, ncol = n_features)
    C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
    trainC[[i]] <- as(C, "CsparseMatrix")
  }

  fcastData <- list()
  fcastC <- list()
  for (i in seq_len(n_folds)) {
    y <- seq(0, 10, length.out = n_samples / 10)
    fcastData[[i]] <- y

    C <- matrix(rnorm((n_samples/10) * n_features),
                nrow = n_samples/10, ncol = n_features)
    C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
    fcastC[[i]] <- as(C, "CsparseMatrix")
  }

  completeData <- seq(0, 200, length.out = n_samples * 2)
  C <- matrix(rnorm((n_samples * 2) * n_features),
              nrow = n_samples * 2, ncol = n_features)
  C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
  completeC <- as(C, "CsparseMatrix")

  lambdas <- list()
  for (i in seq_len(n_folds)) {
    if (i == 1) {
      lambdas[[i]] <- list(lambdas = c(0.1, 0, 0))
    } else {
      lambdas[[i]] <- list(lambdas = c(0, 0, 0))
    }
  }

  seats <- list()
  for (i in seq_len(n_folds)) {
    if (i == 1) {
      seats[[i]] <- list(lengths = c(n_features, 0, 0))
    } else {
      seats[[i]] <- list(lengths = c(0, 0, 0))
    }
  }

  regMatrix <- Diagonal(n_features, x = 0.1)

  list(
    trainData    = trainData,
    fcastData    = fcastData,
    completeData = completeData,
    trainC       = trainC,
    fcastC       = fcastC,
    completeC    = completeC,
    regMatrix    = regMatrix,
    seats        = seats,
    lambdas      = lambdas
  )
}

test_nFoldSTRCV <- function() {
  test_data <- create_dummy_data()

  solver_configs <- list(
    list(type = "Matrix",    method = "cholesky"),
    list(type = "iterative", method = "cg-chol"),
    list(type = "iterative", method = "lsmr-chol"),
    list(type = "iterative", method = "lsmr")
  )

  iterControl <- list(maxiter = 1000, tol = 1e-6)

  for (solver in solver_configs) {
    cat("\nTesting solver:", solver$type, solver$method, "\n")
    result <- tryCatch({
      stR:::nFoldSTRCV(
        n            = 3,
        trainData    = test_data$trainData,
        fcastData    = test_data$fcastData,
        completeData = test_data$completeData,
        trainC       = test_data$trainC,
        fcastC       = test_data$fcastC,
        completeC    = test_data$completeC,
        regMatrix    = test_data$regMatrix,
        regSeats     = test_data$seats,
        lambdas      = test_data$lambdas,
        solver       = solver,
        trace        = TRUE,
        iterControl  = iterControl
      )
    }, error = function(e) e)

    if (inherits(result, "error")) {
      cat("Error with solver", solver$type, solver$method, ":\n", result$message, "\n")
    } else {
      cat("Error:", result, "\n")
    }
  }
}

test_nFoldSTRCV()

I get this error:

Testing solver: Matrix cholesky 
Error: 1519.23 

Testing solver: iterative cg-chol 

Iter: 5 Error: 1.688701e-29   
Iter: 5 Error: 8.062725e-32   
Iter: 5 Error: 9.078954e-33   Error: 1519.23 

Testing solver: iterative lsmr-chol 

Iterations: 5   
Iterations: 5   
Iterations: 5   Error: 1519.23 

Testing solver: iterative lsmr 
Error in colSums(X^2) : 
  argument 'x' must be an array with at least two dimensions

Error in lmSolver...
Error in colSums(X^2) : 
  argument 'x' must be an array with at least two dimensions

Error in lmSolver...
Error in colSums(X^2) : 
  argument 'x' must be an array with at least two dimensions

Error in lmSolver...
Error: Inf 

However when defining lmSolver and nFoldSTRCV externally:

library(Matrix)
library(foreach)
library(stR)

lmSolver <- function(X, y, type = "Matrix", method = "cholesky", env = NULL, iterControl = list(maxiter = 20, tol = 1E-6), trace = TRUE) {
  if (length(y) > nrow(X)) stop("y is too long in lmSolver...")
  y <- c(y, rep(0, nrow(X) - length(y)))
  if (type == "Matrix") {
    if (method == "cholesky") {
      result <- .solve.dgC.chol(t(X), y)
      return(result$coef)
    }
    if (method == "qr") {
      b <- solve(qr(X), y) # Not optimal
      return(b)
    }
    stop("Unknown method in lmSolver...")
  }
  if (type == "iterative") {
    if (method == "cg") {
      f <- function(z) crossprod(X, X %*% z)
      if (is.null(env) || is.null(env$L) || is.null(env$Lt) || is.null(env$P) || is.null(env$Pt)) {
        invm <- 1 / colSums(X^2)
        invf <- function(z) invm * z
      } else {
        invf <- function(z) solve(env$P, solve(env$Lt, solve(env$L, solve(env$Pt, z))))
      }
      Xty <- crossprod(X, y)
      if (is.null(env) || is.null(env$b0)) {
        b0 <- rep(0, length(Xty))
      } else {
        b0 <- env$b0
      }
      result <- stR:::olscg(FUN = f, y = Xty, b = b0, invFUN = invf, iterControl = iterControl, trace = trace)
      return(result$b)
    }
    if (method == "cg-chol") {
      if (is.null(env)) {
        stop("env variable is NULL.")
      }
      if (is.null(env$L) ||
        is.null(env$Lt) ||
        is.null(env$P) ||
        is.null(env$Pt)) {
        result <- .solve.dgC.chol(t(X), y)
        eL <- Matrix::expand(result$L)
        env$L <- eL$L
        env$Lt <- t(eL$L)
        env$P <- eL$P
        env$Pt <- t(eL$P)
        env$b0 <- result$coef
        return(result$coef)
      } else {
        f <- function(z) crossprod(X, X %*% z)
        invf <- function(z) solve(env$P, solve(env$Lt, solve(env$L, solve(env$Pt, z))))
        Xty <- crossprod(X, y)
        if (is.null(env$b0)) {
          b0 <- rep(0, length(Xty))
        } else {
          b0 <- env$b0
        }
        result <- stR:::olscg(FUN = f, y = Xty, b = b0, invFUN = invf, iterControl = iterControl, trace = trace)
        return(result$b)
      }
    }
    if (method == "lsmr-chol") {
      if (is.null(env)) {
        stop("env variable is NULL.")
      }
      if (is.null(env$L) ||
        is.null(env$Lt) ||
        is.null(env$P) ||
        is.null(env$Pt)) {
        result <- .solve.dgC.chol(t(X), y)
        eL <- Matrix::expand(result$L)
        env$L <- eL$L
        env$Lt <- t(eL$L)
        env$P <- eL$P
        env$Pt <- t(eL$P)
        env$b0 <- result$coef
        return(result$coef)
      } else {
        A <- function(x, k) {
          if (k == 1) {
            return(X %*% solve(env$P, solve(env$Lt, x)))
          } else {
            return(solve(env$L, solve(env$Pt, crossprod(X, x))))
          }
        }
        if (!is.null(env$b0)) {
          x0 <- env$b0
        } else {
          x0 <- 0
        }
        result <- stR:::lsmr(A = A, b = y - X %*% x0, atol = iterControl$tol, btol = iterControl$tol, itnlim = iterControl$maxiter)
        if (trace) {
          cat("\nIterations: ")
          cat(result$itn)
          cat("   ")
        }
        return(solve(env$P, solve(env$Lt, result$x)) + x0)
      }
    }
    if (method == "lsmr") {
      D <- sqrt(colSums(X^2))
      invD <- 1 / D
      A <- function(x, k) {
        if (k == 1) {
          return(X %*% (invD * x))
        } else {
          return(invD * crossprod(X, x))
        }
      }
      if (!is.null(env$b0)) {
        x0 <- env$b0
      } else {
        x0 <- rep(0, ncol(X))
      }
      result <- stR:::lsmr(A = A, b = y - X %*% x0, atol = iterControl$tol, btol = iterControl$tol, itnlim = iterControl$maxiter)
      if (trace) {
        cat("\nIterations: ")
        cat(result$itn)
        cat("   ")
      }
      output <- invD * result$x + x0
      if (!is.null(env)) {
        env$b0 <- output
      }
      return(output)
    }
    stop("Unknown method in lmSolver...")
  }
  stop("Unknown type in lmSolver...")
}


nFoldSTRCV <- function(n,
                       trainData, fcastData, completeData = NULL, # parameter completeData is required only for iterative methods
                       trainC, fcastC, completeC,
                       regMatrix, regSeats, lambdas,
                       solver = c("Matrix", "cholesky"), trace = FALSE, iterControl = list(maxiter = 20, tol = 1E-6)) {
  SSE <- 0
  l <- 0
  lm <- stR:::lambdaMatrix(lambdas, regSeats)
  R <- lm %*% regMatrix

  e <- NULL
  if (solver[1] == "iterative") {
    if (solver[2] %in% c("cg-chol", "lsmr-chol", "lsmr")) {
      e <- new.env(parent = .GlobalEnv)
    }
    if (solver[2] %in% c("cg-chol", "lsmr-chol")) {
      noNA <- !is.na(completeData)
      y <- completeData[noNA]
      C <- completeC[noNA, ]
      X <- rbind(C, R)
      coef0 <- try(lmSolver(X, y, type = solver[1], method = solver[2], env = e, iterControl = iterControl, trace = trace), silent = !trace)
      if ("try-error" %in% class(coef0)) {
        if (trace) cat("\nError in lmSolver... iterative solvers without preconditioners will be used...\n")
      }
    }
  }

  resultList <- list()
  # for(i in rev(1:n)) {
  resultList <- foreach(i = 1:n) %dopar% {
    noNA <- !is.na(trainData[[i]])
    y <- (trainData[[i]])[noNA]
    C <- (trainC[[i]])[noNA, ]
    X <- rbind(C, R)
    coef <- try(lmSolver(X, y, type = solver[1], method = solver[2], env = e, iterControl = iterControl, trace = trace), silent = !trace)
    if ("try-error" %in% class(coef)) {
      if (trace) cat("\nError in lmSolver...\n")
      # return(Inf)
      # next
      c(SSE = Inf, l = 1)
      # c(SSE = 0, l = 0)
    } else {
      fcast <- fcastC[[i]] %*% coef
      resid <- fcastData[[i]] - as.vector(fcast)
      # resultList[[length(resultList) + 1]] = c(SSE = sum(resid^2, na.rm = TRUE), l = sum(!is.na(resid)))
      c(SSE = sum(resid^2, na.rm = TRUE), l = sum(!is.na(resid)))
    }
  }
  for (i in seq_along(resultList)) {
    SSE <- SSE + resultList[[i]][1]
    l <- l + resultList[[i]][2]
  }
  if (l == 0) {
    return(Inf)
  }
  return(SSE / l)
}


create_dummy_data <- function(n_samples = 100, n_features = 5, n_folds = 3) {
  set.seed(42)

  trainData <- list()
  trainC <- list()
  for (i in seq_len(n_folds)) {
    y <- seq(0, 100, length.out = n_samples)
    nan_indices <- sample(seq_len(n_samples), 5, replace = FALSE)
    y[nan_indices] <- NA_real_
    trainData[[i]] <- y

    C <- matrix(rnorm(n_samples * n_features), nrow = n_samples, ncol = n_features)
    C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
    trainC[[i]] <- as(C, "CsparseMatrix")
  }

  fcastData <- list()
  fcastC <- list()
  for (i in seq_len(n_folds)) {
    y <- seq(0, 10, length.out = n_samples / 10)
    fcastData[[i]] <- y

    C <- matrix(rnorm((n_samples/10) * n_features),
                nrow = n_samples/10, ncol = n_features)
    C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
    fcastC[[i]] <- as(C, "CsparseMatrix")
  }

  completeData <- seq(0, 200, length.out = n_samples * 2)
  C <- matrix(rnorm((n_samples * 2) * n_features),
              nrow = n_samples * 2, ncol = n_features)
  C <- sweep(C, 2, sqrt(colSums(C^2)), "/")
  completeC <- as(C, "CsparseMatrix")

  lambdas <- list()
  for (i in seq_len(n_folds)) {
    if (i == 1) {
      lambdas[[i]] <- list(lambdas = c(0.1, 0, 0))
    } else {
      lambdas[[i]] <- list(lambdas = c(0, 0, 0))
    }
  }

  seats <- list()
  for (i in seq_len(n_folds)) {
    if (i == 1) {
      seats[[i]] <- list(lengths = c(n_features, 0, 0))
    } else {
      seats[[i]] <- list(lengths = c(0, 0, 0))
    }
  }

  regMatrix <- Diagonal(n_features, x = 0.1)

  list(
    trainData    = trainData,
    fcastData    = fcastData,
    completeData = completeData,
    trainC       = trainC,
    fcastC       = fcastC,
    completeC    = completeC,
    regMatrix    = regMatrix,
    seats        = seats,
    lambdas      = lambdas
  )
}

test_nFoldSTRCV <- function() {
  test_data <- create_dummy_data()

  solver_configs <- list(
    list(type = "Matrix",    method = "cholesky"),
    list(type = "iterative", method = "cg-chol"),
    list(type = "iterative", method = "lsmr-chol"),
    list(type = "iterative", method = "lsmr")
  )

  iterControl <- list(maxiter = 1000, tol = 1e-6)

  for (solver in solver_configs) {
    cat("\nTesting solver:", solver$type, solver$method, "\n")
    result <- tryCatch({
      nFoldSTRCV(
        n            = 3,
        trainData    = test_data$trainData,
        fcastData    = test_data$fcastData,
        completeData = test_data$completeData,
        trainC       = test_data$trainC,
        fcastC       = test_data$fcastC,
        completeC    = test_data$completeC,
        regMatrix    = test_data$regMatrix,
        regSeats     = test_data$seats,
        lambdas      = test_data$lambdas,
        solver       = solver,
        trace        = TRUE,
        iterControl  = iterControl
      )
    }, error = function(e) e)

    if (inherits(result, "error")) {
      cat("Error with solver", solver$type, solver$method, ":\n", result$message, "\n")
    } else {
      cat("Error:", result, "\n")
    }
  }
}

test_nFoldSTRCV()

I do not get any error and all the solvers work:

Testing solver: Matrix cholesky 
Error: 1519.23 

Testing solver: iterative cg-chol 

Iter: 5 Error: 1.688701e-29   
Iter: 5 Error: 8.062725e-32   
Iter: 5 Error: 9.078954e-33   Error: 1519.23 

Testing solver: iterative lsmr-chol 

Iterations: 5   
Iterations: 5   
Iterations: 5   Error: 1519.23 

Testing solver: iterative lsmr 

Iterations: 5   
Iterations: 5   
Iterations: 5   Error: 1519.23 

I think there is some issue with importing the olscg, lsmr or lambdaMatrix internally, since they work when defining externally.