[Lang] Add more functions to math module (#4939)

* add more functions to math module * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add more functions to math module * add more functions to math module * add more functions to math module * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add more functions to math module * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add more functions to math module * add more functions to math module * Update _funcs.py * Update python/taichi/_funcs.py Co-authored-by: pengyu <6712304+FantasyVR@users.noreply.github.com> * Update python/taichi/math/mathimpl.py Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: pengyu <6712304+FantasyVR@users.noreply.github.com>
taichi-dev · May 11, 2022 · c3631c0 · c3631c0
1 parent 6d538e1
commit c3631c0
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 21 deletions.
diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py
@@ -100,22 +100,39 @@ def _matrix_outer_product(self, other):
 @func
 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`.
+        Decomposed 2x2 matrices `U` and `P`. `U` is a 2x2 orthogonal matrix
+        and `P` is a 2x2 positive or semi-positive definite matrix.
     """
-    x, y = A(0, 0) + A(1, 1), A(1, 0) - A(0, 1)
-    scale = (1.0 / ops.sqrt(x * x + y * y))
-    c = x * scale
-    s = y * scale
-    r = Matrix([[c, -s], [s, c]], dt=dt)
-    return r, r.transpose() @ A
+    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
 
 
 @func

diff --git a/python/taichi/math/__init__.py b/python/taichi/math/__init__.py
@@ -3,6 +3,6 @@
 The math module supports glsl-style vectors, matrices and functions.
 """
 from ._complex import *
-from .mathimpl import *
+from .mathimpl import *  # pylint: disable=W0622
 
 del mathimpl
diff --git a/python/taichi/math/mathimpl.py b/python/taichi/math/mathimpl.py
@@ -1,9 +1,12 @@
+# pylint: disable=W0622
 """
 Math functions for glsl-like functions and other stuff.
 """
 from math import e, pi
 
 from taichi.lang import impl
+from taichi.lang.ops import (acos, asin, atan2, ceil, cos, exp, floor, log,
+                             max, min, pow, round, sin, sqrt, tan, tanh)
 
 import taichi as ti
 
@@ -576,7 +579,7 @@ def rot3(axis, ang):
         >>> from taichi.math import *
         >>> @ti.kernel
         >>> def test():
-        >>>     M = rot3(vec3(1, 1, 1), radians(30))
+        >>>     M = rot3(normalize(vec3(1, 1, 1)), radians(30))
         [[0.732051, -0.366025, 0.633975],
          [0.633975, 0.732051, -0.366025],
          [-0.366025, 0.633975, 0.732051]]
@@ -588,10 +591,32 @@ def rot3(axis, ang):
     return I + sa * K + (1.0 - ca) * K @ K
 
 
+@ti.func
+def length(x):
+    """Calculate the length of a vector.
+
+    This function is equivalent to the `length` function in GLSL.
+    Args:
+        x (:class:`~taichi.Matrix`): The vector of which to calculate the length.
+
+    Returns:
+        The Euclidean norm of the vector.
+
+    Example::
+
+        >>> x = ti.Vector([1, 1, 1])
+        >>> length(x)
+        1.732051
+    """
+    return x.norm()
+
+
 __all__ = [
-    "clamp", "cross", "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", "uvec2", "uvec3",
-    "uvec4", "vec2", "vec3", "vec4"
+    "acos", "asin", "atan2", "ceil", "clamp", "cos", "cross", "degrees",
+    "distance", "dot", "e", "exp", "eye", "floor", "fract", "ivec2", "ivec3",
+    "ivec4", "length", "log", "log2", "mat2", "mat3", "mat4", "max", "min",
+    "mix", "mod", "normalize", "pi", "pow", "radians", "reflect", "refract",
+    "rot2", "rot3", "rotate2d", "rotate3d", "round", "sign", "sin",
+    "smoothstep", "sqrt", "step", "tan", "tanh", "uvec2", "uvec3", "uvec4",
+    "vec2", "vec3", "vec4"
 ]
diff --git a/tests/python/test_api.py b/tests/python/test_api.py
@@ -97,12 +97,14 @@ def _get_expected_matrix_apis():
     'dynamic', 'finalize', 'lazy_grad', 'place', 'pointer'
 ]
 user_api[ti.math] = [
-    'cconj', 'cdiv', 'cexp', 'cinv', 'clamp', 'clog', 'cmul', 'cpow', 'cross',
-    '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', 'uvec2', 'uvec3',
-    'uvec4', 'vec2', 'vec3', 'vec4'
+    'acos', 'asin', 'atan2', 'cconj', 'cdiv', 'ceil', 'cexp', 'cinv', 'clamp',
+    'clog', 'cmul', 'cos', 'cpow', 'cross', 'csqrt', 'degrees', 'distance',
+    'dot', 'e', 'exp', 'eye', 'floor', 'fract', 'ivec2', 'ivec3', 'ivec4',
+    'length', 'log', 'log2', 'mat2', 'mat3', 'mat4', 'max', 'min', 'mix',
+    'mod', 'normalize', 'pi', 'pow', 'radians', 'reflect', 'refract', 'rot2',
+    'rot3', 'rotate2d', 'rotate3d', 'round', 'sign', 'sin', 'smoothstep',
+    'sqrt', 'step', 'tan', 'tanh', 'uvec2', 'uvec3', 'uvec4', 'vec2', 'vec3',
+    'vec4'
 ]
 user_api[ti.Matrix] = _get_expected_matrix_apis()
 user_api[ti.MatrixField] = [