def test_obey_kernel_simplicity():
x = ti.var(ti.f32)
y = ti.var(ti.f32)
@ti.layout
def place():
ti.root.dense(ti.i, 1).place(x, y)
ti.root.lazy_grad()
@ti.kernel
def func():
for i in x:
# OK: nested for loop
for j in ti.static(range(3)):
# OK: a series of non-for-loop statements
y[i] += x[i] * 42
y[i] -= x[i] * 5
y.grad[0] = 1.0
x[0] = 0.1
func()
func.grad()
assert x.grad[0] == approx((42 - 5) * 3)
def test_basic_utils():
a = ti.Vector(3, dt=ti.f32)
b = ti.Vector(2, dt=ti.f32)
abT = ti.Matrix.field(3, 2, dtype=ti.f32)
aNormalized = ti.Vector(3, dt=ti.f32)
normA = ti.field(ti.f32)
normSqrA = ti.field(ti.f32)
normInvA = ti.field(ti.f32)
ti.root.place(a, b, abT, aNormalized, normA, normSqrA, normInvA)
@ti.kernel
def init():
a[None] = ti.Vector([1.0, 2.0, 3.0])
b[None] = ti.Vector([4.0, 5.0])
abT[None] = a[None].outer_product(b[None])
normA[None] = a[None].norm()
normSqrA[None] = a[None].norm_sqr()
normInvA[None] = a[None].norm_inv()
aNormalized[None] = a[None].normalized()
init()
for i in range(3):
for j in range(2):
assert abT[None][i, j] == a[None][i] * b[None][j]
sqrt14 = np.sqrt(14.0)
invSqrt14 = 1.0 / sqrt14
assert normSqrA[None] == approx(14.0)
assert normInvA[None] == approx(invSqrt14)
assert normA[None] == approx(sqrt14)
assert aNormalized[None][0] == approx(1.0 * invSqrt14)
assert aNormalized[None][1] == approx(2.0 * invSqrt14)
assert aNormalized[None][2] == approx(3.0 * invSqrt14)
def test_float16():
dtype = ti.float16
x = ti.field(dtype, shape=())
x[None] = 0.3
print(x[None])
assert (x[None] == approx(0.3, rel=1e-3))
def test_snode_read_write():
dtype = ti.f16
x = ti.field(dtype, shape=())
x[None] = 0.3
print(x[None])
assert (x[None] == approx(0.3, rel=1e-3))
def test_to_numpy_as_vector_deprecated():
v = ti.Vector(3, dt=ti.f32, shape=(2))
u = np.array([[2, 3, 4], [5, 6, 7]])
v.from_numpy(u)
assert v.to_numpy(as_vector=True) == approx(u)
assert v.to_numpy() == approx(u)
def run_mpm88_test():
dim = 2
N = 64
n_particles = N * N
n_grid = 128
dx = 1 / n_grid
inv_dx = 1 / dx
dt = 2.0e-4
p_vol = (dx * 0.5)**2
p_rho = 1
p_mass = p_vol * p_rho
E = 400
x = ti.Vector.field(dim, dtype=ti.f32, shape=n_particles)
v = ti.Vector.field(dim, dtype=ti.f32, shape=n_particles)
C = ti.Matrix.field(dim, dim, dtype=ti.f32, shape=n_particles)
J = ti.field(dtype=ti.f32, shape=n_particles)
grid_v = ti.Vector.field(dim, dtype=ti.f32, shape=(n_grid, n_grid))
grid_m = ti.field(dtype=ti.f32, shape=(n_grid, n_grid))
@ti.kernel
def substep():
for p in x:
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]]) + p_mass * C[p]
for i in ti.static(range(3)):
for j in ti.static(range(3)):
offset = ti.Vector([i, j])
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 grid_m:
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 x:
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]).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 * 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
# gui = ti.core.GUI("MPM88", ti.veci(512, 512))
# canvas = gui.get_canvas()
for i in range(n_particles):
x[i] = [i % N / N * 0.4 + 0.2, i / N / N * 0.4 + 0.05]
v[i] = [0, -3]
J[i] = 1
for frame in range(10):
for s in range(50):
grid_v.fill([0, 0])
grid_m.fill(0)
substep()
pos = x.to_numpy()
pos[:, 1] *= 2
regression = [
0.31722742,
0.15826741,
0.10224003,
0.07810827,
]
for i in range(4):
assert (pos**(i + 1)).mean() == approx(regression[i], rel=1e-2)
assert ti.cfg.arch in [ti.cpu]
@ti.test(arch=[ti.cpu, ti.opengl],
require=[ti.extension.sparse, ti.extension.bls])
def test_require_extensions_2():
assert ti.cfg.arch in [ti.cuda]
### `ti.approx` and `ti.allclose`
@ti.test()
@pytest.mark.parametrize('x', [0.1, 3])
@pytest.mark.parametrize('allclose',
[ti.allclose, lambda x, y: x == ti.approx(y)])
def test_allclose_rel(x, allclose):
rel = 1e-3 if ti.cfg.arch == ti.opengl else 1e-6
assert not allclose(x + x * rel * 3.0, x)
assert not allclose(x + x * rel * 1.2, x)
assert allclose(x + x * rel * 0.9, x)
assert allclose(x + x * rel * 0.5, x)
assert allclose(x, x)
assert allclose(x - x * rel * 0.5, x)
assert allclose(x - x * rel * 0.9, x)
assert not allclose(x - x * rel * 1.2, x)
assert not allclose(x - x * rel * 3.0, x)
@pytest.mark.parametrize('x', [0.1, 3])
@ti.test()
def test_atomic_add_global_f32():
run_atomic_add_global_case(ti.f32,
4.2,
valproc=lambda x: approx(x, rel=1e-5))
def test_mpm88_numpy():
import numpy as np
dim = 2
N = 64
n_particles = N * N
n_grid = 128
dx = 1 / n_grid
inv_dx = 1 / dx
dt = 2.0e-4
p_vol = (dx * 0.5)**2
p_rho = 1
p_mass = p_vol * p_rho
E = 400
x = np.zeros((n_particles, dim), dtype=np.float32)
v = np.zeros((n_particles, dim), dtype=np.float32)
C = np.zeros((n_particles, dim, dim), dtype=np.float32)
J = np.zeros(n_particles, dtype=np.float32)
grid_v = np.zeros((n_grid, n_grid, dim), dtype=np.float32)
grid_m = np.zeros((n_grid, n_grid), dtype=np.float32)
@ti.kernel
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
for i in range(n_particles):
x[i] = [i % N / N * 0.4 + 0.2, i / N / N * 0.4 + 0.05]
v[i] = [0, -3]
J[i] = 1
for frame in range(10):
for s in range(50):
grid_v.fill(0)
grid_m.fill(0)
substep(x, v, C, J, grid_v, grid_m)
pos = x
pos[:, 1] *= 2
regression = [
0.31722742,
0.15826741,
0.10224003,
0.07810827,
]
for i in range(4):
assert (pos**(i + 1)).mean() == approx(regression[i], rel=1e-2)