From 6482d08918de7233ba4c27584b2159a2fd58a406 Mon Sep 17 00:00:00 2001
From: Gareth Aneurin Tribello <garethtribello@dhcp-129-242.wifi.sissa.it>
Date: Wed, 5 Feb 2025 13:44:05 +0100
Subject: [PATCH] Added some helper pacakages that will be needed for next step

---
 src/colvar/MultiColvarTemplate.h | 25 +++++++++++++++++++++++++
 src/core/ParallelTaskManager.h   | 14 ++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/colvar/MultiColvarTemplate.h b/src/colvar/MultiColvarTemplate.h
index f24a6ccf43..f58bfaa0bb 100644
--- a/src/colvar/MultiColvarTemplate.h
+++ b/src/colvar/MultiColvarTemplate.h
@@ -62,6 +62,7 @@ class MultiColvarTemplate : public ActionWithVector {
   unsigned getNumberOfAtomsPerTask() const override ;
   void addValueWithDerivatives( const std::vector<unsigned>& shape=std::vector<unsigned>() ) override ;
   void addComponentWithDerivatives( const std::string& name, const std::vector<unsigned>& shape=std::vector<unsigned>() ) override ;
+  void getInputData( std::vector<double>& inputdata ) const override ;
   void performTask( const unsigned&, MultiValue& ) const override ;
   void calculate() override;
   static void performTask( const unsigned& m, const std::vector<std::size_t>& der_indices, const bool noderiv, const bool haspbc, const Pbc& pbc, MultiValue& myvals );
@@ -124,6 +125,13 @@ MultiColvarTemplate<T>::MultiColvarTemplate(const ActionOptions&ao):
 
   // Setup the values
   mode = T::getModeAndSetupValues( this );
+  // This sets up an array in the parallel task manager to hold all the indices
+  std::vector<std::size_t> ind( ablocks.size()*ablocks[0].size() );
+  for(unsigned i=0; i<ablocks[0].size(); ++i) {
+      for(unsigned j=0; j<ablocks.size(); ++j) ind[i*ablocks.size() + j] = ablocks[j][i];
+  }
+  // Sets up the index list in the task manager
+  taskmanager.setupIndexList( ind );
 }
 
 template <class T>
@@ -154,6 +162,23 @@ unsigned MultiColvarTemplate<T>::getNumberOfAtomsPerTask() const {
   return ablocks.size();
 }
 
+template <class T>
+void MultiColvarTemplate<T>::getInputData( std::vector<double>& inputdata ) const {
+  unsigned ntasks = ablocks[0].size(); std::size_t k=0;
+  if( inputdata.size()!=5*ablocks.size()*ntasks ) inputdata.resize( 5*ablocks.size()*ntasks );
+  for(unsigned i=0; i<ntasks; ++i) {
+      for(unsigned j=0; j<ablocks.size(); ++j) { 
+          Vector mypos( getPosition( ablocks[j][i] ) );
+          inputdata[k] = mypos[0];
+          inputdata[k+1] = mypos[1];
+          inputdata[k+2] = mypos[2];
+          inputdata[k+3] = getMass( ablocks[j][i] );
+          inputdata[k+4] = getCharge( ablocks[j][i] );
+          k+=5;
+      }    
+  }
+}
+
 template <class T>
 void MultiColvarTemplate<T>::performTask( const unsigned& task_index, MultiValue& myvals ) const {
   // Retrieve the positions
diff --git a/src/core/ParallelTaskManager.h b/src/core/ParallelTaskManager.h
index 4b35581ec0..b0e5bbe209 100644
--- a/src/core/ParallelTaskManager.h
+++ b/src/core/ParallelTaskManager.h
@@ -42,8 +42,14 @@ class ParallelTaskManager {
   std::vector<double> buffer;
 /// A tempory vector of MultiValue so we can avoid doing lots of resizes
   std::vector<MultiValue> myvals;
+/// This holds indices for creating derivatives
+  std::vector<std::size_t> indices;
+/// This holds all the input data that is required to calculate all values for all tasks
+  std::vector<double> inputdata;
 public:
   ParallelTaskManager(ActionWithVector* av);
+/// Setup an array to hold all the indices that are used for derivatives
+  void setupIndexList( const std::vector<std::size_t>& ind );
 /// This runs all the tasks
   void runAllTasks( const unsigned& natoms=0 );
 /// This runs each of the tasks
@@ -60,6 +66,11 @@ ParallelTaskManager<T>::ParallelTaskManager(ActionWithVector* av):
   if(am) ismatrix=true;
 }
 
+template <class T>
+void ParallelTaskManager<T>::setupIndexList( const std::vector<std::size_t>& ind ) {
+  indices.resize( ind.size() ); for(unsigned i=0; i<ind.size(); ++i) indices[i] = ind[i];
+}
+
 template <class T>
 void ParallelTaskManager<T>::runAllTasks( const unsigned& natoms ) {
   unsigned stride=comm.Get_size();
@@ -86,6 +97,9 @@ void ParallelTaskManager<T>::runAllTasks( const unsigned& natoms ) {
   // Clear buffer
   buffer.assign( buffer.size(), 0.0 );
 
+  // Get all the input data so we can broadcast it to the GPU
+  action->getInputData( inputdata );
+
   #pragma omp parallel num_threads(nt)
   {
     std::vector<double> omp_buffer;