Math Module
Taichi provides a built-in math
module that supports frequently used mathematical functions and utility functions, including:
- Commonly-used mathematical functions that are analogous to those in Python's built-in
math
module. - Small vector and matrix types that are analogous to those in the OpenGL shading language (GLSL).
- Some GLSL-standard functions.
- Complex number operations in the form of 2D vectors.
Mathematical functions
You must call the mathematical functions provided by Taichi's math
module from within the Taichi scope. For example:
import taichi as ti
import taichi.math as tm
ti.init()
@ti.kernel
def test():
a = 1.0
x = tm.sin(a)
y = tm.floor(a)
z = tm.degrees(a)
w = tm.log2(a)
...
These functions also take vectors and matrices as arguments and operate on them element-wise:
@ti.kernel
def test():
a = ti.Vector([1.0, 2.0, 3.0])
x = tm.sin(a) # [0.841471, 0.909297, 0.141120]
y = tm.floor(a) # [1.000000, 2.000000, 3.000000]
z = tm.degrees(a) # [57.295780, 114.591560, 171.887344]
b = ti.Vector([2.0, 3.0, 4.0])
w = tm.atan2(b, a) # [1.107149, 0.982794, 0.927295]
...
note
Taichi's math module overlaps to a large extent with Python's built-in math module. Ensure that you follow a few extra rules when using Taichi's math module:
- You must call the functions provided by Taichi's math module from within the Taichi scope.
- Functions in Taichi's math module also take vectors or matrices as arguments.
- The precision of a function in Taichi's math module depends on the settings of
default_fp
andarch
(backend) inti.init()
.
Small vector and matrix types
Taichi's math module provides a few small vector and matrix types:
vec2/vec3/vec4
: 2D/3D/4D floating-point vector types.ivec2/ivec3/ivec4
: 2D/3D/4D integer vector types.uvec2/uvec3/uvec4
: 2D/3D/4D unsigned integer vector types.mat2/mat3/mat4
: 2D/3D/4D floating-point square matrix types.
To create one of the vector/matrix types above, use template function ti.types.vector()
or ti.types.matrix()
. For example, vec2
is defined as follows:
vec2 = ti.types.vector(2, float)
The number of precision bits of such a type is determined by default_fp
or default_ip
in the ti.init()
method call. For example, if ti.init(default_fp=ti.f64)
is called, then vec2/vec3/vec4
and mat2/mat3/mat4
defined in the Taichi scope all have a 64-bit floating-point precision.
You can use these types to instantiate vectors/matrices or annotate data types for function arguments and struct members. See the Type System for more information. Here we emphasize that they have very flexible initialization routines:
mat2 = ti.math.mat2
vec3 = ti.math.mat3
vec4 = ti.math.vec4
m = mat2(1) # [[1., 1.], [1., 1.]]
m = mat2(1, 2, 3, 4) # [[1., 2.], [3, 4.]]
m = mat2([1, 2], [3, 4]) # [[1., 2.], [3, 4.]]
m = mat2([1, 2, 3, 4]) # [[1., 2.], [3, 4.]]
v = vec3(1, 2, 3)
m = mat2(v, 4) # [[1., 2.], [3, 4.]]
u = vec4([1, 2], [3, 4])
u = vec4(v, 4.0)
Another important feature of vector types created by ti.types.vector()
is that they support vector swizzling just as GLSL vectors do. This means you can use xyzw
, rgba
, stpq
to access their elements with indices ≤ four:
v = ti.math.vec4(1, 2, 3, 4)
u = v.xyz # vec3(1, 2, 3)
u = v.xxx # vec3(1, 1, 1)
u = v.wzyx # vec4(4, 3, 2, 1)
u = v.rraa # vec4(1, 1, 2, 2)
Relations between ti.Vector
, ti.types.vector
and ti.math.vec3
ti.Vector
is a function that accepts a 1D array and returns a matrix instance that has only one column. For example,ti.Vector([1, 2, 3, 4, 5])
.ti.types.vector
is a function that accepts an integer and a primitive type and returns a vector type. For example:vec5f = ti.types.vector(5, float)
.vec5f
can then be used to instantiate 5D vectors or annotate data types of function arguments and struct members:@ti.kernel
def test(v: vec5f):
print(v.xyz)Unlike
ti.Vector
, whose input data must be a 1D array, vector types created byti.types.vector()
have more flexible ways to initialize, as explained above.ti.math.vec3
is created byvec3 = ti.types.vector(3, float)
.
GLSL-standard functions
Taichi's math module also supports a few GLSL standard functions. These functions follow the GLSL standard, except that they accept arbitrary vectors and matrices as arguments and operate on them element-wise. For example:
import taichi as ti
import taichi.math as tm
@ti.kernel
def example():
v = tm.vec3(0., 1., 2.)
w = tm.smoothstep(0.0, 1.0, v)
w = tm.clamp(w, 0.2, 0.8)
w = tm.reflect(v, tm.normalize(tm.vec3(1)))
note
Texture support in Taichi is implemented in the ti.types.texture_types
module.
Complex number operations
Taichi's math module also supports basic complex arithmetic operations on 2D vectors.
You can use a 2D vector of type ti.math.vec2
to represent a complex number. In this way, additions and subtractions of complex numbers come in the form of 2D vector additions and subtractions. You can call ti.math.cmul()
and ti.math.cdiv()
to conduct multiplication and division of complex numbers:
import taichi as ti
import taichi.math as tm
ti.init()
@ti.kernel
def test():
x = tm.vec2(1, 1) # complex number 1+1j
y = tm.vec2(0, 1) # complex number 1j
z = tm.cmul(x, y) # vec2(-1, 1) = -1+1j
w = tm.cdiv(x, y) # vec2(2, 0) = 2+0j
You can also compute the power, logarithm, and exponential of a complex number:
@ti.kernel
def test():
x = tm.vec2(1, 1) # complex number 1 + 1j
y = tm.cpow(x, 2) # complex number (1 + 1j)**2 = 2j
z = tm.clog(x) # complex number (0.346574 + 0.785398j)
w = tm.cexp(x) # complex number (1.468694 + 2.287355j)