Simplifications

FEniCS · Jan 11, 2025 · 5fbe071 · 5fbe071
1 parent fd23f49
commit 5fbe071
Show file tree

Hide file tree

Showing 7 changed files with 27 additions and 12 deletions.
diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
@@ -323,11 +323,14 @@ class MatrixCSR
   /// @note MPI Collective
   double squared_norm() const;
 
-  /// @brief Computes y += Ax for the local CSR matrix and local dense vectors
+  /// @brief Compute the product `y += Ax`.
   ///
-  /// @param[in] x Input Vector
-  /// @param[out] y Output vector
-  void spmv(Vector<value_type>& x, Vector<value_type>& y);
+  /// The vectors `x` and `y` must have parallel layouts that are
+  /// compatible with `A`.
+  ///
+  /// @param[in] x Vector to be apply `A` to.
+  /// @param[in,out] y Vector to accumulate the result into.
+  void mult(Vector<value_type>& x, Vector<value_type>& y);
 
   /// @brief Index maps for the row and column space.
   ///
@@ -765,7 +768,7 @@ double MatrixCSR<U, V, W, X>::squared_norm() const
 //
 /// Computes y += A*x for a parallel CSR matrix A and parallel dense vectors x,y
 template <typename Scalar, typename V, typename W, typename X>
-void MatrixCSR<Scalar, V, W, X>::spmv(la::Vector<Scalar>& x,
+void MatrixCSR<Scalar, V, W, X>::mult(la::Vector<Scalar>& x,
                                       la::Vector<Scalar>& y)
 {
   // start communication (update ghosts)
@@ -787,23 +790,31 @@ void MatrixCSR<Scalar, V, W, X>::spmv(la::Vector<Scalar>& x,
   // First stage:  spmv - diagonal
   // yi[0] += Ai[0] * xi[0]
   if (_bs[1] == 1)
+  {
     impl::spmv<Scalar, 1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
                           _bs[0], 1);
+  }
   else
+  {
     impl::spmv<Scalar, -1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
                            _bs[0], _bs[1]);
+  }
 
   // finalize ghost update
   x.scatter_fwd_end();
 
   // Second stage:  spmv - off-diagonal
   // yi[0] += Ai[1] * xi[1]
   if (_bs[1] == 1)
+  {
     impl::spmv<Scalar, 1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
                           _bs[0], 1);
+  }
   else
+  {
     impl::spmv<Scalar, -1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
                            _bs[0], _bs[1]);
+  }
 }
 
 } // namespace dolfinx::la
diff --git a/cpp/dolfinx/la/matrix_csr_impl.h b/cpp/dolfinx/la/matrix_csr_impl.h
@@ -252,14 +252,18 @@ void spmv(std::span<const T> values, std::span<const std::int64_t> row_begin,
         if constexpr (BS1 == -1)
         {
           for (int k1 = 0; k1 < bs1; ++k1)
+          {
             vi += values[j * bs1 * bs0 + k1 * bs0 + k0]
                   * x[indices[j] * bs1 + k1];
+          }
         }
         else
         {
           for (int k1 = 0; k1 < BS1; ++k1)
+          {
             vi += values[j * BS1 * bs0 + k1 * bs0 + k0]
                   * x[indices[j] * BS1 + k1];
+          }
         }
       }
 

diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
@@ -113,7 +113,7 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
 
   // Matrix A represents the action of the Laplace operator, so when
   // applied to a constant vector the result should be zero
-  A.spmv(x, y);
+  A.mult(x, y);
 
   std::ranges::for_each(y.array(),
                         [](auto a) { REQUIRE(std::abs(a) < 1e-13); });

diff --git a/python/dolfinx/la.py b/python/dolfinx/la.py
@@ -63,14 +63,14 @@ def index_map(self, i: int) -> IndexMap:
         """
         return self._cpp_object.index_map(i)
 
-    def spmv(self, x, y):
-        """Sparse matvec with DOLFINx Vectors
+    def mult(self, x, y):
+        """Compute ``y += Ax``.
 
         Args:
             x: Input Vector
             y: Output Vector
         """
-        self._cpp_object.spmv(x._cpp_object, y._cpp_object)
+        self._cpp_object.mult(x._cpp_object, y._cpp_object)
 
     @property
     def block_size(self):

diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
@@ -132,7 +132,7 @@ void declare_objects(nb::module_& m, const std::string& type)
                &dolfinx::la::MatrixCSR<T>::set),
            nb::arg("x"))
       .def("scatter_reverse", &dolfinx::la::MatrixCSR<T>::scatter_rev)
-      .def("spmv", &dolfinx::la::MatrixCSR<T>::spmv)
+      .def("mult", &dolfinx::la::MatrixCSR<T>::mult)
       .def("to_dense",
            [](const dolfinx::la::MatrixCSR<T>& self)
            {

diff --git a/python/test/unit/fem/test_discrete_operators.py b/python/test/unit/fem/test_discrete_operators.py
@@ -251,7 +251,7 @@ def f(x):
     # Compute global matrix vector product
     w = Function(W)
     w.x.array[:] = 0.0
-    G.spmv(u.x, w.x)
+    G.mult(u.x, w.x)
     w.x.scatter_forward()
 
     atol = 100 * np.finfo(default_real_type).resolution

diff --git a/python/test/unit/la/test_matrix_vector.py b/python/test/unit/la/test_matrix_vector.py
@@ -71,7 +71,7 @@ def test_matvec(dtype):
     b = la.vector(imap, dtype=dtype)
     u = la.vector(imap, dtype=dtype)
     b.array[:] = 1.0
-    A.spmv(b, u)
+    A.mult(b, u)
     u.scatter_forward()
     assert np.allclose(u.array[: imap.size_local], 2.0)