import numpy as np
import pytest
from taichi._testing import make_temp_file
from taichi.lang.misc import get_host_arch_list
from taichi.lang.util import to_numpy_type
import taichi as ti
# jpg is also supported but hard to test here since it's lossy:
@pytest.mark.parametrize('comp,ext', [(3, 'bmp'), (1, 'png'), (3, 'png'),
(4, 'png')])
@pytest.mark.parametrize('resx,resy', [(201, 173)])
@pytest.mark.parametrize('is_field', [False, True])
@pytest.mark.parametrize('dt', [ti.u8])
@ti.test(arch=get_host_arch_list())
def test_image_io(resx, resy, comp, ext, is_field, dt):
if comp != 1:
shape = (resx, resy, comp)
else:
shape = (resx, resy)
if is_field:
pixel_t = ti.field(dt, shape)
pixel = np.random.randint(256, size=shape, dtype=to_numpy_type(dt))
if is_field:
pixel_t.from_numpy(pixel)
fn = make_temp_file(suffix='.' + ext)
if is_field:
ti.imwrite(pixel_t, fn)
else:
ti.imwrite(pixel, fn)
import pytest
from taichi.lang import impl
from taichi.lang.misc import get_host_arch_list
import taichi as ti
from tests import test_utils
data_types = [ti.i32, ti.f32, ti.i64, ti.f64]
field_shapes = [(), 8, (6, 12)]
vector_dims = [3]
matrix_dims = [(1, 2), (2, 3)]
@pytest.mark.parametrize('dtype', data_types)
@pytest.mark.parametrize('shape', field_shapes)
@test_utils.test(arch=get_host_arch_list())
def test_scalar_field(dtype, shape):
x = ti.field(dtype, shape)
if isinstance(shape, tuple):
assert x.shape == shape
else:
assert x.shape == (shape, )
assert x.dtype == dtype
@pytest.mark.parametrize('n', vector_dims)
@pytest.mark.parametrize('dtype', data_types)
@pytest.mark.parametrize('shape', field_shapes)
@test_utils.test(arch=get_host_arch_list())
def _test_scalar_ndarray(dtype, shape):
x = ti.ndarray(dtype, shape)
if isinstance(shape, tuple):
assert x.shape == shape
else:
assert x.shape == (shape, )
assert x.element_shape == ()
assert x.dtype == dtype
@pytest.mark.parametrize('dtype', data_types)
@pytest.mark.parametrize('shape', ndarray_shapes)
@pytest.mark.skipif(not has_pytorch(), reason='Pytorch not installed.')
@ti.test(arch=get_host_arch_list(), ndarray_use_torch=True)
def test_scalar_ndarray_torch(dtype, shape):
_test_scalar_ndarray(dtype, shape)
@pytest.mark.parametrize('dtype', data_types)
@pytest.mark.parametrize('shape', ndarray_shapes)
@ti.test(arch=get_host_arch_list())
def test_scalar_ndarray(dtype, shape):
_test_scalar_ndarray(dtype, shape)
def _test_vector_ndarray(n, dtype, shape):
x = ti.Vector.ndarray(n, dtype, shape)
if isinstance(shape, tuple):
