gstaichi 0.1.18.dev1__cp310-cp310-macosx_15_0_arm64.whl

taichi/tools/diagnose.py

@@ -0,0 +1,124 @@
+# type: ignore
+
+import locale
+import os
+import platform
+import subprocess
+import sys
+
+
+def main():
+    print("Taichi system diagnose:")
+    print("")
+    executable = sys.executable
+
+    print(f"python: {sys.version}")
+    print(f"system: {sys.platform}")
+    print(f"executable: {executable}")
+    print(f"platform: {platform.platform()}")
+    print(f'architecture: {" ".join(platform.architecture())}')
+    print(f"uname: {platform.uname()}")
+
+    print(f'locale: {".".join(locale.getdefaultlocale())}')
+    print(f'PATH: {os.environ.get("PATH")}')
+    print(f"PYTHONPATH: {sys.path}")
+    print("")
+
+    try:
+        lsb_release = subprocess.check_output(["lsb_release", "-a"])
+    except Exception as e:
+        print(f"`lsb_release` not available: {e}")
+    else:
+        print(f"{lsb_release.decode()}")
+
+    for k, v in os.environ.items():
+        if k.startswith("TI_"):
+            print(f"{k}={v}")
+    print("")
+
+    def try_print(tag, expr):
+        try:
+            cmd = f'import taichi as ti; print("===="); print({expr}, end="")'
+            ret = subprocess.check_output([executable, "-c", cmd]).decode()
+            ret = ret.split("====" + os.linesep, maxsplit=1)[1]
+            print(f"{tag}: {ret}")
+        except Exception as e:
+            print(f"{tag}: ERROR {e}")
+
+    print("")
+    try_print("import", "ti")
+    print("")
+    for arch in ["cpu", "metal", "opengl", "cuda", "vulkan"]:
+        try_print(arch, f"ti.lang.misc.is_arch_supported(ti.{arch})")
+    print("")
+
+    try:
+        glewinfo = subprocess.check_output(["glewinfo"])
+    except Exception as e:
+        print(f"`glewinfo` not available: {e}")
+    else:
+        for line in glewinfo.decode().splitlines():
+            if line.startswith("OpenGL version"):
+                print(line)
+                continue
+
+            exts = [
+                "GL_ARB_compute_shader",
+                "GL_ARB_gpu_shader_int64",
+                "GL_NV_shader_atomic_float",
+                "GL_NV_shader_atomic_float64",
+                "GL_NV_shader_atomic_int64",
+            ]
+            if line.split(":")[0] in exts:
+                print(line)
+
+    print("")
+    try:
+        nvidia_smi = subprocess.check_output(["nvidia-smi"])
+    except Exception as e:
+        print(f"`nvidia-smi` not available: {e}")
+    else:
+        print(f"{nvidia_smi.decode()}")
+
+    try:
+        ti_header = subprocess.check_output([executable, "-c", "import taichi"])
+    except Exception as e:
+        print(f"`import taichi` failed: {e}")
+    else:
+        print(f"{ti_header.decode()}")
+
+    try:
+        ti_init_test = subprocess.check_output([executable, "-c", "import taichi as ti; ti.init()"])
+    except Exception as e:
+        print(f"`ti.init()` failed: {e}")
+    else:
+        print(f"{ti_init_test.decode()}")
+
+    try:
+        ti_opengl_test = subprocess.check_output([executable, "-c", "import taichi as ti; ti.init(arch=ti.opengl)"])
+    except Exception as e:
+        print(f"Taichi OpenGL test failed: {e}")
+    else:
+        print(f"{ti_opengl_test.decode()}")
+
+    try:
+        ti_cuda_test = subprocess.check_output([executable, "-c", "import taichi as ti; ti.init(arch=ti.cuda)"])
+    except Exception as e:
+        print(f"Taichi CUDA test failed: {e}")
+    else:
+        print(f"{ti_cuda_test.decode()}")
+
+    try:
+        ti_laplace = subprocess.check_output([executable, "-m", "taichi", "example", "minimal"])
+    except Exception as e:
+        print(f"`python/taichi/examples/algorithm/laplace.py` failed: {e}")
+    else:
+        print(f"{ti_laplace.decode()}")
+
+    print("Consider attaching this log when maintainers ask about system information.")
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = []

taichi/tools/np2ply.py

@@ -0,0 +1,364 @@
+# type: ignore
+
+# convert numpy array to ply files
+import sys
+
+import numpy as np
+
+
+class PLYWriter:
+    """Writes `numpy.array` data to `ply` files.
+
+    Args:
+        num_vertices (int): number of vertices.
+        num_faces (int, optional): number of faces.
+        face_type (str): `tri` or `quad`.
+        comment (str): comment message.
+    """
+
+    def __init__(
+        self,
+        num_vertices: int,
+        num_faces=0,
+        face_type="tri",
+        comment="created by PLYWriter",
+    ):
+        assert num_vertices > 0, "num_vertices should be greater than 0"
+        assert num_faces >= 0, "num_faces shouldn't be less than 0"
+        assert face_type == "tri" or face_type == "quad", "Only tri and quad faces are supported for now"
+
+        self.ply_supported_types = [
+            "char",
+            "uchar",
+            "short",
+            "ushort",
+            "int",
+            "uint",
+            "float",
+            "double",
+        ]
+        self.corresponding_numpy_types = [
+            np.int8,
+            np.uint8,
+            np.int16,
+            np.uint16,
+            np.int32,
+            np.uint32,
+            np.float32,
+            np.float64,
+        ]
+        self.type_map = {}
+        for i, ply_type in enumerate(self.ply_supported_types):
+            self.type_map[ply_type] = self.corresponding_numpy_types[i]
+
+        self.num_vertices = num_vertices
+        self.num_vertex_channels = 0
+        self.vertex_channels = []
+        self.vertex_data_type = []
+        self.vertex_data = []
+        self.num_faces = num_faces
+        self.num_face_channels = 0
+        self.face_channels = []
+        self.face_data_type = []
+        self.face_data = []
+        self.face_type = face_type
+        if face_type == "tri":
+            self.face_indices = -np.ones((self.num_faces, 3), dtype=np.int32)
+        elif face_type == "quad":
+            self.face_indices = -np.ones((self.num_faces, 4), dtype=np.int32)
+        self.comment = comment
+
+    def add_vertex_channel(self, key: str, data_type: str, data: np.array):
+        if data_type not in self.ply_supported_types:
+            print("Unknown type " + data_type + " detected, skipping this channel")
+            return
+        if data.ndim == 1:
+            assert data.size == self.num_vertices, "The dimension of the vertex channel is not correct"
+            self.num_vertex_channels += 1
+            if key in self.vertex_channels:
+                print("WARNING: duplicate key " + key + " detected")
+            self.vertex_channels.append(key)
+            self.vertex_data_type.append(data_type)
+            self.vertex_data.append(self.type_map[data_type](data))
+        else:
+            num_col = data.size // self.num_vertices
+            assert (
+                data.ndim == 2 and data.size == num_col * self.num_vertices
+            ), "The dimension of the vertex channel is not correct"
+            data.shape = (self.num_vertices, num_col)
+            self.num_vertex_channels += num_col
+            for i in range(num_col):
+                item_key = key + "_" + str(i + 1)
+                if item_key in self.vertex_channels:
+                    print("WARNING: duplicate key " + item_key + " detected")
+                self.vertex_channels.append(item_key)
+                self.vertex_data_type.append(data_type)
+                self.vertex_data.append(self.type_map[data_type](data[:, i]))
+
+    def add_vertex_pos(self, x: np.array, y: np.array, z: np.array):
+        """Set the (x, y, z) coordinates of the vertices.
+
+        Args:
+            x (`numpy.array(float)`): x-coordinates of the vertices.
+            y (`numpy.array(float)`): y-coordinates of the vertices.
+            z (`numpy.array(float)`): z-coordinates of the vertices.
+        """
+        self.add_vertex_channel("x", "float", x)
+        self.add_vertex_channel("y", "float", y)
+        self.add_vertex_channel("z", "float", z)
+
+    # TODO active and refactor later if user feedback indicates the necessity for a compact the input list
+    # pass ti vector/matrix field directly
+    # def add_vertex_pos(self, pos):
+    #     assert isinstance(pos, (np.ndarray, ti.Matrix))
+    #     if not isinstance(pos, np.ndarray):
+    #         pos = pos.to_numpy()
+    #     dim = pos.shape[pos.ndim-1]
+    #     assert dim == 2 or dim == 3, "Only 2D and 3D positions are supported"
+    #     n = pos.size // dim
+    #     assert n == self.num_vertices, "Size of the input is not correct"
+    #     pos = np.reshape(pos, (n, dim))
+    #     self.add_vertex_channel("x", "float", pos[:, 0])
+    #     self.add_vertex_channel("y", "float", pos[:, 1])
+    #     if(dim == 3):
+    #         self.add_vertex_channel("z", "float", pos[:, 2])
+    #     if(dim == 2):
+    #         self.add_vertex_channel("z", "float", np.zeros(n))
+
+    def add_vertex_normal(self, nx: np.array, ny: np.array, nz: np.array):
+        """Add normal vectors at the vertices.
+
+        The three arguments are all numpy arrays of float type and have
+        the same length.
+
+        Args:
+            nx (`numpy.array(float)`): x-coordinates of the normal vectors.
+            ny (`numpy.array(float)`): y-coordinates of the normal vectors.
+            nz (`numpy.array(float)`): z-coordinates of the normal vectors.
+        """
+        self.add_vertex_channel("nx", "float", nx)
+        self.add_vertex_channel("ny", "float", ny)
+        self.add_vertex_channel("nz", "float", nz)
+
+    # TODO active and refactor later if user feedback indicates the necessity for a compact the input list
+    # pass ti vector/matrix field directly
+    # def add_vertex_normal(self, normal):
+    #     assert isinstance(normal, (np.ndarray, ti.Matrix))
+    #     if not isinstance(normal, np.ndarray):
+    #         normal = normal.to_numpy()
+    #     dim = normal.shape[normal.ndim-1]
+    #     assert dim == 3, "Only 3D normal is supported"
+    #     n = normal.size // dim
+    #     assert n == self.num_vertices, "Size of the input is not correct"
+    #     normal = np.reshape(normal, (n, dim))
+    #     self.add_vertex_channel("nx", "float", normal[:, 0])
+    #     self.add_vertex_channel("ny", "float", normal[:, 1])
+    #     self.add_vertex_channel("nz", "float", normal[:, 2])
+
+    def add_vertex_vel(self, vx: np.array, vy: np.array, vz: np.array):
+        """Add velocity vectors at the vertices.
+
+        Args:
+            vx (`numpy.array(float)`): x-coordinates of the velocity vectors.
+            vy (`numpy.array(float)`): y-coordinates of the velocity vectors.
+            vz (`numpy.array(float)`): z-coordinates of the velocity vectors.
+        """
+        self.add_vertex_channel("vx", "float", vx)
+        self.add_vertex_channel("vy", "float", vy)
+        self.add_vertex_channel("vz", "float", vz)
+
+    def add_vertex_color(self, r: np.array, g: np.array, b: np.array):
+        """Sets the (r, g, b) channels of the colors at the vertices.
+
+        The three arguments are all numpy arrays of float type and have
+        the same length.
+
+        Args:
+            r (`numpy.array(float)`): the r-channel (red) of the colors.
+            g (`numpy.array(float)`): the g-channel (green) of the color.
+            b (`numpy.array(float)`): the b-channel (blue) of the colors.
+        """
+        self.add_vertex_channel("red", "float", r)
+        self.add_vertex_channel("green", "float", g)
+        self.add_vertex_channel("blue", "float", b)
+
+    def add_vertex_alpha(self, alpha: np.array):
+        """Sets the alpha-channel (transparent) of the vertex colors.
+
+        Args:
+            alpha (`numpy.array(float)`): the alpha-channel (transparent) of the colors.
+        """
+        self.add_vertex_channel("Alpha", "float", alpha)
+
+    def add_vertex_rgba(self, r: np.array, g: np.array, b: np.array, a: np.array):
+        """Sets the (r, g, b, a) channels of the colors at the vertices.
+
+        Args:
+            r (`numpy.array(float)`): the r-channel (red) of the colors.
+            g (`numpy.array(float)`): the g-channel (green) of the color.
+            b (`numpy.array(float)`): the b-channel (blue) of the colors.
+            a (`numpy.array(float)`): the a-channel (alpha) of the colors.
+        """
+        self.add_vertex_channel("red", "float", r)
+        self.add_vertex_channel("green", "float", g)
+        self.add_vertex_channel("blue", "float", b)
+        self.add_vertex_channel("Alpha", "float", a)
+
+    # TODO active and refactor later if user feedback indicates the necessity for a compact the input list
+    # pass ti vector/matrix field directly
+    # def add_vertex_color(self, color):
+    #     assert isinstance(color, (np.ndarray, ti.Matrix))
+    #     if not isinstance(color, np.ndarray):
+    #         color = color.to_numpy()
+    #     channels = color.shape[color.ndim-1]
+    #     assert channels == 3 or channels == 4, "The dimension for color should be either be 3 (rgb) or 4 (rgba)"
+    #     n = color.size // channels
+    #     assert n == self.num_vertices, "Size of the input is not correct"
+    #     color = np.reshape(color, (n, channels))
+    #     self.add_vertex_channel("red", "float", color[:, 0])
+    #     self.add_vertex_channel("green", "float", color[:, 1])
+    #     self.add_vertex_channel("blue", "float", color[:, 2])
+    #     if channels == 4:
+    #         self.add_vertex_channel("Alpha", "float", color[:, 3])
+
+    def add_vertex_id(self):
+        """Sets the ids of the vertices.
+
+        The id of a vertex is equal to its index in the vertex array.
+        """
+        self.add_vertex_channel("id", "int", np.arange(self.num_vertices))
+
+    def add_vertex_piece(self, piece: np.array):
+        self.add_vertex_channel("piece", "int", piece)
+
+    def add_faces(self, indices: np.array):
+        if self.face_type == "tri":
+            vert_per_face = 3
+        else:
+            vert_per_face = 4
+        assert vert_per_face * self.num_faces == indices.size, "The dimension of the face vertices is not correct"
+        self.face_indices = np.reshape(indices, (self.num_faces, vert_per_face))
+        self.face_indices = self.face_indices.astype(np.int32)
+
+    def add_face_channel(self, key: str, data_type: str, data: np.array):
+        if data_type not in self.ply_supported_types:
+            print("Unknown type " + data_type + " detected, skipping this channel")
+            return
+        if data.ndim == 1:
+            assert data.size == self.num_faces, "The dimension of the face channel is not correct"
+            self.num_face_channels += 1
+            if key in self.face_channels:
+                print("WARNING: duplicate key " + key + " detected")
+            self.face_channels.append(key)
+            self.face_data_type.append(data_type)
+            self.face_data.append(self.type_map[data_type](data))
+        else:
+            num_col = data.size // self.num_faces
+            assert (
+                data.ndim == 2 and data.size == num_col * self.num_faces
+            ), "The dimension of the face channel is not correct"
+            data.shape = (self.num_faces, num_col)
+            self.num_face_channels += num_col
+            for i in range(num_col):
+                item_key = key + "_" + str(i + 1)
+                if item_key in self.face_channels:
+                    print("WARNING: duplicate key " + item_key + " detected")
+                self.face_channels.append(item_key)
+                self.face_data_type.append(data_type)
+                self.face_data.append(self.type_map[data_type](data[:, i]))
+
+    def add_face_id(self):
+        self.add_face_channel("id", "int", np.arange(self.num_faces))
+
+    def add_face_piece(self, piece: np.array):
+        self.add_face_channel("piece", "int", piece)
+
+    def sanity_check(self):
+        assert "x" in self.vertex_channels, "The vertex pos channel is missing"
+        assert "y" in self.vertex_channels, "The vertex pos channel is missing"
+        assert "z" in self.vertex_channels, "The vertex pos channel is missing"
+        if self.num_faces > 0:
+            for idx in self.face_indices.flatten():
+                assert idx >= 0 and idx < self.num_vertices, "The face indices are invalid"
+
+    def print_header(self, path: str, _format: str):
+        with open(path, "w") as f:
+            f.writelines(
+                [
+                    "ply\n",
+                    "format " + _format + " 1.0\n",
+                    "comment " + self.comment + "\n",
+                ]
+            )
+            f.write("element vertex " + str(self.num_vertices) + "\n")
+            for i in range(self.num_vertex_channels):
+                f.write("property " + self.vertex_data_type[i] + " " + self.vertex_channels[i] + "\n")
+            if self.num_faces != 0:
+                f.write("element face " + str(self.num_faces) + "\n")
+                f.write("property list uchar int vertex_indices\n")
+                for i in range(self.num_face_channels):
+                    f.write("property " + self.face_data_type[i] + " " + self.face_channels[i] + "\n")
+            f.write("end_header\n")
+
+    def export(self, path):
+        self.sanity_check()
+        self.print_header(path, "binary_" + sys.byteorder + "_endian")
+        with open(path, "ab") as f:
+            for i in range(self.num_vertices):
+                for j in range(self.num_vertex_channels):
+                    f.write(self.vertex_data[j][i])
+            if self.face_type == "tri":
+                vert_per_face = np.uint8(3)
+            else:
+                vert_per_face = np.uint8(4)
+            for i in range(self.num_faces):
+                f.write(vert_per_face)
+                for j in range(vert_per_face):
+                    f.write(self.face_indices[i, j])
+                for j in range(self.num_face_channels):
+                    f.write(self.face_data[j][i])
+
+    def export_ascii(self, path):
+        self.sanity_check()
+        self.print_header(path, "ascii")
+        with open(path, "a") as f:
+            for i in range(self.num_vertices):
+                for j in range(self.num_vertex_channels):
+                    f.write(str(self.vertex_data[j][i]) + " ")
+                f.write("\n")
+            if self.face_type == "tri":
+                vert_per_face = 3
+            else:
+                vert_per_face = 4
+            for i in range(self.num_faces):
+                f.writelines(
+                    [
+                        str(vert_per_face) + " ",
+                        " ".join(map(str, self.face_indices[i, :])),
+                        " ",
+                    ]
+                )
+                for j in range(self.num_face_channels):
+                    f.write(str(self.face_data[j][i]) + " ")
+                f.write("\n")
+
+    def export_frame_ascii(self, series_num: int, path: str):
+        # if path has ply ending
+        last_4_char = path[-4:]
+        if last_4_char == ".ply":
+            path = path[:-4]
+
+        real_path = path + "_" + f"{series_num:0=6d}" + ".ply"
+        self.export_ascii(real_path)
+
+    def export_frame(self, series_num: int, path: str):
+        # if path has ply ending
+        last_4_char = path[-4:]
+        if last_4_char == ".ply":
+            path = path[:-4]
+
+        real_path = path + "_" + f"{series_num:0=6d}" + ".ply"
+        self.export(real_path)
+
+
+__all__ = ["PLYWriter"]

taichi/tools/vtk.py

@@ -0,0 +1,38 @@
+# type: ignore
+
+import numpy as np
+
+
+def write_vtk(scalar_field, filename):
+    try:
+        from pyevtk.hl import gridToVTK  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise RuntimeError(
+            "Failed to import pyevtk. Please install it via /\
+        `pip install pyevtk` first. "
+        )
+
+    scalar_field_np = scalar_field.to_numpy()
+    field_shape = scalar_field_np.shape
+    dimensions = len(field_shape)
+
+    if dimensions not in (2, 3):
+        raise ValueError("The input field must be a 2D or 3D scalar field.")
+
+    if dimensions == 2:
+        scalar_field_np = scalar_field_np[np.newaxis, :, :]
+        zcoords = np.array([0, 1])
+    elif dimensions == 3:
+        zcoords = np.arange(0, field_shape[2])
+    else:
+        raise ValueError("dimensions should be 2 or 3")
+    gridToVTK(
+        filename,
+        x=np.arange(0, field_shape[0]),
+        y=np.arange(0, field_shape[1]),
+        z=zcoords,
+        cellData={filename: scalar_field_np},
+    )
+
+
+__all__ = ["write_vtk"]

taichi/types/__init__.py

@@ -0,0 +1,19 @@
+# type: ignore
+
+"""
+This module defines data types in Taichi:
+
+- primitive: int, float, etc.
+- compound: matrix, vector, struct.
+- template: for reference types.
+- ndarray: for arbitrary arrays.
+- quant: for quantized types, see "https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf"
+"""
+
+from taichi.types import quant
+from taichi.types.annotations import *
+from taichi.types.compound_types import *
+from taichi.types.ndarray_type import *
+from taichi.types.primitive_types import *
+from taichi.types.texture_type import *
+from taichi.types.utils import *

taichi/types/annotations.py

@@ -0,0 +1,47 @@
+# type: ignore
+
+
+class Template:
+    """Type annotation for template kernel parameter.
+    Useful for passing parameters to kernels by reference.
+
+    See also https://docs.taichi-lang.org/docs/meta.
+
+    Args:
+        tensor (Any): unused
+        dim (Any): unused
+
+    Example::
+
+        >>> a = 1
+        >>>
+        >>> @ti.kernel
+        >>> def test():
+        >>>     print(a)
+        >>>
+        >>> @ti.kernel
+        >>> def test_template(a: ti.template()):
+        >>>     print(a)
+        >>>
+        >>> test(a)  # will print 1
+        >>> test_template(a)  # will also print 1
+        >>> a = 2
+        >>> test(a)  # will still print 1
+        >>> test_template(a)  # will print 2
+    """
+
+    def __init__(self, tensor=None, dim=None):
+        self.tensor = tensor
+        self.dim = dim
+
+
+template = Template
+"""Alias for :class:`~taichi.types.annotations.Template`.
+"""
+
+
+class sparse_matrix_builder:
+    pass
+
+
+__all__ = ["template", "sparse_matrix_builder", "Template"]

taichi/types/compound_types.py

@@ -0,0 +1,90 @@
+# type: ignore
+
+import taichi
+from taichi._lib.utils import ti_python_core as _ti_python_core
+
+_type_factory = _ti_python_core.get_type_factory_instance()
+
+
+class CompoundType:
+    def from_kernel_struct_ret(self, launch_ctx, index: tuple):
+        raise NotImplementedError()
+
+
+# TODO: maybe move MatrixType, StructType here to avoid the circular import?
+def matrix(n=None, m=None, dtype=None):
+    """Creates a matrix type with given shape and data type.
+
+    Args:
+        n (int): number of rows of the matrix.
+        m (int): number of columns of the matrix.
+        dtype (:mod:`~taichi.types.primitive_types`): matrix data type.
+
+    Returns:
+        A matrix type.
+
+    Example::
+
+        >>> mat2x2 = ti.types.matrix(2, 2, ti.f32)  # 2x2 matrix type
+        >>> M = mat2x2([[1., 2.], [3., 4.]])  # an instance of this type
+    """
+    return taichi.lang.matrix.MatrixType(n, m, 2, dtype)
+
+
+def vector(n=None, dtype=None):
+    """Creates a vector type with given shape and data type.
+
+    Args:
+        n (int): dimension of the vector.
+        dtype (:mod:`~taichi.types.primitive_types`): vector data type.
+
+    Returns:
+        A vector type.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)  # 3d vector type
+        >>> v = vec3([1., 2., 3.])  # an instance of this type
+    """
+    return taichi.lang.matrix.VectorType(n, dtype)
+
+
+def struct(**kwargs):
+    """Creates a struct type with given members.
+
+    Args:
+        kwargs (dict): a dictionary contains the names and types of the
+            struct members.
+
+    Returns:
+        A struct type.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)
+        >>> sphere = ti.types.struct(center=vec3, radius=float)
+        >>> s = sphere(center=vec3([0., 0., 0.]), radius=1.0)
+    """
+    return taichi.lang.struct.StructType(**kwargs)
+
+
+def argpack(**kwargs):
+    """Creates an argument pack type with given members.
+
+    Args:
+        kwargs (dict): a dictionary contains the names and types of the
+            argument pack members.
+
+    Returns:
+        A argument pack type.
+
+    Example::
+
+        >>> vec3 = ti.types.vector(3, ti.f32)
+        >>> sphere = ti.types.argpack(center=vec3, radius=float)
+        >>> s = sphere(center=vec3([0., 0., 0.]), radius=1.0)
+    """
+    return taichi.lang.argpack.ArgPackType(**kwargs)
+
+
+__all__ = ["matrix", "vector", "struct", "argpack"]