diff --git a/cpp/dolfinx/la/SparsityPattern.cpp b/cpp/dolfinx/la/SparsityPattern.cpp index 1dab9d35d26..ac7ee1b6efc 100644 --- a/cpp/dolfinx/la/SparsityPattern.cpp +++ b/cpp/dolfinx/la/SparsityPattern.cpp @@ -140,6 +140,27 @@ SparsityPattern::SparsityPattern( } } //----------------------------------------------------------------------------- +void SparsityPattern::insert(std::int32_t row, std::int32_t col) +{ + if (!_offsets.empty()) + { + throw std::runtime_error( + "Cannot insert into sparsity pattern. It has already been finalized"); + } + + assert(_index_maps[0]); + const std::int32_t max_row + = _index_maps[0]->size_local() + _index_maps[0]->num_ghosts() - 1; + + if (row > max_row or row < 0) + { + throw std::runtime_error( + "Cannot insert rows that do not exist in the IndexMap."); + } + + _row_cache[row].push_back(col); +} +//----------------------------------------------------------------------------- void SparsityPattern::insert(std::span rows, std::span cols) { @@ -160,7 +181,6 @@ void SparsityPattern::insert(std::span rows, throw std::runtime_error( "Cannot insert rows that do not exist in the IndexMap."); } - _row_cache[row].insert(_row_cache[row].end(), cols.begin(), cols.end()); } } diff --git a/cpp/dolfinx/la/SparsityPattern.h b/cpp/dolfinx/la/SparsityPattern.h index 33632d72813..726455a0c39 100644 --- a/cpp/dolfinx/la/SparsityPattern.h +++ b/cpp/dolfinx/la/SparsityPattern.h @@ -67,6 +67,19 @@ class SparsityPattern /// @brief Insert non-zero locations using local (process-wise) /// indices. + /// @param[in] row local row index + /// @param[in] col local column index + void insert(std::int32_t row, std::int32_t col); + + /// @brief Insert non-zero locations using local (process-wise) + /// indices. + /// + /// This routine inserts non-zero locations at the outer product of rows and + /// cols into the sparsity pattern, i.e. adds the matrix entries at + /// A[row[i], col[j]] for all i, j. + /// + /// @param[in] rows list of the local row indices + /// @param[in] cols list of the local column indices void insert(std::span rows, std::span cols); diff --git a/cpp/test/matrix.cpp b/cpp/test/matrix.cpp index 3c7ed9dba8d..448fdce779d 100644 --- a/cpp/test/matrix.cpp +++ b/cpp/test/matrix.cpp @@ -201,9 +201,9 @@ void test_matrix() { auto map0 = std::make_shared(MPI_COMM_SELF, 8); la::SparsityPattern p(MPI_COMM_SELF, {map0, map0}, {1, 1}); - p.insert(std::vector{0}, std::vector{0}); - p.insert(std::vector{4}, std::vector{5}); - p.insert(std::vector{5}, std::vector{4}); + p.insert(0, 0); + p.insert(4, 5); + p.insert(5, 4); p.finalize(); using T = float; diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp index e84590e31d1..5a1f7865f9a 100644 --- a/python/dolfinx/wrappers/la.cpp +++ b/python/dolfinx/wrappers/la.cpp @@ -8,6 +8,7 @@ #include "caster_mpi.h" #include "numpy_dtype.h" #include +#include #include #include #include @@ -265,6 +266,10 @@ void la(nb::module_& m) std::span(cols.data(), cols.size())); }, nb::arg("rows"), nb::arg("cols")) + .def("insert", + nb::overload_cast( + &dolfinx::la::SparsityPattern::insert), + nb::arg("row"), nb::arg("col")) .def( "insert_diagonal", [](dolfinx::la::SparsityPattern& self, diff --git a/python/test/unit/la/test_matrix_csr.py b/python/test/unit/la/test_matrix_csr.py index 738b2b46b86..1f8e5c35d09 100644 --- a/python/test/unit/la/test_matrix_csr.py +++ b/python/test/unit/la/test_matrix_csr.py @@ -25,9 +25,9 @@ def create_test_sparsity(n, bs): if bs == 1: for i in range(2): for j in range(2): - sp.insert(np.array([2 + i]), np.array([4 + j])) + sp.insert(2 + i, 4 + j) elif bs == 2: - sp.insert(np.array([1]), np.array([2])) + sp.insert(1, 2) sp.finalize() return sp @@ -126,10 +126,10 @@ def test_distributed_csr(dtype): sp = SparsityPattern(MPI.COMM_WORLD, [im, im], [1, 1]) for i in range(n): for j in range(n + nghost): - sp.insert(np.array([i]), np.array([j])) + sp.insert(i, j) for i in range(n, n + nghost): for j in range(n, n + nghost): - sp.insert(np.array([i]), np.array([j])) + sp.insert(i, j) sp.finalize() mat = matrix_csr(sp, dtype=dtype)