[Lang] Matrix 3x3 eigen decomposition

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']

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)