diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py index 0d9e0392e5223..bd865b49b20f7 100644 --- a/python/taichi/_funcs.py +++ b/python/taichi/_funcs.py @@ -306,6 +306,85 @@ 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 using Cardano's method. + + Mathematical concept refers to https://www.mpi-hd.mpg.de/personalhomes/globes/3x3/. + + 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. + """ + 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: + 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 def _svd(A, dt): """Perform singular value decomposition (A=USV^T) for arbitrary size matrix. @@ -431,7 +510,9 @@ def sym_eig(A, dt=None): dt = impl.get_runtime().default_fp if A.n == 2: return sym_eig2x2(A, dt) - raise Exception("Symmetric eigen solver only supports 2D matrices.") + if A.n == 3: + return sym_eig3x3(A, dt) + raise Exception("Symmetric eigen solver only supports 2D and 3D matrices.") __all__ = ['randn', 'polar_decompose', 'eig', 'sym_eig', 'svd'] diff --git a/tests/python/test_eig.py b/tests/python/test_eig.py index 3bbb8dceca540..6a76eef4fb2dd 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, 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] = [[a00, 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,17 @@ def test_sym_eig2x2_f32(): fast_math=False) 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(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(a00): + _test_sym_eig3x3(ti.f64, a00)