def gridToParticles(self):
ti.block_dim(self.n_grid)
for p in self.x:
Xp = self.x[p] * self.inv_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
] # Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
new_V = ti.zero(self.v[p])
new_C = ti.zero(self.C[p])
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = self.real(1)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
g_v = self.grid_v[base + offset]
new_V += weight * g_v
new_C += 4 * self.inv_dx * weight * g_v.outer_product(dpos)
self.v[p] = new_V
self.C[p] = new_C
self.F[p] = (ti.Matrix.identity(self.real, self.dim) + self.dt *
self.C[p]) @ self.F[p] # F' = (I+dt * grad v)F
self.updateIsotropicHelper(p, self.F[p])
self.x[p] += self.dt * self.v[p]
def linearSolverReinitialize(self):
for I in self.mass_matrix:
self.r[I] = ti.zero(self.r[I])
self.p[I] = ti.zero(self.p[I])
self.q[I] = ti.zero(self.q[I])
self.temp[I] = ti.zero(self.temp[I])
self.step_direction[I] = ti.zero(self.step_direction[I])
def init_markers(self):
self.total_mk[None] = 0
for I in ti.grouped(self.cell_type):
if self.cell_type[I] == utils.FLUID:
for offset in ti.static(
ti.grouped(ti.ndrange(*((0, 2), ) * self.dim))):
num = ti.atomic_add(self.total_mk[None], 1)
self.p_x[num] = (
I +
(offset +
[ti.random()
for _ in ti.static(range(self.dim))]) / 2) * self.dx
self.p_v[num] = ti.zero(self.p_v[num])
for k in ti.static(range(self.dim)):
self.p_cp[k][num] = ti.zero(self.p_cp[k][num])
def refract(x, n, eta):
"""Calculate the refraction direction for an incident vector.
This function is equivalent to the `refract` function in GLSL.
Args:
x (:class:`~taichi.Matrix`): The incident vector.
n (:class:`~taichi.Matrix`): The normal vector.
eta (float): The ratio of indices of refraction.
Returns:
:class:`~taichi.Matrix`: The refraction direction vector.
Example::
>>> x = ti.Vector([1., 1., 1.])
>>> n = ti.Vector([0, 1., 0])
>>> refract(x, y, 2.0)
[2., -1., 2]
"""
dxn = x.dot(n)
result = ti.zero(x)
k = 1.0 - eta * eta * (1.0 - dxn * dxn)
if k >= 0.0:
result = eta * x - (eta * dxn + ti.sqrt(k)) * n
return result
def calc_kappa(self, d: ti.template(), phi: ti.template()):
for I in ti.grouped(self.kappa):
self.kappa[I] = ti.zero(self.kappa[I])
for k in ti.static(range(self.dim)):
unit = ti.Vector.unit(self.dim, k)
self.kappa[I] += (utils.sample(self.n, I + unit * 0.5) - \
utils.sample(self.n, I - unit * 0.5))[k] / self.dx
def substep():
for I in ti.grouped(F_grid_m):
F_grid_v[I] = ti.zero(F_grid_v[I])
F_grid_m[I] = 0
ti.block_dim(n_grid)
for p in F_x:
Xp = F_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 * (F_J[p] - 1) / dx**2
affine = ti.Matrix.identity(float, dim) * stress + p_mass * F_C[p]
for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= w[offset[i]][i]
F_grid_v[base +
offset] += weight * (p_mass * F_v[p] + affine @ dpos)
F_grid_m[base + offset] += weight * p_mass
for I in ti.grouped(F_grid_m):
if F_grid_m[I] > 0:
F_grid_v[I] /= F_grid_m[I]
F_grid_v[I][1] -= dt * gravity
cond = (I < bound) & (F_grid_v[I] < 0) | \
(I > n_grid - bound) & (F_grid_v[I] > 0)
F_grid_v[I] = 0 if cond else F_grid_v[I]
ti.block_dim(n_grid)
for p in F_x:
Xp = F_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.zero(F_v[p])
new_C = ti.zero(F_C[p])
for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= w[offset[i]][i]
g_v = F_grid_v[base + offset]
new_v += weight * g_v
new_C += 4 * weight * g_v.outer_product(dpos) / dx**2
F_v[p] = new_v
F_x[p] += dt * F_v[p]
F_J[p] *= 1 + dt * new_C.trace()
F_C[p] = new_C
def substep():
for I in ti.grouped(grid_m):
grid_v[I] = ti.zero(grid_v[I])
grid_m[I] = 0
ti.block_dim(n_grid)
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.identity(float, dim) * stress + p_mass * C[p]
for offset in ti.grouped(ti.ndrange(*neighbour)):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= list_subscript(w, offset[i])[i]
grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
grid_m[base + offset] += weight * p_mass
for I in ti.grouped(grid_m):
if grid_m[I] > 0:
grid_v[I] /= grid_m[I]
grid_v[I][1] -= dt * gravity
cond = I < bound and grid_v[I] < 0 or I > n_grid - bound and grid_v[
I] > 0
grid_v[I] = 0 if cond else grid_v[I]
ti.block_dim(n_grid)
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.zero(v[p])
new_C = ti.zero(C[p])
for offset in ti.grouped(ti.ndrange(*neighbour)):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= list_subscript(w, offset[i])[i]
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
def update_magnetic_field(self):
for k in ti.static(range(self.dim)):
offset = ti.Vector.unit(self.dim, k)
for I in ti.grouped(self.H[k]):
self.H[k][I] = ti.zero(self.H[k][I])
if I[k] - 1 >= 0 and I[k] < self.res[k]:
self.H[k][I] -= (self.psi[I] -
self.psi[I - offset]) / self.dx
self.H[k][I] += self.H_ext[k][I] # H = H_ext - grad psi
def multiply(self, x: ti.template(), b: ti.template()):
for I in b:
b[I] = ti.zero(b[I])
# Note the relationship H dx = - df, where H is the stiffness matrix
# inertia part
for I in x:
b[I] += self.mass_matrix[I] * x[I]
self.computeDvAndGradDv(x)
# scratch_gradV is now temporaraly used for storing gradDV (evaluated at particles)
# scratch_vp is now temporaraly used for storing DV (evaluated at particles)
for p in self.x:
self.scratch_stress[p] = ti.zero(self.scratch_stress[p])
for p in self.x:
self.computeStressDifferential(p, self.scratch_gradV[p],
self.scratch_stress[p],
self.scratch_vp[p])
# scratch_stress is now V_p^0 dP (F_p^n)^T (dP is Ap in snow paper)
ti.block_dim(self.n_grid)
for p in self.x:
Xp = self.x[p] * self.inv_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
] # Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
stress = self.scratch_stress[p]
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = self.real(1)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
b[self.idx(base + offset)] += self.dt * self.dt * (
weight * stress @ dpos
) # fi -= \sum_p (Ap (xi-xp) - fp )w_ip Dp_inv
def computeDvAndGradDv(self, dv: ti.template()):
for p in self.x:
Xp = self.x[p] * self.inv_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]
vp = ti.zero(self.scratch_vp[p])
gradV = ti.zero(self.scratch_gradV[p])
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = ti.cast(1.0, self.real)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
dv0 = dv[self.idx(base + offset)]
vp += weight * dv0
gradV += 4 * self.inv_dx * weight * dv0.outer_product(dpos)
self.scratch_vp[p] = vp
self.scratch_gradV[p] = gradV
def distance_of_aabb(self, x, x0, x1):
phi = ti.cast(0, self.real)
if all(x > x0) and all(x < x1): # (inside)
phi = (ti.max(x0 - x, x - x1)).max()
else: # (outside)
# Find the closest point (p,q,r) on the surface of the box
p = ti.zero(x)
for k in ti.static(range(self.dim)):
if x[k] < x0[k]: p[k] = x0[k]
elif x[k] > x1[k]: p[k] = x1[k]
else: p[k] = x[k]
phi = (x - p).norm()
return phi
def splat(data, weights, v, pos, cp):
tot = data.shape
dim = len(tot)
I0 = ti.Vector.zero(int, len(tot))
I1 = ti.Vector.zero(int, len(tot))
w = ti.zero(pos)
for k in ti.static(range(len(tot))):
I0[k] = ts.clamp(int(pos[k]), 0, tot[k] - 1)
I1[k] = ts.clamp(I0[k] + 1, 0, tot[k] - 1)
w[k] = ts.clamp(pos[k] - I0[k], 0.0, 1.0)
for u in ti.static(ti.grouped(ti.ndrange(*((0, 2), ) * len(tot)))):
dpos = ti.zero(pos)
I = ti.Vector.zero(int, len(tot))
W = 1.0
for k in ti.static(range(len(tot))):
dpos[k] = (pos[k] - I0[k]) if u[k] == 0 else (pos[k] - I1[k])
I[k] = I0[k] if u[k] == 0 else I1[k]
W *= (1 - w[k]) if u[k] == 0 else w[k]
data[I] += (v + cp.dot(dpos)) * W
weights[I] += W
def build_b_kernel(self, cell_type: ti.template(), b: ti.template()):
for I in ti.grouped(cell_type): # solve the entire domain tao
if all(I == 0):
b[I] = ti.zero(
b[I]) # choose a reference point, where psi(p)=0
elif self.in_domain(I):
for k in ti.static(range(self.dim)):
offset = 0.5 * (1 - ti.Vector.unit(self.dim, k))
unit = ti.Vector.unit(self.dim, k)
if self.in_domain(I + unit): # magnetic shielding
b[I] += utils.sample(self.chi, I + unit +
offset) * self.H_ext[k][I + unit]
if self.in_domain(I - unit):
b[I] -= utils.sample(self.chi,
I + offset) * self.H_ext[k][I]
def initialize(self):
for I in ti.grouped(ti.ndrange(*[self.res[_]
for _ in range(self.dim)])):
self.r[0][I] = 0
self.z[0][I] = 0
self.Ap[I] = 0
self.p[I] = 0
self.x[I] = 0
self.b[I] = 0
for l in ti.static(range(self.n_mg_levels)):
for I in ti.grouped(
ti.ndrange(
*[self.res[_] // (2**l) for _ in range(self.dim)])):
self.grid_type[l][I] = 0
self.Adiag[l][I] = 0
self.Ax[l][I] = ti.zero(self.Ax[l][I])
def makePD2d(self, M: ti.template()):
a = M[0, 0]
b = (M[0, 1] + M[1, 0]) / 2
d = M[1, 1]
b2 = b * b
D = a * d - b2
T_div_2 = (a + d) / 2
sqrtTT4D = ti.sqrt(ti.abs(T_div_2 * T_div_2 - D))
L2 = T_div_2 - sqrtTT4D
if L2 < 0.0:
L1 = T_div_2 + sqrtTT4D
if L1 <= 0.0: M = ti.zero(M)
else:
if b2 == 0: M = ti.Matrix([[L1, 0], [0, 0]])
else:
L1md = L1 - d
L1md_div_L1 = L1md / L1
M = ti.Matrix([[L1md_div_L1 * L1md, b * L1md_div_L1],
[b * L1md_div_L1, b2 / L1]])
def updateState(self):
ti.block_dim(self.n_grid)
for p in self.x:
Xp = self.x[p] * self.inv_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_C = ti.zero(self.C[p])
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = ti.cast(1.0, self.real)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
g_v = self.grid_v[base + offset] + self.dv[self.idx(base +
offset)]
new_C += 4 * self.inv_dx * weight * g_v.outer_product(dpos)
self.F[p] = (ti.Matrix.identity(self.real, self.dim) +
self.dt * new_C) @ self.old_F[p]
self.updateIsotropicHelper(p, self.F[p])
self.scratch_xp[p] = self.x[p] + self.dt * self.scratch_vp[p]
def computeResidual(self):
for I in self.dv:
self.residual[I] = self.dt * self.mass_matrix[I] * self.gravity
for I in self.dv:
self.residual[I] -= self.mass_matrix[I] * self.dv[I]
ti.block_dim(self.n_grid)
for p in self.x:
Xp = self.x[p] * self.inv_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_C = ti.zero(self.C[p])
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = ti.cast(1.0, self.real)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
g_v = self.grid_v[base + offset] + self.dv[self.idx(base +
offset)]
new_C += 4 * self.inv_dx * weight * g_v.outer_product(dpos)
F = (ti.Matrix.identity(self.real, self.dim) +
self.dt * new_C) @ self.old_F[p]
stress = (-self.p_vol * 4 * self.inv_dx *
self.inv_dx) * self.dpsi_dF(F) @ F.transpose()
for offset in ti.static(ti.grouped(ti.ndrange(*self.neighbour))):
dpos = (offset - fx) * self.dx
weight = ti.cast(1.0, self.real)
for i in ti.static(range(self.dim)):
weight *= w[offset[i]][i]
force = weight * stress @ dpos
self.residual[self.idx(base + offset)] += self.dt * force
self.project(self.residual)
def reinitializeIsotropicHelper(self, p):
if ti.static(self.dim == 2):
self.psi0[p] = 0 # d_PsiHat_d_sigma0
self.psi1[p] = 0 # d_PsiHat_d_sigma1
self.psi00[p] = 0 # d^2_PsiHat_d_sigma0_d_sigma0
self.psi01[p] = 0 # d^2_PsiHat_d_sigma0_d_sigma1
self.psi11[p] = 0 # d^2_PsiHat_d_sigma1_d_sigma1
self.m01[
p] = 0 # (psi0-psi1)/(sigma0-sigma1), usually can be computed robustly
self.p01[
p] = 0 # (psi0+psi1)/(sigma0+sigma1), need to clamp bottom with 1e-6
self.Aij[p] = ti.zero(self.Aij[p])
self.B01[p] = ti.zero(self.B01[p])
if ti.static(self.dim == 3):
self.psi0[p] = 0 # d_PsiHat_d_sigma0
self.psi1[p] = 0 # d_PsiHat_d_sigma1
self.psi2[p] = 0 # d_PsiHat_d_sigma2
self.psi00[p] = 0 # d^2_PsiHat_d_sigma0_d_sigma0
self.psi11[p] = 0 # d^2_PsiHat_d_sigma1_d_sigma1
self.psi22[p] = 0 # d^2_PsiHat_d_sigma2_d_sigma2
self.psi01[p] = 0 # d^2_PsiHat_d_sigma0_d_sigma1
self.psi02[p] = 0 # d^2_PsiHat_d_sigma0_d_sigma2
self.psi12[p] = 0 # d^2_PsiHat_d_sigma1_d_sigma2
self.m01[
p] = 0 # (psi0-psi1)/(sigma0-sigma1), usually can be computed robustly
self.p01[
p] = 0 # (psi0+psi1)/(sigma0+sigma1), need to clamp bottom with 1e-6
self.m02[
p] = 0 # (psi0-psi2)/(sigma0-sigma2), usually can be computed robustly
self.p02[
p] = 0 # (psi0+psi2)/(sigma0+sigma2), need to clamp bottom with 1e-6
self.m12[
p] = 0 # (psi1-psi2)/(sigma1-sigma2), usually can be computed robustly
self.p12[
p] = 0 # (psi1+psi2)/(sigma1+sigma2), need to clamp bottom with 1e-6
self.Aij[p] = ti.zero(self.Aij[p])
self.B01[p] = ti.zero(self.B01[p])
self.B12[p] = ti.zero(self.B12[p])
self.B20[p] = ti.zero(self.B20[p])
def reduce_tmp() -> dtype:
s = ti.zero(tot[None]) if op == OP_ADD or op == OP_XOR else a[0]
for i in a:
ti_op(s, a[i])
return s
def substep(g_x: float, g_y: float, g_z: float):
for I in ti.grouped(grid_m):
grid_v[I] = ti.zero(grid_v[I])
grid_m[I] = 0
ti.block_dim(n_grid)
for p in x:
if used[p] == 0:
continue
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]
F[p] = (ti.Matrix.identity(float, 3) +
dt * C[p]) @ F[p] # deformation gradient update
h = ti.exp(
10 *
(1.0 -
Jp[p])) # Hardening coefficient: snow gets harder when compressed
if materials[p] == JELLY: # jelly, make it softer
h = 0.3
mu, la = mu_0 * h, lambda_0 * h
if materials[p] == WATER: # liquid
mu = 0.0
U, sig, V = ti.svd(F[p])
J = 1.0
for d in ti.static(range(3)):
new_sig = sig[d, d]
if materials[p] == SNOW: # 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 materials[
p] == WATER: # Reset deformation gradient to avoid numerical instability
new_F = ti.Matrix.identity(float, 3)
new_F[0, 0] = J
F[p] = new_F
elif materials[p] == SNOW:
F[p] = U @ sig @ V.transpose(
) # Reconstruct elastic deformation gradient after plasticity
stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose(
) + ti.Matrix.identity(float, 3) * la * J * (J - 1)
stress = (-dt * p_vol * 4) * stress / dx**2
affine = stress + p_mass * C[p]
for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= w[offset[i]][i]
grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
grid_m[base + offset] += weight * p_mass
for I in ti.grouped(grid_m):
if grid_m[I] > 0:
grid_v[I] /= grid_m[I]
grid_v[I] += dt * ti.Vector([g_x, g_y, g_z])
cond = I < bound and grid_v[I] < 0 or I > n_grid - bound and grid_v[
I] > 0
grid_v[I] = 0 if cond else grid_v[I]
ti.block_dim(n_grid)
for p in x:
if used[p] == 0:
continue
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.zero(v[p])
new_C = ti.zero(C[p])
for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
dpos = (offset - fx) * dx
weight = 1.0
for i in ti.static(range(dim)):
weight *= w[offset[i]][i]
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]
C[p] = new_C
def reduce_tmp() -> dtype:
s = ti.zero(tot[None])
for i in a:
s += a[i]
return s
def zero(x: ti.template()):
for I in ti.grouped(x):
x[I] = ti.zero(x[I])
def project(self, x: ti.template()):
for p in x:
I = self.node(p)
cond = any(I < self.bound and self.grid_v[I] < 0) or any(
I > self.n_grid - self.bound and self.grid_v[I] > 0)
if cond: x[p] = ti.zero(x[p])
def reduce_tmp() -> dtype:
s = ti.zero(tot[None])
for i in a:
ti_op(s, a[i])
return s
