跳转至主要内容
Version: v1.1.3

稀疏矩阵

在科学和工程问题中求解线性系统时,经常会用到稀疏矩阵。 Taichi为用户使用稀疏矩阵提供有用的 API。

在 Taichi 编程中使用稀疏矩阵,需要遵循以下三步:

  1. 使用 ti.linalg.SparseMatrixBuilder()创建一个builder .
  2. 使用矩阵数据填充 builder.
  3. builder 创建稀疏矩阵.
WARNING

The sparse matrix is still under implementation. 以下是一些限制:

  • 仅支持 CPU 后端
  • 稀疏矩阵的数据类型是32位浮点
  • The storage format is column-major

Here's an example:

import taichi as ti
ti.init(arch=ti.x64) # only CPU backend is supported for now

n = 4
# step 1: create sparse matrix builder
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)

@ti.kernel
def fill(A: ti.types.sparse_matrix_builder()):
for i in range(n):
A[i, i] += 1 # Only += and -= operators are supported for now.

# step 2: 填充 builder
fill(K)

print(">>>> K.print_triplets()")
K.print_triplets()
# outputs:
# >>>> K.print_triplets()
# n=4, m=4, num_triplets=4 (max=100)(0, 0) val=1.0(1, 1) val=1.0(2, 2) val=1.0(3, 3) val=1.0

# step 3: 由builder创建稀疏矩阵.
A = K.build()
print(">>>> A = K.build()")
print(A)
# outputs:
# >>>> A = K.build()
# [1, 0, 0, 0]
# [0, 1, 0, 0]
# [0, 0, 1, 0]
# [0, 0, 0, 1]

目前稀疏矩阵支持如+, -, *, @ 这些基本操作和转置操作.

print(">>>> Summation: C = A + A")
C = A + A
print(C)
# outputs:
# >>>> Summation: C = A + A
# [2, 0, 0, 0]
# [0, 2, 0, 0]
# [0, 0, 2, 0]
# [0, 0, 0, 2]

print(">>>> Subtraction: D = A - A")
D = A - A
print(D)
# outputs:
# >>>> Subtraction: D = A - A
# [0, 0, 0, 0]
# [0, 0, 0, 0]
# [0, 0, 0, 0]
# [0, 0, 0, 0]

print(">>>> Multiplication with a scalar on the right: E = A * 3.0")
E = A * 3.0
print(E)
# outputs:
# >>>> Multiplication with a scalar on the right: E = A * 3.0
# [3, 0, 0, 0]
# [0, 3, 0, 0]
# [0, 0, 3, 0]
# [0, 0, 0, 3]

print(">>>> Multiplication with a scalar on the left: E = 3.0 * A")
E = 3.0 * A
print(E)
# outputs:
# >>>> Multiplication with a scalar on the left: E = 3.0 * A
# [3, 0, 0, 0]
# [0, 3, 0, 0]
# [0, 0, 3, 0]
# [0, 0, 0, 3]

print(">>>> Transpose: F = A.transpose()")
F = A.transpose()
print(F)
# outputs:
# >>>> Transpose: F = A.transpose()
# [1, 0, 0, 0]
# [0, 1, 0, 0]
# [0, 0, 1, 0]
# [0, 0, 0, 1]

print(">>>> Matrix multiplication: G = E @ A")
G = E @ A
print(G)
# outputs:
# >>>> Matrix multiplication: G = E @ A
# [3, 0, 0, 0]
# [0, 3, 0, 0]
# [0, 0, 3, 0]
# [0, 0, 0, 3]

print(">>>> Element-wise multiplication: H = E * A")
H = E * A
print(H)
# outputs:
# >>>> Element-wise multiplication: H = E * A
# [3, 0, 0, 0]
# [0, 3, 0, 0]
# [0, 0, 3, 0]
# [0, 0, 0, 3]

print(f">>>> Element Access: A[0,0] = {A[0,0]}")
# outputs:
# >>>> Element Access: A[0,0] = 1.0

稀疏线性求解器

您可能想要使用稀疏矩阵解决一些线性方程。 那么以下步骤可能会有帮助:

  1. 使用 ti.linalg.SparseSolver(solver_type, ordering) 创建一个solver。 目前,稀疏求解器支持 LLT, LDLTLU 因式求解类型, 支持的 ordering 包括 AMD COLAMD
  2. 使用 solver.analyze_pattern(Sparse_matrix)solver.factorize(Sparse_matrix) 分析和因式分解您想要求解的稀疏矩阵。
  3. 调用 solver.solve(b) 求解。 b 是表示线性系统右侧的 NumPy 数组或 Taichi field。
  4. 调用 solver.info() 来检查求解过程是否成功。

下面是一个完整的示例:

import taichi as ti

ti.init(arch=ti.x64)

n = 4

K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
b = ti.field(ti.f32, shape=n)

@ti.kernel
def fill(A: ti.types.sparse_matrix_builder(), b: ti.template(), interval: ti.i32):
for i in range(n):
A[i, i] += 2.0

if i % interval == 0:
b[i] += 1.0

fill(K, b, 3)

A = K.build()
print(">>>> Matrix A:")
print(A)
print(">>>> Vector b:")
print(b)
# outputs:
# >>>> Matrix A:
# [2, 0, 0, 0]
# [0, 2, 0, 0]
# [0, 0, 2, 0]
# [0, 0, 0, 2]
# >>>> Vector b:
# [1. 0. 0. 1.]
solver = ti.linalg.SparseSolver(solver_type="LLT")
solver.analyze_pattern(A)
solver.factorize(A)
x = solver.solve(b)
isSuccess = solver.info()
print(">>>> Solve sparse linear systems Ax = b with the solution x:")
print(x)
print(f">>>> Computation was successful?: {isSuccess}")
# outputs:
# >>>> Solve sparse linear systems Ax = b with the solution x:
# [0.5 0. 0. 0.5]
# >>>> Computation was successful?: True

示例

请查看我们的两个演示示例:

  • Stable fluid: 使用 稀疏 Laplacian 矩阵来求解 Poisson 压力方程的一个二维流体模拟项目。
  • Implicit mass spring: 一个用稀疏矩阵求解线性系统的二维布料仿真项目。
本文有帮助吗?