Fix symeig3x3 for inputs that are not of order 1 (#2045)

Summary:
## Motivation

`symeig3x3` uses a fixed absolute epsilon (`torch.finfo(torch.float).eps`) in several places where the magnitude of the input actually matters:

- `p1.clamp(self._eps)`
- the `diag_soft_cond` test, which decides whether a matrix is treated as diagonal
- the `+= eps * sign(...)` regularizers in `_get_ev0` and `_get_ev1`

Since those thresholds are absolute, the routine only behaves as intended for matrices whose entries are of order 1. Once the entries drop below roughly `1e-3` the regularizers dominate the data and the output stops being an eigendecomposition at all. At the other end, `1e10` is already wrong and `1e30` returns NaN.

This is not an exotic corner. `estimate_pointcloud_normals` and `estimate_pointcloud_local_coord_frames` call `symeig3x3` by default (`use_symeig_workaround=True`) on per-point covariance matrices, and those scale with the square of the local neighbourhood size. The denser the point cloud, the smaller the covariances, and the worse the normals get, silently.

Points sampled on a sphere, `neighborhood_size=50`, median angle between the estimated normal and the true radial normal:

| points | radius | mean abs cov entry | main | this PR | `torch.linalg.eigh` |
|--------|--------|--------------------|------|---------|---------------------|
| 4 000 | 1.0 | 4.5e-03 | 1.60° | 1.61° | 1.61° |
| 20 000 | 1.0 | 9.1e-04 | 1.55° | 0.68° | 0.68° |
| 100 000 | 1.0 | 1.8e-04 | 31.86° | 0.31° | 0.31° |
| 4 000 | 0.1 | 4.5e-05 | 50.24° | 1.54° | 1.54° |

This looks like the cause of https://github.com/facebookresearch/pytorch3d/issues/1577, where the reporter narrowed bad normals down to `symeig3x3` and found that passing `use_symeig_workaround=False` fixed them.

## Fix

The decomposition is homogeneous: `eig(s * A) == s * eig(A)` and the eigenvectors do not change. So preconditioning each matrix to unit maximum magnitude and undoing the scaling on the eigenvalues afterwards is exact, and it puts the absolute thresholds back in the regime they were written for.

The scale is detached. For any fixed scale the expression is exactly the eigendecomposition of the input, so gradients are unaffected, and this avoids differentiating through `amax`.

Max relative reconstruction error `max|V diag(w) Vᵀ - A| / max|A|` over 2 000 random symmetric matrices:

| input scale | main | this PR |
|-------------|---------|---------|
| 1e30 | nan | 9.9e-07 |
| 1e10 | 1.00 | 1.0e-06 |
| 1e0 | 1.6e-06 | 1.0e-06 |
| 1e-10 | 1.74 | 9.7e-07 |
| 1e-30 | 1.73 | 9.7e-07 |

float64 inputs break on main below about `1e-8` and are accurate here down to `1e-60`. Exactly-zero, identity, diagonal and rank-1 inputs, and batches mixing magnitudes within one call, all behave.

## Tests

Added `test_is_eigen_small_magnitude`, which runs the existing eigen check on the same random matrices scaled by `1e-2`, `1e-4` and `1e-6` with a correspondingly scaled `atol`. On main it fails with a max relative difference of 101.6; with the fix it passes.

The 13 existing tests in `tests/test_symeig3x3.py` pass unchanged both before and after, on CPU and on GPU, `gradcheck` included.

Cost is one `amax` and one division: no measurable change on GPU, around 2% on CPU at 1e6 matrices.

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/2045

Reviewed By: MichaelRamamonjisoa

Differential Revision: D119970731

Pulled By: bottler

fbshipit-source-id: 35d85a342513c11faca4e6d8ee98052be90172a6

tests/test_symeig3x3.py

@@ -66,6 +66,51 @@ class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
             rtol=rtol,
         )
 
+    def test_is_eigen_scaled_gpu(self):
+        self._test_is_eigen_scaled(self._gpu)
+
+    def test_is_eigen_scaled_cpu(self):
+        self._test_is_eigen_scaled(self._cpu)
+
+    def _test_is_eigen_scaled(self, device):
+        """
+        The decomposition is homogeneous, so scaling the input must not affect
+        the relative accuracy of the result.
+        """
+        test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
+
+        for scale in (1e-6, 1e-4, 1e-2, 1e2, 1e4, 1e6):
+            scaled_input = test_input * scale
+
+            self._test_is_not_nan_or_inf(scaled_input)
+            self._test_is_eigen(scaled_input, atol=1e-04 * scale, rtol=1e-02)
+            self._test_eigenvectors_are_orthonormal(scaled_input)
+
+    def test_mixed_magnitudes_gpu(self):
+        self._test_mixed_magnitudes(self._gpu)
+
+    def test_mixed_magnitudes_cpu(self):
+        self._test_mixed_magnitudes(self._cpu)
+
+    def _test_mixed_magnitudes(self, device):
+        """
+        Each matrix in a batch must be conditioned on its own magnitude, so
+        matrices spanning many orders of magnitude can be batched together.
+        """
+        test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
+        scales = torch.logspace(-6, 6, self.TEST_BATCH_SIZE, device=device)[
+            :, None, None
+        ]
+        scaled_input = test_input * scales
+
+        eigenvalues, eigenvectors = symeig3x3(scaled_input, eigenvectors=True)
+        reconstruction = (
+            eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1)
+        )
+
+        self.assertClose(test_input, reconstruction / scales, atol=1e-04, rtol=1e-02)
+        self._test_eigenvectors_are_orthonormal(scaled_input)
+
     def test_eigenvectors_are_orthonormal_gpu(self):
         test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
 
@@ -192,21 +237,31 @@ class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
     def test_degenerate_eigenvalues_cpu(self):
         self._test_degenerate_eigenvalues(self._cpu)
 
-    def _test_degenerate_eigenvalues(self, device):
+    def test_degenerate_eigenvalues_scaled_gpu(self):
+        for scale in (1e-6, 1e6):
+            self._test_degenerate_eigenvalues(self._gpu, scale=scale)
+
+    def test_degenerate_eigenvalues_scaled_cpu(self):
+        for scale in (1e-6, 1e6):
+            self._test_degenerate_eigenvalues(self._cpu, scale=scale)
+
+    def _test_degenerate_eigenvalues(self, device, scale: float = 1.0):
         """
-        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity
+        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity.
+        These are the cases held together by the regularizers inside the
+        decomposition, which only bite at the right input magnitude.
         """
         # Error tolerances for degenerate values are increased as things might become
         #  numerically unstable
-        deg_atol = 1e-3
+        deg_atol = 1e-3 * scale
         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
+        test_eigenvals = (
+            torch.randn((self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device) * scale
         )
         self._test_eigenvalues_and_eigenvectors(
             test_eigenvecs, test_eigenvals, atol=deg_atol, rtol=deg_rtol