Repository: dfki-ric/pytransform3d

Issue:
Implement estimate_gaussian_matrix_from_samples
... to estimate distribution of rotation matrices, similarly to estimate_gaussian_transform_from_samples according to https://ethaneade.com/lie.pdf, section 8.5

Discussion:

Patch:
diff --git a/doc/source/api.rst b/doc/source/api.rst
index dafc5cf8..1007b38e 100644
--- a/doc/source/api.rst
+++ b/doc/source/api.rst
@@ -28,6 +28,8 @@ Rotation Matrix
    ~matrix_requires_renormalization
    ~norm_matrix
 
+   ~random_matrix
+
    ~plot_basis
 
    ~assert_rotation_matrix
@@ -595,7 +597,9 @@ Dual Quaternions
    :toctree: _apidoc/
    :template: function.rst
 
+   ~estimate_gaussian_rotation_matrix_from_samples
    ~estimate_gaussian_transform_from_samples
+   ~frechet_mean
    ~invert_uncertain_transform
    ~concat_globally_uncertain_transforms
    ~concat_locally_uncertain_transforms
@@ -606,7 +610,7 @@ Dual Quaternions
 
 
 :mod:`pytransform3d.coordinates`
-=================================
+================================
 
 .. automodule:: pytransform3d.coordinates
     :no-members:
diff --git a/pytransform3d/batch_rotations.py b/pytransform3d/batch_rotations.py
index fd3167c1..1fb6e64c 100644
--- a/pytransform3d/batch_rotations.py
+++ b/pytransform3d/batch_rotations.py
@@ -366,19 +366,24 @@ def axis_angles_from_matrices(Rs, traces=None, out=None):
     # We can still use the original formula to reconstruct the signs of
     # the rotation axis correctly.
     angle_close_to_pi = np.abs(angles - np.pi) < 1e-4
-    angle_not_zero = np.abs(angles) != 0.0
+    angle_zero = angles == 0.0
+    angle_not_zero = np.logical_not(angle_zero)
 
     Rs_diag = np.einsum("nii->ni", Rs.reshape(-1, 3, 3))
     if instances_shape:
         Rs_diag = Rs_diag.reshape(*(instances_shape + (3,)))
     else:
         Rs_diag = Rs_diag[0]
+
     out[angle_close_to_pi, :3] = (
         np.sqrt(0.5 * (Rs_diag[angle_close_to_pi] + 1.0))
         * np.sign(out[angle_close_to_pi, :3]))
     out[angle_not_zero, :3] /= np.linalg.norm(
         out[angle_not_zero, :3], axis=-1)[..., np.newaxis]
 
+    out[angle_zero, 0] = 1.0
+    out[angle_zero, 1:3] = 0.0
+
     out[..., 3] = angles
 
     return out
diff --git a/pytransform3d/camera.py b/pytransform3d/camera.py
index 6bb63b3b..cb1c3abd 100644
--- a/pytransform3d/camera.py
+++ b/pytransform3d/camera.py
@@ -3,9 +3,7 @@
 See :doc:`user_guide/camera` for more information.
 """
 import numpy as np
-from .plot_utils import Camera
-from .transformations import (invert_transform, transform, check_transform,
-                              vectors_to_points)
+from .transformations import invert_transform, transform, check_transform
 
 
 def make_world_grid(n_lines=11, n_points_per_line=51, xlim=(-0.5, 0.5),
@@ -254,6 +252,8 @@ def plot_camera(ax=None, M=None, cam2world=None, virtual_image_distance=1.0,
     ValueError
         If input is not valid
     """
+    from .plot_utils import Camera
+
     if ax is None:
         from .plot_utils import make_3d_axis
         ax = make_3d_axis(ax_s)
diff --git a/pytransform3d/plot_utils/__init__.py b/pytransform3d/plot_utils/__init__.py
index c6bfdb85..63cc6c4c 100644
--- a/pytransform3d/plot_utils/__init__.py
+++ b/pytransform3d/plot_utils/__init__.py
@@ -16,6 +16,9 @@ try:
         "plot_mesh", "plot_ellipsoid", "plot_capsule", "plot_cone",
         "plot_vector", "plot_length_variable"
     ]
-except ImportError:
-    warnings.warn("Matplotlib is not installed, visualization is not "
-                  "available")
+except ImportError as e:
+    if e.name == "matplotlib":
+        warnings.warn("Matplotlib is not installed, visualization is not "
+                      "available")
+    else:
+        raise e
diff --git a/pytransform3d/plot_utils/_artists.py b/pytransform3d/plot_utils/_artists.py
index 5625be8c..06add54f 100644
--- a/pytransform3d/plot_utils/_artists.py
+++ b/pytransform3d/plot_utils/_artists.py
@@ -332,7 +332,6 @@ class Camera(artist.Artist):
     kwargs : dict, optional (default: {})
         Additional arguments for the plotting functions, e.g. alpha.
     """
-
     def __init__(
         self,
         M,
diff --git a/pytransform3d/rotations/__init__.py b/pytransform3d/rotations/__init__.py
index 60dd3087..32d5570f 100644
--- a/pytransform3d/rotations/__init__.py
+++ b/pytransform3d/rotations/__init__.py
@@ -18,7 +18,8 @@ from ._random import (
     random_vector,
     random_axis_angle,
     random_compact_axis_angle,
-    random_quaternion)
+    random_quaternion,
+    random_matrix)
 from ._conversions import (
     quaternion_from_axis_angle, quaternion_from_compact_axis_angle,
     quaternion_from_matrix, quaternion_wxyz_from_xyzw,
@@ -139,6 +140,7 @@ __all__ = [
     "random_axis_angle",
     "random_compact_axis_angle",
     "random_quaternion",
+    "random_matrix",
     "check_skew_symmetric_matrix",
     "check_rot_log",
     "check_matrix",
diff --git a/pytransform3d/rotations/_random.py b/pytransform3d/rotations/_random.py
index fd44cc24..b0507e8f 100644
--- a/pytransform3d/rotations/_random.py
+++ b/pytransform3d/rotations/_random.py
@@ -1,6 +1,7 @@
 import numpy as np
 
-from ._utils import norm_vector
+from ._utils import norm_vector, check_matrix, norm_matrix
+from ._conversions import matrix_from_compact_axis_angle
 
 
 def random_vector(rng=np.random.default_rng(0), n=3):
@@ -84,3 +85,38 @@ def random_quaternion(rng=np.random.default_rng(0)):
         Unit quaternion to represent rotation: (w, x, y, z)
     """
     return norm_vector(rng.standard_normal(size=4))
+
+
+def random_matrix(rng=np.random.default_rng(0), mean=np.eye(3), cov=np.eye(3)):
+    r"""Generate random rotation matrix.
+
+    Generate :math:`\Delta \boldsymbol{R}_{B_{i+1}{B_i}}
+    \boldsymbol{R}_{{B_i}A}`, with :math:`\Delta \boldsymbol{R}_{B_{i+1}{B_i}}
+    = Exp(\hat{\omega} \theta)` and :math:`\hat{\omega}\theta \sim
+    \mathcal{N}(\boldsymbol{0}_3, \boldsymbol{\Sigma}_{3 \times 3})`.
+    The mean :math:`\boldsymbol{R}_{{B_i}A}` and the covariance
+    :math:`\boldsymbol{\Sigma}_{3 \times 3}` are parameters of the function.
+
+    Note that uncertainty is defined in the global frame B, not in the
+    body frame A.
+
+    Parameters
+    ----------
+    rng : np.random.Generator, optional (default: random seed 0)
+        Random number generator.
+
+    mean : array-like, shape (3, 3), optional (default: I)
+        Mean rotation matrix.
+
+    cov : array-like, shape (3, 3), optional (default: I)
+        Covariance of noise in exponential coordinate space.
+
+    Returns
+    -------
+    R : array, shape (3, 3)
+        Rotation matrix
+    """
+    mean = check_matrix(mean)
+    a = rng.multivariate_normal(mean=np.zeros(3), cov=cov)
+    delta = matrix_from_compact_axis_angle(a)
+    return norm_matrix(np.dot(delta, mean))
diff --git a/pytransform3d/rotations/_random.pyi b/pytransform3d/rotations/_random.pyi
index 59df9af2..4aae3404 100644
--- a/pytransform3d/rotations/_random.pyi
+++ b/pytransform3d/rotations/_random.pyi
@@ -12,3 +12,7 @@ def random_compact_axis_angle(rng: np.random.Generator = ...) -> np.ndarray: ...
 
 
 def random_quaternion(rng: np.random.Generator = ...) -> np.ndarray: ...
+
+
+def random_matrix(rng: np.random.Generator = ..., mean: npt.ArrayLike = ...,
+                  cov: npt.ArrayLike = ...) -> np.ndarray: ...
diff --git a/pytransform3d/transformations/_random.py b/pytransform3d/transformations/_random.py
index e4ce503c..54784bc3 100644
--- a/pytransform3d/transformations/_random.py
+++ b/pytransform3d/transformations/_random.py
@@ -32,7 +32,7 @@ def random_transform(
 
     Returns
     -------
-    A2B : array-like, shape (4, 4)
+    A2B : array, shape (4, 4)
         Random transform from frame A to frame B
     """
     mean = check_transform(mean)
diff --git a/pytransform3d/uncertainty.py b/pytransform3d/uncertainty.py
index 09e38362..448879c9 100644
--- a/pytransform3d/uncertainty.py
+++ b/pytransform3d/uncertainty.py
@@ -3,36 +3,74 @@
 See :doc:`user_guide/uncertainty` for more information.
 """
 import numpy as np
-from .transformations import (
-    invert_transform, transform_from, concat, adjoint_from_transform,
-    left_jacobian_SE3_inv, transform_from_exponential_coordinates,
-    exponential_coordinates_from_transform)
-from .trajectories import (exponential_coordinates_from_transforms,
-                           transforms_from_exponential_coordinates,
-                           concat_many_to_one)
+
 from ._geometry import unit_sphere_surface_grid
+from .batch_rotations import (axis_angles_from_matrices,
+                              matrices_from_compact_axis_angles)
+from .trajectories import (concat_many_to_one,
+                           exponential_coordinates_from_transforms,
+                           transforms_from_exponential_coordinates)
+from .transformations import (
+    adjoint_from_transform, concat, exponential_coordinates_from_transform,
+    invert_transform, left_jacobian_SE3_inv, transform_from,
+    transform_from_exponential_coordinates)
 
 
-def estimate_gaussian_transform_from_samples(samples):
-    r"""Estimate Gaussian distribution over transformations from samples.
+def estimate_gaussian_rotation_matrix_from_samples(samples):
+    """Estimate Gaussian distribution over rotations from samples.
 
-    Uses iterative computation of mean described by Eade [1]_ and computes
-    covariance in exponential coordinate space (using an unbiased estimator).
-    The derivation of the procedure is presented by Pennec [2]_.
+    Computes the Fréchet mean of the samples and the covariance in tangent
+    space (exponential coordinates of rotation / rotation vectors) using an
+    unbiased estimator as outlines by Eade [1]_.
 
-    1. For a set of transformations :math:`\{T_1, \ldots, T_N\}`, we initialize
-       the estimated mean :math:`\overline{T}_0 \leftarrow T_1`.
-    2. For a fixed number of steps, in each iteration :math:`k` we improve the
-       estimation of the mean by
+    Parameters
+    ----------
+    samples : array-like, shape (n_samples, 3, 3)
+        Sampled rotations represented by rotation matrices.
 
-       a. Computing the distance of each sample to the current estimate of the
-          mean in exponential coordinates
-          :math:`V_{i,k} \leftarrow T_i \cdot \overline{T}_k^{-1}`.
-       b. Updating the estimate of the mean
-          :math:`\overline{T}_{k+1} \leftarrow Exp(\frac{1}{N}\sum_i V_{i,k}) \cdot \overline{T}_k`.
+    Returns
+    -------
+    mean : array, shape (3, 3)
+        Mean of the Gaussian distribution as rotation matrix.
 
-    3. Estimate the covariance in the space of exponential coordinates with
-       :math:`\Sigma \leftarrow \frac{1}{N - 1} \sum_i \left[ V_{i,k} \cdot V_{i,k}^T \right]`
+    cov : array, shape (3, 3)
+        Covariance of the Gaussian distribution in exponential coordinates.
+
+    See Also
+    --------
+    frechet_mean
+        Algorithm used to compute the mean of the Gaussian.
+
+    References
+    ----------
+    .. [1] Eade, E. (2017). Lie Groups for 2D and 3D Transformations.
+       https://ethaneade.com/lie.pdf
+    """
+    def compact_axis_angles_from_matrices(Rs):
+        A = axis_angles_from_matrices(Rs)
+        return A[:, :3] * A[:, 3, np.newaxis]
+
+    mean, mean_diffs = frechet_mean(
+        samples=samples,
+        mean0=samples[0],
+        exp=matrices_from_compact_axis_angles,
+        log=compact_axis_angles_from_matrices,
+        inv=lambda R: R.T,
+        concat_one_to_one=lambda R1, R2: np.dot(R2, R1),
+        concat_many_to_one=concat_many_to_one,
+        n_iter=20
+    )
+
+    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.
+
+    Computes the Fréchet mean of the samples and the covariance in tangent
+    space (exponential coordinates of transformation) using an unbiased
+    estimator as outlines by Eade [1]_.
 
     Parameters
     ----------
@@ -47,27 +85,117 @@ def estimate_gaussian_transform_from_samples(samples):
     cov : array, shape (6, 6)
         Covariance of distribution in exponential coordinate space.
 
+    See Also
+    --------
+    frechet_mean
+        Algorithm used to compute the mean of the Gaussian.
+
     References
     ----------
     .. [1] Eade, E. (2017). Lie Groups for 2D and 3D Transformations.
        https://ethaneade.com/lie.pdf
+    """
+    mean, mean_diffs = frechet_mean(
+        samples=samples,
+        mean0=samples[0],
+        exp=transforms_from_exponential_coordinates,
+        log=exponential_coordinates_from_transforms,
+        inv=invert_transform,
+        concat_one_to_one=concat,
+        concat_many_to_one=concat_many_to_one,
+        n_iter=20
+    )
+
+    cov = np.cov(mean_diffs, rowvar=False, bias=True)
+    return mean, cov
+
+
+def frechet_mean(
+        samples, mean0, exp, log, inv, concat_one_to_one, concat_many_to_one,
+        n_iter=20):
+    r"""Compute the Fréchet mean of samples on a smooth Riemannian manifold.
+
+    The mean is computed with an iterative optimization algorithm [1]_ [2]_:
+
+    1. For a set of samples :math:`\{x_1, \ldots, x_n\}`, we initialize
+       the estimated mean :math:`\bar{x}_0`, e.g., as :math:`x_1`.
+    2. For a fixed number of steps :math:`K`, in each iteration :math:`k` we
+       improve the estimation of the mean by
+
+       a. Computing the distance of each sample to the current estimate of the
+          mean in tangent space with
+          :math:`d_{i,k} \leftarrow \log (x_i \cdot \bar{x}_k^{-1})`.
+       b. Updating the estimate of the mean with
+          :math:`\bar{x}_{k+1} \leftarrow
+          \exp(\frac{1}{N}\sum_i d_{i,k}) \cdot \bar{x}_k`.
+
+    3. Return :math:`\bar{x}_K`.
+
+    Parameters
+    ----------
+    samples : array-like, shape (n_samples, ...)
+        Samples on a smooth Riemannian manifold.
+
+    mean0 : array-like, shape (...)
+        Initial guess for the mean on the manifold.
+
+    exp : callable
+        Exponential map from the tangent space to the manifold.
+
+    log : callable
+        Logarithmic map from the manifold to the tangent space.
+
+    inv : callable
+        Computes the inverse of an element on the manifold.
+
+    concat_one_to_one : callable
+        Concatenates elements on the manifold.
+
+    concat_many_to_one : callable
+        Concatenates multiple elements on the manifold to one element on the
+        manifold.
+
+    n_iter : int, optional (default: 20)
+        Number of iterations of the optimization algorithm.
+
+    Returns
+    -------
+    mean : array, shape (...)
+        Fréchet mean on the manifold.
+
+    mean_diffs : array, shape (n_samples, n_tangent_space_components)
+        Differences between the mean and the samples in the tangent space.
+        These can be used to compute the covariance. They are returned to
+        avoid recomputing them.
+
+    See Also
+    --------
+    estimate_gaussian_rotation_matrix_from_samples
+        Uses the Frechet mean to compute the mean of a Gaussian distribution
+        over rotations.
+
+    estimate_gaussian_transform_from_samples
+        Uses the Frechet mean to compute the mean of a Gaussian distribution
+        over transformations.
+
+    References
+    ----------
+    .. [1] Fréchet, M. (1948). Les éléments aléatoires de nature quelconque
+       dans un espace distancié. Annales de l’Institut Henri Poincaré, 10(3),
+       215–310.
 
     .. [2] Pennec, X. (2006). Intrinsic Statistics on Riemannian Manifolds:
        Basic Tools for Geometric Measurements. J Math Imaging Vis 25, 127-154.
        https://doi.org/10.1007/s10851-006-6228-4
     """
     assert len(samples) > 0
-    mean = samples[0]
-    for _ in range(20):
-        mean_inv = invert_transform(mean)
-        mean_diffs = exponential_coordinates_from_transforms(
-            concat_many_to_one(samples, mean_inv))
+    samples = np.asarray(samples)
+    mean = np.copy(mean0)
+    for _ in range(n_iter):
+        mean_diffs = log(concat_many_to_one(samples, inv(mean)))
         avg_mean_diff = np.mean(mean_diffs, axis=0)
-        mean = np.dot(
-            transform_from_exponential_coordinates(avg_mean_diff), mean)
-
-    cov = np.cov(mean_diffs, rowvar=False, bias=True)
-    return mean, cov
+        mean = concat_one_to_one(mean, exp(avg_mean_diff))
+    return mean, mean_diffs
 
 
 def invert_uncertain_transform(mean, cov):
@@ -500,11 +628,10 @@ def to_projected_ellipsoid(mean, cov, factor=1.96, n_steps=20):
     radii = factor * np.sqrt(vals[:3])
 
     # Grid on ellipsoid in exponential coordinate space
-    radius_x, radius_y, radius_z = radii
     x, y, z = unit_sphere_surface_grid(n_steps)
-    x *= radius_x
-    y *= radius_y
-    z *= radius_z
+    x *= radii[0]
+    y *= radii[1]
+    z *= radii[2]
     P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
     P = np.dot(P, vecs[:, :3].T)
 
diff --git a/pytransform3d/uncertainty.pyi b/pytransform3d/uncertainty.pyi
index b3719b4e..cace61c4 100644
--- a/pytransform3d/uncertainty.pyi
+++ b/pytransform3d/uncertainty.pyi
@@ -1,13 +1,25 @@
 import numpy as np
 import numpy.typing as npt
-from typing import Tuple, Union
+from typing import Tuple, Union, Callable
 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]: ...
 
 
+def frechet_mean(
+        samples: npt.ArrayLike, mean0: npt.ArrayLike,
+        exp: Callable, log: Callable,
+        inv: Callable, concat_one_to_one: Callable,
+        concat_many_to_one: Callable,
+        n_iter: int = ...) -> Tuple[np.ndarray, np.ndarray]: ...
+
+
 def invert_uncertain_transform(
         mean: npt.ArrayLike,
         cov: npt.ArrayLike) -> Tuple[np.ndarray, np.ndarray]: ...


Test patch:
diff --git a/pytransform3d/test/test_batch_rotations.py b/pytransform3d/test/test_batch_rotations.py
index be7bb58d..c65d7392 100644
--- a/pytransform3d/test/test_batch_rotations.py
+++ b/pytransform3d/test/test_batch_rotations.py
@@ -356,6 +356,23 @@ def test_axis_angles_from_matrices_output_variable():
     assert_array_almost_equal(A, A2_compact)
 
 
+def test_axis_angles_from_matrices_norot():
+    a = pbr.axis_angles_from_matrices(np.eye(3))
+    assert_array_almost_equal(a, [1, 0, 0, 0])
+
+    A = pbr.axis_angles_from_matrices([np.eye(3), np.eye(3)])
+    assert_array_almost_equal(A, [[1, 0, 0, 0], [1, 0, 0, 0]])
+
+    A = pbr.axis_angles_from_matrices(
+        [[np.eye(3), np.eye(3)], [np.eye(3), np.eye(3)]])
+    assert_array_almost_equal(
+        A,
+        [
+            [[1, 0, 0, 0], [1, 0, 0, 0]],
+            [[1, 0, 0, 0], [1, 0, 0, 0]]
+        ])
+
+
 def test_axis_angles_from_quaternions():
     rng = np.random.default_rng(48322)
     n_rotations = 20
diff --git a/pytransform3d/test/test_rotations.py b/pytransform3d/test/test_rotations.py
index 186f6604..923f02c5 100644
--- a/pytransform3d/test/test_rotations.py
+++ b/pytransform3d/test/test_rotations.py
@@ -2458,3 +2458,12 @@ def test_assert_euler_almost_equal():
         [0.2, 0.3, -0.5], [0.2 + np.pi, -0.3, -0.5 - np.pi], 0, 1, 0)
     pr.assert_euler_equal(
         [0.2, 0.3, -0.5], [0.2 + np.pi, np.pi - 0.3, -0.5 - np.pi], 0, 1, 2)
+
+
+def test_random_matrices():
+    rng = np.random.default_rng(33)
+    mean = pr.random_matrix(rng, np.eye(3), np.eye(3))
+    cov = np.eye(3) * 1e-10
+    samples = np.stack([pr.random_matrix(rng, mean, cov) for _ in range(10)])
+    for sample in samples:
+        assert_array_almost_equal(sample, mean, decimal=4)
diff --git a/pytransform3d/test/test_uncertainty.py b/pytransform3d/test/test_uncertainty.py
index d420e3c7..2e8ce5e3 100644
--- a/pytransform3d/test/test_uncertainty.py
+++ b/pytransform3d/test/test_uncertainty.py
@@ -6,6 +6,17 @@ from numpy.testing import assert_array_almost_equal
 import pytest
 
 
+def test_estimate_gaussian_rotation_matrix_from_samples():
+    rng = np.random.default_rng(34)
+    mean = pr.matrix_from_axis_angle([0.3, 0.2, 0.5, 0.25 * np.pi])
+    cov = np.diag([0.1, 0.05, 0.09]) ** 2
+    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=2)
+    assert_array_almost_equal(cov, cov_est, decimal=2)
+
+
 def test_same_fuse_poses():
     mean1 = np.array([
         [0.8573, -0.2854, 0.4285, 3.5368],