Merge pull request #14194 from KratosMultiphysics/core/csr_utilities

[FastPR][Core] Enabling getting indices from connectivities array

Author: rubenzorrilla

kratos/python/add_sparse_matrices_to_python.cpp

@@ -44,14 +44,62 @@ namespace
 
     using EquationIdVectorType = std::vector<std::size_t>;
 
+    /**
+     * @brief Get the Equation Id Csr Indices object
+     * @details This function computes the CSR matrix value vector indices corresponding to the equation ids of the entities in the provided container.
+     * @tparam TContainerType Type of the container of the entities.
+     * @param rCsrMatrix The CSR matrix to get the indices from.
+     * @param rConnectivities The connectivities of the entities.
+     * @return NDData The NDData object containing the indices.
+     */
+    NDData<int> GetEquationIdCsrIndices(
+        const CsrMatrixType& rCsrMatrix,
+        const NDData<int>& rConnectivities)
+    {
+        // Get shapes from the input connectivities
+        const auto& r_shape = rConnectivities.Shape();
+        KRATOS_ERROR_IF(r_shape.size() != 2) << "Input connectivities must have shape (n_entities, local_size)" << std::endl;
+        const std::size_t n_entities = r_shape[0];
+        const std::size_t local_size = r_shape[1];
+
+        // Assign the input NDData to have shape: number of entities * local_size * local_size
+        DenseVector<unsigned int> nd_data_shape(3);
+        nd_data_shape[0] = n_entities;
+        nd_data_shape[1] = local_size;
+        nd_data_shape[2] = local_size;
+        auto eq_ids_data = NDData<int>(nd_data_shape);
+
+        // Loop over the connectivities
+        auto eq_ids_data_view = eq_ids_data.ViewData();
+        auto connectivities_view = rConnectivities.ViewData();
+        IndexPartition<std::size_t>(n_entities).for_each([&](std::size_t i) {
+            // Get current entity position in the connectivities view
+            const std::size_t idx_start = i * local_size;
+            const std::size_t eq_ids_pos = i * (local_size * local_size);
+
+            // Loop over the DOFs
+            for (unsigned int i_local = 0; i_local < local_size; ++i_local) {
+                const unsigned int i_global = connectivities_view[idx_start + i_local]; // Row global equation id
+                for(unsigned int j_local = 0; j_local < local_size; ++j_local) {
+                    const unsigned int j_global = connectivities_view[idx_start + j_local]; // Column global equation id
+                    const unsigned int csr_index = rCsrMatrix.FindValueIndex(i_global,j_global); // Index in the CSR matrix values vector
+                    eq_ids_data_view[eq_ids_pos + i_local * local_size + j_local] = csr_index;
+                }
+            }
+        });
+
+        // Return the data container
+        return eq_ids_data;
+    }
+
     /**
      * @brief Get the Equation Id Csr Indices object
      * @details This function computes the CSR matrix value vector indices corresponding to the equation ids of the entities in the provided container.
      * @tparam TContainerType Type of the container of the entities.
      * @param rCsrMatrix The CSR matrix to get the indices from.
      * @param rContainer The container of the entities.
      * @param rProcessInfo The process info.
-     * @param pNDData Pointer to the NDData object to store the indices.
+     * @return NDData The NDData object containing the indices.
      */
     template<class TContainerType>
     NDData<int> GetEquationIdCsrIndices(
@@ -73,21 +121,21 @@ namespace
         auto eq_ids_data = NDData<int>(nd_data_shape);
 
         // Loop over the container
+        EquationIdVectorType aux_tls;
         auto eq_ids_data_view = eq_ids_data.ViewData();
-        IndexPartition<std::size_t>(n_entities).for_each([&](std::size_t i) {
+        IndexPartition<std::size_t>(n_entities).for_each(aux_tls, [&](std::size_t i, EquationIdVectorType& rTLS) {
             // Get current entity
             auto it = rContainer.begin() + i;
             const std::size_t it_pos = i * (local_size * local_size);
 
             // Get current entity equation ids
-            EquationIdVectorType equation_ids;
-            it->EquationIdVector(equation_ids, rProcessInfo);
+            it->EquationIdVector(rTLS, rProcessInfo);
 
             // Loop over the DOFs
             for (unsigned int i_local = 0; i_local < local_size; ++i_local) {
-                const unsigned int i_global = equation_ids[i_local]; // Row global equation id
+                const unsigned int i_global = rTLS[i_local]; // Row global equation id
                 for(unsigned int j_local = 0; j_local < local_size; ++j_local) {
-                    const unsigned int j_global = equation_ids[j_local]; // Column global equation id
+                    const unsigned int j_global = rTLS[j_local]; // Column global equation id
                     const unsigned int csr_index = rCsrMatrix.FindValueIndex(i_global,j_global); // Index in the CSR matrix values vector
                     eq_ids_data_view[it_pos + i_local * local_size + j_local] = csr_index;
                 }
@@ -132,7 +180,7 @@ namespace
                     const std::size_t aux_idx = entity_pos + i_local * local_size_1 + j_local; // Position in the contributions and equation ids data
                     const int csr_index = r_idx_data[aux_idx]; // Index in the CSR matrix values vector
                     const double lhs_contribution = r_lhs_contribution_data[aux_idx]; // Scalar contribution to the left hand side
-                    r_lhs_data[csr_index] += lhs_contribution;
+                    AtomicAdd(r_lhs_data[csr_index], lhs_contribution);
                 }
             }
         });
@@ -306,6 +354,9 @@ void AddSparseMatricesToPython(pybind11::module& m)
     .def("__matmul__", [](CsrMatrix<double,IndexType>& rA,CsrMatrix<double,IndexType>& rB){
         return AmgclCSRSpMMUtilities::SparseMultiply(rA,rB);
     }, py::is_operator())
+    .def("GetEquationIdCsrIndices", [](CsrMatrix<double,IndexType>& rA, const NDData<int>& rConnectivities) {
+        return GetEquationIdCsrIndices(rA, rConnectivities);
+    }, py::return_value_policy::move)
     .def("GetEquationIdCsrIndices", [](CsrMatrix<double,IndexType>& rA, const ElementsContainerType& rElements, const ProcessInfo& rProcessInfo) {
         return GetEquationIdCsrIndices(rA, rElements, rProcessInfo);
     }, py::return_value_policy::move)

kratos/tests/test_sparse_matrices.py

@@ -136,7 +136,7 @@ def test_matrix_assembly(self):
         for i in range(y.Size()):
             self.assertEqual(c[i], 5.0)
 
-    def test_get_equation_id_csr_indices(self):
+    def test_get_equation_id_csr_indices_1(self):
         # Create model part
         model = KratosMultiphysics.Model()
         mp = model.CreateModelPart("Main")
@@ -187,6 +187,35 @@ def test_get_equation_id_csr_indices(self):
         for i, j in zip(data, expected_data):
             self.assertEqual(i, j)
 
+    def test_get_equation_id_csr_indices_2(self):
+        # Create elements connectivities
+        connectivities = KratosMultiphysics.IntNDData(np.array([[0,1,2],[1,3,2]]))
+
+        # Build sparse matrix graph
+        connectivities_data = connectivities.to_numpy()
+        graph = KratosMultiphysics.SparseContiguousRowGraph(4)
+        for i in range(connectivities_data.shape[0]):
+            graph.AddEntries(connectivities_data[i], connectivities_data[i])
+        graph.Finalize()
+
+        # Create Matrix
+        A = KratosMultiphysics.CsrMatrix(graph)
+
+        # Get the elemental contributions CSR indices
+        elem_csr_indices = A.GetEquationIdCsrIndices(connectivities)
+
+        # Check results
+        shape = elem_csr_indices.to_numpy().shape
+        self.assertEqual(shape[0], 2)
+        self.assertEqual(shape[1], 3)
+        self.assertEqual(shape[2], 3)
+
+        data = elem_csr_indices.to_numpy().flatten()
+        expected_data = [0, 1, 2, 3, 4, 5, 7, 8, 9,
+                         4, 6, 5, 11, 13, 12, 8, 10, 9]
+        for i, j in zip(data, expected_data):
+            self.assertEqual(i, j)
+
     def test_assemble_with_csr_indices(self):
         # Create model part
         model = KratosMultiphysics.Model()