Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Oct 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions include/nuri/core/geometry.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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*
Expand Down Expand Up @@ -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
Expand Down
36 changes: 23 additions & 13 deletions src/core/geometry.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -568,8 +568,9 @@
}

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))
Expand Down Expand Up @@ -598,7 +599,7 @@
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;
Expand All @@ -610,16 +611,16 @@
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 };

Check warning on line 614 in src/core/geometry.cpp

View check run for this annotation

Codecov / codecov/patch

src/core/geometry.cpp#L614

Added line #L614 was not covered by tests
}

eig -= f / df;

if (std::abs(oldeig - eig) < std::abs(evalprec * eig))
return eig;
return { eig, true };
}

return -1;
return { eig, false };

Check warning on line 623 in src/core/geometry.cpp

View check run for this annotation

Codecov / codecov/patch

src/core/geometry.cpp#L623

Added line #L623 was not covered by tests
}

std::pair<Eigen::Quaterniond, bool>
Expand Down Expand Up @@ -720,8 +721,9 @@

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();
Expand All @@ -734,17 +736,22 @@
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;
}

Expand All @@ -758,6 +765,9 @@
}

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();
Expand Down
16 changes: 15 additions & 1 deletion test/core/geometry_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
});
}

Expand Down