gstaichi 2.1.1rc3__cp310-cp310-macosx_11_0_arm64.whl

gstaichi/linalg/sparse_matrix.py

@@ -0,0 +1,303 @@
+# type: ignore
+
+from functools import reduce
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang._ndarray import Ndarray, ScalarNdarray
+from gstaichi.lang.exception import GsTaichiRuntimeError
+from gstaichi.lang.field import Field
+from gstaichi.lang.impl import get_runtime
+from gstaichi.types import f32
+
+
+class SparseMatrix:
+    """GsTaichi's Sparse Matrix class
+
+    A sparse matrix allows the programmer to solve a large linear system.
+
+    Args:
+        n (int): the first dimension of a sparse matrix.
+        m (int): the second dimension of a sparse matrix.
+        sm (SparseMatrix): another sparse matrix that will be built from.
+    """
+
+    def __init__(self, n=None, m=None, sm=None, dtype=f32, storage_format="col_major"):
+        self.dtype = dtype
+        if sm is None:
+            self.n = n
+            self.m = m if m else n
+            self.matrix = get_runtime().prog.create_sparse_matrix(n, m, dtype, storage_format)
+        else:
+            self.n = sm.num_rows()
+            self.m = sm.num_cols()
+            self.matrix = sm
+
+    def __iadd__(self, other):
+        """Addition operation for sparse matrix.
+
+        Returns:
+            The result sparse matrix of the addition.
+        """
+        assert (
+            self.n == other.n and self.m == other.m
+        ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        self.matrix += other.matrix
+        return self
+
+    def __add__(self, other):
+        """Addition operation for sparse matrix.
+
+        Returns:
+            The result sparse matrix of the addition.
+        """
+        assert (
+            self.n == other.n and self.m == other.m
+        ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        sm = self.matrix + other.matrix
+        return SparseMatrix(sm=sm)
+
+    def __isub__(self, other):
+        """Subtraction operation for sparse matrix.
+
+        Returns:
+             The result sparse matrix of the subtraction.
+        """
+        assert (
+            self.n == other.n and self.m == other.m
+        ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        self.matrix -= other.matrix
+        return self
+
+    def __sub__(self, other):
+        """Subtraction operation for sparse matrix.
+
+        Returns:
+             The result sparse matrix of the subtraction.
+        """
+        assert (
+            self.n == other.n and self.m == other.m
+        ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        sm = self.matrix - other.matrix
+        return SparseMatrix(sm=sm)
+
+    def __mul__(self, other):
+        """Sparse matrix's multiplication against real numbers or the hadamard product against another matrix
+
+        Args:
+            other (float or SparseMatrix): the other operand of multiplication.
+        Returns:
+            The result of multiplication.
+        """
+        if isinstance(other, float):
+            sm = other * self.matrix
+            return SparseMatrix(sm=sm)
+        if isinstance(other, SparseMatrix):
+            assert (
+                self.n == other.n and self.m == other.m
+            ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+            sm = self.matrix * other.matrix
+            return SparseMatrix(sm=sm)
+
+        return None
+
+    def __rmul__(self, other):
+        """Right scalar multiplication for sparse matrix.
+
+        Args:
+            other (float): the other operand of scalar multiplication.
+        Returns:
+            The result of multiplication.
+        """
+        if isinstance(other, float):
+            sm = self.matrix * other
+            return SparseMatrix(sm=sm)
+
+        return None
+
+    def transpose(self):
+        """Sparse Matrix transpose.
+
+        Returns:
+            The transposed sparse mastrix.
+        """
+        sm = self.matrix.transpose()
+        return SparseMatrix(sm=sm)
+
+    def __matmul__(self, other):
+        """Matrix multiplication.
+
+        Args:
+            other (SparseMatrix, Field, or numpy.array): the other sparse matrix of the multiplication.
+        Returns:
+            The result of matrix multiplication.
+        """
+        if isinstance(other, SparseMatrix):
+            assert (
+                self.m == other.n
+            ), f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+            sm = self.matrix.matmul(other.matrix)
+            return SparseMatrix(sm=sm)
+        if isinstance(other, Field):
+            assert (
+                self.m == other.shape[0]
+            ), f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
+            return self.matrix.mat_vec_mul(other.to_numpy())
+        if isinstance(other, np.ndarray):
+            assert (
+                self.m == other.shape[0]
+            ), f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
+            return self.matrix.mat_vec_mul(other)
+        if isinstance(other, Ndarray):
+            if self.m != other.shape[0]:
+                raise GsTaichiRuntimeError(
+                    f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
+                )
+            res = ScalarNdarray(dtype=other.dtype, arr_shape=(self.n,))
+            self.matrix.spmv(get_runtime().prog, other.arr, res.arr)
+            return res
+        raise GsTaichiRuntimeError(
+            f"Sparse matrix-matrix/vector multiplication does not support {type(other)} for now. Supported types are SparseMatrix, ti.field, and numpy ndarray."
+        )
+
+    def __getitem__(self, indices):
+        return self.matrix.get_element(indices[0], indices[1])
+
+    def __setitem__(self, indices, value):
+        self.matrix.set_element(indices[0], indices[1], value)
+
+    def __str__(self):
+        """Python scope matrix print support."""
+        return self.matrix.to_string()
+
+    def __repr__(self):
+        return self.matrix.to_string()
+
+    @property
+    def shape(self):
+        """The shape of the sparse matrix."""
+        return (self.n, self.m)
+
+    def build_from_ndarray(self, ndarray):
+        """Build the sparse matrix from a ndarray.
+
+        Args:
+            ndarray (Union[ti.ndarray, ti.Vector.ndarray, ti.Matrix.ndarray]): the ndarray to build the sparse matrix from.
+
+        Raises:
+            GsTaichiRuntimeError: If the input is not a ndarray or the length is not divisible by 3.
+
+        Example::
+            >>> N = 5
+            >>> triplets = ti.Vector.ndarray(n=3, dtype=ti.f32, shape=10, layout=ti.Layout.AOS)
+            >>> @ti.kernel
+            >>> def fill(triplets: ti.types.ndarray()):
+            >>>     for i in range(N):
+            >>>        triplets[i] = ti.Vector([i, (i + 1) % N, i+1], dt=ti.f32)
+            >>> fill(triplets)
+            >>> A = ti.linalg.SparseMatrix(n=N, m=N, dtype=ti.f32)
+            >>> A.build_from_ndarray(triplets)
+            >>> print(A)
+            [0, 1, 0, 0, 0]
+            [0, 0, 2, 0, 0]
+            [0, 0, 0, 3, 0]
+            [0, 0, 0, 0, 4]
+            [5, 0, 0, 0, 0]
+        """
+        if isinstance(ndarray, Ndarray):
+            num_scalars = reduce(lambda x, y: x * y, ndarray.shape + ndarray.element_shape)
+            if num_scalars % 3 != 0:
+                raise GsTaichiRuntimeError("The number of ndarray elements must have a length that is divisible by 3.")
+            get_runtime().prog.make_sparse_matrix_from_ndarray(self.matrix, ndarray.arr)
+        else:
+            raise GsTaichiRuntimeError(
+                "Sparse matrix only supports building from [ti.ndarray, ti.Vector.ndarray, ti.Matrix.ndarray]"
+            )
+
+    def mmwrite(self, filename):
+        """Writes the sparse matrix to Matrix Market file-like target.
+
+        Args:
+            filename (str): the file name to write the sparse matrix to.
+        """
+        self.matrix.mmwrite(filename)
+
+
+class SparseMatrixBuilder:
+    """A python wrap around sparse matrix builder.
+
+    Use this builder to fill the sparse matrix.
+
+    Args:
+        num_rows (int): the first dimension of a sparse matrix.
+        num_cols (int): the second dimension of a sparse matrix.
+        max_num_triplets (int): the maximum number of triplets.
+        dtype (ti.dtype): the data type of the sparse matrix.
+        storage_format (str): the storage format of the sparse matrix.
+    """
+
+    def __init__(
+        self,
+        num_rows=None,
+        num_cols=None,
+        max_num_triplets=0,
+        dtype=f32,
+        storage_format="col_major",
+    ):
+        self.num_rows = num_rows
+        self.num_cols = num_cols if num_cols else num_rows
+        self.dtype = dtype
+        if num_rows is not None:
+            gstaichi_arch = get_runtime().prog.config().arch
+            if gstaichi_arch in [
+                _ti_core.Arch.x64,
+                _ti_core.Arch.arm64,
+                _ti_core.Arch.cuda,
+            ]:
+                self.ptr = _ti_core.SparseMatrixBuilder(
+                    num_rows,
+                    num_cols,
+                    max_num_triplets,
+                    dtype,
+                    storage_format,
+                )
+                self.ptr.create_ndarray(get_runtime().prog)
+            else:
+                raise GsTaichiRuntimeError("SparseMatrix only supports CPU and CUDA for now.")
+
+    def _get_addr(self):
+        """Get the address of the sparse matrix"""
+        return self.ptr.get_addr()
+
+    def _get_ndarray_addr(self):
+        """Get the address of the ndarray"""
+        return self.ptr.get_ndarray_data_ptr()
+
+    def print_triplets(self):
+        """Print the triplets stored in the builder"""
+        gstaichi_arch = get_runtime().prog.config().arch
+        if gstaichi_arch in [_ti_core.Arch.x64, _ti_core.Arch.arm64]:
+            self.ptr.print_triplets_eigen()
+        elif gstaichi_arch == _ti_core.Arch.cuda:
+            self.ptr.print_triplets_cuda()
+
+    def build(self, dtype=f32, _format="CSR"):
+        """Create a sparse matrix using the triplets"""
+        gstaichi_arch = get_runtime().prog.config().arch
+        if gstaichi_arch in [_ti_core.Arch.x64, _ti_core.Arch.arm64]:
+            sm = self.ptr.build()
+            return SparseMatrix(sm=sm, dtype=self.dtype)
+        if gstaichi_arch == _ti_core.Arch.cuda:
+            if self.dtype != f32:
+                raise GsTaichiRuntimeError("CUDA sparse matrix only supports f32.")
+            sm = self.ptr.build_cuda()
+            return SparseMatrix(sm=sm, dtype=self.dtype)
+        raise GsTaichiRuntimeError("Sparse matrix only supports CPU and CUDA backends.")
+
+    def __del__(self):
+        if get_runtime() is not None and get_runtime().prog is not None:
+            self.ptr.delete_ndarray(get_runtime().prog)
+
+
+__all__ = ["SparseMatrix", "SparseMatrixBuilder"]

gstaichi/linalg/sparse_solver.py

@@ -0,0 +1,123 @@
+# type: ignore
+
+import numpy as np
+
+import gstaichi.lang
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang._ndarray import Ndarray, ScalarNdarray
+from gstaichi.lang.exception import GsTaichiRuntimeError
+from gstaichi.lang.field import Field
+from gstaichi.lang.impl import get_runtime
+from gstaichi.linalg.sparse_matrix import SparseMatrix
+from gstaichi.types.primitive_types import f32
+
+
+class SparseSolver:
+    """Sparse linear system solver
+
+    Use this class to solve linear systems represented by sparse matrices.
+
+    Args:
+        solver_type (str): The factorization type.
+        ordering (str): The method for matrices re-ordering.
+    """
+
+    def __init__(self, dtype=f32, solver_type="LLT", ordering="AMD"):
+        self.matrix = None
+        self.dtype = dtype
+        solver_type_list = ["LLT", "LDLT", "LU"]
+        solver_ordering = ["AMD", "COLAMD"]
+        if solver_type in solver_type_list and ordering in solver_ordering:
+            gstaichi_arch = gstaichi.lang.impl.get_runtime().prog.config().arch
+            assert (
+                gstaichi_arch == _ti_core.Arch.x64
+                or gstaichi_arch == _ti_core.Arch.arm64
+                or gstaichi_arch == _ti_core.Arch.cuda
+            ), "SparseSolver only supports CPU and CUDA for now."
+            if gstaichi_arch == _ti_core.Arch.cuda:
+                self.solver = _ti_core.make_cusparse_solver(dtype, solver_type, ordering)
+            else:
+                self.solver = _ti_core.make_sparse_solver(dtype, solver_type, ordering)
+        else:
+            raise GsTaichiRuntimeError(
+                f"The solver type {solver_type} with {ordering} is not supported for now. Only {solver_type_list} with {solver_ordering} are supported."
+            )
+
+    @staticmethod
+    def _type_assert(sparse_matrix):
+        raise GsTaichiRuntimeError(
+            f"The parameter type: {type(sparse_matrix)} is not supported in linear solvers for now."
+        )
+
+    def compute(self, sparse_matrix):
+        """This method is equivalent to calling both `analyze_pattern` and then `factorize`.
+
+        Args:
+            sparse_matrix (SparseMatrix): The sparse matrix to be computed.
+        """
+        if isinstance(sparse_matrix, SparseMatrix):
+            self.matrix = sparse_matrix
+            gstaichi_arch = gstaichi.lang.impl.get_runtime().prog.config().arch
+            if gstaichi_arch == _ti_core.Arch.x64 or gstaichi_arch == _ti_core.Arch.arm64:
+                self.solver.compute(sparse_matrix.matrix)
+            elif gstaichi_arch == _ti_core.Arch.cuda:
+                self.analyze_pattern(self.matrix)
+                self.factorize(self.matrix)
+        else:
+            self._type_assert(sparse_matrix)
+
+    def analyze_pattern(self, sparse_matrix):
+        """Reorder the nonzero elements of the matrix, such that the factorization step creates less fill-in.
+
+        Args:
+            sparse_matrix (SparseMatrix): The sparse matrix to be analyzed.
+        """
+        if isinstance(sparse_matrix, SparseMatrix):
+            self.matrix = sparse_matrix
+            if self.matrix.dtype != self.dtype:
+                raise GsTaichiRuntimeError(
+                    f"The SparseSolver's dtype {self.dtype} is not consistent with the SparseMatrix's dtype {self.matrix.dtype}."
+                )
+            self.solver.analyze_pattern(sparse_matrix.matrix)
+        else:
+            self._type_assert(sparse_matrix)
+
+    def factorize(self, sparse_matrix):
+        """Do the factorization step
+
+        Args:
+            sparse_matrix (SparseMatrix): The sparse matrix to be factorized.
+        """
+        if isinstance(sparse_matrix, SparseMatrix):
+            self.matrix = sparse_matrix
+            self.solver.factorize(sparse_matrix.matrix)
+        else:
+            self._type_assert(sparse_matrix)
+
+    def solve(self, b):  # pylint: disable=R1710
+        """Computes the solution of the linear systems.
+        Args:
+            b (numpy.array or Field): The right-hand side of the linear systems.
+
+        Returns:
+            numpy.array: The solution of linear systems.
+        """
+        if self.matrix is None:
+            raise GsTaichiRuntimeError("Please call compute() before calling solve().")
+        if isinstance(b, Field):
+            return self.solver.solve(b.to_numpy())
+        if isinstance(b, np.ndarray):
+            return self.solver.solve(b)
+        if isinstance(b, Ndarray):
+            x = ScalarNdarray(b.dtype, [self.matrix.m])
+            self.solver.solve_rf(get_runtime().prog, self.matrix.matrix, b.arr, x.arr)
+            return x
+        raise GsTaichiRuntimeError(f"The parameter type: {type(b)} is not supported in linear solvers for now.")
+
+    def info(self):
+        """Check if the linear systems are solved successfully.
+
+        Returns:
+            bool: True if the solving process succeeded, False otherwise.
+        """
+        return self.solver.info()

gstaichi/math/__init__.py

@@ -0,0 +1,11 @@
+# type: ignore
+
+"""GsTaichi math module.
+
+The math module supports glsl-style vectors, matrices and functions.
+"""
+
+from ._complex import *
+from .mathimpl import *  # pylint: disable=W0622
+
+del mathimpl

gstaichi/math/_complex.py

@@ -0,0 +1,205 @@
+# type: ignore
+
+from gstaichi.lang import ops
+from gstaichi.lang.kernel_impl import func
+
+from .mathimpl import dot, vec2
+
+
+@func
+def cmul(z1, z2):
+    """Performs complex multiplication between two 2d vectors.
+
+    This is equivalent to the multiplication in the complex number field
+    when `z1` and `z2` are treated as complex numbers.
+
+    Args:
+        z1 (:class:`~gstaichi.math.vec2`): The first input.
+        z2 (:class:`~gstaichi.math.vec2`): The second input.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z1 = ti.math.vec2(1, 1)
+        >>>     z2 = ti.math.vec2(0, 1)
+        >>>     ti.math.cmul(z1, z2)  # [-1, 1]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: the complex multiplication `z1 * z2`.
+    """
+    x1, y1 = z1[0], z1[1]
+    x2, y2 = z2[0], z2[1]
+    return vec2(x1 * x2 - y1 * y2, x1 * y2 + x2 * y1)
+
+
+@func
+def cconj(z):
+    """Returns the complex conjugate of a 2d vector.
+
+    If `z=(x, y)` then the conjugate of `z` is `(x, -y)`.
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The input.
+
+    Returns:
+       :class:`~gstaichi.math.vec2`: The complex conjugate of `z`.
+    """
+    return vec2(z[0], -z[1])
+
+
+@func
+def cdiv(z1, z2):
+    """Performs complex division between two 2d vectors.
+
+    This is equivalent to the division in the complex number field
+    when `z1` and `z2` are treated as complex numbers.
+
+    Args:
+        z1 (:class:`~gstaichi.math.vec2`): The first input.
+        z2 (:class:`~gstaichi.math.vec2`): The second input.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z1 = ti.math.vec2(1, 1)
+        >>>     z2 = ti.math.vec2(0, 1)
+        >>>     ti.math.cdiv(z1, z2)  # [1, -1]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: the complex division of `z1 / z2`.
+    """
+    x1, y1 = z1[0], z1[1]
+    x2, y2 = z2[0], z2[1]
+    return vec2(x1 * x2 + y1 * y2, -x1 * y2 + x2 * y1) / dot(z2, z2)
+
+
+@func
+def csqrt(z):
+    """Returns the complex square root of a 2d vector `z`, so that
+    if `w^2=z`, then `w = csqrt(z)`.
+
+    Among the two square roots of `z`, if their real parts are non-zero,
+    the one with positive real part is returned. If both their real parts
+    are zero, the one with non-negative imaginary part is returned.
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The input.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z = ti.math.vec2(-1, 0)
+        >>>     w = ti.math.csqrt(z)  # [0, 1]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: The complex square root.
+    """
+    result = vec2(0.0)
+    if any(z):
+        r = ops.sqrt(z.norm())
+        a = ops.atan2(z[1], z[0])
+        result = r * vec2(ops.cos(a / 2.0), ops.sin(a / 2.0))
+
+    return result
+
+
+@func
+def cinv(z):
+    """Computes the reciprocal of a complex `z`.
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The input.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z = ti.math.vec2(1, 1)
+        >>>     w = ti.math.cinv(z)  # [0.5, -0.5]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: The reciprocal of `z`.
+    """
+    return cconj(z) / dot(z, z)
+
+
+@func
+def cpow(z, n):
+    """Computes the power of a complex `z`: :math:`z^a`.
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The base.
+        a (float): The exponent.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z = ti.math.vec2(1, 1)
+        >>>     w = ti.math.cpow(z)  # [-2, 2]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: The power :math:`z^a`.
+    """
+    result = vec2(0.0)
+    if any(z):
+        r2 = dot(z, z)
+        a = ops.atan2(z[1], z[0]) * n
+        result = ops.pow(r2, n / 2.0) * vec2(ops.cos(a), ops.sin(a))
+
+    return result
+
+
+@func
+def cexp(z):
+    """Returns the complex exponential :math:`e^z`.
+
+    `z` is a 2d vector treated as a complex number.
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The exponent.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z = ti.math.vec2(1, 1)
+        >>>     w = ti.math.cexp(z)  # [1.468694, 2.287355]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: The power :math:`exp(z)`
+    """
+    r = ops.exp(z[0])
+    return vec2(r * ops.cos(z[1]), r * ops.sin(z[1]))
+
+
+@func
+def clog(z):
+    """Returns the complex logarithm of `z`, so that if :math:`e^w = z`,
+    then :math:`log(z) = w`.
+
+    `z` is a 2d vector treated as a complex number. The argument of :math:`w`
+    lies in the range (-pi, pi].
+
+    Args:
+        z (:class:`~gstaichi.math.vec2`): The input.
+
+    Example::
+
+        >>> @ti.kernel
+        >>> def test():
+        >>>     z = ti.math.vec2(1, 1)
+        >>>     w = ti.math.clog(z)  # [0.346574, 0.785398]
+
+    Returns:
+        :class:`~gstaichi.math.vec2`: The logarithm of `z`.
+    """
+    ang = ops.atan2(z[1], z[0])
+    r2 = dot(z, z)
+    return vec2(ops.log(r2) / 2.0, ang)
+
+
+__all__ = ["cconj", "cdiv", "cexp", "cinv", "clog", "cmul", "cpow", "csqrt"]