Skip to content

Instantly share code, notes, and snippets.

@bcolloran
Created January 27, 2022 04:20
Show Gist options
  • Save bcolloran/7b313bd514ca6ea9c1f7ee94c30b74de to your computer and use it in GitHub Desktop.
Save bcolloran/7b313bd514ca6ea9c1f7ee94c30b74de to your computer and use it in GitHub Desktop.
# 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()
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