From 3e81b7b742b7275b788b540b979df03fd10118c4 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sun, 23 Aug 2026 04:43:55 +0800
Subject: [PATCH 1/3] Preserve SO3 Bingham subclasses in operations

---
 src/pyrecest/distributions/so3_bingham_distribution.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/pyrecest/distributions/so3_bingham_distribution.py b/src/pyrecest/distributions/so3_bingham_distribution.py
index fc4179c4a1..1ced1bcbf8 100644
--- a/src/pyrecest/distributions/so3_bingham_distribution.py
+++ b/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)
         )
 

From cd6cd98e4224da22c64e5fa3b68c81c16afda8c5 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sun, 23 Aug 2026 04:44:12 +0800
Subject: [PATCH 2/3] Preserve SO3 tangent Gaussian subclasses in factory

---
 .../distributions/so3_tangent_gaussian_distribution.py      | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/pyrecest/distributions/so3_tangent_gaussian_distribution.py b/src/pyrecest/distributions/so3_tangent_gaussian_distribution.py
index b685f8d587..db8b6f95f3 100644
--- a/src/pyrecest/distributions/so3_tangent_gaussian_distribution.py
+++ b/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))

From 7adc8484a1560bf016be0002edfa97c212a20785 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sun, 23 Aug 2026 04:44:24 +0800
Subject: [PATCH 3/3] Add SO3 subclass-preservation regressions

---
 .../test_so3_subclass_preservation.py         | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 tests/distributions/test_so3_subclass_preservation.py

diff --git a/tests/distributions/test_so3_subclass_preservation.py b/tests/distributions/test_so3_subclass_preservation.py
new file mode 100644
index 0000000000..f7d9e246d5
--- /dev/null
+++ b/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)