[Example] A stable fluid demo with sparse matrix

examples/simulation/stable_fluid.py

@@ -4,10 +4,22 @@
 # https://www.bilibili.com/video/BV1ZK411H7Hc?p=4
 # https://github.com/ShaneFX/GAMES201/tree/master/HW01
 
+import argparse
+
 import numpy as np
 
 import taichi as ti
 
+# How to run:
+#   `python stable_fluid.py`: use the jacobi iteration to solve the linear system.
+#   `python stable_fluid.py -s`: use a sparse matrix to do so.
+parser = argparse.ArgumentParser()
+parser.add_argument('-s',
+                    '--use-sp-mat',
+                    action='store_true',
+                    help='Solve Poisson\'s equation by using a sparse matrix')
+args = parser.parse_args()
+
 res = 512
 dt = 0.03
 p_jacobi_iters = 500  # 40 for a quicker but less accurate result
@@ -20,8 +32,16 @@
 gravity = True
 debug = False
 paused = False
+use_sparse_matrix = False
+
+use_sparse_matrix = args.use_sp_mat
 
-ti.init(arch=ti.gpu)
+if use_sparse_matrix:
+    ti.init(arch=ti.x64)
+    print('Using sparse matrix')
+else:
+    ti.init(arch=ti.gpu)
+    print('Using jacobi iteration')
 
 _velocities = ti.Vector.field(2, float, shape=(res, res))
 _new_velocities = ti.Vector.field(2, float, shape=(res, res))
@@ -46,6 +66,37 @@ def swap(self):
 pressures_pair = TexPair(_pressures, _new_pressures)
 dyes_pair = TexPair(_dye_buffer, _new_dye_buffer)
 
+if use_sparse_matrix:
+    # use a sparse matrix to solve Poisson's pressure equation.
+    @ti.kernel
+    def fill_laplacian_matrix(A: ti.sparse_matrix_builder()):
+        for i, j in ti.ndrange(res, res):
+            row = i * res + j
+            center = 0.0
+            if j != 0:
+                A[row, row - 1] += -1.0
+                center += 1.0
+            if j != res - 1:
+                A[row, row + 1] += -1.0
+                center += 1.0
+            if i != 0:
+                A[row, row - res] += -1.0
+                center += 1.0
+            if i != res - 1:
+                A[row, row + res] += -1.0
+                center += 1.0
+            A[row, row] += center
+
+    N = res * res
+    K = ti.SparseMatrixBuilder(N, N, max_num_triplets=N * 6)
+    b = ti.field(ti.f32, shape=N)
+
+    fill_laplacian_matrix(K)
+    L = K.build()
+    solver = ti.SparseSolver(solver_type="LLT")
+    solver.analyze_pattern(L)
+    solver.factorize(L)
+
 
 @ti.func
 def sample(qf, u, v):
@@ -187,6 +238,30 @@ def enhance_vorticity(vf: ti.template(), cf: ti.template()):
         vf[i, j] = min(max(vf[i, j] + force * dt, -1e3), 1e3)
 
 
+@ti.kernel
+def copy_divergence(div_in: ti.template(), div_out: ti.template()):
+    for I in ti.grouped(div_in):
+        div_out[I[0] * res + I[1]] = -div_in[I]
+
+
+@ti.kernel
+def apply_pressure(p_in: ti.ext_arr(), p_out: ti.template()):
+    for I in ti.grouped(p_out):
+        p_out[I] = p_in[I[0] * res + I[1]]
+
+
+def solve_pressure_sp_mat():
+    copy_divergence(velocity_divs, b)
+    x = solver.solve(b)
+    apply_pressure(x, pressures_pair.cur)
+
+
+def solve_pressure_jacobi():
+    for _ in range(p_jacobi_iters):
+        pressure_jacobi(pressures_pair.cur, pressures_pair.nxt)
+        pressures_pair.swap()
+
+
 def step(mouse_data):
     advect(velocities_pair.cur, velocities_pair.cur, velocities_pair.nxt)
     advect(velocities_pair.cur, dyes_pair.cur, dyes_pair.nxt)
@@ -201,9 +276,10 @@ def step(mouse_data):
         vorticity(velocities_pair.cur)
         enhance_vorticity(velocities_pair.cur, velocity_curls)
 
-    for _ in range(p_jacobi_iters):
-        pressure_jacobi(pressures_pair.cur, pressures_pair.nxt)
-        pressures_pair.swap()
+    if use_sparse_matrix:
+        solve_pressure_sp_mat()
+    else:
+        solve_pressure_jacobi()
 
     subtract_gradient(velocities_pair.cur, pressures_pair.cur)
 

taichi/math/linalg.h

@@ -537,8 +537,9 @@ template <typename T, int dim, InstSetExt ISE = default_instruction_set>
 using TVector = VectorND<dim, T, ISE>;
 
 template <int dim, typename T, InstSetExt ISE>
-TI_FORCE_INLINE VectorND<dim, T, ISE> operator
-    *(T a, const VectorND<dim, T, ISE> &v) {
+TI_FORCE_INLINE VectorND<dim, T, ISE> operator*(
+    T a,
+    const VectorND<dim, T, ISE> &v) {
   return VectorND<dim, T, ISE>(a) * v;
 }
 
@@ -891,8 +892,9 @@ struct MatrixND {
 };
 
 template <int dim, typename T, InstSetExt ISE>
-TI_FORCE_INLINE MatrixND<dim, T, ISE> operator
-    *(const T a, const MatrixND<dim, T, ISE> &M) {
+TI_FORCE_INLINE MatrixND<dim, T, ISE> operator*(
+    const T a,
+    const MatrixND<dim, T, ISE> &M) {
   MatrixND<dim, T, ISE> ret;
   for (int i = 0; i < dim; i++) {
     ret[i] = a * M[i];

taichi/runtime/llvm/atomic.h

@@ -31,10 +31,10 @@ DEFINE_ATOMIC_OP_INTRINSIC(or, i32)
 DEFINE_ATOMIC_OP_INTRINSIC(or, i64)
 DEFINE_ATOMIC_OP_INTRINSIC(or, u32)
 DEFINE_ATOMIC_OP_INTRINSIC(or, u64)
-DEFINE_ATOMIC_OP_INTRINSIC (xor, i32)
-DEFINE_ATOMIC_OP_INTRINSIC (xor, i64)
-DEFINE_ATOMIC_OP_INTRINSIC (xor, u32)
-DEFINE_ATOMIC_OP_INTRINSIC (xor, u64)
+DEFINE_ATOMIC_OP_INTRINSIC(xor, i32)
+DEFINE_ATOMIC_OP_INTRINSIC(xor, i64)
+DEFINE_ATOMIC_OP_INTRINSIC(xor, u32)
+DEFINE_ATOMIC_OP_INTRINSIC(xor, u64)
 
 inline f32 add_f32(f32 a, f32 b) {
   return a + b;