def substep():
for i, j in ti.ndrange(n_grid, n_grid):
grid_v[i, j] = [0, 0]
grid_m[i, j] = 0
for p in range(n_particles): # 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 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1), 0.5 * ti.sqr(fx - 0.5)]
F[p] = (ti.Matrix.identity(ti.f32, 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(ti.f32, 2) * ti.sqrt(J)
elif material[p] == 2:
F[p] = U @ sig @ V.T() # Reconstruct elastic deformation gradient after plasticity
stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T() + ti.Matrix.identity(ti.f32, 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 ti.ndrange(n_grid, n_grid):
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 range(n_particles): # grid to particle (G2P)
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0), 0.5 * ti.sqr(fx - 0.5)]
new_v = ti.Vector.zero(ti.f32, 2)
new_C = ti.Matrix.zero(ti.f32, 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 * ti.outer_product(g_v, dpos)
v[p], C[p] = new_v, new_C
x[p] += dt * v[p] # advection
def substep():
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
stress = -dt * p_vol * (J[p] - 1) * 4 * inv_dx * inv_dx * E
affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
dpos = (offset.cast(float) - fx) * dx
weight = w[i][0] * w[j][1]
grid_v[base + offset].atomic_add(
weight * (p_mass * v[p] + affine @ dpos))
grid_m[base + offset].atomic_add(weight * p_mass)
for i, j in grid_m:
if grid_m[i, j] > 0:
bound = 3
inv_m = 1 / grid_m[i, j]
grid_v[i, j] = inv_m * grid_v[i, j]
grid_v[i, j][1] -= dt * 9.8
if i < bound and grid_v[i, j][0] < 0:
grid_v[i, j][0] = 0
if i > n_grid - bound and grid_v[i, j][0] > 0:
grid_v[i, j][0] = 0
if j < bound and grid_v[i, j][1] < 0:
grid_v[i, j][1] = 0
if j > n_grid - bound and grid_v[i, j][1] > 0:
grid_v[i, j][1] = 0
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
]
new_v = ti.Vector.zero(ti.f32, 2)
new_C = ti.Matrix.zero(ti.f32, 2, 2)
for i in ti.static(range(3)):
for j in ti.static(range(3)):
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 = new_v + weight * g_v
new_C = new_C + 4 * weight * ti.outer_product(g_v,
dpos) * inv_dx
v[p] = new_v
x[p] += dt * v[p]
J[p] *= 1 + dt * new_C.trace()
C[p] = new_C
def g2p():
for p in x:
base = ti.cast(x[p] * inv_dx - 0.5, ti.i32)
fx = x[p] * inv_dx - ti.cast(base, ti.f32)
w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2]
new_v = ti.Vector([0.0, 0.0])
new_C = ti.Matrix([[0.0, 0.0], [0.0, 0.0]])
for i in ti.static(range(3)):
for j in ti.static(range(3)):
dpos = ti.cast(ti.Vector([i, j]), ti.f32) - fx
g_v = grid_v[base(0) + i, base(1) + j]
weight = w[i](0) * w[j](1)
new_v += weight * g_v
new_C += 4 * weight * ti.outer_product(g_v, dpos) * inv_dx
v[p] = new_v
x[p] += dt * v[p]
C[p] = new_C
def g2p(f: ti.i32):
for p in range(n_particles):
base = ti.cast(x[f, p] * inv_dx - 0.5, ti.i32)
fx = x[f, p] * inv_dx - ti.cast(base, real)
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)]
new_v = ti.Vector([0.0, 0.0])
new_C = ti.Matrix([[0.0, 0.0], [0.0, 0.0]])
for i in ti.static(range(3)):
for j in ti.static(range(3)):
dpos = ti.cast(ti.Vector([i, j]), real) - fx
g_v = grid_v_out[base(0) + i, base(1) + j]
weight = w[i](0) * w[j](1)
new_v += weight * g_v
new_C += 4 * weight * ti.outer_product(g_v, dpos) * inv_dx
v[f + 1, p] = new_v
x[f + 1, p] = x[f, p] + dt * v[f + 1, p]
C[f + 1, p] = new_C
def g2p(self, dt: ti.f32):
for p in self.x:
base = (self.x[p] * self.inv_dx - 0.5).cast(int)
fx = self.x[p] * self.inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
]
new_v = ti.Vector.zero(ti.f32, self.dim)
new_C = ti.Matrix.zero(ti.f32, self.dim, self.dim)
# loop over 3x3 grid node neighborhood
for I in ti.static(ti.grouped(self.stencil_range())):
dpos = I.cast(float) - fx
g_v = self.grid_v[base + I]
weight = 1.0
for d in ti.static(range(self.dim)):
weight *= w[I[d]][d]
new_v += weight * g_v
new_C += 4 * self.inv_dx * weight * ti.outer_product(g_v, dpos)
self.v[p], self.C[p] = new_v, new_C
self.x[p] += dt * self.v[p] # advection
def substep():
for i, j in grid_m:
grid_v[i, j] = [0, 0]
grid_m[i, j] = 0
for p in x:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1),
0.5 * ti.sqr(fx - 0.5)
]
F[p] = (ti.Matrix.identity(ti.f32, 2) + dt * C[p]) @ F[p]
h = ti.exp(10 * (1.0 - Jp[p]))
if material[p] == 0:
h = 0.3
mu, la = mu_0 * h, lambda_0 * h
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] == 1:
new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)
Jp[p] *= sig[d, d] / new_sig
sig[d, d] = new_sig
J *= new_sig
if material[p] == 1:
F[p] = U @ sig @ V.T()
stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T() + ti.Matrix.identity(
ti.f32, 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)):
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:
grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j]
grid_v[i, j][1] -= dt * 50
if i < 3 and grid_v[i, j][0] < 0: grid_v[i, j][0] = 0
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:
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0),
0.5 * ti.sqr(fx - 0.5)
]
new_v = ti.Vector.zero(ti.f32, 2)
new_C = ti.Matrix.zero(ti.f32, 2, 2)
for i, j in ti.static(ti.ndrange(3, 3)):
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 * ti.outer_product(g_v, dpos)
v[p], C[p] = new_v, new_C
x[p] += dt * v[p]
