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
