refactor: unify the exertion of fitting group forces

This commit introduces the group_force_object that unifies the exertion
of atomic forces in an atom group. If a colvarcomp requires to apply
forces individually on atoms in an atom group, now it can do
```
auto ag_force = atoms->get_group_force_object();
for (size_t ia = 0; ia < atoms->size(); ia++) {
  const f_ia = ...;
  ag_force.set_atom_force(ia, f_ia);
}
```
If the optimal fitting is enabled for the atom group, group_force_object
will "intercept" the forces applied to the main group, and then calculate
the forces required on the fitting group by calling calc_fit_forces
internally, and then apply the forces on destruction.

src/colvaratoms.cpp

@@ -1290,7 +1290,7 @@ void cvm::atom_group::calc_fit_forces_impl(
 }
 
 
-std::vector<cvm::rvector> cvm::atom_group::calc_fit_forces(
+void cvm::atom_group::calc_fit_forces(
   const std::vector<cvm::rvector>& forces_on_main_group,
   std::vector<cvm::rvector>& forces_on_fitting_group) const {
   if (cvm::debug())
@@ -1307,7 +1307,6 @@ std::vector<cvm::rvector> cvm::atom_group::calc_fit_forces(
 
   if (cvm::debug())
     cvm::log("Done calculating fit forces.\n");
-  return std::vector<cvm::rvector>();
 }
 
 
@@ -1481,39 +1480,70 @@ void cvm::atom_group::apply_force(cvm::rvector const &force)
     return;
   }
 
-  if (is_enabled(f_ag_rotate)) {
-    // TODO: What is the best way to avoid repeating the allocation??
-    std::vector<cvm::rvector> forces_on_main_group;
-    const auto rot_inv = rot.inverse().matrix();
-    if (cvm::debug()) {
-      cvm::log("Force on main group:\n");
-    }
-    for (cvm::atom_iter ai = this->begin(); ai != this->end(); ai++) {
-      const auto f = (ai->mass/total_mass) * force;
-      ai->apply_force(rot_inv * f);
-      if (cvm::debug()) {
-        cvm::log(cvm::to_str(rot_inv * f));
-      }
-      forces_on_main_group.push_back(f);
-    }
-    if (cvm::debug()) {
-      cvm::log("Force on fitting group:\n");
+  auto ag_force = get_group_force_object();
+  for (size_t i = 0; i < size(); ++i) {
+    ag_force.set_atom_force(i, atoms[i].mass / total_mass * force);
+  }
+}
+
+cvm::atom_group::group_force_object cvm::atom_group::get_group_force_object() {
+  return cvm::atom_group::group_force_object(this);
+}
+
+cvm::atom_group::group_force_object::group_force_object(cvm::atom_group* ag):
+m_ag(ag), m_group_for_fit(m_ag->fitting_group ? m_ag->fitting_group : m_ag),
+m_has_fitting_force(m_ag->is_enabled(f_ag_center) || m_ag->is_enabled(f_ag_rotate)) {
+  if (m_has_fitting_force) {
+    if (m_ag->group_forces.size() != m_ag->size()) {
+      m_ag->group_forces.assign(m_ag->size(), 0);
+    } else {
+      std::fill(m_ag->group_forces.begin(),
+                m_ag->group_forces.end(), 0);
     }
-    std::vector<cvm::rvector> forces_on_fitting_group;
-    calc_fit_forces(forces_on_main_group, forces_on_fitting_group);
-    atom_group *group_for_fit = fitting_group ? fitting_group : this;
-    if (forces_on_fitting_group.size() == group_for_fit->size()) {
-      for (size_t j = 0; j < group_for_fit->size(); j++) {
-        (*group_for_fit)[j].apply_force(forces_on_fitting_group[j]);
-        if (cvm::debug()) {
-          cvm::log(cvm::to_str(forces_on_fitting_group[j]));
-        }
-      }
+    if (m_ag->fitting_group_forces.size() != m_group_for_fit->size()) {
+      m_ag->fitting_group_forces.assign(m_group_for_fit->size(), 0);
+    } else {
+      std::fill(m_ag->fitting_group_forces.begin(),
+                m_ag->fitting_group_forces.end(), 0);
     }
+  }
+}
+
+cvm::atom_group::group_force_object::~group_force_object() {
+  if (m_has_fitting_force) {
+    apply_force_with_fitting_group();
+  }
+}
+
+void cvm::atom_group::group_force_object::set_atom_force(size_t i, const cvm::rvector& force) {
+  if (m_has_fitting_force) {
+    m_ag->group_forces[i] = force;
   } else {
+    // Apply the force directly if we don't use fitting
+    (*m_ag)[i].apply_force(force);
+  }
+}
 
-    for (cvm::atom_iter ai = this->begin(); ai != this->end(); ai++) {
-      ai->apply_force((ai->mass/total_mass) * force);
+void cvm::atom_group::group_force_object::apply_force_with_fitting_group() {
+  m_ag->calc_fit_forces(m_ag->group_forces, m_ag->fitting_group_forces);
+  const cvm::rmatrix rot_inv = m_ag->rot.inverse().matrix();
+  if (cvm::debug()) {
+    cvm::log("Applying force on main group " + m_ag->name + ":\n");
+  }
+  for (size_t ia = 0; ia < m_ag->size(); ++ia) {
+    const cvm::rvector f_ia = rot_inv * m_ag->group_forces[ia];
+    (*m_ag)[ia].apply_force(f_ia);
+    if (cvm::debug()) {
+      cvm::log(cvm::to_str(f_ia));
+    }
+  }
+  if (cvm::debug()) {
+    cvm::log("Applying force on the fitting group of main group" + m_ag->name + ":\n");
+  }
+  for (size_t ia = 0; ia < m_group_for_fit->size(); ia++) {
+    (*(m_group_for_fit))[ia].apply_force(m_ag->fitting_group_forces[ia]);
+    if (cvm::debug()) {
+      cvm::log(cvm::to_str(m_ag->fitting_group_forces[ia]));
     }
   }
 }

src/colvaratoms.h

@@ -257,8 +257,25 @@ class colvarmodule::atom_group
   /// \brief Index in the colvarproxy arrays (if the group is scalable)
   int index;
 
+  std::vector<cvm::rvector> group_forces;
+  std::vector<cvm::rvector> fitting_group_forces;
+
 public:
 
+  class group_force_object {
+  public:
+    group_force_object(cvm::atom_group* ag);
+    ~group_force_object();
+    void set_atom_force(size_t i, const cvm::rvector& force);
+  private:
+    cvm::atom_group* m_ag;
+    cvm::atom_group* m_group_for_fit;
+    bool m_has_fitting_force;
+    void apply_force_with_fitting_group();
+  };
+
+  group_force_object get_group_force_object();
+
   inline cvm::atom & operator [] (size_t const i)
   {
     return atoms[i];
@@ -497,7 +514,7 @@ class colvarmodule::atom_group
   /// \brief Calculate the derivatives of the fitting transformation
   void calc_fit_gradients();
 
-  std::vector<cvm::rvector> calc_fit_forces(
+  void calc_fit_forces(
     const std::vector<cvm::rvector>& forces_on_main_group,
     std::vector<cvm::rvector>& forces_on_fitting_group) const;
 

src/colvarcomp.h

@@ -1008,11 +1008,6 @@ class colvar::orientation
   struct rotation_derivative_impl_;
   std::unique_ptr<rotation_derivative_impl_> rot_deriv_impl;
 
-  bool atom_rotated;
-  cvm::atom_group *group_for_fit;
-  std::vector<cvm::rvector> main_group_forces;
-  std::vector<cvm::rvector> fitting_group_forces;
-
 public:
 
   orientation();

src/colvarcomp_distances.cpp

@@ -1403,11 +1403,12 @@ void colvar::cartesian::apply_force(colvarvalue const &force)
   size_t ia, j;
   if (!atoms->noforce) {
     cvm::rvector f;
+    auto ag_force = atoms->get_group_force_object();
     for (ia = 0; ia < atoms->size(); ia++) {
       for (j = 0; j < dim; j++) {
         f[axes[j]] = force.vector1d_value[dim*ia + j];
       }
-      (*atoms)[ia].apply_force(f);
+      ag_force.set_atom_force(ia, f);
     }
   }
 }

src/colvarcomp_rotations.cpp

@@ -22,8 +22,7 @@ struct colvar::orientation::rotation_derivative_impl_: public rotation_derivativ
 };
 
 
-colvar::orientation::orientation():
-  atom_rotated(false), group_for_fit(nullptr)
+colvar::orientation::orientation()
 {
   set_function_type("orientation");
   rot_deriv_impl = std::unique_ptr<rotation_derivative_impl_>(new rotation_derivative_impl_(this));
@@ -44,12 +43,6 @@ int colvar::orientation::init(std::string const &conf)
     return error_code | COLVARS_INPUT_ERROR;
   }
   ref_pos.reserve(atoms->size());
-  main_group_forces.resize(atoms->size());
-  atom_rotated = atoms->is_enabled(f_ag_rotate);
-  if (atom_rotated) {
-    group_for_fit = atoms->fitting_group ? atoms->fitting_group : atoms;
-    fitting_group_forces.resize(group_for_fit->size());
-  }
 
   if (get_keyval(conf, "refPositions", ref_pos, ref_pos)) {
     cvm::log("Using reference positions from input file.\n");
@@ -144,45 +137,14 @@ void colvar::orientation::apply_force(colvarvalue const &force)
   if (!atoms->noforce) {
     rot_deriv_impl->prepare_derivative(rotation_derivative_dldq::use_dq);
     cvm::vector1d<cvm::rvector> dq0_2;
-    cvm::rmatrix ag_rot;
-    if (atom_rotated) {
-      ag_rot = atoms->rot.inverse().matrix();
-    }
-    if (cvm::debug()) {
-      cvm::log("Force on main group:\n");
-    }
+    auto ag_force = atoms->get_group_force_object();
     for (size_t ia = 0; ia < atoms->size(); ia++) {
       rot_deriv_impl->calc_derivative_wrt_group2(ia, nullptr, &dq0_2);
       const auto f_ia = FQ[0] * dq0_2[0] +
                         FQ[1] * dq0_2[1] +
                         FQ[2] * dq0_2[2] +
                         FQ[3] * dq0_2[3];
-      if (atom_rotated) {
-        main_group_forces[ia] = f_ia;
-        (*atoms)[ia].apply_force(ag_rot * f_ia);
-        if (cvm::debug()) {
-          cvm::log(cvm::to_str(ag_rot * f_ia));
-        }
-      } else {
-        (*atoms)[ia].apply_force(f_ia);
-        if (cvm::debug()) {
-          cvm::log(cvm::to_str(f_ia));
-        }
-      }
-    }
-    if (atom_rotated) {
-      atoms->calc_fit_forces(main_group_forces, fitting_group_forces);
-      if (cvm::debug()) {
-        cvm::log("Force on fitting group:\n");
-      }
-      for (size_t ia = 0; ia < group_for_fit->size(); ia++) {
-        (*(group_for_fit))[ia].apply_force(fitting_group_forces[ia]);
-        if (cvm::debug()) {
-          cvm::log(cvm::to_str(fitting_group_forces[ia]));
-        }
-        // Clear the fitting group force
-        fitting_group_forces[ia] = 0;
-      }
+      ag_force.set_atom_force(ia, f_ia);
     }
   }
 }