def test_sparse_matrix_elementwise_multiplication(): n = 8 Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) Bbuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder(), Bbuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += i + j Bbuilder[i, j] += i - j fill(Abuilder, Bbuilder) A = Abuilder.build() B = Bbuilder.build() C = A * B for i in range(n): for j in range(n): assert C[i, j] == (i + j) * (i - j)
def test_sparse_matrix_subtraction(): n = 8 Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) Bbuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder(), Bbuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += i + j Bbuilder[i, j] += i - j fill(Abuilder, Bbuilder) A = Abuilder.build() B = Bbuilder.build() C = A - B for i in range(n): for j in range(n): assert C[i, j] == 2 * j
def test_sparse_matrix_multiplication(): n = 2 Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) Bbuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder(), Bbuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += i + j Bbuilder[i, j] += i - j fill(Abuilder, Bbuilder) A = Abuilder.build() B = Bbuilder.build() C = A @ B assert C[0, 0] == 1.0 assert C[0, 1] == 0.0 assert C[1, 0] == 2.0 assert C[1, 1] == -1.0
def test_sparse_matrix_nonsymmetric_multiplication(): n, k, m = 2, 3, 4 Abuilder = ti.SparseMatrixBuilder(n, k, max_num_triplets=100) Bbuilder = ti.SparseMatrixBuilder(k, m, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder(), Bbuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, k): Abuilder[i, j] += i + j for i, j in ti.ndrange(k, m): Bbuilder[i, j] -= i + j fill(Abuilder, Bbuilder) A = Abuilder.build() B = Bbuilder.build() C = A @ B GT = [[-5, -8, -11, -14], [-8, -14, -20, -26]] for i in range(n): for j in range(m): assert C[i, j] == GT[i][j]
def test_sparse_matrix_element_access(): 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() for i in range(n): assert A[i, i] == i
def test_sparse_matrix_builder(): n = 8 Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += i + j fill(Abuilder) A = Abuilder.build() for i in range(n): for j in range(n): assert A[i, j] == i + j
def test_sparse_matrix_scalar_multiplication(): n = 8 Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) @ti.kernel def fill(Abuilder: ti.sparse_matrix_builder()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += i + j fill(Abuilder) A = Abuilder.build() B = A * 3.0 for i in range(n): for j in range(n): assert B[i, j] == 3 * (i + j)
import taichi as ti ti.init(arch=ti.x64) n = 8 K = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) f = ti.SparseMatrixBuilder(n, 1, max_num_triplets=100) @ti.kernel def fill(A: ti.sparse_matrix_builder(), b: ti.sparse_matrix_builder(), interval: ti.i32): for i in range(n): if i > 0: A[i - 1, i] += -1.0 A[i, i] += 1 if i < n - 1: A[i + 1, i] += -1.0 A[i, i] += 1.0 if i % interval == 0: b[i, 0] += 1.0 fill(K, f, 3) print(">>>> K.print_triplets()") K.print_triplets() A = K.build()