mpm88_parallel_p2g.py

# MPM-MLS in 88 lines of Taichi code, originally created by @yuanming-hu
# P2G is parallel in this script.
import taichi as ti

ti.init(arch=ti.gpu)

n_particles = 8192
n_grid = 128
dx = 1 / n_grid
dt = 2e-4

p_rho = 1
p_vol = (dx * 0.5)**2
p_mass = p_vol * p_rho
gravity = 9.8
bound = 3
E = 400

x = ti.Vector.field(2, float, n_particles)
v = ti.Vector.field(2, float, n_particles)
C = ti.Matrix.field(2, 2, float, n_particles)
J = ti.field(float, n_particles)

grid_v = ti.Vector.field(2, float, (n_grid, n_grid))
grid_m = ti.field(float, (n_grid, n_grid))

grid_num_particles = ti.field(int)
max_num_particles_per_cell = 100
grid2particles = ti.field(int)
grid_weight = ti.field(float)
grid_snode = ti.root.dense(ti.ij, n_grid)
grid_snode.place(grid_num_particles)
grid_snode.dense(ti.k, max_num_particles_per_cell).place(grid2particles)
grid_snode.dense(ti.k, max_num_particles_per_cell).place(grid_weight)

@ti.func
def get_base(pos):
    return int(pos / dx - 0.5)

@ti.func
def is_in_simulation_grid(c):
    # @c: Vector(i32)
    return 0 <= c[0] and c[0] < n_grid and 0 <= c[1] and c[1] < n_grid

@ti.kernel
def substep():
    # reset grid
    for i, j in grid_m:
        grid_v[i, j] = [0, 0]
        grid_m[i, j] = 0

    # # P2G
    # for p in x:
    #     Xp = x[p] / dx
    #     base = int(Xp - 0.5)
    #     fx = Xp - base
    #     w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
    #     stress = -dt * 4 * E * p_vol * (J[p] - 1) / dx**2
    #     affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
    #     for i, j in ti.static(ti.ndrange(3, 3)):
    #         offset = ti.Vector([i, j])
    #         dpos = (offset - fx) * dx
    #         weight = w[i].x * w[j].y
    #         grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos) # momentum
    #         grid_m[base + offset] += weight * p_mass
    
    # clear neighbor lookup table
    for I in ti.grouped(grid_num_particles):
        grid_num_particles[I] = 0
    # update grid
    for p_i in x:
        Xp = x[p_i] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        for offs in ti.static(ti.grouped(ti.ndrange(3, 3))):
            grid_index = base + offs
            if is_in_simulation_grid(grid_index):
                # TODO: check if this is correct
                particle_index = ti.atomic_add(grid_num_particles[grid_index], 1)
                grid2particles[grid_index, particle_index] = p_i
                grid_weight[grid_index, particle_index] = w[offs.x].x * w[offs.y].y

    # Traverse grid
    for grid_index in ti.grouped(ti.ndrange(n_grid, n_grid)):
        for k in range(grid_num_particles[grid_index]): # Neighbor search
            p = grid2particles[grid_index, k]
            weight = grid_weight[grid_index, k]
            stress = -dt * 4 * E * p_vol * (J[p] - 1) / dx**2
            affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
            dpos = grid_index * dx - x[p]
            grid_v[grid_index] = grid_v[grid_index] + weight * (p_mass * v[p] + affine @ dpos) # momentum
            grid_m[grid_index] = grid_m[grid_index] + weight * p_mass
    

    for i, j in grid_m:
        if grid_m[i, j] > 0:
            grid_v[i, j] /= grid_m[i, j]
        grid_v[i, j].y -= dt * gravity
        if i < bound and grid_v[i, j].x < 0:
            grid_v[i, j].x = 0
        if i > n_grid - bound and grid_v[i, j].x > 0:
            grid_v[i, j].x = 0
        if j < bound and grid_v[i, j].y < 0:
            grid_v[i, j].y = 0
        if j > n_grid - bound and grid_v[i, j].y > 0:
            grid_v[i, j].y = 0
    # G2P
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**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)):
            offset = ti.Vector([i, j])
            dpos = (offset - fx) * dx
            weight = w[i].x * w[j].y
            g_v = grid_v[base + offset]
            new_v += weight * g_v
            new_C += 4 * weight * g_v.outer_product(dpos) / dx**2
        v[p] = new_v
        x[p] += dt * v[p]
        J[p] *= 1 + dt * new_C.trace()
        C[p] = new_C


@ti.kernel
def init():
    for i in range(n_particles):
        x[i] = [ti.random() * 0.4 + 0.2, ti.random() * 0.4 + 0.2]
        v[i] = [0, -1]
        J[i] = 1


init()
gui = ti.GUI('MPM88')
pause = False 
while gui.running:

    if gui.get_event(ti.GUI.SPACE):
        pause = not pause

    if not pause:
        for s in range(50):
            substep()
    gui.clear(0x112F41)
    gui.circles(x.to_numpy(), radius=1.5, color=0x068587)
    gui.show()