Add estimate_gaussian_rotation_matrix_from_samples

doc/source/api.rst

@@ -534,6 +534,7 @@ Operations and Utility Functions
    :toctree: _apidoc/
    :template: function.rst
 
+   ~pytransform3d.uncertainty.estimate_gaussian_rotation_matrix_from_samples
    ~pytransform3d.uncertainty.estimate_gaussian_transform_from_samples
    ~pytransform3d.uncertainty.invert_uncertain_transform
    ~pytransform3d.uncertainty.concat_globally_uncertain_transforms
@@ -545,7 +546,7 @@ Operations and Utility Functions
 
 
 :mod:`pytransform3d.coordinates`
-=================================
+================================
 
 .. automodule:: pytransform3d.coordinates
     :no-members:

pytransform3d/test/test_uncertainty.py

@@ -6,6 +6,17 @@
 import pytest
 
 
+def test_estimate_gaussian_rotation_matrix_from_samples():
+    mean = pr.matrix_from_axis_angle([0.3, 0.2, 0.5, 0.25 * np.pi])
+    cov = np.diag([1.3, 1.5, 2.9])
+    rng = np.random.default_rng(34)
+    samples = np.array([pr.random_matrix(rng, mean, cov) for _ in range(1000)])
+    mean_est, cov_est = pu.estimate_gaussian_rotation_matrix_from_samples(
+        samples)
+    assert_array_almost_equal(mean, mean_est, decimal=1)
+    assert_array_almost_equal(cov, cov_est, decimal=0)
+
+
 def test_same_fuse_poses():
     mean1 = np.array([
         [0.8573, -0.2854, 0.4285, 3.5368],

pytransform3d/uncertainty.py

@@ -1,5 +1,7 @@
 """Operations related to uncertain transformations."""
 import numpy as np
+from .batch_rotations import (matrices_from_compact_axis_angles,
+                              axis_angles_from_matrices)
 from .transformations import (
     invert_transform, transform_from, concat, adjoint_from_transform,
     left_jacobian_SE3_inv, transform_from_exponential_coordinates,
@@ -10,6 +12,51 @@
 from ._geometry import unit_sphere_surface_grid
 
 
+def estimate_gaussian_rotation_matrix_from_samples(samples):
+    """Estimate Gaussian distribution over rotations from samples.
+
+    Uses iterative approximation of mean described by Eade [1]_ and computes
+    covariance in exponential coordinate space (using an unbiased estimator).
+
+    Parameters
+    ----------
+    samples : array-like, shape (n_samples, 3, 3)
+        Sampled rotations represented by rotation matrices.
+
+    Returns
+    -------
+    mean : array, shape (3, 3)
+        Mean as rotation matrix.
+
+    cov : array, shape (3, 3)
+        Covariance of distribution in exponential coordinate space.
+
+    References
+    ----------
+    .. [1] Eade, E. (2017). Lie Groups for 2D and 3D Transformations.
+       https://ethaneade.com/lie.pdf
+    """
+    assert len(samples) > 0
+    samples = np.asarray(samples)
+    mean = samples[0]
+    for _ in range(20):
+        mean_inv = mean.T
+        mean_diffs_axis_angle = axis_angles_from_matrices(
+            # Uses implementation detail concat_many_to_one that is not part of
+            # the public interface. It is supposed to support only homgeneous
+            # transformation matrices, but it works for rotation matrices, too!
+            concat_many_to_one(samples, mean_inv)
+        )
+        mean_diffs = (mean_diffs_axis_angle[:, :3]
+                      * mean_diffs_axis_angle[:, 3, np.newaxis])
+        avg_mean_diff = np.mean(mean_diffs, axis=0)
+        print(f"{np.linalg.norm(avg_mean_diff)=}")
+        mean = np.dot(matrices_from_compact_axis_angles(avg_mean_diff), mean)
+
+    cov = np.cov(mean_diffs, rowvar=False, bias=True)
+    return mean, cov
+
+
 def estimate_gaussian_transform_from_samples(samples):
     """Estimate Gaussian distribution over transformations from samples.
 
@@ -35,6 +82,7 @@ def estimate_gaussian_transform_from_samples(samples):
        https://ethaneade.com/lie.pdf
     """
     assert len(samples) > 0
+    samples = np.asarray(samples)
     mean = samples[0]
     for _ in range(20):
         mean_inv = invert_transform(mean)

pytransform3d/uncertainty.pyi

@@ -4,6 +4,10 @@ from typing import Tuple, Union
 from mpl_toolkits.mplot3d import Axes3D
 
 
+def estimate_gaussian_rotation_matrix_from_samples(
+        samples: npt.ArrayLike) -> Tuple[np.ndarray, np.ndarray]: ...
+
+
 def estimate_gaussian_transform_from_samples(
         samples: npt.ArrayLike) -> Tuple[np.ndarray, np.ndarray]: ...