fix(axes): make principal axes reproducible across LAPACK builds

Author: harryswift01

CodeEntropy/levels/axes.py

@@ -837,13 +837,7 @@ def get_vanilla_axes(self, molecule):
         """
         make_whole(molecule.atoms)
         moment_of_inertia_tensor = molecule.moment_of_inertia()
-        principal_axes = self.get_principal_axes(molecule)
-
-        eigenvalues, _ = np.linalg.eigh(moment_of_inertia_tensor)
-        order = np.argsort(np.abs(eigenvalues))[::-1]
-        moment_of_inertia = eigenvalues[order]
-
-        return principal_axes, moment_of_inertia
+        return self.get_reproducible_axes(moment_of_inertia_tensor)
 
     def get_custom_axes(
         self,
@@ -1082,67 +1076,91 @@ def get_custom_principal_axes(
     ) -> tuple[np.ndarray, np.ndarray]:
         """Compute principal axes and moments from a custom MOI tensor.
 
-        Principal axes and centre of axes from the ordered eigenvalues and
-        eigenvectors of a moment of inertia tensor. This function allows for a
-        custom moment of inertia tensor to be used, which isn't possible with the
-        built-in MDAnalysis principal_axes() function.
-
-        Original behaviour preserved:
-
-        - Eigenvalues are sorted by descending absolute magnitude.
-        - Eigenvectors are transposed so axes are returned as rows.
-        - Z axis is flipped to enforce the same handedness convention as the
-          original implementation.
+        Principal axes and moments from the ordered eigenvalues and eigenvectors of
+        a moment of inertia tensor. This function allows for a custom moment of
+        inertia tensor to be used. The axes are made reproducible by
+        :meth:`get_reproducible_axes`.
 
         Args:
             moment_of_inertia_tensor: (3, 3) custom inertia tensor.
 
         Returns:
             Tuple[np.ndarray, np.ndarray]:
                 - principal_axes: (3, 3) principal axes (rows).
-                - moment_of_inertia: (3,) principal moments.
+                - moment_of_inertia: (3,) principal moments sorted descending by
+                  absolute value.
         """
-        eigenvalues, eigenvectors = np.linalg.eigh(moment_of_inertia_tensor)
-        order = np.abs(eigenvalues).argsort()[::-1]  # descending order
-        transposed = np.transpose(eigenvectors)  # columns -> rows
-        moment_of_inertia = eigenvalues[order]
-        principal_axes = transposed[order]
-
-        # point z axis in correct direction, as per original code
-        cross_xy = np.cross(principal_axes[0], principal_axes[1])
-        dot_z = float(np.dot(cross_xy, principal_axes[2]))
-        if dot_z < 0:
-            principal_axes[2] *= -1
-
-        return principal_axes, moment_of_inertia
+        return self.get_reproducible_axes(moment_of_inertia_tensor)
 
     def get_principal_axes(self, group, wrap: bool = False) -> np.ndarray:
         """Compute the principal axes of an atom group from its moment of inertia.
 
-        Uses ``np.linalg.eigh`` rather than ``np.linalg.eig``. The moment of inertia
-        tensor is symmetric, so ``eigh`` returns orthonormal axes, whereas ``eig``
-        can return skewed axes for (near-)degenerate moments.
-
-        - Axes are sorted by descending (signed) eigenvalue and returned as rows.
-        - If the axes are left-handed, all three are flipped.
-
         Args:
             group: Atom group (or anything with ``.atoms``).
             wrap: Whether to wrap atoms into the primary unit cell before computing
                 the moment of inertia.
 
         Returns:
-            np.ndarray: (3, 3) principal axes (rows).
+            np.ndarray: (3, 3) principal axes (rows), as given by
+            :meth:`get_reproducible_axes`.
         """
-        atomgroup = group.atoms
-        e_val, e_vec = np.linalg.eigh(atomgroup.moment_of_inertia(wrap=wrap))
+        tensor = group.atoms.moment_of_inertia(wrap=wrap)
+        principal_axes, _ = self.get_reproducible_axes(tensor)
+        return principal_axes
 
-        e_vec = e_vec[:, np.argsort(e_val)[::-1]].T
+    def get_reproducible_axes(
+        self, tensor: np.ndarray, rtol: float = 1e-5
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Diagonalise a symmetric tensor into reproducible principal axes.
+
+        ``np.linalg.eigh`` returns orthonormal axes, but the sign of each axis, and
+        the choice of axes within a set of (near-)equal moments, can differ between
+        LAPACK builds. These axes are used to average covariances over molecules and
+        frames, so that choice changes the entropy. This method fixes it:
+
+        - Moments are sorted descending by absolute value.
+        - Axes whose moments are equal to within ``rtol`` (relative to the largest
+          moment) are re-chosen from the eigenvectors of a fixed reference tensor
+          projected onto the shared subspace. Axes with distinct moments are left as
+          they are.
+        - The first two axes are signed to point along the lab-frame direction
+          (1, 2, 3), and the third is their cross product, so the axes are
+          right-handed.
 
-        if np.dot(np.cross(e_vec[0], e_vec[1]), e_vec[2]) < 0:
-            e_vec *= -1
+        The result is reproducible but is a function of the lab-frame orientation,
+        so it is not rotation covariant.
+
+        Args:
+            tensor: (3, 3) symmetric tensor, e.g. a moment of inertia tensor.
+            rtol: Relative tolerance for treating two moments as equal.
 
-        return e_vec
+        Returns:
+            Tuple[np.ndarray, np.ndarray]:
+                - principal_axes: (3, 3) principal axes (rows).
+                - moments: (3,) principal moments sorted descending by absolute
+                  value.
+        """
+        moments, vectors = np.linalg.eigh(tensor)
+        order = np.argsort(np.abs(moments))[::-1]
+        moments, vectors = moments[order], vectors[:, order]
+
+        reference_direction = np.array([1.0, 2.0, 3.0])
+        gaps = np.abs(np.diff(np.abs(moments)))
+        group_starts = np.flatnonzero(gaps > rtol * np.abs(moments[0])) + 1
+        for group in np.split(np.arange(3), group_starts):
+            if len(group) > 1:
+                shared = vectors[:, group]
+                reference = shared.T @ np.diag(reference_direction) @ shared
+                _, rotation = np.linalg.eigh(reference)
+                vectors[:, group] = shared @ rotation
+
+        for i in (0, 1):
+            if vectors[:, i] @ reference_direction < 0:
+                vectors[:, i] *= -1
+
+        third_axis = np.cross(vectors[:, 0], vectors[:, 1])
+        principal_axes = np.array([vectors[:, 0], vectors[:, 1], third_axis])
+        return principal_axes, moments
 
     def get_UA_masses(self, molecule) -> list[float]:
         """Return united-atom (UA) masses for a molecule.

tests/unit/CodeEntropy/levels/test_axes.py

@@ -293,7 +293,7 @@ def test_get_flipped_axes_flips_negative_dot():
     assert np.allclose(flipped[0], np.array([1.0, 0.0, 0.0]))
 
 
-def test_get_custom_principal_axes_flips_z_when_left_handed():
+def test_get_custom_principal_axes_returns_axes_and_moments():
     ax = AxesCalculator()
     moi = np.eye(3)
     axes, vals = ax.get_custom_principal_axes(moi)
@@ -368,19 +368,97 @@ def test_get_principal_axes_orthonormal_and_right_handed_for_near_degenerate():
     assert np.isclose(abs(np.dot(axes[0], q[:, 2])), 1.0)
 
 
-def test_get_principal_axes_flips_all_axes_when_left_handed(monkeypatch):
+def _scrambled_eigh(seed):
+    """Return an ``eigh`` that gives valid but arbitrary signs and degenerate bases."""
+    rng = np.random.default_rng(seed)
+    real_eigh = np.linalg.eigh
+
+    def scrambled(matrix):
+        values, vectors = real_eigh(matrix)
+        n = len(values)
+        for i in range(n):
+            for j in range(i + 1, n):
+                if np.isclose(values[i], values[j], rtol=1e-5):
+                    angle = rng.uniform(0, 2 * np.pi)
+                    c, s = np.cos(angle), np.sin(angle)
+                    vi, vj = vectors[:, i].copy(), vectors[:, j].copy()
+                    vectors[:, i] = c * vi + s * vj
+                    vectors[:, j] = -s * vi + c * vj
+        return values, vectors * rng.choice([-1.0, 1.0], size=n)
+
+    return scrambled
+
+
+def _example_tensors():
+    rng = np.random.default_rng(1)
+    q, _ = np.linalg.qr(rng.normal(size=(3, 3)))
+    return {
+        "asymmetric_top": q @ np.diag([5.0, 3.0, 1.0]) @ q.T,
+        "linear": q @ np.diag([4.0, 4.0, 0.0]) @ q.T,
+        "symmetric_top": q @ np.diag([5.0, 2.0, 2.0]) @ q.T,
+        "spherical_top": 3.0 * np.eye(3),
+    }
+
+
+@pytest.mark.parametrize("name", list(_example_tensors()))
+def test_get_reproducible_axes_orthonormal_right_handed_and_diagonalising(name):
     ax = AxesCalculator()
-    group = MagicMock()
-    group.atoms.moment_of_inertia.return_value = np.eye(3)
-    monkeypatch.setattr(
-        np.linalg, "eigh", lambda _m: (np.array([1.0, 2.0, 3.0]), np.eye(3))
-    )
+    tensor = _example_tensors()[name]
 
-    axes = ax.get_principal_axes(group)
+    axes, moments = ax.get_reproducible_axes(tensor)
+
+    assert np.allclose(axes @ axes.T, np.eye(3), atol=1e-12)
+    assert np.isclose(np.linalg.det(axes), 1.0)
+    assert np.allclose(axes @ tensor @ axes.T, np.diag(moments), atol=1e-12)
+    assert np.all(np.diff(np.abs(moments)) <= 1e-12)
+
+
+@pytest.mark.parametrize("name", list(_example_tensors()))
+def test_get_reproducible_axes_independent_of_eigensolver_choices(monkeypatch, name):
+    ax = AxesCalculator()
+    tensor = _example_tensors()[name]
+    expected_axes, expected_moments = ax.get_reproducible_axes(tensor)
+
+    for seed in range(20):
+        monkeypatch.setattr(np.linalg, "eigh", _scrambled_eigh(seed))
+        axes, moments = ax.get_reproducible_axes(tensor)
+        assert np.allclose(axes, expected_axes, atol=1e-9)
+        assert np.allclose(moments, expected_moments)
+
+
+@pytest.mark.parametrize("name", list(_example_tensors()))
+def test_get_reproducible_axes_stable_under_rounding_level_noise(name):
+    ax = AxesCalculator()
+    tensor = _example_tensors()[name]
+    expected_axes, _ = ax.get_reproducible_axes(tensor)
+    rng = np.random.default_rng(2)
+
+    for _ in range(50):
+        noise = rng.normal(scale=1e-7, size=(3, 3))
+        axes, _ = ax.get_reproducible_axes(tensor + (noise + noise.T) / 2)
+        assert np.allclose(axes, expected_axes, atol=1e-4)
+
+
+def test_get_reproducible_axes_keeps_unique_axis_of_linear_tensor_exact():
+    ax = AxesCalculator()
+    direction = np.array([1.0, -2.0, 0.5])
+    direction /= np.linalg.norm(direction)
+    tensor = 4.0 * (np.eye(3) - np.outer(direction, direction))
+
+    axes, moments = ax.get_reproducible_axes(tensor)
+
+    assert np.isclose(moments[2], 0.0, atol=1e-12)
+    assert np.isclose(abs(axes[2] @ direction), 1.0, atol=1e-12)
+
+
+def test_get_reproducible_axes_signs_first_two_axes_along_reference_direction():
+    ax = AxesCalculator()
+    tensor = np.diag([5.0, 3.0, 1.0])
+
+    axes, _ = ax.get_reproducible_axes(tensor)
 
-    # descending order gives rows e3, e2, e1 (left-handed), so all are negated
-    expected = -np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=float)
-    assert np.allclose(axes, expected)
+    assert axes[0] @ np.array([1.0, 2.0, 3.0]) > 0
+    assert axes[1] @ np.array([1.0, 2.0, 3.0]) > 0
 
 
 def test_get_principal_axes_passes_wrap_to_moment_of_inertia():
@@ -653,7 +731,7 @@ def test_get_custom_axes_raises_when_normalization_degenerate():
         ac.get_custom_axes(a=a, b_list=b_list, c=c, dimensions=dims)
 
 
-def test_get_custom_principal_axes_flips_z_for_handedness():
+def test_get_custom_principal_axes_is_right_handed():
     ac = AxesCalculator()
 
     moi = np.diag([3.0, 2.0, 1.0])