Preserve SO(3) distribution subclasses in factories and operations

src/pyrecest/distributions/so3_bingham_distribution.py

@@ -223,7 +223,7 @@ def multiply(self, B2):
         if not isinstance(B2, SO3BinghamDistribution):
             raise ValueError("B2 must be an SO3BinghamDistribution.")
         product = self.distFullSphere.multiply(B2.distFullSphere)
-        return SO3BinghamDistribution(product.Z, product.M)
+        return type(self)(product.Z, product.M)
 
     def compose(self, B2):
         """Approximate the distribution of the composed rotation ``self * other``."""
@@ -243,7 +243,7 @@ def compose(self, B2):
             ] * right_matrix @ first_moment @ transpose(right_matrix)
 
         composed_moment = 0.5 * (composed_moment + transpose(composed_moment))
-        return SO3BinghamDistribution.from_bingham_distribution(
+        return type(self).from_bingham_distribution(
             BinghamDistribution.fit_to_moment(composed_moment)
         )
 

src/pyrecest/distributions/so3_tangent_gaussian_distribution.py

@@ -228,7 +228,7 @@ def is_valid(self, tolerance=1e-6):
 
         return _to_python_bool(all(linalg.eigvalsh(self.C) > 0.0))
 
-    @staticmethod
-    def from_covariance_diagonal(mu, covariance_diagonal):
+    @classmethod
+    def from_covariance_diagonal(cls, mu, covariance_diagonal):
         """Create a tangent Gaussian from a diagonal covariance vector."""
-        return SO3TangentGaussianDistribution(mu, diag(covariance_diagonal))
+        return cls(mu, diag(covariance_diagonal))

tests/distributions/test_so3_subclass_preservation.py

@@ -0,0 +1,52 @@
+# pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import array
+from pyrecest.distributions import (
+    SO3BinghamDistribution,
+    SO3TangentGaussianDistribution,
+)
+
+
+class _SO3BinghamSubclass(SO3BinghamDistribution):
+    pass
+
+
+class _SO3TangentGaussianSubclass(SO3TangentGaussianDistribution):
+    pass
+
+
+def _identity_quaternion():
+    return array([0.0, 0.0, 0.0, 1.0])
+
+
+def test_so3_bingham_multiply_preserves_left_subclass():
+    left = _SO3BinghamSubclass.from_mode_and_concentration(
+        _identity_quaternion(), 2.0
+    )
+    right = SO3BinghamDistribution.from_mode_and_concentration(
+        _identity_quaternion(), 3.0
+    )
+
+    product = left.multiply(right)
+
+    assert isinstance(product, _SO3BinghamSubclass)
+
+
+def test_so3_bingham_compose_preserves_left_subclass():
+    left = _SO3BinghamSubclass.from_mode_and_concentration(
+        _identity_quaternion(), 2.0
+    )
+    right = SO3BinghamDistribution.from_mode_and_concentration(
+        _identity_quaternion(), 3.0
+    )
+
+    composed = left.compose(right)
+
+    assert isinstance(composed, _SO3BinghamSubclass)
+
+
+def test_so3_tangent_gaussian_diagonal_factory_preserves_requested_subclass():
+    distribution = _SO3TangentGaussianSubclass.from_covariance_diagonal(
+        _identity_quaternion(), array([0.1, 0.2, 0.3])
+    )
+
+    assert isinstance(distribution, _SO3TangentGaussianSubclass)