Skip to content

Commit

Permalink
[doc] Update contents in linear_solver.md (#7967)
Browse files Browse the repository at this point in the history
WIP. Don't merge.

Issue: #7837 
Issue: #7921 

### Brief Summary
Update contents in `linear_solver.md` to reflect the latest changes and
development progress on the linear solver.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
houkensjtu and pre-commit-ci[bot] authored May 10, 2023
1 parent 4d645f5 commit 346f03c
Showing 1 changed file with 126 additions and 5 deletions.
131 changes: 126 additions & 5 deletions docs/lang/articles/math/linear_solver.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,19 @@ sidebar_position: 3

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`
2. Solvers built for `ti.field`
1. Solvers built for `SparseMatrix`, including:
- Direct solver `SparseSolver`
- Iterative (conjugate-gradient method) solver `SparseCG`
2. 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
You may want to solve some linear equations using sparse matrices.
Then, the following steps could help:
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. Create a `solver` using `ti.linalg.SparseSolver(solver_type, ordering)`. 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. Analyze and factorize the sparse matrix you want to solve using `solver.analyze_pattern(sparse_matrix)` and `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.
Expand Down Expand Up @@ -66,8 +73,122 @@ print(f">>>> Computation succeed: {success}")
# [0.5 0. 0. 0.5]
# >>>> Computation was successful?: True
```
## Examples

Please have a look at our two demos for more information:
+ [Stable fluid](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/stable_fluid.py): A 2D fluid simulation using a sparse Laplacian matrix to solve Poisson's pressure equation.
+ [Implicit mass spring](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/implicit_mass_spring.py): 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:
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:

```python
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:
```python
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:
```python cont
# >>>> 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:
```python cont
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:

```python
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())
```

0 comments on commit 346f03c

Please sign in to comment.