test_symeig3x3.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


import unittest

import torch
from common_testing import get_random_cuda_device, TestCaseMixin
from pytorch3d.common.workaround import symeig3x3
from pytorch3d.transforms.rotation_conversions import random_rotations


class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
    TEST_BATCH_SIZE = 1024

    @staticmethod
    def create_random_sym3x3(device, n):
        random_3x3 = torch.randn((n, 3, 3), device=device)
        random_3x3_T = torch.transpose(random_3x3, 1, 2)
        random_sym_3x3 = (random_3x3 * random_3x3_T).contiguous()

        return random_sym_3x3

    @staticmethod
    def create_diag_sym3x3(device, n, noise=0.0):
        # Create purly diagonal matrices
        random_diag_3x3 = torch.randn((n, 3), device=device).diag_embed()

        # Make them 'almost' diagonal
        random_diag_3x3 += noise * TestSymEig3x3.create_random_sym3x3(device, n)

        return random_diag_3x3

    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(42)

        self._gpu = get_random_cuda_device()
        self._cpu = torch.device("cpu")

    def test_is_eigen_gpu(self):
        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)

        self._test_is_eigen(test_input)

    def test_is_eigen_cpu(self):
        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)

        self._test_is_eigen(test_input)

    def _test_is_eigen(self, test_input, atol=1e-04, rtol=1e-02):
        """
        Verify that values and vectors produced are really eigenvalues and eigenvectors
        and can restore the original input matrix with good precision
        """
        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)

        self.assertClose(
            test_input,
            eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1),
            atol=atol,
            rtol=rtol,
        )

    def test_eigenvectors_are_orthonormal_gpu(self):
        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)

        self._test_eigenvectors_are_orthonormal(test_input)

    def test_eigenvectors_are_orthonormal_cpu(self):
        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)

        self._test_eigenvectors_are_orthonormal(test_input)

    def _test_eigenvectors_are_orthonormal(self, test_input):
        """
        Verify that eigenvectors are an orthonormal set
        """
        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)

        batched_eye = torch.zeros_like(test_input)
        batched_eye[..., :, :] = torch.eye(3, device=batched_eye.device)

        self.assertClose(
            batched_eye, eigenvectors @ eigenvectors.transpose(-2, -1), atol=1e-06
        )

    def test_is_not_nan_or_inf_gpu(self):
        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)

        self._test_is_not_nan_or_inf(test_input)

    def test_is_not_nan_or_inf_cpu(self):
        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)

        self._test_is_not_nan_or_inf(test_input)

    def _test_is_not_nan_or_inf(self, test_input):
        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)

        self.assertTrue(torch.isfinite(eigenvalues).all())
        self.assertTrue(torch.isfinite(eigenvectors).all())

    def test_degenerate_inputs_gpu(self):
        self._test_degenerate_inputs(self._gpu)

    def test_degenerate_inputs_cpu(self):
        self._test_degenerate_inputs(self._cpu)

    def _test_degenerate_inputs(self, device):
        """
        Test degenerate case when input matrices are diagonal or near-diagonal
        """

        # Purely diagonal case
        test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE)

        self._test_is_not_nan_or_inf(test_input)
        self._test_is_eigen(test_input)
        self._test_eigenvectors_are_orthonormal(test_input)

        # Almost-diagonal case
        test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE, noise=1e-4)

        self._test_is_not_nan_or_inf(test_input)
        self._test_is_eigen(test_input)
        self._test_eigenvectors_are_orthonormal(test_input)

    def test_gradients_cpu(self):
        self._test_gradients(self._cpu)

    def test_gradients_gpu(self):
        self._test_gradients(self._gpu)

    def _test_gradients(self, device):
        """
        Tests if gradients pass though without any problems (infs, nans etc) and
        also performs gradcheck (compares numerical and analytical gradients)
        """
        test_random_input = self.create_random_sym3x3(device, n=16)
        test_diag_input = self.create_diag_sym3x3(device, n=16)
        test_almost_diag_input = self.create_diag_sym3x3(device, n=16, noise=1e-4)

        test_input = torch.cat(
            (test_random_input, test_diag_input, test_almost_diag_input)
        )
        test_input.requires_grad = True

        with torch.autograd.detect_anomaly():
            eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)

            loss = eigenvalues.mean() + eigenvectors.mean()
            loss.backward()

        test_random_input.requires_grad = True
        # Inputs are converted to double to increase the precision of gradcheck.
        torch.autograd.gradcheck(
            symeig3x3, test_random_input.double(), eps=1e-6, atol=1e-2, rtol=1e-2
        )

    def _test_eigenvalues_and_eigenvectors(
        self, test_eigenvectors, test_eigenvalues, atol=1e-04, rtol=1e-04
    ):
        test_input = (
            test_eigenvectors.transpose(-2, -1)
            @ test_eigenvalues.diag_embed()
            @ test_eigenvectors
        )

        test_eigenvalues_sorted, _ = torch.sort(test_eigenvalues, dim=-1)

        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)

        self.assertClose(
            test_eigenvalues_sorted,
            eigenvalues,
            atol=atol,
            rtol=rtol,
        )

        self._test_is_not_nan_or_inf(test_input)
        self._test_is_eigen(test_input, atol=atol, rtol=rtol)
        self._test_eigenvectors_are_orthonormal(test_input)

    def test_degenerate_eigenvalues_gpu(self):
        self._test_degenerate_eigenvalues(self._gpu)

    def test_degenerate_eigenvalues_cpu(self):
        self._test_degenerate_eigenvalues(self._cpu)

    def _test_degenerate_eigenvalues(self, device):
        """
        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity
        """
        # Error tolerances for degenerate values are increased as things might become
        #  numerically unstable
        deg_atol = 1e-3
        deg_rtol = 1.0

        # Construct random orthonormal sets
        test_eigenvecs = random_rotations(n=self.TEST_BATCH_SIZE, device=device)

        # Construct random eigenvalues
        test_eigenvals = torch.randn(
            (self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device
        )
        self._test_eigenvalues_and_eigenvectors(
            test_eigenvecs, test_eigenvals, atol=deg_atol, rtol=deg_rtol
        )

        # First eigenvalue is always 0.0 here: [0.0 X Y]
        test_eigenvals_with_zero = test_eigenvals.clone()
        test_eigenvals_with_zero[..., 0] = 0.0
        self._test_eigenvalues_and_eigenvectors(
            test_eigenvecs, test_eigenvals_with_zero, atol=deg_atol, rtol=deg_rtol
        )

        # First two eigenvalues are always the same here: [X X Y]
        test_eigenvals_with_multiplicity2 = test_eigenvals.clone()
        test_eigenvals_with_multiplicity2[..., 1] = test_eigenvals_with_multiplicity2[
            ..., 0
        ]
        self._test_eigenvalues_and_eigenvectors(
            test_eigenvecs,
            test_eigenvals_with_multiplicity2,
            atol=deg_atol,
            rtol=deg_rtol,
        )

        # All three eigenvalues are the same here: [X X X]
        test_eigenvals_with_multiplicity3 = test_eigenvals_with_multiplicity2.clone()
        test_eigenvals_with_multiplicity3[..., 2] = test_eigenvals_with_multiplicity2[
            ..., 0
        ]
        self._test_eigenvalues_and_eigenvectors(
            test_eigenvecs,
            test_eigenvals_with_multiplicity3,
            atol=deg_atol,
            rtol=deg_rtol,
        )

    def test_more_dimensions(self):
        """
        Tests if function supports arbitrary leading dimensions
        """
        repeat = 4

        test_input = self.create_random_sym3x3(self._cpu, n=16)
        test_input_4d = test_input[None, ...].expand((repeat,) + test_input.shape)

        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
        eigenvalues_4d, eigenvectors_4d = symeig3x3(test_input_4d, eigenvectors=True)

        eigenvalues_4d_gt = eigenvalues[None, ...].expand((repeat,) + eigenvalues.shape)
        eigenvectors_4d_gt = eigenvectors[None, ...].expand(
            (repeat,) + eigenvectors.shape
        )

        self.assertClose(eigenvalues_4d_gt, eigenvalues_4d)
        self.assertClose(eigenvectors_4d_gt, eigenvectors_4d)