test_api.py

import sys

import pytest

import taichi as ti
from tests import test_utils


def _get_matrix_swizzle_apis():
    swizzle_gen = ti.lang.matrix._generate_swizzle_patterns
    KEMAP_SET = ["xyzw", "rgba", "stpq"]
    res = []
    for key_group in KEMAP_SET:
        sw_patterns = swizzle_gen(key_group, required_length=4)
        res += sw_patterns
    return sorted(res)


def _get_expected_matrix_apis():
    base = [
        "all",
        "any",
        "cast",
        "cols",
        "cross",
        "determinant",
        "diag",
        "dot",
        "field",
        "fill",
        "identity",
        "inverse",
        "max",
        "min",
        "ndarray",
        "norm",
        "norm_inv",
        "norm_sqr",
        "normalized",
        "one",
        "outer_product",
        "rows",
        "sum",
        "to_list",
        "to_numpy",
        "trace",
        "transpose",
        "unit",
        "zero",
        "get_shape",
    ]
    res = base + _get_matrix_swizzle_apis()
    return sorted(res)


user_api = {}
user_api[ti] = [
    "ArgPack",
    "BitpackedFields",
    "CRITICAL",
    "DEBUG",
    "DeviceCapability",
    "ERROR",
    "Field",
    "FieldsBuilder",
    "Format",
    "GUI",
    "INFO",
    "Layout",
    "Matrix",
    "MatrixField",
    "MatrixNdarray",
    "Mesh",
    "MeshInstance",
    "Ndarray",
    "SNode",
    "ScalarField",
    "ScalarNdarray",
    "Struct",
    "StructField",
    "TRACE",
    "TaichiAssertionError",
    "TaichiCompilationError",
    "TaichiNameError",
    "TaichiRuntimeError",
    "TaichiRuntimeTypeError",
    "TaichiSyntaxError",
    "TaichiTypeError",
    "Texture",
    "Vector",
    "VectorNdarray",
    "WARN",
    "abs",
    "acos",
    "activate",
    "ad",
    "algorithms",
    "amdgpu",
    "aot",
    "append",
    "argpack",
    "arm64",
    "asin",
    "assume_in_range",
    "atan2",
    "atomic_add",
    "atomic_and",
    "atomic_max",
    "atomic_min",
    "atomic_mul",
    "atomic_or",
    "atomic_sub",
    "atomic_xor",
    "axes",
    "bit_cast",
    "bit_shr",
    "block_local",
    "cache_read_only",
    "cast",
    "ceil",
    "cos",
    "cpu",
    "cuda",
    "data_oriented",
    "dataclass",
    "deactivate",
    "deactivate_all_snodes",
    "dx11",
    "dx12",
    "eig",
    "exp",
    "experimental",
    "extension",
    "f16",
    "f32",
    "f64",
    "field",
    "float16",
    "float32",
    "float64",
    "floor",
    "frexp",
    "func",
    "get_addr",
    "gles",
    "global_thread_idx",
    "gpu",
    "graph",
    "grouped",
    "hex_to_rgb",
    "i",
    "i16",
    "i32",
    "i64",
    "i8",
    "ij",
    "ijk",
    "ijkl",
    "ijl",
    "ik",
    "ikl",
    "il",
    "init",
    "int16",
    "int32",
    "int64",
    "int8",
    "is_active",
    "is_logging_effective",
    "j",
    "jk",
    "jkl",
    "jl",
    "k",
    "kernel",
    "kl",
    "l",
    "lang",
    "length",
    "linalg",
    "log",
    "loop_config",
    "math",
    "max",
    "mesh_local",
    "mesh_patch_idx",
    "metal",
    "min",
    "ndarray",
    "ndrange",
    "no_activate",
    "one",
    "opengl",
    "polar_decompose",
    "pow",
    "profiler",
    "pyfunc",
    "randn",
    "random",
    "raw_div",
    "raw_mod",
    "real_func",
    "ref",
    "rescale_index",
    "reset",
    "rgb_to_hex",
    "root",
    "round",
    "rsqrt",
    "select",
    "set_logging_level",
    "simt",
    "sin",
    "solve",
    "sparse",
    "sparse_matrix_builder",
    "sqrt",
    "static",
    "static_assert",
    "static_print",
    "stop_grad",
    "svd",
    "sym_eig",
    "sync",
    "tan",
    "tanh",
    "template",
    "tools",
    "types",
    "u1",
    "u16",
    "u32",
    "u64",
    "u8",
    "ui",
    "uint1",
    "uint16",
    "uint32",
    "uint64",
    "uint8",
    "vulkan",
    "x64",
    "x86_64",
    "zero",
]
user_api[ti.ad] = [
    "FwdMode",
    "Tape",
    "clear_all_gradients",
    "grad_for",
    "grad_replaced",
    "no_grad",
]
user_api[ti.algorithms] = ["PrefixSumExecutor", "parallel_sort"]
user_api[ti.Field] = [
    "copy_from",
    "dtype",
    "fill",
    "from_numpy",
    "from_paddle",
    "from_torch",
    "parent",
    "shape",
    "snode",
    "to_numpy",
    "to_paddle",
    "to_torch",
]
user_api[ti.FieldsBuilder] = [
    "bitmasked",
    "deactivate_all",
    "dense",
    "dynamic",
    "finalize",
    "lazy_dual",
    "lazy_grad",
    "place",
    "pointer",
    "quant_array",
]
user_api[ti.math] = [
    "acos",
    "asin",
    "atan2",
    "cconj",
    "cdiv",
    "ceil",
    "cexp",
    "cinv",
    "clamp",
    "clog",
    "clz",
    "cmul",
    "cos",
    "cpow",
    "cross",
    "csqrt",
    "degrees",
    "determinant",
    "distance",
    "dot",
    "e",
    "exp",
    "eye",
    "floor",
    "fract",
    "inf",
    "inverse",
    "isinf",
    "isnan",
    "ivec2",
    "ivec3",
    "ivec4",
    "length",
    "log",
    "log2",
    "mat2",
    "mat3",
    "mat4",
    "max",
    "min",
    "mix",
    "mod",
    "nan",
    "normalize",
    "pi",
    "popcnt",
    "pow",
    "radians",
    "reflect",
    "refract",
    "rot_by_axis",
    "rot_yaw_pitch_roll",
    "rotation2d",
    "rotation3d",
    "round",
    "scale",
    "sign",
    "sin",
    "smoothstep",
    "sqrt",
    "step",
    "tan",
    "tanh",
    "translate",
    "uvec2",
    "uvec3",
    "uvec4",
    "vdir",
    "vec2",
    "vec3",
    "vec4",
]
user_api[ti.Matrix] = _get_expected_matrix_apis()
user_api[ti.MatrixField] = [
    "copy_from",
    "dtype",
    "fill",
    "from_numpy",
    "from_paddle",
    "from_torch",
    "get_scalar_field",
    "parent",
    "shape",
    "snode",
    "to_numpy",
    "to_paddle",
    "to_torch",
]
user_api[ti.MatrixNdarray] = [
    "copy_from",
    "element_shape",
    "fill",
    "from_numpy",
    "get_type",
    "to_numpy",
]
user_api[ti.Ndarray] = ["copy_from", "element_shape", "fill", "get_type"]
user_api[ti.Texture] = ["from_field", "from_image", "from_ndarray", "to_image"]
user_api[ti.SNode] = [
    "bitmasked",
    "deactivate_all",
    "dense",
    "dynamic",
    "lazy_dual",
    "lazy_grad",
    "parent",
    "place",
    "pointer",
    "quant_array",
    "shape",
]
user_api[ti.ScalarField] = [
    "copy_from",
    "dtype",
    "fill",
    "from_numpy",
    "from_paddle",
    "from_torch",
    "parent",
    "shape",
    "snode",
    "to_numpy",
    "to_paddle",
    "to_torch",
]
user_api[ti.ScalarNdarray] = [
    "copy_from",
    "element_shape",
    "fill",
    "from_numpy",
    "get_type",
    "to_numpy",
]
user_api[ti.Struct] = ["entries", "field", "items", "keys", "methods", "to_dict"]
user_api[ti.ArgPack] = ["items", "keys", "to_dict"]
user_api[ti.StructField] = [
    "copy_from",
    "dtype",
    "fill",
    "from_numpy",
    "from_paddle",
    "from_torch",
    "get_member_field",
    "keys",
    "parent",
    "shape",
    "snode",
    "to_numpy",
    "to_paddle",
    "to_torch",
]
user_api[ti.VectorNdarray] = [
    "copy_from",
    "element_shape",
    "fill",
    "from_numpy",
    "get_type",
    "to_numpy",
]
user_api[ti.sparse] = ["grid", "usage"]


@pytest.mark.parametrize("src", user_api.keys())
@test_utils.test(arch=ti.cpu)
def test_api(src):
    # When Python version is below 3.7, deprecated names are
    # handled as normal names, which will fail this test.
    expected = user_api[src]
    actual = [s for s in dir(src) if not s.startswith("_")]
    assert (
        sys.version_info < (3, 7) or actual == expected
    ), f"Failed for API={src}:\n  expected={expected}\n  actual={actual}"