def paint(): for i, j in ti.ndrange(n * 4, n * 4): # 4x4 super sampling: ret = ti.taichi_logo(ti.Vector([i, j]) / (n * 4)) x[i // 4, j // 4] += ret / 16
def func(a: ti.types.ndarray()): for i in a: for j, k in ti.ndrange(2, 2): a[i][j, k] = j * j + k * k
def func(arr_src: ti.template(), arr_dst: ti.template()): for i, j in ti.ndrange(*arr_src.shape): arr_dst[i, j] = arr_src[i, j]
def func(): for i in ti.ndrange(3, 4): print(i)
def func(): for i in ti.grouped(ti.ndrange(3, 4)): print(i)
def func(): for i, j in ti.ndrange((4, 10), (3, t)): val = i + j * 10 x[i, j] = val
def func(): for I in ti.static(ti.grouped(ti.ndrange((4, 5), (3, 5), 5))): x[I] = I[0] + I[1] * 10 + I[2] * 100
def populate(): for I in ti.grouped(ti.ndrange(*ndrange)): s = 0 for i in ti.static(range(dim)): s += I[i]**(i + 1) x[I] = s
def fill(Abuilder: ti.linalg.sparse_matrix_builder(), InputArray: ti.ext_arr(), b: ti.template()): for i, j in ti.ndrange(n, n): Abuilder[i, j] += InputArray[i, j] for i in range(n): b[i] = i + 1
def compute_Ax(self): for i, j in ti.ndrange((self.N_ext, self.N_tot - self.N_ext), (self.N_ext, self.N_tot - self.N_ext)): self.Ax[i, j] = (4.0 + self.offset) * self.x[i, j] - self.x[ i + 1, j] - self.x[i - 1, j] - self.x[i, j + 1] - self.x[i, j - 1]
def fill_m1(): for i, j in ti.ndrange(N, N): m1[i, j] = ti.random()
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 paint(): for i, j in ti.ndrange(n * 4, n * 4): ret = 1 - inside_taichi(Vector2(i / n / 4, j / n / 4)) x[i // 4, j // 4] += ret / 16
def verify(dx: ti.template(), dy: ti.template()): for i, j in ti.ndrange((boundary_offset, N - boundary_offset), (boundary_offset, N - boundary_offset)): assert y[i, j] == x[i + dx, j + dy]
def init(): for i, j in ti.ndrange(n, n): for k, l in ti.ndrange(n, n): r = i * n**3 + j * n**2 + k * n + l A[i, j, k, l] = r
def update_vertices(): for i, j in ti.ndrange(n, n): vertices[i * n + j] = x[i, j]
def func(): # `__getitem__ cannot be called from Python-scope` will be raised if # `a[None]` is not transformed to `ti.subscript(a, None)` in ti.ndrange: for i, j in ti.ndrange(a[None], b[None]): r = i * n + j + 1 A[i, j] = r
def march(self) -> int: r_n = 0 for I in ti.grouped(ti.ndrange(*[[1, self.N - 1]] * self.dim)): id = 0 if ti.static(self.dim == 2): i, j = I if self.m[i, j] > 1: id |= 1 if self.m[i + 1, j] > 1: id |= 2 if self.m[i, j + 1] > 1: id |= 4 if self.m[i + 1, j + 1] > 1: id |= 8 else: i, j, k = I if self.m[i, j, k] > 1: id |= 1 if self.m[i + 1, j, k] > 1: id |= 2 if self.m[i + 1, j + 1, k] > 1: id |= 4 if self.m[i, j + 1, k] > 1: id |= 8 if self.m[i, j, k + 1] > 1: id |= 16 if self.m[i + 1, j, k + 1] > 1: id |= 32 if self.m[i + 1, j + 1, k + 1] > 1: id |= 64 if self.m[i, j + 1, k + 1] > 1: id |= 128 for m in range(self.et.shape[1]): if self.et[id, m][0] == -1: break n = ti.atomic_add(r_n, 1) for l in ti.static(range(self.dim)): e = self.et[id, m][l] R = float(I) if ti.static(self.dim == 2): i, j = I # (.5, 0), (0, .5), (1, .5), (.5, 1) if e == 0: p = (1 - self.m[i, j]) / (self.m[i + 1, j] - self.m[i, j]) R = [i + p, j] elif e == 1: p = (1 - self.m[i, j]) / (self.m[i, j + 1] - self.m[i, j]) R = [i, j + p] elif e == 2: p = (1 - self.m[i + 1, j]) / ( self.m[i + 1, j + 1] - self.m[i + 1, j]) R = [i + 1, j + p] elif e == 3: p = (1 - self.m[i, j + 1]) / ( self.m[i + 1, j + 1] - self.m[i, j + 1]) R = [i + p, j + 1] else: i, j, k = I # (.5, 0, 0), (1, .5, 0), (.5, 1, 0), (0, .5, 0) # (.5, 0, 1), (1, .5, 1), (.5, 1, 1), (0, .5, 1) # (0, 0, .5), (1, 0, .5), (1, 1, .5), (0, 1, .5) if e == 0: p = (1 - self.m[i, j, k]) / (self.m[i + 1, j, k] - self.m[i, j, k]) R = [i + p, j, k] elif e == 1: p = (1 - self.m[i + 1, j, k]) / ( self.m[i + 1, j + 1, k] - self.m[i + 1, j, k]) R = [i + 1, j + p, k] elif e == 2: p = (1 - self.m[i, j + 1, k]) / ( self.m[i + 1, j + 1, k] - self.m[i, j + 1, k]) R = [i + p, j + 1, k] elif e == 3: p = (1 - self.m[i, j, k]) / (self.m[i, j + 1, k] - self.m[i, j, k]) R = [i, j + p, k] elif e == 4: p = (1 - self.m[i, j, k + 1]) / ( self.m[i + 1, j, k + 1] - self.m[i, j, k + 1]) R = [i + p, j, k + 1] elif e == 5: p = (1 - self.m[i + 1, j, k + 1]) / ( self.m[i + 1, j + 1, k + 1] - self.m[i + 1, j, k + 1]) R = [i + 1, j + p, k + 1] elif e == 6: p = (1 - self.m[i, j + 1, k + 1]) / ( self.m[i + 1, j + 1, k + 1] - self.m[i, j + 1, k + 1]) R = [i + p, j + 1, k + 1] elif e == 7: p = (1 - self.m[i, j, k + 1]) / ( self.m[i, j + 1, k + 1] - self.m[i, j, k + 1]) R = [i, j + p, k + 1] elif e == 8: p = (1 - self.m[i, j, k]) / (self.m[i, j, k + 1] - self.m[i, j, k]) R = [i, j, k + p] elif e == 9: p = (1 - self.m[i + 1, j, k]) / ( self.m[i + 1, j, k + 1] - self.m[i + 1, j, k]) R = [i + 1, j, k + p] elif e == 10: p = (1 - self.m[i + 1, j + 1, k]) / ( self.m[i + 1, j + 1, k + 1] - self.m[i + 1, j + 1, k]) R = [i + 1, j + 1, k + p] elif e == 11: p = (1 - self.m[i, j + 1, k]) / ( self.m[i, j + 1, k + 1] - self.m[i, j + 1, k]) R = [i, j + 1, k + p] self.r[n, l] = R return r_n
def func(): for i, j, k in ti.ndrange((4, 10), (3, 8), 17): x[i, j, k] = i + j * 10 + k * 100
def init(): for i, j in ti.ndrange((boundary_offset, N - boundary_offset), (boundary_offset, N - boundary_offset)): x[i, j] = ti.random(dtype=ti.i32) % 2
def func(): for i, j in ti.ndrange(16, 4): for I in ti.static(ti.grouped(ti.ndrange(2, 3))): x[i, j][I] = I[0] + I[1] * 10 + i + j * 4
def verify(): for i, j in ti.ndrange((boundary_offset, N - boundary_offset), (boundary_offset, N - boundary_offset)): assert y[i, j] == x[i, j]
def func(): for i, j, k in ti.ndrange(3, 4): print(i, j, k)
def func(): for i in ti.ndrange((4, 10)): x[i] = i
def func(): for i in ti.static(ti.ndrange(3, 4)): print(i)
def init(): for i, j in ti.ndrange(n, n): if i == j: A[i, j] = 1
def func(a: ti.types.ndarray()): for i in ti.grouped(a): for j, k in ti.ndrange(2, 2): a[i][j, k] = j * j
def init(): for i, j in ti.ndrange(n, n): if i // 2 == 0: A[i, j] = i
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 __iter__(self): for I in ti.grouped(ti.ndrange(*self.meta.shape)): yield I