def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): if static(as_vector): mat[I][p] = arr[I, p] else: mat[I][p, q] = arr[I, p, q]
def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): arr[I, p] = ndarray[I][p] else: arr[I, p, q] = ndarray[I][p, q]
def matrix_to_ext_arr(mat: template(), arr: ndarray_type.ndarray(), as_vector: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): if static(as_vector): arr[I, p] = mat[I][p] else: arr[I, p, q] = mat[I][p, q]
def outer_product(self, other): impl.static( impl.static_assert(self.m == 1, "lhs for outer_product is not a vector")) impl.static( impl.static_assert(other.m == 1, "rhs for outer_product is not a vector")) ret = Matrix([[self[i] * other[j] for j in range(other.n)] for i in range(self.n)]) return ret
def _gauss_elimination_3x3(Ab, dt): for i in static(range(3)): max_row = i max_v = ops.abs(Ab[i, i]) for j in static(range(i + 1, 3)): if ops.abs(Ab[j, i]) > max_v: max_row = j max_v = ops.abs(Ab[j, i]) assert max_v != 0.0, "Matrix is singular in linear solve." if i != max_row: if max_row == 1: for col in static(range(4)): Ab[i, col], Ab[1, col] = Ab[1, col], Ab[i, col] else: for col in static(range(4)): Ab[i, col], Ab[2, col] = Ab[2, col], Ab[i, col] assert Ab[i, i] != 0.0, "Matrix is singular in linear solve." for j in static(range(i + 1, 3)): scale = Ab[j, i] / Ab[i, i] Ab[j, i] = 0.0 for k in static(range(i + 1, 4)): Ab[j, k] -= Ab[i, k] * scale # Back substitution x = Vector.zero(dt, 3) for i in static(range(2, -1, -1)): x[i] = Ab[i, 3] for k in static(range(i + 1, 3)): x[i] -= Ab[i, k] * x[k] x[i] = x[i] / Ab[i, i] return x
def dot(self, other): """Perform the dot product with the input Vector (1-D Matrix). Args: other (:class:`~taichi.lang.matrix.Matrix`): The input Vector (1-D Matrix) to perform the dot product. Returns: DataType: The dot product result (scalar) of the two Vectors. """ impl.static( impl.static_assert(self.m == 1, "lhs for dot is not a vector")) impl.static( impl.static_assert(other.m == 1, "rhs for dot is not a vector")) return (self * other).sum()
def solve(A, b, dt=None): """Solve a matrix using Gauss elimination method. Args: A (ti.Matrix(n, n)): input nxn matrix `A`. b (ti.Vector(n, 1)): input nx1 vector `b`. dt (DataType): The datatype for the `A` and `b`. Returns: x (ti.Vector(n, 1)): the solution of Ax=b. """ assert A.n == A.m, "Only sqaure matrix is supported" assert A.n >= 2 and A.n <= 3, "Only 2D and 3D matrices are supported" assert A.m == b.n, "Matrix and Vector dimension dismatch" if dt is None: dt = impl.get_runtime().default_fp nrow, ncol = static(A.n, A.n + 1) Ab = expr_init(Matrix.zero(dt, nrow, ncol)) lhs = tuple([e.ptr for e in A.entries]) rhs = tuple([e.ptr for e in b.entries]) for i in range(nrow): for j in range(nrow): Ab(i, j)._assign(lhs[nrow * i + j]) for i in range(nrow): Ab(i, nrow)._assign(rhs[i]) if A.n == 2: return _gauss_elimination_2x2(Ab, dt) if A.n == 3: return _gauss_elimination_3x3(Ab, dt) raise Exception("Solver only supports 2D and 3D matrices.")
def ndarray_matrix_to_ext_arr(ndarray: ndarray_type.ndarray(), arr: ndarray_type.ndarray(), layout_is_aos: template(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): if static(layout_is_aos): arr[I, p] = ndarray[I][p] else: arr[p, I] = ndarray[I][p] else: if static(layout_is_aos): arr[I, p, q] = ndarray[I][p, q] else: arr[p, q, I] = ndarray[I][p, q]
def _matrix_outer_product(self, other): """Perform the outer product with the input Vector (1-D Matrix). Args: other (:class:`~taichi.lang.matrix.Matrix`): The input Vector (1-D Matrix) to perform the outer product. Returns: :class:`~taichi.lang.matrix.Matrix`: The outer product result (Matrix) of the two Vectors. """ impl.static( impl.static_assert(self.m == 1, "lhs for outer_product is not a vector")) impl.static( impl.static_assert(other.m == 1, "rhs for outer_product is not a vector")) return matrix.Matrix([[self[i] * other[j] for j in range(other.n)] for i in range(self.n)])
def normalized(self, eps=0): """Normalize a vector. Args: eps (Number): a safe-guard value for sqrt, usually 0. Examples:: a = ti.Vector([3, 4]) a.normalized() # [3 / 5, 4 / 5] # `a.normalized()` is equivalent to `a / a.norm()`. Note: Only vector normalization is supported. """ impl.static( impl.static_assert(self.m == 1, "normalized() only works on vector")) invlen = 1 / (self.norm() + eps) return invlen * self
def _svd(A, dt): """Perform singular value decomposition (A=USV^T) for arbitrary size matrix. Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition. 2D implementation refers to :func:`taichi.svd2d`. 3D implementation refers to :func:`taichi.svd3d`. Args: A (ti.Matrix(n, n)): input nxn matrix `A`. dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64. Returns: Decomposed nxn matrices `U`, 'S' and `V`. """ if static(A.n == 2): # pylint: disable=R1705 ret = svd2d(A, dt) return ret elif static(A.n == 3): return svd3d(A, dt) else: raise Exception("SVD only supports 2D and 3D matrices.")
def _polar_decompose(A, dt): """Perform polar decomposition (A=UP) for arbitrary size matrix. Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition. 2D implementation refers to :func:`taichi.polar_decompose2d`. 3D implementation refers to :func:`taichi.polar_decompose3d`. Args: A (ti.Matrix(n, n)): input nxn matrix `A`. dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64. Returns: Decomposed nxn matrices `U` and `P`. """ if static(A.n == 2): # pylint: disable=R1705 ret = polar_decompose2d(A, dt) return ret elif static(A.n == 3): return polar_decompose3d(A, dt) else: raise Exception( "Polar decomposition only supports 2D and 3D matrices.")
def vector_to_fast_image(img: template(), out: ndarray_type.ndarray()): # FIXME: Why is ``for i, j in img:`` slower than: for i, j in ndrange(*img.shape): r, g, b = 0, 0, 0 color = img[i, img.shape[1] - 1 - j] if static(img.dtype in [f16, f32, f64]): r, g, b = min(255, max(0, int(color * 255))) else: static_assert(img.dtype == u8) r, g, b = color idx = j * img.shape[0] + i # We use i32 for |out| since OpenGL and Metal doesn't support u8 types if static(get_os_name() != 'osx'): out[idx] = (r << 16) + (g << 8) + b else: # What's -16777216? # # On Mac, we need to set the alpha channel to 0xff. Since Mac's GUI # is big-endian, the color is stored in ABGR order, and we need to # add 0xff000000, which is -16777216 in I32's legit range. (Albeit # the clarity, adding 0xff000000 doesn't work.) alpha = -16777216 out[idx] = (b << 16) + (g << 8) + r + alpha
def _gauss_elimination_2x2(Ab, dt): if ops.abs(Ab[0, 0]) < ops.abs(Ab[1, 0]): Ab[0, 0], Ab[1, 0] = Ab[1, 0], Ab[0, 0] Ab[0, 1], Ab[1, 1] = Ab[1, 1], Ab[0, 1] Ab[0, 2], Ab[1, 2] = Ab[1, 2], Ab[0, 2] assert Ab[0, 0] != 0.0, "Matrix is singular in linear solve." scale = Ab[1, 0] / Ab[0, 0] Ab[1, 0] = 0.0 for k in static(range(1, 3)): Ab[1, k] -= Ab[0, k] * scale x = Vector.zero(dt, 2) # Back substitution x[1] = Ab[1, 2] / Ab[1, 1] x[0] = (Ab[0, 2] - Ab[0, 1] * x[1]) / Ab[0, 0] return x
def vector_to_image(mat: template(), arr: ndarray_type.ndarray()): for I in grouped(mat): for p in static(range(mat.n)): arr[I, p] = ops.cast(mat[I][p], f32) if static(mat.n <= 2): arr[I, 2] = 0
def normalized(self, eps=0): impl.static( impl.static_assert(self.m == 1, "normalized() only works on vector")) invlen = 1 / (self.norm() + eps) return invlen * self
def dot(self, other): impl.static( impl.static_assert(self.m == 1, "lhs for dot is not a vector")) impl.static( impl.static_assert(other.m == 1, "rhs for dot is not a vector")) return (self * other).sum()
def fill_matrix(mat: template(), vals: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): mat[I][p, q] = vals[p][q]
def clear_gradients(_vars: template()): for I in grouped(ScalarField(Expr(_vars[0]))): for s in static(_vars): ScalarField(Expr(s))[I] = 0