Created
January 27, 2022 04:20
-
-
Save bcolloran/7b313bd514ca6ea9c1f7ee94c30b74de 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
# from https://github.com/taichi-dev/taichi/issues/2102 | |
import taichi as ti | |
ti.init(default_fp=ti.f64, arch=ti.gpu) | |
# Sparse Grids | |
# ---Params | |
dim = 2 | |
invDx = 1.0 / 0.003 | |
nGrid = ti.ceil(invDx) | |
grid_size = 4096 | |
grid_block_size = 128 | |
leaf_block_size = 16 if dim == 2 else 8 | |
indices = ti.ij if dim == 2 else ti.ijk | |
# offset = tuple(-grid_size // 2 for _ in range(dim)) | |
offset = tuple(0 for _ in range(dim)) | |
# ---Grid Shapes and Pointers | |
grid = ti.root.pointer(indices, grid_size // grid_block_size) # 32 | |
block = grid.pointer(indices, grid_block_size // leaf_block_size) # 8 | |
pid = ti.field(int) | |
block.dynamic(ti.axes(dim), 1024 * 1024, chunk_size=leaf_block_size ** dim * 8).place( | |
pid, offset=offset + (0,) | |
) | |
useSparse = True | |
grid_m1 = ti.field( | |
dtype=float, shape=(nGrid, nGrid) | |
) # grid node field 1 mass is nGrid x nGrid, each grid node has a mass for each field | |
if useSparse: | |
grid_m1 = ti.field(dtype=float) | |
grid2 = ti.root.pointer(indices, grid_size // grid_block_size) # 32 | |
block2 = grid2.pointer(indices, grid_block_size // leaf_block_size) # 8 | |
block2.dense(indices, leaf_block_size).place(grid_m1, offset=offset) | |
# Particle Structures | |
x = ti.Vector.field(dim, dtype=float) # position | |
mp = ti.field(dtype=float) # particle masses | |
particle = ti.root.dynamic( | |
ti.i, 2 ** 27, 2 ** 19 | |
) # 2**20 causes problems in CUDA (maybe asking for too much space) | |
particle.place(x, mp) | |
@ti.func | |
def stencil_range(): | |
return ti.ndrange(*((3,) * dim)) | |
def addParticles(): | |
N = 10 | |
w = 0.2 | |
dw = w / N | |
for i in range(N): | |
for j in range(N): | |
x[i * 10 + j] = [0.4 + (i * dw), 0.4 + (j * dw)] | |
mp[i * 10 + j] = 0.001 | |
@ti.kernel | |
def build_pid(): | |
ti.block_dim(64) # this sets the number of threads per block / block dimension | |
for p in x: | |
base = int(ti.floor(x[p] * invDx - 0.5)) | |
ti.append(pid.parent(), base - ti.Vector(offset), p) | |
@ti.kernel | |
def p2g(): | |
ti.block_dim(256) | |
ti.no_activate(particle) | |
particleCount = 0 | |
for I in ti.grouped(pid): | |
# print(I) | |
p = pid[I] | |
particleCount += 1 | |
base = ti.floor(x[p] * invDx - 0.5).cast(int) | |
for D in ti.static(range(dim)): | |
base[D] = ti.assume_in_range(base[D], I[D], 0, 1) | |
fx = x[p] * invDx - base.cast(float) | |
w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1) ** 2, 0.5 * (fx - 0.5) ** 2] | |
# P2G for mass | |
for offset in ti.static(ti.grouped(stencil_range())): | |
# ti.Vector(offset).cast(float) | |
# offset.cast(float) | |
# Loop over grid node stencil | |
gridIdx = base + offset | |
weight = 1.0 | |
for d in ti.static(range(dim)): | |
weight *= w[offset[d]][d] | |
grid_m1[gridIdx] += ( | |
weight * mp[p] | |
) # add mass to active field for this particle | |
print("particleCount", particleCount) | |
@ti.kernel | |
def momentumToVelocity(): | |
activeNodes = 0 | |
for I in ti.grouped(grid_m1): | |
if grid_m1[I] > 0.0001: | |
# print("I:", I, " m1:", grid_m1[I]) | |
activeNodes += 1 | |
print("activeNodes:", activeNodes) | |
addParticles() | |
grid.deactivate_all() | |
build_pid() | |
p2g() | |
momentumToVelocity() |
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
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 | |
particle = ti.root.dense(ti.i, n_particles) | |
particle.place(x, v, C, F, material, Jp) | |
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=()) | |
group_size = n_particles // 3 | |
water = ti.Vector.field(2, dtype=float, shape=group_size) # position | |
jelly = ti.Vector.field(2, dtype=float, shape=group_size) # position | |
snow = ti.Vector.field(2, dtype=float, shape=group_size) # position | |
mouse_circle = ti.Vector.field(2, dtype=float, shape=(1,)) | |
grid_block_size = 32 | |
leaf_block_size = 8 | |
# grid = ti.root.pointer(ti.ij, n_grid // grid_block_size) # 32 | |
# block = grid.pointer(ti.ij, grid_block_size // leaf_block_size) # 8 | |
block = ti.root.pointer(ti.ij, n_grid // leaf_block_size) # 32 | |
pid = ti.field(int) | |
dim = 2 | |
# offset = tuple(0 for _ in range(dim)) | |
block.dynamic(ti.axes(dim), 1024 * 1024, chunk_size=leaf_block_size ** 2 * 8).place( | |
pid, | |
# offset=offset + (0,) | |
) | |
@ti.kernel | |
def build_pid(): | |
ti.block_dim(64) # this sets the number of threads per block / block dimension | |
for p in x: | |
# base = int(ti.floor(x[p] * inv_dx - 0.5)) | |
base = (x[p] * inv_dx - 0.5).cast(int) | |
ti.append(pid.parent(), base, p) | |
@ti.kernel | |
def substep(): | |
for i, j in grid_m: | |
grid_v[i, j] = [0, 0] | |
grid_m[i, j] = 0 | |
ti.no_activate(particle) | |
for I in ti.grouped(pid): | |
p = pid[I] | |
# 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: | |
# Reconstruct elastic deformation gradient after plasticity | |
F[p] = U @ sig @ V.transpose() | |
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 I in ti.grouped(pid): | |
p = pid[I] | |
# 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(): | |
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) | |
@ti.kernel | |
def render(): | |
for i in range(group_size): | |
water[i] = x[i] | |
jelly[i] = x[i + group_size] | |
snow[i] = x[i + 2 * group_size] | |
print( | |
"[Hint] Use WSAD/arrow keys to control gravity. Use left/right mouse bottons to attract/repel. Press R to reset." | |
) | |
res = (512, 512) | |
window = ti.ui.Window("Taichi MLS-MPM-128", res=res, vsync=True) | |
canvas = window.get_canvas() | |
radius = 0.003 | |
reset() | |
gravity[None] = [0, -1] | |
while window.running: | |
if window.get_event(ti.ui.PRESS): | |
if window.event.key == "r": | |
reset() | |
elif window.event.key in [ti.ui.ESCAPE]: | |
break | |
if window.event is not None: | |
gravity[None] = [0, 0] # if had any event | |
if window.is_pressed(ti.ui.LEFT, "a"): | |
gravity[None][0] = -1 | |
if window.is_pressed(ti.ui.RIGHT, "d"): | |
gravity[None][0] = 1 | |
if window.is_pressed(ti.ui.UP, "w"): | |
gravity[None][1] = 1 | |
if window.is_pressed(ti.ui.DOWN, "s"): | |
gravity[None][1] = -1 | |
mouse = window.get_cursor_pos() | |
mouse_circle[0] = ti.Vector([mouse[0], mouse[1]]) | |
canvas.circles(mouse_circle, color=(0.2, 0.4, 0.6), radius=0.05) | |
attractor_pos[None] = [mouse[0], mouse[1]] | |
attractor_strength[None] = 0 | |
if window.is_pressed(ti.ui.LMB): | |
attractor_strength[None] = 1 | |
if window.is_pressed(ti.ui.RMB): | |
attractor_strength[None] = -1 | |
for s in range(int(2e-3 // dt)): | |
build_pid() | |
substep() | |
render() | |
canvas.set_background_color((0.067, 0.184, 0.255)) | |
canvas.circles(water, radius=radius, color=(0, 0.5, 0.5)) | |
canvas.circles(jelly, radius=radius, color=(0.93, 0.33, 0.23)) | |
canvas.circles(snow, radius=radius, color=(1, 1, 1)) | |
window.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment