[Lang] Sparse matrix element modify

taichi-dev · Oct 8, 2021 · 2468b24 · 2468b24
1 parent 0bd85c3
commit 2468b24
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 0 deletions.
diff --git a/python/taichi/lang/sparse_matrix.py b/python/taichi/lang/sparse_matrix.py
@@ -59,6 +59,9 @@ def __matmul__(self, other):
     def __getitem__(self, indices):
         return self.matrix.get_element(indices[0], indices[1])
 
+    def __setitem__(self, indices, value):
+        self.matrix.set_element(indices[0], indices[1], value)
+
     def __str__(self):
         return self.matrix.to_string()
 

diff --git a/taichi/program/sparse_matrix.cpp b/taichi/program/sparse_matrix.cpp
@@ -118,5 +118,9 @@ float32 SparseMatrix::get_element(int row, int col) {
   return matrix_.coeff(row, col);
 }
 
+void SparseMatrix::set_element(int row, int col, float32 value){
+  matrix_.coeffRef(row,col) = value;
+}
+
 }  // namespace lang
 }  // namespace taichi
diff --git a/taichi/program/sparse_matrix.h b/taichi/program/sparse_matrix.h
@@ -41,6 +41,7 @@ class SparseMatrix {
   Eigen::SparseMatrix<float32> &get_matrix();
   const Eigen::SparseMatrix<float32> &get_matrix() const;
   float32 get_element(int row, int col);
+  void set_element(int row, int col, float32 value);
 
   friend SparseMatrix operator+(const SparseMatrix &sm1,
                                 const SparseMatrix &sm2);

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -987,6 +987,7 @@ void export_lang(py::module &m) {
       .def("transpose", &SparseMatrix::transpose,
            py::return_value_policy::reference_internal)
       .def("get_element", &SparseMatrix::get_element)
+      .def("set_element", &SparseMatrix::set_element)
       .def("num_rows", &SparseMatrix::num_rows)
       .def("num_cols", &SparseMatrix::num_cols);
 

diff --git a/tests/python/test_sparse_matrix.py b/tests/python/test_sparse_matrix.py
@@ -33,6 +33,20 @@ def fill(Abuilder: ti.sparse_matrix_builder()):
     for i in range(n):
         assert A[i, i] == i
 
+@ti.test(arch=ti.cpu)
+def test_sparse_matrix_element_modify():
+    n = 8
+    Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100)
+
+    @ti.kernel
+    def fill(Abuilder: ti.sparse_matrix_builder()):
+        for i in range(n):
+            Abuilder[i, i] += i
+
+    fill(Abuilder)
+    A = Abuilder.build()
+    A[0,0] = 1024.0
+    assert A[0,0] == 1024.0
 
 @ti.test(arch=ti.cpu)
 def test_sparse_matrix_addition():