Initial commit

fbshipit-source-id: ad58e416e3ceeca85fae0583308968d04e78fe0d

tests/test_so3.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+
+import numpy as np
+import unittest
+import torch
+
+from pytorch3d.transforms.so3 import (
+    hat,
+    so3_exponential_map,
+    so3_log_map,
+    so3_relative_angle,
+)
+
+
+class TestSO3(unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+        np.random.seed(42)
+
+    @staticmethod
+    def init_log_rot(batch_size: int = 10):
+        """
+        Initialize a list of `batch_size` 3-dimensional vectors representing
+        randomly generated logarithms of rotation matrices.
+        """
+        device = torch.device("cuda:0")
+        log_rot = torch.randn(
+            (batch_size, 3), dtype=torch.float32, device=device
+        )
+        return log_rot
+
+    @staticmethod
+    def init_rot(batch_size: int = 10):
+        """
+        Randomly generate a batch of `batch_size` 3x3 rotation matrices.
+        """
+        device = torch.device("cuda:0")
+
+        # TODO(dnovotny): replace with random_rotation from random_rotation.py
+        rot = []
+        for _ in range(batch_size):
+            r = torch.qr(torch.randn((3, 3), device=device))[0]
+            f = torch.randint(2, (3,), device=device, dtype=torch.float32)
+            if f.sum() % 2 == 0:
+                f = 1 - f
+            rot.append(r * (2 * f - 1).float())
+        rot = torch.stack(rot)
+
+        return rot
+
+    def test_determinant(self):
+        """
+        Tests whether the determinants of 3x3 rotation matrices produced
+        by `so3_exponential_map` are (almost) equal to 1.
+        """
+        log_rot = TestSO3.init_log_rot(batch_size=30)
+        Rs = so3_exponential_map(log_rot)
+        for R in Rs:
+            det = np.linalg.det(R.cpu().numpy())
+            self.assertAlmostEqual(float(det), 1.0, 5)
+
+    def test_cross(self):
+        """
+        For a pair of randomly generated 3-dimensional vectors `a` and `b`,
+        tests whether a matrix product of `hat(a)` and `b` equals the result
+        of a cross product between `a` and `b`.
+        """
+        device = torch.device("cuda:0")
+        a, b = torch.randn((2, 100, 3), dtype=torch.float32, device=device)
+        hat_a = hat(a)
+        cross = torch.bmm(hat_a, b[:, :, None])[:, :, 0]
+        torch_cross = torch.cross(a, b, dim=1)
+        max_df = (cross - torch_cross).abs().max()
+        self.assertAlmostEqual(float(max_df), 0.0, 5)
+
+    def test_bad_so3_input_value_err(self):
+        """
+        Tests whether `so3_exponential_map` and `so3_log_map` correctly return
+        a ValueError if called with an argument of incorrect shape or, in case
+        of `so3_exponential_map`, unexpected trace.
+        """
+        device = torch.device("cuda:0")
+        log_rot = torch.randn(size=[5, 4], device=device)
+        with self.assertRaises(ValueError) as err:
+            so3_exponential_map(log_rot)
+        self.assertTrue(
+            "Input tensor shape has to be Nx3." in str(err.exception)
+        )
+
+        rot = torch.randn(size=[5, 3, 5], device=device)
+        with self.assertRaises(ValueError) as err:
+            so3_log_map(rot)
+        self.assertTrue(
+            "Input has to be a batch of 3x3 Tensors." in str(err.exception)
+        )
+
+        # trace of rot definitely bigger than 3 or smaller than -1
+        rot = torch.cat(
+            (
+                torch.rand(size=[5, 3, 3], device=device) + 4.0,
+                torch.rand(size=[5, 3, 3], device=device) - 3.0,
+            )
+        )
+        with self.assertRaises(ValueError) as err:
+            so3_log_map(rot)
+        self.assertTrue(
+            "A matrix has trace outside valid range [-1-eps,3+eps]."
+            in str(err.exception)
+        )
+
+    def test_so3_exp_singularity(self, batch_size: int = 100):
+        """
+        Tests whether the `so3_exponential_map` is robust to the input vectors
+        the norms of which are close to the numerically unstable region
+        (vectors with low l2-norms).
+        """
+        # generate random log-rotations with a tiny angle
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        log_rot_small = log_rot * 1e-6
+        R = so3_exponential_map(log_rot_small)
+        # tests whether all outputs are finite
+        R_sum = float(R.sum())
+        self.assertEqual(R_sum, R_sum)
+
+    def test_so3_log_singularity(self, batch_size: int = 100):
+        """
+        Tests whether the `so3_log_map` is robust to the input matrices
+        who's rotation angles are close to the numerically unstable region
+        (i.e. matrices with low rotation angles).
+        """
+        # generate random rotations with a tiny angle
+        device = torch.device("cuda:0")
+        r = torch.eye(3, device=device)[None].repeat((batch_size, 1, 1))
+        r += torch.randn((batch_size, 3, 3), device=device) * 1e-3
+        r = torch.stack([torch.qr(r_)[0] for r_ in r])
+        # the log of the rotation matrix r
+        r_log = so3_log_map(r)
+        # tests whether all outputs are finite
+        r_sum = float(r_log.sum())
+        self.assertEqual(r_sum, r_sum)
+
+    def test_so3_log_to_exp_to_log(self, batch_size: int = 100):
+        """
+        Check that `so3_log_map(so3_exponential_map(log_rot))==log_rot` for
+        a randomly generated batch of rotation matrix logarithms `log_rot`.
+        """
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        log_rot_ = so3_log_map(so3_exponential_map(log_rot))
+        max_df = (log_rot - log_rot_).abs().max()
+        self.assertAlmostEqual(float(max_df), 0.0, 4)
+
+    def test_so3_exp_to_log_to_exp(self, batch_size: int = 100):
+        """
+        Check that `so3_exponential_map(so3_log_map(R))==R` for
+        a batch of randomly generated rotation matrices `R`.
+        """
+        rot = TestSO3.init_rot(batch_size=batch_size)
+        rot_ = so3_exponential_map(so3_log_map(rot))
+        angles = so3_relative_angle(rot, rot_)
+        max_angle = angles.max()
+        # a lot of precision lost here :(
+        # TODO: fix this test??
+        self.assertTrue(np.allclose(float(max_angle), 0.0, atol=0.1))
+
+    def test_so3_cos_angle(self, batch_size: int = 100):
+        """
+        Check that `so3_relative_angle(R1, R2, cos_angle=False).cos()`
+        is the same as `so3_relative_angle(R1, R2, cos_angle=True)`
+        batches of randomly generated rotation matrices `R1` and `R2`.
+        """
+        rot1 = TestSO3.init_rot(batch_size=batch_size)
+        rot2 = TestSO3.init_rot(batch_size=batch_size)
+        angles = so3_relative_angle(rot1, rot2, cos_angle=False).cos()
+        angles_ = so3_relative_angle(rot1, rot2, cos_angle=True)
+        self.assertTrue(torch.allclose(angles, angles_))
+
+    @staticmethod
+    def so3_expmap(batch_size: int = 10):
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        torch.cuda.synchronize()
+
+        def compute_rots():
+            so3_exponential_map(log_rot)
+            torch.cuda.synchronize()
+
+        return compute_rots
+
+    @staticmethod
+    def so3_logmap(batch_size: int = 10):
+        log_rot = TestSO3.init_rot(batch_size=batch_size)
+        torch.cuda.synchronize()
+
+        def compute_logs():
+            so3_log_map(log_rot)
+            torch.cuda.synchronize()
+
+        return compute_logs