Added mean axis for dirac on hypersphere

Author: FlorianPfaff

pyrecest/distributions/hypersphere_subset/abstract_hypersphere_subset_dirac_distribution.py

@@ -1,6 +1,6 @@
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import log, outer, sum, zeros
+from pyrecest.backend import log, outer, sum, zeros, reshape, linalg
 
 from ..abstract_dirac_distribution import AbstractDiracDistribution
 from .abstract_hypersphere_subset_distribution import (
@@ -31,3 +31,36 @@ def entropy(self):
 
     def integrate(self, integration_boundaries=None):
         raise NotImplementedError()
+    
+    def mean_axis(self):
+        """
+        Returns the principal axis of the Dirac mixture on the hypersphere.
+        Because ±v represent the same axis, the sign of the returned vector
+        is arbitrary.
+        """
+        # Column vector of weights, shape (n, 1)
+        w_col = reshape(self.w, (-1, 1))  # or self.w[:, None]
+
+        # Weighted second-moment matrix: S = Σ w_i d_i d_i^T
+        # d has shape (n, D), so (d * w_col) is (n, D), then transpose @ d -> (D, D)
+        S = (self.d * w_col).T @ self.d
+
+        # Normalize in case weights don't sum to 1
+        S = S / sum(self.w)
+
+        # Enforce symmetry (numerical safety)
+        S = 0.5 * (S + S.T)
+
+        # Eigen-decomposition of symmetric S
+        D, V = linalg.eig(S)
+
+        # Index of largest eigenvalue
+        # If you don't have argmax in the backend, use argsort instead
+        idx = D.argmax()  # or idx = argsort(D)[-1]
+
+        axis = V[:, idx]
+
+        # Normalize to unit length (should already be, but just in case)
+        axis = axis / linalg.norm(axis)
+
+        return axis

pyrecest/distributions/hypersphere_subset/hyperhemispherical_dirac_distribution.py

@@ -9,4 +9,8 @@
 class HyperhemisphericalDiracDistribution(
     AbstractHypersphereSubsetDiracDistribution, AbstractHyperhemisphericalDistribution
 ):
-    pass
+    def mean_axis(self):
+        axis = super().mean_axis()
+        if axis[-1] < 0:
+            axis = -axis
+        return axis

pyrecest/tests/distributions/test_hyperspherical_dirac_distribution.py

@@ -8,7 +8,7 @@
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import array, linalg, mod, ones, random, sqrt, sum
+from pyrecest.backend import array, linalg, mod, ones, random, sqrt, sum, allclose
 from pyrecest.distributions import VonMisesFisherDistribution
 from pyrecest.distributions.hypersphere_subset.hyperspherical_dirac_distribution import (
     HypersphericalDiracDistribution,
@@ -79,6 +79,25 @@ def test_from_distribution(self):
             dirac_dist.mean_direction(), vmf.mean_direction(), decimal=2
         )
 
+    def test_mean_axis_symmetric_two_point_distribution(self):
+        # Two antipodal points on S^2: ±e_x
+        d = array([
+            [1.0, 0.0, 0.0],
+            [-1.0, 0.0, 0.0],
+        ])
+        w = array([0.5, 0.5])
+
+        dist = HypersphericalDiracDistribution(d, w)
+
+        axis = dist.mean_axis()
+
+        # 1) axis should be unit length
+        assert allclose(linalg.norm(axis), 1.0, atol=1e-7)
+
+        # 2) axis should be parallel to (1, 0, 0), i.e. |dot(axis, e_x)| ≈ 1
+        v = array([1.0, 0.0, 0.0])
+        dot = float(axis @ v)
+        assert abs(dot) > 1.0 - 1e-6
 
 if __name__ == "__main__":
     unittest.main()