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
