diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml
index c214932d62cc0..9c78bb7c32b56 100644
--- a/.github/workflows/presubmit.yml
+++ b/.github/workflows/presubmit.yml
@@ -85,7 +85,7 @@ jobs:
           export PATH=$TAICHI_REPO_DIR/bin:$PATH
           export PATH=$TAICHI_REPO_DIR/taichi-llvm/bin/:$PATH
           export PYTHONPATH=$TAICHI_REPO_DIR/python
-          python examples/laplace.py
+          python examples/algorithm/laplace.py
           ti diagnose
           ./build/taichi_cpp_tests
           ti test -vr2 -t2
@@ -139,7 +139,7 @@ jobs:
           export PATH=$TAICHI_REPO_DIR/bin:$PATH
           export PATH=$TAICHI_REPO_DIR/taichi-llvm/bin/:$PATH
           export PYTHONPATH=$TAICHI_REPO_DIR/python
-          python examples/laplace.py
+          python examples/algorithm/laplace.py
           ti diagnose
           [ "$RUN_CPP_TESTS" = "ON" ] && ./build/taichi_cpp_tests
           ti test -vr2 -t2
@@ -173,7 +173,7 @@ jobs:
           export PYTHONPATH=$TAICHI_REPO_DIR/python
           export DISPLAY=:1
           glewinfo
-          $PYTHON examples/laplace.py
+          $PYTHON examples/algorithm/laplace.py
           ti diagnose
           ti test -vr2 -t2
 
@@ -230,7 +230,7 @@ jobs:
           $env:PATH += ";C:\taichi_clang\bin"
           $env:PATH += ";$env:TAICHI_REPO_DIR\bin"
           python -c "import taichi"
-          python examples/laplace.py
+          python examples/algorithm/laplace.py
           python bin/taichi diagnose
           python bin/taichi test -Cvr2 -t2
         env:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 4bc810e3e2ff8..235f0d9f9121a 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -51,7 +51,7 @@ jobs:
           export PATH=$TAICHI_REPO_DIR/bin:$PATH
           export PATH=$TAICHI_REPO_DIR/taichi-llvm/bin/:$PATH
           export PYTHONPATH=$TAICHI_REPO_DIR/python
-          python examples/laplace.py
+          python examples/algorithm/laplace.py
           ti diagnose
           ti test -vr2 -t2
 
diff --git a/examples/revert_image_color.py b/examples/algorithm/invert_image_color.py
similarity index 100%
rename from examples/revert_image_color.py
rename to examples/algorithm/invert_image_color.py
diff --git a/examples/laplace.py b/examples/algorithm/laplace.py
similarity index 52%
rename from examples/laplace.py
rename to examples/algorithm/laplace.py
index 07f75e893b80f..fbbb2c16c9c49 100644
--- a/examples/laplace.py
+++ b/examples/algorithm/laplace.py
@@ -1,18 +1,20 @@
 import taichi as ti
 
 ti.init(arch=ti.cpu)
+
+N = 16
+
 x, y = ti.field(ti.f32), ti.field(ti.f32)
 
-ti.root.dense(ti.ij, 16).place(x, y)
+ti.root.dense(ti.ij, N).place(x, y)
 
 
 @ti.kernel
 def laplace():
     for i, j in x:
-        if (i + j) % 3 == 0:
-            y[i, j] = 4.0 * x[i, j] - x[i - 1,
-                                        j] - x[i + 1,
-                                               j] - x[i, j - 1] - x[i, j + 1]
+        if (i + j) % 3 == 0 and 1 <= i < N - 1 and 1 <= j < N - 1:
+            y[i, j] = 4.0 * x[i, j] \
+                    - x[i - 1, j] - x[i + 1, j] - x[i, j - 1] - x[i, j + 1]
         else:
             y[i, j] = 0.0
 
diff --git a/examples/algorithm/marching_squares.py b/examples/algorithm/marching_squares.py
new file mode 100644
index 0000000000000..8e1b8c1f9550e
--- /dev/null
+++ b/examples/algorithm/marching_squares.py
@@ -0,0 +1,107 @@
+# Marching squares algorithm
+# https://en.wikipedia.org/wiki/Marching_squares
+
+import numpy as np
+
+import taichi as ti
+
+ti.init(arch=ti.cpu)
+
+N = 128
+
+edge_table_np = np.array(
+    [
+        [[-1, -1], [-1, -1]],  # Case 0
+        [[3, 0], [-1, -1]],  # Case 1
+        [[0, 1], [-1, -1]],  # Case 2
+        [[1, 3], [-1, -1]],  # Case 3
+        [[1, 2], [-1, -1]],  # Case 4
+        [[0, 1], [2, 3]],  # Case 5
+        [[0, 2], [-1, -1]],  # Case 6
+        [[2, 3], [-1, -1]],  # Case 7
+        [[2, 3], [-1, -1]],  # Case 8
+        [[0, 2], [-1, -1]],  # Case 9
+        [[0, 3], [1, 2]],  # Case 10
+        [[1, 2], [-1, -1]],  # Case 11
+        [[1, 3], [-1, -1]],  # Case 12
+        [[0, 1], [-1, -1]],  # Case 13
+        [[3, 0], [-1, -1]],  # Case 14
+        [[-1, -1], [-1, -1]],  # Case 15
+    ],
+    dtype=np.int32)
+
+pixels = ti.field(float, (N, N))
+
+edge_coords = ti.Vector.field(2, float, (N**2, 2))  # generated edges
+edge_table = ti.Vector.field(2, int, (16, 2))  # edge table (constant)
+edge_table.from_numpy(edge_table_np)
+
+
+@ti.func
+def gauss(x, sigma):
+    # Un-normalized Gaussian distribution
+    return ti.exp(-x**2 / (2 * sigma**2))
+
+
+@ti.kernel
+def touch(mx: float, my: float, size: float):
+    for I in ti.grouped(pixels):
+        mouse_pos = ti.Vector([mx, my]) + 0.5 / N
+        peak_center = gauss((I / N - 0.5).norm(), size)
+        peak_mouse = gauss((I / N - mouse_pos).norm(), size)
+        pixels[I] = peak_center + peak_mouse
+
+
+@ti.func
+def get_vertex(vertex_id):
+    # bottom edge
+    v = ti.Vector([0.5, 0.0])
+    # right edge
+    if vertex_id == 1:
+        v = ti.Vector([1.0, 0.5])
+    # top edge
+    elif vertex_id == 2:
+        v = ti.Vector([0.5, 1.0])
+    # left edge
+    elif vertex_id == 3:
+        v = ti.Vector([0.0, 0.5])
+    return v
+
+
+@ti.kernel
+def march(level: float) -> int:
+    n_edges = 0
+
+    for i, j in ti.ndrange(N - 1, N - 1):
+        case_id = 0
+        if pixels[i, j] > level: case_id |= 1
+        if pixels[i + 1, j] > level: case_id |= 2
+        if pixels[i + 1, j + 1] > level: case_id |= 4
+        if pixels[i, j + 1] > level: case_id |= 8
+
+        for k in range(2):
+            if edge_table[case_id, k][0] == -1:
+                break
+
+            n = ti.atomic_add(n_edges, 1)
+            for l in ti.static(range(2)):
+                vertex_id = edge_table[case_id, k][l]
+                edge_coords[n, l] = ti.Vector([i, j
+                                               ]) + get_vertex(vertex_id) + 0.5
+
+    return n_edges
+
+
+level = 0.2
+
+gui = ti.GUI('Marching squares')
+while gui.running and not gui.get_event(gui.ESCAPE):
+    touch(*gui.get_cursor_pos(), 0.05)
+    n_edges = march(level)
+    edge_coords_np = edge_coords.to_numpy()[:n_edges] / N
+    gui.set_image(ti.imresize(pixels, *gui.res) / level)
+    gui.lines(edge_coords_np[:, 0],
+              edge_coords_np[:, 1],
+              color=0xff66cc,
+              radius=1.5)
+    gui.show()
diff --git a/examples/mciso_advanced.py b/examples/algorithm/mciso_advanced.py
similarity index 100%
rename from examples/mciso_advanced.py
rename to examples/algorithm/mciso_advanced.py
diff --git a/examples/mgpcg.py b/examples/algorithm/mgpcg.py
similarity index 100%
rename from examples/mgpcg.py
rename to examples/algorithm/mgpcg.py
diff --git a/examples/mgpcg_advanced.py b/examples/algorithm/mgpcg_advanced.py
similarity index 100%
rename from examples/mgpcg_advanced.py
rename to examples/algorithm/mgpcg_advanced.py
diff --git a/examples/print_offset.py b/examples/algorithm/print_offset.py
similarity index 100%
rename from examples/print_offset.py
rename to examples/algorithm/print_offset.py
diff --git a/examples/quadtree.py b/examples/algorithm/quadtree.py
similarity index 96%
rename from examples/quadtree.py
rename to examples/algorithm/quadtree.py
index 28cfbec9b4ddd..d13b8989f419d 100644
--- a/examples/quadtree.py
+++ b/examples/algorithm/quadtree.py
@@ -9,10 +9,12 @@
 R = 7
 N = K**R
 
-Broot = ti.root
+Broot = None
 B = ti.root
 for r in range(R):
     B = B.bitmasked(ti.ij, (K, K))
+    if r == 0:
+        Broot = B
 
 qt = ti.field(ti.f32)
 B.place(qt)
diff --git a/examples/autodiff_minimization.py b/examples/autodiff/minimization.py
similarity index 97%
rename from examples/autodiff_minimization.py
rename to examples/autodiff/minimization.py
index a3917306cc651..c59af83e0cdec 100644
--- a/examples/autodiff_minimization.py
+++ b/examples/autodiff/minimization.py
@@ -2,6 +2,8 @@
 
 import taichi as ti
 
+ti.init(arch=ti.cpu)
+
 n = 8
 x = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
 y = ti.field(dtype=ti.f32, shape=n)
diff --git a/examples/mnist.py b/examples/autodiff/mnist.py
similarity index 99%
rename from examples/mnist.py
rename to examples/autodiff/mnist.py
index 19369c1e5b5ab..f96da54ab475c 100644
--- a/examples/mnist.py
+++ b/examples/autodiff/mnist.py
@@ -5,6 +5,8 @@
 
 import taichi as ti
 
+ti.init(arch=ti.gpu)
+
 # ti.runtime.print_preprocessed = True
 # ti.cfg.print_ir = True
 
diff --git a/examples/mnist_download_data.py b/examples/autodiff/mnist_download_data.py
similarity index 99%
rename from examples/mnist_download_data.py
rename to examples/autodiff/mnist_download_data.py
index 8701b82f68892..49014a3aa500f 100644
--- a/examples/mnist_download_data.py
+++ b/examples/autodiff/mnist_download_data.py
@@ -1,9 +1,10 @@
 # https://github.com/hsjeong5/MNIST-for-Numpy
 
-import numpy as np
-from urllib import request
 import gzip
 import pickle
+from urllib import request
+
+import numpy as np
 
 filename = [["training_images", "train-images-idx3-ubyte.gz"],
             ["test_images", "t10k-images-idx3-ubyte.gz"],
diff --git a/examples/autodiff_regression.py b/examples/autodiff/regression.py
similarity index 98%
rename from examples/autodiff_regression.py
rename to examples/autodiff/regression.py
index 1a27a003e4f1d..2692b8afafa87 100644
--- a/examples/autodiff_regression.py
+++ b/examples/autodiff/regression.py
@@ -6,6 +6,8 @@
 import taichi as ti
 import taichi as tc
 
+ti.init(arch=ti.cpu)
+
 tc.set_gdb_trigger(True)
 
 number_coeffs = 4
diff --git a/examples/taichi_autodiff.py b/examples/autodiff/simple_derivative.py
similarity index 97%
rename from examples/taichi_autodiff.py
rename to examples/autodiff/simple_derivative.py
index 579245f32ff4b..32d3e2bf93eb9 100644
--- a/examples/taichi_autodiff.py
+++ b/examples/autodiff/simple_derivative.py
@@ -2,6 +2,8 @@
 
 import taichi as ti
 
+ti.init(arch=ti.cpu)
+
 N = 2048
 x, y = ti.field(ti.f32), ti.field(ti.f32)
 
diff --git a/examples/fullscreen.py b/examples/features/gui/fullscreen.py
similarity index 100%
rename from examples/fullscreen.py
rename to examples/features/gui/fullscreen.py
diff --git a/examples/gui_image_io.py b/examples/features/gui/gui_image_io.py
similarity index 95%
rename from examples/gui_image_io.py
rename to examples/features/gui/gui_image_io.py
index b4893a463e9ec..150475be4e95f 100644
--- a/examples/gui_image_io.py
+++ b/examples/features/gui/gui_image_io.py
@@ -2,6 +2,8 @@
 
 import taichi as ti
 
+ti.init(arch=ti.cpu)
+
 pixel = ti.field(ti.u8, shape=(512, 512, 3))
 
 
diff --git a/examples/gui_widgets.py b/examples/features/gui/gui_widgets.py
similarity index 100%
rename from examples/gui_widgets.py
rename to examples/features/gui/gui_widgets.py
diff --git a/examples/keyboard.py b/examples/features/gui/keyboard.py
similarity index 100%
rename from examples/keyboard.py
rename to examples/features/gui/keyboard.py
diff --git a/examples/export_mesh.py b/examples/features/io/export_mesh.py
similarity index 99%
rename from examples/export_mesh.py
rename to examples/features/io/export_mesh.py
index 89f8ee0afbf38..7927cc5e4c2d8 100644
--- a/examples/export_mesh.py
+++ b/examples/features/io/export_mesh.py
@@ -1,6 +1,9 @@
-import taichi as ti
-import numpy as np
 import os
+
+import numpy as np
+
+import taichi as ti
+
 # A 2D grid with quad faces
 #     y
 #     |
diff --git a/examples/export_ply.py b/examples/features/io/export_ply.py
similarity index 100%
rename from examples/export_ply.py
rename to examples/features/io/export_ply.py
diff --git a/examples/export_videos.py b/examples/features/io/export_videos.py
similarity index 100%
rename from examples/export_videos.py
rename to examples/features/io/export_videos.py
diff --git a/examples/taichi_bitmasked.py b/examples/features/sparse/taichi_bitmasked.py
similarity index 100%
rename from examples/taichi_bitmasked.py
rename to examples/features/sparse/taichi_bitmasked.py
diff --git a/examples/taichi_dynamic.py b/examples/features/sparse/taichi_dynamic.py
similarity index 100%
rename from examples/taichi_dynamic.py
rename to examples/features/sparse/taichi_dynamic.py
diff --git a/examples/taichi_sparse.py b/examples/features/sparse/taichi_sparse.py
similarity index 100%
rename from examples/taichi_sparse.py
rename to examples/features/sparse/taichi_sparse.py
diff --git a/examples/mciso.py b/examples/mciso.py
deleted file mode 100644
index 4bea2d4df26a0..0000000000000
--- a/examples/mciso.py
+++ /dev/null
@@ -1,95 +0,0 @@
-import numpy as np
-
-import taichi as ti
-
-ti.init()
-
-N = 128
-
-_et = np.array(
-    [
-        [[-1, -1], [-1, -1]],  #
-        [[0, 1], [-1, -1]],  #a
-        [[0, 2], [-1, -1]],  #b
-        [[1, 2], [-1, -1]],  #ab
-        [[1, 3], [-1, -1]],  #c
-        [[0, 3], [-1, -1]],  #ca
-        [[2, 3], [0, 1]],  #cb
-        [[2, 3], [-1, -1]],  #cab
-        [[2, 3], [-1, -1]],  #d
-        [[2, 3], [0, 1]],  #da
-        [[0, 3], [-1, -1]],  #db
-        [[1, 3], [-1, -1]],  #dab
-        [[1, 2], [-1, -1]],  #dc
-        [[0, 2], [-1, -1]],  #dca
-        [[0, 1], [-1, -1]],  #dcb
-        [[-1, -1], [-1, -1]],  #dcab
-    ],
-    np.int32)
-
-m = ti.field(float, (N, N))  # field value of voxels
-
-r = ti.Vector.field(2, float, (N**2, 2))  # generated edges
-et = ti.Vector.field(2, int, _et.shape[:2])  # edge table (constant)
-et.from_numpy(_et)
-
-
-@ti.func
-def gauss(x):
-    return ti.exp(-6 * x**2)
-
-
-@ti.kernel
-def touch(mx: float, my: float):
-    for i, j in m:
-        p = ti.Vector([i, j]) / N
-        a = gauss((p - 0.5).norm() / 0.25)
-        p.x -= mx - 0.5 / N
-        p.y -= my - 0.5 / N
-        b = gauss(p.norm() / 0.25)
-        m[i, j] = (a + b) * 3
-
-
-@ti.func
-def list_subscript(a,
-                   i):  # magic method to subscript a list with dynamic index
-    ret = sum(a) * 0
-    for j in ti.static(range(len(a))):
-        if i == j:
-            ret = a[j]
-    return ret
-
-
-@ti.kernel
-def march() -> int:
-    r_n = 0
-
-    for i, j in ti.ndrange(N - 1, N - 1):
-        id = 0
-        if m[i, j] > 1: id |= 1
-        if m[i + 1, j] > 1: id |= 2
-        if m[i, j + 1] > 1: id |= 4
-        if m[i + 1, j + 1] > 1: id |= 8
-
-        E = [ti.Vector(_) + .5 for _ in [(.5, 0), (0, .5), (1, .5), (.5, 1)]]
-
-        for k in range(2):
-            if et[id, k][0] == -1:
-                break
-
-            n = ti.atomic_add(r_n, 1)
-            for l in ti.static(range(2)):
-                e = et[id, k][l]
-                r[n, l] = ti.Vector([i, j]) + list_subscript(E, e)
-
-    return r_n
-
-
-gui = ti.GUI('Marching rect')
-while gui.running and not gui.get_event(gui.ESCAPE):
-    touch(*gui.get_cursor_pos())
-    ret_len = march()
-    ret = r.to_numpy()[:ret_len] / N
-    gui.set_image(ti.imresize(m, *gui.res))
-    gui.lines(ret[:, 0], ret[:, 1], color=0xff66cc, radius=1.5)
-    gui.show()
diff --git a/examples/mpm128.py b/examples/mpm128.py
deleted file mode 100644
index 1795b3ad6f43d..0000000000000
--- a/examples/mpm128.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import taichi as ti
-import numpy as np
-
-ti.init(arch=ti.gpu) # Try to run on GPU
-
-quality = 1 # Use a larger value for higher-res simulations
-n_particles, n_grid = 9000 * quality ** 2, 128 * quality
-dx, inv_dx = 1 / n_grid, float(n_grid)
-dt = 1e-4 / quality
-p_vol, p_rho = (dx * 0.5)**2, 1
-p_mass = p_vol * p_rho
-E, nu = 5e3, 0.2 # Young's modulus and Poisson's ratio
-mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters
-
-x = ti.Vector.field(2, dtype=float, shape=n_particles) # position
-v = ti.Vector.field(2, dtype=float, shape=n_particles) # velocity
-C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # affine velocity field
-F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # deformation gradient
-material = ti.field(dtype=int, shape=n_particles) # material id
-Jp = ti.field(dtype=float, shape=n_particles) # plastic deformation
-grid_v = ti.Vector.field(2, dtype=float, shape=(n_grid, n_grid)) # grid node momentum/velocity
-grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) # grid node mass
-gravity = ti.Vector.field(2, dtype=float, shape=())
-attractor_strength = ti.field(dtype=float, shape=())
-attractor_pos = ti.Vector.field(2, dtype=float, shape=())
-
-@ti.kernel
-def substep():
-  for i, j in grid_m:
-    grid_v[i, j] = [0, 0]
-    grid_m[i, j] = 0
-  for p in x: # Particle state update and scatter to grid (P2G)
-    base = (x[p] * inv_dx - 0.5).cast(int)
-    fx = x[p] * inv_dx - base.cast(float)
-    # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
-    w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1) ** 2, 0.5 * (fx - 0.5) ** 2]
-    F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p] # deformation gradient update
-    h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p])))) # Hardening coefficient: snow gets harder when compressed
-    if material[p] == 1: # jelly, make it softer
-      h = 0.3
-    mu, la = mu_0 * h, lambda_0 * h
-    if material[p] == 0: # liquid
-      mu = 0.0
-    U, sig, V = ti.svd(F[p])
-    J = 1.0
-    for d in ti.static(range(2)):
-      new_sig = sig[d, d]
-      if material[p] == 2:  # Snow
-        new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)  # Plasticity
-      Jp[p] *= sig[d, d] / new_sig
-      sig[d, d] = new_sig
-      J *= new_sig
-    if material[p] == 0:  # Reset deformation gradient to avoid numerical instability
-      F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
-    elif material[p] == 2:
-      F[p] = U @ sig @ V.transpose() # Reconstruct elastic deformation gradient after plasticity
-    stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose() + ti.Matrix.identity(float, 2) * la * J * (J - 1)
-    stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
-    affine = stress + p_mass * C[p]
-    for i, j in ti.static(ti.ndrange(3, 3)): # Loop over 3x3 grid node neighborhood
-      offset = ti.Vector([i, j])
-      dpos = (offset.cast(float) - fx) * dx
-      weight = w[i][0] * w[j][1]
-      grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
-      grid_m[base + offset] += weight * p_mass
-  for i, j in grid_m:
-    if grid_m[i, j] > 0: # No need for epsilon here
-      grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
-      grid_v[i, j] += dt * gravity[None] * 30 # gravity
-      dist = attractor_pos[None] - dx * ti.Vector([i, j])
-      grid_v[i, j] += dist / (0.01 + dist.norm()) * attractor_strength[None] * dt * 100
-      if i < 3 and grid_v[i, j][0] < 0:          grid_v[i, j][0] = 0 # Boundary conditions
-      if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
-      if j < 3 and grid_v[i, j][1] < 0:          grid_v[i, j][1] = 0
-      if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
-  for p in x: # grid to particle (G2P)
-    base = (x[p] * inv_dx - 0.5).cast(int)
-    fx = x[p] * inv_dx - base.cast(float)
-    w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1.0) ** 2, 0.5 * (fx - 0.5) ** 2]
-    new_v = ti.Vector.zero(float, 2)
-    new_C = ti.Matrix.zero(float, 2, 2)
-    for i, j in ti.static(ti.ndrange(3, 3)): # loop over 3x3 grid node neighborhood
-      dpos = ti.Vector([i, j]).cast(float) - fx
-      g_v = grid_v[base + ti.Vector([i, j])]
-      weight = w[i][0] * w[j][1]
-      new_v += weight * g_v
-      new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
-    v[p], C[p] = new_v, new_C
-    x[p] += dt * v[p] # advection
-
-@ti.kernel
-def reset():
-  group_size = n_particles // 3
-  for i in range(n_particles):
-    x[i] = [ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size), ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)]
-    material[i] = i // group_size # 0: fluid 1: jelly 2: snow
-    v[i] = [0, 0]
-    F[i] = ti.Matrix([[1, 0], [0, 1]])
-    Jp[i] = 1
-    C[i] = ti.Matrix.zero(float, 2, 2)
-  
-print("[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset.")
-gui = ti.GUI("Taichi MLS-MPM-128", res=512, background_color=0x112F41)
-reset()
-gravity[None] = [0, -1]
-
-for frame in range(20000):
-  if gui.get_event(ti.GUI.PRESS):
-    if gui.event.key == 'r': reset()
-    elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break
-  if gui.event is not None: gravity[None] = [0, 0] # if had any event
-  if gui.is_pressed(ti.GUI.LEFT,  'a'): gravity[None][0] = -1
-  if gui.is_pressed(ti.GUI.RIGHT, 'd'): gravity[None][0] = 1
-  if gui.is_pressed(ti.GUI.UP,    'w'): gravity[None][1] = 1
-  if gui.is_pressed(ti.GUI.DOWN,  's'): gravity[None][1] = -1
-  mouse = gui.get_cursor_pos()
-  gui.circle((mouse[0], mouse[1]), color=0x336699, radius=15)
-  attractor_pos[None] = [mouse[0], mouse[1]]
-  attractor_strength[None] = 0
-  if gui.is_pressed(ti.GUI.LMB):
-    attractor_strength[None] = 1
-  if gui.is_pressed(ti.GUI.RMB):
-    attractor_strength[None] = -1
-  for s in range(int(2e-3 // dt)):
-    substep()
-  colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
-  gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
-  gui.show() # Change to gui.show(f'{frame:06d}.png') to write images to disk
diff --git a/examples/mpm99.py b/examples/mpm99.py
deleted file mode 100644
index f47fce34be6ba..0000000000000
--- a/examples/mpm99.py
+++ /dev/null
@@ -1,99 +0,0 @@
-import taichi as ti
-import numpy as np
-ti.init(arch=ti.gpu) # Try to run on GPU
-quality = 1 # Use a larger value for higher-res simulations
-n_particles, n_grid = 9000 * quality ** 2, 128 * quality
-dx, inv_dx = 1 / n_grid, float(n_grid)
-dt = 1e-4 / quality
-p_vol, p_rho = (dx * 0.5)**2, 1
-p_mass = p_vol * p_rho
-E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
-mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters
-x = ti.Vector.field(2, dtype=float, shape=n_particles) # position
-v = ti.Vector.field(2, dtype=float, shape=n_particles) # velocity
-C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # affine velocity field
-F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles) # deformation gradient
-material = ti.field(dtype=int, shape=n_particles) # material id
-Jp = ti.field(dtype=float, shape=n_particles) # plastic deformation
-grid_v = ti.Vector.field(2, dtype=float, shape=(n_grid, n_grid)) # grid node momentum/velocity
-grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) # grid node mass
-
-@ti.kernel
-def substep():
-  for i, j in grid_m:
-    grid_v[i, j] = [0, 0]
-    grid_m[i, j] = 0
-  for p in x: # Particle state update and scatter to grid (P2G)
-    base = (x[p] * inv_dx - 0.5).cast(int)
-    fx = x[p] * inv_dx - base.cast(float)
-    # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
-    w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1) ** 2, 0.5 * (fx - 0.5) ** 2]
-    F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p] # deformation gradient update
-    h = ti.exp(10 * (1.0 - Jp[p])) # Hardening coefficient: snow gets harder when compressed
-    if material[p] == 1: # jelly, make it softer
-      h = 0.3
-    mu, la = mu_0 * h, lambda_0 * h
-    if material[p] == 0: # liquid
-      mu = 0.0
-    U, sig, V = ti.svd(F[p])
-    J = 1.0
-    for d in ti.static(range(2)):
-      new_sig = sig[d, d]
-      if material[p] == 2:  # Snow
-        new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)  # Plasticity
-      Jp[p] *= sig[d, d] / new_sig
-      sig[d, d] = new_sig
-      J *= new_sig
-    if material[p] == 0:  # Reset deformation gradient to avoid numerical instability
-      F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
-    elif material[p] == 2:
-      F[p] = U @ sig @ V.transpose() # Reconstruct elastic deformation gradient after plasticity
-    stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose() + ti.Matrix.identity(float, 2) * la * J * (J - 1)
-    stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
-    affine = stress + p_mass * C[p]
-    for i, j in ti.static(ti.ndrange(3, 3)): # Loop over 3x3 grid node neighborhood
-      offset = ti.Vector([i, j])
-      dpos = (offset.cast(float) - fx) * dx
-      weight = w[i][0] * w[j][1]
-      grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
-      grid_m[base + offset] += weight * p_mass
-  for i, j in grid_m:
-    if grid_m[i, j] > 0: # No need for epsilon here
-      grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
-      grid_v[i, j][1] -= dt * 50 # gravity
-      if i < 3 and grid_v[i, j][0] < 0:          grid_v[i, j][0] = 0 # Boundary conditions
-      if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
-      if j < 3 and grid_v[i, j][1] < 0:          grid_v[i, j][1] = 0
-      if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
-  for p in x: # grid to particle (G2P)
-    base = (x[p] * inv_dx - 0.5).cast(int)
-    fx = x[p] * inv_dx - base.cast(float)
-    w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1.0) ** 2, 0.5 * (fx - 0.5) ** 2]
-    new_v = ti.Vector.zero(float, 2)
-    new_C = ti.Matrix.zero(float, 2, 2)
-    for i, j in ti.static(ti.ndrange(3, 3)): # loop over 3x3 grid node neighborhood
-      dpos = ti.Vector([i, j]).cast(float) - fx
-      g_v = grid_v[base + ti.Vector([i, j])]
-      weight = w[i][0] * w[j][1]
-      new_v += weight * g_v
-      new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
-    v[p], C[p] = new_v, new_C
-    x[p] += dt * v[p] # advection
-
-group_size = n_particles // 3
-@ti.kernel
-def initialize():
-  for i in range(n_particles):
-    x[i] = [ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size), ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)]
-    material[i] = i // group_size # 0: fluid 1: jelly 2: snow
-    v[i] = ti.Matrix([0, 0])
-    F[i] = ti.Matrix([[1, 0], [0, 1]])
-    Jp[i] = 1
-initialize()
-gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
-while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
-  for s in range(int(2e-3 // dt)):
-    substep()
-  colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
-  gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
-  gui.show() # Change to gui.show(f'{frame:06d}.png') to write images to disk
diff --git a/examples/cornell_box.py b/examples/rendering/cornell_box.py
similarity index 76%
rename from examples/cornell_box.py
rename to examples/rendering/cornell_box.py
index e74c76462ed84..94e916733ca85 100644
--- a/examples/cornell_box.py
+++ b/examples/rendering/cornell_box.py
@@ -1,9 +1,6 @@
-import math
 import time
 
 import numpy as np
-from renderer_utils import (intersect_sphere, ray_aabb_intersection,
-                            ray_plane_intersect, reflect, refract)
 
 import taichi as ti
 
@@ -40,7 +37,7 @@
 light_normal = ti.Vector([0.0, -1.0, 0.0])
 
 # No absorbtion, integrates over a unit hemisphere
-lambertian_brdf = 1.0 / math.pi
+lambertian_brdf = 1.0 / np.pi
 # diamond!
 refr_idx = 2.4
 
@@ -50,8 +47,8 @@
 
 
 def make_box_transform_matrices():
-    rad = math.pi / 8.0
-    c, s = math.cos(rad), math.sin(rad)
+    rad = np.pi / 8.0
+    c, s = np.cos(rad), np.sin(rad)
     rot = np.array([[c, 0, s, 0], [0, 1, 0, 0], [-s, 0, c, 0], [0, 0, 0, 1]])
     translate = np.array([
         [1, 0, 0, -0.7],
@@ -72,20 +69,91 @@ def make_box_transform_matrices():
 
 
 @ti.func
-def intersect_light(pos, d, tmax):
-    hit, t, _ = ray_aabb_intersection(light_min_pos, light_max_pos, pos, d)
-    if hit and 0 < t < tmax:
-        hit = 1
+def reflect(d, n):
+    # Assuming |d| and |n| are normalized
+    return d - 2.0 * d.dot(n) * n
+
+
+@ti.func
+def refract(d, n, ni_over_nt):
+    # Assuming |d| and |n| are normalized
+    has_r, rd = 0, d
+    dt = d.dot(n)
+    discr = 1.0 - ni_over_nt * ni_over_nt * (1.0 - dt * dt)
+    if discr > 0.0:
+        has_r = 1
+        rd = (ni_over_nt * (d - n * dt) - n * ti.sqrt(discr)).normalized()
     else:
-        hit = 0
-        t = inf
-    return hit, t
+        rd *= 0.0
+    return has_r, rd
+
+
+@ti.func
+def mat_mul_point(m, p):
+    hp = ti.Vector([p[0], p[1], p[2], 1.0])
+    hp = m @ hp
+    hp /= hp[3]
+    return ti.Vector([hp[0], hp[1], hp[2]])
 
 
 @ti.func
-def ray_aabb_intersection2(box_min, box_max, o, d):
-    # Compared to ray_aabb_intersection2(), this also returns the normal of
-    # the nearest t.
+def mat_mul_vec(m, v):
+    hv = ti.Vector([v[0], v[1], v[2], 0.0])
+    hv = m @ hv
+    return ti.Vector([hv[0], hv[1], hv[2]])
+
+
+@ti.func
+def intersect_sphere(pos, d, center, radius):
+    T = pos - center
+    A = 1.0
+    B = 2.0 * T.dot(d)
+    C = T.dot(T) - radius * radius
+    delta = B * B - 4.0 * A * C
+    dist = inf
+    hit_pos = ti.Vector([0.0, 0.0, 0.0])
+
+    if delta > -1e-4:
+        delta = ti.max(delta, 0)
+        sdelta = ti.sqrt(delta)
+        ratio = 0.5 / A
+        ret1 = ratio * (-B - sdelta)
+        dist = ret1
+        if dist < inf:
+            # refinement
+            old_dist = dist
+            new_pos = pos + d * dist
+            T = new_pos - center
+            A = 1.0
+            B = 2.0 * T.dot(d)
+            C = T.dot(T) - radius * radius
+            delta = B * B - 4 * A * C
+            if delta > 0:
+                sdelta = ti.sqrt(delta)
+                ratio = 0.5 / A
+                ret1 = ratio * (-B - sdelta) + old_dist
+                if ret1 > 0:
+                    dist = ret1
+                    hit_pos = new_pos + ratio * (-B - sdelta) * d
+            else:
+                dist = inf
+
+    return dist, hit_pos
+
+
+@ti.func
+def intersect_plane(pos, d, pt_on_plane, norm):
+    dist = inf
+    hit_pos = ti.Vector([0.0, 0.0, 0.0])
+    denom = d.dot(norm)
+    if abs(denom) > eps:
+        dist = norm.dot(pt_on_plane - pos) / denom
+        hit_pos = pos + d * dist
+    return dist, hit_pos
+
+
+@ti.func
+def intersect_aabb(box_min, box_max, o, d):
     intersect = 1
 
     near_t = -inf
@@ -115,39 +183,19 @@ def ray_aabb_intersection2(box_min, box_max, o, d):
     if near_t > far_t:
         intersect = 0
     if intersect:
-        # TODO: Issue#1004...
-        if near_face == 0:
-            near_norm[0] = -1 + near_is_max * 2
-        elif near_face == 1:
-            near_norm[1] = -1 + near_is_max * 2
-        else:
-            near_norm[2] = -1 + near_is_max * 2
-
+        for i in ti.static(range(2)):
+            if near_face == i:
+                near_norm[i] = -1 + near_is_max * 2
     return intersect, near_t, far_t, near_norm
 
 
 @ti.func
-def mat_mul_point(m, p):
-    hp = ti.Vector([p[0], p[1], p[2], 1.0])
-    hp = m @ hp
-    hp /= hp[3]
-    return ti.Vector([hp[0], hp[1], hp[2]])
-
-
-@ti.func
-def mat_mul_vec(m, v):
-    hv = ti.Vector([v[0], v[1], v[2], 0.0])
-    hv = m @ hv
-    return ti.Vector([hv[0], hv[1], hv[2]])
-
-
-@ti.func
-def ray_aabb_intersection2_transformed(box_min, box_max, o, d):
+def intersect_aabb_transformed(box_min, box_max, o, d):
     # Transform the ray to the box's local space
     obj_o = mat_mul_point(box_m_inv, o)
     obj_d = mat_mul_vec(box_m_inv, d)
-    intersect, near_t, _, near_norm = ray_aabb_intersection2(
-        box_min, box_max, obj_o, obj_d)
+    intersect, near_t, _, near_norm = intersect_aabb(box_min, box_max, obj_o,
+                                                     obj_d)
     if intersect and 0 < near_t:
         # Transform the normal in the box's local space to world space
         near_norm = mat_mul_vec(box_m_inv_t, near_norm)
@@ -156,6 +204,18 @@ def ray_aabb_intersection2_transformed(box_min, box_max, o, d):
     return intersect, near_t, near_norm
 
 
+@ti.func
+def intersect_light(pos, d, tmax):
+    hit, t, far_t, near_norm = intersect_aabb(light_min_pos, light_max_pos,
+                                              pos, d)
+    if hit and 0 < t < tmax:
+        hit = 1
+    else:
+        hit = 0
+        t = inf
+    return hit, t
+
+
 @ti.func
 def intersect_scene(pos, ray_dir):
     closest, normal = inf, ti.Vector.zero(ti.f32, 3)
@@ -168,8 +228,8 @@ def intersect_scene(pos, ray_dir):
         normal = (hit_pos - sp1_center).normalized()
         c, mat = ti.Vector([1.0, 1.0, 1.0]), mat_glass
     # left box
-    hit, cur_dist, pnorm = ray_aabb_intersection2_transformed(
-        box_min, box_max, pos, ray_dir)
+    hit, cur_dist, pnorm = intersect_aabb_transformed(box_min, box_max, pos,
+                                                      ray_dir)
     if hit and 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
@@ -177,16 +237,16 @@ def intersect_scene(pos, ray_dir):
 
     # left
     pnorm = ti.Vector([1.0, 0.0, 0.0])
-    cur_dist, _ = ray_plane_intersect(pos, ray_dir, ti.Vector([-1.1, 0.0,
-                                                               0.0]), pnorm)
+    cur_dist, _ = intersect_plane(pos, ray_dir, ti.Vector([-1.1, 0.0, 0.0]),
+                                  pnorm)
     if 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
         c, mat = ti.Vector([0.65, 0.05, 0.05]), mat_lambertian
     # right
     pnorm = ti.Vector([-1.0, 0.0, 0.0])
-    cur_dist, _ = ray_plane_intersect(pos, ray_dir, ti.Vector([1.1, 0.0, 0.0]),
-                                      pnorm)
+    cur_dist, _ = intersect_plane(pos, ray_dir, ti.Vector([1.1, 0.0, 0.0]),
+                                  pnorm)
     if 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
@@ -194,24 +254,24 @@ def intersect_scene(pos, ray_dir):
     # bottom
     gray = ti.Vector([0.93, 0.93, 0.93])
     pnorm = ti.Vector([0.0, 1.0, 0.0])
-    cur_dist, _ = ray_plane_intersect(pos, ray_dir, ti.Vector([0.0, 0.0, 0.0]),
-                                      pnorm)
+    cur_dist, _ = intersect_plane(pos, ray_dir, ti.Vector([0.0, 0.0, 0.0]),
+                                  pnorm)
     if 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
         c, mat = gray, mat_lambertian
     # top
     pnorm = ti.Vector([0.0, -1.0, 0.0])
-    cur_dist, _ = ray_plane_intersect(pos, ray_dir, ti.Vector([0.0, 2.0, 0.0]),
-                                      pnorm)
+    cur_dist, _ = intersect_plane(pos, ray_dir, ti.Vector([0.0, 2.0, 0.0]),
+                                  pnorm)
     if 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
         c, mat = gray, mat_lambertian
     # far
     pnorm = ti.Vector([0.0, 0.0, 1.0])
-    cur_dist, _ = ray_plane_intersect(pos, ray_dir, ti.Vector([0.0, 0.0, 0.0]),
-                                      pnorm)
+    cur_dist, _ = intersect_plane(pos, ray_dir, ti.Vector([0.0, 0.0, 0.0]),
+                                  pnorm)
     if 0 < cur_dist < closest:
         closest = cur_dist
         normal = pnorm
@@ -264,7 +324,7 @@ def compute_area_light_pdf(pos, ray_dir):
 
 @ti.func
 def compute_brdf_pdf(normal, sample_dir):
-    return dot_or_zero(normal, sample_dir) / math.pi
+    return dot_or_zero(normal, sample_dir) / np.pi
 
 
 @ti.func
@@ -279,51 +339,23 @@ def sample_area_light(hit_pos, pos_normal):
 @ti.func
 def sample_brdf(normal):
     # cosine hemisphere sampling
-    # first, uniformly sample on a disk (r, theta)
+    # Uniformly sample on a disk using concentric sampling(r, theta)
+    # https://www.pbr-book.org/3ed-2018/Monte_Carlo_Integration/2D_Sampling_with_Multidimensional_Transformations#CosineSampleHemisphere
     r, theta = 0.0, 0.0
     sx = ti.random() * 2.0 - 1.0
     sy = ti.random() * 2.0 - 1.0
-    if sx >= -sy:
-        if sx > sy:
-            # first region
+    if sx != 0 or sy != 0:
+        if abs(sx) > abs(sy):
             r = sx
-            div = abs(sy / r)
-            if sy > 0.0:
-                theta = div
-            else:
-                theta = 7.0 + div
+            theta = np.pi / 4 * (sy / sx)
         else:
-            # second region
             r = sy
-            div = abs(sx / r)
-            if sx > 0.0:
-                theta = 1.0 + sx / r
-            else:
-                theta = 2.0 + sx / r
-    else:
-        if sx <= sy:
-            # third region
-            r = -sx
-            div = abs(sy / r)
-            if sy > 0.0:
-                theta = 3.0 + div
-            else:
-                theta = 4.0 + div
-        else:
-            # fourth region
-            r = -sy
-            div = abs(sx / r)
-            if sx < 0.0:
-                theta = 5.0 + div
-            else:
-                theta = 6.0 + div
-    # Malley's method
+            theta = np.pi / 4 * (2 - sx / sy)
+    # Apply Malley's method to project disk to hemisphere
     u = ti.Vector([1.0, 0.0, 0.0])
     if abs(normal[1]) < 1 - eps:
         u = normal.cross(ti.Vector([0.0, 1.0, 0.0]))
     v = normal.cross(u)
-
-    theta = theta * math.pi * 0.25
     costt, sintt = ti.cos(theta), ti.sin(theta)
     xy = (u * costt + v * sintt) * r
     zlen = ti.sqrt(max(0.0, 1.0 - xy.dot(xy)))
diff --git a/examples/particle_renderer.py b/examples/rendering/particle_renderer.py
similarity index 100%
rename from examples/particle_renderer.py
rename to examples/rendering/particle_renderer.py
diff --git a/examples/rasterizer.py b/examples/rendering/rasterizer.py
similarity index 100%
rename from examples/rasterizer.py
rename to examples/rendering/rasterizer.py
diff --git a/examples/renderer_utils.py b/examples/rendering/renderer_utils.py
similarity index 100%
rename from examples/renderer_utils.py
rename to examples/rendering/renderer_utils.py
diff --git a/examples/sdf2d.py b/examples/rendering/sdf2d.py
similarity index 100%
rename from examples/sdf2d.py
rename to examples/rendering/sdf2d.py
diff --git a/examples/sdf_renderer.py b/examples/rendering/sdf_renderer.py
similarity index 100%
rename from examples/sdf_renderer.py
rename to examples/rendering/sdf_renderer.py
diff --git a/examples/simple_uv.py b/examples/rendering/simple_uv.py
similarity index 100%
rename from examples/simple_uv.py
rename to examples/rendering/simple_uv.py
diff --git a/examples/taichi_logo.py b/examples/rendering/taichi_logo.py
similarity index 94%
rename from examples/taichi_logo.py
rename to examples/rendering/taichi_logo.py
index f09918f71afae..1a9b59232366b 100644
--- a/examples/taichi_logo.py
+++ b/examples/rendering/taichi_logo.py
@@ -1,5 +1,7 @@
 import taichi as ti
 
+ti.init(arch=ti.cpu)
+
 n = 512
 x = ti.field(ti.f32, shape=(n, n))
 
diff --git a/examples/ad_gravity.py b/examples/simulation/ad_gravity.py
similarity index 100%
rename from examples/ad_gravity.py
rename to examples/simulation/ad_gravity.py
diff --git a/examples/cg_possion.py b/examples/simulation/cg_possion.py
similarity index 100%
rename from examples/cg_possion.py
rename to examples/simulation/cg_possion.py
diff --git a/examples/comet.py b/examples/simulation/comet.py
similarity index 60%
rename from examples/comet.py
rename to examples/simulation/comet.py
index 9fe8442dc2909..ea89ab6e55c08 100644
--- a/examples/comet.py
+++ b/examples/simulation/comet.py
@@ -1,4 +1,4 @@
-import taichi_glsl as tl
+import math
 
 import taichi as ti
 
@@ -11,10 +11,10 @@
 sun = ti.Vector([0.5, 0.5] if dim == 2 else [0.5, 0.5, 0.0])
 gravity = 0.5
 pressure = 0.3
-pps_scale = 0.4
-colorinit = 0.3
-colordeg = 1.6
-initvel = 0.07
+tail_paticle_scale = 0.4
+color_init = 0.3
+color_decay = 1.6
+vel_init = 0.07
 res = 640
 
 inv_m = ti.field(ti.f32)
@@ -22,22 +22,34 @@
 x = ti.Vector.field(dim, ti.f32)
 v = ti.Vector.field(dim, ti.f32)
 ti.root.bitmasked(ti.i, N).place(x, v, inv_m, color)
-xcnt = ti.field(ti.i32, ())
+count = ti.field(ti.i32, ())
 img = ti.field(ti.f32, (res, res))
 
 
+@ti.func
+def rand_unit_2d():
+    a = ti.random() * 2 * math.pi
+    return ti.Vector([ti.cos(a), ti.sin(a)])
+
+
+@ti.func
+def rand_unit_3d():
+    u = rand_unit_2d()
+    s = ti.random() * 2 - 1
+    c = ti.sqrt(1 - s**2)
+    return ti.Vector([c * u[0], c * u[1], s])
+
+
 @ti.kernel
 def substep():
     ti.no_activate(x)
     for i in x:
         r = x[i] - sun
-        ir = r / r.norm(1e-3)**3
-        acci = pressure * ir * inv_m[i]
-        acci += -gravity * ir
-        v[i] += acci * dt
-
+        r_sq_inverse = r / r.norm(1e-3)**3
+        acceleration = (pressure * inv_m[i] - gravity) * r_sq_inverse
+        v[i] += acceleration * dt
         x[i] += v[i] * dt
-        color[i] *= ti.exp(-dt * colordeg)
+        color[i] *= ti.exp(-dt * color_decay)
 
         if not all(-0.1 <= x[i] <= 1.1):
             ti.deactivate(x.snode.parent(), [i])
@@ -46,19 +58,18 @@ def substep():
 @ti.kernel
 def generate():
     r = x[0] - sun
-    ir = 1 / r.norm(1e-3)**2
-    pps = int(pps_scale * ir)
-    for _ in range(pps):
+    n_tail_paticles = int(tail_paticle_scale / r.norm(1e-3)**2)
+    for _ in range(n_tail_paticles):
         r = x[0]
         if ti.static(dim == 3):
-            r = tl.randUnit3D()
+            r = rand_unit_3d()
         else:
-            r = tl.randUnit2D()
-        xi = ti.atomic_add(xcnt[None], 1) % (N - 1) + 1
+            r = rand_unit_2d()
+        xi = ti.atomic_add(count[None], 1) % (N - 1) + 1
         x[xi] = x[0]
-        v[xi] = r * initvel + v[0]
+        v[xi] = r * vel_init + v[0]
         inv_m[xi] = 0.5 + ti.random()
-        color[xi] = colorinit
+        color[xi] = color_init
 
 
 @ti.kernel
diff --git a/examples/euler.py b/examples/simulation/euler.py
similarity index 100%
rename from examples/euler.py
rename to examples/simulation/euler.py
diff --git a/examples/fem128.py b/examples/simulation/fem128.py
similarity index 86%
rename from examples/fem128.py
rename to examples/simulation/fem128.py
index 289f186dca901..07a2d3b3a56a1 100644
--- a/examples/fem128.py
+++ b/examples/simulation/fem128.py
@@ -1,14 +1,15 @@
 import taichi as ti
+
 ti.init(arch=ti.gpu)
 
 N = 12
 dt = 5e-5
 dx = 1 / N
 rho = 4e1
-NF = 2 * N ** 2   # number of faces
-NV = (N + 1) ** 2 # number of vertices
+NF = 2 * N**2  # number of faces
+NV = (N + 1)**2  # number of vertices
 E, nu = 4e4, 0.2  # Young's modulus and Poisson's ratio
-mu, lam = E / 2 / (1 + nu), E * nu / (1 + nu) / (1 - 2 * nu) # Lame parameters
+mu, lam = E / 2 / (1 + nu), E * nu / (1 + nu) / (1 - 2 * nu)  # Lame parameters
 ball_pos, ball_radius = ti.Vector([0.5, 0.0]), 0.31
 damping = 14.5
 
@@ -25,6 +26,7 @@
 attractor_pos = ti.Vector.field(2, float, ())
 attractor_strength = ti.field(float, ())
 
+
 @ti.kernel
 def update_U():
     for i in range(NF):
@@ -38,30 +40,33 @@ def update_U():
         log_J_i = ti.log(F_i.determinant())
         phi_i = mu / 2 * ((F_i.transpose() @ F_i).trace() - 2)
         phi_i -= mu * log_J_i
-        phi_i += lam / 2 * log_J_i ** 2
+        phi_i += lam / 2 * log_J_i**2
         phi[i] = phi_i
         U[None] += V[i] * phi_i
 
+
 @ti.kernel
 def advance():
     for i in range(NV):
-        acc = -pos.grad[i] / (rho * dx ** 2)
-        g = gravity[None] * 0.8 + attractor_strength[None] * (attractor_pos[None] - pos[i]).normalized(1e-5)
+        acc = -pos.grad[i] / (rho * dx**2)
+        g = gravity[None] * 0.8 + attractor_strength[None] * (
+            attractor_pos[None] - pos[i]).normalized(1e-5)
         vel[i] += dt * (acc + g * 40)
         vel[i] *= ti.exp(-dt * damping)
     for i in range(NV):
         # ball boundary condition:
         disp = pos[i] - ball_pos
         disp2 = disp.norm_sqr()
-        if disp2 <= ball_radius ** 2:
+        if disp2 <= ball_radius**2:
             NoV = vel[i].dot(disp)
             if NoV < 0: vel[i] -= NoV * disp / disp2
         cond = pos[i] < 0 and vel[i] < 0 or pos[i] > 1 and vel[i] > 0
         # rect boundary condition:
         for j in ti.static(range(pos.n)):
-           if cond[j]: vel[i][j] = 0
+            if cond[j]: vel[i][j] = 0
         pos[i] += dt * vel[i]
 
+
 @ti.kernel
 def init_pos():
     for i, j in ti.ndrange(N + 1, N + 1):
@@ -74,6 +79,7 @@ def init_pos():
         B_i_inv = ti.Matrix.cols([a - c, b - c])
         B[i] = B_i_inv.inverse()
 
+
 @ti.kernel
 def init_mesh():
     for i, j in ti.ndrange(N, N):
@@ -85,6 +91,7 @@ def init_mesh():
         f2v[k + 0] = [a, b, c]
         f2v[k + 1] = [c, d, a]
 
+
 def paint_phi(gui):
     pos_ = pos.to_numpy()
     phi_ = phi.to_numpy()
@@ -94,12 +101,15 @@ def paint_phi(gui):
     gb = (1 - k) * 0.5
     gui.triangles(a, b, c, color=ti.rgb_to_hex([k + gb, gb, gb]))
 
+
 init_mesh()
 init_pos()
 gravity[None] = [0, -1]
 
 gui = ti.GUI('FEM128')
-print("[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset.")
+print(
+    "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset."
+)
 while gui.running:
     for e in gui.get_events(gui.PRESS):
         if e.key == gui.ESCAPE:
@@ -116,7 +126,8 @@ def paint_phi(gui):
             gravity[None] = [0, +1]
     mouse_pos = gui.get_cursor_pos()
     attractor_pos[None] = mouse_pos
-    attractor_strength[None] = gui.is_pressed(gui.LMB) - gui.is_pressed(gui.RMB)
+    attractor_strength[None] = gui.is_pressed(gui.LMB) - gui.is_pressed(
+        gui.RMB)
     for i in range(50):
         with ti.Tape(loss=U):
             update_U()
diff --git a/examples/fem99.py b/examples/simulation/fem99.py
similarity index 89%
rename from examples/fem99.py
rename to examples/simulation/fem99.py
index 8361fc58cd40c..304d816c606cb 100644
--- a/examples/fem99.py
+++ b/examples/simulation/fem99.py
@@ -1,14 +1,15 @@
 import taichi as ti
+
 ti.init(arch=ti.gpu)
 
 N = 32
 dt = 1e-4
 dx = 1 / N
 rho = 4e1
-NF = 2 * N ** 2   # number of faces
-NV = (N + 1) ** 2 # number of vertices
+NF = 2 * N**2  # number of faces
+NV = (N + 1)**2  # number of vertices
 E, nu = 4e4, 0.2  # Young's modulus and Poisson's ratio
-mu, lam = E / 2 / (1 + nu), E * nu / (1 + nu) / (1 - 2 * nu) # Lame parameters
+mu, lam = E / 2 / (1 + nu), E * nu / (1 + nu) / (1 - 2 * nu)  # Lame parameters
 ball_pos, ball_radius = ti.Vector([0.5, 0.0]), 0.32
 gravity = ti.Vector([0, -40])
 damping = 12.5
@@ -22,6 +23,7 @@
 phi = ti.field(float, NF)  # potential energy of each face (Neo-Hookean)
 U = ti.field(float, (), needs_grad=True)  # total potential energy
 
+
 @ti.kernel
 def update_U():
     for i in range(NF):
@@ -35,29 +37,31 @@ def update_U():
         log_J_i = ti.log(F_i.determinant())
         phi_i = mu / 2 * ((F_i.transpose() @ F_i).trace() - 2)
         phi_i -= mu * log_J_i
-        phi_i += lam / 2 * log_J_i ** 2
+        phi_i += lam / 2 * log_J_i**2
         phi[i] = phi_i
         U[None] += V[i] * phi_i
 
+
 @ti.kernel
 def advance():
     for i in range(NV):
-        acc = -pos.grad[i] / (rho * dx ** 2)
+        acc = -pos.grad[i] / (rho * dx**2)
         vel[i] += dt * (acc + gravity)
         vel[i] *= ti.exp(-dt * damping)
     for i in range(NV):
         # ball boundary condition:
         disp = pos[i] - ball_pos
         disp2 = disp.norm_sqr()
-        if disp2 <= ball_radius ** 2:
+        if disp2 <= ball_radius**2:
             NoV = vel[i].dot(disp)
             if NoV < 0: vel[i] -= NoV * disp / disp2
         # rect boundary condition:
         cond = pos[i] < 0 and vel[i] < 0 or pos[i] > 1 and vel[i] > 0
         for j in ti.static(range(pos.n)):
-           if cond[j]: vel[i][j] = 0
+            if cond[j]: vel[i][j] = 0
         pos[i] += dt * vel[i]
 
+
 @ti.kernel
 def init_pos():
     for i, j in ti.ndrange(N + 1, N + 1):
@@ -70,6 +74,7 @@ def init_pos():
         B_i_inv = ti.Matrix.cols([a - c, b - c])
         B[i] = B_i_inv.inverse()
 
+
 @ti.kernel
 def init_mesh():
     for i, j in ti.ndrange(N, N):
@@ -81,6 +86,7 @@ def init_mesh():
         f2v[k + 0] = [a, b, c]
         f2v[k + 1] = [c, d, a]
 
+
 init_mesh()
 init_pos()
 gui = ti.GUI('FEM99')
diff --git a/examples/fractal.py b/examples/simulation/fractal.py
similarity index 100%
rename from examples/fractal.py
rename to examples/simulation/fractal.py
diff --git a/examples/game_of_life.py b/examples/simulation/game_of_life.py
similarity index 100%
rename from examples/game_of_life.py
rename to examples/simulation/game_of_life.py
diff --git a/examples/mass_spring_3d.py b/examples/simulation/mass_spring_3d.py
similarity index 100%
rename from examples/mass_spring_3d.py
rename to examples/simulation/mass_spring_3d.py
diff --git a/examples/mass_spring_game.py b/examples/simulation/mass_spring_game.py
similarity index 100%
rename from examples/mass_spring_game.py
rename to examples/simulation/mass_spring_game.py
diff --git a/examples/simulation/math_utils.py b/examples/simulation/math_utils.py
new file mode 100644
index 0000000000000..781d8f541808d
--- /dev/null
+++ b/examples/simulation/math_utils.py
@@ -0,0 +1,70 @@
+import math
+
+import taichi as ti
+
+
+@ti.func
+def rand_vector(n):
+    '''
+    Samples a n-dimensional random uniform vector.
+    '''
+    return ti.Vector([ti.random() for _ in range(n)])
+
+
+@ti.func
+def clamp(x, minval, maxval):
+    return min(max(x, minval), maxval)
+
+
+@ti.func
+def rand_unit_2d():
+    '''
+    Uniformly samples a vector on a 2D unit circle.
+    '''
+    a = ti.random() * math.tau
+    return ti.Vector([ti.cos(a), ti.sin(a)])
+
+
+@ti.func
+def rand_unit_3d():
+    '''
+    Uniformly samples a vector on a 3D unit sphere.
+    '''
+    u = rand_unit_2d()
+    s = ti.random() * 2 - 1
+    c = ti.sqrt(1 - s**2)
+    return ti.Vector([c * u[0], c * u[1], s])
+
+
+@ti.func
+def rand_disk_2d():
+    '''
+    Uniformly samples a point within the area of a 2D unit disk.
+    '''
+    x = 2 * ti.random() - 1
+    y = 2 * ti.random() - 1
+    while x * x + y * y > 1:
+        x = 2 * ti.random() - 1
+        y = 2 * ti.random() - 1
+    return ti.Vector([x, y])
+
+
+@ti.func
+def reflect_boundary(pos,
+                     vel,
+                     pmin=0,
+                     pmax=1,
+                     rebound=1,
+                     rebound_perpendicular=1):
+    '''
+    Reflects particle velocity from a rectangular boundary (if collides).
+    `boundaryReflect` takes particle position, velocity and other parameters.
+    '''
+    cond = pos < pmin and vel < 0 or pos > pmax and vel > 0
+    for j in ti.static(range(pos.n)):
+        if cond[j]:
+            vel[j] *= -rebound
+            for k in ti.static(range(pos.n)):
+                if k != j:
+                    vel[k] *= rebound_perpendicular
+    return vel
diff --git a/examples/simulation/mpm128.py b/examples/simulation/mpm128.py
new file mode 100644
index 0000000000000..a32bf420ee679
--- /dev/null
+++ b/examples/simulation/mpm128.py
@@ -0,0 +1,154 @@
+import numpy as np
+
+import taichi as ti
+
+ti.init(arch=ti.gpu)  # Try to run on GPU
+
+quality = 1  # Use a larger value for higher-res simulations
+n_particles, n_grid = 9000 * quality**2, 128 * quality
+dx, inv_dx = 1 / n_grid, float(n_grid)
+dt = 1e-4 / quality
+p_vol, p_rho = (dx * 0.5)**2, 1
+p_mass = p_vol * p_rho
+E, nu = 5e3, 0.2  # Young's modulus and Poisson's ratio
+mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / (
+    (1 + nu) * (1 - 2 * nu))  # Lame parameters
+
+x = ti.Vector.field(2, dtype=float, shape=n_particles)  # position
+v = ti.Vector.field(2, dtype=float, shape=n_particles)  # velocity
+C = ti.Matrix.field(2, 2, dtype=float,
+                    shape=n_particles)  # affine velocity field
+F = ti.Matrix.field(2, 2, dtype=float,
+                    shape=n_particles)  # deformation gradient
+material = ti.field(dtype=int, shape=n_particles)  # material id
+Jp = ti.field(dtype=float, shape=n_particles)  # plastic deformation
+grid_v = ti.Vector.field(2, dtype=float,
+                         shape=(n_grid, n_grid))  # grid node momentum/velocity
+grid_m = ti.field(dtype=float, shape=(n_grid, n_grid))  # grid node mass
+gravity = ti.Vector.field(2, dtype=float, shape=())
+attractor_strength = ti.field(dtype=float, shape=())
+attractor_pos = ti.Vector.field(2, dtype=float, shape=())
+
+
+@ti.kernel
+def substep():
+    for i, j in grid_m:
+        grid_v[i, j] = [0, 0]
+        grid_m[i, j] = 0
+    for p in x:  # Particle state update and scatter to grid (P2G)
+        base = (x[p] * inv_dx - 0.5).cast(int)
+        fx = x[p] * inv_dx - base.cast(float)
+        # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
+        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
+        F[p] = (ti.Matrix.identity(float, 2) +
+                dt * C[p]) @ F[p]  # deformation gradient update
+        h = max(0.1, min(5, ti.exp(10 * (1.0 - Jp[p])))
+                )  # Hardening coefficient: snow gets harder when compressed
+        if material[p] == 1:  # jelly, make it softer
+            h = 0.3
+        mu, la = mu_0 * h, lambda_0 * h
+        if material[p] == 0:  # liquid
+            mu = 0.0
+        U, sig, V = ti.svd(F[p])
+        J = 1.0
+        for d in ti.static(range(2)):
+            new_sig = sig[d, d]
+            if material[p] == 2:  # Snow
+                new_sig = min(max(sig[d, d], 1 - 2.5e-2),
+                              1 + 4.5e-3)  # Plasticity
+            Jp[p] *= sig[d, d] / new_sig
+            sig[d, d] = new_sig
+            J *= new_sig
+        if material[
+                p] == 0:  # Reset deformation gradient to avoid numerical instability
+            F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
+        elif material[p] == 2:
+            F[p] = U @ sig @ V.transpose(
+            )  # Reconstruct elastic deformation gradient after plasticity
+        stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose(
+        ) + ti.Matrix.identity(float, 2) * la * J * (J - 1)
+        stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
+        affine = stress + p_mass * C[p]
+        for i, j in ti.static(ti.ndrange(
+                3, 3)):  # Loop over 3x3 grid node neighborhood
+            offset = ti.Vector([i, j])
+            dpos = (offset.cast(float) - fx) * dx
+            weight = w[i][0] * w[j][1]
+            grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
+            grid_m[base + offset] += weight * p_mass
+    for i, j in grid_m:
+        if grid_m[i, j] > 0:  # No need for epsilon here
+            grid_v[i,
+                   j] = (1 / grid_m[i, j]) * grid_v[i,
+                                                    j]  # Momentum to velocity
+            grid_v[i, j] += dt * gravity[None] * 30  # gravity
+            dist = attractor_pos[None] - dx * ti.Vector([i, j])
+            grid_v[i, j] += dist / (
+                0.01 + dist.norm()) * attractor_strength[None] * dt * 100
+            if i < 3 and grid_v[i, j][0] < 0:
+                grid_v[i, j][0] = 0  # Boundary conditions
+            if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
+            if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
+            if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
+    for p in x:  # grid to particle (G2P)
+        base = (x[p] * inv_dx - 0.5).cast(int)
+        fx = x[p] * inv_dx - base.cast(float)
+        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2]
+        new_v = ti.Vector.zero(float, 2)
+        new_C = ti.Matrix.zero(float, 2, 2)
+        for i, j in ti.static(ti.ndrange(
+                3, 3)):  # loop over 3x3 grid node neighborhood
+            dpos = ti.Vector([i, j]).cast(float) - fx
+            g_v = grid_v[base + ti.Vector([i, j])]
+            weight = w[i][0] * w[j][1]
+            new_v += weight * g_v
+            new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
+        v[p], C[p] = new_v, new_C
+        x[p] += dt * v[p]  # advection
+
+
+@ti.kernel
+def reset():
+    group_size = n_particles // 3
+    for i in range(n_particles):
+        x[i] = [
+            ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
+            ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)
+        ]
+        material[i] = i // group_size  # 0: fluid 1: jelly 2: snow
+        v[i] = [0, 0]
+        F[i] = ti.Matrix([[1, 0], [0, 1]])
+        Jp[i] = 1
+        C[i] = ti.Matrix.zero(float, 2, 2)
+
+
+print(
+    "[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset."
+)
+gui = ti.GUI("Taichi MLS-MPM-128", res=512, background_color=0x112F41)
+reset()
+gravity[None] = [0, -1]
+
+for frame in range(20000):
+    if gui.get_event(ti.GUI.PRESS):
+        if gui.event.key == 'r': reset()
+        elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break
+    if gui.event is not None: gravity[None] = [0, 0]  # if had any event
+    if gui.is_pressed(ti.GUI.LEFT, 'a'): gravity[None][0] = -1
+    if gui.is_pressed(ti.GUI.RIGHT, 'd'): gravity[None][0] = 1
+    if gui.is_pressed(ti.GUI.UP, 'w'): gravity[None][1] = 1
+    if gui.is_pressed(ti.GUI.DOWN, 's'): gravity[None][1] = -1
+    mouse = gui.get_cursor_pos()
+    gui.circle((mouse[0], mouse[1]), color=0x336699, radius=15)
+    attractor_pos[None] = [mouse[0], mouse[1]]
+    attractor_strength[None] = 0
+    if gui.is_pressed(ti.GUI.LMB):
+        attractor_strength[None] = 1
+    if gui.is_pressed(ti.GUI.RMB):
+        attractor_strength[None] = -1
+    for s in range(int(2e-3 // dt)):
+        substep()
+    colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
+    gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
+    gui.show(
+    )  # Change to gui.show(f'{frame:06d}.png') to write images to disk
diff --git a/examples/mpm3d.py b/examples/simulation/mpm3d.py
similarity index 100%
rename from examples/mpm3d.py
rename to examples/simulation/mpm3d.py
diff --git a/examples/mpm88.py b/examples/simulation/mpm88.py
similarity index 99%
rename from examples/mpm88.py
rename to examples/simulation/mpm88.py
index e7adefc026cc9..581b8e37ca389 100644
--- a/examples/mpm88.py
+++ b/examples/simulation/mpm88.py
@@ -1,5 +1,6 @@
 # MPM-MLS in 88 lines of Taichi code, originally created by @yuanming-hu
 import taichi as ti
+
 ti.init(arch=ti.gpu)
 
 n_particles = 8192
@@ -22,6 +23,7 @@
 grid_v = ti.Vector.field(2, float, (n_grid, n_grid))
 grid_m = ti.field(float, (n_grid, n_grid))
 
+
 @ti.kernel
 def substep():
     for i, j in grid_m:
@@ -71,6 +73,7 @@ def substep():
         J[p] *= 1 + dt * new_C.trace()
         C[p] = new_C
 
+
 @ti.kernel
 def init():
     for i in range(n_particles):
@@ -78,6 +81,7 @@ def init():
         v[i] = [0, -1]
         J[i] = 1
 
+
 init()
 gui = ti.GUI('MPM88')
 while gui.running and not gui.get_event(gui.ESCAPE):
diff --git a/examples/simulation/mpm99.py b/examples/simulation/mpm99.py
new file mode 100644
index 0000000000000..9a9d8fe59062c
--- /dev/null
+++ b/examples/simulation/mpm99.py
@@ -0,0 +1,128 @@
+import numpy as np
+
+import taichi as ti
+
+ti.init(arch=ti.gpu)  # Try to run on GPU
+quality = 1  # Use a larger value for higher-res simulations
+n_particles, n_grid = 9000 * quality**2, 128 * quality
+dx, inv_dx = 1 / n_grid, float(n_grid)
+dt = 1e-4 / quality
+p_vol, p_rho = (dx * 0.5)**2, 1
+p_mass = p_vol * p_rho
+E, nu = 0.1e4, 0.2  # Young's modulus and Poisson's ratio
+mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / (
+    (1 + nu) * (1 - 2 * nu))  # Lame parameters
+x = ti.Vector.field(2, dtype=float, shape=n_particles)  # position
+v = ti.Vector.field(2, dtype=float, shape=n_particles)  # velocity
+C = ti.Matrix.field(2, 2, dtype=float,
+                    shape=n_particles)  # affine velocity field
+F = ti.Matrix.field(2, 2, dtype=float,
+                    shape=n_particles)  # deformation gradient
+material = ti.field(dtype=int, shape=n_particles)  # material id
+Jp = ti.field(dtype=float, shape=n_particles)  # plastic deformation
+grid_v = ti.Vector.field(2, dtype=float,
+                         shape=(n_grid, n_grid))  # grid node momentum/velocity
+grid_m = ti.field(dtype=float, shape=(n_grid, n_grid))  # grid node mass
+
+
+@ti.kernel
+def substep():
+    for i, j in grid_m:
+        grid_v[i, j] = [0, 0]
+        grid_m[i, j] = 0
+    for p in x:  # Particle state update and scatter to grid (P2G)
+        base = (x[p] * inv_dx - 0.5).cast(int)
+        fx = x[p] * inv_dx - base.cast(float)
+        # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
+        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
+        F[p] = (ti.Matrix.identity(float, 2) +
+                dt * C[p]) @ F[p]  # deformation gradient update
+        h = ti.exp(
+            10 *
+            (1.0 -
+             Jp[p]))  # Hardening coefficient: snow gets harder when compressed
+        if material[p] == 1:  # jelly, make it softer
+            h = 0.3
+        mu, la = mu_0 * h, lambda_0 * h
+        if material[p] == 0:  # liquid
+            mu = 0.0
+        U, sig, V = ti.svd(F[p])
+        J = 1.0
+        for d in ti.static(range(2)):
+            new_sig = sig[d, d]
+            if material[p] == 2:  # Snow
+                new_sig = min(max(sig[d, d], 1 - 2.5e-2),
+                              1 + 4.5e-3)  # Plasticity
+            Jp[p] *= sig[d, d] / new_sig
+            sig[d, d] = new_sig
+            J *= new_sig
+        if material[
+                p] == 0:  # Reset deformation gradient to avoid numerical instability
+            F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
+        elif material[p] == 2:
+            F[p] = U @ sig @ V.transpose(
+            )  # Reconstruct elastic deformation gradient after plasticity
+        stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose(
+        ) + ti.Matrix.identity(float, 2) * la * J * (J - 1)
+        stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
+        affine = stress + p_mass * C[p]
+        for i, j in ti.static(ti.ndrange(
+                3, 3)):  # Loop over 3x3 grid node neighborhood
+            offset = ti.Vector([i, j])
+            dpos = (offset.cast(float) - fx) * dx
+            weight = w[i][0] * w[j][1]
+            grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
+            grid_m[base + offset] += weight * p_mass
+    for i, j in grid_m:
+        if grid_m[i, j] > 0:  # No need for epsilon here
+            grid_v[i,
+                   j] = (1 / grid_m[i, j]) * grid_v[i,
+                                                    j]  # Momentum to velocity
+            grid_v[i, j][1] -= dt * 50  # gravity
+            if i < 3 and grid_v[i, j][0] < 0:
+                grid_v[i, j][0] = 0  # Boundary conditions
+            if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
+            if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
+            if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
+    for p in x:  # grid to particle (G2P)
+        base = (x[p] * inv_dx - 0.5).cast(int)
+        fx = x[p] * inv_dx - base.cast(float)
+        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2]
+        new_v = ti.Vector.zero(float, 2)
+        new_C = ti.Matrix.zero(float, 2, 2)
+        for i, j in ti.static(ti.ndrange(
+                3, 3)):  # loop over 3x3 grid node neighborhood
+            dpos = ti.Vector([i, j]).cast(float) - fx
+            g_v = grid_v[base + ti.Vector([i, j])]
+            weight = w[i][0] * w[j][1]
+            new_v += weight * g_v
+            new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
+        v[p], C[p] = new_v, new_C
+        x[p] += dt * v[p]  # advection
+
+
+group_size = n_particles // 3
+
+
+@ti.kernel
+def initialize():
+    for i in range(n_particles):
+        x[i] = [
+            ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
+            ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)
+        ]
+        material[i] = i // group_size  # 0: fluid 1: jelly 2: snow
+        v[i] = ti.Matrix([0, 0])
+        F[i] = ti.Matrix([[1, 0], [0, 1]])
+        Jp[i] = 1
+
+
+initialize()
+gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
+while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
+    for s in range(int(2e-3 // dt)):
+        substep()
+    colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
+    gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
+    gui.show(
+    )  # Change to gui.show(f'{frame:06d}.png') to write images to disk
diff --git a/examples/mpm_lagrangian_forces.py b/examples/simulation/mpm_lagrangian_forces.py
similarity index 100%
rename from examples/mpm_lagrangian_forces.py
rename to examples/simulation/mpm_lagrangian_forces.py
diff --git a/examples/nbody_oscillator.py b/examples/simulation/nbody_oscillator.py
similarity index 100%
rename from examples/nbody_oscillator.py
rename to examples/simulation/nbody_oscillator.py
diff --git a/examples/odop_solar.py b/examples/simulation/odop_solar.py
similarity index 100%
rename from examples/odop_solar.py
rename to examples/simulation/odop_solar.py
diff --git a/examples/pbf2d.py b/examples/simulation/pbf2d.py
similarity index 100%
rename from examples/pbf2d.py
rename to examples/simulation/pbf2d.py
diff --git a/examples/physarum.py b/examples/simulation/physarum.py
similarity index 100%
rename from examples/physarum.py
rename to examples/simulation/physarum.py
diff --git a/examples/stable_fluid.py b/examples/simulation/stable_fluid.py
similarity index 100%
rename from examples/stable_fluid.py
rename to examples/simulation/stable_fluid.py
diff --git a/examples/tree_gravity.py b/examples/simulation/tree_gravity.py
similarity index 86%
rename from examples/tree_gravity.py
rename to examples/simulation/tree_gravity.py
index 07e661b70cedc..acabe6f0f55eb 100644
--- a/examples/tree_gravity.py
+++ b/examples/simulation/tree_gravity.py
@@ -1,12 +1,11 @@
 # N-body gravity simulation in 300 lines of Taichi, tree method, no multipole, O(N log N)
 # Author: archibate <1931127624@qq.com>, all left reserved
-import taichi_glsl as tl
+from math_utils import (clamp, rand_disk_2d, rand_unit_3d, rand_vector,
+                        reflect_boundary)
 
 import taichi as ti
 
-ti.init()
-if not hasattr(ti, 'jkl'):
-    ti.jkl = ti.indices(1, 2, 3)
+ti.init(arch=ti.cpu)
 
 kUseTree = True
 #kDisplay = 'tree mouse pixels cmap save_result'
@@ -46,7 +45,8 @@
     node_children = ti.field(ti.i32)
     node_table = ti.root.dense(ti.i, kMaxNodes)
     node_table.place(node_mass, node_particle_id, node_weighted_pos)
-    node_table.dense({2: ti.jk, 3: ti.jkl}[kDim], 2).place(node_children)
+    node_table.dense(ti.indices(*list(range(1, 1 + kDim))),
+                     2).place(node_children)
     node_table_len = ti.field(ti.i32, ())
 
 if 'mouse' in kDisplay:
@@ -131,13 +131,14 @@ def alloc_a_node_for_particle(particle_id, parent, parent_geo_center,
 
 @ti.kernel
 def add_particle_at(mx: ti.f32, my: ti.f32, mass: ti.f32):
-    mouse_pos = tl.vec(mx, my) + tl.randND(2) * (0.05 / kResolution)
+    mouse_pos = ti.Vector([mx, my]) + rand_vector(2) * (0.05 / kResolution)
 
     particle_id = alloc_particle()
     if ti.static(kDim == 2):
         particle_pos[particle_id] = mouse_pos
     else:
-        particle_pos[particle_id] = tl.vec(mouse_pos, 0.0)
+        particle_pos[particle_id] = ti.Vector(
+            [mouse_pos[0], mouse_pos[1], 0.0])
     particle_mass[particle_id] = mass
 
 
@@ -146,15 +147,16 @@ def add_random_particles(angular_velocity: ti.f32):
     num = ti.static(1)
     particle_id = alloc_particle()
     if ti.static(kDim == 2):
-        particle_pos[particle_id] = tl.randSolid2D() * 0.2 + 0.5
+        particle_pos[particle_id] = rand_disk_2d() * 0.2 + 0.5
         velocity = (particle_pos[particle_id] - 0.5) * angular_velocity * 250
-        particle_vel[particle_id] = tl.vec(-velocity.y, velocity.x)
+        particle_vel[particle_id] = ti.Vector([-velocity.y, velocity.x])
     else:
-        particle_pos[particle_id] = tl.randUnit3D() * 0.2 + 0.5
-        velocity = (particle_pos[particle_id].xy -
+        particle_pos[particle_id] = rand_unit_3d() * 0.2 + 0.5
+        velocity = (ti.Vector(
+            [particle_pos[particle_id].x, particle_pos[particle_id].y]) -
                     0.5) * angular_velocity * 180
-        particle_vel[particle_id] = tl.vec(-velocity.y, velocity.x, 0.0)
-    particle_mass[particle_id] = tl.randRange(0.0, 1.5)
+        particle_vel[particle_id] = ti.Vector([-velocity.y, velocity.x, 0.0])
+    particle_mass[particle_id] = 1.5 * ti.random()
 
 
 @ti.kernel
@@ -182,7 +184,8 @@ def build_tree():
 
 @ti.func
 def gravity_func(distance):
-    return tl.normalizePow(distance, -2, 1e-3)
+    norm_sqr = distance.norm_sqr() + 1e-3
+    return distance / norm_sqr**(3 / 2)
 
 
 @ti.func
@@ -248,35 +251,26 @@ def substep_tree():
     while particle_id < particle_table_len[None]:
         acceleration = get_tree_gravity_at(particle_pos[particle_id])
         particle_vel[particle_id] += acceleration * dt
-        # well... seems our tree inserter will break if particle out-of-bound:
-        particle_vel[particle_id] = tl.boundReflect(particle_pos[particle_id],
-                                                    particle_vel[particle_id],
-                                                    0, 1)
+        # well... seems our tree inserter will break if a particle is out-of-bound:
+        particle_vel[particle_id] = reflect_boundary(particle_pos[particle_id],
+                                                     particle_vel[particle_id],
+                                                     0, 1)
         particle_id = particle_id + 1
     for i in range(particle_table_len[None]):
         particle_pos[i] += particle_vel[i] * dt
 
 
-@ti.kernel
-def render_arrows(mx: ti.f32, my: ti.f32):
-    pos = tl.vec(mx, my)
-    acc = get_raw_gravity_at(pos) * 0.001
-    tl.paintArrow(display_image, pos, acc, tl.D.yyx)
-    acc_tree = get_tree_gravity_at(pos) * 0.001
-    tl.paintArrow(display_image, pos, acc_tree, tl.D.yxy)
-
-
 @ti.kernel
 def render_pixels():
     for i in range(particle_table_len[None]):
-        position = particle_pos[i].xy
+        position = ti.Vector([particle_pos[i].x, particle_pos[i].y])
         pix = int(position * kResolution)
-        display_image[tl.clamp(pix, 0, kResolution - 1)] += 0.3
+        display_image[clamp(pix, 0, kResolution - 1)] += 0.3
 
 
 def render_tree(gui,
                 parent=0,
-                parent_geo_center=tl.vec(0.5, 0.5),
+                parent_geo_center=ti.Vector([0.5, 0.5]),
                 parent_geo_size=1.0):
     child_geo_size = parent_geo_size * 0.5
     if node_particle_id[parent] >= 0:
@@ -324,8 +318,6 @@ def render_tree(gui,
         substep_raw()
     if len(kDisplay) and 'trace' not in kDisplay:
         display_image.fill(0)
-    if 'mouse' in kDisplay:
-        render_arrows(*gui.get_cursor_pos())
     if 'pixels' in kDisplay:
         render_pixels()
     if 'cmap' in kDisplay:
diff --git a/examples/vortex_rings.py b/examples/simulation/vortex_rings.py
similarity index 100%
rename from examples/vortex_rings.py
rename to examples/simulation/vortex_rings.py
diff --git a/examples/waterwave.py b/examples/simulation/waterwave.py
similarity index 52%
rename from examples/waterwave.py
rename to examples/simulation/waterwave.py
index 76845ff15cf2a..717e3bc001b25 100644
--- a/examples/waterwave.py
+++ b/examples/simulation/waterwave.py
@@ -1,14 +1,13 @@
-import numpy as np
+# Water wave effect partially based on shallow water equations
+# https://en.wikipedia.org/wiki/Shallow_water_equations#Non-conservative_form
 
 import taichi as ti
 
 ti.init(arch=ti.gpu)
 
 light_color = 1
-kappa = 2
-gamma = 0.2
-eta = 1.333
-depth = 4
+gravity = 2.0  # larger gravity makes wave propagates faster
+damping = 0.2  # larger damping makes wave vanishes faster when propagating
 dx = 0.02
 dt = 0.01
 shape = 512, 512
@@ -16,17 +15,18 @@
 background = ti.field(dtype=float, shape=shape)
 height = ti.field(dtype=float, shape=shape)
 velocity = ti.field(dtype=float, shape=shape)
-acceleration = ti.field(dtype=float, shape=shape)
 
 
 @ti.kernel
 def reset():
     for i, j in height:
         t = i // 16 + j // 16
-        background[i, j] = (t * 0.5) % 1.0
+        if t % 2 == 0:
+            background[i, j] = 0.0
+        else:
+            background[i, j] = 0.25
         height[i, j] = 0
         velocity[i, j] = 0
-        acceleration[i, j] = 0
 
 
 @ti.func
@@ -43,52 +43,35 @@ def gradient(i, j):
     ]) * (0.5 / dx)
 
 
-@ti.func
-def take_linear(i, j):
-    m, n = int(i), int(j)
-    i, j = i - m, j - n
-    ret = 0.0
-    if 0 <= i < shape[0] and 0 <= i < shape[1]:
-        ret = (i * j * background[m + 1, n + 1] +
-               (1 - i) * j * background[m, n + 1] + i *
-               (1 - j) * background[m + 1, n] + (1 - i) *
-               (1 - j) * background[m, n])
-    return ret
-
-
 @ti.kernel
-def touch_at(hurt: ti.f32, x: ti.f32, y: ti.f32):
+def create_wave(amplitude: ti.f32, x: ti.f32, y: ti.f32):
     for i, j in ti.ndrange((1, shape[0] - 1), (1, shape[1] - 1)):
         r2 = (i - x)**2 + (j - y)**2
-        height[i, j] = height[i, j] + hurt * ti.exp(-0.02 * r2)
+        height[i, j] = height[i, j] + amplitude * ti.exp(-0.02 * r2)
 
 
 @ti.kernel
 def update():
     for i, j in ti.ndrange((1, shape[0] - 1), (1, shape[1] - 1)):
-        acceleration[i, j] = kappa * laplacian(i, j) - gamma * velocity[i, j]
-
-    for i, j in ti.ndrange((1, shape[0] - 1), (1, shape[1] - 1)):
-        velocity[i, j] = velocity[i, j] + acceleration[i, j] * dt
+        acceleration = gravity * laplacian(i, j) - damping * velocity[i, j]
+        velocity[i, j] = velocity[i, j] + acceleration * dt
         height[i, j] = height[i, j] + velocity[i, j] * dt
 
 
 @ti.kernel
-def paint():
+def visualize_wave():
+    # visualizes the wave using a fresnel-like shading
+    # a brighter color indicates a steeper wave
+    # (closer to grazing angle when looked from above)
     for i, j in pixels:
         g = gradient(i, j)
-        # https://www.jianshu.com/p/66a40b06b436
         cos_i = 1 / ti.sqrt(1 + g.norm_sqr())
-        cos_o = ti.sqrt(1 - (1 - (cos_i)**2) * (1 / eta**2))
-        fr = pow(1 - cos_i, 2)
-        coh = cos_o * depth
-        g = g * coh
-        k, l = g[0], g[1]
-        color = take_linear(i + k, j + l)
-        pixels[i, j] = (1 - fr) * color + fr * light_color
+        brightness = pow(1 - cos_i, 2)
+        color = background[i, j]
+        pixels[i, j] = (1 - brightness) * color + brightness * light_color
 
 
-print("[Hint] click on the window to create wavelet")
+print("[Hint] click on the window to create waves")
 
 reset()
 gui = ti.GUI('Water Wave', shape)
@@ -100,8 +83,8 @@ def paint():
             reset()
         elif e.key == ti.GUI.LMB:
             x, y = e.pos
-            touch_at(3, x * shape[0], y * shape[1])
+            create_wave(3, x * shape[0], y * shape[1])
     update()
-    paint()
+    visualize_wave()
     gui.set_image(pixels)
     gui.show()
diff --git a/python/taichi/main.py b/python/taichi/main.py
index 12da831c8d9e0..aed5e49ecdc56 100644
--- a/python/taichi/main.py
+++ b/python/taichi/main.py
@@ -132,7 +132,10 @@ def _get_available_examples() -> set:
         """Get a set of all available example names."""
         examples_dir = TaichiMain._get_examples_dir()
         all_examples = examples_dir.rglob('*.py')
-        all_example_names = {Path(f).stem for f in all_examples}
+        all_example_names = {
+            Path(f).stem: Path(f).parent
+            for f in all_examples
+        }
         return all_example_names
 
     @staticmethod
@@ -156,8 +159,8 @@ def example(self, arguments: list = sys.argv[2:]):
         parser.add_argument(
             "name",
             help=f"Name of an example (supports .py extension too)\n",
-            type=TaichiMain._example_choices_type(choices),
-            choices=sorted(choices))
+            type=TaichiMain._example_choices_type(choices.keys()),
+            choices=sorted(choices.keys()))
         parser.add_argument(
             '-p',
             '--print',
@@ -182,9 +185,10 @@ def example(self, arguments: list = sys.argv[2:]):
         args = parser.parse_args(arguments)
 
         examples_dir = TaichiMain._get_examples_dir()
-        target = str((examples_dir / f"{args.name}.py").resolve())
+        target = str(
+            (examples_dir / choices[args.name] / f"{args.name}.py").resolve())
         # path for examples needs to be modified for implicit relative imports
-        sys.path.append(str(examples_dir.resolve()))
+        sys.path.append(str((examples_dir / choices[args.name]).resolve()))
 
         # Short circuit for testing
         if self.test_mode: return args