Extend to_dense for parallel use cases (#3354)

* Add parallel support to `to_dense`

Previously only the diagonal (local-to-local) interaction block of the CSR matrix was outputted,
this is extended to the complete locally owned row structure, allowing for paralell debugging.

When run in sequential the behavior of `to_dense` is not altered

Extend  test case for parallel to dense usage

Adapt  export

Adapt test

Apply suggestion

* Switch to array instead of vector

* Make use of new insert function

---------

Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>

cpp/dolfinx/la/MatrixCSR.h

@@ -630,15 +630,17 @@ 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
-      = _index_maps[1]->size_local() + _index_maps[1]->num_ghosts();
+  const std::size_t ncols = _index_maps[1]->size_global();
   std::vector<value_type> A(nrows * ncols * _bs[0] * _bs[1], 0.0);
   for (std::size_t r = 0; r < nrows; ++r)
     for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j)
       for (int i0 = 0; i0 < _bs[0]; ++i0)
         for (int i1 = 0; i1 < _bs[1]; ++i1)
         {
-          A[(r * _bs[1] + i0) * ncols * _bs[0] + _cols[j] * _bs[1] + i1]
+          std::array<std::int32_t, 1> local_col {_cols[j]};
+          std::array<std::int64_t, 1> global_col{0};
+          _index_maps[1]->local_to_global(local_col, global_col);
+          A[(r * _bs[1] + i0) * ncols * _bs[0] + global_col[0] * _bs[1] + i1]
               = _data[j * _bs[0] * _bs[1] + i0 * _bs[1] + i1];
         }
 

cpp/test/matrix.cpp

@@ -16,6 +16,8 @@
 #include <dolfinx/la/MatrixCSR.h>
 #include <dolfinx/la/SparsityPattern.h>
 #include <dolfinx/la/Vector.h>
+#include <mpi.h>
+#include <span>
 
 using namespace dolfinx;
 
@@ -199,8 +201,8 @@ la::MatrixCSR<double> create_operator(MPI_Comm comm)
 
 void test_matrix()
 {
-  auto map0 = std::make_shared<common::IndexMap>(MPI_COMM_SELF, 8);
-  la::SparsityPattern p(MPI_COMM_SELF, {map0, map0}, {1, 1});
+  auto map0 = std::make_shared<common::IndexMap>(MPI_COMM_WORLD, 8);
+  la::SparsityPattern p(MPI_COMM_WORLD, {map0, map0}, {1, 1});
   p.insert(0, 0);
   p.insert(4, 5);
   p.insert(5, 4);
@@ -215,23 +217,36 @@ void test_matrix()
 
   const std::vector Adense0 = A.to_dense();
 
+  // Note: we cut off the ghost rows by intent here! But therefore we are not
+  // able to work with the dimensions of Adense0 to compute indices, these
+  // contain the ghost rows, which also vary between processes.
   MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
-      const T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<std::size_t, 8, 8>>
-      Adense(Adense0.data(), 8, 8);
+      const T,
+      MDSPAN_IMPL_STANDARD_NAMESPACE::extents<
+          std::size_t, 8, MDSPAN_IMPL_STANDARD_NAMESPACE::dynamic_extent>>
+      Adense(Adense0.data(), 8, A.index_map(1)->size_global());
 
-  std::vector<T> Aref_data(8 * 8, 0);
+  std::vector<T> Aref_data(8 * A.index_map(1)->size_global(), 0);
   MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
-      T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<std::size_t, 8, 8>>
-      Aref(Aref_data.data(), 8, 8);
-  Aref(0, 0) = 1;
-  Aref(4, 5) = 2.3;
+      T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<
+             std::size_t, 8, MDSPAN_IMPL_STANDARD_NAMESPACE::dynamic_extent>>
+      Aref(Aref_data.data(), 8, A.index_map(1)->size_global());
+
+  auto to_global_col = [&](auto col)
+  {
+    std::array<std::int64_t, 1> tmp;
+    A.index_map(1)->local_to_global(std::vector<std::int32_t>{col}, tmp);
+    return tmp[0];
+  };
+  Aref(0, to_global_col(0)) = 1;
+  Aref(4, to_global_col(5)) = 2.3;
 
   for (std::size_t i = 0; i < Adense.extent(0); ++i)
     for (std::size_t j = 0; j < Adense.extent(1); ++j)
       CHECK(Adense(i, j) == Aref(i, j));
 
-  Aref(4, 4) = 2.3;
-  CHECK(Adense(4, 4) != Aref(4, 4));
+  Aref(4, to_global_col(4)) = 2.3;
+  CHECK(Adense(4, to_global_col(4)) != Aref(4, to_global_col(4)));
 }
 
 } // namespace

python/dolfinx/wrappers/la.cpp

@@ -136,10 +136,10 @@ void declare_objects(nb::module_& m, const std::string& type)
            {
              const std::array<int, 2> bs = self.block_size();
              std::size_t nrows = self.num_all_rows() * bs[0];
-             auto map_col = self.index_map(1);
-             std::size_t ncols
-                 = (map_col->size_local() + map_col->num_ghosts()) * bs[1];
-             return dolfinx_wrappers::as_nbarray(self.to_dense(),
+             std::size_t ncols = self.index_map(1)->size_global() * bs[1];
+             auto dense = self.to_dense();
+             assert(nrows*ncols == dense.size());
+             return dolfinx_wrappers::as_nbarray(std::move(self.to_dense()),
                                                  {nrows, ncols});
            })
       .def_prop_ro(

python/test/unit/la/test_matrix_vector.py

@@ -40,11 +40,8 @@ def test_create_matrix_csr():
     A = la.matrix_csr(pattern, dtype=np.complex128)
     assert A.data.dtype == np.complex128
 
-    cmap = pattern.column_index_map()
-    num_cols = cmap.size_local + cmap.num_ghosts
-    num_rows = bs * (map.size_local + map.num_ghosts)
-    zero = np.zeros((num_rows, bs * num_cols), dtype=np.complex128)
-    assert np.allclose(A.to_dense(), zero)
+    dense = A.to_dense()
+    assert np.allclose(dense, np.zeros(dense.shape, dtype=np.complex128))
 
 
 @pytest.mark.parametrize(