Created
March 4, 2023 14:08
-
-
Save FantasyVR/587e633d9f9140a3d917cedbb42f7a2c to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 | |
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment