Optimize the fit gradients

doc/doxygen/Doxyfile

@@ -1420,7 +1420,7 @@ FORMULA_TRANSPARENT    = YES
 # The default value is: NO.
 # This tag requires that the tag GENERATE_HTML is set to YES.
 
-USE_MATHJAX            = NO
+USE_MATHJAX            = YES
 
 # When MathJax is enabled you can set the default output format to be used for
 # the MathJax output. See the MathJax site (see:

src/colvar_rotation_derivative.h

@@ -5,11 +5,21 @@
 #include <type_traits>
 #include <cstring>
 
+#ifndef _noalias
+#if defined(__INTEL_COMPILER) || (defined(__PGI) && !defined(__NVCOMPILER))
+#define _noalias restrict
+#elif defined(__GNUC__) || defined(__INTEL_LLVM_COMPILER) || defined(__NVCOMPILER)
+#define _noalias __restrict
+#else
+#define _noalias
+#endif
+#endif
+
 /// \brief Helper function for loading the ia-th atom in the vector pos to x, y and z (C++11 SFINAE is used)
 template <typename T, typename std::enable_if<std::is_same<T, cvm::atom_pos>::value, bool>::type = true>
 inline void read_atom_coord(
   size_t ia, const std::vector<T>& pos,
-  cvm::real* x, cvm::real* y, cvm::real* z) {
+  cvm::real* _noalias x, cvm::real* _noalias y, cvm::real* _noalias z) {
   *x = pos[ia].x;
   *y = pos[ia].y;
   *z = pos[ia].z;
@@ -18,7 +28,7 @@ inline void read_atom_coord(
 template <typename T, typename std::enable_if<std::is_same<T, cvm::atom>::value, bool>::type = true>
 inline void read_atom_coord(
   size_t ia, const std::vector<T>& pos,
-  cvm::real* x, cvm::real* y, cvm::real* z) {
+  cvm::real* _noalias x, cvm::real* _noalias y, cvm::real* _noalias z) {
   *x = pos[ia].pos.x;
   *y = pos[ia].pos.y;
   *z = pos[ia].pos.z;
@@ -327,12 +337,13 @@ struct rotation_derivative {
     *  @param[out] dq0_out The output of derivative of Q
     *  @param[out] ds_out  The output of derivative of overlap matrix S
     */
+  template <bool use_dl, bool use_dq, bool use_ds>
   void calc_derivative_impl(
     const cvm::rvector (&ds)[4][4],
-    cvm::rvector* const dl0_out,
-    cvm::vector1d<cvm::rvector>* const dq0_out,
-    cvm::matrix2d<cvm::rvector>* const ds_out) const {
-    if (ds_out != nullptr) {
+    cvm::rvector* _noalias const dl0_out,
+    cvm::vector1d<cvm::rvector>* _noalias const dq0_out,
+    cvm::matrix2d<cvm::rvector>* _noalias const ds_out) const {
+    if (use_ds) {
       // this code path is for debug_gradients, so not necessary to unroll the loop
       *ds_out = cvm::matrix2d<cvm::rvector>(4, 4);
       for (int i = 0; i < 4; ++i) {
@@ -341,7 +352,7 @@ struct rotation_derivative {
         }
       }
     }
-    if (dl0_out != nullptr) {
+    if (use_dl) {
       /* manually loop unrolling of the following loop:
         dl0_1.reset();
         for (size_t i = 0; i < 4; i++) {
@@ -367,7 +378,7 @@ struct rotation_derivative {
                  tmp_Q0Q0[3][2] * ds[3][2] +
                  tmp_Q0Q0[3][3] * ds[3][3];
     }
-    if (dq0_out != nullptr) {
+    if (use_dq) {
       // we can skip this check if a fixed-size array is used
       if (dq0_out->size() != 4) dq0_out->resize(4);
       /* manually loop unrolling of the following loop:
@@ -462,32 +473,21 @@ struct rotation_derivative {
     *  @param[out] ds_1_out  The output of derivative of overlap matrix S with
     *                        respect to ia-th atom of group 1
     */
+  template <bool use_dl, bool use_dq, bool use_ds>
   void calc_derivative_wrt_group1(
-    size_t ia, cvm::rvector* const dl0_1_out = nullptr,
-    cvm::vector1d<cvm::rvector>* const dq0_1_out = nullptr,
-    cvm::matrix2d<cvm::rvector>* const ds_1_out = nullptr) const {
-      if (dl0_1_out == nullptr && dq0_1_out == nullptr) return;
+    size_t ia, cvm::rvector* _noalias const dl0_1_out = nullptr,
+    cvm::vector1d<cvm::rvector>* _noalias const dq0_1_out = nullptr,
+    cvm::matrix2d<cvm::rvector>* _noalias const ds_1_out = nullptr) const {
+      // if (dl0_1_out == nullptr && dq0_1_out == nullptr) return;
       cvm::real a2x, a2y, a2z;
       // we can get rid of the helper function read_atom_coord if C++17 (constexpr) is available
       read_atom_coord(ia, m_pos2, &a2x, &a2y, &a2z);
-      cvm::rvector ds_1[4][4];
-      ds_1[0][0].set( a2x,  a2y,  a2z);
-      ds_1[1][0].set( 0.0,  a2z, -a2y);
-      ds_1[0][1] = ds_1[1][0];
-      ds_1[2][0].set(-a2z,  0.0,  a2x);
-      ds_1[0][2] = ds_1[2][0];
-      ds_1[3][0].set( a2y, -a2x,  0.0);
-      ds_1[0][3] = ds_1[3][0];
-      ds_1[1][1].set( a2x, -a2y, -a2z);
-      ds_1[2][1].set( a2y,  a2x,  0.0);
-      ds_1[1][2] = ds_1[2][1];
-      ds_1[3][1].set( a2z,  0.0,  a2x);
-      ds_1[1][3] = ds_1[3][1];
-      ds_1[2][2].set(-a2x,  a2y, -a2z);
-      ds_1[3][2].set( 0.0,  a2z,  a2y);
-      ds_1[2][3] = ds_1[3][2];
-      ds_1[3][3].set(-a2x, -a2y,  a2z);
-      calc_derivative_impl(ds_1, dl0_1_out, dq0_1_out, ds_1_out);
+      const cvm::rvector ds_1[4][4] = {
+        {{ a2x,  a2y,  a2z}, { 0.0, a2z,  -a2y}, {-a2z,  0.0,  a2x}, { a2y, -a2x,  0.0}},
+        {{ 0.0,  a2z, -a2y}, { a2x, -a2y, -a2z}, { a2y,  a2x,  0.0}, { a2z,  0.0,  a2x}},
+        {{-a2z,  0.0,  a2x}, { a2y,  a2x,  0.0}, {-a2x,  a2y, -a2z}, { 0.0,  a2z,  a2y}},
+        {{ a2y, -a2x,  0.0}, { a2z,  0.0,  a2x}, { 0.0,  a2z,  a2y}, {-a2x, -a2y,  a2z}}};
+      calc_derivative_impl<use_dl, use_dq, use_ds>(ds_1, dl0_1_out, dq0_1_out, ds_1_out);
     }
   /*! @brief Calculate the derivatives of S, the leading eigenvalue L and
     *         the leading eigenvector Q with respect to `m_pos2`
@@ -499,32 +499,21 @@ struct rotation_derivative {
     *  @param[out] ds_2_out  The output of derivative of overlap matrix S with
     *                        respect to ia-th atom of group 2
     */
+  template <bool use_dl, bool use_dq, bool use_ds>
   void calc_derivative_wrt_group2(
-    size_t ia, cvm::rvector* const dl0_2_out = nullptr,
-    cvm::vector1d<cvm::rvector>* const dq0_2_out = nullptr,
-    cvm::matrix2d<cvm::rvector>* const ds_2_out = nullptr) const {
-    if (dl0_2_out == nullptr && dq0_2_out == nullptr) return;
+    size_t ia, cvm::rvector* _noalias const dl0_2_out = nullptr,
+    cvm::vector1d<cvm::rvector>* _noalias const dq0_2_out = nullptr,
+    cvm::matrix2d<cvm::rvector>* _noalias const ds_2_out = nullptr) const {
+    // if (dl0_2_out == nullptr && dq0_2_out == nullptr) return;
     cvm::real a1x, a1y, a1z;
     // we can get rid of the helper function read_atom_coord if C++17 (constexpr) is available
     read_atom_coord(ia, m_pos1, &a1x, &a1y, &a1z);
-    cvm::rvector ds_2[4][4];
-    ds_2[0][0].set( a1x,  a1y,  a1z);
-    ds_2[1][0].set( 0.0, -a1z,  a1y);
-    ds_2[0][1] = ds_2[1][0];
-    ds_2[2][0].set( a1z,  0.0, -a1x);
-    ds_2[0][2] = ds_2[2][0];
-    ds_2[3][0].set(-a1y,  a1x,  0.0);
-    ds_2[0][3] = ds_2[3][0];
-    ds_2[1][1].set( a1x, -a1y, -a1z);
-    ds_2[2][1].set( a1y,  a1x,  0.0);
-    ds_2[1][2] = ds_2[2][1];
-    ds_2[3][1].set( a1z,  0.0,  a1x);
-    ds_2[1][3] = ds_2[3][1];
-    ds_2[2][2].set(-a1x,  a1y, -a1z);
-    ds_2[3][2].set( 0.0,  a1z,  a1y);
-    ds_2[2][3] = ds_2[3][2];
-    ds_2[3][3].set(-a1x, -a1y,  a1z);
-    calc_derivative_impl(ds_2, dl0_2_out, dq0_2_out, ds_2_out);
+    const cvm::rvector ds_2[4][4] = {
+      {{ a1x,  a1y,  a1z}, { 0.0, -a1z,  a1y}, { a1z,  0.0, -a1x}, {-a1y,  a1x,  0.0}},
+      {{ 0.0, -a1z,  a1y}, { a1x, -a1y, -a1z}, { a1y,  a1x,  0.0}, { a1z,  0.0,  a1x}},
+      {{ a1z,  0.0, -a1x}, { a1y,  a1x,  0.0}, {-a1x,  a1y, -a1z}, { 0.0,  a1z,  a1y}},
+      {{-a1y,  a1x,  0.0}, { a1z,  0.0,  a1x}, { 0.0,  a1z,  a1y}, {-a1x, -a1y,  a1z}}};
+    calc_derivative_impl<use_dl, use_dq, use_ds>(ds_2, dl0_2_out, dq0_2_out, ds_2_out);
   }
 };
 
@@ -588,7 +577,7 @@ void debug_gradients(
     // cvm::real const &a1x = pos1[ia].x;
     // cvm::real const &a1y = pos1[ia].y;
     // cvm::real const &a1z = pos1[ia].z;
-    deriv.calc_derivative_wrt_group2(ia, &dl0_2, &dq0_2, &ds_2);
+    deriv.template calc_derivative_wrt_group2<true, true, true>(ia, &dl0_2, &dq0_2, &ds_2);
     // make an infitesimal move along each cartesian coordinate of
     // this atom, and solve again the eigenvector problem
     for (size_t comp = 0; comp < 3; comp++) {

src/colvaratoms.cpp

@@ -1055,6 +1055,14 @@ void cvm::atom_group::calc_apply_roto_translation()
     }
   }
 
+  if (is_enabled(f_ag_fit_gradients) && !b_dummy) {
+    // Save the unrotated frame for fit gradients
+    pos_unrotated.resize(size());
+    for (size_t i = 0; i < size(); ++i) {
+      pos_unrotated[i] = atoms[i].pos;
+    }
+  }
+
   if (is_enabled(f_ag_rotate)) {
     // rotate the group (around the center of geometry if f_ag_center is
     // enabled, around the origin otherwise)
@@ -1250,17 +1258,13 @@ void cvm::atom_group::calc_fit_forces_impl(
   cvm::vector1d<cvm::rvector> dq0_1(4);
   // loop 1: iterate over the current atom group
   for (size_t i = 0; i < size(); i++) {
-    cvm::atom_pos pos_orig;
     if (B_ag_center) {
       atom_grad += accessor_main(i);
-      if (B_ag_rotate) pos_orig = rot_inv * (atoms[i].pos - ref_pos_cog);
-    } else {
-      if (B_ag_rotate) pos_orig = rot_inv * atoms[i].pos;
     }
     if (B_ag_rotate) {
       // calculate \partial(R(q) \vec{x}_i)/\partial q) \cdot \partial\xi/\partial\vec{x}_i
       cvm::quaternion const dxdq =
-        rot.q.position_derivative_inner(pos_orig, accessor_main(i));
+        rot.q.position_derivative_inner(pos_unrotated[i], accessor_main(i));
       sum_dxdq[0] += dxdq[0];
       sum_dxdq[1] += dxdq[1];
       sum_dxdq[2] += dxdq[2];
@@ -1279,7 +1283,7 @@ void cvm::atom_group::calc_fit_forces_impl(
       fitting_force_grad += atom_grad;
     }
     if (B_ag_rotate) {
-      rot_deriv->calc_derivative_wrt_group1(j, nullptr, &dq0_1);
+      rot_deriv->calc_derivative_wrt_group1<false, true, false>(j, nullptr, &dq0_1);
       // multiply by {\partial q}/\partial\vec{x}_j and add it to the fit gradients
       fitting_force_grad += sum_dxdq[0] * dq0_1[0] +
                             sum_dxdq[1] * dq0_1[1] +
@@ -1296,7 +1300,7 @@ void cvm::atom_group::calc_fit_forces_impl(
 template <typename main_force_accessor_T, typename fitting_force_accessor_T>
 void cvm::atom_group::calc_fit_forces(
   main_force_accessor_T accessor_main,
-    fitting_force_accessor_T accessor_fitting) const {
+  fitting_force_accessor_T accessor_fitting) const {
   if (is_enabled(f_ag_center) && is_enabled(f_ag_rotate))
     calc_fit_forces_impl<true, true, main_force_accessor_T, fitting_force_accessor_T>(accessor_main, accessor_fitting);
   if (is_enabled(f_ag_center) && !is_enabled(f_ag_rotate))
@@ -1533,7 +1537,7 @@ void cvm::atom_group::group_force_object::apply_force_with_fitting_group() {
   // computed. For a vector component, we can only know the forces on the fitting
   // group, but checking this flag can mimic results that the users expect (if
   // "enableFitGradients no" then there is no force on the fitting group).
-  if (m_ag->is_enabled(f_ag_fit_gradients)) {
+  if (!m_ag->b_dummy && m_ag->is_enabled(f_ag_fit_gradients)) {
     auto accessor_main = [this](size_t i){return m_ag->group_forces[i];};
     auto accessor_fitting = [this](size_t j, const cvm::rvector& fitting_force){
       (*(m_group_for_fit))[j].apply_force(fitting_force);

src/colvaratoms.h

@@ -264,10 +264,46 @@ class colvarmodule::atom_group
 
 public:
 
+  /*! @class group_force_object
+   *  @brief A helper class for applying forces on an atom group in a way that
+   *         is aware of the fitting group. NOTE: you are encouraged to use
+   *         get_group_force_object() to get an instance of group_force_object
+   *         instead of constructing directly.
+   */
   class group_force_object {
   public:
+    /*! @brief Constructor of group_force_object
+     *  @param ag The pointer to the atom group that forces will be applied on.
+     */
     group_force_object(cvm::atom_group* ag);
+    /*! @brief Destructor of group_force_object
+     */
     ~group_force_object();
+    /*! @brief Apply force to atom i
+     *  @param i The i-th of atom in the atom group.
+     *  @param force The force being added to atom i.
+     *
+     * The function can be used as follows,
+     * @code
+     *       // In your colvar::cvc::apply_force() loop of a component:
+     *       auto ag_force = atoms->get_group_force_object();
+     *       for (ia = 0; ia < atoms->size(); ia++) {
+     *         const cvm::rvector f = compute_force_on_atom_ia();
+     *         ag_force.add_atom_force(ia, f);
+     *       }
+     * @endcode
+     * There are actually two scenarios under the hood:
+     * (i) If the atom group does not have a fitting group, then the force is
+     *     added to atom i directly;
+     * (ii) If the atom group has a fitting group, the force on atom i will just
+     *      be temporary stashed into ag->group_forces. At the end of the loop
+     *      of apply_force(), the destructor ~group_force_object() will be called,
+     *      which then call apply_force_with_fitting_group(). The forces on the
+     *      main group will be rotated back by multiplying ag->group_forces with
+     *      the inverse rotation. The forces on the fitting group (if
+     *      enableFitGradients is on) will be calculated by calling
+     *      calc_fit_forces.
+     */
     void add_atom_force(size_t i, const cvm::rvector& force);
   private:
     cvm::atom_group* m_ag;
@@ -442,6 +478,9 @@ class colvarmodule::atom_group
   /// \brief Center of geometry before any fitting
   cvm::atom_pos cog_orig;
 
+  /// \brief Unrotated atom positions for fit gradients
+  std::vector<cvm::atom_pos> pos_unrotated;
+
 public:
 
   /// \brief Return the center of geometry of the atomic positions
@@ -525,15 +564,25 @@ class colvarmodule::atom_group
  *  @tparam B_ag_rotate Calculate the optimal rotation? This should follow
  *          the value of `is_enabled(f_ag_rotate)`.
  *  @tparam main_force_accessor_T The type of accessor of the main
- *          group forces or gradients.
+ *          group forces or gradients acting on the rotated frame.
  *  @tparam fitting_force_accessor_T The type of accessor of the fitting group
  *          forces or gradients.
  *  @param accessor_main The accessor of the main group forces or gradients.
  *         accessor_main(i) should return the i-th force or gradient of the
- *         main group.
+ *         rotated main group.
  *  @param accessor_fitting The accessor of the fitting group forces or gradients.
  *         accessor_fitting(j, v) should store/apply the j-th atom gradient or
  *         force in the fitting group.
+ *
+ *  This function is used to (i) project the gradients of CV with respect to
+ *  rotated main group atoms to fitting group atoms, or (ii) project the forces
+ *  on rotated main group atoms to fitting group atoms, by the following two steps
+ *  (using the goal (ii) for example):
+ *  (1) Loop over the positions of main group atoms and call cvm::quaternion::position_derivative_inner
+ *      to project the forces on rotated main group atoms to the forces on quaternion.
+ *  (2) Loop over the positions of fitting group atoms, compute the gradients of
+ *      \f$\mathbf{q}\f$ with respect to the position of each atom, and then multiply
+ *      that with the force on \f$\mathbf{q}\f$ (chain rule).
  */
   template <bool B_ag_center, bool B_ag_rotate,
             typename main_force_accessor_T, typename fitting_force_accessor_T>
@@ -552,6 +601,9 @@ class colvarmodule::atom_group
  *  @param accessor_fitting The accessor of the fitting group forces or gradients.
  *         accessor_fitting(j, v) should store/apply the j-th atom gradient or
  *         force in the fitting group.
+ *
+ *  This function just dispatches the parameters to calc_fit_forces_impl that really
+ *  performs the calculations.
  */
   template <typename main_force_accessor_T, typename fitting_force_accessor_T>
   void calc_fit_forces(

src/colvarcomp.cpp

@@ -533,8 +533,10 @@ void colvar::cvc::calc_Jacobian_derivative()
 
 void colvar::cvc::calc_fit_gradients()
 {
-  for (size_t ig = 0; ig < atom_groups.size(); ig++) {
-    atom_groups[ig]->calc_fit_gradients();
+  if (is_enabled(f_cvc_explicit_gradient)) {
+    for (size_t ig = 0; ig < atom_groups.size(); ig++) {
+      atom_groups[ig]->calc_fit_gradients();
+    }
   }
 }
 

src/colvarcomp_distances.cpp

@@ -1003,7 +1003,7 @@ void colvar::rmsd::calc_Jacobian_derivative()
     for (size_t ia = 0; ia < atoms->size(); ia++) {
 
       // Gradient of optimal quaternion wrt current Cartesian position
-      atoms->rot_deriv->calc_derivative_wrt_group1(ia, nullptr, &dq);
+      atoms->rot_deriv->calc_derivative_wrt_group1<false, true, false>(ia, nullptr, &dq);
 
       g11 = 2.0 * (atoms->rot.q)[1]*dq[1];
       g22 = 2.0 * (atoms->rot.q)[2]*dq[2];
@@ -1302,7 +1302,7 @@ void colvar::eigenvector::calc_Jacobian_derivative()
     // Gradient of optimal quaternion wrt current Cartesian position
     // trick: d(R^-1)/dx = d(R^t)/dx = (dR/dx)^t
     // we can just transpose the derivatives of the direct matrix
-    atoms->rot_deriv->calc_derivative_wrt_group1(ia, nullptr, &dq_1);
+    atoms->rot_deriv->calc_derivative_wrt_group1<false, true, false>(ia, nullptr, &dq_1);
 
     g11 = 2.0 * quat0[1]*dq_1[1];
     g22 = 2.0 * quat0[2]*dq_1[2];