[Lang] Add 2d and 3d rotation functions to math module

python/taichi/math/mathimpl.py

@@ -5,6 +5,8 @@
 
 import taichi as ti
 
+from .vectypes import vec2
+
 mat2 = ti.types.matrix(2, 2, float)  # pylint: disable=E1101
 """2x2 float matrix type
 """
@@ -425,8 +427,125 @@ def mod(x, y):
     return x - y * ti.floor(x / y)
 
 
+@ti.func
+def rotate2d(p, ang):
+    """Rotates a 2d vector by a given angle in counter-clockwise.
+
+    Args:
+        p (:class:`~taichi.math.vec2`): The 2d vector to rotate.
+        ang (float): Angle of rotation, in radians.
+
+    Returns:
+        :class:`~taichi.math.vec2`: The vector after rotation.
+
+    Example::
+
+        >>> from taichi.math import *
+        >>> @ti.kernel
+        >>> def test():
+        >>>     v = vec2(1, 0)
+        >>>     print(rotate2d(v, radians(30)))
+        [0.866025, 0.500000]
+    """
+    ca, sa = ti.cos(ang), ti.sin(ang)
+    x, y = p
+    return vec2(x * ca - p.y * sa, x * sa + y * ca)
+
+
+@ti.func
+def rotate3d(p, axis, ang):
+    """Rotates a vector in 3d space, given an axis and angle of rotation.
+
+    The vector `axis` should be a unit vector.
+
+    See "https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula"
+
+    Args:
+        p (:class:`~taichi.math.vec3`): The 3d vector to rotate.
+        axis (:class:`~taichi.math.vec3`): Axis of rotation.
+        ang (float): Angle of rotation, in radians.
+
+    Example::
+
+        >>> from taichi.math import *
+        >>> @ti.kernel
+        >>> def test():
+        >>>     v = vec3(1, 0, 0)
+        >>>     axis = normalize(vec3(1, 1, 1))
+        >>>     print(rotate3d(v, axis, radians(30)))
+        [0.910684, 0.333333, -0.244017]
+
+    Returns:
+        :class:`~taichi.math.vec3`: The vector after rotation.
+    """
+    ca, sa = ti.cos(ang), ti.sin(ang)
+    return mix(dot(p, axis) * axis, p, ca) + cross(axis, p) * sa
+
+
+@ti.func
+def eye(n: ti.template()):
+    """Returns the nxn identiy matrix.
+
+    Alias for :func:`~taichi.Matrix.identity`.
+    """
+    return ti.Matrix.identity(float, n)
+
+
+@ti.func
+def rot2(ang):
+    """Returns the matrix representation of a 2d counter-clockwise rotation,
+    given the angle of rotation.
+
+    Args:
+        ang (float): Angle of rotation in radians.
+
+    Returns:
+        :class:`~taichi.math.mat2`: 2x2 rotation matrix.
+
+    Example::
+
+        >>> from taichi.math import *
+        >>> @ti.kernel
+        >>> def test():
+        >>>     M = rot2(radians(30))
+        [[0.866025, -0.500000], [0.500000, 0.866025]]
+    """
+    ca, sa = ti.cos(ang), ti.sin(ang)
+    return mat2([[ca, -sa], [sa, ca]])
+
+
+@ti.func
+def rot3(axis, ang):
+    """Returns the matrix representation of a 3d rotation,
+    given the axis and angle of rotation.
+
+    Args:
+        axis (:class:`~taichi.math.vec3`): Axis of rotation.
+        ang (float): Angle of rotation in radians.
+
+    Returns:
+        :class:`~taichi.math.mat3`: 3x3 rotation matrix.
+
+    Example::
+
+        >>> from taichi.math import *
+        >>> @ti.kernel
+        >>> def test():
+        >>>     M = rot3(vec3(1, 1, 1), radians(30))
+        [[0.732051, -0.366025, 0.633975],
+         [0.633975, 0.732051, -0.366025],
+         [-0.366025, 0.633975, 0.732051]]
+    """
+    ca, sa = ti.cos(ang), ti.sin(ang)
+    x, y, z = axis
+    I = eye(3)
+    K = mat3([[0, -z, y], [z, 0, -x], [-y, x, 0]])
+    return I + sa * K + (1.0 - ca) * K @ K
+
+
 __all__ = [
-    "clamp", "cross", "degrees", "distance", "dot", "e", "fract", "log2",
-    "mat2", "mat3", "mat4", "mix", "mod", "normalize", "pi", "radians",
-    "reflect", "refract", "sign", "smoothstep", "step"
+    "clamp", "cross", "degrees", "distance", "dot", "e", "eye", "fract",
+    "log2", "mat2", "mat3", "mat4", "mix", "mod", "normalize", "pi", "radians",
+    "reflect", "refract", "rot2", "rot3", "rotate2d", "rotate3d", "sign",
+    "smoothstep", "step"
 ]

tests/python/test_api.py

@@ -45,10 +45,11 @@
 ]
 user_api[ti.math] = [
     'cconj', 'cdiv', 'cexp', 'cinv', 'clamp', 'clog', 'cmul', 'cpow', 'cross',
-    'csqrt', 'degrees', 'distance', 'dot', 'e', 'fract', 'ivec2', 'ivec3',
-    'ivec4', 'log2', 'mat2', 'mat3', 'mat4', 'mix', 'mod', 'normalize', 'pi',
-    'radians', 'reflect', 'refract', 'sign', 'smoothstep', 'step', 'vec2',
-    'vec3', 'vec4'
+    'csqrt', 'degrees', 'distance', 'dot', 'e', 'eye', 'fract', 'ivec2',
+    'ivec3', 'ivec4', 'log2', 'mat2', 'mat3', 'mat4', 'mix', 'mod',
+    'normalize', 'pi', 'radians', 'reflect', 'refract', 'rot2', 'rot3',
+    'rotate2d', 'rotate3d', 'sign', 'smoothstep', 'step', 'vec2', 'vec3',
+    'vec4'
 ]
 user_api[ti.Matrix] = [
     'all', 'any', 'cast', 'cols', 'cross', 'determinant', 'diag', 'dot',