From ecd252a2708fa80a7e3dc02957434e6382f8ae37 Mon Sep 17 00:00:00 2001
From: Hs293Go <hellston20a@gmail.com>
Date: Mon, 4 Apr 2022 13:27:40 -0400
Subject: [PATCH] Implement point cloud alignment for SE3

This is equivalent to `Eigen::umeyama` but faster and stores the results
more compactly.
---
 sophus/se3.hpp         | 47 ++++++++++++++++++++++++++++++++++++++++++
 test/core/test_se3.cpp | 44 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/sophus/se3.hpp b/sophus/se3.hpp
index c3c4290f7..96a666744 100644
--- a/sophus/se3.hpp
+++ b/sophus/se3.hpp
@@ -1014,6 +1014,53 @@ class SE3 : public SE3Base<SE3<Scalar_, Options>> {
     return upsilon_omega;
   }
 
+  template <typename Derived, typename OtherDerived>
+  SOPHUS_FUNC void setFromPointClouds(
+      Eigen::MatrixBase<Derived> const& src,
+      Eigen::MatrixBase<OtherDerived> const& dst) {
+    using Eigen::Index;
+    using MatrixType =
+        typename Eigen::internal::plain_matrix_type<Derived>::type;
+    const Index N = src.cols();
+
+    SOPHUS_ENSURE(
+        src.rows() == 3 && dst.rows() == 3 && N == dst.cols(),
+        "Point clouds are stored in 3xN matrices, one point per column. Source "
+        "and destination matrices are {}x{} and {}x{} respectively",
+        src.rows(), N, dst.rows(), dst.cols());
+
+    Scalar w(1. / N);
+    Vector3<Scalar> const src_mean = src.rowwise().sum() * w;
+    Vector3<Scalar> const dst_mean = dst.rowwise().sum() * w;
+
+    MatrixType m1 = src.colwise() - src_mean;
+    MatrixType m2 = dst.colwise() - dst_mean;
+
+    m1 += m2;
+    m2 = m1 - Scalar(2) * m2;
+    Matrix4<Scalar> W = Matrix4<Scalar>::Zero();
+    for (Index i(0); i < N; ++i) {
+      Matrix3<Scalar> const S = SO3<Scalar>::hat(m1.col(i));
+      Vector3<Scalar> const v = m2.col(i);
+      Matrix4<Scalar> V;
+      V << S, -v, v.transpose(), Scalar(0);
+      W += w * V.transpose() * V;
+    }
+    const Eigen::SelfAdjointEigenSolver<Matrix4<Scalar>> eig(W);
+    const Eigen::Quaternion<Scalar> rotation(eig.eigenvectors().real().col(0));
+    so3_.setQuaternion(rotation);
+    translation_ = dst_mean - rotation * src_mean;
+  }
+
+  template <typename Derived, typename OtherDerived>
+  SOPHUS_FUNC static SE3<Scalar> FromPointClouds(
+      Eigen::MatrixBase<Derived> const& src,
+      Eigen::MatrixBase<OtherDerived> const& dst) {
+    SE3<Scalar> ret;
+    ret.setFromPointClouds(src, dst);
+    return ret;
+  }
+
  protected:
   SO3Member so3_;
   TranslationMember translation_;
diff --git a/test/core/test_se3.cpp b/test/core/test_se3.cpp
index 5dffc5599..16ca9dd7f 100644
--- a/test/core/test_se3.cpp
+++ b/test/core/test_se3.cpp
@@ -57,6 +57,7 @@ class Tests {
     passed &= testMutatingAccessors();
     passed &= testConstructors();
     passed &= testFit();
+    passed &= testPointCloudAlignment();
     processTestResult(passed);
   }
 
@@ -243,6 +244,49 @@ class Tests {
     return true;
   }
 
+  template <class S = Scalar>
+  enable_if_t<std::is_floating_point<S>::value, bool>
+  testPointCloudAlignment() {
+    bool passed = true;
+
+    using MatrixType = Eigen::Matrix<S, 3, Eigen::Dynamic>;
+
+    const int sz = point_vec_.size();
+    Eigen::Map<const MatrixType> src(point_vec_.front().data(), 3, sz);
+    MatrixType dst(3, sz);
+
+    for (const auto& it : se3_vec_) {
+      for (int i = 0; i < sz; ++i) {
+        dst.col(i) = it * src.col(i);
+      }
+      SE3Type result;
+      result.setFromPointClouds(src, dst);
+      SOPHUS_TEST_APPROX(passed, it.translation(), result.translation(),
+                         Constants<Scalar>::epsilon(),
+                         "Expected translation: {}; Got: {}", it.translation(),
+                         result.translation());
+      SOPHUS_TEST_APPROX(
+          passed, it.matrix(), result.matrix(), Constants<Scalar>::epsilon(),
+          "Expected rotation: {}; Got: {}", it.matrix(), result.matrix());
+
+      result = SE3Type::FromPointClouds(src, dst);
+      SOPHUS_TEST_APPROX(passed, it.translation(), result.translation(),
+                         Constants<Scalar>::epsilon(),
+                         "Expected translation: {}; Got: {}", it.translation(),
+                         result.translation());
+      SOPHUS_TEST_APPROX(
+          passed, it.matrix(), result.matrix(), Constants<Scalar>::epsilon(),
+          "Expected rotation: {}; Got: {}", it.matrix(), result.matrix());
+    }
+    return passed;
+  }
+
+  template <class S = Scalar>
+  enable_if_t<!std::is_floating_point<S>::value, bool>
+  testPointCloudAlignment() {
+    return true;
+  }
+
   std::vector<SE3Type, Eigen::aligned_allocator<SE3Type>> se3_vec_;
   std::vector<Tangent, Eigen::aligned_allocator<Tangent>> tangent_vec_;
   std::vector<Point, Eigen::aligned_allocator<Point>> point_vec_;