def stencil_2d_array(y: ti.any_arr(), x: ti.any_arr()):
for I in ti.grouped(x):
if ti.static(scatter):
for offset in ti.static(stencil_common):
y[I + ti.Vector(offset)] += x[I]
else: # gather
s = ti.cast(0.0, dtype)
for offset in ti.static(stencil_common):
s = s + x[I + ti.Vector(offset)]
y[I] = s
def substep_update_grid_v(grid_v: ti.any_arr(field_dim=2),
grid_m: ti.any_arr(field_dim=2)):
for i, j in grid_m:
if grid_m[i, j] > 0:
grid_v[i, j] /= grid_m[i, j]
grid_v[i, j].y -= dt * gravity
if i < bound and grid_v[i, j].x < 0:
grid_v[i, j].x = 0
if i > n_grid - bound and grid_v[i, j].x > 0:
grid_v[i, j].x = 0
if j < bound and grid_v[i, j].y < 0:
grid_v[i, j].y = 0
if j > n_grid - bound and grid_v[i, j].y > 0:
grid_v[i, j].y = 0
def init(d: ti.i32, density1: ti.any_arr(), density2: ti.any_arr(),
density3: ti.any_arr(), density4: ti.any_arr(),
density5: ti.any_arr(), density6: ti.any_arr(),
density7: ti.any_arr(), density8: ti.any_arr()):
for i, j in density1:
density1[i, j] = d + 1
density2[i, j] = d + 2
density3[i, j] = d + 3
density4[i, j] = d + 4
density5[i, j] = d + 5
density6[i, j] = d + 6
density7[i, j] = d + 7
density8[i, j] = d + 8
def substep_g2p(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1),
C: ti.any_arr(field_dim=1), J: ti.any_arr(field_dim=1),
grid_v: ti.any_arr(field_dim=2)):
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.Vector.zero(float, 2)
new_C = ti.Matrix.zero(float, 2, 2)
for i, j in ti.static(ti.ndrange(3, 3)):
offset = ti.Vector([i, j])
dpos = (offset - fx) * dx
weight = w[i].x * w[j].y
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 substep(x: ti.any_arr(element_shape=(dim, )),
v: ti.any_arr(element_shape=(dim, )),
C: ti.any_arr(element_shape=(dim, dim)), J: ti.any_arr(),
grid_v: ti.any_arr(element_shape=(dim, )),
grid_m: ti.any_arr()):
for p in range(n_particles):
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
stress = -dt * p_vol * (J[p] - 1) * 4 * inv_dx * inv_dx * E
affine = ti.Matrix([[stress, 0], [0, stress]],
dt=ti.f32) + p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j], dt=ti.i32)
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 ti.ndrange(n_grid, n_grid):
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 range(n_particles):
base = (x[p] * inv_dx - 0.5).cast(int)
fx = x[p] * inv_dx - base.cast(float)
w = [
0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2
]
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], dt=ti.i32).cast(float) - fx
g_v = grid_v[base + ti.Vector([i, j], dt=ti.i32)]
weight = w[i][0] * w[j][1]
new_v += weight * g_v
new_C += 4 * weight * g_v.outer_product(dpos) * inv_dx
v[p] = new_v
x[p] += dt * v[p]
J[p] *= 1 + dt * new_C.trace()
C[p] = new_C
def substep_p2g(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1),
C: ti.any_arr(field_dim=1), J: ti.any_arr(field_dim=1),
grid_v: ti.any_arr(field_dim=2),
grid_m: ti.any_arr(field_dim=2)):
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([[stress, 0], [0, stress]]) + p_mass * C[p]
for i, j in ti.static(ti.ndrange(3, 3)):
offset = ti.Vector([i, j])
dpos = (offset - fx) * dx
weight = w[i].x * w[j].y
grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
grid_m[base + offset] += weight * p_mass
def run(x: ti.any_arr(), y: ti.any_arr()):
for i in range(n):
for j in range(m):
x[i, j] += i + j + y[i, j]
def fill_2d_array(dst: ti.any_arr()):
for i, j in dst:
dst[i, j] = ti.random(dtype)
def fill_1d_array(dst: ti.any_arr()):
for i in dst:
dst[i] = ti.random(dtype)
def init(d: ti.i32, arr: ti.any_arr()):
for i, j in arr:
arr[i, j] = d
def func2(a: ti.any_arr()):
for I in ti.grouped(a):
a[I] = I.sum()
def func(a: ti.any_arr()):
for i in a:
for j, k in ti.ndrange(2, 2):
a[i][j, k] = j * j + k * k
def reduce(a: ti.any_arr()) -> ti.i32:
s = 0
for i in a:
ti.atomic_add(s, a[i])
ti.atomic_sub(s, 2)
return s
def func4(a: ti.any_arr(layout=ti.Layout.SOA)):
pass
def func3(a: ti.any_arr(layout=ti.Layout.AOS)):
pass
def func2(a: ti.any_arr(element_dim=2)):
pass
def func1(a: ti.any_arr(element_dim=1)):
pass
def func(a: ti.any_arr(element_dim=1)):
for i in range(5):
for j in range(4):
a[i][j * j] = j * j
def func(a: ti.any_arr()):
for i in range(5):
for j in range(4):
a[i][j * j] = j * j
def substep_reset_grid(grid_v: ti.any_arr(field_dim=2),
grid_m: ti.any_arr(field_dim=2)):
for i, j in grid_m:
grid_v[i, j] = [0, 0]
grid_m[i, j] = 0
def func5(a: ti.any_arr(element_shape=(2, 3))):
pass
def run(x: ti.any_arr(), y: ti.any_arr()):
for i in range(n):
x[i] += i + y[i]
def run(x: ti.any_arr(element_dim=2, layout=ti.Layout.AOS),
y: ti.any_arr()):
for i in ti.static(range(n)):
for j in ti.static(range(m)):
x[i, j][0, 0] += i + j + y[i, j]
def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1),
J: ti.any_arr(field_dim=1)):
for i in range(n_particles):
x[i] = [ti.random() * 0.4 + 0.2, ti.random() * 0.4 + 0.2]
v[i] = [0, -1]
J[i] = 1
def func6(a: ti.any_arr(element_dim=1, element_shape=(2, 3))):
pass
def func2(b: ti.any_arr(element_dim=1, layout=ti.Layout.SOA)):
for i in range(5):
for j in range(4):
b[i][j * j] = j * j
def func1(a: ti.any_arr(element_dim=2)):
for i in range(5):
for j, k in ti.ndrange(2, 2):
a[i][j, k] = j * j + k * k
def func7(a: ti.any_arr(field_dim=2)):
pass
def any_array(x: ti.any_arr()):
b = x[3, 1.1]
def init(x: ti.any_arr()):
for i, j in x:
x[i, j] = i + j
