def polar_decompose2d(A, dt):
"""Perform polar decomposition (A=UP) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U` and `P`. `U` is a 2x2 orthogonal matrix
and `P` is a 2x2 positive or semi-positive definite matrix.
"""
U = Matrix.identity(dt, 2)
P = ops.cast(A, dt)
zero = ops.cast(0.0, dt)
# if A is a zero matrix we simply return the pair (I, A)
if (A[0, 0] == zero and A[0, 1] == zero and A[1, 0] == zero
and A[1, 1] == zero):
pass
else:
detA = A[0, 0] * A[1, 1] - A[1, 0] * A[0, 1]
adetA = abs(detA)
B = Matrix([[A[0, 0] + A[1, 1], A[0, 1] - A[1, 0]],
[A[1, 0] - A[0, 1], A[1, 1] + A[0, 0]]], dt)
if detA < zero:
B = Matrix([[A[0, 0] - A[1, 1], A[0, 1] + A[1, 0]],
[A[1, 0] + A[0, 1], A[1, 1] - A[0, 0]]], dt)
# here det(B) != 0 if A is not the zero matrix
adetB = abs(B[0, 0] * B[1, 1] - B[1, 0] * B[0, 1])
k = ops.cast(1.0, dt) / ops.sqrt(adetB)
U = B * k
P = (A.transpose() @ A + adetA * Matrix.identity(dt, 2)) * k
return U, P
def build_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = impl.expr_init(build_stmt(ctx, node.iter))
ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
ndrange_end = ti_ops.cast(
expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
primitive_types.i32)
ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
_ti_core.begin_frontend_range_for(ndrange_loop_var.ptr,
ndrange_begin.ptr,
ndrange_end.ptr)
I = impl.expr_init(ndrange_loop_var)
targets = ASTTransformer.get_for_loop_targets(node)
for i, target in enumerate(targets):
if i + 1 < len(targets):
target_tmp = impl.expr_init(
I // ndrange_var.acc_dimensions[i + 1])
else:
target_tmp = impl.expr_init(I)
ctx.create_variable(
target,
impl.expr_init(target_tmp + impl.subscript(
impl.subscript(ndrange_var.bounds, i), 0)))
if i + 1 < len(targets):
I.assign(I -
target_tmp * ndrange_var.acc_dimensions[i + 1])
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def build_range_for(ctx, node):
with ctx.variable_scope_guard():
loop_name = node.target.id
ctx.check_loop_var(loop_name)
loop_var = expr.Expr(_ti_core.make_id_expr(''))
ctx.create_variable(loop_name, loop_var)
if len(node.iter.args) not in [1, 2]:
raise TaichiSyntaxError(
f"Range should have 1 or 2 arguments, found {len(node.iter.args)}"
)
if len(node.iter.args) == 2:
begin = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[0])),
primitive_types.i32)
end = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[1])),
primitive_types.i32)
else:
begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
end = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[0])),
primitive_types.i32)
_ti_core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr)
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def _augassign(self, value, op):
call_func = f"insert_triplet_{self.dtype}"
if op == 'Add':
taichi.lang.impl.call_internal(call_func, self.ptr, self.i, self.j,
ops.cast(value, self.dtype))
elif op == 'Sub':
taichi.lang.impl.call_internal(call_func, self.ptr, self.i, self.j,
-ops.cast(value, self.dtype))
else:
assert False, "Only operations '+=' and '-=' are supported on sparse matrices."
def build_assign_annotated(ctx, target, value, is_static_assign,
annotation):
"""Build an annotated assginment like this: target: annotation = value.
Args:
ctx (ast_builder_utils.BuilderContext): The builder context.
target (ast.Name): A variable name. `target.id` holds the name as
a string.
annotation: A type we hope to assign to the target
value: A node representing the value.
is_static_assign: A boolean value indicating whether this is a static assignment
"""
is_local = isinstance(target, ast.Name)
anno = impl.expr_init(annotation)
if is_static_assign:
raise TaichiSyntaxError(
"Static assign cannot be used on annotated assignment")
if is_local and not ctx.is_var_declared(target.id):
var = ti_ops.cast(value, anno)
var = impl.expr_init(var)
ctx.create_variable(target.id, var)
else:
var = build_stmt(ctx, target)
if var.ptr.get_ret_type() != anno:
raise TaichiSyntaxError(
"Static assign cannot have type overloading")
var.assign(value)
return var
def build_Return(ctx, node):
if not impl.get_runtime().experimental_real_function:
if ctx.is_in_non_static():
raise TaichiSyntaxError(
"Return inside non-static if/for is not supported")
build_stmt(ctx, node.value)
if ctx.is_kernel or impl.get_runtime().experimental_real_function:
# TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
if node.value is not None:
if ctx.func.return_type is None:
raise TaichiSyntaxError(
f'A {"kernel" if ctx.is_kernel else "function"} '
'with a return value must be annotated '
'with a return type, e.g. def func() -> ti.f32')
_ti_core.create_kernel_return(
ti_ops.cast(expr.Expr(node.value.ptr),
ctx.func.return_type).ptr)
# For args[0], it is an ast.Attribute, because it loads the
# attribute, |ptr|, of the expression |ret_expr|. Therefore we
# only need to replace the object part, i.e. args[0].value
else:
ctx.return_data = node.value.ptr
if not impl.get_runtime().experimental_real_function:
ctx.returned = True
return None
def _randn(dt):
"""
Generate a random float sampled from univariate standard normal
(Gaussian) distribution of mean 0 and variance 1, using the
Box-Muller transformation.
"""
assert dt == f32 or dt == f64
u1 = ops.cast(1.0, dt) - ops.random(dt)
u2 = ops.random(dt)
r = ops.sqrt(-2 * ops.log(u1))
c = ops.cos(math.tau * u2)
return r * c
def build_Return(ctx, node):
if not ctx.is_real_function:
if ctx.is_in_non_static_control_flow():
raise TaichiSyntaxError(
"Return inside non-static if/for is not supported")
if node.value is not None:
build_stmt(ctx, node.value)
if node.value is None or node.value.ptr is None:
if not ctx.is_real_function:
ctx.returned = ReturnStatus.ReturnedVoid
return None
if ctx.is_kernel or ctx.is_real_function:
# TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
if ctx.func.return_type is None:
raise TaichiSyntaxError(
f'A {"kernel" if ctx.is_kernel else "function"} '
'with a return value must be annotated '
'with a return type, e.g. def func() -> ti.f32')
if id(ctx.func.return_type) in primitive_types.type_ids:
ctx.ast_builder.create_kernel_exprgroup_return(
expr.make_expr_group(
ti_ops.cast(expr.Expr(node.value.ptr),
ctx.func.return_type).ptr))
elif isinstance(ctx.func.return_type, MatrixType):
ctx.ast_builder.create_kernel_exprgroup_return(
expr.make_expr_group([
ti_ops.cast(exp, ctx.func.return_type.dtype) for exp in
itertools.chain.from_iterable(node.value.ptr.to_list())
]))
else:
raise TaichiSyntaxError(
"The return type is not supported now!")
# For args[0], it is an ast.Attribute, because it loads the
# attribute, |ptr|, of the expression |ret_expr|. Therefore we
# only need to replace the object part, i.e. args[0].value
else:
ctx.return_data = node.value.ptr
if not ctx.is_real_function:
ctx.returned = ReturnStatus.ReturnedValue
return None
def func_call_rvalue(self, key, args):
# Skip the template args, e.g., |self|
assert self.is_real_function
non_template_args = []
for i, anno in enumerate(self.argument_annotations):
if not isinstance(anno, template):
if id(anno) in primitive_types.type_ids:
non_template_args.append(ops.cast(args[i], anno))
else:
non_template_args.append(args[i])
non_template_args = impl.make_expr_group(non_template_args)
return Expr(
_ti_core.make_func_call_expr(
self.taichi_functions[key.instance_id], non_template_args))
def cast(self, dtype):
"""Cast the matrix element data type.
Args:
dtype (DataType): the data type of the casted matrix element.
Returns:
A new matrix with each element's type is dtype.
"""
_taichi_skip_traceback = 1
ret = self.copy()
for i in range(len(self.entries)):
ret.entries[i] = ops_mod.cast(ret.entries[i], dtype)
return ret
def build_grouped_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = impl.expr_init(build_stmt(ctx, node.iter.args[0]))
ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
ndrange_end = ti_ops.cast(
expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
primitive_types.i32)
ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
_ti_core.begin_frontend_range_for(ndrange_loop_var.ptr,
ndrange_begin.ptr,
ndrange_end.ptr)
targets = ASTTransformer.get_for_loop_targets(node)
if len(targets) != 1:
raise TaichiSyntaxError(
f"Group for should have 1 loop target, found {len(targets)}"
)
target = targets[0]
target_var = impl.expr_init(
matrix.Vector([0] * len(ndrange_var.dimensions),
dt=primitive_types.i32))
ctx.create_variable(target, target_var)
I = impl.expr_init(ndrange_loop_var)
for i in range(len(ndrange_var.dimensions)):
if i + 1 < len(ndrange_var.dimensions):
target_tmp = I // ndrange_var.acc_dimensions[i + 1]
else:
target_tmp = I
impl.subscript(target_var,
i).assign(target_tmp + ndrange_var.bounds[i][0])
if i + 1 < len(ndrange_var.dimensions):
I.assign(I -
target_tmp * ndrange_var.acc_dimensions[i + 1])
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def cast(self, struct):
# sanity check members
if self.members.keys() != struct.entries.keys():
raise TaichiSyntaxError(
"Incompatible arguments for custom struct members!")
entries = {}
for k, dtype in self.members.items():
if isinstance(dtype, CompoundType):
entries[k] = dtype.cast(struct.entries[k])
else:
if in_python_scope():
v = struct.entries[k]
entries[k] = int(v) if dtype in integer_types else float(v)
else:
entries[k] = ops.cast(struct.entries[k], dtype)
return Struct(entries)
def cast(self, struct, in_place=False):
if not in_place:
struct = struct.copy()
# sanity check members
if self.members.keys() != struct.entries.keys():
raise TaichiSyntaxError(
"Incompatible arguments for custom struct members!")
for k, dtype in self.members.items():
if isinstance(dtype, CompoundType):
struct.entries[k] = dtype.cast(struct.entries[k])
else:
if in_python_scope():
v = struct.entries[k]
struct.entries[k] = int(
v) if dtype in ti.integer_types else float(v)
else:
struct.entries[k] = cast(struct.entries[k], dtype)
return struct
def svd2d(A, dt):
"""Perform singular value decomposition (A=USV^T) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U`, 'S' and `V`.
"""
R, S = polar_decompose2d(A, dt)
c, s = ops.cast(0.0, dt), ops.cast(0.0, dt)
s1, s2 = ops.cast(0.0, dt), ops.cast(0.0, dt)
if abs(S[0, 1]) < 1e-5:
c, s = 1, 0
s1, s2 = S[0, 0], S[1, 1]
else:
tao = ops.cast(0.5, dt) * (S[0, 0] - S[1, 1])
w = ops.sqrt(tao**2 + S[0, 1]**2)
t = ops.cast(0.0, dt)
if tao > 0:
t = S[0, 1] / (tao + w)
else:
t = S[0, 1] / (tao - w)
c = 1 / ops.sqrt(t**2 + 1)
s = -t * c
s1 = c**2 * S[0, 0] - 2 * c * s * S[0, 1] + s**2 * S[1, 1]
s2 = s**2 * S[0, 0] + 2 * c * s * S[0, 1] + c**2 * S[1, 1]
V = Matrix.zero(dt, 2, 2)
if s1 < s2:
tmp = s1
s1 = s2
s2 = tmp
V = Matrix([[-s, c], [-c, -s]], dt=dt)
else:
V = Matrix([[c, s], [-s, c]], dt=dt)
U = R @ V
return U, Matrix([[s1, ops.cast(0, dt)], [ops.cast(0, dt), s2]], dt=dt), V
def tensor_to_image(tensor: template(), arr: ext_arr()):
for I in grouped(tensor):
t = ops.cast(tensor[I], f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
def __ti_int__(self):
return ops.cast(self, int)
def __ti_float__(self):
return ops.cast(self, float)
def tensor_to_image(tensor: template(), arr: ndarray_type.ndarray()):
for I in grouped(tensor):
t = ops.cast(tensor[I], f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
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 cast(self, dtype):
_taichi_skip_traceback = 1
ret = self.copy()
for i in range(len(self.entries)):
ret.entries[i] = ops_mod.cast(ret.entries[i], dtype)
return ret
