add solve_jacobi and mvp_nodiag

gagiuntoli · Jan 19, 2025 · 32c7cfa · 32c7cfa
1 parent c576d9f
commit 32c7cfa
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 0 deletions.
diff --git a/src/ellpack.cc b/src/ellpack.cc
@@ -39,6 +39,22 @@ int Ellpack::mvp(std::vector<double> &y, std::vector<double> x) const {
   return 0;
 }
 
+int Ellpack::mvp_nodiag(std::vector<double> &y, std::vector<double> x) const {
+  for (size_t row = 0; row < nrows; row++) {
+    double tmp = 0;
+    for (size_t j = 0; j < non_zeros_per_row; j++) {
+      int index = cols[row * non_zeros_per_row + j];
+      if (index < 0) break;
+
+      if (row != index) {
+        tmp += vals[row * non_zeros_per_row + j] * x[index];
+      }
+    }
+    y[row] = tmp;
+  }
+  return 0;
+}
+
 int Ellpack::solve_cg(std::vector<double> &x, std::vector<double> b) const {
   int max_iters = 10000000;
   int n = nrows;
@@ -79,6 +95,36 @@ int Ellpack::solve_cg(std::vector<double> &x, std::vector<double> b) const {
   return 0;
 }
 
+int Ellpack::solve_jacobi(std::vector<double> &x, std::vector<double> b) const {
+  int max_iters = 10000000;
+  int n = nrows;
+  std::fill(x.begin(), x.end(), 0.0);
+
+  for (int iter = 0; iter < max_iters; iter++) {
+    std::vector<double> x_new(n);
+
+    mvp_nodiag(x_new, x);
+
+    double norm = 0;
+    for (int i = 0; i < n; i++) {
+      double aii;
+      get(aii, i, i);
+      if (aii < 1.0e-10) {
+        return 2;
+      }
+      x_new[i] = (b[i] - x_new[i]) / aii;
+
+      norm += std::abs(x_new[i] - x[i]);
+    }
+
+    if (norm < 1.0e-7) break;
+
+    x = x_new;
+  }
+
+  return 0;
+}
+
 size_t Ellpack::getIndex(size_t rowTarget, size_t colTarget) const {
   size_t index = rowTarget * non_zeros_per_row;
   for (; index < (rowTarget + 1) * non_zeros_per_row; index++) {

diff --git a/src/ellpack.h b/src/ellpack.h
@@ -49,7 +49,9 @@ struct Ellpack {
   bool get(double &value, size_t row, size_t col) const;
 
   int mvp(std::vector<double> &y, std::vector<double> x) const;
+  int mvp_nodiag(std::vector<double> &y, std::vector<double> x) const;
   int solve_cg(std::vector<double> &x, std::vector<double> b) const;
+  int solve_jacobi(std::vector<double> &x, std::vector<double> b) const;
 
   std::string toString() const;
 };

diff --git a/src/ellpack_test.cc b/src/ellpack_test.cc
@@ -129,3 +129,46 @@ TEST(EllpackTest, delete_row) {
   found = matrix.get(value, 0, 2);
   EXPECT_FALSE(found);
 }
+
+TEST(EllpackTest, solve_jacobi_symmetric_system) {
+  Ellpack matrix(3, 3, 3);
+  matrix.insert(0, 0, 4.0);
+  matrix.insert(0, 1, -1.0);
+  matrix.insert(1, 0, -1.0);
+  matrix.insert(1, 1, 4.0);
+  matrix.insert(1, 2, -1.0);
+  matrix.insert(2, 1, -1.0);
+  matrix.insert(2, 2, 4.0);
+
+  std::vector<double> b = {15.0, 10.0, 10.0};
+  std::vector<double> x(3);
+  std::vector<double> b1(3);
+
+  matrix.solve_jacobi(x, b);
+  matrix.mvp(b1, x);
+
+  for (int i = 0; i < b.size(); i++) {
+    EXPECT_NEAR(b1[i], b[i], 1.0e-6);
+  }
+}
+
+TEST(EllpackTest, solve_jacobi_non_symmetric_system) {
+  Ellpack matrix(3, 3, 3);
+  matrix.insert(0, 0, 4.0);
+  matrix.insert(0, 1, -1.0);
+  matrix.insert(1, 0, -1.0);
+  matrix.insert(1, 1, 4.0);
+  matrix.insert(1, 2, -1.0);
+  matrix.insert(2, 2, 4.0);
+
+  std::vector<double> b = {15.0, 10.0, 10.0};
+  std::vector<double> x(3);
+  std::vector<double> b1(3);
+
+  matrix.solve_jacobi(x, b);
+  matrix.mvp(b1, x);
+
+  for (int i = 0; i < b.size(); i++) {
+    EXPECT_NEAR(b1[i], b[i], 1.0e-6);
+  }
+}