Linear Solver
Solving linear equations is a common task in scientific computing. Taichi provides basic direct and iterative linear solvers for various simulation scenarios. Currently, there are two categories of linear solvers available:
- Solvers built for
SparseMatrix
, including:
- Direct solver
SparseSolver
- Iterative (conjugate-gradient method) solver
SparseCG
- Solvers built for
LinearOperator
- Iterative (matrix-free conjugate-gradient method) solver
MatrixfreeCG
It's important to understand that those solvers are built for specific types of matrices. For example, if you built a coefficient matrix A
as a SparseMatrix
, then you can only use SparseSolver
or SparseCG
to solve the corresponding linear system. Below we will explain the usage of each type of solvers.
稀疏线性求解器
There are two types of linear solvers available for SparseMatrix
, direct solver and iterative solver.
Sparse direct solver
To solve a linear system whose coefficient matrix is a SparseMatrix
using a direct method, follow the steps below:
- 使用
ti.linalg.SparseSolver(solver_type, ordering)
创建一个solver
。 Currently, the factorization types supported on CPU backends areLLT
,LDLT
, andLU
, and supported orderings includeAMD
andCOLAMD
. The sparse solver on CUDA supports theLLT
factorization type only. - 使用
solver.analyze_pattern(Sparse_matrix)
和solver.factorize(Sparse_matrix)
分析和因式分解您想要求解的稀疏矩阵。 - Call
x = solver.solve(b)
, wherex
is the solution andb
is the right-hand side of the linear system. On CPU backends,x
andb
can be NumPy arrays, Taichi Ndarrays, or Taichi fields. On the CUDA backend,x
andb
must be Taichi Ndarrays. - 调用
solver.info()
来检查求解过程是否成功。
下面是一个完整的示例:
import taichi as ti
arch = ti.cpu # or ti.cuda
ti.init(arch=arch)
n = 4
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
b = ti.ndarray(ti.f32, shape=n)
@ti.kernel
def fill(A: ti.types.sparse_matrix_builder(), b: ti.types.ndarray(), 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)
success = solver.info()
print(">>>> Solve sparse linear systems Ax = b with the solution x:")
print(x)
print(f">>>> Computation succeed: {success}")
# 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: 一个用稀疏矩阵求解线性系统的二维布料仿真项目。
Sparse iterative solver
To solve a linear system whose coefficient matrix is a SparseMatrix
using a iterative (conjugate-gradient) method, follow the steps below:
- Create a
solver
usingti.linalg.SparseCG(A, b, x0, max_iter, atol)
, whereA
is aSparseMatrix
that stores the coefficient matrix of the linear system,b
is the right-hand side of the equations,x0
is the initial guess andatol
is the absolute tolerance threshold. - Call
x, exit_code = solver.solve()
to obtain the solutionx
along with theexit_code
that indicates the status of the solution.exit_code
should beTrue
if the solving was successful. Here is an example:
import taichi as ti
ti.init(arch=ti.cpu)
n = 4
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
b = ti.ndarray(ti.f32, shape=n)
@ti.kernel
def fill(A: ti.types.sparse_matrix_builder(), b: ti.types.ndarray(), 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.to_numpy())
# 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.SparseCG(A, b)
x, exit_code = solver.solve()
print(">>>> Solve sparse linear systems Ax = b with the solution x:")
print(x)
print(f">>>> Computation was successful?: {exit_code}")
# outputs:
# >>>> Solve sparse linear systems Ax = b with the solution x:
# [0.5 0. 0. 0.5]
# >>>> Computation was successful?: True
Note that the building process of SparseMatrix
A
is exactly the same as in the case of SparseSolver
, the only difference here is that we created a solver
whose type is SparseCG
instead of SparseSolver
.
Matrix-free iterative solver
Apart from SparseMatrix
as an efficient representation of matrices, Taichi also support the LinearOperator
type, which is a matrix-free representation of matrices. Keep in mind that matrices can be seen as a linear transformation from an input vector to a output vector, it is possible to encapsulate the information of a matrice as a LinearOperator
.
To create a LinearOperator
in Taichi, we first need to define a kernel that represent the linear transformation:
import taichi as ti
from taichi.linalg import LinearOperator
ti.init(arch=ti.cpu)
@ti.kernel
def compute_matrix_vector(v:ti.template(), mv:ti.template()):
for i in v:
mv[i] = 2 * v[i]
In this case, compute_matrix_vector
kernel accepts an input vector v
and calculates the corresponding matrix-vector product mv
. It is mathematically equal to a matrice whose diagonal elements are all 2. In the case of n=4
, the equivalent matrice A
is:
# >>>> Matrix A:
# [2, 0, 0, 0]
# [0, 2, 0, 0]
# [0, 0, 2, 0]
# [0, 0, 0, 2]
Then we can create the LinearOperator
as follows:
A = LinearOperator(compute_matrix_vector)
To solve a system of linear equations represented by this LinearOperator
, we can use the built-in matrix-free solver MatrixFreeCG
. Here is a full example:
import taichi as ti
import math
from taichi.linalg import MatrixFreeCG, LinearOperator
ti.init(arch=ti.cpu)
GRID = 4
Ax = ti.field(dtype=ti.f32, shape=(GRID, GRID))
x = ti.field(dtype=ti.f32, shape=(GRID, GRID))
b = ti.field(dtype=ti.f32, shape=(GRID, GRID))
@ti.kernel
def init():
for i, j in ti.ndrange(GRID, GRID):
xl = i / (GRID - 1)
yl = j / (GRID - 1)
b[i, j] = ti.sin(2 * math.pi * xl) * ti.sin(2 * math.pi * yl)
x[i, j] = 0.0
@ti.kernel
def compute_Ax(v: ti.template(), mv: ti.template()):
for i, j in v:
l = v[i - 1, j] if i - 1 >= 0 else 0.0
r = v[i + 1, j] if i + 1 <= GRID - 1 else 0.0
t = v[i, j + 1] if j + 1 <= GRID - 1 else 0.0
b = v[i, j - 1] if j - 1 >= 0 else 0.0
# Avoid ill-conditioned matrix A
mv[i, j] = 20 * v[i, j] - l - r - t - b
A = LinearOperator(compute_Ax)
init()
MatrixFreeCG(A, b, x, maxiter=10 * GRID * GRID, tol=1e-18, quiet=True)
print(x.to_numpy())