def test_python_scope_linalg():
ti.init()
a = np.array([3, 4, -2])
b = np.array([-5, 0, 6])
x = ti.Vector(a)
y = ti.Vector(b)
assert test_utils.allclose(x.dot(y), np.dot(a, b))
assert test_utils.allclose(x.norm(), np.sqrt(np.dot(a, a)))
assert test_utils.allclose(x.normalized(), a / np.sqrt(np.dot(a, a)))
assert x.any() == 1 # To match that of Taichi IR, we return -1 for True
assert y.all() == 0
def test_writeback_binary_i(rhs_is_mat):
x = ti.Matrix.field(3, 2, ti.i32, 12)
y = ti.Matrix.field(3, 2, ti.i32, ())
if rhs_is_mat:
z = ti.Matrix.field(3, 2, ti.i32, ())
else:
z = ti.field(ti.i32, ())
y.from_numpy(np.array([[0, 2], [9, 3], [7, 4]], np.int32))
if rhs_is_mat:
z.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.int32))
else:
z[None] = 5
@ti.kernel
def func():
for i in x:
x[i] = y[None]
x[0] = z[None]
x[1] += z[None]
x[2] -= z[None]
x[3] *= z[None]
x[4] //= z[None]
x[5] %= z[None]
x[6] &= z[None]
x[7] |= z[None]
x[8] ^= z[None]
ti.atomic_min(x[10], z[None])
ti.atomic_max(x[11], z[None])
func()
x = x.to_numpy()
y = y.to_numpy()
z = z.to_numpy()
assert test_utils.allclose(x[1], y + z)
assert test_utils.allclose(x[2], y - z)
assert test_utils.allclose(x[3], y * z)
assert test_utils.allclose(x[4], y // z)
assert test_utils.allclose(x[5], y % z)
assert test_utils.allclose(x[6], y & z)
assert test_utils.allclose(x[7], y | z)
assert test_utils.allclose(x[8], y ^ z)
assert test_utils.allclose(x[10], np.minimum(y, z))
assert test_utils.allclose(x[11], np.maximum(y, z))
def test_python_scope_matmul():
ti.init()
a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])
x = ti.Vector(a)
y = ti.Vector(b)
result = (x @ y).to_numpy()
expected = a @ b
assert test_utils.allclose(result, expected)
def test_python_scope_vector_unary(ti_func, np_func):
ti.init()
x = ti.Vector([2, 3] if ti_func in
[ops.invert, ti.lang.ops.logical_not] else [0.2, 0.3])
result = ti_func(x).to_numpy()
if ti_func in [ti.lang.ops.logical_not]:
result = result.astype(bool)
expected = np_func(x.to_numpy())
assert test_utils.allclose(result, expected)
def test_python_scope_vector_binary(ti_func, np_func):
ti.init()
x = ti.Vector([2, 3])
y = ti.Vector([5, 4])
result = ti_func(x, y).to_numpy()
if ti_func in [ops.eq, ops.ne, ops.lt, ops.le, ops.gt, ops.ge]:
result = result.astype(bool)
expected = np_func(x.to_numpy(), y.to_numpy())
assert test_utils.allclose(result, expected)
def test_writeback_binary_f(rhs_is_mat):
x = ti.Matrix.field(3, 2, ti.f32, 9)
y = ti.Matrix.field(3, 2, ti.f32, ())
if rhs_is_mat:
z = ti.Matrix.field(3, 2, ti.f32, ())
else:
z = ti.field(ti.f32, ())
y.from_numpy(np.array([[0, 2], [9, 3.1], [7, 4]], np.float32))
if rhs_is_mat:
z.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.float32))
else:
z[None] = 5
@ti.kernel
def func():
for i in x:
x[i] = y[None]
if ti.static(rhs_is_mat):
x[0] = z[None]
else:
x[0].fill(z[None])
x[1] += z[None]
x[2] -= z[None]
x[3] *= z[None]
x[4] /= z[None]
x[5] //= z[None]
x[6] %= z[None]
ti.atomic_min(x[7], z[None])
ti.atomic_max(x[8], z[None])
func()
x = x.to_numpy()
y = y.to_numpy()
z = z.to_numpy()
assert test_utils.allclose(x[1], y + z)
assert test_utils.allclose(x[2], y - z)
assert test_utils.allclose(x[3], y * z)
assert test_utils.allclose(x[4], y / z)
assert test_utils.allclose(x[5], y // z)
assert test_utils.allclose(x[6], y % z)
assert test_utils.allclose(x[7], np.minimum(y, z))
assert test_utils.allclose(x[8], np.maximum(y, z))
def test_ternary_i(is_mat):
cond_is_mat, lhs_is_mat, rhs_is_mat = is_mat
x = ti.Matrix.field(3, 2, ti.i32, 1)
if cond_is_mat:
y = ti.Matrix.field(3, 2, ti.i32, ())
else:
y = ti.field(ti.i32, ())
if lhs_is_mat:
z = ti.Matrix.field(3, 2, ti.i32, ())
else:
z = ti.field(ti.i32, ())
if rhs_is_mat:
w = ti.Matrix.field(3, 2, ti.i32, ())
else:
w = ti.field(ti.i32, ())
if cond_is_mat:
y.from_numpy(np.array([[0, 2], [9, 0], [7, 4]], np.int32))
else:
y[None] = 0
if lhs_is_mat:
z.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.int32))
else:
z[None] = 5
if rhs_is_mat:
w.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.int32))
else:
w[None] = 4
@ti.kernel
def func():
x[0] = ti.select(y[None], z[None], w[None])
func()
x = x.to_numpy()
y = y.to_numpy()
z = z.to_numpy()
w = w.to_numpy()
assert test_utils.allclose(
x[0],
np.int32(np.bool_(y)) * z + np.int32(1 - np.bool_(y)) * w)
def test_texture():
res = (256, 256)
@ti.kernel
def make_texture(tex: ti.types.rw_texture(num_dimensions=2,
num_channels=1,
channel_format=ti.f32,
lod=0)):
for i, j in ti.ndrange(128, 128):
tex.store(ti.Vector([i, j]), ti.Vector([0.1, 0.0, 0.0, 0.0]))
@ti.kernel
def paint(t: ti.f32, pixels: ti.types.ndarray(field_dim=2),
tex: ti.types.texture(num_dimensions=2)):
for i, j in pixels:
uv = ti.Vector([i / res[0], j / res[1]])
warp_uv = uv + ti.Vector(
[ti.cos(t + uv.x * 5.0),
ti.sin(t + uv.y * 5.0)]) * 0.1
c = ti.math.vec4(0.0)
if uv.x > 0.5:
c = tex.sample_lod(warp_uv, 0.0)
else:
c = tex.fetch(ti.cast(warp_uv * 128, ti.i32), 0)
pixels[i, j] = [c.r, c.r, c.r, 1.0]
_t = ti.graph.Arg(ti.graph.ArgKind.SCALAR, 't', ti.f32)
_pixels_arr = ti.graph.Arg(ti.graph.ArgKind.NDARRAY,
'pixels_arr',
ti.f32,
field_dim=2,
element_shape=(4, ))
_rw_tex = ti.graph.Arg(ti.graph.ArgKind.RWTEXTURE,
'rw_tex',
channel_format=ti.f32,
shape=(128, 128),
num_channels=1)
_tex = ti.graph.Arg(ti.graph.ArgKind.TEXTURE,
'tex',
channel_format=ti.f32,
shape=(128, 128),
num_channels=1)
g_builder = ti.graph.GraphBuilder()
g_builder.dispatch(make_texture, _rw_tex)
g_builder.dispatch(paint, _t, _pixels_arr, _tex)
g = g_builder.compile()
pixels_arr = ti.Vector.ndarray(4, dtype=float, shape=res)
texture = ti.Texture(ti.f32, 1, (128, 128))
t = 1
g.run({
'rw_tex': texture,
't': t,
'pixels_arr': pixels_arr,
'tex': texture
})
pixels = pixels_arr.to_numpy()
for i in range(res[0]):
for j in range(res[1]):
assert test_utils.allclose(pixels[i, j], [0.1, 0.1, 0.1, 1.])
def test_binary_i(is_mat):
lhs_is_mat, rhs_is_mat = is_mat
x = ti.Matrix.field(3, 2, ti.i32, 20)
if lhs_is_mat:
y = ti.Matrix.field(3, 2, ti.i32, ())
else:
y = ti.field(ti.i32, ())
if rhs_is_mat:
z = ti.Matrix.field(3, 2, ti.i32, ())
else:
z = ti.field(ti.i32, ())
if lhs_is_mat:
y.from_numpy(np.array([[0, 2], [9, 3], [7, 4]], np.int32))
else:
y[None] = 6
if rhs_is_mat:
z.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.int32))
else:
z[None] = 5
@ti.kernel
def func():
x[0] = y[None] + z[None]
x[1] = y[None] - z[None]
x[2] = y[None] * z[None]
x[3] = y[None] // z[None]
x[4] = ti.raw_div(y[None], z[None])
x[5] = y[None] % z[None]
x[6] = ti.raw_mod(y[None], z[None])
x[7] = y[None]**z[None]
x[8] = y[None] == z[None]
x[9] = y[None] != z[None]
x[10] = y[None] > z[None]
x[11] = y[None] >= z[None]
x[12] = y[None] < z[None]
x[13] = y[None] <= z[None]
x[14] = y[None] & z[None]
x[15] = y[None] ^ z[None]
x[16] = y[None] | z[None]
x[17] = ti.min(y[None], z[None])
x[18] = ti.max(y[None], z[None])
x[19] = y[None] << z[None]
func()
x = x.to_numpy()
y = y.to_numpy()
z = z.to_numpy()
assert test_utils.allclose(x[0], y + z)
assert test_utils.allclose(x[1], y - z)
assert test_utils.allclose(x[2], y * z)
assert test_utils.allclose(x[3], y // z)
assert test_utils.allclose(x[4], y // z)
assert test_utils.allclose(x[5], y % z)
assert test_utils.allclose(x[6], y % z)
assert test_utils.allclose(x[7], y**z, rel=1e-5)
assert test_utils.allclose(x[8], y == z)
assert test_utils.allclose(x[9], y != z)
assert test_utils.allclose(x[10], y > z)
assert test_utils.allclose(x[11], y >= z)
assert test_utils.allclose(x[12], y < z)
assert test_utils.allclose(x[13], y <= z)
assert test_utils.allclose(x[14], y & z)
assert test_utils.allclose(x[15], y ^ z)
assert test_utils.allclose(x[16], y | z)
assert test_utils.allclose(x[17], np.minimum(y, z))
assert test_utils.allclose(x[18], np.maximum(y, z))
assert test_utils.allclose(x[19], y << z)
def test_unary():
xi = ti.Matrix.field(3, 2, ti.i32, 4)
yi = ti.Matrix.field(3, 2, ti.i32, ())
xf = ti.Matrix.field(3, 2, ti.f32, 15)
yf = ti.Matrix.field(3, 2, ti.f32, ())
yi.from_numpy(np.array([[3, 2], [9, 0], [7, 4]], np.int32))
yf.from_numpy(np.array([[0.3, 0.2], [0.9, 0.1], [0.7, 0.4]], np.float32))
@ti.kernel
def func():
xi[0] = -yi[None]
xi[1] = ~yi[None]
xi[2] = not yi[None]
xi[3] = abs(yi[None])
xf[0] = -yf[None]
xf[1] = abs(yf[None])
xf[2] = ti.sqrt(yf[None])
xf[3] = ti.sin(yf[None])
xf[4] = ti.cos(yf[None])
xf[5] = ti.tan(yf[None])
xf[6] = ti.asin(yf[None])
xf[7] = ti.acos(yf[None])
xf[8] = ti.tanh(yf[None])
xf[9] = ti.floor(yf[None])
xf[10] = ti.ceil(yf[None])
xf[11] = ti.exp(yf[None])
xf[12] = ti.log(yf[None])
xf[13] = ti.rsqrt(yf[None])
xf[14] = ti.round(yf[None])
func()
xi = xi.to_numpy()
yi = yi.to_numpy()
xf = xf.to_numpy()
yf = yf.to_numpy()
assert test_utils.allclose(xi[0], -yi)
assert test_utils.allclose(xi[1], ~yi)
assert test_utils.allclose(xi[3], np.abs(yi))
assert test_utils.allclose(xf[0], -yf)
assert test_utils.allclose(xf[1], np.abs(yf))
assert test_utils.allclose(xf[2], np.sqrt(yf), rel=1e-5)
assert test_utils.allclose(xf[3], np.sin(yf), rel=1e-4)
assert test_utils.allclose(xf[4], np.cos(yf), rel=1e-4)
assert test_utils.allclose(xf[5], np.tan(yf), rel=1e-4)
# vulkan need 1e-3
assert test_utils.allclose(xf[6], np.arcsin(yf), rel=1e-3)
assert test_utils.allclose(xf[7], np.arccos(yf), rel=1e-3)
assert test_utils.allclose(xf[8], np.tanh(yf), rel=1e-4)
assert test_utils.allclose(xf[9], np.floor(yf), rel=1e-5)
assert test_utils.allclose(xf[10], np.ceil(yf), rel=1e-5)
assert test_utils.allclose(xf[11], np.exp(yf), rel=1e-5)
assert test_utils.allclose(xf[12], np.log(yf), rel=1e-5)
assert test_utils.allclose(xf[13], 1 / np.sqrt(yf), rel=1e-5)
assert test_utils.allclose(xf[14], np.round(yf), rel=1e-5)
def test_binary_f(lhs_is_mat, rhs_is_mat):
x = ti.Matrix.field(3, 2, ti.f32, 16)
if lhs_is_mat:
y = ti.Matrix.field(3, 2, ti.f32, ())
else:
y = ti.field(ti.f32, ())
if rhs_is_mat:
z = ti.Matrix.field(3, 2, ti.f32, ())
else:
z = ti.field(ti.f32, ())
if lhs_is_mat:
y.from_numpy(np.array([[0, 2], [9, 3.1], [7, 4]], np.float32))
else:
y[None] = 6.1
if rhs_is_mat:
z.from_numpy(np.array([[4, 5], [6, 3], [9, 2]], np.float32))
else:
z[None] = 5
@ti.kernel
def func():
x[0] = y[None] + z[None]
x[1] = y[None] - z[None]
x[2] = y[None] * z[None]
x[3] = y[None] / z[None]
x[4] = y[None] // z[None]
x[5] = y[None] % z[None]
x[6] = y[None]**z[None]
x[7] = y[None] == z[None]
x[8] = y[None] != z[None]
x[9] = y[None] > z[None]
x[10] = y[None] >= z[None]
x[11] = y[None] < z[None]
x[12] = y[None] <= z[None]
x[13] = ti.atan2(y[None], z[None])
x[14] = ti.min(y[None], z[None])
x[15] = ti.max(y[None], z[None])
func()
x = x.to_numpy()
y = y.to_numpy()
z = z.to_numpy()
assert test_utils.allclose(x[0], y + z)
assert test_utils.allclose(x[1], y - z)
assert test_utils.allclose(x[2], y * z)
assert test_utils.allclose(x[3], y / z)
assert test_utils.allclose(x[4], y // z)
assert test_utils.allclose(x[5], y % z)
assert test_utils.allclose(x[6], y**z)
assert test_utils.allclose(x[7], y == z)
assert test_utils.allclose(x[8], y != z)
assert test_utils.allclose(x[9], y > z)
assert test_utils.allclose(x[10], y >= z)
assert test_utils.allclose(x[11], y < z)
assert test_utils.allclose(x[12], y <= z)
assert test_utils.allclose(x[13], np.arctan2(y, z))
assert test_utils.allclose(x[14], np.minimum(y, z))
assert test_utils.allclose(x[15], np.maximum(y, z))
