From 636898f386932000bc3fa0458af7f9f292cee4b6 Mon Sep 17 00:00:00 2001
From: Gareth Aneurin Tribello <garethtribello@Gareths-MacBook-Pro.local>
Date: Tue, 24 Sep 2024 15:06:06 +0100
Subject: [PATCH] Reduced the amount of memory that is used in Q6 so you can
 now run 4 threads

---
 src/core/ActionWithVector.cpp         | 22 ++++++++++++--------
 src/core/ActionWithVector.h           |  2 ++
 src/matrixtools/MatrixTimesVector.cpp | 30 ++++++++++++++++++++++++---
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/core/ActionWithVector.cpp b/src/core/ActionWithVector.cpp
index a71d2098b9..8eed7e7771 100644
--- a/src/core/ActionWithVector.cpp
+++ b/src/core/ActionWithVector.cpp
@@ -332,6 +332,16 @@ bool ActionWithVector::checkChainForNonScalarForces() const {
   return false;
 }
 
+void ActionWithVector::getNumberOfForceDerivatives( unsigned& nforces, unsigned& nderiv ) const {
+  nforces=0; unsigned nargs = getNumberOfArguments(); int nmasks = getNumberOfMasks();
+  if( nargs>=nmasks && nmasks>0 ) nargs = nargs - nmasks;
+  if( getNumberOfAtoms()>0 ) nforces += 3*getNumberOfAtoms() + 9;
+  for(unsigned i=0; i<nargs; ++i) {
+    nforces += getPntrToArgument(i)->getNumberOfStoredValues();
+  }
+  nderiv = nforces;
+}
+
 bool ActionWithVector::checkForForces() {
   if( getPntrToComponent(0)->getRank()==0 ) return ActionWithValue::checkForForces();
 
@@ -357,13 +367,8 @@ bool ActionWithVector::checkForForces() {
   if( omp_forces.size()!=nt ) omp_forces.resize(nt);
 
   // Recover the number of derivatives we require (this should be equal to the number of forces)
-  unsigned nderiv=0, nargs = getNumberOfArguments(); int nmasks = getNumberOfMasks();
-  if( nargs>=nmasks && nmasks>0 ) nargs = nargs - nmasks;
-  if( getNumberOfAtoms()>0 ) nderiv += 3*getNumberOfAtoms() + 9;
-  for(unsigned i=0; i<nargs; ++i) {
-    nderiv += getPntrToArgument(i)->getNumberOfStoredValues();
-  }
-  if( forcesForApply.size()!=nderiv ) forcesForApply.resize( nderiv );
+  unsigned nderiv, nforces; getNumberOfForceDerivatives( nforces, nderiv );
+  if( forcesForApply.size()!=nforces ) forcesForApply.resize( nforces );
   // Clear force buffer
   forcesForApply.assign( forcesForApply.size(), 0.0 );
 
@@ -408,8 +413,7 @@ bool ActionWithVector::checkForTaskForce( const unsigned& itask, const Value* my
 }
 
 void ActionWithVector::gatherForcesOnStoredValue( const unsigned& ival, const unsigned& itask, const MultiValue& myvals, std::vector<double>& forces ) const {
-  const Value* myval = getConstPntrToComponent(ival);
-  double fforce = myval->getForce(itask);
+  const Value* myval = getConstPntrToComponent(ival); double fforce = myval->getForce(itask);
   for(unsigned j=0; j<myvals.getNumberActive(ival); ++j) {
     unsigned jder=myvals.getActiveIndex(ival, j); plumed_dbg_assert( jder<forces.size() );
     forces[jder] += fforce*myvals.getDerivative( ival, jder );
diff --git a/src/core/ActionWithVector.h b/src/core/ActionWithVector.h
index 9daefd8724..fa6e63c0c7 100644
--- a/src/core/ActionWithVector.h
+++ b/src/core/ActionWithVector.h
@@ -110,6 +110,8 @@ class ActionWithVector:
   virtual void gatherForces( const unsigned& i, const MultiValue& myvals, std::vector<double>& forces ) const ;
 /// This is to transfer data from the buffer to the final value
   void finishComputations( const std::vector<double>& buf );
+/// Get the number of forces to use
+  virtual void getNumberOfForceDerivatives( unsigned& nforces, unsigned& nderiv ) const ;
 /// Apply the forces on this data
   virtual void apply();
 };
diff --git a/src/matrixtools/MatrixTimesVector.cpp b/src/matrixtools/MatrixTimesVector.cpp
index 462542a8fa..b6e21eb7d3 100644
--- a/src/matrixtools/MatrixTimesVector.cpp
+++ b/src/matrixtools/MatrixTimesVector.cpp
@@ -46,6 +46,8 @@ class MatrixTimesVector : public ActionWithVector {
   void calculate() override ;
   void performTask( const unsigned& task_index, MultiValue& myvals ) const override ;
   int checkTaskIsActive( const unsigned& itask ) const override ;
+  void getNumberOfForceDerivatives( unsigned& nforces, unsigned& nderiv ) const override ;
+  void gatherForces( const unsigned& itask, const MultiValue& myvals, std::vector<double>& forces ) const override ;
 };
 
 PLUMED_REGISTER_ACTION(MatrixTimesVector,"MATRIX_VECTOR_PRODUCT")
@@ -173,7 +175,7 @@ int MatrixTimesVector::checkTaskIsActive( const unsigned& itask ) const {
 
 void MatrixTimesVector::performTask( const unsigned& task_index, MultiValue& myvals ) const {
   if( sumrows ) {
-    unsigned base=0, n=getNumberOfArguments()-1; Value* myvec = getPntrToArgument(n);
+    unsigned n=getNumberOfArguments()-1; Value* myvec = getPntrToArgument(n);
     for(unsigned i=0; i<n; ++i) {
       Value* mymat = getPntrToArgument(i);
       unsigned ncol = mymat->getNumberOfColumns();
@@ -184,11 +186,10 @@ void MatrixTimesVector::performTask( const unsigned& task_index, MultiValue& myv
       // And the derivatives
       if( doNotCalculateDerivatives() ) continue;
 
-      unsigned dloc = base + task_index*ncol;
+      unsigned dloc = task_index*ncol;
       for(unsigned j=0; j<nmat; ++j) {
         myvals.addDerivative( i, dloc + j, 1.0 ); myvals.updateIndex( i, dloc + j );
       }
-      base += mymat->getNumberOfStoredValues();
     }
   } else if( getPntrToArgument(1)->getRank()==1 ) {
     Value* mymat = getPntrToArgument(0);
@@ -238,5 +239,28 @@ void MatrixTimesVector::performTask( const unsigned& task_index, MultiValue& myv
   }
 }
 
+void MatrixTimesVector::getNumberOfForceDerivatives( unsigned& nforces, unsigned& nderiv ) const {
+  ActionWithVector::getNumberOfForceDerivatives( nforces, nderiv );
+  if( sumrows ) nderiv = getPntrToArgument(0)->getNumberOfStoredValues() + getPntrToArgument(getNumberOfArguments()-1)->getNumberOfStoredValues();
+}
+
+void MatrixTimesVector::gatherForces( const unsigned& itask, const MultiValue& myvals, std::vector<double>& forces ) const {
+  if( !sumrows ) { ActionWithVector::gatherForces( itask, myvals, forces ); return; }
+  if( checkComponentsForForce() ) {
+    unsigned base = 0;
+    for(unsigned ival=0; ival<getNumberOfComponents(); ++ival) {
+      const Value* myval=getConstPntrToComponent(ival);
+      if( myval->forcesWereAdded() ) {
+        double fforce = myval->getForce(itask);
+        for(unsigned j=0; j<myvals.getNumberActive(ival); ++j) {
+          unsigned jder=myvals.getActiveIndex(ival, j); plumed_dbg_assert( jder<forces.size() );
+          forces[base+jder] += fforce*myvals.getDerivative( ival, jder );
+        }
+      }
+      base += getPntrToArgument(ival)->getNumberOfStoredValues();
+    }
+  }
+}
+
 }
 }