diff --git a/docs/lang/articles/math/linear_solver.md b/docs/lang/articles/math/linear_solver.md index cd852e4d48a93..bc98328cfe868 100644 --- a/docs/lang/articles/math/linear_solver.md +++ b/docs/lang/articles/math/linear_solver.md @@ -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. @@ -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()) +```