R/solve_ols.R

In sophiaycl/6520HW2: STSCI6520HW

#' Solve linear system
solve_ols = function(A, b,   iterNum = 10000, type ='GS',ncores = 1){



  JCSeqIter = function(JC, invD, x0, iterNum, b){
    xJTemp = x0
    b0 = invD %*% b
    for (i in 1:iterNum) {
      xJ = JC %*% xJTemp + b0
      xJTemp = xJ
    }
    xJ
  }
  GSIter = function(GS, invDL, x0, iterNum, b){
    xGTemp = x0
    b0 = invDL %*% b
    for (i in 1:iterNum) {
      xG = GS %*% xGTemp + b0
      xGTemp = xG
    }
    xG
  }
  JCParaIter = function(A, x0, iterNum, b, ncores){
    cl <- parallel::makeCluster(mc <- getOption("cl.cores", ncores))  # Establish cluster of detectCores()-1 child processes
    xJTemp = x0
    for (i in 1:iterNum) {
      update_xJ<-function(k){  #k-th chain
        x = (b[k] - crossprod(A[k,-k], xJTemp[-k]))/A[k,k]
      }
      xJ = parallel::parLapply(cl=cl,X=1:length(xJTemp),fun=update_xJ)
      xJTemp = xJ
    }
    parallel::stopCluster(cl)
    rm(cl)
    xJ
  }

  n = length(b)

  D = diag(diag(A))
  L = matrix(data = 0, nrow = n, ncol = n)
  U = matrix(data = 0, nrow = n, ncol = n)
  L[lower.tri(L)] = A[lower.tri(A)]
  U[upper.tri(U)] = A[upper.tri(A)]
  x0 = array(data = 0, dim = n)

  if (type == 'JC') {
    if (ncores > 1) {
      x = JCParaIter(A, x0, iterNum, b, ncores)
    } else {
      invD = solve(D)
      Jacobi = - invD %*% (L + U)
      x = JCSeqIter(Jacobi, invD, x0, iterNum, b)
    }
  } else {
    invDL = solve(D + L)
    GS = - invDL %*% U
    x = GSIter(GS, invDL, x0, iterNum, b)
  }
  x
}
#' @param A Matrix.
#' @param b Vector.
#' @param ncores Number of cores
#' @param type Type of iteration method
#' @param iterNum Number of iterations
#' @return The solution vector
#' @export