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.
Sparse linear solver
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:
- Create a
solver
usingti.linalg.SparseSolver(solver_type, ordering)
. 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. - Analyze and factorize the sparse matrix you want to solve using
solver.analyze_pattern(sparse_matrix)
andsolver.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. - Call
solver.info()
to check if the solving process succeeds.
Here's a full example.
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
Please have a look at our two demos for more information:
- Stable fluid: A 2D fluid simulation using a sparse Laplacian matrix to solve Poisson's pressure equation.
- Implicit mass spring: A 2D cloth simulation demo using sparse matrices to solve the linear systems.
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())