From 261dffcbb879f865c949490c74df6634a91e5be5 Mon Sep 17 00:00:00 2001
From: Stephane Del Pino <stephane.delpino44@gmail.com>
Date: Thu, 1 Aug 2019 11:39:12 +0200
Subject: [PATCH] Add Vector class and CRSMatrix*Vector

---
 src/algebra/CRSMatrix.hpp |  21 +++++++
 src/algebra/Vector.hpp    | 129 ++++++++++++++++++++++++++++++++++++++
 src/main.cpp              |  14 ++++-
 3 files changed, 162 insertions(+), 2 deletions(-)
 create mode 100644 src/algebra/Vector.hpp

diff --git a/src/algebra/CRSMatrix.hpp b/src/algebra/CRSMatrix.hpp
index 112ceae31..1c6b4f689 100644
--- a/src/algebra/CRSMatrix.hpp
+++ b/src/algebra/CRSMatrix.hpp
@@ -5,6 +5,8 @@
 #include <Kokkos_StaticCrsGraph.hpp>
 #include <PugsAssert.hpp>
 
+#include <Vector.hpp>
+
 #include <iostream>
 
 class FlexibleMatrix
@@ -131,6 +133,25 @@ class CRSMatrix
   Array<double> m_values;
 
  public:
+  Vector<double> operator*(const Vector<double>& x) const
+  {
+    Vector<double> Ax{m_host_matrix.numRows()};
+    auto host_row_map = m_host_matrix.row_map;
+
+    parallel_for(
+      m_host_matrix.numRows(), PUGS_LAMBDA(const size_t& i_row) {
+        const auto row_begin = host_row_map(i_row);
+        const auto row_end   = host_row_map(i_row + 1);
+        double sum{0};
+        for (size_t j = row_begin; j < row_end; ++j) {
+          sum += m_values[j] * x[m_host_matrix.entries(j)];
+        }
+        Ax[i_row] = sum;
+      });
+
+    return Ax;
+  }
+
   friend std::ostream&
   operator<<(std::ostream& os, const CRSMatrix& M)
   {
diff --git a/src/algebra/Vector.hpp b/src/algebra/Vector.hpp
new file mode 100644
index 000000000..bed915b78
--- /dev/null
+++ b/src/algebra/Vector.hpp
@@ -0,0 +1,129 @@
+#ifndef VECTOR_HPP
+#define VECTOR_HPP
+
+#include <PugsMacros.hpp>
+#include <PugsUtils.hpp>
+
+#include <PugsAssert.hpp>
+
+#include <Array.hpp>
+
+template <typename DataType>
+class Vector
+{
+ public:
+  using data_type  = DataType;
+  using index_type = size_t;
+
+ private:
+  Array<DataType> m_values;
+  static_assert(std::is_same_v<typename decltype(m_values)::index_type, index_type>);
+
+ public:
+  friend std::ostream&
+  operator<<(std::ostream& os, const Vector<DataType>& x)
+  {
+    for (index_type i = 0; i < x.size(); ++i) {
+      os << i << " : " << x[i] << '\n';
+    }
+    return os;
+  }
+
+  PUGS_INLINE
+  DataType
+  operator,(const Vector& y)
+  {
+    DataType sum = 0;
+
+    // Does not use parallel_for to preserve sum order
+    for (index_type i = 0; i < this->size(); ++i) {
+      sum += m_values[i] + y.m_values[i];
+    }
+
+    return sum;
+  }
+
+  PUGS_INLINE
+  Vector&
+  operator/=(const DataType& a)
+  {
+    const DataType inv_a = 1. / a;
+    return (*this) *= inv_a;
+  }
+
+  PUGS_INLINE
+  Vector&
+  operator*=(const DataType& a)
+  {
+    parallel_for(
+      this->size(), PUGS_LAMBDA(const index_type& i) { m_values[i] *= a; });
+
+    return *this;
+  }
+
+  PUGS_INLINE
+  Vector&
+  operator-=(const Vector& y)
+  {
+    Assert(this->size() == y.size());
+
+    parallel_for(
+      this->size(), PUGS_LAMBDA(const index_type& i) { m_values[i] -= y.m_values[i]; });
+
+    return *this;
+  }
+
+  PUGS_INLINE
+  Vector&
+  operator+=(const Vector& y)
+  {
+    Assert(this->size() == y.size());
+
+    parallel_for(
+      this->size(), PUGS_LAMBDA(const index_type& i) { m_values[i] += y.m_values[i]; });
+
+    return *this;
+  }
+
+  PUGS_INLINE
+  DataType& operator[](const index_type& i) const noexcept(NO_ASSERT)
+  {
+    return m_values[i];
+  }
+
+  PUGS_INLINE
+  size_t
+  size() const noexcept
+  {
+    return m_values.size();
+  }
+
+  PUGS_INLINE Vector&
+  operator=(const DataType& value) noexcept
+  {
+    m_values.fill(value);
+  }
+
+  template <typename DataType2>
+  PUGS_INLINE Vector&
+  operator=(const Vector<DataType2>& vector) noexcept
+  {
+    m_values = vector.m_values;
+    return *this;
+  }
+
+  PUGS_INLINE
+  Vector& operator=(const Vector&) = default;
+
+  PUGS_INLINE
+  Vector& operator=(Vector&&) = default;
+
+  Vector(const Vector&) = default;
+
+  Vector(Vector&&) = default;
+
+  Vector(const size_t& size) : m_values(size) {}
+  ~Vector() = default;
+};
+
+#endif   // VECTOR_HPP
diff --git a/src/main.cpp b/src/main.cpp
index 00de1f4cc..bcd9161bd 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -51,9 +51,19 @@ main(int argc, char* argv[])
       }
       std::cout << F << '\n';
 
-      CRSMatrix M{F};
+      CRSMatrix A{F};
 
-      std::cout << M << '\n';
+      std::cout << A << '\n';
+
+      Vector<double> x{10};
+      for (size_t i = 0; i < x.size(); ++i) {
+        x[i] = 2 * i + 1;
+      }
+
+      std::cout << "x=\n" << x << '\n';
+
+      Vector<double> Ax = A * x;
+      std::cout << "Ax=\n" << Ax << '\n';
     }
 
     Kokkos::finalize();
-- 
GitLab