gstaichi 0.1.18.dev1__cp310-cp310-win_amd64.whl

taichi/lang/_ndrange.py

@@ -0,0 +1,152 @@
+# type: ignore
+
+import collections.abc
+from typing import Iterable
+
+import numpy as np
+
+from taichi.lang import ops
+from taichi.lang.exception import TaichiSyntaxError, TaichiTypeError
+from taichi.lang.expr import Expr
+from taichi.lang.matrix import Matrix
+from taichi.types.utils import is_integral
+
+
+class _Ndrange:
+    def __init__(self, *args):
+        args = list(args)
+        for i, arg in enumerate(args):
+            if not isinstance(arg, collections.abc.Sequence):
+                args[i] = (0, arg)
+            if len(args[i]) != 2:
+                raise TaichiSyntaxError(
+                    "Every argument of ndrange should be a scalar or a tuple/list like (begin, end)"
+                )
+            args[i] = (args[i][0], ops.max(args[i][0], args[i][1]))
+        for arg in args:
+            for bound in arg:
+                if not isinstance(bound, (int, np.integer)) and not (
+                    isinstance(bound, Expr) and is_integral(bound.ptr.get_rvalue_type())
+                ):
+                    raise TaichiTypeError(
+                        "Every argument of ndrange should be an integer scalar or a tuple/list of (int, int)"
+                    )
+        self.bounds = args
+
+        self.dimensions = [None] * len(args)
+        for i, bound in enumerate(self.bounds):
+            self.dimensions[i] = bound[1] - bound[0]
+
+        self.acc_dimensions = self.dimensions.copy()
+        for i in reversed(range(len(self.bounds) - 1)):
+            self.acc_dimensions[i] = self.acc_dimensions[i] * self.acc_dimensions[i + 1]
+        if len(self.acc_dimensions) == 0:  # for the empty case, e.g. ti.ndrange()
+            self.acc_dimensions = [1]
+
+    def __iter__(self):
+        def gen(d, prefix):
+            if d == len(self.bounds):
+                yield prefix
+            else:
+                for t in range(self.bounds[d][0], self.bounds[d][1]):
+                    yield from gen(d + 1, prefix + (t,))
+
+        yield from gen(0, ())
+
+    def grouped(self):
+        return GroupedNDRange(self)
+
+
+def ndrange(*args) -> Iterable:
+    """Return an immutable iterator object for looping over multi-dimensional indices.
+
+    This returned set of multi-dimensional indices is the direct product (in the set-theory sense)
+    of n groups of integers, where n equals the number of arguments in the input list, and looks like
+
+    range(x1, y1) x range(x2, y2) x ... x range(xn, yn)
+
+    The k-th argument corresponds to the k-th `range()` factor in the above product, and each
+    argument must be an integer or a pair of two integers. An integer argument n will be interpreted
+    as `range(0, n)`, and a pair of two integers (start, end) will be interpreted as `range(start, end)`.
+
+    You can loop over these multi-dimensonal indices in different ways, see the examples below.
+
+    Args:
+        entries: (int, tuple): Must be either an integer, or a tuple/list of two integers.
+
+    Returns:
+        An immutable iterator object.
+
+    Example::
+
+        You can loop over 1-D integers in range [start, end), as in native Python
+
+            >>> @ti.kernel
+            >>> def loop_1d():
+            >>>     start = 2
+            >>>     end = 5
+            >>>     for i in ti.ndrange((start, end)):
+            >>>         print(i)  # will print 2 3 4
+
+        Note the braces around `(start, end)` in the above code. If without them,
+        the parameter `2` will be interpreted as `range(0, 2)`, `5` will be
+        interpreted as `range(0, 5)`, and you will get a set of 2-D indices which
+        contains 2x5=10 elements, and need two indices i, j to loop over them:
+
+            >>> @ti.kernel
+            >>> def loop_2d():
+            >>>     for i, j in ti.ndrange(2, 5):
+            >>>         print(i, j)
+            0 0
+            ...
+            0 4
+            ...
+            1 4
+
+        But you do can use a single index i to loop over these 2-D indices, in this case
+        the indices are returned as a 1-D array `(0, 1, ..., 9)`:
+
+            >>> @ti.kernel
+            >>> def loop_2d_as_1d():
+            >>>     for i in ti.ndrange(2, 5):
+            >>>         print(i)
+            will print 0 1 2 3 4 5 6 7 8 9
+
+        In general, you can use any `1 <= k <= n` iterators to loop over a set of n-D
+        indices. For `k=n` all the indices are n-dimensional, and they are returned in
+        lexical order, but for `k<n` iterators the last n-k+1 dimensions will be collapsed into
+        a 1-D array of consecutive integers `(0, 1, 2, ...)` whose length equals the
+        total number of indices in the last n-k+1 dimensions:
+
+            >>> @ti.kernel
+            >>> def loop_3d_as_2d():
+            >>>     # use two iterators to loop over a set of 3-D indices
+            >>>     # the last two dimensions for 4, 5 will collapse into
+            >>>     # the array [0, 1, 2, ..., 19]
+            >>>     for i, j in ti.ndrange(3, 4, 5):
+            >>>         print(i, j)
+            will print 0 0, 0 1, ..., 0 19, ..., 2 19.
+
+        A typical usage of `ndrange` is when you want to loop over a tensor and process
+        its entries in parallel. You should avoid writing nested `for` loops here since
+        only top level `for` loops are paralleled in taichi, instead you can use `ndrange`
+        to hold all entries in one top level loop:
+
+            >>> @ti.kernel
+            >>> def loop_tensor():
+            >>>     for row, col, channel in ti.ndrange(image_height, image_width, channels):
+            >>>         image[row, col, channel] = ...
+    """
+    return _Ndrange(*args)
+
+
+class GroupedNDRange:
+    def __init__(self, r):
+        self.r = r
+
+    def __iter__(self):
+        for ind in self.r:
+            yield Matrix(list(ind))
+
+
+__all__ = ["ndrange"]

taichi/lang/_texture.py

@@ -0,0 +1,172 @@
+# type: ignore
+
+import numpy as np
+
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.expr import Expr, make_expr_group
+from taichi.lang.matrix import Matrix
+from taichi.lang.util import taichi_scope
+from taichi.types import vector
+from taichi.types.primitive_types import f32
+
+
+def _get_entries(mat):
+    if isinstance(mat, Matrix):
+        return mat.entries
+    return [mat]
+
+
+class TextureSampler:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @taichi_scope
+    def sample_lod(self, uv, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(uv), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kSampleLod, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @taichi_scope
+    def fetch(self, index, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kFetchTexel, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+
+class RWTextureAccessor:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        # taichi_python.TexturePtrExpression.
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @taichi_scope
+    def load(self, index):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index))
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kLoad, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @taichi_scope
+    def store(self, index, value):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), *_get_entries(value))
+        impl.expr_init(
+            ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kStore, self.ptr_expr, args_group, dbg_info)
+        )
+
+    @property
+    @taichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr_expr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        ret = [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr_expr, i, dbg_info)) for i in range(dim)]
+        return ret
+
+    @taichi_scope
+    def _loop_range(self):
+        """Gets the corresponding taichi_python.Expr to serve as loop range.
+
+        Returns:
+            taichi_python.Expr: See above.
+        """
+        return self.ptr_expr
+
+
+class Texture:
+    """Taichi Texture class.
+
+    Args:
+        fmt (ti.Format): Color format of the texture.
+        shape (Tuple[int]): Shape of the Texture.
+    """
+
+    def __init__(self, fmt, arr_shape):
+        self.tex = impl.get_runtime().prog.create_texture(fmt, arr_shape)
+        self.fmt = fmt
+        self.num_dims = len(arr_shape)
+        self.shape = arr_shape
+
+    def from_ndarray(self, ndarray):
+        """Loads an ndarray to texture.
+
+        Args:
+            ndarray (ti.Ndarray): Source ndarray to load from.
+        """
+        self.tex.from_ndarray(ndarray.arr)
+
+    def from_field(self, field):
+        """Loads a field to texture.
+
+        Args:
+            field (ti.Field): Source field to load from.
+        """
+        self.tex.from_snode(field.snode.ptr)
+
+    def _device_allocation_ptr(self):
+        return self.tex.device_allocation_ptr()
+
+    def from_image(self, image):
+        """Loads a PIL image to texture. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Args:
+            image (PIL.Image.Image): Source PIL image to load from.
+
+        """
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        assert isinstance(image, Image.Image)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        assert image.size == tuple(self.shape)
+
+        assert self.num_dims == 2
+        # Don't use transpose method since its enums are too new
+        image = image.rotate(90, expand=True)
+        arr = np.asarray(image)
+        from taichi._kernels import (  # pylint: disable=import-outside-toplevel
+            load_texture_from_numpy,
+        )
+
+        load_texture_from_numpy(self, arr)
+
+    def to_image(self):
+        """Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Returns:
+            img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture.
+        """
+        assert self.num_dims == 2
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        res = np.zeros(self.shape + (3,), np.uint8)
+        from taichi._kernels import (  # pylint: disable=import-outside-toplevel
+            save_texture_to_numpy,
+        )
+
+        save_texture_to_numpy(self, res)
+        return Image.fromarray(res).rotate(270, expand=True)

taichi/lang/_wrap_inspect.py

@@ -0,0 +1,189 @@
+# type: ignore
+
+# Taichi's custom inspect module.
+# This module is used by Taichi's ast transformer to parse the source code.
+# Currently this module is aimed for working in the following modes:
+# 1. Usual Python/IPython mode, e.g. python script.py
+#    In this case we mainly rely on the built-in `inspect` module, except
+#    we need some hacks when we are in IPython mode and there is a cell magic.
+# 2. Blender's scripting mode, e.g. Users write Taichi code in the scripting
+#    window in Blender and press the run button. In this case we need to
+#    retrieve the source using Blender's `bpy.data.texts` and write it to a temp
+#    file so that the inspect module can parse.
+# 3. The interactive shell mode, e.g. Users directly type their code in the
+#    interactive shell. In this case we use `dill` to get the source.
+#
+# NB: Running Taichi in other modes are likely not supported.
+
+import atexit
+import inspect
+import os
+import tempfile
+
+import dill
+
+_builtin_getfile = inspect.getfile
+_builtin_findsource = inspect.findsource
+
+
+def _find_source_with_custom_getfile_func(func, obj):
+    """Use a custom function `func` to replace inspect's `getfile`, return the
+    source found by the new routine and restore the original `getfile` back.
+    """
+    inspect.getfile = func  # replace with our custom func
+    source = inspect.findsource(obj)
+    inspect.getfile = _builtin_getfile  # restore
+    return source
+
+
+def _blender_get_text_name(filename: str):
+    """Extract filename from path in the Blender mode."""
+    # In Blender's scripting mode, unsaved files are named
+    # like `/Text`, `/Text.001`, `/test.py`, etc.
+    # We simply remove this path seperator.
+    if filename.startswith(os.path.sep) and filename.count(os.path.sep) == 1:
+        return filename[1:]  # "/Text.001" --> "Text.001"
+
+    # Saved text files are named like `some-path/xxx.blend/Text` or
+    # `some-path/xxx.blend/test.py`
+    # We drop the path and extract the filename with extension.
+    index = filename.rfind(".blend" + os.path.sep)
+    if index != -1:
+        return filename[index + 7 :]  # "xxx.blend/test.py" --> "test.py"
+
+    return None
+
+
+def _blender_findsource(obj):
+    try:
+        import bpy  # pylint: disable=import-outside-toplevel
+    except:
+        raise ImportError("Not in Blender environment!")
+
+    # Inspect's built-in `getfile` returns the filename like
+    # `/Text`, `/Text.001`, `some-path/xxx.blend/test.py`
+    # This filename may not be a full valid path.
+    filename = _builtin_getfile(obj)
+    # Extract the text name without path
+    text_name = _blender_get_text_name(filename)
+    if text_name is None:
+        raise IOError("Object `{obj.__name__}` is not defined in a .blend file!")
+    # Get the lines of code via text_name
+    lines = bpy.data.texts[text_name].as_string()
+    # Now we have found the lines of code.
+    # We first check if they are already cached, to avoid file io in each query.
+    try:
+        filename = _blender_findsource._saved_inspect_cache[lines]  # pylint: disable=no-member
+    except KeyError:
+        # Save the code to a valid path.
+        fd, filename = tempfile.mkstemp(prefix="_Blender_", suffix=f"_{text_name}.py")
+        os.close(fd)
+
+        with open(filename, "w") as f:
+            f.write(lines)
+
+        _blender_findsource._saved_inspect_cache[lines] = filename  # pylint: disable=no-member
+        atexit.register(os.unlink, filename)  # Remove file when program exits
+
+    # Our custom getfile function
+    def wrapped_getfile(ob):
+        if id(ob) == id(obj):
+            return filename
+
+        return _builtin_getfile(ob)
+
+    return _find_source_with_custom_getfile_func(wrapped_getfile, obj)
+
+
+_blender_findsource._saved_inspect_cache = {}
+
+
+def _Python_IPython_findsource(obj):
+    try:
+        # In Python and IPython the builtin inspect would suffice in most cases
+        return _builtin_findsource(obj)
+    except IOError:
+        # Except that the cell has a magic command like %%time or %%timeit
+        # In this case the filename returned by the built-in's getfile is wrong,
+        # it becomes something like `<timed exec>` or `<magic-timeit>`.
+        filename = _builtin_getfile(obj)
+        if filename in {"<timed exec>", "<magic-timeit>"}:
+            try:
+                ip = get_ipython()
+                if ip is not None:
+                    # So we are in IPython's cell magic
+                    session_id = ip.history_manager.get_last_session_id()
+                    fd, filename = tempfile.mkstemp(prefix="_IPython_", suffix=f"_{session_id}.py")
+                    os.close(fd)
+                    # The latest lines of code can be retrived from here
+                    lines = ip.history_manager._i00
+
+                    # `lines` is a string that also contains the cell magic
+                    # command, we need to remove the magic command
+                    # (and spaces/sep around it) to obtain a valid Python code
+                    # snippet before saving it to a file
+                    index = lines.find("%time")
+                    lines_stripped = lines[index:]
+                    lines_stripped = lines_stripped.split(maxsplit=1)[1]
+
+                    with open(filename, "w") as f:
+                        f.write(lines_stripped)
+
+                    atexit.register(os.unlink, filename)  # Remove the file after the program exits
+                    func = lambda obj: filename
+                    return _find_source_with_custom_getfile_func(func, obj)
+
+            except ImportError:
+                pass
+        raise IOError(
+            f"Cannot find source code for Object: {obj}, it's likely \
+you are not running Taichi from command line or IPython."
+        )
+
+
+def _REPL_findsource(obj):
+    """Findsource in the interactive shell mode."""
+    return dill.source.findsource(obj)
+
+
+def _custom_findsource(obj):
+    try:
+        return _Python_IPython_findsource(obj)
+    except IOError:
+        try:
+            return _REPL_findsource(obj)
+        except:
+            try:
+                return _blender_findsource(obj)
+            except:
+                raise IOError(
+                    f"Cannot find source code for Object: {obj}, this \
+is possibly because of you are running Taichi in an environment that Taichi's own \
+inspect module cannot find the source. Please report an issue to help us fix: \
+https://github.com/taichi-dev/taichi/issues"
+                )
+
+
+class _InspectContextManager:
+    def __enter__(self):
+        inspect.findsource = _custom_findsource
+        return self
+
+    def __exit__(self, *_):
+        inspect.findsource = _builtin_findsource
+
+
+def getsourcelines(obj):
+    with _InspectContextManager():
+        return inspect.getsourcelines(obj)
+
+
+def getsourcefile(obj):
+    with _InspectContextManager():
+        ret = inspect.getsourcefile(obj)
+        if ret is None:
+            ret = inspect.getfile(obj)
+        return ret
+
+
+__all__ = ["getsourcelines", "getsourcefile"]

taichi/lang/any_array.py

@@ -0,0 +1,99 @@
+# type: ignore
+
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.expr import Expr, make_expr_group
+from taichi.lang.util import taichi_scope
+from taichi.types.enums import Layout
+from taichi.types.ndarray_type import NdarrayTypeMetadata
+
+
+class AnyArray:
+    """Class for arbitrary arrays in Python AST.
+
+    Args:
+        ptr (taichi_python.Expr): A taichi_python.Expr wrapping a taichi_python.ExternalTensorExpression.
+        element_shape (Tuple[Int]): () if scalar elements (default), (n) if vector elements, and (n, m) if matrix elements.
+        layout (Layout): Memory layout.
+    """
+
+    def __init__(self, ptr):
+        assert ptr.is_external_tensor_expr()
+        self.ptr = ptr
+        self.ptr.type_check(impl.get_runtime().prog.config())
+
+    def element_shape(self):
+        return _ti_core.get_external_tensor_element_shape(self.ptr)
+
+    def layout(self):
+        # 0: scalar; 1: vector (SOA); 2: matrix (SOA); -1: vector
+        # (AOS); -2: matrix (AOS)
+        element_dim = _ti_core.get_external_tensor_element_dim(self.ptr)
+        if element_dim == 1 or element_dim == 2:
+            return Layout.SOA
+        return Layout.AOS
+
+    def get_type(self):
+        return NdarrayTypeMetadata(
+            _ti_core.get_external_tensor_element_type(self.ptr), None, _ti_core.get_external_tensor_needs_grad(self.ptr)
+        )  # AnyArray can take any shape
+
+    @property
+    @taichi_scope
+    def grad(self):
+        """Returns the gradient of this array."""
+        return AnyArray(_ti_core.make_external_tensor_grad_expr(self.ptr))
+
+    @property
+    @taichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        return [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr, i, dbg_info)) for i in range(dim)]
+
+    @taichi_scope
+    def _loop_range(self):
+        """Gets the corresponding taichi_python.Expr to serve as loop range.
+
+        Returns:
+            taichi_python.Expr: See above.
+        """
+        return self.ptr
+
+
+class AnyArrayAccess:
+    """Class for first-level access to AnyArray with Vector/Matrix elements in Python AST.
+
+    Args:
+        arr (AnyArray): See above.
+        indices_first (Tuple[Int]): Indices of first-level access.
+    """
+
+    def __init__(self, arr, indices_first):
+        self.arr = arr
+        self.indices_first = indices_first
+
+    @taichi_scope
+    def subscript(self, i, j):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+
+        indices_second = (i,) if len(self.arr.element_shape()) == 1 else (i, j)
+        if self.arr.layout() == Layout.SOA:
+            indices = indices_second + self.indices_first
+        else:
+            indices = self.indices_first + indices_second
+        return Expr(
+            ast_builder.expr_subscript(
+                self.arr.ptr,
+                make_expr_group(*indices),
+                _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()),
+            )
+        )
+
+
+__all__ = []