diff --git a/Src/Base/AMReX_SmallMatrix.H b/Src/Base/AMReX_SmallMatrix.H index add7c2fa88..4fe673816d 100644 --- a/Src/Base/AMReX_SmallMatrix.H +++ b/Src/Base/AMReX_SmallMatrix.H @@ -20,16 +20,16 @@ namespace amrex { /** * \brief Matrix class with compile-time size * - * The starting index for both rows and columns is always zero. Also - * note that column vectors and row vectors are special cases of a + * Note that column vectors and row vectors are special cases of a * Matrix. * * \tparam T Matrix element data type. * \tparam NRows Number of rows. * \tparam NCols Number of columns. * \tparam ORDER Memory layout order. Order::F (i.e., column-major) by default. + * \tparam StartIndex Starting index. Either 0 or 1. */ - template + template struct SmallMatrix { using value_type = T; @@ -37,6 +37,7 @@ namespace amrex { static constexpr int row_size = NRows; static constexpr int column_size = NCols; static constexpr Order ordering = ORDER; + static constexpr int starting_index = StartIndex; /** * \brief Default constructor @@ -78,10 +79,10 @@ namespace amrex { explicit SmallMatrix (std::initializer_list> const& init) { AMREX_ASSERT(NRows == init.size()); - int i = 0; + int i = StartIndex; for (auto const& row : init) { AMREX_ASSERT(NCols == row.size()); - int j = 0; + int j = StartIndex; for (auto const& x : row) { (*this)(i,j) = x; ++j; @@ -93,6 +94,11 @@ namespace amrex { //! Returns a const reference to the element at row i and column j. [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE const T& operator() (int i, int j) const noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + --j; + } AMREX_ASSERT(i < NRows && j < NCols); if constexpr (ORDER == Order::F) { return m_mat[i+j*NRows]; @@ -104,6 +110,11 @@ namespace amrex { //! Returns a reference to the element at row i and column j. [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE T& operator() (int i, int j) noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + --j; + } AMREX_ASSERT(i < NRows && j < NCols); if constexpr (ORDER == Order::F) { return m_mat[i+j*NRows]; @@ -116,6 +127,10 @@ namespace amrex { template = 0> [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE const T& operator() (int i) const noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + } AMREX_ASSERT(i < NRows*NCols); return m_mat[i]; } @@ -124,6 +139,10 @@ namespace amrex { template = 0> [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE T& operator() (int i) noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + } AMREX_ASSERT(i < NRows*NCols); return m_mat[i]; } @@ -132,6 +151,10 @@ namespace amrex { template = 0> [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE const T& operator[] (int i) const noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + } AMREX_ASSERT(i < NRows*NCols); return m_mat[i]; } @@ -140,6 +163,10 @@ namespace amrex { template = 0> [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE T& operator[] (int i) noexcept { + static_assert(StartIndex == 0 || StartIndex == 1); + if constexpr (StartIndex == 1) { + --i; + } AMREX_ASSERT(i < NRows*NCols); return m_mat[i]; } @@ -174,7 +201,7 @@ namespace amrex { //! Set all elements in the matrix to the given value AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix& + SmallMatrix& setVal (T val) { for (auto& x : m_mat) { x = val; } @@ -185,30 +212,32 @@ namespace amrex { template = 0> static constexpr AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix + SmallMatrix Identity () noexcept { - SmallMatrix I{}; - constexpr_for<0,NRows>([&] (int i) { I(i,i) = T(1); }); + static_assert(StartIndex == 0 || StartIndex == 1); + SmallMatrix I{}; + constexpr_for( + [&] (int i) { I(i,i) = T(1); }); return I; } //! Returns a matrix initialized with zeros static constexpr AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix + SmallMatrix Zero () noexcept { - SmallMatrix Z{}; + SmallMatrix Z{}; return Z; } //! Returns transposed matrix [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix + SmallMatrix transpose () const { - SmallMatrix r; - for (int j = 0; j < NRows; ++j) { - for (int i = 0; i < NCols; ++i) { + SmallMatrix r; + for (int j = StartIndex; j < NRows+StartIndex; ++j) { + for (int i = StartIndex; i < NCols+StartIndex; ++i) { r(i,j) = (*this)(j,i); } } @@ -218,11 +247,12 @@ namespace amrex { //! Transposes a square matrix in-place. template = 0> AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix& + SmallMatrix& transposeInPlace () { - for (int j = 1; j < NCols; ++j) { - for (int i = 0; i < j; ++i) { + static_assert(StartIndex == 0 || StartIndex == 1); + for (int j = 1+StartIndex; j < NCols+StartIndex; ++j) { + for (int i = StartIndex; i < j; ++i) { amrex::Swap((*this)(i,j), (*this)(j,i)); } } @@ -257,14 +287,14 @@ namespace amrex { T trace () const { T t = 0; - constexpr_for<0,MM>([&] (int i) { t += (*this)(i,i); }); + constexpr_for([&] (int i) { t += (*this)(i,i); }); return t; } //! Operator += performing matrix addition as in (*this) += rhs AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix& - operator += (SmallMatrix const& rhs) + SmallMatrix& + operator += (SmallMatrix const& rhs) { for (int n = 0; n < NRows*NCols; ++n) { m_mat[n] += rhs.m_mat[n]; @@ -274,9 +304,9 @@ namespace amrex { //! Binary operator + returning the result of maxtrix addition, lhs+rhs AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - friend SmallMatrix - operator+ (SmallMatrix lhs, - SmallMatrix const& rhs) + friend SmallMatrix + operator+ (SmallMatrix lhs, + SmallMatrix const& rhs) { lhs += rhs; return lhs; @@ -284,8 +314,8 @@ namespace amrex { //! Operator -= performing matrix subtraction as in (*this) -= rhs AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix& - operator -= (SmallMatrix const& rhs) + SmallMatrix& + operator -= (SmallMatrix const& rhs) { for (int n = 0; n < NRows*NCols; ++n) { m_mat[n] -= rhs.m_mat[n]; @@ -295,9 +325,9 @@ namespace amrex { //! Binary operator - returning the result of maxtrix subtraction, lhs-rhs AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - friend SmallMatrix - operator- (SmallMatrix lhs, - SmallMatrix const& rhs) + friend SmallMatrix + operator- (SmallMatrix lhs, + SmallMatrix const& rhs) { lhs -= rhs; return lhs; @@ -305,7 +335,7 @@ namespace amrex { //! Unary minus operator [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix + SmallMatrix operator- () const { return (*this) * T(-1); @@ -313,7 +343,7 @@ namespace amrex { //! Operator *= that scales this matrix in place by a scalar. AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix& + SmallMatrix& operator *= (T a) { for (auto& x : m_mat) { @@ -324,8 +354,8 @@ namespace amrex { //! Returns the product of a matrix and a scalar AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - friend SmallMatrix - operator* (SmallMatrix m, T a) + friend SmallMatrix + operator* (SmallMatrix m, T a) { m *= a; return m; @@ -333,23 +363,23 @@ namespace amrex { //! Returns the product of a scalar and a matrix AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - friend SmallMatrix - operator* (T a, SmallMatrix m) + friend SmallMatrix + operator* (T a, SmallMatrix m) { m *= a; return m; } //! Returns matrix product of two matrices - template + template AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - friend SmallMatrix - operator* (SmallMatrix const& lhs, - SmallMatrix const& rhs); + friend SmallMatrix + operator* (SmallMatrix const& lhs, + SmallMatrix const& rhs); //! Returns the dot product of two vectors [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - T dot (SmallMatrix const& rhs) const + T dot (SmallMatrix const& rhs) const { T r = 0; for (int n = 0; n < NRows*NCols; ++n) { @@ -362,30 +392,31 @@ namespace amrex { T m_mat[NRows*NCols]; }; - template + template [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - SmallMatrix - operator* (SmallMatrix const& lhs, - SmallMatrix const& rhs) + SmallMatrix + operator* (SmallMatrix const& lhs, + SmallMatrix const& rhs) { - SmallMatrix r; + static_assert(SI == 0 || SI == 1); + SmallMatrix r; if constexpr (Ord == Order::F) { - for (int j = 0; j < N3; ++j) { - constexpr_for<0,N1>([&] (int i) { r(i,j) = U(0); }); - for (int k = 0; k < N2; ++k) { + for (int j = SI; j < N3+SI; ++j) { + constexpr_for([&] (int i) { r(i,j) = U(0); }); + for (int k = SI; k < N2+SI; ++k) { auto b = rhs(k,j); - constexpr_for<0,N1>([&] (int i) + constexpr_for([&] (int i) { r(i,j) += lhs(i,k) * b; }); } } } else { - for (int i = 0; i < N1; ++i) { - constexpr_for<0,N3>([&] (int j) { r(i,j) = U(0); }); - for (int k = 0; k < N2; ++k) { + for (int i = SI; i < N1+SI; ++i) { + constexpr_for([&] (int j) { r(i,j) = U(0); }); + for (int k = SI; k < N2+SI; ++k) { auto a = lhs(i,k); - constexpr_for<0,N3>([&] (int j) + constexpr_for([&] (int j) { r(i,j) += a * rhs(k,j); }); @@ -395,13 +426,13 @@ namespace amrex { return r; } - template + template std::ostream& operator<< (std::ostream& os, - SmallMatrix const& mat) + SmallMatrix const& mat) { - for (int i = 0; i < NRows; ++i) { - os << mat(i,0); - for (int j = 1; j < NCols; ++j) { + for (int i = SI; i < NRows+SI; ++i) { + os << mat(i,SI); + for (int j = 1+SI; j < NCols+SI; ++j) { os << " " << mat(i,j); } os << "\n"; @@ -409,11 +440,11 @@ namespace amrex { return os; } - template - using SmallVector = SmallMatrix; + template + using SmallVector = SmallMatrix; - template - using SmallRowVector = SmallMatrix; + template + using SmallRowVector = SmallMatrix; } #endif diff --git a/Tests/SmallMatrix/main.cpp b/Tests/SmallMatrix/main.cpp index 2dc502b044..de5732101f 100644 --- a/Tests/SmallMatrix/main.cpp +++ b/Tests/SmallMatrix/main.cpp @@ -12,6 +12,7 @@ int main (int argc, char* argv[]) Order::F == Order::ColumnMajor); amrex::Initialize(argc, argv); + // 0-based indexing { SmallMatrix m34{}; for (int j = 0; j < 4; ++j) { @@ -140,5 +141,135 @@ int main (int argc, char* argv[]) b.setVal(-1); AMREX_ALWAYS_ASSERT(a.dot(b) == -30); } + + // 1-based indexing + { + SmallMatrix m34{}; + for (int j = 1; j <= 4; ++j) { + for (int i = 1; i <= 3; ++i) { + AMREX_ALWAYS_ASSERT(m34(i,j) == 0.0_rt); + } + } + } + { + SmallVector cv{}; + SmallRowVector rv{}; + SmallVector cv2{1,2,3}; + SmallRowVector rv2{0,10,20}; + SmallVector cv3{0,1,2}; + for (int j = 0; j < 3; ++j) { + AMREX_ALWAYS_ASSERT(cv(j+1) == 0.0_rt && + rv(j+1) == 0.0_rt && + cv2(j+1) == j+1 && + rv2(j+1) == j*10 && + cv3(j+1) == j); + } + AMREX_ALWAYS_ASSERT(cv3(4) == 0 && cv3(5) == 0); + } + { + SmallMatrix m34{{0,3,6,9}, + {1,4,7,10}, + {2,5,8,11}}; + int v = 0; + for (int j = 1; j <= 4; ++j) { + for (int i = 1; i <= 3; ++i) { + AMREX_ALWAYS_ASSERT(m34(i,j) == v++); + } + } + std::cout << m34; + } + { + SmallMatrix m34{{0,1,2,3}, + {4,5,6,7}, + {8,9,10,11}}; + int v = 0; + for (int i = 1; i <= 3; ++i) { + for (int j = 1; j <= 4; ++j) { + AMREX_ALWAYS_ASSERT(m34(i,j) == v++); + } + } + } + { + auto v3 = SmallVector::Zero(); + v3[1] = 1.; + v3(2) = 2.; + v3[3] = 3.; + auto m33 = SmallMatrix::Identity(); + auto r = m33*v3; + AMREX_ALWAYS_ASSERT(almostEqual(r[1],v3[1]) && + almostEqual(r[2],v3[2]) && + almostEqual(r[3],v3[3])); + } + { + SmallMatrix A{{1, 0, 1}, + {2, 1, 1}, + {0, 1, 1}, + {1, 1, 2}}; + SmallMatrix B{{1, 2, 1}, + {2, 3, 1}, + {4, 2, 2}}; + SmallMatrix C{10, 8, 6}; + auto ABC = A*B*C; + AMREX_ALWAYS_ASSERT(ABC(1,1) == 100 && + ABC(2,1) == 182 && + ABC(3,1) == 118 && + ABC(4,1) == 218); + } + { + SmallMatrix A{{1, 2, 0, 1}, + {0, 1, 1, 1}, + {1, 1, 1, 2}}; + SmallMatrix B{{1, 2, 4}, + {2, 3, 2}, + {1, 1, 2}}; + SmallMatrix C{10, 8, 6}; + auto ABC = A.transpose()*B.transposeInPlace()*C.transpose(); + AMREX_ALWAYS_ASSERT(ABC(1,1) == 100 && + ABC(2,1) == 182 && + ABC(3,1) == 118 && + ABC(4,1) == 218); + } + { + SmallMatrix m; + m.setVal(2); + using M = decltype(m); + AMREX_ALWAYS_ASSERT(m.product() == Math::powi(2)); + AMREX_ALWAYS_ASSERT(m.sum() == 2*m.row_size*m.column_size); + } + { + SmallMatrix m{{1.0, 3.4, 4.5, 5.6, 6.7}, + {1.3, 2.0, 4.5, 5.6, 6.7}, + {1.3, 1.0, 3.0, 5.6, 6.7}, + {1.3, 1.4, 4.5, 4.0, 6.7}, + {1.3, 1.0, 4.5, 5.6, 5.0}}; + AMREX_ALWAYS_ASSERT(m.trace() == double(1+2+3+4+5)); + } + { + SmallMatrix a{{+1, +2, +3}, + {+7, +8, +9}}; + SmallMatrix b{{-1, -2, -3}, + {-7, -8, -9}}; + auto c = a*2 + 2*b; + for (auto const& x : c) { + AMREX_ALWAYS_ASSERT(x == 0); + } + } + { + SmallMatrix a{{+1, +2, +3}, + {+7, +8, +9}}; + SmallMatrix b{{-1, -2, -3}, + {-7, -8, -9}}; + auto c = -a - b; + for (auto const& x : c) { + AMREX_ALWAYS_ASSERT(x == 0); + } + } + { + SmallMatrix a{{+1, +2, +3}, + {+7, +8, +9}}; + SmallMatrix b; + b.setVal(-1); + AMREX_ALWAYS_ASSERT(a.dot(b) == -30); + } amrex::Finalize(); }