From 901ac7de05d91517dcca55d3502d8faf8f863b4f Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Fri, 18 Mar 2022 16:26:48 +0800 Subject: [PATCH 1/8] compute eig3x3 eigen values --- misc/test_eig3x3.py | 10 +++++ python/taichi/_funcs.py | 29 ++++++++++++ taichi/math/eig.h | 83 +++++++++++++++++++++++++++++++++++ taichi/python/export_lang.cpp | 3 ++ 4 files changed, 125 insertions(+) create mode 100644 misc/test_eig3x3.py create mode 100644 taichi/math/eig.h diff --git a/misc/test_eig3x3.py b/misc/test_eig3x3.py new file mode 100644 index 0000000000000..12a64233c5a29 --- /dev/null +++ b/misc/test_eig3x3.py @@ -0,0 +1,10 @@ +import taichi as ti +ti.init() + +@ti.kernel +def test(): + A = ti.Matrix([[3, 1, 1],[1, 2, 2], [1, 2, 2]]) + diagnal = ti.sym_eig(A) + print(diagnal) + +test() diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py index 0d9e0392e5223..ed145607c711d 100644 --- a/python/taichi/_funcs.py +++ b/python/taichi/_funcs.py @@ -306,6 +306,33 @@ def sym_eig2x2(A, dt): return eigenvalues, eigenvectors +def sym_eig3x3(A, dt): + """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 3x3 real symmetric matrix. + + Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix. + + Args: + A (ti.Matrix(3, 3)): input 3x3 symmetric matrix `A`. + dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64. + + Returns: + eigenvalues (ti.Vector(3)): The eigenvalues. Each entry store one eigen value. + eigenvectors (ti.Matrix(3, 3)): The eigenvectors. Each column stores one eigenvector. + """ + assert A.n == 3 and A.m == 3 + inputs = tuple([e.ptr for e in A.entries]) + assert dt in [f32, f64] + if dt == f32: + rets = get_runtime().prog.current_ast_builder().eig_3x3_f32(*inputs) + else: + rets = get_runtime().prog.current_ast_builder().eig_3x3_f64(*inputs) + assert len(rets) == 3 + D = expr_init(Matrix.zero(dt, 3, 1)) + for i in range(3): + D(i)._assign(rets[i]) + return D + + @func def _svd(A, dt): """Perform singular value decomposition (A=USV^T) for arbitrary size matrix. @@ -431,6 +458,8 @@ def sym_eig(A, dt=None): dt = impl.get_runtime().default_fp if A.n == 2: return sym_eig2x2(A, dt) + elif A.n == 3: + return sym_eig3x3(A, dt) raise Exception("Symmetric eigen solver only supports 2D matrices.") diff --git a/taichi/math/eig.h b/taichi/math/eig.h new file mode 100644 index 0000000000000..a9446299372b7 --- /dev/null +++ b/taichi/math/eig.h @@ -0,0 +1,83 @@ +#include +#include "taichi/ir/frontend_ir.h" +#include "taichi/ir/ir.h" +#include "taichi/ir/expression_ops.h" + +TLANG_NAMESPACE_BEGIN + +template +std::tuple eig_3x3_export(ASTBuilder *ast_builder, + const Expr &a00, + const Expr &a01, + const Expr &a02, + const Expr &a10, + const Expr &a11, + const Expr &a12, + const Expr &a20, + const Expr &a21, + const Expr &a22) { + static_assert(sizeof(Tf) == sizeof(Ti), ""); + constexpr Tf M_SQRT3 = 1.73205080756887729352744634151; + auto Var = + std::bind(&ASTBuilder::make_var, ast_builder, std::placeholders::_1); + + auto A00 = Var(Expr(Tf(0.0))); + auto A01 = Var(Expr(Tf(0.0))); + auto A02 = Var(Expr(Tf(0.0))); + auto A10 = Var(Expr(Tf(0.0))); + auto A11 = Var(Expr(Tf(0.0))); + auto A12 = Var(Expr(Tf(0.0))); + auto A20 = Var(Expr(Tf(0.0))); + auto A21 = Var(Expr(Tf(0.0))); + auto A22 = Var(Expr(Tf(0.0))); + ast_builder->insert_assignment(A00, a00); + ast_builder->insert_assignment(A01, a01); + ast_builder->insert_assignment(A02, a02); + ast_builder->insert_assignment(A10, a10); + ast_builder->insert_assignment(A11, a11); + ast_builder->insert_assignment(A12, a12); + ast_builder->insert_assignment(A20, a20); + ast_builder->insert_assignment(A21, a21); + ast_builder->insert_assignment(A22, a22); + + auto w0 = Var(Expr(Tf(0.0))); + auto w1 = Var(Expr(Tf(0.0))); + auto w2 = Var(Expr(Tf(0.0))); + + auto m = Var(Expr(Tf(0.0))); + auto c1 = Var(Expr(Tf(0.0))); + auto c0 = Var(Expr(Tf(0.0))); + auto p = Var(Expr(Tf(0.0))); + auto sqrt_p = Var(Expr(Tf(0.0))); + auto q = Var(Expr(Tf(0.0))); + auto c = Var(Expr(Tf(0.0))); + auto s = Var(Expr(Tf(0.0))); + auto phi = Var(Expr(Tf(0.0))); + auto de = Var(Expr(Tf(0.0))); + auto dd = Var(Expr(Tf(0.0))); + auto ee = Var(Expr(Tf(0.0))); + auto ff = Var(Expr(Tf(0.0))); + + ast_builder->insert_assignment(de, A01 * A12); + ast_builder->insert_assignment(dd, A01 * A01); + ast_builder->insert_assignment(ee, A12 * A12); + ast_builder->insert_assignment(ff, A02 * A02); + ast_builder->insert_assignment(m, A00 + A11 + A22); + ast_builder->insert_assignment(c1, A00 * A11 + A00 * A22 + A11 * A22 - (dd + ee + ff)); + ast_builder->insert_assignment(c0, A22 * dd + A00 * ee + A11 * ff - A00 * A11 * A22 - Expr(Tf(2.0)) * A02 * de); + ast_builder->insert_assignment(p, m * m - Expr(Tf(3.0)) * c1); + ast_builder->insert_assignment(q, m * (p - Expr(Tf(1.5)) * c1) - Expr(Tf(13.5)) * c0); + ast_builder->insert_assignment(sqrt_p, sqrt(abs(p))); + ast_builder->insert_assignment(phi, Expr(Tf(27.0)) * (Expr(Tf(0.25)) * c1 * c1 * (p - c1) + c0 * (q + Expr(Tf(6.75)) * c0))); + ast_builder->insert_assignment(phi, Expr(Tf(1.0))/ Expr(Tf(3.0)) * atan2(sqrt(abs(phi)), q)); + ast_builder->insert_assignment(c, sqrt_p * cos(phi)); + ast_builder->insert_assignment(s, Expr(Tf(1.0)) / Expr(M_SQRT3) * sqrt_p * sin(phi)); + + ast_builder->insert_assignment(w2, Expr(Tf(1.0)) / Expr(Tf(3.0)) * (m - c)); + ast_builder->insert_assignment(w1, w2 + s); + ast_builder->insert_assignment(w0, w2 + c); + ast_builder->insert_assignment(w2, w2 - s); + return std::make_tuple(w0, w1, w2); +} + +TLANG_NAMESPACE_END diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index 5a6acabc655ce..1e1214010867a 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -21,6 +21,7 @@ #include "taichi/program/ndarray.h" #include "taichi/python/export.h" #include "taichi/math/svd.h" +#include "taichi/math/eig.h" #include "taichi/util/statistics.h" #include "taichi/util/action_recorder.h" #include "taichi/system/timeline.h" @@ -287,6 +288,8 @@ void export_lang(py::module &m) { .def("insert_patch_idx_expr", &ASTBuilder::insert_patch_idx_expr) .def("sifakis_svd_f32", sifakis_svd_export) .def("sifakis_svd_f64", sifakis_svd_export) + .def("eig_3x3_f32", eig_3x3_export) + .def("eig_3x3_f64", eig_3x3_export) .def("expr_var", &ASTBuilder::make_var) .def("bit_vectorize", &ASTBuilder::bit_vectorize) .def("parallelize", &ASTBuilder::parallelize) From bafbb9ed816b8246702b784309a585fa0a8b8d7f Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Fri, 18 Mar 2022 16:28:48 +0800 Subject: [PATCH 2/8] format --- misc/test_eig3x3.py | 7 ++++-- taichi/math/eig.h | 56 ++++++++++++++++++++++++++------------------- 2 files changed, 37 insertions(+), 26 deletions(-) diff --git a/misc/test_eig3x3.py b/misc/test_eig3x3.py index 12a64233c5a29..876277106daac 100644 --- a/misc/test_eig3x3.py +++ b/misc/test_eig3x3.py @@ -1,10 +1,13 @@ -import taichi as ti +import taichi as ti + ti.init() + @ti.kernel def test(): - A = ti.Matrix([[3, 1, 1],[1, 2, 2], [1, 2, 2]]) + A = ti.Matrix([[3, 1, 1], [1, 2, 2], [1, 2, 2]]) diagnal = ti.sym_eig(A) print(diagnal) + test() diff --git a/taichi/math/eig.h b/taichi/math/eig.h index a9446299372b7..f5238bb4cf01c 100644 --- a/taichi/math/eig.h +++ b/taichi/math/eig.h @@ -21,24 +21,24 @@ std::tuple eig_3x3_export(ASTBuilder *ast_builder, auto Var = std::bind(&ASTBuilder::make_var, ast_builder, std::placeholders::_1); - auto A00 = Var(Expr(Tf(0.0))); - auto A01 = Var(Expr(Tf(0.0))); - auto A02 = Var(Expr(Tf(0.0))); - auto A10 = Var(Expr(Tf(0.0))); - auto A11 = Var(Expr(Tf(0.0))); - auto A12 = Var(Expr(Tf(0.0))); - auto A20 = Var(Expr(Tf(0.0))); - auto A21 = Var(Expr(Tf(0.0))); - auto A22 = Var(Expr(Tf(0.0))); - ast_builder->insert_assignment(A00, a00); - ast_builder->insert_assignment(A01, a01); - ast_builder->insert_assignment(A02, a02); - ast_builder->insert_assignment(A10, a10); - ast_builder->insert_assignment(A11, a11); - ast_builder->insert_assignment(A12, a12); - ast_builder->insert_assignment(A20, a20); - ast_builder->insert_assignment(A21, a21); - ast_builder->insert_assignment(A22, a22); + auto A00 = Var(Expr(Tf(0.0))); + auto A01 = Var(Expr(Tf(0.0))); + auto A02 = Var(Expr(Tf(0.0))); + auto A10 = Var(Expr(Tf(0.0))); + auto A11 = Var(Expr(Tf(0.0))); + auto A12 = Var(Expr(Tf(0.0))); + auto A20 = Var(Expr(Tf(0.0))); + auto A21 = Var(Expr(Tf(0.0))); + auto A22 = Var(Expr(Tf(0.0))); + ast_builder->insert_assignment(A00, a00); + ast_builder->insert_assignment(A01, a01); + ast_builder->insert_assignment(A02, a02); + ast_builder->insert_assignment(A10, a10); + ast_builder->insert_assignment(A11, a11); + ast_builder->insert_assignment(A12, a12); + ast_builder->insert_assignment(A20, a20); + ast_builder->insert_assignment(A21, a21); + ast_builder->insert_assignment(A22, a22); auto w0 = Var(Expr(Tf(0.0))); auto w1 = Var(Expr(Tf(0.0))); @@ -63,15 +63,23 @@ std::tuple eig_3x3_export(ASTBuilder *ast_builder, ast_builder->insert_assignment(ee, A12 * A12); ast_builder->insert_assignment(ff, A02 * A02); ast_builder->insert_assignment(m, A00 + A11 + A22); - ast_builder->insert_assignment(c1, A00 * A11 + A00 * A22 + A11 * A22 - (dd + ee + ff)); - ast_builder->insert_assignment(c0, A22 * dd + A00 * ee + A11 * ff - A00 * A11 * A22 - Expr(Tf(2.0)) * A02 * de); + ast_builder->insert_assignment( + c1, A00 * A11 + A00 * A22 + A11 * A22 - (dd + ee + ff)); + ast_builder->insert_assignment(c0, A22 * dd + A00 * ee + A11 * ff - + A00 * A11 * A22 - + Expr(Tf(2.0)) * A02 * de); ast_builder->insert_assignment(p, m * m - Expr(Tf(3.0)) * c1); - ast_builder->insert_assignment(q, m * (p - Expr(Tf(1.5)) * c1) - Expr(Tf(13.5)) * c0); + ast_builder->insert_assignment( + q, m * (p - Expr(Tf(1.5)) * c1) - Expr(Tf(13.5)) * c0); ast_builder->insert_assignment(sqrt_p, sqrt(abs(p))); - ast_builder->insert_assignment(phi, Expr(Tf(27.0)) * (Expr(Tf(0.25)) * c1 * c1 * (p - c1) + c0 * (q + Expr(Tf(6.75)) * c0))); - ast_builder->insert_assignment(phi, Expr(Tf(1.0))/ Expr(Tf(3.0)) * atan2(sqrt(abs(phi)), q)); + ast_builder->insert_assignment( + phi, Expr(Tf(27.0)) * (Expr(Tf(0.25)) * c1 * c1 * (p - c1) + + c0 * (q + Expr(Tf(6.75)) * c0))); + ast_builder->insert_assignment( + phi, Expr(Tf(1.0)) / Expr(Tf(3.0)) * atan2(sqrt(abs(phi)), q)); ast_builder->insert_assignment(c, sqrt_p * cos(phi)); - ast_builder->insert_assignment(s, Expr(Tf(1.0)) / Expr(M_SQRT3) * sqrt_p * sin(phi)); + ast_builder->insert_assignment( + s, Expr(Tf(1.0)) / Expr(M_SQRT3) * sqrt_p * sin(phi)); ast_builder->insert_assignment(w2, Expr(Tf(1.0)) / Expr(Tf(3.0)) * (m - c)); ast_builder->insert_assignment(w1, w2 + s); From d9bcb948750f7120e18bc9811389332463f6f180 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Fri, 18 Mar 2022 16:40:48 +0800 Subject: [PATCH 3/8] pylint --- python/taichi/_funcs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py index ed145607c711d..46b6bc248443e 100644 --- a/python/taichi/_funcs.py +++ b/python/taichi/_funcs.py @@ -458,7 +458,7 @@ def sym_eig(A, dt=None): dt = impl.get_runtime().default_fp if A.n == 2: return sym_eig2x2(A, dt) - elif A.n == 3: + if A.n == 3: return sym_eig3x3(A, dt) raise Exception("Symmetric eigen solver only supports 2D matrices.") From 9b4e36abea1bb9ec2cff911ec6634a9d796709e6 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Mon, 21 Mar 2022 10:40:00 +0800 Subject: [PATCH 4/8] implement eig3x3 in python --- misc/test_eig3x3.py | 7 +-- python/taichi/_funcs.py | 78 +++++++++++++++++++++++++----- taichi/math/eig.h | 91 ----------------------------------- taichi/python/export_lang.cpp | 3 -- 4 files changed, 69 insertions(+), 110 deletions(-) delete mode 100644 taichi/math/eig.h diff --git a/misc/test_eig3x3.py b/misc/test_eig3x3.py index 876277106daac..c2a1e3940f732 100644 --- a/misc/test_eig3x3.py +++ b/misc/test_eig3x3.py @@ -5,9 +5,10 @@ @ti.kernel def test(): - A = ti.Matrix([[3, 1, 1], [1, 2, 2], [1, 2, 2]]) - diagnal = ti.sym_eig(A) - print(diagnal) + A = ti.Matrix([[3.0, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 2.0]], + dt=ti.f32) + diagnal, Q = ti.sym_eig(A) + print(diagnal, Q) test() diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py index 46b6bc248443e..bcdf0c86e12b3 100644 --- a/python/taichi/_funcs.py +++ b/python/taichi/_funcs.py @@ -306,10 +306,11 @@ def sym_eig2x2(A, dt): return eigenvalues, eigenvectors +@func def sym_eig3x3(A, dt): - """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 3x3 real symmetric matrix. + """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 3x3 real symmetric matrix using Cardano's method. - Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix. + Mathematical concept refers to https://www.mpi-hd.mpg.de/personalhomes/globes/3x3/. Args: A (ti.Matrix(3, 3)): input 3x3 symmetric matrix `A`. @@ -319,18 +320,69 @@ def sym_eig3x3(A, dt): eigenvalues (ti.Vector(3)): The eigenvalues. Each entry store one eigen value. eigenvectors (ti.Matrix(3, 3)): The eigenvectors. Each column stores one eigenvector. """ - assert A.n == 3 and A.m == 3 - inputs = tuple([e.ptr for e in A.entries]) - assert dt in [f32, f64] - if dt == f32: - rets = get_runtime().prog.current_ast_builder().eig_3x3_f32(*inputs) + M_SQRT3 = 1.73205080756887729352744634151 + m = A.trace() + dd = A[0, 1] * A[0, 1] + ee = A[1, 2] * A[1, 2] + ff = A[0, 2] * A[0, 2] + c1 = A[0, 0] * A[1, 1] + A[0, 0] * A[2, 2] + A[1, 1] * A[2, 2] - (dd + ee + + ff) + c0 = A[2, 2] * dd + A[0, 0] * ee + A[1, 1] * ff - A[0, 0] * A[1, 1] * A[ + 2, 2] - 2.0 * A[0, 2] * A[0, 1] * A[1, 2] + + p = m * m - 3.0 * c1 + q = m * (p - 1.5 * c1) - 13.5 * c0 + sqrt_p = ops.sqrt(ops.abs(p)) + phi = 27.0 * (0.25 * c1 * c1 * (p - c1) + c0 * (q + 6.75 * c0)) + phi = (1.0 / 3.0) * ops.atan2(ops.sqrt(ops.abs(phi)), q) + + c = sqrt_p * ops.cos(phi) + s = (1.0 / M_SQRT3) * sqrt_p * ops.sin(phi) + eigenvalues = Vector([0.0, 0.0, 0.0], dt=dt) + eigenvalues[2] = (1.0 / 3.0) * (m - c) + eigenvalues[1] = eigenvalues[2] + s + eigenvalues[0] = eigenvalues[2] + c + eigenvalues[2] = eigenvalues[2] - s + + t = ops.abs(eigenvalues[0]) + u = ops.abs(eigenvalues[1]) + if u > t: + t = u + u = ops.abs(eigenvalues[2]) + if u > t: + t = u + if t < 1.0: + u = t else: - rets = get_runtime().prog.current_ast_builder().eig_3x3_f64(*inputs) - assert len(rets) == 3 - D = expr_init(Matrix.zero(dt, 3, 1)) - for i in range(3): - D(i)._assign(rets[i]) - return D + u = t * t + Q = Matrix.zero(dt, 3, 3) + Q[0, 1] = A[0, 1] * A[1, 2] - A[0, 2] * A[1, 1] + Q[1, 1] = A[0, 2] * A[0, 1] - A[1, 2] * A[0, 0] + Q[2, 1] = A[0, 1] * A[0, 1] + + Q[0, 0] = Q[0, 1] + A[0, 2] * eigenvalues[0] + Q[1, 0] = Q[1, 1] + A[1, 2] * eigenvalues[0] + Q[2, 0] = (A[0, 0] - eigenvalues[0]) * (A[1, 1] - eigenvalues[0]) - Q[2, 1] + norm = Q[0, 0] * Q[0, 0] + Q[1, 0] * Q[1, 0] + Q[2, 0] * Q[2, 0] + norm = ops.sqrt(1.0 / norm) + Q[0, 0] *= norm + Q[1, 0] *= norm + Q[2, 0] *= norm + + Q[0, 1] = Q[0, 1] + A[0, 2] * eigenvalues[1] + Q[1, 1] = Q[1, 1] + A[1, 2] * eigenvalues[1] + Q[2, 1] = (A[0, 0] - eigenvalues[1]) * (A[1, 1] - eigenvalues[1]) - Q[2, 1] + norm = Q[0, 1] * Q[0, 1] + Q[1, 1] * Q[1, 1] + Q[2, 1] * Q[2, 1] + + norm = ops.sqrt(1.0 / norm) + Q[0, 1] *= norm + Q[1, 1] *= norm + Q[2, 1] *= norm + + Q[0, 2] = Q[1, 0] * Q[2, 1] - Q[2, 0] * Q[1, 1] + Q[1, 2] = Q[2, 0] * Q[0, 1] - Q[0, 0] * Q[2, 1] + Q[2, 2] = Q[0, 0] * Q[1, 1] - Q[1, 0] * Q[0, 1] + return eigenvalues, Q @func diff --git a/taichi/math/eig.h b/taichi/math/eig.h deleted file mode 100644 index f5238bb4cf01c..0000000000000 --- a/taichi/math/eig.h +++ /dev/null @@ -1,91 +0,0 @@ -#include -#include "taichi/ir/frontend_ir.h" -#include "taichi/ir/ir.h" -#include "taichi/ir/expression_ops.h" - -TLANG_NAMESPACE_BEGIN - -template -std::tuple eig_3x3_export(ASTBuilder *ast_builder, - const Expr &a00, - const Expr &a01, - const Expr &a02, - const Expr &a10, - const Expr &a11, - const Expr &a12, - const Expr &a20, - const Expr &a21, - const Expr &a22) { - static_assert(sizeof(Tf) == sizeof(Ti), ""); - constexpr Tf M_SQRT3 = 1.73205080756887729352744634151; - auto Var = - std::bind(&ASTBuilder::make_var, ast_builder, std::placeholders::_1); - - auto A00 = Var(Expr(Tf(0.0))); - auto A01 = Var(Expr(Tf(0.0))); - auto A02 = Var(Expr(Tf(0.0))); - auto A10 = Var(Expr(Tf(0.0))); - auto A11 = Var(Expr(Tf(0.0))); - auto A12 = Var(Expr(Tf(0.0))); - auto A20 = Var(Expr(Tf(0.0))); - auto A21 = Var(Expr(Tf(0.0))); - auto A22 = Var(Expr(Tf(0.0))); - ast_builder->insert_assignment(A00, a00); - ast_builder->insert_assignment(A01, a01); - ast_builder->insert_assignment(A02, a02); - ast_builder->insert_assignment(A10, a10); - ast_builder->insert_assignment(A11, a11); - ast_builder->insert_assignment(A12, a12); - ast_builder->insert_assignment(A20, a20); - ast_builder->insert_assignment(A21, a21); - ast_builder->insert_assignment(A22, a22); - - auto w0 = Var(Expr(Tf(0.0))); - auto w1 = Var(Expr(Tf(0.0))); - auto w2 = Var(Expr(Tf(0.0))); - - auto m = Var(Expr(Tf(0.0))); - auto c1 = Var(Expr(Tf(0.0))); - auto c0 = Var(Expr(Tf(0.0))); - auto p = Var(Expr(Tf(0.0))); - auto sqrt_p = Var(Expr(Tf(0.0))); - auto q = Var(Expr(Tf(0.0))); - auto c = Var(Expr(Tf(0.0))); - auto s = Var(Expr(Tf(0.0))); - auto phi = Var(Expr(Tf(0.0))); - auto de = Var(Expr(Tf(0.0))); - auto dd = Var(Expr(Tf(0.0))); - auto ee = Var(Expr(Tf(0.0))); - auto ff = Var(Expr(Tf(0.0))); - - ast_builder->insert_assignment(de, A01 * A12); - ast_builder->insert_assignment(dd, A01 * A01); - ast_builder->insert_assignment(ee, A12 * A12); - ast_builder->insert_assignment(ff, A02 * A02); - ast_builder->insert_assignment(m, A00 + A11 + A22); - ast_builder->insert_assignment( - c1, A00 * A11 + A00 * A22 + A11 * A22 - (dd + ee + ff)); - ast_builder->insert_assignment(c0, A22 * dd + A00 * ee + A11 * ff - - A00 * A11 * A22 - - Expr(Tf(2.0)) * A02 * de); - ast_builder->insert_assignment(p, m * m - Expr(Tf(3.0)) * c1); - ast_builder->insert_assignment( - q, m * (p - Expr(Tf(1.5)) * c1) - Expr(Tf(13.5)) * c0); - ast_builder->insert_assignment(sqrt_p, sqrt(abs(p))); - ast_builder->insert_assignment( - phi, Expr(Tf(27.0)) * (Expr(Tf(0.25)) * c1 * c1 * (p - c1) + - c0 * (q + Expr(Tf(6.75)) * c0))); - ast_builder->insert_assignment( - phi, Expr(Tf(1.0)) / Expr(Tf(3.0)) * atan2(sqrt(abs(phi)), q)); - ast_builder->insert_assignment(c, sqrt_p * cos(phi)); - ast_builder->insert_assignment( - s, Expr(Tf(1.0)) / Expr(M_SQRT3) * sqrt_p * sin(phi)); - - ast_builder->insert_assignment(w2, Expr(Tf(1.0)) / Expr(Tf(3.0)) * (m - c)); - ast_builder->insert_assignment(w1, w2 + s); - ast_builder->insert_assignment(w0, w2 + c); - ast_builder->insert_assignment(w2, w2 - s); - return std::make_tuple(w0, w1, w2); -} - -TLANG_NAMESPACE_END diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index 1e1214010867a..5a6acabc655ce 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -21,7 +21,6 @@ #include "taichi/program/ndarray.h" #include "taichi/python/export.h" #include "taichi/math/svd.h" -#include "taichi/math/eig.h" #include "taichi/util/statistics.h" #include "taichi/util/action_recorder.h" #include "taichi/system/timeline.h" @@ -288,8 +287,6 @@ void export_lang(py::module &m) { .def("insert_patch_idx_expr", &ASTBuilder::insert_patch_idx_expr) .def("sifakis_svd_f32", sifakis_svd_export) .def("sifakis_svd_f64", sifakis_svd_export) - .def("eig_3x3_f32", eig_3x3_export) - .def("eig_3x3_f64", eig_3x3_export) .def("expr_var", &ASTBuilder::make_var) .def("bit_vectorize", &ASTBuilder::bit_vectorize) .def("parallelize", &ASTBuilder::parallelize) From 496d685638479f109efa241bb9930dd1523a921d Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Mon, 21 Mar 2022 11:14:19 +0800 Subject: [PATCH 5/8] add eig3x3 decomposition test --- tests/python/test_eig.py | 41 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/python/test_eig.py b/tests/python/test_eig.py index 3bbb8dceca540..d0ffd33f65fe6 100644 --- a/tests/python/test_eig.py +++ b/tests/python/test_eig.py @@ -108,6 +108,35 @@ def eigen_solve(): _eigen_vector_equal(w_ti[:, idx_ti[1]], w_np[:, idx_np[1]], tol) +def _test_sym_eig3x3(dt): + A = ti.Matrix.field(3, 3, dtype=dt, shape=()) + v = ti.Vector.field(3, dtype=dt, shape=()) + w = ti.Matrix.field(3, 3, dtype=dt, shape=()) + + A[None] = [[3.0, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 2.0]] + + @ti.kernel + def eigen_solve(): + v[None], w[None] = ti.sym_eig(A[None]) + + tol = 1e-5 if dt == ti.f32 else 1e-12 + dtype = np.float32 if dt == ti.f32 else np.float64 + + eigen_solve() + v_np, w_np = np.linalg.eig(A.to_numpy().astype(dtype)) + v_ti = v.to_numpy().astype(dtype) + w_ti = w.to_numpy().astype(dtype) + + # sort by eigenvalues + idx_np = np.argsort(v_np) + idx_ti = np.argsort(v_ti) + + np.testing.assert_allclose(v_ti[idx_ti], v_np[idx_np], atol=tol, rtol=tol) + _eigen_vector_equal(w_ti[:, idx_ti[0]], w_np[:, idx_np[0]], tol) + _eigen_vector_equal(w_ti[:, idx_ti[1]], w_np[:, idx_np[1]], tol) + _eigen_vector_equal(w_ti[:, idx_ti[2]], w_np[:, idx_np[2]], tol) + + @pytest.mark.parametrize("func", [_test_eig2x2_real, _test_eig2x2_complex]) @test_utils.test(default_fp=ti.f32, fast_math=False) def test_eig2x2_f32(func): @@ -132,3 +161,15 @@ def test_sym_eig2x2_f32(): fast_math=False) def test_sym_eig2x2_f64(): _test_sym_eig2x2(ti.f64) + + +@test_utils.test(default_fp=ti.f32, fast_math=False) +def test_sym_eig3x3_f32(): + _test_sym_eig3x3(ti.f32) + + +@test_utils.test(require=ti.extension.data64, + default_fp=ti.f64, + fast_math=False) +def test_sym_eig3x3_f64(): + _test_sym_eig3x3(ti.f64) From ec2c9190a20a8464a13a4a1ea1cd8f88e4524460 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Mon, 21 Mar 2022 11:18:11 +0800 Subject: [PATCH 6/8] rm misc --- misc/test_eig3x3.py | 14 -------------- 1 file changed, 14 deletions(-) delete mode 100644 misc/test_eig3x3.py diff --git a/misc/test_eig3x3.py b/misc/test_eig3x3.py deleted file mode 100644 index c2a1e3940f732..0000000000000 --- a/misc/test_eig3x3.py +++ /dev/null @@ -1,14 +0,0 @@ -import taichi as ti - -ti.init() - - -@ti.kernel -def test(): - A = ti.Matrix([[3.0, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 2.0]], - dt=ti.f32) - diagnal, Q = ti.sym_eig(A) - print(diagnal, Q) - - -test() From d62458b955b42461703e62f03ac0b87f6fae89c4 Mon Sep 17 00:00:00 2001 From: Peng Yu <6712304+FantasyVR@users.noreply.github.com> Date: Wed, 23 Mar 2022 15:09:26 +0800 Subject: [PATCH 7/8] Update python/taichi/_funcs.py Co-authored-by: Yuanming Hu --- python/taichi/_funcs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py index bcdf0c86e12b3..bd865b49b20f7 100644 --- a/python/taichi/_funcs.py +++ b/python/taichi/_funcs.py @@ -512,7 +512,7 @@ def sym_eig(A, dt=None): return sym_eig2x2(A, dt) if A.n == 3: return sym_eig3x3(A, dt) - raise Exception("Symmetric eigen solver only supports 2D matrices.") + raise Exception("Symmetric eigen solver only supports 2D and 3D matrices.") __all__ = ['randn', 'polar_decompose', 'eig', 'sym_eig', 'svd'] From d6fa2d9666567602db436af2247125016e843944 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Wed, 23 Mar 2022 15:35:26 +0800 Subject: [PATCH 8/8] Add more test cases --- tests/python/test_eig.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/python/test_eig.py b/tests/python/test_eig.py index d0ffd33f65fe6..6a76eef4fb2dd 100644 --- a/tests/python/test_eig.py +++ b/tests/python/test_eig.py @@ -108,12 +108,12 @@ def eigen_solve(): _eigen_vector_equal(w_ti[:, idx_ti[1]], w_np[:, idx_np[1]], tol) -def _test_sym_eig3x3(dt): +def _test_sym_eig3x3(dt, a00): A = ti.Matrix.field(3, 3, dtype=dt, shape=()) v = ti.Vector.field(3, dtype=dt, shape=()) w = ti.Matrix.field(3, 3, dtype=dt, shape=()) - A[None] = [[3.0, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 2.0]] + A[None] = [[a00, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 2.0]] @ti.kernel def eigen_solve(): @@ -163,13 +163,15 @@ def test_sym_eig2x2_f64(): _test_sym_eig2x2(ti.f64) +@pytest.mark.parametrize('a00', [i for i in range(10)]) @test_utils.test(default_fp=ti.f32, fast_math=False) -def test_sym_eig3x3_f32(): - _test_sym_eig3x3(ti.f32) +def test_sym_eig3x3_f32(a00): + _test_sym_eig3x3(ti.f32, a00) +@pytest.mark.parametrize('a00', [i for i in range(10)]) @test_utils.test(require=ti.extension.data64, default_fp=ti.f64, fast_math=False) -def test_sym_eig3x3_f64(): - _test_sym_eig3x3(ti.f64) +def test_sym_eig3x3_f64(a00): + _test_sym_eig3x3(ti.f64, a00)