Merge pull request #390 from seoklab/qcp-reflection

feat(core/geometry): add option to consider reflection in qcp algorithm

include/nuri/core/geometry.h

@@ -487,6 +487,7 @@ kabsch(const Eigen::Ref<const Matrix3Xd> &query,
  * @param mode Selects the return value. Defaults to AlignMode::kBoth. Note that
  *        even if AlignMode::kXformOnly is selected, the MSD value will report a
  *        negative value if the calculation fails.
+ * @param reflection Whether to allow reflection. Defaults to false.
  * @param evalprec The precision of eigenvalue calculation. Defaults to 1e-11.
  * @param evecprec The precision of eigenvector calculation. Defaults to 1e-6.
  * @param maxiter The maximum number of Newton-Raphson iterations. Defaults
@@ -499,9 +500,13 @@ kabsch(const Eigen::Ref<const Matrix3Xd> &query,
  *         guarantee convergence.
  *
  * This implementation is based on the reference implementation by P Liu and DL
- * Theobald, but modified for better stability and error handling.
+ * Theobald, but modified for better stability and error handling. Also, an
+ * option to allow reflection is added based on observations of EA Coutsias, C
+ * Seok, and KA Dill (see more details in the following references).
  *
  * References:
+ * - EA Coutsias, C Seok, and KA Dill. *J. Comput. Chem.* **2004**, *25* (15),
+ *   1849-1857. DOI:[10.1002/jcc.20110](https://doi.org/10.1002/jcc.20110)
  * - P Liu, DK Agrafiotis, and DL Theobald. *J. Comput. Chem.* **2011**, *32*
  *   (1), 185-186. DOI:[10.1002/jcc.21607](https://doi.org/10.1002/jcc.21607)
  * - P Liu, DK Agrafiotis, and DL Theobald. *J. Comput. Chem.* **2010**, *31*
@@ -547,7 +552,8 @@ kabsch(const Eigen::Ref<const Matrix3Xd> &query,
 extern std::pair<Affine3d, double>
 qcp(const Eigen::Ref<const Matrix3Xd> &query,
     const Eigen::Ref<const Matrix3Xd> &templ, AlignMode mode = AlignMode::kBoth,
-    double evalprec = 1e-11, double evecprec = 1e-6, int maxiter = 50);
+    bool reflection = false, double evalprec = 1e-11, double evecprec = 1e-6,
+    int maxiter = 50);
 
 /**
  * @brief A routine for converting squared pairwise distances to cartesian

src/core/geometry.cpp

@@ -568,8 +568,9 @@ std::pair<Affine3d, double> kabsch(const Eigen::Ref<const Matrix3Xd> &query,
 }
 
 namespace {
-  double qcp_find_largest_eig(const Matrix3d &R, const double e0,
-                              const double evalprec, const int maxiter) {
+  std::pair<double, bool>
+  qcp_find_nearest_eig(const Matrix3d &R, const double det, const double e0,
+                       const double evalprec, const int maxiter) {
     const Matrix3d Rsq = R.cwiseAbs2();
 
     double F = (-(R(0, 2) + R(2, 0)) * (R(1, 2) - R(2, 1))
@@ -598,7 +599,7 @@ namespace {
            E2 = 2 * (R(1, 1) * R(2, 2) - R(1, 2) * R(2, 1)),
            E = (E1 - E2) * (E1 + E2);
 
-    const double c2 = -2 * Rsq.sum(), c1 = -8 * R.determinant(),
+    const double c2 = -2 * Rsq.sum(), c1 = -8 * det,
                  c0 = Dsqrt * Dsqrt + E + F + G + H + I;
 
     double eig = e0, oldeig;
@@ -610,16 +611,16 @@ namespace {
       const double f = esq * esq + c2 * esq + c1 * eig + c0;
       if (ABSL_PREDICT_FALSE(std::abs(df) <= kEps)) {
         // singular; test for convergence
-        return std::abs(f) <= evalprec ? eig : -1;
+        return { eig, std::abs(f) <= evalprec };
       }
 
       eig -= f / df;
 
       if (std::abs(oldeig - eig) < std::abs(evalprec * eig))
-        return eig;
+        return { eig, true };
     }
 
-    return -1;
+    return { eig, false };
   }
 
   std::pair<Eigen::Quaterniond, bool>
@@ -720,8 +721,9 @@ namespace {
 
 std::pair<Affine3d, double> qcp(const Eigen::Ref<const Matrix3Xd> &query,
                                 const Eigen::Ref<const Matrix3Xd> &templ,
-                                AlignMode mode, const double evalprec,
-                                const double evecprec, const int maxiter) {
+                                AlignMode mode, bool reflection,
+                                const double evalprec, const double evecprec,
+                                const int maxiter) {
   std::pair<Affine3d, double> ret { {}, 0.0 };
 
   MatrixX3d qt = query.transpose();
@@ -734,17 +736,22 @@ std::pair<Affine3d, double> qcp(const Eigen::Ref<const Matrix3Xd> &query,
   tt.rowwise() -= tm;
 
   const Matrix3d R = tt.transpose() * qt;
-  const double GA = tt.cwiseAbs2().sum(), GB = qt.cwiseAbs2().sum();
+  const double GA = tt.cwiseAbs2().sum(), GB = qt.cwiseAbs2().sum(),
+               det = R.determinant();
 
-  const double eig = qcp_find_largest_eig(R, (GA + GB) / 2, evalprec, maxiter);
-  if (eig < 0) {
+  double e0 = (GA + GB) / 2;
+  if (reflection)
+    e0 = std::copysign(e0, det);
+
+  const auto [eig, ok] = qcp_find_nearest_eig(R, det, e0, evalprec, maxiter);
+  if (!ok) {
     ret.second = -1;
     return ret;
   }
 
   if (mode != AlignMode::kXformOnly) {
-    double msd =
-        nonnegative(GA + GB - 2 * eig) / static_cast<double>(templ.cols());
+    double msd = nonnegative(GA + GB - 2 * std::abs(eig))
+                 / static_cast<double>(templ.cols());
     ret.second = msd;
   }
 
@@ -758,6 +765,9 @@ std::pair<Affine3d, double> qcp(const Eigen::Ref<const Matrix3Xd> &query,
   }
 
   ret.first.linear() = qhat.toRotationMatrix();
+  if (reflection)
+    ret.first.linear() *= std::copysign(1, det);
+
   ret.first.translation().noalias() =
       -1 * (ret.first.linear() * qm.transpose());
   ret.first.translation() += tm.transpose();

test/core/geometry_test.cpp

@@ -346,19 +346,33 @@ TEST_F(AlignTest, KabschBoth) {
 TEST_F(AlignTest, QcpMSDOnly) {
   auto [_, msd] = qcp(query_, templ_, AlignMode::kMsdOnly);
   EXPECT_NEAR(msd, msd_reflected_, 1e-6);
+
+  std::tie(_, msd) = qcp(query_, templ_, AlignMode::kMsdOnly, true);
+  EXPECT_NEAR(msd, msd_, 1e-6);
 }
 
 TEST_F(AlignTest, QcpXformOnly) {
   auto [xform, flag] = qcp(query_, templ_, AlignMode::kXformOnly);
   ASSERT_GE(flag, 0);
   NURI_EXPECT_EIGEN_EQ_TOL(xform.matrix(), xform_reflected_.matrix(), 1e-3);
+
+  std::tie(xform, flag) = qcp(query_, templ_, AlignMode::kXformOnly, true);
+  ASSERT_GE(flag, 0);
+  NURI_EXPECT_EIGEN_EQ_TOL(xform.linear(), -xform_.linear(), 1e-3);
+  NURI_EXPECT_EIGEN_EQ_TOL(xform.translation(), xform_.translation(), 1e-3);
 }
 
 TEST_F(AlignTest, QcpBoth) {
   auto [xform, msd] = qcp(query_, templ_, AlignMode::kBoth);
   ASSERT_GE(msd, 0);
   NURI_EXPECT_EIGEN_EQ_TOL(xform.matrix(), xform_reflected_.matrix(), 1e-3);
   EXPECT_NEAR(msd, msd_reflected_, 1e-6);
+
+  std::tie(xform, msd) = qcp(query_, templ_, AlignMode::kBoth, true);
+  ASSERT_GE(msd, 0);
+  NURI_EXPECT_EIGEN_EQ_TOL(xform.linear(), -xform_.linear(), 1e-3);
+  NURI_EXPECT_EIGEN_EQ_TOL(xform.translation(), xform_.translation(), 1e-3);
+  EXPECT_NEAR(msd, msd_, 1e-6);
 }
 
 class AlignSingularTest: public ::testing::Test {
@@ -420,7 +434,7 @@ TEST_F(AlignSingularTest, Kabsch) {
 TEST_F(AlignSingularTest, Qcp) {
   run_test([](const auto &q, const auto &t, auto mode) {
     // might fail on optimized builds if evecprec is too high
-    return qcp(q, t, mode, 1e-11, 1e-8);
+    return qcp(q, t, mode, false, 1e-11, 1e-8);
   });
 }