diff --git a/src/pyrecest/distributions/so3_bingham_distribution.py b/src/pyrecest/distributions/so3_bingham_distribution.py index fc4179c4a..1ced1bcbf 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) ) diff --git a/src/pyrecest/distributions/so3_tangent_gaussian_distribution.py b/src/pyrecest/distributions/so3_tangent_gaussian_distribution.py index b685f8d58..db8b6f95f 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)) diff --git a/tests/distributions/test_so3_subclass_preservation.py b/tests/distributions/test_so3_subclass_preservation.py new file mode 100644 index 000000000..f7d9e246d --- /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)