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)