R/solve_ols.R

In DavidJamesKent/stsci6520_hw2: STSCI6520 HW2

#' Solves Ax=b using Gauss-Seidel or Jacobi iteration.
#'
#' @param A \code{n x n} matrix.
#' @param b Length \code{n} vector.
#' @param x0 (Optional) length \code{n} initial value for iterations. All zeroes if not provided.
#' @param iter (Optional) number of iterations to perform.
#' @param method (Optional) method, one of "Gauss-Seidel" or "Jacobi".
#' @param ncores (Optional) number of cores to use. Ignored if method is "Gauss-Seidel".
#' @param verbose (Optional). If \code{TRUE}, print information about progress to the screen.
#' @return Value of \code{x} after \code{iter} iterations of the method.
#' @export
#' @import doParallel
#'
solve_ols <- function(A,b,x0=NULL,niter=100,method="Gauss-Seidel",ncores=1,verbose=T) {

  n <- nrow(A)

  if( is.null(x0) )
    x <- rep(0,n)
  else
    x <- x0

  if( method=="Jacobi" && ncores > 1 ) {
    # Parallel Jacobi
    if(verbose)
      cat(paste("Performing",method,"iteration.\n"))
    cl <- makeCluster(ncores)
    registerDoParallel(cl,cores=ncores)

    for(j in 1:niter) {
      if(verbose) {
        if((j %% 100) == 0)
          cat(paste("Iter ",j,"/",niter,"\n",sep=""))
      }
      x <- foreach(i=1:n,.combine="c") %dopar% {
        (b[i] - A[i,-i]%*% x[-i])/A[i,i]
      }
    }
    stopCluster(cl)
  } else if( method=="Jacobi" ) {
    # Serial Jacobi
    if(verbose)
      cat(paste("Performing",method,"iteration.\n"))

    for(j in 1:niter) {
      if(verbose) {
        if((j %% 100) == 0)
          cat(paste("Iter ",j,"/",niter,"\n",sep=""))
      }
      xnew <- rep(NA,n)
      for(i in 1:n) {
        xnew[i] <- (b[i] - A[i,-i]%*%x[-i])/A[i,i]
      }
      x <- xnew
    }
  } else if( method=="Gauss-Seidel" ) {
    # (Serial) Gauss-Seidel
    if(verbose)
      cat(paste("Performing",method,"iteration.\n"))
    for(j in 1:niter) {
      if(verbose) {
        if((j %% 100) == 0)
          cat(paste("Iter ",j,"/",niter,"\n",sep=""))
      }
      for(i in 1:n) {
        x[i] <- (b[i] - A[i,-i]%*%x[-i])/A[i,i]
      }
    }
  } else
    cat(paste("Method",method,"not recognized.\n"))
  return(x)
}