From 1967f5e5d35c8767b1a83687b83a609204474735 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 12:19:52 +0100
Subject: [PATCH 01/21] Use container not allocator as template

---
 cpp/dolfinx/fem/Function.h      | 16 +++++-----
 cpp/dolfinx/la/Vector.h         | 52 +++++++++++++++------------------
 cpp/dolfinx/la/petsc.h          |  4 +--
 python/dolfinx/wrappers/fem.cpp |  2 +-
 python/dolfinx/wrappers/la.cpp  | 31 +++++++++++---------
 5 files changed, 52 insertions(+), 53 deletions(-)

diff --git a/cpp/dolfinx/fem/Function.h b/cpp/dolfinx/fem/Function.h
index 4ca9c1acb36..3d0efee4369 100644
--- a/cpp/dolfinx/fem/Function.h
+++ b/cpp/dolfinx/fem/Function.h
@@ -51,8 +51,8 @@ class Function
   /// @param[in] V The function space
   explicit Function(std::shared_ptr<const FunctionSpace<U>> V)
       : _function_space(V),
-        _x(std::make_shared<la::Vector<T>>(V->dofmap()->index_map,
-                                           V->dofmap()->index_map_bs()))
+        _x(std::make_shared<la::Vector<std::vector<T>>>(
+            V->dofmap()->index_map, V->dofmap()->index_map_bs()))
   {
     if (!V->component().empty())
     {
@@ -70,7 +70,7 @@ class Function
   /// @param[in] V The function space
   /// @param[in] x The vector
   Function(std::shared_ptr<const FunctionSpace<U>> V,
-           std::shared_ptr<la::Vector<T>> x)
+           std::shared_ptr<la::Vector<std::vector<T>>> x)
       : _function_space(V), _x(x)
   {
     // We do not check for a subspace since this constructor is used for
@@ -116,8 +116,8 @@ class Function
     auto [V, map] = _function_space->collapse();
 
     // Create new vector
-    auto x = std::make_shared<la::Vector<T>>(V.dofmap()->index_map,
-                                             V.dofmap()->index_map_bs());
+    auto x = std::make_shared<la::Vector<std::vector<T>>>(
+        V.dofmap()->index_map, V.dofmap()->index_map_bs());
 
     // Copy values into new vector
     std::span<const T> x_old = _x->array();
@@ -140,10 +140,10 @@ class Function
   }
 
   /// @brief Underlying vector
-  std::shared_ptr<const la::Vector<T>> x() const { return _x; }
+  std::shared_ptr<const la::Vector<std::vector<T>>> x() const { return _x; }
 
   /// @brief Underlying vector
-  std::shared_ptr<la::Vector<T>> x() { return _x; }
+  std::shared_ptr<la::Vector<std::vector<T>>> x() { return _x; }
 
   /// @brief Interpolate a provided Function.
   /// @param[in] v The function to be interpolated
@@ -626,7 +626,7 @@ class Function
   std::shared_ptr<const FunctionSpace<U>> _function_space;
 
   // The vector of expansion coefficients (local)
-  std::shared_ptr<la::Vector<T>> _x;
+  std::shared_ptr<la::Vector<std::vector<T>>> _x;
 };
 
 } // namespace dolfinx::fem
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 20df13a2b0a..db8148a4c96 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -21,23 +21,18 @@ namespace dolfinx::la
 
 /// Distributed vector
 
-template <typename T, class Allocator = std::allocator<T>>
+template <class V>
 class Vector
 {
-public:
-  /// The value type
-  using value_type = T;
-
-  /// The allocator type
-  using allocator_type = Allocator;
+  using T = typename V::value_type;
 
+public:
   /// Create a distributed vector
-  Vector(std::shared_ptr<const common::IndexMap> map, int bs,
-         const Allocator& alloc = Allocator())
+  Vector(std::shared_ptr<const common::IndexMap> map, int bs)
       : _map(map), _scatterer(std::make_shared<common::Scatterer<>>(*_map, bs)),
-        _bs(bs), _buffer_local(_scatterer->local_buffer_size(), alloc),
-        _buffer_remote(_scatterer->remote_buffer_size(), alloc),
-        _x(bs * (map->size_local() + map->num_ghosts()), alloc)
+        _bs(bs), _buffer_local(_scatterer->local_buffer_size()),
+        _buffer_remote(_scatterer->remote_buffer_size()),
+        _x(bs * (map->size_local() + map->num_ghosts()))
   {
   }
 
@@ -179,9 +174,6 @@ class Vector
   /// Get local part of the vector
   std::span<T> mutable_array() { return std::span(_x); }
 
-  /// Get the allocator associated with the container
-  constexpr allocator_type allocator() const { return _x.get_allocator(); }
-
 private:
   // Map describing the data layout
   std::shared_ptr<const common::IndexMap> _map;
@@ -196,10 +188,10 @@ class Vector
   std::vector<MPI_Request> _request = {MPI_REQUEST_NULL};
 
   // Buffers for ghost scatters
-  std::vector<T, Allocator> _buffer_local, _buffer_remote;
+  V _buffer_local, _buffer_remote;
 
   // Vector data
-  std::vector<T, Allocator> _x;
+  V _x;
 };
 
 /// Compute the inner product of two vectors. The two vectors must have
@@ -208,9 +200,10 @@ class Vector
 /// @param a A vector
 /// @param b A vector
 /// @return Returns `a^{H} b` (`a^{T} b` if `a` and `b` are real)
-template <typename T, class Allocator>
-T inner_product(const Vector<T, Allocator>& a, const Vector<T, Allocator>& b)
+template <class V>
+auto inner_product(const Vector<V>& a, const Vector<V>& b)
 {
+  using T = typename V::value_type;
   const std::int32_t local_size = a.bs() * a.map()->size_local();
   if (local_size != b.bs() * b.map()->size_local())
     throw std::runtime_error("Incompatible vector sizes");
@@ -238,9 +231,10 @@ T inner_product(const Vector<T, Allocator>& a, const Vector<T, Allocator>& b)
 
 /// Compute the squared L2 norm of vector
 /// @note Collective MPI operation
-template <typename T, class Allocator>
-auto squared_norm(const Vector<T, Allocator>& a)
+template <class V>
+auto squared_norm(const Vector<V>& a)
 {
+  using T = typename V::value_type;
   T result = inner_product(a, a);
   return std::real(result);
 }
@@ -249,9 +243,10 @@ auto squared_norm(const Vector<T, Allocator>& a)
 /// @note Collective MPI operation
 /// @param a A vector
 /// @param type Norm type (supported types are \f$L^2\f$ and \f$L^\infty\f$)
-template <typename T, class Allocator>
-auto norm(const Vector<T, Allocator>& a, Norm type = Norm::l2)
+template <class V>
+auto norm(const Vector<V>& a, Norm type = Norm::l2)
 {
+  using T = typename V::value_type;
   switch (type)
   {
   case Norm::l2:
@@ -279,9 +274,10 @@ auto norm(const Vector<T, Allocator>& a, Norm type = Norm::l2)
 /// vectors must have identical parallel layouts. The vectors are
 /// modified in-place.
 /// @param[in] tol The tolerance used to detect a linear dependency
-template <typename T, class Allocator>
-void orthonormalize(std::span<Vector<T, Allocator>> basis, double tol = 1.0e-10)
+template <class V>
+void orthonormalize(std::span<Vector<V>> basis, double tol = 1.0e-10)
 {
+  using T = typename V::value_type;
   // Loop over each vector in basis
   for (std::size_t i = 0; i < basis.size(); ++i)
   {
@@ -313,10 +309,10 @@ void orthonormalize(std::span<Vector<T, Allocator>> basis, double tol = 1.0e-10)
 /// @param[in] basis The set of vectors to check
 /// @param[in] tol The tolerance used to test for orthonormality
 /// @return True is basis is orthonormal, otherwise false
-template <typename T, class Allocator>
-bool is_orthonormal(std::span<const Vector<T, Allocator>> basis,
-                    double tol = 1.0e-10)
+template <class V>
+bool is_orthonormal(std::span<const Vector<V>> basis, double tol = 1.0e-10)
 {
+  using T = typename V::value_type;
   for (std::size_t i = 0; i < basis.size(); i++)
   {
     for (std::size_t j = i; j < basis.size(); j++)
diff --git a/cpp/dolfinx/la/petsc.h b/cpp/dolfinx/la/petsc.h
index 7d8dd78f389..ab42a02c3a2 100644
--- a/cpp/dolfinx/la/petsc.h
+++ b/cpp/dolfinx/la/petsc.h
@@ -78,8 +78,8 @@ Vec create_vector_wrap(const common::IndexMap& map, int bs,
 /// Create a PETSc Vec that wraps the data in an array
 /// @param[in] x The vector to be wrapped
 /// @return A PETSc Vec object that shares the data in @p x
-template <typename Allocator>
-Vec create_vector_wrap(const la::Vector<PetscScalar, Allocator>& x)
+template <class V>
+Vec create_vector_wrap(const la::Vector<V>& x)
 {
   assert(x.map());
   return create_vector_wrap(*x.map(), x.bs(), x.array());
diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp
index 39162e77f01..7dd3878c6cc 100644
--- a/python/dolfinx/wrappers/fem.cpp
+++ b/python/dolfinx/wrappers/fem.cpp
@@ -142,7 +142,7 @@ void declare_objects(py::module& m, const std::string& type)
                std::shared_ptr<const dolfinx::fem::FunctionSpace<double>>>(),
            "Create a function on the given function space")
       .def(py::init<std::shared_ptr<dolfinx::fem::FunctionSpace<double>>,
-                    std::shared_ptr<dolfinx::la::Vector<T>>>())
+                    std::shared_ptr<dolfinx::la::Vector<std::vector<T>>>>())
       .def_readwrite("name", &dolfinx::fem::Function<T, double>::name)
       .def("sub", &dolfinx::fem::Function<T, double>::sub,
            "Return sub-function (view into parent Function")
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index 408e00aaf71..762bb982994 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -39,35 +39,38 @@ void declare_objects(py::module& m, const std::string& type)
 {
   // dolfinx::la::Vector
   std::string pyclass_vector_name = std::string("Vector_") + type;
-  py::class_<dolfinx::la::Vector<T>, std::shared_ptr<dolfinx::la::Vector<T>>>(
+  py::class_<dolfinx::la::Vector<std::vector<T>>,
+             std::shared_ptr<dolfinx::la::Vector<std::vector<T>>>>(
       m, pyclass_vector_name.c_str())
-      .def(py::init([](std::shared_ptr<const dolfinx::common::IndexMap> map,
-                       int bs) { return dolfinx::la::Vector<T>(map, bs); }),
+      .def(py::init(
+               [](std::shared_ptr<const dolfinx::common::IndexMap> map, int bs)
+               { return dolfinx::la::Vector<std::vector<T>>(map, bs); }),
            py::arg("map"), py::arg("bs"))
-      .def(py::init([](const dolfinx::la::Vector<T>& vec)
-                    { return dolfinx::la::Vector<T>(vec); }),
+      .def(py::init([](const dolfinx::la::Vector<std::vector<T>>& vec)
+                    { return dolfinx::la::Vector<std::vector<T>>(vec); }),
            py::arg("vec"))
-      .def_property_readonly("dtype", [](const dolfinx::la::Vector<T>& self)
+      .def_property_readonly("dtype",
+                             [](const dolfinx::la::Vector<std::vector<T>>& self)
                              { return py::dtype::of<T>(); })
-      .def("set", &dolfinx::la::Vector<T>::set, py::arg("v"))
+      .def("set", &dolfinx::la::Vector<std::vector<T>>::set, py::arg("v"))
       .def(
           "norm",
-          [](dolfinx::la::Vector<T>& self, dolfinx::la::Norm type)
+          [](dolfinx::la::Vector<std::vector<T>>& self, dolfinx::la::Norm type)
           { return dolfinx::la::norm(self, type); },
           py::arg("type") = dolfinx::la::Norm::l2)
-      .def_property_readonly("map", &dolfinx::la::Vector<T>::map)
-      .def_property_readonly("bs", &dolfinx::la::Vector<T>::bs)
+      .def_property_readonly("map", &dolfinx::la::Vector<std::vector<T>>::map)
+      .def_property_readonly("bs", &dolfinx::la::Vector<std::vector<T>>::bs)
       .def_property_readonly("array",
-                             [](dolfinx::la::Vector<T>& self)
+                             [](dolfinx::la::Vector<std::vector<T>>& self)
                              {
                                std::span<T> array = self.mutable_array();
                                return py::array_t<T>(array.size(), array.data(),
                                                      py::cast(self));
                              })
-      .def("scatter_forward", &dolfinx::la::Vector<T>::scatter_fwd)
+      .def("scatter_forward", &dolfinx::la::Vector<std::vector<T>>::scatter_fwd)
       .def(
           "scatter_reverse",
-          [](dolfinx::la::Vector<T>& self, PyScatterMode mode)
+          [](dolfinx::la::Vector<std::vector<T>>& self, PyScatterMode mode)
           {
             switch (mode)
             {
@@ -146,7 +149,7 @@ void petsc_module(py::module& m)
         py::arg("bs"), "Create a ghosted PETSc Vec for index map.");
   m.def(
       "create_vector_wrap",
-      [](dolfinx::la::Vector<PetscScalar, std::allocator<PetscScalar>>& x)
+      [](dolfinx::la::Vector<std::vector<PetscScalar>>& x)
       { return dolfinx::la::petsc::create_vector_wrap(x); },
       py::return_value_policy::take_ownership, py::arg("x"),
       "Create a ghosted PETSc Vec that wraps a DOLFINx Vector");

From 4ccd29a6fa9a0240e43110c12eb7f5dd8efb191b Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 12:28:25 +0100
Subject: [PATCH 02/21] Update test

---
 cpp/test/matrix.cpp | 7 ++++---
 cpp/test/vector.cpp | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
index a234cb0fe5b..6b1f833280f 100644
--- a/cpp/test/matrix.cpp
+++ b/cpp/test/matrix.cpp
@@ -71,7 +71,8 @@ void spmv_impl(std::span<const T> values,
 /// @param[in] x Input vector
 /// @param[in, out] y Output vector
 template <typename T>
-void spmv(la::MatrixCSR<T>& A, la::Vector<T>& x, la::Vector<T>& y)
+void spmv(la::MatrixCSR<T>& A, la::Vector<std::vector<T>>& x,
+          la::Vector<std::vector<T>>& y)
 
 {
   // start communication (update ghosts)
@@ -169,8 +170,8 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
   // Get compatible vectors
   auto maps = A.index_maps();
 
-  la::Vector<double> x(maps[1], 1);
-  la::Vector<double> y(maps[1], 1);
+  la::Vector<std::vector<double>> x(maps[1], 1);
+  la::Vector<std::vector<double>> y(maps[1], 1);
 
   std::size_t col_size = maps[1]->size_local() + maps[1]->num_ghosts();
   CHECK(x.array().size() == col_size);
diff --git a/cpp/test/vector.cpp b/cpp/test/vector.cpp
index 073e9bf585e..6baf7583ce8 100644
--- a/cpp/test/vector.cpp
+++ b/cpp/test/vector.cpp
@@ -36,7 +36,7 @@ void test_vector()
   auto index_map = std::make_shared<common::IndexMap>(
       MPI_COMM_WORLD, size_local, ghosts, global_ghost_owner);
 
-  la::Vector<T> v(index_map, 1);
+  la::Vector<std::vector<T>> v(index_map, 1);
   std::fill(v.mutable_array().begin(), v.mutable_array().end(), 1.0);
 
   double norm2 = la::squared_norm(v);

From 38cd4e67b98bde7d04ef7086571ef1627f66b3b0 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 13:48:01 +0100
Subject: [PATCH 03/21] Updates for demos

---
 cpp/demo/biharmonic/main.cpp          |  5 +++--
 cpp/demo/hyperelasticity/main.cpp     |  2 +-
 cpp/demo/poisson/main.cpp             |  5 +++--
 cpp/demo/poisson_matrix_free/main.cpp | 22 ++++++++++++++--------
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/cpp/demo/biharmonic/main.cpp b/cpp/demo/biharmonic/main.cpp
index 9a50787f403..a08d9d0033a 100644
--- a/cpp/demo/biharmonic/main.cpp
+++ b/cpp/demo/biharmonic/main.cpp
@@ -236,8 +236,9 @@ int main(int argc, char* argv[])
     // Compute solution
     fem::Function<T> u(V);
     auto A = la::petsc::Matrix(fem::petsc::create_matrix(*a), false);
-    la::Vector<T> b(L->function_spaces()[0]->dofmap()->index_map,
-                    L->function_spaces()[0]->dofmap()->index_map_bs());
+    la::Vector<std::vector<T>> b(
+        L->function_spaces()[0]->dofmap()->index_map,
+        L->function_spaces()[0]->dofmap()->index_map_bs());
 
     MatZeroEntries(A.mat());
     fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES),
diff --git a/cpp/demo/hyperelasticity/main.cpp b/cpp/demo/hyperelasticity/main.cpp
index 99f27490066..0f460488b59 100644
--- a/cpp/demo/hyperelasticity/main.cpp
+++ b/cpp/demo/hyperelasticity/main.cpp
@@ -100,7 +100,7 @@ class HyperElasticProblem
 private:
   std::shared_ptr<fem::Form<T>> _l, _j;
   std::vector<std::shared_ptr<const fem::DirichletBC<T>>> _bcs;
-  la::Vector<T> _b;
+  la::Vector<std::vector<T>> _b;
   Vec _b_petsc = nullptr;
   la::petsc::Matrix _matA;
 };
diff --git a/cpp/demo/poisson/main.cpp b/cpp/demo/poisson/main.cpp
index 9e5eb5db6b0..55aa9f0d090 100644
--- a/cpp/demo/poisson/main.cpp
+++ b/cpp/demo/poisson/main.cpp
@@ -216,8 +216,9 @@ int main(int argc, char* argv[])
     // Compute solution
     fem::Function<T> u(V);
     auto A = la::petsc::Matrix(fem::petsc::create_matrix(*a), false);
-    la::Vector<T> b(L->function_spaces()[0]->dofmap()->index_map,
-                    L->function_spaces()[0]->dofmap()->index_map_bs());
+    la::Vector<std::vector<T>> b(
+        L->function_spaces()[0]->dofmap()->index_map,
+        L->function_spaces()[0]->dofmap()->index_map_bs());
 
     MatZeroEntries(A.mat());
     fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES),
diff --git a/cpp/demo/poisson_matrix_free/main.cpp b/cpp/demo/poisson_matrix_free/main.cpp
index 4ecbf9b3122..9637a3554ac 100644
--- a/cpp/demo/poisson_matrix_free/main.cpp
+++ b/cpp/demo/poisson_matrix_free/main.cpp
@@ -42,8 +42,8 @@ namespace linalg
 /// @param[in] x
 /// @param[in] y
 template <typename U>
-void axpy(la::Vector<U>& r, U alpha, const la::Vector<U>& x,
-          const la::Vector<U>& y)
+void axpy(la::Vector<U>& r, typename U::value_type alpha,
+          const la::Vector<U>& x, const la::Vector<U>& y)
 {
   std::transform(x.array().begin(), x.array().end(), y.array().begin(),
                  r.mutable_array().begin(),
@@ -65,12 +65,14 @@ template <typename U, typename ApplyFunction>
 int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
        int kmax = 50, double rtol = 1e-8)
 {
+  using T = typename U::value_type;
+
   // Create working vectors
   la::Vector<U> r(b), y(b);
 
   // Compute initial residual r0 = b - Ax0
   action(x, y);
-  axpy(r, U(-1), y, b);
+  axpy(r, T(-1), y, b);
 
   // Create p work vector
   la::Vector<U> p(r);
@@ -88,7 +90,7 @@ int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
     action(p, y);
 
     // Compute alpha = r.r/p.y
-    const U alpha = rnorm / la::inner_product(p, y);
+    const T alpha = rnorm / la::inner_product(p, y);
 
     // Update x (x <- x + alpha*p)
     axpy(x, alpha, p, x);
@@ -98,7 +100,7 @@ int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
 
     // Update residual norm
     const auto rnorm_new = la::squared_norm(r);
-    const U beta = rnorm_new / rnorm;
+    const T beta = rnorm_new / rnorm;
     rnorm = rnorm_new;
 
     if (rnorm / rnorm0 < rtol2)
@@ -161,7 +163,8 @@ int main(int argc, char* argv[])
     auto bc = std::make_shared<const fem::DirichletBC<T>>(u_D, bdofs);
 
     // Assemble RHS vector
-    la::Vector<T> b(V->dofmap()->index_map, V->dofmap()->index_map_bs());
+    la::Vector<std::vector<T>> b(V->dofmap()->index_map,
+                                 V->dofmap()->index_map_bs());
     fem::assemble_vector(b.mutable_array(), *L);
 
     // Apply lifting to account for Dirichlet boundary condition
@@ -182,8 +185,11 @@ int main(int argc, char* argv[])
     const std::vector<T> constants = fem::pack_constants(*M);
 
     // Create function for computing the action of A on x (y = Ax)
-    std::function<void(la::Vector<T>&, la::Vector<T>&)> action
-        = [&M, &ui, &bc, &coeff, &constants](la::Vector<T>& x, la::Vector<T>& y)
+    std::function<void(la::Vector<std::vector<T>>&,
+                       la::Vector<std::vector<T>>&)>
+        action
+        = [&M, &ui, &bc, &coeff, &constants](la::Vector<std::vector<T>>& x,
+                                             la::Vector<std::vector<T>>& y)
     {
       // Zero y
       y.set(0.0);

From 4d8665c9bd3b5daaa63111def9da26e81be04a5a Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 14:18:40 +0100
Subject: [PATCH 04/21] Also for MatrixCSR

---
 cpp/dolfinx/fem/utils.h              | 13 ++---
 cpp/dolfinx/graph/AdjacencyList.h    |  8 ++--
 cpp/dolfinx/la/MatrixCSR.h           | 45 ++++++++----------
 cpp/test/matrix.cpp                  | 14 +++---
 python/dolfinx/wrappers/assemble.cpp |  6 ++-
 python/dolfinx/wrappers/la.cpp       | 71 ++++++++++++++--------------
 6 files changed, 77 insertions(+), 80 deletions(-)

diff --git a/cpp/dolfinx/fem/utils.h b/cpp/dolfinx/fem/utils.h
index 1739166dc5e..1f4476190b7 100644
--- a/cpp/dolfinx/fem/utils.h
+++ b/cpp/dolfinx/fem/utils.h
@@ -814,12 +814,13 @@ void pack(std::span<T> coeffs, std::int32_t cell, int bs, std::span<const T> v,
 /// @private
 /// @brief  Concepts for function that returns cell index
 template <typename F>
-concept FetchCells = requires(F&& f, std::span<const std::int32_t> v) {
-  requires std::invocable<F, std::span<const std::int32_t>>;
-  {
-    f(v)
-  } -> std::convertible_to<std::int32_t>;
-};
+concept FetchCells
+    = requires(F&& f, std::span<const std::int32_t> v) {
+        requires std::invocable<F, std::span<const std::int32_t>>;
+        {
+          f(v)
+          } -> std::convertible_to<std::int32_t>;
+      };
 
 /// @brief Pack a single coefficient for a set of active entities.
 ///
diff --git a/cpp/dolfinx/graph/AdjacencyList.h b/cpp/dolfinx/graph/AdjacencyList.h
index a79055c5801..30a47d78fdf 100644
--- a/cpp/dolfinx/graph/AdjacencyList.h
+++ b/cpp/dolfinx/graph/AdjacencyList.h
@@ -172,10 +172,10 @@ AdjacencyList(T, U) -> AdjacencyList<typename T::value_type>;
 /// @return An adjacency list
 template <typename U>
   requires requires {
-    typename std::decay_t<U>::value_type;
-    requires std::convertible_to<
-        U, std::vector<typename std::decay_t<U>::value_type>>;
-  }
+             typename std::decay_t<U>::value_type;
+             requires std::convertible_to<
+                 U, std::vector<typename std::decay_t<U>::value_type>>;
+           }
 AdjacencyList<typename std::decay_t<U>::value_type>
 regular_adjacency_list(U&& data, int degree)
 {
diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index a324e02099c..808559bfc73 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -113,8 +113,7 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// The matrix storage format is compressed sparse row. The matrix is
 /// partitioned row-wise across MPI rank.
 ///
-/// @tparam T The data type for the matrix
-/// @tparam Allocator The memory allocator type for the data storage
+/// @tparam V The data container type for the matrix
 ///
 /// @note Highly "experimental" storage of a matrix in CSR format which
 /// can be assembled into using the usual dolfinx assembly routines
@@ -122,15 +121,12 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// code.
 ///
 /// @todo Handle block sizes
-template <typename T, class Allocator = std::allocator<T>>
+template <typename V>
 class MatrixCSR
 {
 public:
   /// The value type
-  using value_type = T;
-
-  /// The allocator type
-  using allocator_type = Allocator;
+  using value_type = typename V::value_type;
 
   /// Insertion functor for setting values in matrix. It is typically
   /// used in finite element assembly functions.
@@ -140,7 +136,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const T>& data) -> int
+               const std::span<const value_type>& data) -> int
     {
       this->set(data, rows, cols);
       return 0;
@@ -154,7 +150,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const T>& data) -> int
+               const std::span<const value_type>& data) -> int
     {
       this->add(data, rows, cols);
       return 0;
@@ -165,11 +161,10 @@ class MatrixCSR
   /// @param[in] p The sparsty pattern the describes the parallel
   /// distribution and the non-zero structure
   /// @param[in] alloc The memory allocator for the data storafe
-  MatrixCSR(const SparsityPattern& p, const Allocator& alloc = Allocator())
+  MatrixCSR(const SparsityPattern& p)
       : _index_maps({p.index_map(0),
                      std::make_shared<common::IndexMap>(p.column_index_map())}),
-        _bs({p.block_size(0), p.block_size(1)}),
-        _data(p.num_nonzeros(), 0, alloc),
+        _bs({p.block_size(0), p.block_size(1)}), _data(p.num_nonzeros(), 0),
         _cols(p.graph().array().begin(), p.graph().array().end()),
         _row_ptr(p.graph().offsets().begin(), p.graph().offsets().end()),
         _comm(MPI_COMM_NULL)
@@ -336,7 +331,7 @@ class MatrixCSR
   /// Set all non-zero local entries to a value
   /// including entries in ghost rows
   /// @param[in] x The value to set non-zero matrix entries to
-  void set(T x) { std::fill(_data.begin(), _data.end(), x); }
+  void set(value_type x) { std::fill(_data.begin(), _data.end(), x); }
 
   /// Set values in the matrix
   /// @note Only entries included in the sparsity pattern used to
@@ -350,7 +345,7 @@ class MatrixCSR
   /// set in the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void set(const std::span<const T>& x,
+  void set(const std::span<const value_type>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -370,7 +365,7 @@ class MatrixCSR
   /// add to the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void add(const std::span<const T>& x,
+  void add(const std::span<const value_type>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -390,12 +385,12 @@ class MatrixCSR
   /// manually by using num_owned_rows() if required.
   /// @return Dense copy of the part of the matrix on the calling rank.
   /// Storage is row-major.
-  std::vector<T> to_dense() const
+  std::vector<value_type> to_dense() const
   {
     const std::size_t nrows = num_all_rows();
     const std::size_t ncols
         = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
-    std::vector<T> A(nrows * ncols);
+    std::vector<value_type> A(nrows * ncols);
     for (std::size_t r = 0; r < nrows; ++r)
       for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
         A[r * ncols + _cols[j]] = _data[j];
@@ -426,7 +421,7 @@ class MatrixCSR
 
     // For each ghost row, pack and send values to send to neighborhood
     std::vector<int> insert_pos = _val_send_disp;
-    std::vector<T> ghost_value_data(_val_send_disp.back());
+    std::vector<value_type> ghost_value_data(_val_send_disp.back());
     for (int i = 0; i < num_ghosts0; ++i)
     {
       const int rank = _ghost_row_to_rank[i];
@@ -454,9 +449,9 @@ class MatrixCSR
 
     int status = MPI_Ineighbor_alltoallv(
         ghost_value_data.data(), val_send_count.data(), _val_send_disp.data(),
-        dolfinx::MPI::mpi_type<T>(), _ghost_value_data_in.data(),
+        dolfinx::MPI::mpi_type<value_type>(), _ghost_value_data_in.data(),
         val_recv_count.data(), _val_recv_disp.data(),
-        dolfinx::MPI::mpi_type<T>(), _comm.comm(), &_request);
+        dolfinx::MPI::mpi_type<value_type>(), _comm.comm(), &_request);
     assert(status == MPI_SUCCESS);
   }
 
@@ -488,7 +483,7 @@ class MatrixCSR
 
     const double norm_sq_local = std::accumulate(
         _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
-        double(0), [](auto norm, T y) { return norm + std::norm(y); });
+        double(0), [](auto norm, value_type y) { return norm + std::norm(y); });
     double norm_sq;
     MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM,
                   _comm.comm());
@@ -509,11 +504,11 @@ class MatrixCSR
 
   /// Get local data values
   /// @note Includes ghost values
-  std::vector<T>& values() { return _data; }
+  V& values() { return _data; }
 
   /// Get local values (const version)
   /// @note Includes ghost values
-  const std::vector<T>& values() const { return _data; }
+  const V& values() const { return _data; }
 
   /// Get local row pointers
   /// @note Includes pointers to ghost rows
@@ -543,7 +538,7 @@ class MatrixCSR
   std::array<int, 2> _bs;
 
   // Matrix data
-  std::vector<T, Allocator> _data;
+  V _data;
   std::vector<std::int32_t> _cols, _row_ptr;
 
   // Start of off-diagonal (unowned columns) on each row
@@ -569,7 +564,7 @@ class MatrixCSR
   std::vector<int> _ghost_row_to_rank;
 
   // Temporary store for finalize data during non-blocking communication
-  std::vector<T> _ghost_value_data_in;
+  V _ghost_value_data_in;
 };
 
 } // namespace dolfinx::la
diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
index 6b1f833280f..872e7b96ffc 100644
--- a/cpp/test/matrix.cpp
+++ b/cpp/test/matrix.cpp
@@ -71,7 +71,7 @@ void spmv_impl(std::span<const T> values,
 /// @param[in] x Input vector
 /// @param[in, out] y Output vector
 template <typename T>
-void spmv(la::MatrixCSR<T>& A, la::Vector<std::vector<T>>& x,
+void spmv(la::MatrixCSR<std::vector<T>>& A, la::Vector<std::vector<T>>& x,
           la::Vector<std::vector<T>>& y)
 
 {
@@ -106,7 +106,7 @@ void spmv(la::MatrixCSR<T>& A, la::Vector<std::vector<T>>& x,
 /// @brief Create a matrix operator
 /// @param comm The communicator to builf the matrix on
 /// @return The assembled matrix
-la::MatrixCSR<double> create_operator(MPI_Comm comm)
+la::MatrixCSR<std::vector<double>> create_operator(MPI_Comm comm)
 {
   auto part = mesh::create_cell_partitioner(mesh::GhostMode::none);
   auto mesh = std::make_shared<mesh::Mesh<double>>(
@@ -123,7 +123,7 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
 
   la::SparsityPattern sp = fem::create_sparsity_pattern(*a);
   sp.assemble();
-  la::MatrixCSR<double> A(sp);
+  la::MatrixCSR<std::vector<double>> A(sp);
   fem::assemble_matrix(A.mat_add_values(), *a, {});
   A.finalize();
 
@@ -132,8 +132,8 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
 
 [[maybe_unused]] void test_matrix_norm()
 {
-  la::MatrixCSR<double> A0 = create_operator(MPI_COMM_SELF);
-  la::MatrixCSR<double> A1 = create_operator(MPI_COMM_WORLD);
+  la::MatrixCSR A0 = create_operator(MPI_COMM_SELF);
+  la::MatrixCSR A1 = create_operator(MPI_COMM_WORLD);
   CHECK(A1.norm_squared() == Approx(A0.norm_squared()).epsilon(1e-8));
 }
 
@@ -162,7 +162,7 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
   sp.assemble();
 
   // Assemble matrix
-  la::MatrixCSR<double> A(sp);
+  la::MatrixCSR<std::vector<double>> A(sp);
   fem::assemble_matrix(A.mat_add_values(), *a, {});
   A.finalize();
   CHECK((V->dofmap()->index_map->size_local() == A.num_owned_rows()));
@@ -197,7 +197,7 @@ void test_matrix()
   p.assemble();
 
   using T = float;
-  la::MatrixCSR<T> A(p);
+  la::MatrixCSR<std::vector<T>> A(p);
   A.add(std::vector<decltype(A)::value_type>{1}, std::vector{0},
         std::vector{0});
   A.add(std::vector<decltype(A)::value_type>{2.3}, std::vector{4},
diff --git a/python/dolfinx/wrappers/assemble.cpp b/python/dolfinx/wrappers/assemble.cpp
index 967fc7e287d..9bc3e49dca4 100644
--- a/python/dolfinx/wrappers/assemble.cpp
+++ b/python/dolfinx/wrappers/assemble.cpp
@@ -139,7 +139,8 @@ void declare_assembly_functions(py::module& m)
   // MatrixCSR
   m.def(
       "assemble_matrix",
-      [](dolfinx::la::MatrixCSR<T>& A, const dolfinx::fem::Form<T, U>& a,
+      [](dolfinx::la::MatrixCSR<std::vector<T>>& A,
+         const dolfinx::fem::Form<T, U>& a,
          const py::array_t<T, py::array::c_style>& constants,
          const std::map<std::pair<dolfinx::fem::IntegralType, int>,
                         py::array_t<T, py::array::c_style>>& coefficients,
@@ -161,7 +162,8 @@ void declare_assembly_functions(py::module& m)
       py::arg("bcs"), "Experimental.");
   m.def(
       "insert_diagonal",
-      [](dolfinx::la::MatrixCSR<T>& A, const dolfinx::fem::FunctionSpace<U>& V,
+      [](dolfinx::la::MatrixCSR<std::vector<T>>& A,
+         const dolfinx::fem::FunctionSpace<U>& V,
          const std::vector<
              std::shared_ptr<const dolfinx::fem::DirichletBC<T, U>>>& bcs,
          T diagonal)
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index 762bb982994..64ae6a1240b 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -34,43 +34,42 @@ enum class PyScatterMode
 };
 
 // Declare objects that have multiple scalar types
-template <typename T>
+template <typename V>
 void declare_objects(py::module& m, const std::string& type)
 {
+  using T = typename V::value_type;
+
   // dolfinx::la::Vector
   std::string pyclass_vector_name = std::string("Vector_") + type;
-  py::class_<dolfinx::la::Vector<std::vector<T>>,
-             std::shared_ptr<dolfinx::la::Vector<std::vector<T>>>>(
+  py::class_<dolfinx::la::Vector<V>, std::shared_ptr<dolfinx::la::Vector<V>>>(
       m, pyclass_vector_name.c_str())
-      .def(py::init(
-               [](std::shared_ptr<const dolfinx::common::IndexMap> map, int bs)
-               { return dolfinx::la::Vector<std::vector<T>>(map, bs); }),
+      .def(py::init([](std::shared_ptr<const dolfinx::common::IndexMap> map,
+                       int bs) { return dolfinx::la::Vector<V>(map, bs); }),
            py::arg("map"), py::arg("bs"))
-      .def(py::init([](const dolfinx::la::Vector<std::vector<T>>& vec)
-                    { return dolfinx::la::Vector<std::vector<T>>(vec); }),
+      .def(py::init([](const dolfinx::la::Vector<V>& vec)
+                    { return dolfinx::la::Vector<V>(vec); }),
            py::arg("vec"))
-      .def_property_readonly("dtype",
-                             [](const dolfinx::la::Vector<std::vector<T>>& self)
+      .def_property_readonly("dtype", [](const dolfinx::la::Vector<V>& self)
                              { return py::dtype::of<T>(); })
-      .def("set", &dolfinx::la::Vector<std::vector<T>>::set, py::arg("v"))
+      .def("set", &dolfinx::la::Vector<V>::set, py::arg("v"))
       .def(
           "norm",
-          [](dolfinx::la::Vector<std::vector<T>>& self, dolfinx::la::Norm type)
+          [](dolfinx::la::Vector<V>& self, dolfinx::la::Norm type)
           { return dolfinx::la::norm(self, type); },
           py::arg("type") = dolfinx::la::Norm::l2)
-      .def_property_readonly("map", &dolfinx::la::Vector<std::vector<T>>::map)
-      .def_property_readonly("bs", &dolfinx::la::Vector<std::vector<T>>::bs)
+      .def_property_readonly("map", &dolfinx::la::Vector<V>::map)
+      .def_property_readonly("bs", &dolfinx::la::Vector<V>::bs)
       .def_property_readonly("array",
-                             [](dolfinx::la::Vector<std::vector<T>>& self)
+                             [](dolfinx::la::Vector<V>& self)
                              {
                                std::span<T> array = self.mutable_array();
                                return py::array_t<T>(array.size(), array.data(),
                                                      py::cast(self));
                              })
-      .def("scatter_forward", &dolfinx::la::Vector<std::vector<T>>::scatter_fwd)
+      .def("scatter_forward", &dolfinx::la::Vector<V>::scatter_fwd)
       .def(
           "scatter_reverse",
-          [](dolfinx::la::Vector<std::vector<T>>& self, PyScatterMode mode)
+          [](dolfinx::la::Vector<V>& self, PyScatterMode mode)
           {
             switch (mode)
             {
@@ -89,23 +88,23 @@ void declare_objects(py::module& m, const std::string& type)
 
   // dolfinx::la::MatrixCSR
   std::string pyclass_matrix_name = std::string("MatrixCSR_") + type;
-  py::class_<dolfinx::la::MatrixCSR<T>,
-             std::shared_ptr<dolfinx::la::MatrixCSR<T>>>(
+  py::class_<dolfinx::la::MatrixCSR<V>,
+             std::shared_ptr<dolfinx::la::MatrixCSR<V>>>(
       m, pyclass_matrix_name.c_str())
       .def(py::init([](const dolfinx::la::SparsityPattern& p)
-                    { return dolfinx::la::MatrixCSR<T>(p); }),
+                    { return dolfinx::la::MatrixCSR<V>(p); }),
            py::arg("p"))
-      .def_property_readonly("dtype", [](const dolfinx::la::MatrixCSR<T>& self)
+      .def_property_readonly("dtype", [](const dolfinx::la::MatrixCSR<V>& self)
                              { return py::dtype::of<T>(); })
-      .def("norm_squared", &dolfinx::la::MatrixCSR<T>::norm_squared)
-      .def("mat_add_values", &dolfinx::la::MatrixCSR<T>::mat_add_values)
+      .def("norm_squared", &dolfinx::la::MatrixCSR<V>::norm_squared)
+      .def("mat_add_values", &dolfinx::la::MatrixCSR<V>::mat_add_values)
       .def("set",
-           static_cast<void (dolfinx::la::MatrixCSR<T>::*)(T)>(
-               &dolfinx::la::MatrixCSR<T>::set),
+           static_cast<void (dolfinx::la::MatrixCSR<V>::*)(T)>(
+               &dolfinx::la::MatrixCSR<V>::set),
            py::arg("x"))
-      .def("finalize", &dolfinx::la::MatrixCSR<T>::finalize)
+      .def("finalize", &dolfinx::la::MatrixCSR<V>::finalize)
       .def("to_dense",
-           [](const dolfinx::la::MatrixCSR<T>& self)
+           [](const dolfinx::la::MatrixCSR<V>& self)
            {
              std::size_t nrows = self.num_all_rows();
              auto map_col = self.index_maps()[1];
@@ -114,14 +113,14 @@ void declare_objects(py::module& m, const std::string& type)
                                                  std::array{nrows, ncols});
            })
       .def_property_readonly("data",
-                             [](dolfinx::la::MatrixCSR<T>& self)
+                             [](dolfinx::la::MatrixCSR<V>& self)
                              {
                                std::span<T> array = self.values();
                                return py::array_t<T>(array.size(), array.data(),
                                                      py::cast(self));
                              })
       .def_property_readonly("indices",
-                             [](dolfinx::la::MatrixCSR<T>& self)
+                             [](dolfinx::la::MatrixCSR<V>& self)
                              {
                                std::span<const std::int32_t> array
                                    = self.cols();
@@ -129,15 +128,15 @@ void declare_objects(py::module& m, const std::string& type)
                                    array.size(), array.data(), py::cast(self));
                              })
       .def_property_readonly("indptr",
-                             [](dolfinx::la::MatrixCSR<T>& self)
+                             [](dolfinx::la::MatrixCSR<V>& self)
                              {
                                std::span<const std::int32_t> array
                                    = self.row_ptr();
                                return py::array_t<const std::int32_t>(
                                    array.size(), array.data(), py::cast(self));
                              })
-      .def("finalize_begin", &dolfinx::la::MatrixCSR<T>::finalize_begin)
-      .def("finalize_end", &dolfinx::la::MatrixCSR<T>::finalize_end);
+      .def("finalize_begin", &dolfinx::la::MatrixCSR<V>::finalize_begin)
+      .def("finalize_end", &dolfinx::la::MatrixCSR<V>::finalize_end);
 }
 
 void petsc_module(py::module& m)
@@ -274,9 +273,9 @@ void la(py::module& m)
                              py::return_value_policy::reference_internal);
 
   // Declare objects that are templated over type
-  declare_objects<float>(m, "float32");
-  declare_objects<double>(m, "float64");
-  declare_objects<std::complex<float>>(m, "complex64");
-  declare_objects<std::complex<double>>(m, "complex128");
+  declare_objects<std::vector<float>>(m, "float32");
+  declare_objects<std::vector<double>>(m, "float64");
+  declare_objects<std::vector<std::complex<float>>>(m, "complex64");
+  declare_objects<std::vector<std::complex<double>>>(m, "complex128");
 }
 } // namespace dolfinx_wrappers

From 1cde2835e2ce9d456a9d42410adf063228ddaea0 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 14:41:37 +0100
Subject: [PATCH 05/21] Fix documentation

---
 cpp/dolfinx/la/MatrixCSR.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 808559bfc73..c27c3fcdb49 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -160,7 +160,6 @@ class MatrixCSR
   /// Create a distributed matrix
   /// @param[in] p The sparsty pattern the describes the parallel
   /// distribution and the non-zero structure
-  /// @param[in] alloc The memory allocator for the data storafe
   MatrixCSR(const SparsityPattern& p)
       : _index_maps({p.index_map(0),
                      std::make_shared<common::IndexMap>(p.column_index_map())}),

From 5ac089fb3c1e3153412cb5455fd271547960f16f Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 15:07:19 +0100
Subject: [PATCH 06/21] Add index container

---
 cpp/dolfinx/la/MatrixCSR.h | 7 ++++---
 cpp/dolfinx/la/Vector.h    | 6 +++++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index c27c3fcdb49..4fe931eb3bd 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -114,6 +114,7 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// partitioned row-wise across MPI rank.
 ///
 /// @tparam V The data container type for the matrix
+/// @tparam Idx The index container type
 ///
 /// @note Highly "experimental" storage of a matrix in CSR format which
 /// can be assembled into using the usual dolfinx assembly routines
@@ -121,7 +122,7 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// code.
 ///
 /// @todo Handle block sizes
-template <typename V>
+template <class V, class Idx = std::vector<std::int32_t>>
 class MatrixCSR
 {
 public:
@@ -538,10 +539,10 @@ class MatrixCSR
 
   // Matrix data
   V _data;
-  std::vector<std::int32_t> _cols, _row_ptr;
+  Idx _cols, _row_ptr;
 
   // Start of off-diagonal (unowned columns) on each row
-  std::vector<std::int32_t> _off_diagonal_offset;
+  Idx _off_diagonal_offset;
 
   // Neighborhood communicator (ghost->owner communicator for rows)
   dolfinx::MPI::Comm _comm;
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index db8148a4c96..55a228ff44b 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -20,7 +20,9 @@ namespace dolfinx::la
 {
 
 /// Distributed vector
-
+///
+/// @tparam V data container type
+///
 template <class V>
 class Vector
 {
@@ -28,6 +30,8 @@ class Vector
 
 public:
   /// Create a distributed vector
+  /// @param map IndexMap for parallel distribution of the data
+  /// @param bs Block size
   Vector(std::shared_ptr<const common::IndexMap> map, int bs)
       : _map(map), _scatterer(std::make_shared<common::Scatterer<>>(*_map, bs)),
         _bs(bs), _buffer_local(_scatterer->local_buffer_size()),

From d276f0a0d40611b7f1eaaf6bbccc02c076ed1471 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Tue, 25 Apr 2023 17:53:39 +0100
Subject: [PATCH 07/21] Support scalar as template arg to la::Vector

---
 cpp/demo/poisson/main.cpp | 16 +++++++++++++---
 cpp/dolfinx/la/Vector.h   | 22 +++++++++++++++++++---
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/cpp/demo/poisson/main.cpp b/cpp/demo/poisson/main.cpp
index 55aa9f0d090..faf2d177bdd 100644
--- a/cpp/demo/poisson/main.cpp
+++ b/cpp/demo/poisson/main.cpp
@@ -113,11 +113,22 @@ using T = PetscScalar;
 //
 // .. code-block:: cpp
 
+// template <typename T>
+// struct is_complex_t : public std::false_type
+// {
+// };
+
+// template <typename T>
+// struct is_complex_t<std::complex<T>> : public std::true_type
+// {
+// };
+
 int main(int argc, char* argv[])
 {
   dolfinx::init_logging(argc, argv);
   PetscInitialize(&argc, &argv, nullptr, nullptr);
 
+  // std::cout << "test 0: " << is_complex_t<double> << std::endl;
   {
     // Create mesh and function space
     auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet);
@@ -216,9 +227,8 @@ int main(int argc, char* argv[])
     // Compute solution
     fem::Function<T> u(V);
     auto A = la::petsc::Matrix(fem::petsc::create_matrix(*a), false);
-    la::Vector<std::vector<T>> b(
-        L->function_spaces()[0]->dofmap()->index_map,
-        L->function_spaces()[0]->dofmap()->index_map_bs());
+    la::Vector<T> b(L->function_spaces()[0]->dofmap()->index_map,
+                    L->function_spaces()[0]->dofmap()->index_map_bs());
 
     MatZeroEntries(A.mat());
     fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES),
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 55a228ff44b..18b8b71a816 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -14,6 +14,7 @@
 #include <memory>
 #include <numeric>
 #include <span>
+#include <type_traits>
 #include <vector>
 
 namespace dolfinx::la
@@ -26,7 +27,22 @@ namespace dolfinx::la
 template <class V>
 class Vector
 {
-  using T = typename V::value_type;
+
+  template <typename T>
+  struct is_complex_t : public std::false_type
+  {
+  };
+  template <typename T>
+  struct is_complex_t<std::complex<T>> : public std::true_type
+  {
+  };
+
+  using V_t = typename std::conditional<std::is_arithmetic_v<V>
+                                            || is_complex_t<V>::value,
+                                        std::vector<V>, V>::type;
+  using T = typename std::conditional<std::is_arithmetic_v<V>
+                                          || is_complex_t<V>::value,
+                                      V, typename V_t::value_type>::type;
 
 public:
   /// Create a distributed vector
@@ -192,10 +208,10 @@ class Vector
   std::vector<MPI_Request> _request = {MPI_REQUEST_NULL};
 
   // Buffers for ghost scatters
-  V _buffer_local, _buffer_remote;
+  V_t _buffer_local, _buffer_remote;
 
   // Vector data
-  V _x;
+  V_t _x;
 };
 
 /// Compute the inner product of two vectors. The two vectors must have

From 22abc6a6fb2357978cc13313e73366ba5a6a41f1 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Tue, 25 Apr 2023 17:54:35 +0100
Subject: [PATCH 08/21] Remove commented code

---
 cpp/demo/poisson/main.cpp | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/cpp/demo/poisson/main.cpp b/cpp/demo/poisson/main.cpp
index faf2d177bdd..e9a04476861 100644
--- a/cpp/demo/poisson/main.cpp
+++ b/cpp/demo/poisson/main.cpp
@@ -113,16 +113,6 @@ using T = PetscScalar;
 //
 // .. code-block:: cpp
 
-// template <typename T>
-// struct is_complex_t : public std::false_type
-// {
-// };
-
-// template <typename T>
-// struct is_complex_t<std::complex<T>> : public std::true_type
-// {
-// };
-
 int main(int argc, char* argv[])
 {
   dolfinx::init_logging(argc, argv);

From 0f04e489a35948371cbdfdad5357088e2af41694 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Tue, 25 Apr 2023 18:07:48 +0100
Subject: [PATCH 09/21] Fix bindings

---
 cpp/demo/poisson/main.cpp       |  1 -
 cpp/dolfinx/fem/Function.h      | 16 ++++++++--------
 python/dolfinx/wrappers/fem.cpp |  2 +-
 python/dolfinx/wrappers/la.cpp  |  2 +-
 4 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/cpp/demo/poisson/main.cpp b/cpp/demo/poisson/main.cpp
index e9a04476861..9e5eb5db6b0 100644
--- a/cpp/demo/poisson/main.cpp
+++ b/cpp/demo/poisson/main.cpp
@@ -118,7 +118,6 @@ int main(int argc, char* argv[])
   dolfinx::init_logging(argc, argv);
   PetscInitialize(&argc, &argv, nullptr, nullptr);
 
-  // std::cout << "test 0: " << is_complex_t<double> << std::endl;
   {
     // Create mesh and function space
     auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet);
diff --git a/cpp/dolfinx/fem/Function.h b/cpp/dolfinx/fem/Function.h
index 3d0efee4369..4ca9c1acb36 100644
--- a/cpp/dolfinx/fem/Function.h
+++ b/cpp/dolfinx/fem/Function.h
@@ -51,8 +51,8 @@ class Function
   /// @param[in] V The function space
   explicit Function(std::shared_ptr<const FunctionSpace<U>> V)
       : _function_space(V),
-        _x(std::make_shared<la::Vector<std::vector<T>>>(
-            V->dofmap()->index_map, V->dofmap()->index_map_bs()))
+        _x(std::make_shared<la::Vector<T>>(V->dofmap()->index_map,
+                                           V->dofmap()->index_map_bs()))
   {
     if (!V->component().empty())
     {
@@ -70,7 +70,7 @@ class Function
   /// @param[in] V The function space
   /// @param[in] x The vector
   Function(std::shared_ptr<const FunctionSpace<U>> V,
-           std::shared_ptr<la::Vector<std::vector<T>>> x)
+           std::shared_ptr<la::Vector<T>> x)
       : _function_space(V), _x(x)
   {
     // We do not check for a subspace since this constructor is used for
@@ -116,8 +116,8 @@ class Function
     auto [V, map] = _function_space->collapse();
 
     // Create new vector
-    auto x = std::make_shared<la::Vector<std::vector<T>>>(
-        V.dofmap()->index_map, V.dofmap()->index_map_bs());
+    auto x = std::make_shared<la::Vector<T>>(V.dofmap()->index_map,
+                                             V.dofmap()->index_map_bs());
 
     // Copy values into new vector
     std::span<const T> x_old = _x->array();
@@ -140,10 +140,10 @@ class Function
   }
 
   /// @brief Underlying vector
-  std::shared_ptr<const la::Vector<std::vector<T>>> x() const { return _x; }
+  std::shared_ptr<const la::Vector<T>> x() const { return _x; }
 
   /// @brief Underlying vector
-  std::shared_ptr<la::Vector<std::vector<T>>> x() { return _x; }
+  std::shared_ptr<la::Vector<T>> x() { return _x; }
 
   /// @brief Interpolate a provided Function.
   /// @param[in] v The function to be interpolated
@@ -626,7 +626,7 @@ class Function
   std::shared_ptr<const FunctionSpace<U>> _function_space;
 
   // The vector of expansion coefficients (local)
-  std::shared_ptr<la::Vector<std::vector<T>>> _x;
+  std::shared_ptr<la::Vector<T>> _x;
 };
 
 } // namespace dolfinx::fem
diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp
index 7dd3878c6cc..39162e77f01 100644
--- a/python/dolfinx/wrappers/fem.cpp
+++ b/python/dolfinx/wrappers/fem.cpp
@@ -142,7 +142,7 @@ void declare_objects(py::module& m, const std::string& type)
                std::shared_ptr<const dolfinx::fem::FunctionSpace<double>>>(),
            "Create a function on the given function space")
       .def(py::init<std::shared_ptr<dolfinx::fem::FunctionSpace<double>>,
-                    std::shared_ptr<dolfinx::la::Vector<std::vector<T>>>>())
+                    std::shared_ptr<dolfinx::la::Vector<T>>>())
       .def_readwrite("name", &dolfinx::fem::Function<T, double>::name)
       .def("sub", &dolfinx::fem::Function<T, double>::sub,
            "Return sub-function (view into parent Function")
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index 64ae6a1240b..aa5ffcc7e8d 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -148,7 +148,7 @@ void petsc_module(py::module& m)
         py::arg("bs"), "Create a ghosted PETSc Vec for index map.");
   m.def(
       "create_vector_wrap",
-      [](dolfinx::la::Vector<std::vector<PetscScalar>>& x)
+      [](dolfinx::la::Vector<PetscScalar>& x)
       { return dolfinx::la::petsc::create_vector_wrap(x); },
       py::return_value_policy::take_ownership, py::arg("x"),
       "Create a ghosted PETSc Vec that wraps a DOLFINx Vector");

From 3e64ffdd21bf0c3c9da10296bbee714b2187fdc8 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Tue, 25 Apr 2023 18:27:52 +0100
Subject: [PATCH 10/21] Demo update

---
 cpp/demo/poisson_matrix_free/main.cpp | 25 +++++++++++--------------
 cpp/dolfinx/la/Vector.h               |  2 ++
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/cpp/demo/poisson_matrix_free/main.cpp b/cpp/demo/poisson_matrix_free/main.cpp
index 9637a3554ac..837635588b6 100644
--- a/cpp/demo/poisson_matrix_free/main.cpp
+++ b/cpp/demo/poisson_matrix_free/main.cpp
@@ -41,9 +41,7 @@ namespace linalg
 /// @param[in] alpha
 /// @param[in] x
 /// @param[in] y
-template <typename U>
-void axpy(la::Vector<U>& r, typename U::value_type alpha,
-          const la::Vector<U>& x, const la::Vector<U>& y)
+void axpy(auto& r, auto alpha, const auto& x, const auto& y)
 {
   std::transform(x.array().begin(), x.array().end(), y.array().begin(),
                  r.mutable_array().begin(),
@@ -61,21 +59,24 @@ void axpy(la::Vector<U>& r, typename U::value_type alpha,
 /// @return The number if iterations
 /// @pre It is required that the ghost values of `x` and `b` have been
 /// updated before this function is called
-template <typename U, typename ApplyFunction>
-int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
-       int kmax = 50, double rtol = 1e-8)
+// template <typename U, typename ApplyFunction>
+// int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
+template <typename ApplyFunction>
+int cg(auto& x, auto& b, ApplyFunction&& action, int kmax = 50,
+       double rtol = 1e-8)
 {
-  using T = typename U::value_type;
+  // using T = typename U::value_type;
+  using T = typename std::decay_t<decltype(x)>::value_type;
 
   // Create working vectors
-  la::Vector<U> r(b), y(b);
+  la::Vector r(b), y(b);
 
   // Compute initial residual r0 = b - Ax0
   action(x, y);
   axpy(r, T(-1), y, b);
 
   // Create p work vector
-  la::Vector<U> p(r);
+  la::Vector p(r);
 
   // Iterations of CG
   auto rnorm0 = la::squared_norm(r);
@@ -185,11 +186,7 @@ int main(int argc, char* argv[])
     const std::vector<T> constants = fem::pack_constants(*M);
 
     // Create function for computing the action of A on x (y = Ax)
-    std::function<void(la::Vector<std::vector<T>>&,
-                       la::Vector<std::vector<T>>&)>
-        action
-        = [&M, &ui, &bc, &coeff, &constants](la::Vector<std::vector<T>>& x,
-                                             la::Vector<std::vector<T>>& y)
+    auto action = [&M, &ui, &bc, &coeff, &constants](auto& x, auto& y)
     {
       // Zero y
       y.set(0.0);
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 18b8b71a816..e656437a07c 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -45,6 +45,8 @@ class Vector
                                       V, typename V_t::value_type>::type;
 
 public:
+  using value_type = T;
+
   /// Create a distributed vector
   /// @param map IndexMap for parallel distribution of the data
   /// @param bs Block size

From d60a5b2f189b6a3377b997f65bb0e9a297dcafc9 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Tue, 25 Apr 2023 18:31:18 +0100
Subject: [PATCH 11/21] Simplify traits

---
 cpp/dolfinx/la/Vector.h | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index e656437a07c..6b3088272cf 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -29,20 +29,18 @@ class Vector
 {
 
   template <typename T>
-  struct is_complex_t : public std::false_type
+  struct is_scalar_type_t : public std::is_arithmetic<T>
   {
   };
   template <typename T>
-  struct is_complex_t<std::complex<T>> : public std::true_type
+  struct is_scalar_type_t<std::complex<T>> : public std::true_type
   {
   };
 
-  using V_t = typename std::conditional<std::is_arithmetic_v<V>
-                                            || is_complex_t<V>::value,
-                                        std::vector<V>, V>::type;
-  using T = typename std::conditional<std::is_arithmetic_v<V>
-                                          || is_complex_t<V>::value,
-                                      V, typename V_t::value_type>::type;
+  using V_t = typename std::conditional_t<is_scalar_type_t<V>::value,
+                                          std::vector<V>, V>;
+  using T = typename std::conditional_t<is_scalar_type_t<V>::value, V,
+                                        typename V_t::value_type>;
 
 public:
   using value_type = T;

From d9630a9edc98217457b4d1725f192c49e174ee33 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Tue, 25 Apr 2023 22:02:35 +0100
Subject: [PATCH 12/21] Make doxygen happy

---
 cpp/dolfinx/la/Vector.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 6b3088272cf..bb66711878d 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -43,6 +43,7 @@ class Vector
                                         typename V_t::value_type>;
 
 public:
+  /// value_type Scalar type of Vector
   using value_type = T;
 
   /// Create a distributed vector

From 3646e8e5d022d959f67a8c68a6c3afe62d97bb60 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 08:25:56 +0100
Subject: [PATCH 13/21] Interface change

---
 cpp/demo/biharmonic/main.cpp          |  5 +-
 cpp/demo/hyperelasticity/main.cpp     |  2 +-
 cpp/demo/poisson_matrix_free/main.cpp |  3 +-
 cpp/dolfinx/la/Vector.h               | 72 ++++++++++++++-------------
 cpp/test/matrix.cpp                   |  8 ++-
 cpp/test/vector.cpp                   |  2 +-
 python/dolfinx/wrappers/la.cpp        | 25 +++++-----
 7 files changed, 58 insertions(+), 59 deletions(-)

diff --git a/cpp/demo/biharmonic/main.cpp b/cpp/demo/biharmonic/main.cpp
index a08d9d0033a..9a50787f403 100644
--- a/cpp/demo/biharmonic/main.cpp
+++ b/cpp/demo/biharmonic/main.cpp
@@ -236,9 +236,8 @@ int main(int argc, char* argv[])
     // Compute solution
     fem::Function<T> u(V);
     auto A = la::petsc::Matrix(fem::petsc::create_matrix(*a), false);
-    la::Vector<std::vector<T>> b(
-        L->function_spaces()[0]->dofmap()->index_map,
-        L->function_spaces()[0]->dofmap()->index_map_bs());
+    la::Vector<T> b(L->function_spaces()[0]->dofmap()->index_map,
+                    L->function_spaces()[0]->dofmap()->index_map_bs());
 
     MatZeroEntries(A.mat());
     fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES),
diff --git a/cpp/demo/hyperelasticity/main.cpp b/cpp/demo/hyperelasticity/main.cpp
index 0f460488b59..99f27490066 100644
--- a/cpp/demo/hyperelasticity/main.cpp
+++ b/cpp/demo/hyperelasticity/main.cpp
@@ -100,7 +100,7 @@ class HyperElasticProblem
 private:
   std::shared_ptr<fem::Form<T>> _l, _j;
   std::vector<std::shared_ptr<const fem::DirichletBC<T>>> _bcs;
-  la::Vector<std::vector<T>> _b;
+  la::Vector<T> _b;
   Vec _b_petsc = nullptr;
   la::petsc::Matrix _matA;
 };
diff --git a/cpp/demo/poisson_matrix_free/main.cpp b/cpp/demo/poisson_matrix_free/main.cpp
index 837635588b6..fa74a6cdda0 100644
--- a/cpp/demo/poisson_matrix_free/main.cpp
+++ b/cpp/demo/poisson_matrix_free/main.cpp
@@ -164,8 +164,7 @@ int main(int argc, char* argv[])
     auto bc = std::make_shared<const fem::DirichletBC<T>>(u_D, bdofs);
 
     // Assemble RHS vector
-    la::Vector<std::vector<T>> b(V->dofmap()->index_map,
-                                 V->dofmap()->index_map_bs());
+    la::Vector<T> b(V->dofmap()->index_map, V->dofmap()->index_map_bs());
     fem::assemble_vector(b.mutable_array(), *L);
 
     // Apply lifting to account for Dirichlet boundary condition
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index bb66711878d..334e15e54b0 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -24,27 +24,29 @@ namespace dolfinx::la
 ///
 /// @tparam V data container type
 ///
-template <class V>
+template <typename Scalar, typename Container = std::vector<Scalar>>
 class Vector
 {
-
-  template <typename T>
-  struct is_scalar_type_t : public std::is_arithmetic<T>
-  {
-  };
-  template <typename T>
-  struct is_scalar_type_t<std::complex<T>> : public std::true_type
-  {
-  };
-
-  using V_t = typename std::conditional_t<is_scalar_type_t<V>::value,
-                                          std::vector<V>, V>;
-  using T = typename std::conditional_t<is_scalar_type_t<V>::value, V,
-                                        typename V_t::value_type>;
+  // template <typename T>
+  // struct is_scalar_type_t : public std::is_arithmetic<T>
+  // {
+  // };
+  // template <typename T>
+  // struct is_scalar_type_t<std::complex<T>> : public std::true_type
+  // {
+  // };
+
+  // using V_t = typename std::conditional_t<is_scalar_type_t<V>::value,
+  //                                         std::vector<V>, V>;
+  // using T = typename std::conditional_t<is_scalar_type_t<V>::value, V,
+  // typename V_t::value_type > ;
 
 public:
-  /// value_type Scalar type of Vector
-  using value_type = T;
+  /// Scalar type of Vector
+  using value_type = Scalar;
+
+  /// Container type
+  using container_type = Container;
 
   /// Create a distributed vector
   /// @param map IndexMap for parallel distribution of the data
@@ -83,14 +85,14 @@ class Vector
 
   /// Set all entries (including ghosts)
   /// @param[in] v The value to set all entries to (on calling rank)
-  void set(T v) { std::fill(_x.begin(), _x.end(), v); }
+  void set(Scalar v) { std::fill(_x.begin(), _x.end(), v); }
 
   /// Begin scatter of local data from owner to ghosts on other ranks
   /// @note Collective MPI operation
   void scatter_fwd_begin()
   {
     const std::int32_t local_size = _bs * _map->size_local();
-    std::span<const T> x_local(_x.data(), local_size);
+    std::span<const Scalar> x_local(_x.data(), local_size);
 
     auto pack = [](const auto& in, const auto& idx, auto& out)
     {
@@ -99,8 +101,8 @@ class Vector
     };
     pack(x_local, _scatterer->local_indices(), _buffer_local);
 
-    _scatterer->scatter_fwd_begin(std::span<const T>(_buffer_local),
-                                  std::span<T>(_buffer_remote),
+    _scatterer->scatter_fwd_begin(std::span<const Scalar>(_buffer_local),
+                                  std::span<Scalar>(_buffer_remote),
                                   std::span<MPI_Request>(_request));
   }
 
@@ -110,7 +112,7 @@ class Vector
   {
     const std::int32_t local_size = _bs * _map->size_local();
     const std::int32_t num_ghosts = _bs * _map->num_ghosts();
-    std::span<T> x_remote(_x.data() + local_size, num_ghosts);
+    std::span<Scalar> x_remote(_x.data() + local_size, num_ghosts);
     _scatterer->scatter_fwd_end(std::span<MPI_Request>(_request));
 
     auto unpack = [](const auto& in, const auto& idx, auto& out, auto op)
@@ -137,7 +139,7 @@ class Vector
   {
     const std::int32_t local_size = _bs * _map->size_local();
     const std::int32_t num_ghosts = _bs * _map->num_ghosts();
-    std::span<T> x_remote(_x.data() + local_size, num_ghosts);
+    std::span<Scalar> x_remote(_x.data() + local_size, num_ghosts);
 
     auto pack = [](const auto& in, const auto& idx, auto& out)
     {
@@ -146,8 +148,8 @@ class Vector
     };
     pack(x_remote, _scatterer->remote_indices(), _buffer_remote);
 
-    _scatterer->scatter_rev_begin(std::span<const T>(_buffer_remote),
-                                  std::span<T>(_buffer_local), _request);
+    _scatterer->scatter_rev_begin(std::span<const Scalar>(_buffer_remote),
+                                  std::span<Scalar>(_buffer_local), _request);
   }
 
   /// End scatter of ghost data to owner. This process may receive data
@@ -160,7 +162,7 @@ class Vector
   void scatter_rev_end(BinaryOperation op)
   {
     const std::int32_t local_size = _bs * _map->size_local();
-    std::span<T> x_local(_x.data(), local_size);
+    std::span<Scalar> x_local(_x.data(), local_size);
     _scatterer->scatter_rev_end(_request);
 
     auto unpack = [](const auto& in, const auto& idx, auto& out, auto op)
@@ -190,10 +192,10 @@ class Vector
   constexpr int bs() const { return _bs; }
 
   /// Get local part of the vector (const version)
-  std::span<const T> array() const { return std::span<const T>(_x); }
+  std::span<const Scalar> array() const { return std::span<const Scalar>(_x); }
 
   /// Get local part of the vector
-  std::span<T> mutable_array() { return std::span(_x); }
+  std::span<Scalar> mutable_array() { return std::span(_x); }
 
 private:
   // Map describing the data layout
@@ -209,10 +211,10 @@ class Vector
   std::vector<MPI_Request> _request = {MPI_REQUEST_NULL};
 
   // Buffers for ghost scatters
-  V_t _buffer_local, _buffer_remote;
+  Container _buffer_local, _buffer_remote;
 
   // Vector data
-  V_t _x;
+  Container _x;
 };
 
 /// Compute the inner product of two vectors. The two vectors must have
@@ -222,7 +224,7 @@ class Vector
 /// @param b A vector
 /// @return Returns `a^{H} b` (`a^{T} b` if `a` and `b` are real)
 template <class V>
-auto inner_product(const Vector<V>& a, const Vector<V>& b)
+auto inner_product(const V& a, const V& b)
 {
   using T = typename V::value_type;
   const std::int32_t local_size = a.bs() * a.map()->size_local();
@@ -253,7 +255,7 @@ auto inner_product(const Vector<V>& a, const Vector<V>& b)
 /// Compute the squared L2 norm of vector
 /// @note Collective MPI operation
 template <class V>
-auto squared_norm(const Vector<V>& a)
+auto squared_norm(const V& a)
 {
   using T = typename V::value_type;
   T result = inner_product(a, a);
@@ -265,7 +267,7 @@ auto squared_norm(const Vector<V>& a)
 /// @param a A vector
 /// @param type Norm type (supported types are \f$L^2\f$ and \f$L^\infty\f$)
 template <class V>
-auto norm(const Vector<V>& a, Norm type = Norm::l2)
+auto norm(const V& a, Norm type = Norm::l2)
 {
   using T = typename V::value_type;
   switch (type)
@@ -296,7 +298,7 @@ auto norm(const Vector<V>& a, Norm type = Norm::l2)
 /// modified in-place.
 /// @param[in] tol The tolerance used to detect a linear dependency
 template <class V>
-void orthonormalize(std::span<Vector<V>> basis, double tol = 1.0e-10)
+void orthonormalize(std::span<V> basis, double tol = 1.0e-10)
 {
   using T = typename V::value_type;
   // Loop over each vector in basis
@@ -331,7 +333,7 @@ void orthonormalize(std::span<Vector<V>> basis, double tol = 1.0e-10)
 /// @param[in] tol The tolerance used to test for orthonormality
 /// @return True is basis is orthonormal, otherwise false
 template <class V>
-bool is_orthonormal(std::span<const Vector<V>> basis, double tol = 1.0e-10)
+bool is_orthonormal(std::span<const V> basis, double tol = 1.0e-10)
 {
   using T = typename V::value_type;
   for (std::size_t i = 0; i < basis.size(); i++)
diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
index 872e7b96ffc..a9372907bb8 100644
--- a/cpp/test/matrix.cpp
+++ b/cpp/test/matrix.cpp
@@ -71,9 +71,7 @@ void spmv_impl(std::span<const T> values,
 /// @param[in] x Input vector
 /// @param[in, out] y Output vector
 template <typename T>
-void spmv(la::MatrixCSR<std::vector<T>>& A, la::Vector<std::vector<T>>& x,
-          la::Vector<std::vector<T>>& y)
-
+void spmv(la::MatrixCSR<std::vector<T>>& A, la::Vector<T>& x, la::Vector<T>& y)
 {
   // start communication (update ghosts)
   x.scatter_fwd_begin();
@@ -170,8 +168,8 @@ la::MatrixCSR<std::vector<double>> create_operator(MPI_Comm comm)
   // Get compatible vectors
   auto maps = A.index_maps();
 
-  la::Vector<std::vector<double>> x(maps[1], 1);
-  la::Vector<std::vector<double>> y(maps[1], 1);
+  la::Vector<double> x(maps[1], 1);
+  la::Vector<double> y(maps[1], 1);
 
   std::size_t col_size = maps[1]->size_local() + maps[1]->num_ghosts();
   CHECK(x.array().size() == col_size);
diff --git a/cpp/test/vector.cpp b/cpp/test/vector.cpp
index 6baf7583ce8..073e9bf585e 100644
--- a/cpp/test/vector.cpp
+++ b/cpp/test/vector.cpp
@@ -36,7 +36,7 @@ void test_vector()
   auto index_map = std::make_shared<common::IndexMap>(
       MPI_COMM_WORLD, size_local, ghosts, global_ghost_owner);
 
-  la::Vector<std::vector<T>> v(index_map, 1);
+  la::Vector<T> v(index_map, 1);
   std::fill(v.mutable_array().begin(), v.mutable_array().end(), 1.0);
 
   double norm2 = la::squared_norm(v);
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index aa5ffcc7e8d..dbf5575dd26 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -38,38 +38,39 @@ template <typename V>
 void declare_objects(py::module& m, const std::string& type)
 {
   using T = typename V::value_type;
+  // using V = typename V::value_type;
 
   // dolfinx::la::Vector
   std::string pyclass_vector_name = std::string("Vector_") + type;
-  py::class_<dolfinx::la::Vector<V>, std::shared_ptr<dolfinx::la::Vector<V>>>(
+  py::class_<dolfinx::la::Vector<T>, std::shared_ptr<dolfinx::la::Vector<T>>>(
       m, pyclass_vector_name.c_str())
       .def(py::init([](std::shared_ptr<const dolfinx::common::IndexMap> map,
-                       int bs) { return dolfinx::la::Vector<V>(map, bs); }),
+                       int bs) { return dolfinx::la::Vector<T>(map, bs); }),
            py::arg("map"), py::arg("bs"))
-      .def(py::init([](const dolfinx::la::Vector<V>& vec)
-                    { return dolfinx::la::Vector<V>(vec); }),
+      .def(py::init([](const dolfinx::la::Vector<T>& vec)
+                    { return dolfinx::la::Vector<T>(vec); }),
            py::arg("vec"))
-      .def_property_readonly("dtype", [](const dolfinx::la::Vector<V>& self)
+      .def_property_readonly("dtype", [](const dolfinx::la::Vector<T>& self)
                              { return py::dtype::of<T>(); })
-      .def("set", &dolfinx::la::Vector<V>::set, py::arg("v"))
+      .def("set", &dolfinx::la::Vector<T>::set, py::arg("v"))
       .def(
           "norm",
-          [](dolfinx::la::Vector<V>& self, dolfinx::la::Norm type)
+          [](dolfinx::la::Vector<T>& self, dolfinx::la::Norm type)
           { return dolfinx::la::norm(self, type); },
           py::arg("type") = dolfinx::la::Norm::l2)
-      .def_property_readonly("map", &dolfinx::la::Vector<V>::map)
-      .def_property_readonly("bs", &dolfinx::la::Vector<V>::bs)
+      .def_property_readonly("map", &dolfinx::la::Vector<T>::map)
+      .def_property_readonly("bs", &dolfinx::la::Vector<T>::bs)
       .def_property_readonly("array",
-                             [](dolfinx::la::Vector<V>& self)
+                             [](dolfinx::la::Vector<T>& self)
                              {
                                std::span<T> array = self.mutable_array();
                                return py::array_t<T>(array.size(), array.data(),
                                                      py::cast(self));
                              })
-      .def("scatter_forward", &dolfinx::la::Vector<V>::scatter_fwd)
+      .def("scatter_forward", &dolfinx::la::Vector<T>::scatter_fwd)
       .def(
           "scatter_reverse",
-          [](dolfinx::la::Vector<V>& self, PyScatterMode mode)
+          [](dolfinx::la::Vector<T>& self, PyScatterMode mode)
           {
             switch (mode)
             {

From 3698311ab2b343661cbd25f052366af25b81b482 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 08:43:54 +0100
Subject: [PATCH 14/21] Updates

---
 cpp/dolfinx/la/MatrixCSR.h           | 69 +++++++++++++++-------------
 cpp/test/matrix.cpp                  | 10 ++--
 python/dolfinx/wrappers/assemble.cpp |  6 +--
 python/dolfinx/wrappers/la.cpp       | 43 ++++++++---------
 4 files changed, 64 insertions(+), 64 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 4fe931eb3bd..4d3242ea808 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -113,24 +113,28 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// The matrix storage format is compressed sparse row. The matrix is
 /// partitioned row-wise across MPI rank.
 ///
-/// @tparam V The data container type for the matrix
-/// @tparam Idx The index container type
-///
-/// @note Highly "experimental" storage of a matrix in CSR format which
-/// can be assembled into using the usual dolfinx assembly routines
+/// @warning Highly experimental storage of a matrix in CSR format which
+/// can be assembled into using the usual DOLFINx assembly routines
 /// Matrix internal data can be accessed for interfacing with other
 /// code.
-///
 /// @todo Handle block sizes
-template <class V, class Idx = std::vector<std::int32_t>>
+///
+/// @tparam Scalar Scalar type
+/// @tparam Container Sequence container type to store matrix entries
+/// @tparam Idx Index container type
+template <class Scalar, class Container = std::vector<Scalar>,
+          class Idx = std::vector<std::int32_t>>
 class MatrixCSR
 {
 public:
   /// The value type
-  using value_type = typename V::value_type;
+  using value_type = Scalar;
 
-  /// Insertion functor for setting values in matrix. It is typically
-  /// used in finite element assembly functions.
+  /// Container type
+  using container_type = Container;
+
+  /// @brief Insertion functor for setting values in matrix. It is
+  /// typically used in finite element assembly functions.
   /// @return Function for inserting values into `A`
   /// @todo clarify setting on non-owned entries
   auto mat_set_values()
@@ -144,7 +148,7 @@ class MatrixCSR
     };
   }
 
-  /// Insertion functor for accumulating values in matrix. It is
+  /// @brief Insertion functor for accumulating values in matrix. It is
   /// typically used in finite element assembly functions.
   /// @return Function for inserting values into `A`
   auto mat_add_values()
@@ -158,9 +162,9 @@ class MatrixCSR
     };
   }
 
-  /// Create a distributed matrix
-  /// @param[in] p The sparsty pattern the describes the parallel
-  /// distribution and the non-zero structure
+  /// @brief Create a distributed matrix.
+  /// @param[in] p The sparsity pattern the describes the parallel
+  /// distribution and the non-zero structure.
   MatrixCSR(const SparsityPattern& p)
       : _index_maps({p.index_map(0),
                      std::make_shared<common::IndexMap>(p.column_index_map())}),
@@ -328,12 +332,12 @@ class MatrixCSR
   /// @todo Check handling of MPI_Request
   MatrixCSR(MatrixCSR&& A) = default;
 
-  /// Set all non-zero local entries to a value
-  /// including entries in ghost rows
+  /// @brief Set all non-zero local entries to a value including entries
+  /// in ghost rows.
   /// @param[in] x The value to set non-zero matrix entries to
   void set(value_type x) { std::fill(_data.begin(), _data.end(), x); }
 
-  /// Set values in the matrix
+  /// @brief Set values in the matrix.
   /// @note Only entries included in the sparsity pattern used to
   /// initialize the matrix can be set
   /// @note All indices are local to the calling MPI rank and entries
@@ -353,7 +357,7 @@ class MatrixCSR
                   _index_maps[0]->size_local());
   }
 
-  /// Accumulate values in the matrix
+  /// @brief Accumulate values in the matrix
   /// @note Only entries included in the sparsity pattern used to
   /// initialize the matrix can be accumulated in to
   /// @note All indices are local to the calling MPI rank and entries
@@ -398,17 +402,17 @@ class MatrixCSR
     return A;
   }
 
-  /// Transfer ghost row data to the owning ranks
-  /// accumulating received values on the owned rows, and zeroing any existing
-  /// data in ghost rows.
+  /// @brief Transfer ghost row data to the owning ranks accumulating
+  /// received values on the owned rows, and zeroing any existing data
+  /// in ghost rows.
   void finalize()
   {
     finalize_begin();
     finalize_end();
   }
 
-  /// Begin transfer of ghost row data to owning ranks, where it will be
-  /// accumulated into existing owned rows.
+  /// @brief Begin transfer of ghost row data to owning ranks, where it
+  /// will be accumulated into existing owned rows.
   /// @note Calls to this function must be followed by
   /// MatrixCSR::finalize_end(). Between the two calls matrix values
   /// must not be changed.
@@ -455,7 +459,7 @@ class MatrixCSR
     assert(status == MPI_SUCCESS);
   }
 
-  /// End transfer of ghost row data to owning ranks
+  /// @brief End transfer of ghost row data to owning ranks.
   /// @note Must be preceded by MatrixCSR::finalize_begin()
   /// @note Matrix data received from other processes will be
   /// accumulated into locally owned rows, and ghost rows will be
@@ -490,10 +494,11 @@ class MatrixCSR
     return norm_sq;
   }
 
-  /// Index maps for the row and column space. The row IndexMap contains
-  /// ghost entries for rows which may be inserted into and the column
-  /// IndexMap contains all local and ghost columns that may exist in
-  /// the owned rows.
+  /// @brief Index maps for the row and column space.
+  ///
+  /// The row IndexMap contains ghost entries for rows which may be
+  /// inserted into and the column IndexMap contains all local and ghost
+  /// columns that may exist in the owned rows.
   ///
   /// @return Row (0) and column (1) index maps
   const std::array<std::shared_ptr<const common::IndexMap>, 2>&
@@ -504,11 +509,11 @@ class MatrixCSR
 
   /// Get local data values
   /// @note Includes ghost values
-  V& values() { return _data; }
+  container_type& values() { return _data; }
 
   /// Get local values (const version)
   /// @note Includes ghost values
-  const V& values() const { return _data; }
+  const container_type& values() const { return _data; }
 
   /// Get local row pointers
   /// @note Includes pointers to ghost rows
@@ -538,7 +543,7 @@ class MatrixCSR
   std::array<int, 2> _bs;
 
   // Matrix data
-  V _data;
+  Container _data;
   Idx _cols, _row_ptr;
 
   // Start of off-diagonal (unowned columns) on each row
@@ -564,7 +569,7 @@ class MatrixCSR
   std::vector<int> _ghost_row_to_rank;
 
   // Temporary store for finalize data during non-blocking communication
-  V _ghost_value_data_in;
+  container_type _ghost_value_data_in;
 };
 
 } // namespace dolfinx::la
diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
index a9372907bb8..e69cacf811f 100644
--- a/cpp/test/matrix.cpp
+++ b/cpp/test/matrix.cpp
@@ -71,7 +71,7 @@ void spmv_impl(std::span<const T> values,
 /// @param[in] x Input vector
 /// @param[in, out] y Output vector
 template <typename T>
-void spmv(la::MatrixCSR<std::vector<T>>& A, la::Vector<T>& x, la::Vector<T>& y)
+void spmv(la::MatrixCSR<T>& A, la::Vector<T>& x, la::Vector<T>& y)
 {
   // start communication (update ghosts)
   x.scatter_fwd_begin();
@@ -104,7 +104,7 @@ void spmv(la::MatrixCSR<std::vector<T>>& A, la::Vector<T>& x, la::Vector<T>& y)
 /// @brief Create a matrix operator
 /// @param comm The communicator to builf the matrix on
 /// @return The assembled matrix
-la::MatrixCSR<std::vector<double>> create_operator(MPI_Comm comm)
+la::MatrixCSR<double> create_operator(MPI_Comm comm)
 {
   auto part = mesh::create_cell_partitioner(mesh::GhostMode::none);
   auto mesh = std::make_shared<mesh::Mesh<double>>(
@@ -121,7 +121,7 @@ la::MatrixCSR<std::vector<double>> create_operator(MPI_Comm comm)
 
   la::SparsityPattern sp = fem::create_sparsity_pattern(*a);
   sp.assemble();
-  la::MatrixCSR<std::vector<double>> A(sp);
+  la::MatrixCSR<double> A(sp);
   fem::assemble_matrix(A.mat_add_values(), *a, {});
   A.finalize();
 
@@ -160,7 +160,7 @@ la::MatrixCSR<std::vector<double>> create_operator(MPI_Comm comm)
   sp.assemble();
 
   // Assemble matrix
-  la::MatrixCSR<std::vector<double>> A(sp);
+  la::MatrixCSR<double> A(sp);
   fem::assemble_matrix(A.mat_add_values(), *a, {});
   A.finalize();
   CHECK((V->dofmap()->index_map->size_local() == A.num_owned_rows()));
@@ -195,7 +195,7 @@ void test_matrix()
   p.assemble();
 
   using T = float;
-  la::MatrixCSR<std::vector<T>> A(p);
+  la::MatrixCSR<T> A(p);
   A.add(std::vector<decltype(A)::value_type>{1}, std::vector{0},
         std::vector{0});
   A.add(std::vector<decltype(A)::value_type>{2.3}, std::vector{4},
diff --git a/python/dolfinx/wrappers/assemble.cpp b/python/dolfinx/wrappers/assemble.cpp
index 9bc3e49dca4..967fc7e287d 100644
--- a/python/dolfinx/wrappers/assemble.cpp
+++ b/python/dolfinx/wrappers/assemble.cpp
@@ -139,8 +139,7 @@ void declare_assembly_functions(py::module& m)
   // MatrixCSR
   m.def(
       "assemble_matrix",
-      [](dolfinx::la::MatrixCSR<std::vector<T>>& A,
-         const dolfinx::fem::Form<T, U>& a,
+      [](dolfinx::la::MatrixCSR<T>& A, const dolfinx::fem::Form<T, U>& a,
          const py::array_t<T, py::array::c_style>& constants,
          const std::map<std::pair<dolfinx::fem::IntegralType, int>,
                         py::array_t<T, py::array::c_style>>& coefficients,
@@ -162,8 +161,7 @@ void declare_assembly_functions(py::module& m)
       py::arg("bcs"), "Experimental.");
   m.def(
       "insert_diagonal",
-      [](dolfinx::la::MatrixCSR<std::vector<T>>& A,
-         const dolfinx::fem::FunctionSpace<U>& V,
+      [](dolfinx::la::MatrixCSR<T>& A, const dolfinx::fem::FunctionSpace<U>& V,
          const std::vector<
              std::shared_ptr<const dolfinx::fem::DirichletBC<T, U>>>& bcs,
          T diagonal)
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index dbf5575dd26..1d9e51b52b7 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -34,12 +34,9 @@ enum class PyScatterMode
 };
 
 // Declare objects that have multiple scalar types
-template <typename V>
+template <typename T>
 void declare_objects(py::module& m, const std::string& type)
 {
-  using T = typename V::value_type;
-  // using V = typename V::value_type;
-
   // dolfinx::la::Vector
   std::string pyclass_vector_name = std::string("Vector_") + type;
   py::class_<dolfinx::la::Vector<T>, std::shared_ptr<dolfinx::la::Vector<T>>>(
@@ -89,23 +86,23 @@ void declare_objects(py::module& m, const std::string& type)
 
   // dolfinx::la::MatrixCSR
   std::string pyclass_matrix_name = std::string("MatrixCSR_") + type;
-  py::class_<dolfinx::la::MatrixCSR<V>,
-             std::shared_ptr<dolfinx::la::MatrixCSR<V>>>(
+  py::class_<dolfinx::la::MatrixCSR<T>,
+             std::shared_ptr<dolfinx::la::MatrixCSR<T>>>(
       m, pyclass_matrix_name.c_str())
       .def(py::init([](const dolfinx::la::SparsityPattern& p)
-                    { return dolfinx::la::MatrixCSR<V>(p); }),
+                    { return dolfinx::la::MatrixCSR<T>(p); }),
            py::arg("p"))
-      .def_property_readonly("dtype", [](const dolfinx::la::MatrixCSR<V>& self)
+      .def_property_readonly("dtype", [](const dolfinx::la::MatrixCSR<T>& self)
                              { return py::dtype::of<T>(); })
-      .def("norm_squared", &dolfinx::la::MatrixCSR<V>::norm_squared)
-      .def("mat_add_values", &dolfinx::la::MatrixCSR<V>::mat_add_values)
+      .def("norm_squared", &dolfinx::la::MatrixCSR<T>::norm_squared)
+      .def("mat_add_values", &dolfinx::la::MatrixCSR<T>::mat_add_values)
       .def("set",
-           static_cast<void (dolfinx::la::MatrixCSR<V>::*)(T)>(
-               &dolfinx::la::MatrixCSR<V>::set),
+           static_cast<void (dolfinx::la::MatrixCSR<T>::*)(T)>(
+               &dolfinx::la::MatrixCSR<T>::set),
            py::arg("x"))
-      .def("finalize", &dolfinx::la::MatrixCSR<V>::finalize)
+      .def("finalize", &dolfinx::la::MatrixCSR<T>::finalize)
       .def("to_dense",
-           [](const dolfinx::la::MatrixCSR<V>& self)
+           [](const dolfinx::la::MatrixCSR<T>& self)
            {
              std::size_t nrows = self.num_all_rows();
              auto map_col = self.index_maps()[1];
@@ -114,14 +111,14 @@ void declare_objects(py::module& m, const std::string& type)
                                                  std::array{nrows, ncols});
            })
       .def_property_readonly("data",
-                             [](dolfinx::la::MatrixCSR<V>& self)
+                             [](dolfinx::la::MatrixCSR<T>& self)
                              {
                                std::span<T> array = self.values();
                                return py::array_t<T>(array.size(), array.data(),
                                                      py::cast(self));
                              })
       .def_property_readonly("indices",
-                             [](dolfinx::la::MatrixCSR<V>& self)
+                             [](dolfinx::la::MatrixCSR<T>& self)
                              {
                                std::span<const std::int32_t> array
                                    = self.cols();
@@ -129,15 +126,15 @@ void declare_objects(py::module& m, const std::string& type)
                                    array.size(), array.data(), py::cast(self));
                              })
       .def_property_readonly("indptr",
-                             [](dolfinx::la::MatrixCSR<V>& self)
+                             [](dolfinx::la::MatrixCSR<T>& self)
                              {
                                std::span<const std::int32_t> array
                                    = self.row_ptr();
                                return py::array_t<const std::int32_t>(
                                    array.size(), array.data(), py::cast(self));
                              })
-      .def("finalize_begin", &dolfinx::la::MatrixCSR<V>::finalize_begin)
-      .def("finalize_end", &dolfinx::la::MatrixCSR<V>::finalize_end);
+      .def("finalize_begin", &dolfinx::la::MatrixCSR<T>::finalize_begin)
+      .def("finalize_end", &dolfinx::la::MatrixCSR<T>::finalize_end);
 }
 
 void petsc_module(py::module& m)
@@ -274,9 +271,9 @@ void la(py::module& m)
                              py::return_value_policy::reference_internal);
 
   // Declare objects that are templated over type
-  declare_objects<std::vector<float>>(m, "float32");
-  declare_objects<std::vector<double>>(m, "float64");
-  declare_objects<std::vector<std::complex<float>>>(m, "complex64");
-  declare_objects<std::vector<std::complex<double>>>(m, "complex128");
+  declare_objects<float>(m, "float32");
+  declare_objects<double>(m, "float64");
+  declare_objects<std::complex<float>>(m, "complex64");
+  declare_objects<std::complex<double>>(m, "complex128");
 }
 } // namespace dolfinx_wrappers

From bc54276551372985c7f197472c8549de8ad3ae46 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Wed, 26 Apr 2023 11:16:41 +0100
Subject: [PATCH 15/21] Tidy up and make consistent

---
 cpp/dolfinx/la/MatrixCSR.h | 26 ++++++++++++++------------
 cpp/dolfinx/la/Vector.h    | 16 ++--------------
 2 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 4d3242ea808..a476ed9a8f6 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -126,8 +126,10 @@ template <class Scalar, class Container = std::vector<Scalar>,
           class Idx = std::vector<std::int32_t>>
 class MatrixCSR
 {
+  static_assert(std::is_same_v<typename Container::value_type, Scalar>);
+
 public:
-  /// The value type
+  /// Scalar type
   using value_type = Scalar;
 
   /// Container type
@@ -141,7 +143,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const value_type>& data) -> int
+               const std::span<const Scalar>& data) -> int
     {
       this->set(data, rows, cols);
       return 0;
@@ -155,7 +157,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const value_type>& data) -> int
+               const std::span<const Scalar>& data) -> int
     {
       this->add(data, rows, cols);
       return 0;
@@ -335,7 +337,7 @@ class MatrixCSR
   /// @brief Set all non-zero local entries to a value including entries
   /// in ghost rows.
   /// @param[in] x The value to set non-zero matrix entries to
-  void set(value_type x) { std::fill(_data.begin(), _data.end(), x); }
+  void set(Scalar x) { std::fill(_data.begin(), _data.end(), x); }
 
   /// @brief Set values in the matrix.
   /// @note Only entries included in the sparsity pattern used to
@@ -349,7 +351,7 @@ class MatrixCSR
   /// set in the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void set(const std::span<const value_type>& x,
+  void set(const std::span<const Scalar>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -369,7 +371,7 @@ class MatrixCSR
   /// add to the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void add(const std::span<const value_type>& x,
+  void add(const std::span<const Scalar>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -389,12 +391,12 @@ class MatrixCSR
   /// manually by using num_owned_rows() if required.
   /// @return Dense copy of the part of the matrix on the calling rank.
   /// Storage is row-major.
-  std::vector<value_type> to_dense() const
+  std::vector<Scalar> to_dense() const
   {
     const std::size_t nrows = num_all_rows();
     const std::size_t ncols
         = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
-    std::vector<value_type> A(nrows * ncols);
+    std::vector<Scalar> A(nrows * ncols);
     for (std::size_t r = 0; r < nrows; ++r)
       for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
         A[r * ncols + _cols[j]] = _data[j];
@@ -425,7 +427,7 @@ class MatrixCSR
 
     // For each ghost row, pack and send values to send to neighborhood
     std::vector<int> insert_pos = _val_send_disp;
-    std::vector<value_type> ghost_value_data(_val_send_disp.back());
+    std::vector<Scalar> ghost_value_data(_val_send_disp.back());
     for (int i = 0; i < num_ghosts0; ++i)
     {
       const int rank = _ghost_row_to_rank[i];
@@ -453,9 +455,9 @@ class MatrixCSR
 
     int status = MPI_Ineighbor_alltoallv(
         ghost_value_data.data(), val_send_count.data(), _val_send_disp.data(),
-        dolfinx::MPI::mpi_type<value_type>(), _ghost_value_data_in.data(),
+        dolfinx::MPI::mpi_type<Scalar>(), _ghost_value_data_in.data(),
         val_recv_count.data(), _val_recv_disp.data(),
-        dolfinx::MPI::mpi_type<value_type>(), _comm.comm(), &_request);
+        dolfinx::MPI::mpi_type<Scalar>(), _comm.comm(), &_request);
     assert(status == MPI_SUCCESS);
   }
 
@@ -487,7 +489,7 @@ class MatrixCSR
 
     const double norm_sq_local = std::accumulate(
         _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
-        double(0), [](auto norm, value_type y) { return norm + std::norm(y); });
+        double(0), [](auto norm, Scalar y) { return norm + std::norm(y); });
     double norm_sq;
     MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM,
                   _comm.comm());
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 334e15e54b0..632581acf91 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -27,22 +27,10 @@ namespace dolfinx::la
 template <typename Scalar, typename Container = std::vector<Scalar>>
 class Vector
 {
-  // template <typename T>
-  // struct is_scalar_type_t : public std::is_arithmetic<T>
-  // {
-  // };
-  // template <typename T>
-  // struct is_scalar_type_t<std::complex<T>> : public std::true_type
-  // {
-  // };
-
-  // using V_t = typename std::conditional_t<is_scalar_type_t<V>::value,
-  //                                         std::vector<V>, V>;
-  // using T = typename std::conditional_t<is_scalar_type_t<V>::value, V,
-  // typename V_t::value_type > ;
+  static_assert(std::is_same_v<typename Container::value_type, Scalar>);
 
 public:
-  /// Scalar type of Vector
+  /// Scalar type
   using value_type = Scalar;
 
   /// Container type

From c8c4901b7cdb3faf21caa7c381902ce4d6796383 Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Wed, 26 Apr 2023 11:20:11 +0100
Subject: [PATCH 16/21] Remove commented code

---
 cpp/demo/poisson_matrix_free/main.cpp | 3 ---
 cpp/dolfinx/la/MatrixCSR.h            | 6 +++---
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/cpp/demo/poisson_matrix_free/main.cpp b/cpp/demo/poisson_matrix_free/main.cpp
index fa74a6cdda0..0c2512f429a 100644
--- a/cpp/demo/poisson_matrix_free/main.cpp
+++ b/cpp/demo/poisson_matrix_free/main.cpp
@@ -59,13 +59,10 @@ void axpy(auto& r, auto alpha, const auto& x, const auto& y)
 /// @return The number if iterations
 /// @pre It is required that the ghost values of `x` and `b` have been
 /// updated before this function is called
-// template <typename U, typename ApplyFunction>
-// int cg(la::Vector<U>& x, const la::Vector<U>& b, ApplyFunction&& action,
 template <typename ApplyFunction>
 int cg(auto& x, auto& b, ApplyFunction&& action, int kmax = 50,
        double rtol = 1e-8)
 {
-  // using T = typename U::value_type;
   using T = typename std::decay_t<decltype(x)>::value_type;
 
   // Create working vectors
diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index a476ed9a8f6..d9807e7a8ea 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -511,11 +511,11 @@ class MatrixCSR
 
   /// Get local data values
   /// @note Includes ghost values
-  container_type& values() { return _data; }
+  Container& values() { return _data; }
 
   /// Get local values (const version)
   /// @note Includes ghost values
-  const container_type& values() const { return _data; }
+  const Container& values() const { return _data; }
 
   /// Get local row pointers
   /// @note Includes pointers to ghost rows
@@ -571,7 +571,7 @@ class MatrixCSR
   std::vector<int> _ghost_row_to_rank;
 
   // Temporary store for finalize data during non-blocking communication
-  container_type _ghost_value_data_in;
+  Container _ghost_value_data_in;
 };
 
 } // namespace dolfinx::la

From 8fc87556dc8617cfdab12537e92bb897b23e0f57 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 11:29:55 +0100
Subject: [PATCH 17/21] Add static type checks

---
 cpp/dolfinx/la/MatrixCSR.h | 3 +++
 cpp/dolfinx/la/Vector.h    | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 4d3242ea808..24e7efa1839 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -133,6 +133,9 @@ class MatrixCSR
   /// Container type
   using container_type = Container;
 
+  static_assert(std::is_same_v<value_type, typename container_type::value_type>,
+                "Scalar type and container value type must be the same.");
+
   /// @brief Insertion functor for setting values in matrix. It is
   /// typically used in finite element assembly functions.
   /// @return Function for inserting values into `A`
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 334e15e54b0..69449fedf9e 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -48,6 +48,9 @@ class Vector
   /// Container type
   using container_type = Container;
 
+  static_assert(std::is_same_v<value_type, typename container_type::value_type>,
+                "Scalar type and container value type must be the same.");
+
   /// Create a distributed vector
   /// @param map IndexMap for parallel distribution of the data
   /// @param bs Block size

From ca6b7e0935bf8df031e3be4b6e395a079eb62cac Mon Sep 17 00:00:00 2001
From: Chris Richardson <chris@bpi.cam.ac.uk>
Date: Wed, 26 Apr 2023 11:50:43 +0100
Subject: [PATCH 18/21] Change some naming

---
 cpp/dolfinx/la/MatrixCSR.h | 38 +++++++++++++++++++-------------------
 cpp/dolfinx/la/Vector.h    | 30 +++++++++++++++++-------------
 2 files changed, 36 insertions(+), 32 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index d9807e7a8ea..e7142de4a46 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -121,9 +121,9 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 ///
 /// @tparam Scalar Scalar type
 /// @tparam Container Sequence container type to store matrix entries
-/// @tparam Idx Index container type
+/// @tparam IdxContainer Index container type
 template <class Scalar, class Container = std::vector<Scalar>,
-          class Idx = std::vector<std::int32_t>>
+          class IdxContainer = std::vector<std::int32_t>>
 class MatrixCSR
 {
   static_assert(std::is_same_v<typename Container::value_type, Scalar>);
@@ -143,7 +143,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const Scalar>& data) -> int
+               const std::span<const value_type>& data) -> int
     {
       this->set(data, rows, cols);
       return 0;
@@ -157,7 +157,7 @@ class MatrixCSR
   {
     return [&](const std::span<const std::int32_t>& rows,
                const std::span<const std::int32_t>& cols,
-               const std::span<const Scalar>& data) -> int
+               const std::span<const value_type>& data) -> int
     {
       this->add(data, rows, cols);
       return 0;
@@ -337,7 +337,7 @@ class MatrixCSR
   /// @brief Set all non-zero local entries to a value including entries
   /// in ghost rows.
   /// @param[in] x The value to set non-zero matrix entries to
-  void set(Scalar x) { std::fill(_data.begin(), _data.end(), x); }
+  void set(value_type x) { std::fill(_data.begin(), _data.end(), x); }
 
   /// @brief Set values in the matrix.
   /// @note Only entries included in the sparsity pattern used to
@@ -351,7 +351,7 @@ class MatrixCSR
   /// set in the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void set(const std::span<const Scalar>& x,
+  void set(const std::span<const value_type>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -371,7 +371,7 @@ class MatrixCSR
   /// add to the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void add(const std::span<const Scalar>& x,
+  void add(const std::span<const value_type>& x,
            const std::span<const std::int32_t>& rows,
            const std::span<const std::int32_t>& cols)
   {
@@ -391,12 +391,12 @@ class MatrixCSR
   /// manually by using num_owned_rows() if required.
   /// @return Dense copy of the part of the matrix on the calling rank.
   /// Storage is row-major.
-  std::vector<Scalar> to_dense() const
+  std::vector<value_type> to_dense() const
   {
     const std::size_t nrows = num_all_rows();
     const std::size_t ncols
         = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
-    std::vector<Scalar> A(nrows * ncols);
+    std::vector<value_type> A(nrows * ncols);
     for (std::size_t r = 0; r < nrows; ++r)
       for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
         A[r * ncols + _cols[j]] = _data[j];
@@ -427,7 +427,7 @@ class MatrixCSR
 
     // For each ghost row, pack and send values to send to neighborhood
     std::vector<int> insert_pos = _val_send_disp;
-    std::vector<Scalar> ghost_value_data(_val_send_disp.back());
+    std::vector<value_type> ghost_value_data(_val_send_disp.back());
     for (int i = 0; i < num_ghosts0; ++i)
     {
       const int rank = _ghost_row_to_rank[i];
@@ -455,9 +455,9 @@ class MatrixCSR
 
     int status = MPI_Ineighbor_alltoallv(
         ghost_value_data.data(), val_send_count.data(), _val_send_disp.data(),
-        dolfinx::MPI::mpi_type<Scalar>(), _ghost_value_data_in.data(),
+        dolfinx::MPI::mpi_type<value_type>(), _ghost_value_data_in.data(),
         val_recv_count.data(), _val_recv_disp.data(),
-        dolfinx::MPI::mpi_type<Scalar>(), _comm.comm(), &_request);
+        dolfinx::MPI::mpi_type<value_type>(), _comm.comm(), &_request);
     assert(status == MPI_SUCCESS);
   }
 
@@ -489,7 +489,7 @@ class MatrixCSR
 
     const double norm_sq_local = std::accumulate(
         _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
-        double(0), [](auto norm, Scalar y) { return norm + std::norm(y); });
+        double(0), [](auto norm, value_type y) { return norm + std::norm(y); });
     double norm_sq;
     MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM,
                   _comm.comm());
@@ -511,11 +511,11 @@ class MatrixCSR
 
   /// Get local data values
   /// @note Includes ghost values
-  Container& values() { return _data; }
+  container_type& values() { return _data; }
 
   /// Get local values (const version)
   /// @note Includes ghost values
-  const Container& values() const { return _data; }
+  const container_type& values() const { return _data; }
 
   /// Get local row pointers
   /// @note Includes pointers to ghost rows
@@ -545,11 +545,11 @@ class MatrixCSR
   std::array<int, 2> _bs;
 
   // Matrix data
-  Container _data;
-  Idx _cols, _row_ptr;
+  container_type _data;
+  IdxContainer _cols, _row_ptr;
 
   // Start of off-diagonal (unowned columns) on each row
-  Idx _off_diagonal_offset;
+  IdxContainer _off_diagonal_offset;
 
   // Neighborhood communicator (ghost->owner communicator for rows)
   dolfinx::MPI::Comm _comm;
@@ -571,7 +571,7 @@ class MatrixCSR
   std::vector<int> _ghost_row_to_rank;
 
   // Temporary store for finalize data during non-blocking communication
-  Container _ghost_value_data_in;
+  container_type _ghost_value_data_in;
 };
 
 } // namespace dolfinx::la
diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
index 632581acf91..52ec0104ed5 100644
--- a/cpp/dolfinx/la/Vector.h
+++ b/cpp/dolfinx/la/Vector.h
@@ -73,14 +73,14 @@ class Vector
 
   /// Set all entries (including ghosts)
   /// @param[in] v The value to set all entries to (on calling rank)
-  void set(Scalar v) { std::fill(_x.begin(), _x.end(), v); }
+  void set(value_type v) { std::fill(_x.begin(), _x.end(), v); }
 
   /// Begin scatter of local data from owner to ghosts on other ranks
   /// @note Collective MPI operation
   void scatter_fwd_begin()
   {
     const std::int32_t local_size = _bs * _map->size_local();
-    std::span<const Scalar> x_local(_x.data(), local_size);
+    std::span<const value_type> x_local(_x.data(), local_size);
 
     auto pack = [](const auto& in, const auto& idx, auto& out)
     {
@@ -89,8 +89,8 @@ class Vector
     };
     pack(x_local, _scatterer->local_indices(), _buffer_local);
 
-    _scatterer->scatter_fwd_begin(std::span<const Scalar>(_buffer_local),
-                                  std::span<Scalar>(_buffer_remote),
+    _scatterer->scatter_fwd_begin(std::span<const value_type>(_buffer_local),
+                                  std::span<value_type>(_buffer_remote),
                                   std::span<MPI_Request>(_request));
   }
 
@@ -100,7 +100,7 @@ class Vector
   {
     const std::int32_t local_size = _bs * _map->size_local();
     const std::int32_t num_ghosts = _bs * _map->num_ghosts();
-    std::span<Scalar> x_remote(_x.data() + local_size, num_ghosts);
+    std::span<value_type> x_remote(_x.data() + local_size, num_ghosts);
     _scatterer->scatter_fwd_end(std::span<MPI_Request>(_request));
 
     auto unpack = [](const auto& in, const auto& idx, auto& out, auto op)
@@ -127,7 +127,7 @@ class Vector
   {
     const std::int32_t local_size = _bs * _map->size_local();
     const std::int32_t num_ghosts = _bs * _map->num_ghosts();
-    std::span<Scalar> x_remote(_x.data() + local_size, num_ghosts);
+    std::span<value_type> x_remote(_x.data() + local_size, num_ghosts);
 
     auto pack = [](const auto& in, const auto& idx, auto& out)
     {
@@ -136,8 +136,9 @@ class Vector
     };
     pack(x_remote, _scatterer->remote_indices(), _buffer_remote);
 
-    _scatterer->scatter_rev_begin(std::span<const Scalar>(_buffer_remote),
-                                  std::span<Scalar>(_buffer_local), _request);
+    _scatterer->scatter_rev_begin(std::span<const value_type>(_buffer_remote),
+                                  std::span<value_type>(_buffer_local),
+                                  _request);
   }
 
   /// End scatter of ghost data to owner. This process may receive data
@@ -150,7 +151,7 @@ class Vector
   void scatter_rev_end(BinaryOperation op)
   {
     const std::int32_t local_size = _bs * _map->size_local();
-    std::span<Scalar> x_local(_x.data(), local_size);
+    std::span<value_type> x_local(_x.data(), local_size);
     _scatterer->scatter_rev_end(_request);
 
     auto unpack = [](const auto& in, const auto& idx, auto& out, auto op)
@@ -180,10 +181,13 @@ class Vector
   constexpr int bs() const { return _bs; }
 
   /// Get local part of the vector (const version)
-  std::span<const Scalar> array() const { return std::span<const Scalar>(_x); }
+  std::span<const value_type> array() const
+  {
+    return std::span<const value_type>(_x);
+  }
 
   /// Get local part of the vector
-  std::span<Scalar> mutable_array() { return std::span(_x); }
+  std::span<value_type> mutable_array() { return std::span(_x); }
 
 private:
   // Map describing the data layout
@@ -199,10 +203,10 @@ class Vector
   std::vector<MPI_Request> _request = {MPI_REQUEST_NULL};
 
   // Buffers for ghost scatters
-  Container _buffer_local, _buffer_remote;
+  container_type _buffer_local, _buffer_remote;
 
   // Vector data
-  Container _x;
+  container_type _x;
 };
 
 /// Compute the inner product of two vectors. The two vectors must have

From 61b8d903ae4557c9f9890d6974ec2e96946e05ef Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 11:56:40 +0100
Subject: [PATCH 19/21] Move implementation

---
 cpp/dolfinx/la/MatrixCSR.h | 520 +++++++++++++++++++------------------
 1 file changed, 267 insertions(+), 253 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index f61d7043392..795660dc4f5 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -144,9 +144,9 @@ class MatrixCSR
   /// @todo clarify setting on non-owned entries
   auto mat_set_values()
   {
-    return [&](const std::span<const std::int32_t>& rows,
-               const std::span<const std::int32_t>& cols,
-               const std::span<const Scalar>& data) -> int
+    return [&](std::span<const std::int32_t> rows,
+               std::span<const std::int32_t> cols,
+               std::span<const value_type> data) -> int
     {
       this->set(data, rows, cols);
       return 0;
@@ -158,9 +158,9 @@ class MatrixCSR
   /// @return Function for inserting values into `A`
   auto mat_add_values()
   {
-    return [&](const std::span<const std::int32_t>& rows,
-               const std::span<const std::int32_t>& cols,
-               const std::span<const Scalar>& data) -> int
+    return [&](std::span<const std::int32_t> rows,
+               std::span<const std::int32_t> cols,
+               std::span<const value_type> data) -> int
     {
       this->add(data, rows, cols);
       return 0;
@@ -170,168 +170,7 @@ class MatrixCSR
   /// @brief Create a distributed matrix.
   /// @param[in] p The sparsity pattern the describes the parallel
   /// distribution and the non-zero structure.
-  MatrixCSR(const SparsityPattern& p)
-      : _index_maps({p.index_map(0),
-                     std::make_shared<common::IndexMap>(p.column_index_map())}),
-        _bs({p.block_size(0), p.block_size(1)}), _data(p.num_nonzeros(), 0),
-        _cols(p.graph().array().begin(), p.graph().array().end()),
-        _row_ptr(p.graph().offsets().begin(), p.graph().offsets().end()),
-        _comm(MPI_COMM_NULL)
-  {
-    // TODO: handle block sizes
-    if (_bs[0] > 1 or _bs[1] > 1)
-      throw std::runtime_error("Block size not yet supported");
-
-    // Compute off-diagonal offset for each row
-    std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offset();
-    _off_diagonal_offset.reserve(num_diag_nnz.size());
-    std::transform(num_diag_nnz.begin(), num_diag_nnz.end(), _row_ptr.begin(),
-                   std::back_inserter(_off_diagonal_offset), std::plus{});
-
-    // Some short-hand
-    const std::array local_size
-        = {_index_maps[0]->size_local(), _index_maps[1]->size_local()};
-    const std::array local_range
-        = {_index_maps[0]->local_range(), _index_maps[1]->local_range()};
-    const std::vector<std::int64_t>& ghosts1 = _index_maps[1]->ghosts();
-
-    const std::vector<std::int64_t>& ghosts0 = _index_maps[0]->ghosts();
-    const std::vector<int>& src_ranks = _index_maps[0]->src();
-    const std::vector<int>& dest_ranks = _index_maps[0]->dest();
-
-    // Create neighbourhood communicator (owner <- ghost)
-    MPI_Comm comm;
-    MPI_Dist_graph_create_adjacent(_index_maps[0]->comm(), dest_ranks.size(),
-                                   dest_ranks.data(), MPI_UNWEIGHTED,
-                                   src_ranks.size(), src_ranks.data(),
-                                   MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm);
-    _comm = dolfinx::MPI::Comm(comm, false);
-
-    // Build map from ghost row index position to owning (neighborhood)
-    // rank
-    _ghost_row_to_rank.reserve(_index_maps[0]->owners().size());
-    for (int r : _index_maps[0]->owners())
-    {
-      auto it = std::lower_bound(src_ranks.begin(), src_ranks.end(), r);
-      assert(it != src_ranks.end() and *it == r);
-      int pos = std::distance(src_ranks.begin(), it);
-      _ghost_row_to_rank.push_back(pos);
-    }
-
-    // Compute size of data to send to each neighbor
-    std::vector<std::int32_t> data_per_proc(src_ranks.size(), 0);
-    for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
-    {
-      assert(_ghost_row_to_rank[i] < data_per_proc.size());
-      std::size_t pos = local_size[0] + i;
-      data_per_proc[_ghost_row_to_rank[i]] += _row_ptr[pos + 1] - _row_ptr[pos];
-    }
-
-    // Compute send displacements
-    _val_send_disp.resize(src_ranks.size() + 1, 0);
-    std::partial_sum(data_per_proc.begin(), data_per_proc.end(),
-                     std::next(_val_send_disp.begin()));
-
-    // For each ghost row, pack and send indices to neighborhood
-    std::vector<std::int64_t> ghost_index_data(2 * _val_send_disp.back());
-    {
-      std::vector<int> insert_pos = _val_send_disp;
-      for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
-      {
-        const int rank = _ghost_row_to_rank[i];
-        int row_id = local_size[0] + i;
-        for (int j = _row_ptr[row_id]; j < _row_ptr[row_id + 1]; ++j)
-        {
-          // Get position in send buffer
-          const std::int32_t idx_pos = 2 * insert_pos[rank];
-
-          // Pack send data (row, col) as global indices
-          ghost_index_data[idx_pos] = ghosts0[i];
-          if (std::int32_t col_local = _cols[j]; col_local < local_size[1])
-            ghost_index_data[idx_pos + 1] = col_local + local_range[1][0];
-          else
-            ghost_index_data[idx_pos + 1] = ghosts1[col_local - local_size[1]];
-
-          insert_pos[rank] += 1;
-        }
-      }
-    }
-
-    // Communicate data with neighborhood
-    std::vector<std::int64_t> ghost_index_array;
-    std::vector<int> recv_disp;
-    {
-      std::vector<int> send_sizes;
-      std::transform(data_per_proc.begin(), data_per_proc.end(),
-                     std::back_inserter(send_sizes),
-                     [](auto x) { return 2 * x; });
-
-      std::vector<int> recv_sizes(dest_ranks.size());
-      send_sizes.reserve(1);
-      recv_sizes.reserve(1);
-      MPI_Neighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
-                            MPI_INT, _comm.comm());
-
-      // Build send/recv displacement
-      std::vector<int> send_disp = {0};
-      std::partial_sum(send_sizes.begin(), send_sizes.end(),
-                       std::back_inserter(send_disp));
-      recv_disp = {0};
-      std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
-                       std::back_inserter(recv_disp));
-
-      ghost_index_array.resize(recv_disp.back());
-      MPI_Neighbor_alltoallv(ghost_index_data.data(), send_sizes.data(),
-                             send_disp.data(), MPI_INT64_T,
-                             ghost_index_array.data(), recv_sizes.data(),
-                             recv_disp.data(), MPI_INT64_T, _comm.comm());
-    }
-
-    // Store receive displacements for future use, when transferring
-    // data values
-    _val_recv_disp.resize(recv_disp.size());
-    std::transform(recv_disp.begin(), recv_disp.end(), _val_recv_disp.begin(),
-                   [](int d) { return d / 2; });
-
-    // Global-to-local map for ghost columns
-    std::vector<std::pair<std::int64_t, std::int32_t>> global_to_local;
-    global_to_local.reserve(ghosts1.size());
-    std::int32_t local_i = local_size[1];
-    for (std::int64_t idx : ghosts1)
-      global_to_local.push_back({idx, global_to_local.size() + local_size[1]});
-    std::sort(global_to_local.begin(), global_to_local.end());
-
-    // Compute location in which data for each index should be stored
-    // when received
-    for (std::size_t i = 0; i < ghost_index_array.size(); i += 2)
-    {
-      // Row must be on this process
-      const std::int32_t local_row = ghost_index_array[i] - local_range[0][0];
-      assert(local_row >= 0 and local_row < local_size[0]);
-
-      // Column may be owned or unowned
-      std::int32_t local_col = ghost_index_array[i + 1] - local_range[1][0];
-      if (local_col < 0 or local_col >= local_size[1])
-      {
-        auto it = std::lower_bound(
-            global_to_local.begin(), global_to_local.end(),
-            std::pair(ghost_index_array[i + 1], -1),
-            [](auto& a, auto& b) { return a.first < b.first; });
-        assert(it != global_to_local.end()
-               and it->first == ghost_index_array[i + 1]);
-        local_col = it->second;
-      }
-      auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]);
-      auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]);
-
-      // Find position of column index and insert data
-      auto cit = std::lower_bound(cit0, cit1, local_col);
-      assert(cit != cit1);
-      assert(*cit == local_col);
-      std::size_t d = std::distance(_cols.begin(), cit);
-      _unpack_pos.push_back(d);
-    }
-  }
+  MatrixCSR(const SparsityPattern& p);
 
   /// Move constructor
   /// @todo Check handling of MPI_Request
@@ -354,9 +193,8 @@ class MatrixCSR
   /// set in the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void set(const std::span<const Scalar>& x,
-           const std::span<const std::int32_t>& rows,
-           const std::span<const std::int32_t>& cols)
+  void set(std::span<const value_type> x, std::span<const std::int32_t> rows,
+           std::span<const std::int32_t> cols)
   {
     impl::set_csr(_data, _cols, _row_ptr, x, rows, cols,
                   _index_maps[0]->size_local());
@@ -374,9 +212,8 @@ class MatrixCSR
   /// add to the matrix
   /// @param[in] rows The row indices of `x`
   /// @param[in] cols The column indices of `x`
-  void add(const std::span<const Scalar>& x,
-           const std::span<const std::int32_t>& rows,
-           const std::span<const std::int32_t>& cols)
+  void add(std::span<const value_type> x, std::span<const std::int32_t> rows,
+           std::span<const std::int32_t> cols)
   {
     impl::add_csr(_data, _cols, _row_ptr, x, rows, cols);
   }
@@ -394,18 +231,7 @@ class MatrixCSR
   /// manually by using num_owned_rows() if required.
   /// @return Dense copy of the part of the matrix on the calling rank.
   /// Storage is row-major.
-  std::vector<Scalar> to_dense() const
-  {
-    const std::size_t nrows = num_all_rows();
-    const std::size_t ncols
-        = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
-    std::vector<Scalar> A(nrows * ncols);
-    for (std::size_t r = 0; r < nrows; ++r)
-      for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
-        A[r * ncols + _cols[j]] = _data[j];
-
-    return A;
-  }
+  std::vector<value_type> to_dense() const;
 
   /// @brief Transfer ghost row data to the owning ranks accumulating
   /// received values on the owned rows, and zeroing any existing data
@@ -423,81 +249,17 @@ class MatrixCSR
   /// must not be changed.
   /// @note This function does not change the matrix data. Data update only
   /// occurs with `finalize_end()`.
-  void finalize_begin()
-  {
-    const std::int32_t local_size0 = _index_maps[0]->size_local();
-    const std::int32_t num_ghosts0 = _index_maps[0]->num_ghosts();
-
-    // For each ghost row, pack and send values to send to neighborhood
-    std::vector<int> insert_pos = _val_send_disp;
-    std::vector<Scalar> ghost_value_data(_val_send_disp.back());
-    for (int i = 0; i < num_ghosts0; ++i)
-    {
-      const int rank = _ghost_row_to_rank[i];
-
-      // Get position in send buffer to place data to send to this
-      // neighbour
-      const std::int32_t val_pos = insert_pos[rank];
-      std::copy(std::next(_data.data(), _row_ptr[local_size0 + i]),
-                std::next(_data.data(), _row_ptr[local_size0 + i + 1]),
-                std::next(ghost_value_data.begin(), val_pos));
-      insert_pos[rank]
-          += _row_ptr[local_size0 + i + 1] - _row_ptr[local_size0 + i];
-    }
-
-    _ghost_value_data_in.resize(_val_recv_disp.back());
-
-    // Compute data sizes for send and receive from displacements
-    std::vector<int> val_send_count(_val_send_disp.size() - 1);
-    std::adjacent_difference(std::next(_val_send_disp.begin()),
-                             _val_send_disp.end(), val_send_count.begin());
-
-    std::vector<int> val_recv_count(_val_recv_disp.size() - 1);
-    std::adjacent_difference(std::next(_val_recv_disp.begin()),
-                             _val_recv_disp.end(), val_recv_count.begin());
-
-    int status = MPI_Ineighbor_alltoallv(
-        ghost_value_data.data(), val_send_count.data(), _val_send_disp.data(),
-        dolfinx::MPI::mpi_type<Scalar>(), _ghost_value_data_in.data(),
-        val_recv_count.data(), _val_recv_disp.data(),
-        dolfinx::MPI::mpi_type<Scalar>(), _comm.comm(), &_request);
-    assert(status == MPI_SUCCESS);
-  }
+  void finalize_begin();
 
   /// @brief End transfer of ghost row data to owning ranks.
   /// @note Must be preceded by MatrixCSR::finalize_begin()
   /// @note Matrix data received from other processes will be
   /// accumulated into locally owned rows, and ghost rows will be
   /// zeroed.
-  void finalize_end()
-  {
-    int status = MPI_Wait(&_request, MPI_STATUS_IGNORE);
-    assert(status == MPI_SUCCESS);
-
-    // Add to local rows
-    assert(_ghost_value_data_in.size() == _unpack_pos.size());
-    for (std::size_t i = 0; i < _ghost_value_data_in.size(); ++i)
-      _data[_unpack_pos[i]] += _ghost_value_data_in[i];
-
-    // Set ghost row data to zero
-    const std::int32_t local_size0 = _index_maps[0]->size_local();
-    std::fill(std::next(_data.begin(), _row_ptr[local_size0]), _data.end(), 0);
-  }
+  void finalize_end();
 
   /// Compute the Frobenius norm squared
-  double norm_squared() const
-  {
-    const std::size_t num_owned_rows = _index_maps[0]->size_local();
-    assert(num_owned_rows < _row_ptr.size());
-
-    const double norm_sq_local = std::accumulate(
-        _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
-        double(0), [](auto norm, Scalar y) { return norm + std::norm(y); });
-    double norm_sq;
-    MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM,
-                  _comm.comm());
-    return norm_sq;
-  }
+  double norm_squared() const;
 
   /// @brief Index maps for the row and column space.
   ///
@@ -577,4 +339,256 @@ class MatrixCSR
   Container _ghost_value_data_in;
 };
 
+//-----------------------------------------------------------------------------
+template <class U, class V, class W>
+MatrixCSR<U, V, W>::MatrixCSR(const SparsityPattern& p)
+    : _index_maps({p.index_map(0),
+                   std::make_shared<common::IndexMap>(p.column_index_map())}),
+      _bs({p.block_size(0), p.block_size(1)}), _data(p.num_nonzeros(), 0),
+      _cols(p.graph().array().begin(), p.graph().array().end()),
+      _row_ptr(p.graph().offsets().begin(), p.graph().offsets().end()),
+      _comm(MPI_COMM_NULL)
+{
+  // TODO: handle block sizes
+  if (_bs[0] > 1 or _bs[1] > 1)
+    throw std::runtime_error("Block size not yet supported");
+
+  // Compute off-diagonal offset for each row
+  std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offset();
+  _off_diagonal_offset.reserve(num_diag_nnz.size());
+  std::transform(num_diag_nnz.begin(), num_diag_nnz.end(), _row_ptr.begin(),
+                 std::back_inserter(_off_diagonal_offset), std::plus{});
+
+  // Some short-hand
+  const std::array local_size
+      = {_index_maps[0]->size_local(), _index_maps[1]->size_local()};
+  const std::array local_range
+      = {_index_maps[0]->local_range(), _index_maps[1]->local_range()};
+  const std::vector<std::int64_t>& ghosts1 = _index_maps[1]->ghosts();
+
+  const std::vector<std::int64_t>& ghosts0 = _index_maps[0]->ghosts();
+  const std::vector<int>& src_ranks = _index_maps[0]->src();
+  const std::vector<int>& dest_ranks = _index_maps[0]->dest();
+
+  // Create neighbourhood communicator (owner <- ghost)
+  MPI_Comm comm;
+  MPI_Dist_graph_create_adjacent(_index_maps[0]->comm(), dest_ranks.size(),
+                                 dest_ranks.data(), MPI_UNWEIGHTED,
+                                 src_ranks.size(), src_ranks.data(),
+                                 MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm);
+  _comm = dolfinx::MPI::Comm(comm, false);
+
+  // Build map from ghost row index position to owning (neighborhood)
+  // rank
+  _ghost_row_to_rank.reserve(_index_maps[0]->owners().size());
+  for (int r : _index_maps[0]->owners())
+  {
+    auto it = std::lower_bound(src_ranks.begin(), src_ranks.end(), r);
+    assert(it != src_ranks.end() and *it == r);
+    int pos = std::distance(src_ranks.begin(), it);
+    _ghost_row_to_rank.push_back(pos);
+  }
+
+  // Compute size of data to send to each neighbor
+  std::vector<std::int32_t> data_per_proc(src_ranks.size(), 0);
+  for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
+  {
+    assert(_ghost_row_to_rank[i] < data_per_proc.size());
+    std::size_t pos = local_size[0] + i;
+    data_per_proc[_ghost_row_to_rank[i]] += _row_ptr[pos + 1] - _row_ptr[pos];
+  }
+
+  // Compute send displacements
+  _val_send_disp.resize(src_ranks.size() + 1, 0);
+  std::partial_sum(data_per_proc.begin(), data_per_proc.end(),
+                   std::next(_val_send_disp.begin()));
+
+  // For each ghost row, pack and send indices to neighborhood
+  std::vector<std::int64_t> ghost_index_data(2 * _val_send_disp.back());
+  {
+    std::vector<int> insert_pos = _val_send_disp;
+    for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
+    {
+      const int rank = _ghost_row_to_rank[i];
+      int row_id = local_size[0] + i;
+      for (int j = _row_ptr[row_id]; j < _row_ptr[row_id + 1]; ++j)
+      {
+        // Get position in send buffer
+        const std::int32_t idx_pos = 2 * insert_pos[rank];
+
+        // Pack send data (row, col) as global indices
+        ghost_index_data[idx_pos] = ghosts0[i];
+        if (std::int32_t col_local = _cols[j]; col_local < local_size[1])
+          ghost_index_data[idx_pos + 1] = col_local + local_range[1][0];
+        else
+          ghost_index_data[idx_pos + 1] = ghosts1[col_local - local_size[1]];
+
+        insert_pos[rank] += 1;
+      }
+    }
+  }
+
+  // Communicate data with neighborhood
+  std::vector<std::int64_t> ghost_index_array;
+  std::vector<int> recv_disp;
+  {
+    std::vector<int> send_sizes;
+    std::transform(data_per_proc.begin(), data_per_proc.end(),
+                   std::back_inserter(send_sizes),
+                   [](auto x) { return 2 * x; });
+
+    std::vector<int> recv_sizes(dest_ranks.size());
+    send_sizes.reserve(1);
+    recv_sizes.reserve(1);
+    MPI_Neighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
+                          MPI_INT, _comm.comm());
+
+    // Build send/recv displacement
+    std::vector<int> send_disp = {0};
+    std::partial_sum(send_sizes.begin(), send_sizes.end(),
+                     std::back_inserter(send_disp));
+    recv_disp = {0};
+    std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
+                     std::back_inserter(recv_disp));
+
+    ghost_index_array.resize(recv_disp.back());
+    MPI_Neighbor_alltoallv(ghost_index_data.data(), send_sizes.data(),
+                           send_disp.data(), MPI_INT64_T,
+                           ghost_index_array.data(), recv_sizes.data(),
+                           recv_disp.data(), MPI_INT64_T, _comm.comm());
+  }
+
+  // Store receive displacements for future use, when transferring
+  // data values
+  _val_recv_disp.resize(recv_disp.size());
+  std::transform(recv_disp.begin(), recv_disp.end(), _val_recv_disp.begin(),
+                 [](int d) { return d / 2; });
+
+  // Global-to-local map for ghost columns
+  std::vector<std::pair<std::int64_t, std::int32_t>> global_to_local;
+  global_to_local.reserve(ghosts1.size());
+  std::int32_t local_i = local_size[1];
+  for (std::int64_t idx : ghosts1)
+    global_to_local.push_back({idx, global_to_local.size() + local_size[1]});
+  std::sort(global_to_local.begin(), global_to_local.end());
+
+  // Compute location in which data for each index should be stored
+  // when received
+  for (std::size_t i = 0; i < ghost_index_array.size(); i += 2)
+  {
+    // Row must be on this process
+    const std::int32_t local_row = ghost_index_array[i] - local_range[0][0];
+    assert(local_row >= 0 and local_row < local_size[0]);
+
+    // Column may be owned or unowned
+    std::int32_t local_col = ghost_index_array[i + 1] - local_range[1][0];
+    if (local_col < 0 or local_col >= local_size[1])
+    {
+      auto it = std::lower_bound(global_to_local.begin(), global_to_local.end(),
+                                 std::pair(ghost_index_array[i + 1], -1),
+                                 [](auto& a, auto& b)
+                                 { return a.first < b.first; });
+      assert(it != global_to_local.end()
+             and it->first == ghost_index_array[i + 1]);
+      local_col = it->second;
+    }
+    auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]);
+    auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]);
+
+    // Find position of column index and insert data
+    auto cit = std::lower_bound(cit0, cit1, local_col);
+    assert(cit != cit1);
+    assert(*cit == local_col);
+    std::size_t d = std::distance(_cols.begin(), cit);
+    _unpack_pos.push_back(d);
+  }
+}
+//-----------------------------------------------------------------------------
+template <class U, class V, class W>
+std::vector<typename MatrixCSR<U, V, W>::value_type>
+MatrixCSR<U, V, W>::to_dense() const
+{
+  const std::size_t nrows = num_all_rows();
+  const std::size_t ncols
+      = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
+  std::vector<value_type> A(nrows * ncols);
+  for (std::size_t r = 0; r < nrows; ++r)
+    for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
+      A[r * ncols + _cols[j]] = _data[j];
+
+  return A;
+}
+//-----------------------------------------------------------------------------
+template <class U, class V, class W>
+void MatrixCSR<U, V, W>::finalize_begin()
+{
+  const std::int32_t local_size0 = _index_maps[0]->size_local();
+  const std::int32_t num_ghosts0 = _index_maps[0]->num_ghosts();
+
+  // For each ghost row, pack and send values to send to neighborhood
+  std::vector<int> insert_pos = _val_send_disp;
+  std::vector<value_type> ghost_value_data(_val_send_disp.back());
+  for (int i = 0; i < num_ghosts0; ++i)
+  {
+    const int rank = _ghost_row_to_rank[i];
+
+    // Get position in send buffer to place data to send to this
+    // neighbour
+    const std::int32_t val_pos = insert_pos[rank];
+    std::copy(std::next(_data.data(), _row_ptr[local_size0 + i]),
+              std::next(_data.data(), _row_ptr[local_size0 + i + 1]),
+              std::next(ghost_value_data.begin(), val_pos));
+    insert_pos[rank]
+        += _row_ptr[local_size0 + i + 1] - _row_ptr[local_size0 + i];
+  }
+
+  _ghost_value_data_in.resize(_val_recv_disp.back());
+
+  // Compute data sizes for send and receive from displacements
+  std::vector<int> val_send_count(_val_send_disp.size() - 1);
+  std::adjacent_difference(std::next(_val_send_disp.begin()),
+                           _val_send_disp.end(), val_send_count.begin());
+
+  std::vector<int> val_recv_count(_val_recv_disp.size() - 1);
+  std::adjacent_difference(std::next(_val_recv_disp.begin()),
+                           _val_recv_disp.end(), val_recv_count.begin());
+
+  int status = MPI_Ineighbor_alltoallv(
+      ghost_value_data.data(), val_send_count.data(), _val_send_disp.data(),
+      dolfinx::MPI::mpi_type<value_type>(), _ghost_value_data_in.data(),
+      val_recv_count.data(), _val_recv_disp.data(),
+      dolfinx::MPI::mpi_type<value_type>(), _comm.comm(), &_request);
+  assert(status == MPI_SUCCESS);
+}
+//-----------------------------------------------------------------------------
+template <class U, class V, class W>
+void MatrixCSR<U, V, W>::finalize_end()
+{
+  int status = MPI_Wait(&_request, MPI_STATUS_IGNORE);
+  assert(status == MPI_SUCCESS);
+
+  // Add to local rows
+  assert(_ghost_value_data_in.size() == _unpack_pos.size());
+  for (std::size_t i = 0; i < _ghost_value_data_in.size(); ++i)
+    _data[_unpack_pos[i]] += _ghost_value_data_in[i];
+
+  // Set ghost row data to zero
+  const std::int32_t local_size0 = _index_maps[0]->size_local();
+  std::fill(std::next(_data.begin(), _row_ptr[local_size0]), _data.end(), 0);
+}
+//-----------------------------------------------------------------------------
+template <class Scalar, class Container, class Idx>
+double MatrixCSR<Scalar, Container, Idx>::norm_squared() const
+{
+  const std::size_t num_owned_rows = _index_maps[0]->size_local();
+  assert(num_owned_rows < _row_ptr.size());
+  double norm_sq_local = std::accumulate(
+      _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
+      double(0), [](auto norm, Scalar y) { return norm + std::norm(y); });
+  double norm_sq;
+  MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM, _comm.comm());
+  return norm_sq;
+}
+//-----------------------------------------------------------------------------
+
 } // namespace dolfinx::la

From 832bbbba9f2bf38dc02640882ab9a44aca2470b1 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 12:45:28 +0100
Subject: [PATCH 20/21] Extend MatrixCSR container types

---
 cpp/dolfinx/la/MatrixCSR.h     | 186 ++++++++++++++++++---------------
 cpp/test/matrix.cpp            |  12 +--
 python/dolfinx/wrappers/la.cpp |  14 ++-
 3 files changed, 113 insertions(+), 99 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 87272f106e2..24723b69a5f 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -18,10 +18,9 @@
 
 namespace dolfinx::la
 {
-
 namespace impl
 {
-/// @brief Set data in a CSR matrix
+/// @brief Set data in a CSR matrix.
 ///
 /// @param[out] data The CSR matrix data
 /// @param[in] cols The CSR column indices
@@ -33,38 +32,9 @@ namespace impl
 /// @param[in] local_size The maximum row index that can be set. Used
 /// when debugging is own to check that rows beyond a permitted range
 /// are not being set.
-template <typename U, typename V, typename W, typename X>
-void set_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
-             const X& xrows, const X& xcols,
-             [[maybe_unused]] typename X::value_type local_size)
-{
-  assert(x.size() == xrows.size() * xcols.size());
-  for (std::size_t r = 0; r < xrows.size(); ++r)
-  {
-    // Row index and current data row
-    auto row = xrows[r];
-    using T = typename W::value_type;
-    const T* xr = x.data() + r * xcols.size();
-
-#ifndef NDEBUG
-    if (row >= local_size)
-      throw std::runtime_error("Local row out of range");
-#endif
-
-    // Columns indices for row
-    auto cit0 = std::next(cols.begin(), row_ptr[row]);
-    auto cit1 = std::next(cols.begin(), row_ptr[row + 1]);
-    for (std::size_t c = 0; c < xcols.size(); ++c)
-    {
-      // Find position of column index
-      auto it = std::lower_bound(cit0, cit1, xcols[c]);
-      assert(it != cit1);
-      std::size_t d = std::distance(cols.begin(), it);
-      assert(d < data.size());
-      data[d] = xr[c];
-    }
-  }
-}
+template <typename U, typename V, typename W, typename X, typename Y>
+void set_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
+             const Y& xrows, const Y& xcols, typename Y::value_type local_size);
 
 /// @brief Add data to a CSR matrix
 ///
@@ -75,37 +45,9 @@ void set_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 /// to the matrix
 /// @param[in] xrows The row indices of `x`
 /// @param[in] xcols The column indices of `x`
-template <typename U, typename V, typename W, typename X>
-void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
-             const X& xrows, const X& xcols)
-{
-  assert(x.size() == xrows.size() * xcols.size());
-  for (std::size_t r = 0; r < xrows.size(); ++r)
-  {
-    // Row index and current data row
-    auto row = xrows[r];
-    using T = typename W::value_type;
-    const T* xr = x.data() + r * xcols.size();
-
-#ifndef NDEBUG
-    if (row >= (int)row_ptr.size())
-      throw std::runtime_error("Local row out of range");
-#endif
-
-    // Columns indices for row
-    auto cit0 = std::next(cols.begin(), row_ptr[row]);
-    auto cit1 = std::next(cols.begin(), row_ptr[row + 1]);
-    for (std::size_t c = 0; c < xcols.size(); ++c)
-    {
-      // Find position of column index
-      auto it = std::lower_bound(cit0, cit1, xcols[c]);
-      assert(it != cit1);
-      std::size_t d = std::distance(cols.begin(), it);
-      assert(d < data.size());
-      data[d] += xr[c];
-    }
-  }
-}
+template <typename U, typename V, typename W, typename X, typename Y>
+void add_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
+             const Y& xrows, const Y& xcols);
 } // namespace impl
 
 /// Distributed sparse matrix
@@ -121,9 +63,11 @@ void add_csr(U&& data, const V& cols, const V& row_ptr, const W& x,
 ///
 /// @tparam Scalar Scalar type
 /// @tparam Container Sequence container type to store matrix entries
-/// @tparam IdxContainer Index container type
+/// @tparam ColContainer Index container types
+/// @tparam RowPtrContainer Index container types
 template <class Scalar, class Container = std::vector<Scalar>,
-          class IdxContainer = std::vector<std::int32_t>>
+          class ColContainer = std::vector<std::int32_t>,
+          class RowPtrContainer = std::vector<std::int64_t>>
 class MatrixCSR
 {
   static_assert(std::is_same_v<typename Container::value_type, Scalar>);
@@ -132,9 +76,15 @@ class MatrixCSR
   /// Scalar type
   using value_type = Scalar;
 
-  /// Container type
+  /// Matrix entries container type
   using container_type = Container;
 
+  /// Column index container type
+  using column_container_type = ColContainer;
+
+  /// Row pointer container type
+  using rowptr_container_type = RowPtrContainer;
+
   static_assert(std::is_same_v<value_type, typename container_type::value_type>,
                 "Scalar type and container value type must be the same.");
 
@@ -284,11 +234,11 @@ class MatrixCSR
 
   /// Get local row pointers
   /// @note Includes pointers to ghost rows
-  const std::vector<std::int32_t>& row_ptr() const { return _row_ptr; }
+  const rowptr_container_type& row_ptr() const { return _row_ptr; }
 
   /// Get local column indices
   /// @note Includes columns in ghost rows
-  const std::vector<std::int32_t>& cols() const { return _cols; }
+  const column_container_type& cols() const { return _cols; }
 
   /// Get the start of off-diagonal (unowned columns) on each row,
   /// allowing the matrix to be split (virtually) into two parts.
@@ -297,7 +247,7 @@ class MatrixCSR
   /// from operations on the unowned parts.
   /// @note Includes ghost rows, which should be truncated manually if
   /// not required.
-  const std::vector<std::int32_t>& off_diag_offset() const
+  const rowptr_container_type& off_diag_offset() const
   {
     return _off_diagonal_offset;
   }
@@ -311,10 +261,11 @@ class MatrixCSR
 
   // Matrix data
   container_type _data;
-  IdxContainer _cols, _row_ptr;
+  column_container_type _cols;
+  rowptr_container_type _row_ptr;
 
   // Start of off-diagonal (unowned columns) on each row
-  IdxContainer _off_diagonal_offset;
+  rowptr_container_type _off_diagonal_offset;
 
   // Neighborhood communicator (ghost->owner communicator for rows)
   dolfinx::MPI::Comm _comm;
@@ -340,8 +291,73 @@ class MatrixCSR
 };
 
 //-----------------------------------------------------------------------------
-template <class U, class V, class W>
-MatrixCSR<U, V, W>::MatrixCSR(const SparsityPattern& p)
+template <typename U, typename V, typename W, typename X, typename Y>
+void impl::set_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
+                   const Y& xrows, const Y& xcols,
+                   [[maybe_unused]] typename Y::value_type local_size)
+{
+  assert(x.size() == xrows.size() * xcols.size());
+  for (std::size_t r = 0; r < xrows.size(); ++r)
+  {
+    // Row index and current data row
+    auto row = xrows[r];
+    using T = typename X::value_type;
+    const T* xr = x.data() + r * xcols.size();
+
+#ifndef NDEBUG
+    if (row >= local_size)
+      throw std::runtime_error("Local row out of range");
+#endif
+
+    // Columns indices for row
+    auto cit0 = std::next(cols.begin(), row_ptr[row]);
+    auto cit1 = std::next(cols.begin(), row_ptr[row + 1]);
+    for (std::size_t c = 0; c < xcols.size(); ++c)
+    {
+      // Find position of column index
+      auto it = std::lower_bound(cit0, cit1, xcols[c]);
+      assert(it != cit1);
+      std::size_t d = std::distance(cols.begin(), it);
+      assert(d < data.size());
+      data[d] = xr[c];
+    }
+  }
+}
+//-----------------------------------------------------------------------------
+template <typename U, typename V, typename W, typename X, typename Y>
+void impl::add_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
+                   const Y& xrows, const Y& xcols)
+{
+  assert(x.size() == xrows.size() * xcols.size());
+  for (std::size_t r = 0; r < xrows.size(); ++r)
+  {
+    // Row index and current data row
+    auto row = xrows[r];
+    using T = typename X::value_type;
+    const T* xr = x.data() + r * xcols.size();
+
+#ifndef NDEBUG
+    if (row >= (int)row_ptr.size())
+      throw std::runtime_error("Local row out of range");
+#endif
+
+    // Columns indices for row
+    auto cit0 = std::next(cols.begin(), row_ptr[row]);
+    auto cit1 = std::next(cols.begin(), row_ptr[row + 1]);
+    for (std::size_t c = 0; c < xcols.size(); ++c)
+    {
+      // Find position of column index
+      auto it = std::lower_bound(cit0, cit1, xcols[c]);
+      assert(it != cit1);
+      std::size_t d = std::distance(cols.begin(), it);
+      assert(d < data.size());
+      data[d] += xr[c];
+    }
+  }
+}
+//-----------------------------------------------------------------------------
+template <class U, class V, class W, class X>
+MatrixCSR<U, V, W, X>::MatrixCSR(const SparsityPattern& p)
     : _index_maps({p.index_map(0),
                    std::make_shared<common::IndexMap>(p.column_index_map())}),
       _bs({p.block_size(0), p.block_size(1)}), _data(p.num_nonzeros(), 0),
@@ -504,9 +520,9 @@ MatrixCSR<U, V, W>::MatrixCSR(const SparsityPattern& p)
   }
 }
 //-----------------------------------------------------------------------------
-template <class U, class V, class W>
-std::vector<typename MatrixCSR<U, V, W>::value_type>
-MatrixCSR<U, V, W>::to_dense() const
+template <typename U, typename V, typename W, typename X>
+std::vector<typename MatrixCSR<U, V, W, X>::value_type>
+MatrixCSR<U, V, W, X>::to_dense() const
 {
   const std::size_t nrows = num_all_rows();
   const std::size_t ncols
@@ -519,8 +535,8 @@ MatrixCSR<U, V, W>::to_dense() const
   return A;
 }
 //-----------------------------------------------------------------------------
-template <class U, class V, class W>
-void MatrixCSR<U, V, W>::finalize_begin()
+template <typename U, typename V, typename W, typename X>
+void MatrixCSR<U, V, W, X>::finalize_begin()
 {
   const std::int32_t local_size0 = _index_maps[0]->size_local();
   const std::int32_t num_ghosts0 = _index_maps[0]->num_ghosts();
@@ -561,8 +577,8 @@ void MatrixCSR<U, V, W>::finalize_begin()
   assert(status == MPI_SUCCESS);
 }
 //-----------------------------------------------------------------------------
-template <class U, class V, class W>
-void MatrixCSR<U, V, W>::finalize_end()
+template <typename U, typename V, typename W, typename X>
+void MatrixCSR<U, V, W, X>::finalize_end()
 {
   int status = MPI_Wait(&_request, MPI_STATUS_IGNORE);
   assert(status == MPI_SUCCESS);
@@ -577,14 +593,14 @@ void MatrixCSR<U, V, W>::finalize_end()
   std::fill(std::next(_data.begin(), _row_ptr[local_size0]), _data.end(), 0);
 }
 //-----------------------------------------------------------------------------
-template <class Scalar, class Container, class Idx>
-double MatrixCSR<Scalar, Container, Idx>::norm_squared() const
+template <typename U, typename V, typename W, typename X>
+double MatrixCSR<U, V, W, X>::norm_squared() const
 {
   const std::size_t num_owned_rows = _index_maps[0]->size_local();
   assert(num_owned_rows < _row_ptr.size());
   double norm_sq_local = std::accumulate(
       _data.cbegin(), std::next(_data.cbegin(), _row_ptr[num_owned_rows]),
-      double(0), [](auto norm, Scalar y) { return norm + std::norm(y); });
+      double(0), [](auto norm, value_type y) { return norm + std::norm(y); });
   double norm_sq;
   MPI_Allreduce(&norm_sq_local, &norm_sq, 1, MPI_DOUBLE, MPI_SUM, _comm.comm());
   return norm_sq;
diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp
index e69cacf811f..97a884bb8b4 100644
--- a/cpp/test/matrix.cpp
+++ b/cpp/test/matrix.cpp
@@ -29,8 +29,8 @@ namespace
 /// @param[in, out] x Output vector
 template <typename T>
 void spmv_impl(std::span<const T> values,
-               std::span<const std::int32_t> row_begin,
-               std::span<const std::int32_t> row_end,
+               std::span<const std::int64_t> row_begin,
+               std::span<const std::int64_t> row_end,
                std::span<const std::int32_t> indices, std::span<const T> x,
                std::span<T> y)
 {
@@ -77,17 +77,17 @@ void spmv(la::MatrixCSR<T>& A, la::Vector<T>& x, la::Vector<T>& y)
   x.scatter_fwd_begin();
 
   const std::int32_t nrowslocal = A.num_owned_rows();
-  std::span<const std::int32_t> row_ptr(A.row_ptr().data(), nrowslocal + 1);
+  std::span<const std::int64_t> row_ptr(A.row_ptr().data(), nrowslocal + 1);
   std::span<const std::int32_t> cols(A.cols().data(), row_ptr[nrowslocal]);
-  std::span<const std::int32_t> off_diag_offset(A.off_diag_offset().data(),
+  std::span<const std::int64_t> off_diag_offset(A.off_diag_offset().data(),
                                                 nrowslocal);
   std::span<const T> values(A.values().data(), row_ptr[nrowslocal]);
 
   std::span<const T> _x = x.array();
   std::span<T> _y = y.mutable_array();
 
-  std::span<const std::int32_t> row_begin(row_ptr.data(), nrowslocal);
-  std::span<const std::int32_t> row_end(row_ptr.data() + 1, nrowslocal);
+  std::span<const std::int64_t> row_begin(row_ptr.data(), nrowslocal);
+  std::span<const std::int64_t> row_end(row_ptr.data() + 1, nrowslocal);
 
   // First stage:  spmv - diagonal
   // yi[0] += Ai[0] * xi[0]
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
index 1d9e51b52b7..a44ba365895 100644
--- a/python/dolfinx/wrappers/la.cpp
+++ b/python/dolfinx/wrappers/la.cpp
@@ -120,18 +120,16 @@ void declare_objects(py::module& m, const std::string& type)
       .def_property_readonly("indices",
                              [](dolfinx::la::MatrixCSR<T>& self)
                              {
-                               std::span<const std::int32_t> array
-                                   = self.cols();
-                               return py::array_t<const std::int32_t>(
-                                   array.size(), array.data(), py::cast(self));
+                               auto& array = self.cols();
+                               return py::array_t(array.size(), array.data(),
+                                                  py::cast(self));
                              })
       .def_property_readonly("indptr",
                              [](dolfinx::la::MatrixCSR<T>& self)
                              {
-                               std::span<const std::int32_t> array
-                                   = self.row_ptr();
-                               return py::array_t<const std::int32_t>(
-                                   array.size(), array.data(), py::cast(self));
+                               auto& array = self.row_ptr();
+                               return py::array_t(array.size(), array.data(),
+                                                  py::cast(self));
                              })
       .def("finalize_begin", &dolfinx::la::MatrixCSR<T>::finalize_begin)
       .def("finalize_end", &dolfinx::la::MatrixCSR<T>::finalize_end);

From 55fb1bc8f28634e5ad1fd1d4d87ff3feaf98c963 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Wed, 26 Apr 2023 12:48:21 +0100
Subject: [PATCH 21/21] Doc update

---
 cpp/dolfinx/la/MatrixCSR.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
index 24723b69a5f..d88268837a0 100644
--- a/cpp/dolfinx/la/MatrixCSR.h
+++ b/cpp/dolfinx/la/MatrixCSR.h
@@ -61,10 +61,10 @@ void add_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
 /// code.
 /// @todo Handle block sizes
 ///
-/// @tparam Scalar Scalar type
+/// @tparam Scalar Scalar type of matrix entries
 /// @tparam Container Sequence container type to store matrix entries
-/// @tparam ColContainer Index container types
-/// @tparam RowPtrContainer Index container types
+/// @tparam ColContainer Column index container type
+/// @tparam RowPtrContainer Row pointer container type
 template <class Scalar, class Container = std::vector<Scalar>,
           class ColContainer = std::vector<std::int32_t>,
           class RowPtrContainer = std::vector<std::int64_t>>