diff --git a/pytorch3d/transforms/rotation_conversions.py b/pytorch3d/transforms/rotation_conversions.py index eb1a42cc..d11fa49f 100644 --- a/pytorch3d/transforms/rotation_conversions.py +++ b/pytorch3d/transforms/rotation_conversions.py @@ -489,15 +489,18 @@ def axis_angle_to_matrix(axis_angle: torch.Tensor, fast: bool = False) -> torch. cross_product_matrix = torch.stack( [zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=-1 ).view(shape + (3,)) - cross_product_matrix_sqrd = cross_product_matrix @ cross_product_matrix + outer_product = axis_angle.unsqueeze(-1) * axis_angle.unsqueeze(-2) + # This is the Rodrigues formula written with the identities + # cross_product_matrix^2 == outer_product - angle^2 * I + # (1 - cos(angle)) / angle^2 == sinc(angle / (2*pi))^2 / 2 + # which save a matrix multiplication and, as sinc is defined at zero, + # remove the need to special-case a zero angle. identity = torch.eye(3, dtype=dtype, device=device) - angles_sqrd = angles * angles - angles_sqrd = torch.where(angles_sqrd == 0, 1, angles_sqrd) return ( - identity.expand(cross_product_matrix.shape) + torch.cos(angles) * identity + torch.sinc(angles / torch.pi) * cross_product_matrix - + ((1 - torch.cos(angles)) / angles_sqrd) * cross_product_matrix_sqrd + + 0.5 * torch.sinc(angles / (2 * torch.pi)) ** 2 * outer_product ) @@ -536,24 +539,52 @@ def matrix_to_axis_angle(matrix: torch.Tensor, fast: bool = False) -> torch.Tens traces = torch.diagonal(matrix, dim1=-2, dim2=-1).sum(-1).unsqueeze(-1) angles = torch.atan2(norms, traces - 1) - zeros = torch.zeros(3, dtype=matrix.dtype, device=matrix.device) - omegas = torch.where(torch.isclose(angles, torch.zeros_like(angles)), zeros, omegas) + # omegas is 2*sin(angle) times the axis, so the relative precision of the + # axis it gives decays like 1/(pi - angle). Above the following angle the + # axis is taken from the symmetric part of the matrix instead. Switching + # over 0.01 radians before pi keeps the error close to that of the slow + # implementation, while leaving under 1% of uniformly random rotations on + # the more expensive branch. + near_pi = angles > torch.pi - 1e-2 - near_pi = angles.isclose(angles.new_full((1,), torch.pi)).squeeze(-1) + # An angle of zero needs no special case: torch.sinc(0) is 1, and omegas + # is then zero, which is the right answer. torch.sinc is instead zero at + # an angle of pi, so those entries are given a harmless denominator here + # and overwritten below. + sincs = torch.sinc(angles / torch.pi) + axis_angles = omegas * (0.5 / torch.where(near_pi, 1.0, sincs)) - axis_angles = torch.empty_like(omegas) - axis_angles[~near_pi] = ( - 0.5 * omegas[~near_pi] / torch.sinc(angles[~near_pi] / torch.pi) + # Rotations by nearly pi are a small minority, so it is worth selecting + # them rather than evaluating the following densely. The selection index is + # computed once, because each boolean mask index would recompute it (and, + # on the GPU, synchronize). + flat_axis_angles = axis_angles.reshape(-1, 3) + index = near_pi.reshape(-1).nonzero().squeeze(-1) + near_pi_matrix = matrix.reshape(-1, 3, 3).index_select(0, index) + near_pi_omegas = omegas.reshape(-1, 3).index_select(0, index) + + # this derives from: R + R^T - 2*cos(angle)*I = 2*(1 - cos(angle))*nnT, + # together with trace(R) - 1 == 2*cos(angle) + double_cosines = traces.reshape(-1, 1, 1).index_select(0, index) - 1 + nnT = ( + near_pi_matrix + + near_pi_matrix.transpose(-1, -2) + - double_cosines * torch.eye(3, dtype=matrix.dtype, device=matrix.device) ) + # Every column of nnT is a multiple of n. Taking the one whose diagonal + # entry is largest means the multiplier has absolute value at least + # 1/sqrt(3), so the column is never degenerate. + largest = torch.diagonal(nnT, dim1=-2, dim2=-1).argmax(dim=-1) + n = nnT.gather(-1, largest.reshape(-1, 1, 1).expand(-1, 3, 1)).squeeze(-1) + n = n / torch.linalg.vector_norm(n, dim=-1, keepdim=True) + # nnT determines n only up to sign, which omegas resolves. (At exactly pi + # the sign is arbitrary, and there omegas is zero.) + n = torch.where((n * near_pi_omegas).sum(-1, keepdim=True) < 0, -n, n) - # this derives from: nnT = (R + 1) / 2 - n = 0.5 * ( - matrix[near_pi][..., 0, :] - + torch.eye(1, 3, dtype=matrix.dtype, device=matrix.device) + near_pi_angles = angles.reshape(-1, 1).index_select(0, index) + return torch.index_copy(flat_axis_angles, 0, index, near_pi_angles * n).reshape( + axis_angles.shape ) - axis_angles[near_pi] = angles[near_pi] * n / torch.norm(n) - - return axis_angles def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: diff --git a/tests/test_rotation_conversions.py b/tests/test_rotation_conversions.py index 1ed99453..211d2a48 100644 --- a/tests/test_rotation_conversions.py +++ b/tests/test_rotation_conversions.py @@ -228,6 +228,165 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase): self.assertClose(data, axis_angle_to_matrix(euler_angles_fast)) self.assertClose(data, axis_angle_to_matrix(euler_angles, fast=True)) + def test_axis_angle_by_pi(self): + """Rotations by exactly pi are recovered, and don't depend on the batch.""" + options = [0.0, -1.0, 1.0] + axes = torch.nn.functional.normalize( + torch.stack( + [ + torch.tensor(vec, dtype=torch.float64) + for vec in itertools.islice( # exclude [0, 0, 0] + itertools.product(options, options, options), 1, None + ) + ] + ), + dim=-1, + ) + # Rotation by pi around unit vector x is given by 2 x x^T - Id. + R = 2 * torch.matmul(axes[..., None], axes[..., None, :]) - torch.eye( + 3, dtype=torch.float64 + ) + for fast in [False, True]: + axis_angles = matrix_to_axis_angle(R, fast=fast) + self.assertClose( + axis_angles.norm(dim=-1), + torch.full(axes.shape[:1], math.pi, dtype=torch.float64), + ) + self.assertClose(axis_angle_to_matrix(axis_angles), R) + # The axis of one rotation cannot depend on the others. + singly = torch.stack( + [matrix_to_axis_angle(r, fast=fast) for r in R.unbind()] + ) + self.assertClose(singly, axis_angles) + + def test_axis_angle_near_pi(self): + """The fast path stays accurate as the angle approaches pi.""" + for dtype, atol in [(torch.float64, 1e-9), (torch.float32, 1e-4)]: + axes = torch.nn.functional.normalize( + torch.randn(50, 3, dtype=dtype), dim=-1 + ) + for angle in [ + math.pi, + math.pi - 1e-7, + math.pi - 1e-5, + math.pi - 1e-3, + math.pi - 1e-2, + math.pi - 0.1, + 2.0, + ]: + data = axes * angle + R = axis_angle_to_matrix(data) + axis_angles = matrix_to_axis_angle(R, fast=True) + self.assertClose(axis_angle_to_matrix(axis_angles), R, atol=atol) + self.assertClose( + axis_angles.norm(dim=-1), + torch.full(axes.shape[:1], angle, dtype=dtype), + atol=atol, + ) + if angle < math.pi - 1e-2: + # Closer to pi than this, the sign of the axis is not + # determined to float32 precision (and both signs describe + # the same rotation to within the tolerance above). + self.assertClose(axis_angles, data, atol=atol) + + def test_axis_angle_to_matrix_fast(self): + """The Rodrigues implementation agrees with the quaternion one.""" + data = torch.randn(100, 3, dtype=torch.float64) + normalized = torch.nn.functional.normalize(data, dim=-1) + data[:5] = 0.0 + data[5:10] = normalized[5:10] * math.pi + data[10:15] = normalized[10:15] * 1e-12 + self.assertClose( + axis_angle_to_matrix(data, fast=True), + axis_angle_to_matrix(data), + atol=1e-14, + ) + + # A zero angle is not a special case, but must still have gradients. + zeros = torch.zeros(1, 3, dtype=torch.float64, requires_grad=True) + matrices = axis_angle_to_matrix(zeros, fast=True) + self.assertClose(matrices, torch.eye(3, dtype=torch.float64)[None]) + (grad,) = torch.autograd.grad(matrices.sum(), zeros) + self.assertTrue(torch.isfinite(grad).all()) + + def test_axis_angle_tiny(self): + """Tiny angles are preserved, not flattened to zero, and have grads.""" + axis = torch.nn.functional.normalize( + torch.tensor([[0.3, -0.5, 0.81]], dtype=torch.float64), dim=-1 + ) + for angle in [1e-3, 1e-8, 1e-9, 1e-12]: + data = axis * angle + R = axis_angle_to_matrix(data) + for fast in [False, True]: + # atol must be 0 here: the default would accept zero. + self.assertClose(matrix_to_axis_angle(R, fast=fast), data, atol=0) + + identity = torch.eye(3, dtype=torch.float64) + for fast in [False, True]: + R = identity.clone().requires_grad_(True) + axis_angles = matrix_to_axis_angle(R, fast=fast) + self.assertClose(axis_angles, torch.zeros(3, dtype=torch.float64)) + (grad,) = torch.autograd.grad(axis_angles.sum(), R) + self.assertTrue(torch.isfinite(grad).all()) + + def test_axis_angle_batch_independence(self): + """Every matrix in a batch is converted independently of the others.""" + angles = torch.tensor( + [math.pi, math.pi, math.pi - 1e-9, math.pi - 1e-3, 3.0, 1.0, 0.5, 0.0], + dtype=torch.float64, + ) + axes = torch.nn.functional.normalize( + torch.randn(angles.shape[0], 3, dtype=torch.float64), dim=-1 + ) + R = axis_angle_to_matrix(axes * angles[:, None]) + for fast in [False, True]: + batched = matrix_to_axis_angle(R, fast=fast) + singly = torch.stack( + [matrix_to_axis_angle(r, fast=fast) for r in R.unbind()] + ) + self.assertClose(batched, singly) + + def test_axis_angle_shapes(self): + """Arbitrary leading dimensions, including none and zero.""" + R = random_rotations(24, dtype=torch.float64).reshape(2, 3, 4, 3, 3) + # A rotation by pi, so that the near-pi branch is used too. + axis = torch.tensor([0.0, 0.6, -0.8], dtype=torch.float64) + R[1, 2, 0] = 2 * torch.outer(axis, axis) - torch.eye(3, dtype=torch.float64) + for fast in [False, True]: + axis_angles = matrix_to_axis_angle(R, fast=fast) + self.assertEqual(axis_angles.shape, (2, 3, 4, 3)) + self.assertClose(axis_angle_to_matrix(axis_angles), R) + + self.assertEqual(matrix_to_axis_angle(R[0, 0, 0], fast=fast).shape, (3,)) + empty = torch.zeros(0, 3, 3, dtype=torch.float64) + self.assertEqual(matrix_to_axis_angle(empty, fast=fast).shape, (0, 3)) + + def test_axis_angle_near_pi_grad(self): + """Round tripping is the identity, so its derivative is too.""" + axis = torch.nn.functional.normalize( + torch.tensor([0.3, -0.5, 0.81], dtype=torch.float64), dim=-1 + ) + eye = torch.eye(3, dtype=torch.float64) + for angle in [0.5, 2.0, math.pi - 0.02, math.pi - 1e-4]: + data = axis * angle + for fast in [False, True]: + jacobian = torch.autograd.functional.jacobian( + lambda x, fast=fast: matrix_to_axis_angle( + axis_angle_to_matrix(x), fast=fast + ), + data, + ) + self.assertClose(jacobian, eye, atol=1e-8) + + # At exactly pi the axis, and so the derivative, is ambiguous, but the + # result must still be usable. + R = (2 * torch.outer(axis, axis) - eye).requires_grad_(True) + for fast in [False, True]: + (grad,) = torch.autograd.grad( + matrix_to_axis_angle(R, fast=fast).sum(), R, retain_graph=True + ) + self.assertTrue(torch.isfinite(grad).all()) + def test_quaternion_application(self): """Applying a quaternion is the same as applying the matrix.""" quaternions = random_quaternions(3, torch.float64)