conjugate_gradient.cc

#include <Eigen/Core>
#include <Eigen/Sparse>
#include <iostream>

using Vec = Eigen::VectorXd;
using Mat = Eigen::SparseMatrix<double,Eigen::RowMajor>;

void jacobi(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration);
void sor(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration,double omega);
void conjugate_gradient(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration);

int main(int argc,char **argv){
  enum opt_e{
    JACOBI,
    SOR,
    CG
  };

  if(argc < 6){
    std::cout << argv[0] << " [matrix size] [parameter c] [epsilon] [max iteration] [solver option]" << std::endl;
    std::cout << "solver option:" << std::endl;
    std::cout << "jacobi:" << JACOBI << std::endl;
    std::cout << "sor:" << SOR << std::endl;
    std::cout << "cg:" << CG << std::endl;
    return -1;
  }

  int n = std::stoi(argv[1]);
  double c = std::stod(argv[2]);
  double eps = std::stod(argv[3]);
  int max_iteration = std::stoi(argv[4]);

  Vec x(n),b(n);
  Mat mat(n,n);
  for(int i=0;i<n;++i){
    x[i] = 0.0;
    b[i] = 1.0;
    mat.insert(i,i) = c;
  }
  for(int i=0;i<n-1;++i){
    mat.insert(i+1,i) = -1.0;
    mat.insert(i,i+1) = -1.0;
  }
  
  opt_e opt = (opt_e)std::stoi(argv[5]);
  switch (opt) {
    case JACOBI:
      jacobi(mat,b,x,eps,max_iteration);
      break;
    case SOR:
      {
        if(argc<7){
          std::cout << "add omega parameter." << std::endl;
          return -1;
        }
        double omega = std::stod(argv[6]);
        sor(mat,b,x,eps,max_iteration,omega);
        break;
      }
    case CG:
      conjugate_gradient(mat,b,x,eps,max_iteration);
      break;
    default:
      break;
  }

  return 0;
}

void jacobi(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration){
  int n = mat.rows();
  for(int k=0;k<max_iteration;++k){
    Vec pre_x = x;
    for(int i=0;i<n;++i){
      x[i] = (-mat.row(i)*pre_x+mat.coeff(i,i)*pre_x[i]+b[i])/mat.coeff(i,i); // 対角要素は必ず非ゼロと仮定
    }
    double rnorm = (mat*x-b).norm();
    std::cout << k+1 << " " << rnorm << std::endl;
    if( rnorm < eps ) break;
  }
}

void sor(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration,double omega){
  int n = mat.rows();
  Vec y(n);
  for(int k=0;k<max_iteration;++k){
    for(int i=0;i<n;++i){
      y[i] = 0;
      Mat::InnerIterator it(mat,i);   // sparse matrixなので，イテレータを使って要素にアクセスする
      for(;it;++it){
        if(it.col()==i) continue;
        y[i] -= it.value()*x[it.col()];
      }
      y[i] += b[i];
      y[i] /= mat.coeff(i,i); // 対角要素は必ず非ゼロと仮定

      x[i] = x[i] + omega*(y[i] - x[i]);
    }
    double rnorm = (mat*x-b).norm();
    std::cout << k+1 << " " << rnorm << std::endl;
    if( rnorm < eps ) break;
  }
}

void conjugate_gradient(const Mat &mat,const Vec &b,Vec &x,double eps,int max_iteration){
  Vec r = b , p = r;
  double rnorm = r.norm();
  for(int k=0;k < max_iteration&&rnorm >= eps;++k){
    Vec mul = mat*p;

    double quad = p.dot(mul);
    double alpha = r.dot(p)/quad;

    x = x + alpha*p;
    r = r - alpha*mul;
    rnorm = r.norm();
    std::cout << k+1 << " " << rnorm << std::endl;

    double beta = -r.dot(mul)/quad;
    p = r +beta*p;
  } 
}