def _test_polar_decomp(dim, dt):
m = ti.Matrix.field(dim, dim, dt)
r = ti.Matrix.field(dim, dim, dt)
s = ti.Matrix.field(dim, dim, dt)
I = ti.Matrix.field(dim, dim, dt)
D = ti.Matrix.field(dim, dim, dt)
ti.root.place(m, r, s, I, D)
@ti.kernel
def polar():
R, S = ti.polar_decompose(m[None], dt)
r[None] = R
s[None] = S
m[None] = R @ S
I[None] = R @ R.transpose()
D[None] = S - S.transpose()
def V(i, j):
return i * 2 + j * 7 + int(i == j) * 3
for i in range(dim):
for j in range(dim):
m[None][i, j] = V(i, j)
polar()
tol = 5e-5 if dt == ti.f32 else 1e-12
for i in range(dim):
for j in range(dim):
assert m[None][i, j] == approx(V(i, j), abs=tol)
assert I[None][i, j] == approx(int(i == j), abs=tol)
assert D[None][i, j] == approx(0, abs=tol)
def test_matrix_factories():
a = ti.Vector.field(3, dtype=ti.i32, shape=3)
b = ti.Matrix.field(2, 2, dtype=ti.f32, shape=2)
c = ti.Matrix.field(2, 3, dtype=ti.f32, shape=2)
@ti.kernel
def fill():
b[0] = ti.Matrix.identity(ti.f32, 2)
b[1] = ti.Matrix.rotation2d(math.pi / 3)
c[0] = ti.Matrix.zero(ti.f32, 2, 3)
c[1] = ti.Matrix.one(ti.f32, 2, 3)
for i in ti.static(range(3)):
a[i] = ti.Vector.unit(3, i)
fill()
for i in range(3):
for j in range(3):
assert a[i][j] == int(i == j)
sqrt3o2 = math.sqrt(3) / 2
assert b[0].value.to_numpy() == approx(np.eye(2))
assert b[1].value.to_numpy() == approx(
np.array([[0.5, -sqrt3o2], [sqrt3o2, 0.5]]))
assert c[0].value.to_numpy() == approx(np.zeros((2, 3)))
assert c[1].value.to_numpy() == approx(np.ones((2, 3)))
def test_precision():
u = ti.field(ti.f64, shape=())
v = ti.field(ti.f64, shape=())
w = ti.field(ti.f64, shape=())
@ti.kernel
def forward():
v[None] = ti.sqrt(ti.cast(u[None] + 3.25, ti.f64))
w[None] = ti.cast(u[None] + 7, ti.f64) / ti.cast(u[None] + 3, ti.f64)
forward()
assert v[None]**2 == approx(3.25, abs=1e-12)
assert w[None] * 3 == approx(7, abs=1e-12)
def test_unary_op():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo():
x[None] = -y[None]
x[None] = ti.floor(x[None])
y[None] = ti.ceil(y[None])
y[None] = -1.4
foo()
assert (x[None] == approx(1, rel=1e-3))
assert (y[None] == approx(-1, rel=1e-3))
def test_ad_frac():
@ti.func
def frac(x):
fractional = x - ti.floor(x) if x > 0. else x - ti.ceil(x)
return fractional
@ti.kernel
def ti_frac(input_field: ti.template(), output_field: ti.template()):
for i in input_field:
output_field[i] = frac(input_field[i])**2
@ti.kernel
def calc_loss(input_field: ti.template(), loss: ti.template()):
for i in input_field:
loss[None] += input_field[i]
n = 10
field0 = ti.field(dtype=ti.f32, shape=(n, ), needs_grad=True)
randoms = np.random.randn(10).astype(np.float32)
field0.from_numpy(randoms)
field1 = ti.field(dtype=ti.f32, shape=(n, ), needs_grad=True)
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
with ti.Tape(loss):
ti_frac(field0, field1)
calc_loss(field1, loss)
grads = field0.grad.to_numpy()
expected = np.modf(randoms)[0] * 2
for i in range(n):
assert grads[i] == approx(expected[i], rel=1e-4)
def _test_svd(dt, n):
print(
f'arch={ti.cfg.arch} default_fp={ti.cfg.default_fp} fast_math={ti.cfg.fast_math} dim={n}'
)
A = ti.Matrix.field(n, n, dtype=dt, shape=())
A_reconstructed = ti.Matrix.field(n, n, dtype=dt, shape=())
U = ti.Matrix.field(n, n, dtype=dt, shape=())
UtU = ti.Matrix.field(n, n, dtype=dt, shape=())
sigma = ti.Matrix.field(n, n, dtype=dt, shape=())
V = ti.Matrix.field(n, n, dtype=dt, shape=())
VtV = ti.Matrix.field(n, n, dtype=dt, shape=())
@ti.kernel
def run():
U[None], sigma[None], V[None] = ti.svd(A[None], dt)
UtU[None] = U[None].transpose() @ U[None]
VtV[None] = V[None].transpose() @ V[None]
A_reconstructed[None] = U[None] @ sigma[None] @ V[None].transpose()
if n == 3:
A[None] = [[1, 1, 3], [9, -3, 2], [-3, 4, 2]]
else:
A[None] = [[1, 1], [2, 3]]
run()
tol = 1e-5 if dt == ti.f32 else 1e-12
assert mat_equal(UtU.to_numpy(), np.eye(n), tol=tol)
assert mat_equal(VtV.to_numpy(), np.eye(n), tol=tol)
assert mat_equal(A_reconstructed.to_numpy(), A.to_numpy(), tol=tol)
for i in range(n):
for j in range(n):
if i != j:
assert sigma[None][i, j] == approx(0)
def test_scalar_argument():
@ti.kernel
def add(a: ti.f32, b: ti.f32) -> ti.f32:
a = a + b
return a
assert add(1.0, 2.0) == approx(3.0)
def test_random_vector_dup_eval():
a = ti.Vector.field(2, ti.f32, ())
@ti.kernel
def func():
a[None] = ti.Vector([ti.random(), 1]).normalized()
for i in range(4):
func()
assert a[None].value.norm_sqr() == approx(1)
def test_grouped_static_for_cast():
@ti.kernel
def foo() -> ti.f32:
ret = 0.
for I in ti.static(ti.grouped(ti.ndrange((4, 5), (3, 5), 5))):
tmp = I.cast(float)
ret += tmp[2] / 2
return ret
assert foo() == approx(10)
def test_func_random_argument_dup_eval():
@ti.func
def func(a):
return ti.Vector([ti.cos(a), ti.sin(a)])
@ti.kernel
def kern() -> ti.f32:
return func(ti.random()).norm_sqr()
for i in range(4):
assert kern() == approx(1.0, rel=5e-5)
def impl():
print(f'arch={ti.cfg.arch} default_fp={ti.cfg.default_fp}')
x = ti.field(default_fp)
y = ti.field(default_fp)
ti.root.dense(ti.i, 1).place(x, x.grad, y, y.grad)
@ti.kernel
def func():
for i in x:
y[i] = tifunc(x[i])
v = 0.234
y.grad[0] = 1
x[0] = v
func()
func.grad()
assert y[0] == approx(npfunc(v))
assert x.grad[0] == approx(grad(npfunc)(v))
def test_extra_unary_promote():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo():
x[None] = abs(y[None])
y[None] = -0.3
foo()
assert (x[None] == approx(0.3, rel=1e-3))
def test_arg_f16():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo(a: ti.f16):
x[None] = y[None] + a
y[None] = -0.3
foo(1.2)
assert (x[None] == approx(0.9, rel=1e-3))
def test_const_func_ret():
ti.init()
@ti.kernel
def func1() -> ti.f32:
return 3
@ti.kernel
def func2() -> ti.i32:
return 3.3 # return type mismatch, will be auto-casted into ti.i32
assert func1() == approx(3)
assert func2() == 3
def test_binary_extra_promote():
x = ti.field(dtype=ti.f16, shape=())
y = ti.field(dtype=ti.f16, shape=())
z = ti.field(dtype=ti.f16, shape=())
@ti.kernel
def foo():
y[None] = x[None]**2
z[None] = ti.atan2(y[None], 0.3)
x[None] = 0.1
foo()
assert (z[None] == approx(math.atan2(0.1**2, 0.3), rel=1e-3))
def grad_test(tifunc, npfunc=None):
npfunc = npfunc or tifunc
print(f'arch={ti.cfg.arch} default_fp={ti.cfg.default_fp}')
x = ti.field(ti.cfg.default_fp)
y = ti.field(ti.cfg.default_fp)
ti.root.dense(ti.i, 1).place(x, x.grad, y, y.grad)
@ti.kernel
def func():
for i in x:
y[i] = tifunc(x[i])
v = 0.234
y.grad[0] = 1
x[0] = v
func()
func.grad()
assert y[0] == approx(npfunc(v), rel=1e-4)
assert x.grad[0] == approx(grad(npfunc)(v), rel=1e-4)
def test_ad_reduce():
N = 16
x = ti.field(dtype=ti.f32, shape=N, needs_grad=True)
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
@ti.kernel
def func():
for i in x:
loss[None] += x[i]**2
total_loss = 0
for i in range(N):
x[i] = i
total_loss += i * i
loss.grad[None] = 1
func()
func.grad()
assert total_loss == approx(loss[None])
for i in range(N):
assert x.grad[i] == approx(i * 2)
def test_random_independent_product():
n = 1024
x = ti.field(ti.f32, shape=n * n)
@ti.kernel
def fill():
for i in range(n * n):
a = ti.random()
b = ti.random()
x[i] = a * b
fill()
X = x.to_numpy()
for i in range(4):
assert X.mean() == approx(1 / 4, rel=1e-2)
def test_diag():
m1 = ti.Matrix.field(3, 3, dtype=ti.f32, shape=())
@ti.kernel
def fill():
m1[None] = ti.Matrix.diag(dim=3, val=1.4)
fill()
for i in range(3):
for j in range(3):
if i == j:
assert m1[None][i, j] == approx(1.4)
else:
assert m1[None][i, j] == 0.0
def test_random_float():
for precision in [ti.f32, ti.f64]:
ti.init()
n = 1024
x = ti.field(ti.f32, shape=(n, n))
@ti.kernel
def fill():
for i in range(n):
for j in range(n):
x[i, j] = ti.random(precision)
fill()
X = x.to_numpy()
for i in range(1, 4):
assert (X**i).mean() == approx(1 / (i + 1), rel=1e-2)
def test_atomic_min_f16():
f = ti.field(dtype=ti.f16, shape=(2))
@ti.kernel
def foo():
# Parallel min
for i in range(1000):
ti.atomic_min(f[0], -3.13 * i)
# Serial min
for _ in range(1):
for i in range(1000):
f[1] = ti.min(-3.13 * i, f[1])
foo()
assert (f[0] == approx(f[1], rel=1e-3))
def test_atomic_add_f16():
f = ti.field(dtype=ti.f16, shape=(2))
@ti.kernel
def foo():
# Parallel sum
for i in range(1000):
f[0] += 1.12
# Serial sum
for _ in range(1):
for i in range(1000):
f[1] = f[1] + 1.12
foo()
assert (f[0] == approx(f[1], rel=1e-3))
def test_binary_op():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
z = ti.field(dtype, shape=())
@ti.kernel
def add():
x[None] = y[None] + z[None]
x[None] = x[None] * z[None]
y[None] = 0.2
z[None] = 0.72
add()
u = x.to_numpy()
assert (u[None] == approx(0.6624, rel=1e-3))
def _test_binary_func_ret(dt1, dt2, dt3, castor):
@ti.kernel
def func(a: dt1, b: dt2) -> dt3:
return a * b
if ti.types.is_integral(dt1):
xs = list(range(4))
else:
xs = [0.2, 0.4, 0.8, 1.0]
if ti.types.is_integral(dt2):
ys = list(range(4))
else:
ys = [0.2, 0.4, 0.8, 1.0]
for x, y in zip(xs, ys):
assert func(x, y) == approx(castor(x * y))
def test_augassign():
@ti.kernel
def foo(x: ti.i32, y: ti.i32, a: ti.template(), b: ti.template()):
for i in a:
a[i] = x
a[0] += y
a[1] -= y
a[2] *= y
a[3] //= y
a[4] %= y
a[5] **= y
a[6] <<= y
a[7] >>= y
a[8] |= y
a[9] ^= y
a[10] &= y
b[0] = x
b[0] /= y
x = 37
y = 3
a = ti.field(ti.i32, shape=(11, ))
b = ti.field(ti.i32, shape=(11, ))
c = ti.field(ti.f32, shape=(1, ))
d = ti.field(ti.f32, shape=(1, ))
a[0] = x + y
a[1] = x - y
a[2] = x * y
a[3] = x // y
a[4] = x % y
a[5] = x**y
a[6] = x << y
a[7] = x >> y
a[8] = x | y
a[9] = x ^ y
a[10] = x & y
c[0] = x / y
foo(x, y, b, d)
for i in range(11):
assert a[i] == b[i]
assert c[0] == approx(d[0])
def test_transpose():
dim = 3
m = ti.Matrix.field(dim, dim, ti.f32)
ti.root.place(m)
@ti.kernel
def transpose():
mat = m[None].transpose()
m[None] = mat
for i in range(dim):
for j in range(dim):
m[None][i, j] = i * 2 + j * 7
transpose()
for i in range(dim):
for j in range(dim):
assert m[None][j, i] == approx(i * 2 + j * 7)
def test_random_int():
for precision in [ti.i32, ti.i64]:
ti.init()
n = 1024
x = ti.field(ti.f32, shape=(n, n))
@ti.kernel
def fill():
for i in range(n):
for j in range(n):
v = ti.random(precision)
if precision == ti.i32:
x[i, j] = (float(v) + float(2**31)) / float(2**32)
else:
x[i, j] = (float(v) + float(2**63)) / float(2**64)
fill()
X = x.to_numpy()
for i in range(1, 4):
assert (X**i).mean() == approx(1 / (i + 1), rel=1e-2)
def test_random_2d_dist():
n = 8192
x = ti.Vector.field(2, dtype=ti.f32, shape=n)
@ti.kernel
def gen():
for i in range(n):
x[i] = ti.Vector([ti.random(), ti.random()])
gen()
X = x.to_numpy()
counters = [0 for _ in range(4)]
for i in range(n):
c = int(X[i, 0] < 0.5) * 2 + int(X[i, 1] < 0.5)
counters[c] += 1
for c in range(4):
assert counters[c] / n == approx(1 / 4, rel=0.2)
def test_randn():
'''
Tests the generation of Gaussian random numbers.
'''
for precision in [ti.f32, ti.f64]:
ti.init()
n = 1024
x = ti.field(ti.f32, shape=(n, n))
@ti.kernel
def fill():
for i in range(n):
for j in range(n):
x[i, j] = ti.randn(precision)
fill()
X = x.to_numpy()
# https://en.wikipedia.org/wiki/Normal_distribution#Moments
moments = [0.0, 1.0, 0.0, 3.0]
for i in range(4):
assert (X**(i + 1)).mean() == approx(moments[i], abs=3e-2)
def run_test(x, v, C, J, grid_v, grid_m):
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 if isinstance(x, np.ndarray) else 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)
