Line data    Source code

       1             : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
       2             :    Copyright (c) 2011-2023 The plumed team
       3             :    (see the PEOPLE file at the root of the distribution for a list of names)
       4             : 
       5             :    See http://www.plumed.org for more information.
       6             : 
       7             :    This file is part of plumed, version 2.
       8             : 
       9             :    plumed is free software: you can redistribute it and/or modify
      10             :    it under the terms of the GNU Lesser General Public License as published by
      11             :    the Free Software Foundation, either version 3 of the License, or
      12             :    (at your option) any later version.
      13             : 
      14             :    plumed is distributed in the hope that it will be useful,
      15             :    but WITHOUT ANY WARRANTY; without even the implied warranty of
      16             :    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
      17             :    GNU Lesser General Public License for more details.
      18             : 
      19             :    You should have received a copy of the GNU Lesser General Public License
      20             :    along with plumed.  If not, see <http://www.gnu.org/licenses/>.
      21             : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
      22             : #ifndef __PLUMED_matrixtools_OuterProduct_h
      23             : #define __PLUMED_matrixtools_OuterProduct_h
      24             : 
      25             : #include "core/ActionWithMatrix.h"
      26             : #include "core/ParallelTaskManager.h"
      27             : 
      28             : namespace PLMD {
      29             : namespace matrixtools {
      30             : 
      31             : template <class T>
      32           9 : class OuterProductInput {
      33             : public:
      34             :   T funcinput;
      35             :   bool no_thread_gather;
      36             :   bool gatherForceOnColumns;
      37             :   RequiredMatrixElements outmat;
      38             : };
      39             : 
      40             : template <class T>
      41             : class OuterProductBase : public ActionWithMatrix {
      42             : public:
      43             :   using input_type = OuterProductInput<T>;
      44             :   using PTM = ParallelTaskManager<OuterProductBase<T>>;
      45             :   typedef typename PTM::ParallelActionsInput ParallelActionsInput;
      46             :   typedef typename PTM::ParallelActionsOutput ParallelActionsOutput;
      47             : private:
      48             :   bool isproduct;
      49             :   bool secondArgumentIsConstant;
      50             :   PTM taskmanager;
      51             : public:
      52             :   static void registerKeywords( Keywords& keys );
      53             :   explicit OuterProductBase(const ActionOptions&);
      54             :   unsigned getNumberOfDerivatives() override;
      55             :   void prepare() override ;
      56             :   int checkTaskIsActive( const unsigned& itask ) const override ;
      57             :   void calculate() override ;
      58             :   void applyNonZeroRankForces( std::vector<double>& outforces ) override ;
      59             :   static void performTask( std::size_t task_index,
      60             :                            const OuterProductInput<T>& actiondata,
      61             :                            ParallelActionsInput& input,
      62             :                            ParallelActionsOutput& output );
      63             :   static int getNumberOfValuesPerTask( std::size_t task_index,
      64             :                                        const OuterProductInput<T>& actiondata );
      65             :   static void getForceIndices( std::size_t task_index,
      66             :                                std::size_t colno,
      67             :                                std::size_t ntotal_force,
      68             :                                const OuterProductInput<T>& actiondata,
      69             :                                const ParallelActionsInput& input,
      70             :                                ForceIndexHolder force_indices );
      71             : };
      72             : 
      73             : template <class T>
      74         225 : void OuterProductBase<T>::registerKeywords( Keywords& keys ) {
      75         225 :   ActionWithMatrix::registerKeywords(keys);
      76         225 :   T::registerKeywords( keys );
      77         225 :   PTM::registerKeywords( keys );
      78         225 :   keys.add("hidden","NO_ACTION_LOG","suppresses printing from action on the log");
      79         225 : }
      80             : 
      81             : template <class T>
      82         108 : OuterProductBase<T>::OuterProductBase(const ActionOptions&ao):
      83             :   Action(ao),
      84             :   ActionWithMatrix(ao),
      85         108 :   isproduct(false),
      86         108 :   taskmanager(this) {
      87             :   unsigned nargs=getNumberOfArguments();
      88         108 :   if( getNumberOfMasks()>0 ) {
      89          25 :     nargs = nargs - getNumberOfMasks();
      90             :   }
      91         108 :   if( nargs==0 || nargs%2!=0 ) {
      92           1 :     error("should be an even number of arguments to this action, they should all be vectors");
      93             :   }
      94         107 :   std::size_t nvals = nargs / 2;
      95         235 :   for(unsigned i=0; i<nvals; ++i) {
      96         128 :     if( getPntrToArgument(i)->getRank()!=1 || getPntrToArgument(i)->hasDerivatives() ) {
      97           0 :       error("first argument to this action should be a vector");
      98             :     }
      99         128 :     if( getPntrToArgument(0)->getShape()[0]!=getPntrToArgument(i)->getShape()[0] ) {
     100           0 :       error("mismatch between sizes of input vectors");
     101             :     }
     102         128 :     if( getPntrToArgument(nvals+i)->getRank()!=1 || getPntrToArgument(nvals+i)->hasDerivatives() ) {
     103           0 :       error("first argument to this action should be a vector");
     104             :     }
     105         128 :     if( getPntrToArgument(nvals)->getShape()[0]!=getPntrToArgument(nvals+i)->getShape()[0] ) {
     106           0 :       error("mismatch between sizes of input vectors");
     107             :     }
     108             :   }
     109         107 :   if( getNumberOfMasks()==1 ) {
     110          24 :     if( getPntrToArgument(nargs)->getRank()!=2 || getPntrToArgument(nargs)->hasDerivatives() ) {
     111           0 :       error("mask argument should be a matrix");
     112             :     }
     113          63 :     for(unsigned i=0; i<nvals; ++i) {
     114          39 :       if( getPntrToArgument(nargs)->getShape()[0]!=getPntrToArgument(i)->getShape()[0] ) {
     115           0 :         error("mask argument has wrong size");
     116             :       }
     117          39 :       if( getPntrToArgument(nargs)->getShape()[1]!=getPntrToArgument(nvals+i)->getShape()[0] ) {
     118           0 :         error("mask argument has wrong size");
     119             :       }
     120             :     }
     121             :   }
     122             : 
     123         107 :   std::vector<std::size_t> shape(2);
     124         107 :   shape[0]=getPntrToArgument(0)->getShape()[0];
     125         107 :   shape[1]=getPntrToArgument(nvals)->getShape()[0];
     126             : 
     127             :   std::string func;
     128         107 :   if( keywords.exists("FUNC") ) {
     129         100 :     parse("FUNC",func);
     130         100 :     isproduct=(func=="x*y");
     131             :   }
     132          98 :   OuterProductInput<T> actiondata;
     133         107 :   actiondata.gatherForceOnColumns = false;
     134         107 :   actiondata.no_thread_gather = no_thread_gather;
     135         107 :   actiondata.funcinput.setup( shape, func, this );
     136             : 
     137         107 :   if( getNumberOfComponents()==0 ) {
     138         100 :     addValue( shape );
     139         100 :     setNotPeriodic();
     140             :   }
     141         107 :   if( getPntrToArgument(0)->isDerivativeZeroWhenValueIsZero() || getPntrToArgument(nvals)->isDerivativeZeroWhenValueIsZero() ) {
     142         150 :     for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     143          75 :       getPntrToComponent(i)->setDerivativeIsZeroWhenValueIsZero();
     144             :     }
     145             :   }
     146         107 :   secondArgumentIsConstant=true;
     147         149 :   for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     148         107 :     if( !getPntrToArgument(getNumberOfComponents()+i)->isConstant() ) {
     149          65 :       secondArgumentIsConstant=false;
     150          65 :       break;
     151             :     }
     152             :   }
     153         107 :   taskmanager.setActionInput( actiondata );
     154         208 : }
     155             : 
     156             : template <class T>
     157         218 : unsigned OuterProductBase<T>::getNumberOfDerivatives() {
     158         218 :   unsigned nc = getNumberOfComponents();
     159         218 :   return nc*(getPntrToArgument(0)->getNumberOfStoredValues() + getPntrToArgument(nc)->getNumberOfStoredValues());
     160             : }
     161             : 
     162             : template <class T>
     163         284 : void OuterProductBase<T>::prepare() {
     164         284 :   ActionWithVector::prepare();
     165         284 :   Value* myval=getPntrToComponent(0);
     166         284 :   if( myval->getShape()[0]==getPntrToArgument(0)->getShape()[0] && myval->getShape()[1]==getPntrToArgument(getNumberOfComponents())->getShape()[0] ) {
     167         267 :     return;
     168             :   }
     169          17 :   std::vector<std::size_t> shape(2);
     170          17 :   shape[0] = getPntrToArgument(0)->getShape()[0];
     171          17 :   shape[1] = getPntrToArgument(getNumberOfComponents())->getShape()[0];
     172          34 :   for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     173          17 :     getPntrToComponent(i)->setShape( shape );
     174             :   }
     175          17 : }
     176             : 
     177             : template <class T>
     178      367827 : int OuterProductBase<T>::checkTaskIsActive( const unsigned& itask ) const {
     179      367827 :   if( getNumberOfMasks()>0 || !isproduct ) {
     180      108333 :     return ActionWithVector::checkTaskIsActive( itask );
     181             :   }
     182      259494 :   if( fabs( getPntrToArgument(0)->get(itask))>epsilon ) {
     183      173238 :     return 1;
     184             :   }
     185             :   return -1;
     186             : }
     187             : 
     188             : template <class T>
     189         282 : void OuterProductBase<T>::calculate() {
     190         282 :   updateBookeepingArrays( taskmanager.getActionInput().outmat );
     191         282 :   if( no_thread_gather ) {
     192          14 :     taskmanager.setupParallelTaskManager( getNumberOfComponents(), 0 );
     193          14 :     taskmanager.setWorkspaceSize( (2+2*getNumberOfComponents())*getNumberOfComponents() );
     194             :   } else {
     195         268 :     taskmanager.setupParallelTaskManager( 2*getNumberOfComponents(), getNumberOfComponents()*getPntrToComponent(0)->getShape()[1] );
     196         268 :     taskmanager.setWorkspaceSize( 2*getNumberOfComponents() );
     197             :   }
     198         282 :   taskmanager.runAllTasks();
     199         282 : }
     200             : 
     201             : template <class T>
     202      186354 : void OuterProductBase<T>::performTask( std::size_t task_index,
     203             :                                        const OuterProductInput<T>& actiondata,
     204             :                                        ParallelActionsInput& input,
     205             :                                        ParallelActionsOutput& output ) {
     206      186354 :   auto args = output.buffer.subview(0, 2*input.ncomponents);
     207      186354 :   if( actiondata.no_thread_gather ) {
     208         712 :     unsigned first_arg_base = 0, second_arg_base = input.ncomponents, ind = 0;
     209         712 :     if( actiondata.gatherForceOnColumns ) {
     210             :       first_arg_base = input.ncomponents;
     211             :       second_arg_base = 0;
     212             :       ind = input.ncomponents;
     213             :     }
     214        1667 :     for(unsigned i=0; i<input.ncomponents; ++i) {
     215         955 :       args[first_arg_base+i] = input.inputdata[input.argstarts[first_arg_base + i] + task_index];
     216             :     }
     217         712 :     unsigned fstart = task_index*(1+actiondata.outmat.ncols);
     218         712 :     unsigned nelements = actiondata.outmat[fstart];
     219       18767 :     for(unsigned i=0; i<nelements; ++i) {
     220       18055 :       std::size_t argpos = actiondata.outmat[fstart+1+i];
     221       36596 :       for(unsigned j=0; j<input.ncomponents; ++j) {
     222       18541 :         args[second_arg_base+j] = input.inputdata[input.argstarts[second_arg_base+j] + argpos];
     223             :       }
     224       18055 :       MatrixElementOutput matout( input.ncomponents,
     225             :                                   2*input.ncomponents,
     226       18055 :                                   output.values.data()+i*input.ncomponents,
     227             :                                   output.buffer.data()+2*input.ncomponents );
     228       18055 :       T::calculate( input.noderiv, actiondata.funcinput, {args.data(),args.size()}, matout );
     229       18055 :       if( !input.noderiv ) {
     230       23288 :         for(unsigned j=0; j<input.ncomponents; ++j) {
     231       24908 :           for(unsigned k=0; k<input.ncomponents; ++k) {
     232       13102 :             output.derivatives[i*input.ncomponents*input.ncomponents+input.ncomponents*j+k] = matout.derivs[j][ind+k];
     233             :           }
     234             :         }
     235             :       }
     236             :     }
     237             :   } else {
     238      382216 :     for(unsigned i=0; i<input.ncomponents; ++i) {
     239      196574 :       args[i] = input.inputdata[input.argstarts[i] + task_index];
     240             :     }
     241      185642 :     unsigned fstart = task_index*(1+actiondata.outmat.ncols);
     242      185642 :     unsigned nelements = actiondata.outmat[fstart];
     243    10554348 :     for(unsigned i=0; i<nelements; ++i) {
     244    10368706 :       std::size_t argpos = actiondata.outmat[fstart+1+i];
     245    22479755 :       for(unsigned j=0; j<input.ncomponents; ++j) {
     246    12111049 :         args[input.ncomponents+j] = input.inputdata[input.argstarts[input.ncomponents+j] + argpos];
     247             :       }
     248    10368706 :       MatrixElementOutput matout( input.ncomponents,
     249             :                                   2*input.ncomponents,
     250    10368706 :                                   output.values.data()+i*input.ncomponents,
     251    10368706 :                                   output.derivatives.data() + 2*i*input.ncomponents*input.ncomponents );
     252    10368706 :       T::calculate( input.noderiv, actiondata.funcinput, {args.data(),args.size()}, matout );
     253             :     }
     254             :   }
     255      186354 : }
     256             : 
     257             : template <class T>
     258          95 : void OuterProductBase<T>::applyNonZeroRankForces( std::vector<double>& outforces ) {
     259          95 :   taskmanager.applyForces( outforces );
     260          95 :   if( no_thread_gather && !secondArgumentIsConstant ) {
     261          13 :     getColumnBookeepingArrays( taskmanager.getActionInput().outmat );
     262          13 :     taskmanager.getActionInput().gatherForceOnColumns = gatherForceOnColumns = true;
     263          13 :     taskmanager.setNForceScalars( getNumberOfComponents()*maxcolsize );
     264          13 :     taskmanager.applyForces( outforces, false );
     265          13 :     taskmanager.setNForceScalars( getNumberOfComponents()*getPntrToComponent(0)->getNumberOfColumns() );
     266          13 :     taskmanager.getActionInput().gatherForceOnColumns = gatherForceOnColumns = false;
     267             :   }
     268          95 : }
     269             : 
     270             : template <class T>
     271        9370 : int OuterProductBase<T>::getNumberOfValuesPerTask( std::size_t task_index,
     272             :     const OuterProductInput<T>& actiondata ) {
     273        9370 :   unsigned fstart = task_index*(1+actiondata.outmat.ncols);
     274        9370 :   return actiondata.outmat[fstart];
     275             : }
     276             : 
     277             : template <class T>
     278      556576 : void OuterProductBase<T>::getForceIndices( std::size_t task_index,
     279             :     std::size_t colno,
     280             :     std::size_t ntotal_force,
     281             :     const OuterProductInput<T>& actiondata,
     282             :     const ParallelActionsInput& input,
     283             :     ForceIndexHolder force_indices ) {
     284      556576 :   if( actiondata.no_thread_gather ) {
     285             :     unsigned argbase = 0;
     286       11358 :     if( actiondata.gatherForceOnColumns ) {
     287        4847 :       argbase = input.ncomponents;
     288             :     }
     289       23040 :     for(unsigned j=0; j<input.ncomponents; ++j) {
     290       24660 :       for(unsigned k=0; k<input.ncomponents; ++k) {
     291       12978 :         force_indices.indices[j][k] = input.argstarts[argbase+k] + task_index;
     292             :       }
     293       11682 :       force_indices.threadsafe_derivatives_end[j] = input.ncomponents;
     294       11682 :       force_indices.tot_indices[j] = input.ncomponents;
     295             :     }
     296             :   } else {
     297      545218 :     unsigned fstart = task_index*(1+actiondata.outmat.ncols);
     298     1676906 :     for(unsigned j=0; j<input.ncomponents; ++j) {
     299     4609256 :       for(unsigned k=0; k<input.ncomponents; ++k) {
     300     3477568 :         force_indices.indices[j][k] = input.argstarts[k] + task_index;
     301     3477568 :         force_indices.indices[j][input.ncomponents+k] = input.argstarts[input.ncomponents+k] + actiondata.outmat[fstart+1+colno];
     302             :       }
     303     1131688 :       force_indices.threadsafe_derivatives_end[j] = input.ncomponents;
     304     1131688 :       force_indices.tot_indices[j] = 2*input.ncomponents;
     305             :     }
     306             :   }
     307      556576 : }
     308             : 
     309             : }
     310             : }
     311             : #endif