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