Expose lower level insert_diagonal function (#2684)

* enable setting diagonal of local part from python

* add test

* fix importing issue

* remove checks

* update test

* minor updates

* fix docuimentation

* Skip doc

* Revert change

* Updates

---------

Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>

cpp/dolfinx/la/MatrixCSR.h

@@ -204,25 +204,27 @@ class MatrixCSR
   {
     auto set_fn = [](value_type& y, const value_type& x) { y = x; };
 
+    std::int32_t num_rows
+        = _index_maps[0]->size_local() + _index_maps[0]->num_ghosts();
     assert(x.size() == rows.size() * cols.size() * BS0 * BS1);
     if (_bs[0] == BS0 and _bs[1] == BS1)
     {
       impl::insert_csr<BS0, BS1>(_data, _cols, _row_ptr, x, rows, cols, set_fn,
-                                 _index_maps[0]->size_local());
+                                 num_rows);
     }
     else if (_bs[0] == 1 and _bs[1] == 1)
     {
       // Set blocked data in a regular CSR matrix (_bs[0]=1, _bs[1]=1) with
       // correct sparsity
       impl::insert_blocked_csr<BS0, BS1>(_data, _cols, _row_ptr, x, rows, cols,
-                                         set_fn, _index_maps[0]->size_local());
+                                         set_fn, num_rows);
     }
     else
     {
       assert(BS0 == 1 and BS1 == 1);
       // Set non-blocked data in a blocked CSR matrix (BS0=1, BS1=1)
       impl::insert_nonblocked_csr(_data, _cols, _row_ptr, x, rows, cols, set_fn,
-                                  _index_maps[0]->size_local(), _bs[0], _bs[1]);
+                                  num_rows, _bs[0], _bs[1]);
     }
   }
 

cpp/dolfinx/la/matrix_csr_impl.h

@@ -29,8 +29,8 @@ namespace impl
 /// @param[in] xrows The row indices of `x`
 /// @param[in] xcols The column indices of `x`
 /// @param[in] op The operation (set or add)
-/// @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
+/// @param[in] num_rows The maximum row index that can be set. Used
+/// when debugging to check that rows beyond a permitted range
 /// are not being set.
 ///
 /// @note In the case of block data, where BS0 or BS1 are greater than
@@ -49,7 +49,7 @@ template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
           typename X, typename Y>
 void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
                 const Y& xrows, const Y& xcols, OP op,
-                typename Y::value_type local_size);
+                typename Y::value_type num_rows);
 
 /// @brief Incorporate blocked data with given block sizes into a non-blocked
 /// MatrixCSR
@@ -68,14 +68,14 @@ void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
 /// @param[in] xrows The row indices of `x`
 /// @param[in] xcols The column indices of `x`
 /// @param[in] op The operation (set or add)
-/// @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
+/// @param[in] num_rows The maximum row index that can be set. Used
+/// when debugging to check that rows beyond a permitted range
 /// are not being set.
 template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
           typename X, typename Y>
 void insert_blocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
                         const Y& xrows, const Y& xcols, OP op,
-                        typename Y::value_type local_size);
+                        typename Y::value_type num_rows);
 
 /// @brief Incorporate non-blocked data into a blocked matrix (Data block size =
 /// 1)
@@ -89,16 +89,16 @@ void insert_blocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
 /// @param[in] xrows The row indices of `x`
 /// @param[in] xcols The column indices of `x`
 /// @param[in] op The operation (set or add)
-/// @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
+/// @param[in] num_rows The maximum row index that can be set. Used
+/// when debugging to check that rows beyond a permitted range
 /// are not being set.
 /// @param[in] bs0 Row block size of Matrix
 /// @param[in] bs1 Column block size of Matrix
 template <typename OP, typename U, typename V, typename W, typename X,
           typename Y>
 void insert_nonblocked_csr(U&& data, const V& cols, const W& row_ptr,
                            const X& x, const Y& xrows, const Y& xcols, OP op,
-                           typename Y::value_type local_size, int bs0, int bs1);
+                           typename Y::value_type num_rows, int bs0, int bs1);
 
 } // namespace impl
 
@@ -107,7 +107,7 @@ template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
           typename X, typename Y>
 void impl::insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
                       const Y& xrows, const Y& xcols, OP op,
-                      [[maybe_unused]] typename Y::value_type local_size)
+                      [[maybe_unused]] typename Y::value_type num_rows)
 {
   const std::size_t nc = xcols.size();
   assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1);
@@ -119,7 +119,7 @@ void impl::insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x,
     const T* xr = x.data() + r * nc * BS0 * BS1;
 
 #ifndef NDEBUG
-    if (row >= local_size)
+    if (row >= num_rows)
       throw std::runtime_error("Local row out of range");
 #endif
     // Columns indices for row
@@ -152,8 +152,7 @@ template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
           typename X, typename Y>
 void impl::insert_blocked_csr(U&& data, const V& cols, const W& row_ptr,
                               const X& x, const Y& xrows, const Y& xcols, OP op,
-                              [[maybe_unused]]
-                              typename Y::value_type local_size)
+                              [[maybe_unused]] typename Y::value_type num_rows)
 {
   const std::size_t nc = xcols.size();
   assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1);
@@ -163,7 +162,7 @@ void impl::insert_blocked_csr(U&& data, const V& cols, const W& row_ptr,
     auto row = xrows[r] * BS0;
 
 #ifndef NDEBUG
-    if (row >= local_size)
+    if (row >= num_rows)
       throw std::runtime_error("Local row out of range");
 #endif
 
@@ -197,7 +196,7 @@ void impl::insert_nonblocked_csr(U&& data, const V& cols, const W& row_ptr,
                                  const X& x, const Y& xrows, const Y& xcols,
                                  OP op,
                                  [[maybe_unused]]
-                                 typename Y::value_type local_size,
+                                 typename Y::value_type num_rows,
                                  int bs0, int bs1)
 {
   const std::size_t nc = xcols.size();
@@ -212,7 +211,7 @@ void impl::insert_nonblocked_csr(U&& data, const V& cols, const W& row_ptr,
     const T* xr = x.data() + r * nc;
 
 #ifndef NDEBUG
-    if (rdiv.quot >= local_size)
+    if (rdiv.quot >= num_rows)
       throw std::runtime_error("Local row out of range");
 #endif
     // Columns indices for row

python/dolfinx/fem/__init__.py

@@ -11,7 +11,8 @@
 from dolfinx.cpp.fem import create_sparsity_pattern as _create_sparsity_pattern
 from dolfinx.fem import petsc
 from dolfinx.fem.assemble import (apply_lifting, assemble_matrix,
-                                  assemble_scalar, assemble_vector, set_bc)
+                                  assemble_scalar, assemble_vector,
+                                  create_matrix, set_bc)
 from dolfinx.fem.bcs import (DirichletBC, bcs_by_block, dirichletbc,
                              locate_dofs_geometrical, locate_dofs_topological)
 from dolfinx.fem.dofmap import DofMap
@@ -35,7 +36,7 @@ def create_sparsity_pattern(a: Form):
 
 
 __all__ = [
-    "Constant", "Expression", "Function",
+    "Constant", "Expression", "Function", "create_matrix",
     "FunctionSpace", "TensorFunctionSpace",
     "VectorFunctionSpace", "create_sparsity_pattern",
     "assemble_scalar", "assemble_matrix", "assemble_vector", "apply_lifting", "set_bc",

python/dolfinx/wrappers/assemble.cpp

@@ -278,6 +278,15 @@ void declare_assembly_functions(py::module& m)
       },
       py::arg("A"), py::arg("V"), py::arg("bcs"), py::arg("diagonal"),
       "Experimental.");
+  m.def(
+      "insert_diagonal",
+      [](dolfinx::la::MatrixCSR<T>& A, const py::array_t<std::int32_t>& rows,
+         T diagonal)
+      {
+        dolfinx::fem::set_diagonal(
+            A.mat_set_values(), std::span(rows.data(), rows.size()), diagonal);
+      },
+      py::arg("A"), py::arg("rows"), py::arg("diagonal"), "Experimental.");
   m.def(
       "assemble_matrix",
       [](const std::function<int(const py::array_t<std::int32_t>&,

python/dolfinx/wrappers/common.cpp

@@ -97,6 +97,8 @@ void common(py::module& m)
       .def_property_readonly("local_range",
                              &dolfinx::common::IndexMap::local_range,
                              "Range of indices owned by this map")
+      .def_property_readonly("index_to_dest_ranks",
+                             &dolfinx::common::IndexMap::index_to_dest_ranks)
       .def_property_readonly(
           "ghosts",
           [](const dolfinx::common::IndexMap& self)

python/test/unit/la/test_matrix_csr.py

@@ -7,10 +7,15 @@
 
 import numpy as np
 import pytest
-
+import scipy.sparse as sps
+import ufl
 from dolfinx.common import IndexMap
 from dolfinx.cpp.la import BlockMode, SparsityPattern
 from dolfinx.la import matrix_csr
+from dolfinx.mesh import GhostMode, create_unit_square
+
+from dolfinx import cpp as _cpp
+from dolfinx import fem
 
 from mpi4py import MPI
 
@@ -138,3 +143,76 @@ def test_distributed_csr(dtype):
     pre_final_sum = mat.data.sum()
     mat.finalize()
     assert np.isclose(mat.data.sum(), pre_final_sum)
+
+
+@pytest.mark.parametrize('dtype', [np.float32, np.float64, np.complex64, np.complex128])
+def test_set_diagonal_distributed(dtype):
+    mesh_dtype = np.real(dtype(0)).dtype
+    ghost_mode = GhostMode.shared_facet
+    mesh = create_unit_square(MPI.COMM_WORLD, 5, 5, ghost_mode=ghost_mode, dtype=mesh_dtype)
+    V = fem.FunctionSpace(mesh, ("Lagrange", 1))
+
+    tdim = mesh.topology.dim
+    cellmap = mesh.topology.index_map(tdim)
+    num_cells = cellmap.size_local + cellmap.num_ghosts
+
+    # Integration domain includes ghost cells
+    cells_domains = [(1, np.arange(0, num_cells))]
+    dx = ufl.Measure("dx", subdomain_data=cells_domains, domain=mesh)
+
+    u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
+    a = fem.form(ufl.inner(u, v) * dx(1), dtype=dtype)
+
+    # get index map from function space
+    index_map = V.dofmap.index_map
+    num_dofs = index_map.size_local + index_map.num_ghosts
+
+    # list of dofs including ghost dofs
+    dofs = np.arange(0, num_dofs, dtype=np.int32)
+
+    # create matrix
+    A = fem.create_matrix(a)
+    As = sps.csr_matrix((A.data, A.indices, A.indptr))
+
+    # set diagonal values
+    value = dtype(1.0)
+    _cpp.fem.insert_diagonal(A._cpp_object, dofs, value)
+
+    # check diagonal values: they should be 1.0, including ghost dofs
+    diag = As.diagonal()
+    reference = np.full_like(diag, value, dtype=dtype)
+    assert np.allclose(diag, reference)
+
+    # Finalize matrix: this will remove ghost rows and diagonal values of
+    # ghost rows will be added to diagonal of corresponding process
+    A.finalize()
+
+    diag = As.diagonal()
+    nlocal = index_map.size_local
+    assert (diag[nlocal:] == dtype(0.0)).all()
+
+    shared_dofs = index_map.index_to_dest_ranks
+    for dof in range(nlocal):
+        owners = shared_dofs.links(dof)
+        assert diag[dof] == len(owners) + 1
+
+    # create matrix
+    A = fem.create_matrix(a)
+    As = sps.csr_matrix((A.data, A.indices, A.indptr))
+
+    # set diagonal values using dirichlet bc: this will set diagonal values of
+    # owned rows only
+    bc = fem.dirichletbc(dtype(0.0), dofs, V)
+    _cpp.fem.insert_diagonal(A._cpp_object, a.function_spaces[0], [bc._cpp_object], value)
+
+    # check diagonal values: they should be 1.0, except ghost dofs
+    diag = As.diagonal()
+    reference = np.full_like(diag, value, dtype=dtype)
+    assert np.allclose(diag[:nlocal], reference[:nlocal])
+    assert np.allclose(diag[nlocal:], np.zeros_like(diag[nlocal:]))
+
+    # Finalize matrix:
+    # this will zero ghost rows and diagonal values are already zero.
+    A.finalize()
+    assert (As.diagonal()[nlocal:] == dtype(0.)).all()
+    assert (As.diagonal()[:nlocal] == dtype(1.)).all()