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