Version: develop

# 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:

1. Solvers built for `SparseMatrix`, including:
• Direct solver `SparseSolver`
• Iterative (conjugate-gradient method) solver `SparseCG`
1. 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:

1. 使用 `ti.linalg.SparseSolver(solver_type, ordering)` 创建一个`solver`。 Currently, the factorization types supported on CPU backends are `LLT`, `LDLT`, and `LU`, and supported orderings include `AMD` and `COLAMD`. The sparse solver on CUDA supports the `LLT` factorization type only.
2. 使用 `solver.analyze_pattern(Sparse_matrix)``solver.factorize(Sparse_matrix)` 分析和因式分解您想要求解的稀疏矩阵。
3. Call `x = solver.solve(b)`, where `x` is the solution and `b` is the right-hand side of the linear system. On CPU backends, `x` and `b` can be NumPy arrays, Taichi Ndarrays, or Taichi fields. On the CUDA backend, `x` and `b` must be Taichi Ndarrays.
4. 调用 `solver.info()` 来检查求解过程是否成功。

``import taichi as tiarch = ti.cpu # or ti.cudati.init(arch=arch)n = 4K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)b = ti.ndarray(ti.f32, shape=n)@ti.kerneldef 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.0fill(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:

1. Create a `solver` using `ti.linalg.SparseCG(A, b, x0, max_iter, atol)`, where `A` is a `SparseMatrix` that stores the coefficient matrix of the linear system, `b` is the right-hand side of the equations, `x0` is the initial guess and `atol` is the absolute tolerance threshold.
2. Call `x, exit_code = solver.solve()` to obtain the solution `x` along with the `exit_code` that indicates the status of the solution. `exit_code` should be `True` if the solving was successful. Here is an example:
``import taichi as titi.init(arch=ti.cpu)n = 4K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)b = ti.ndarray(ti.f32, shape=n)@ti.kerneldef 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.0fill(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 tifrom taichi.linalg import LinearOperatorti.init(arch=ti.cpu)@ti.kerneldef 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 tiimport mathfrom taichi.linalg import MatrixFreeCG, LinearOperatorti.init(arch=ti.cpu)GRID = 4Ax = 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.kerneldef 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.kerneldef 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 - bA = LinearOperator(compute_Ax)init()MatrixFreeCG(A, b, x, maxiter=10 * GRID * GRID, tol=1e-18, quiet=True)print(x.to_numpy())``