Line data    Source code

       1             : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
       2             :    Copyright (c) 2011-2020 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_gridtools_FunctionOfGrid_h
      23             : #define __PLUMED_gridtools_FunctionOfGrid_h
      24             : 
      25             : #include "ActionWithGrid.h"
      26             : #include "core/ParallelTaskManager.h"
      27             : #include "function/FunctionSetup.h"
      28             : #include "function/FunctionOfVector.h"
      29             : #include "function/Custom.h"
      30             : #include "EvaluateGridFunction.h"
      31             : 
      32             : namespace PLMD {
      33             : namespace gridtools {
      34             : 
      35             : template <class T>
      36             : class FunctionOfGrid : public ActionWithGrid {
      37             : public:
      38             :   using input_type = function::FunctionData<T>;
      39             :   using PTM = ParallelTaskManager<FunctionOfGrid<T>>;
      40             :   typedef typename PTM::ParallelActionsInput ParallelActionsInput;
      41             :   typedef typename PTM::ParallelActionsOutput ParallelActionsOutput;
      42             : private:
      43             : /// The parallel task manager
      44             :   PTM taskmanager;
      45             : //// Ensures we setup on first step
      46             :   bool firststep;
      47             : /// The number of scalars that appear in the input for this calculation
      48             :   unsigned nscalars;
      49             : /// Set equal to one if we are doing EvaluateGridFunction
      50             :   unsigned argstart;
      51             : public:
      52             :   static void registerKeywords(Keywords&);
      53             :   explicit FunctionOfGrid(const ActionOptions&);
      54             : /// Get the number of derivatives for this action
      55             :   unsigned getNumberOfDerivatives() override ;
      56             : /// Get the label to write in the graph
      57             :   std::string writeInGraph() const override ;
      58             : /// Get the underlying names
      59             :   std::vector<std::string> getGridCoordinateNames() const override ;
      60             : /// Get the underlying grid coordinates object
      61             :   const GridCoordinatesObject& getGridCoordinatesObject() const override ;
      62             : /// Get the input data for doing the parallel calculation
      63             :   void getArgumentInputData( std::vector<double>& inputdata ) const override ;
      64             : /// Do the calculation
      65             :   void calculate() override ;
      66             : // Calculate the value of the function at a grid point
      67             :   static void performTask( std::size_t task_index,
      68             :                            const function::FunctionData<T>& actiondata,
      69             :                            ParallelActionsInput& input,
      70             :                            ParallelActionsOutput& output );
      71             : /// Add the forces
      72             : //  void apply() override;
      73             :   void applyNonZeroRankForces( std::vector<double>& outforces ) override ;
      74             : /// This just returns one
      75             :   static int getNumberOfValuesPerTask( std::size_t task_index,
      76             :                                        const function::FunctionData<T>& actiondata );
      77             : /// Get the indexes where the forces are added
      78             :   static void getForceIndices( std::size_t task_index,
      79             :                                std::size_t colno,
      80             :                                std::size_t ntotal_force,
      81             :                                const function::FunctionData<T>& actiondata,
      82             :                                const ParallelActionsInput& input,
      83             :                                ForceIndexHolder force_indices );
      84             : };
      85             : 
      86             : template <class T>
      87         339 : void FunctionOfGrid<T>::registerKeywords(Keywords& keys ) {
      88         339 :   ActionWithGrid::registerKeywords(keys);
      89         339 :   std::string name = keys.getDisplayName();
      90         339 :   std::size_t und=name.find("_GRID");
      91         678 :   keys.setDisplayName( name.substr(0,und) );
      92         339 :   keys.reserve("compulsory","PERIODIC","if the output of your function is periodic then you should specify the periodicity of the function.  If the output is not periodic you must state this using PERIODIC=NO");
      93         339 :   T tfunc;
      94         339 :   T::registerKeywords( keys );
      95         339 :   if( typeid(tfunc)==typeid(function::Custom()) ) {
      96           0 :     keys.add("hidden","NO_ACTION_LOG","suppresses printing from action on the log");
      97             :   }
      98         678 :   if( keys.outputComponentExists(".#!value") ) {
      99        1017 :     keys.setValueDescription("grid","the grid obtained by doing an element-wise application of " + keys.getOutputComponentDescription(".#!value") + " to the input grid");
     100         678 :     keys.addInputKeyword("compulsory","ARG","scalar/grid","the labels of the scalars and functions on a grid that we are using to compute the required function");
     101             :   }
     102         339 :   PTM::registerKeywords( keys );
     103         678 : }
     104             : 
     105             : template <class T>
     106         165 : FunctionOfGrid<T>::FunctionOfGrid(const ActionOptions&ao):
     107             :   Action(ao),
     108             :   ActionWithGrid(ao),
     109         165 :   taskmanager(this),
     110         165 :   firststep(true),
     111         165 :   nscalars(0),
     112         165 :   argstart(0) {
     113         165 :   if( getNumberOfArguments()==0 ) {
     114           0 :     error("found no arguments");
     115             :   }
     116         165 :   if( typeid(taskmanager.getActionInput().f)==typeid(EvaluateGridFunction) ) {
     117           0 :     argstart=1;
     118             :   }
     119             :   // This will require a fix
     120         165 :   if( getPntrToArgument(argstart)->getRank()==0 || !getPntrToArgument(argstart)->hasDerivatives() ) {
     121           0 :     error("first input to this action must be a grid");
     122             :   }
     123             :   // Get the shape of the input grid
     124         165 :   std::vector<std::size_t> shape( getPntrToArgument(argstart)->getShape() );
     125         310 :   for(unsigned i=argstart+1; i<getNumberOfArguments(); ++i ) {
     126         145 :     if( getPntrToArgument(i)->getRank()==0 ) {
     127          89 :       nscalars++;
     128          89 :       continue;
     129          56 :     } else if( getPntrToArgument(i)->hasDerivatives() ) {
     130          56 :       std::vector<std::size_t> s( getPntrToArgument(i)->getShape() );
     131          56 :       if( s.size()!=shape.size() ) {
     132           0 :         error("mismatch between dimensionalities of input grids");
     133          56 :       } else if( nscalars>0 ) {
     134           0 :         error("scalars should be specified in argument list after all functions on grid");
     135             :       }
     136          56 :     } else {
     137           0 :       error("input arguments should be functions on grid or scalars");
     138             :     }
     139             :   }
     140         165 :   auto & myfunc = taskmanager.getActionInput();
     141         165 :   myfunc.argstart = argstart;
     142         165 :   myfunc.nscalars = nscalars;
     143             :   // Create the values for this grid
     144         165 :   function::FunctionData<T>::setup( myfunc.f, keywords.getOutputComponents(), shape, true, this  );
     145             :   // Setup the task manager
     146         165 :   taskmanager.setupParallelTaskManager( getNumberOfArguments()-argstart, nscalars );
     147         165 : }
     148             : 
     149             : template <class T>
     150             : std::string FunctionOfGrid<T>::writeInGraph() const {
     151             :   std::size_t und = getName().find_last_of("_");
     152             :   return getName().substr(0,und);
     153             : }
     154             : 
     155             : template <class T>
     156      314740 : const GridCoordinatesObject& FunctionOfGrid<T>::getGridCoordinatesObject() const {
     157      314740 :   ActionWithGrid* ag=ActionWithGrid::getInputActionWithGrid( getPntrToArgument(0)->getPntrToAction() );
     158      314740 :   plumed_assert( ag );
     159      314740 :   return ag->getGridCoordinatesObject();
     160             : }
     161             : 
     162             : template <class T>
     163         241 : std::vector<std::string> FunctionOfGrid<T>::getGridCoordinateNames() const {
     164         241 :   ActionWithGrid* ag=ActionWithGrid::getInputActionWithGrid( getPntrToArgument(0)->getPntrToAction() );
     165         241 :   plumed_assert( ag );
     166         241 :   return ag->getGridCoordinateNames();
     167             : }
     168             : 
     169             : template <class T>
     170        3805 : unsigned FunctionOfGrid<T>::getNumberOfDerivatives() {
     171        3805 :   return getGridCoordinatesObject().getDimension();
     172             : }
     173             : 
     174             : template <class T>
     175        2762 : void FunctionOfGrid<T>::calculate() {
     176        2762 :   if( firststep ) {
     177         162 :     const GridCoordinatesObject& mygrid = getGridCoordinatesObject();
     178         162 :     unsigned npoints = getPntrToArgument(0)->getNumberOfValues();
     179         162 :     if( mygrid.getGridType()=="flat" ) {
     180         144 :       std::vector<std::size_t> shape( mygrid.getNbin(true) );
     181         317 :       for(unsigned i=1; i<getNumberOfArguments(); ++i ) {
     182         124 :         if( getPntrToArgument(i)->getRank()==0 ) {
     183          75 :           continue;
     184             :         }
     185          49 :         std::vector<std::size_t> s( getPntrToArgument(i)->getShape() );
     186         109 :         for(unsigned j=0; j<shape.size(); ++j) {
     187          60 :           if( shape[j]!=s[j] ) {
     188           0 :             error("mismatch between sizes of input grids");
     189             :           }
     190             :         }
     191             :       }
     192         288 :       for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     193         144 :         if( getPntrToComponent(i)->getRank()>0 ) {
     194         144 :           getPntrToComponent(i)->setShape(shape);
     195             :         }
     196             :       }
     197         144 :     }
     198             :     // This resizes the scalars
     199         324 :     for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     200         162 :       if( getPntrToComponent(i)->getRank()==0 ) {
     201           0 :         getPntrToComponent(i)->resizeDerivatives( npoints );
     202             :       }
     203             :     }
     204         162 :     taskmanager.setupParallelTaskManager( getNumberOfArguments()-argstart, nscalars );
     205         162 :     firststep=false;
     206             :   }
     207        2762 :   taskmanager.runAllTasks();
     208        2762 : }
     209             : 
     210             : template <class T>
     211        2762 : void FunctionOfGrid<T>::getArgumentInputData( std::vector<double>& inputdata ) const {
     212        2762 :   unsigned nargs = getNumberOfArguments();
     213             : 
     214             :   std::size_t ntasks = 0, ngder = 0;
     215        2762 :   for(unsigned i=argstart; i<nargs; ++i) {
     216        2762 :     if( getPntrToArgument(i)->getRank()>0 ) {
     217        2762 :       ntasks = getPntrToArgument(i)->getNumberOfStoredValues();
     218        2762 :       ngder = getPntrToArgument(i)->getNumberOfDerivatives();
     219        2762 :       break;
     220             :     }
     221             :   }
     222             : 
     223        2762 :   std::size_t ndata = static_cast<std::size_t>(nargs-argstart)*ntasks*(1+ngder);
     224        2762 :   if( inputdata.size()!=ndata ) {
     225         162 :     inputdata.resize( ndata );
     226             :   }
     227             : 
     228        8242 :   for(unsigned j=argstart; j<nargs; ++j) {
     229             :     const Value* myarg =  getPntrToArgument(j);
     230        5480 :     if( myarg->getRank()==0 ) {
     231         508 :       double val = myarg->get();
     232      553205 :       for(unsigned i=0; i<ntasks; ++i) {
     233      552697 :         inputdata[(nargs-argstart)*(1+ngder)*i + j-argstart] = val;
     234             :       }
     235             :     } else {
     236     1104112 :       for(unsigned i=0; i<ntasks; ++i) {
     237     1099140 :         inputdata[(nargs-argstart)*(1+ngder)*i + j-argstart] = myarg->get(i);
     238     3504781 :         for(unsigned k=0; k<ngder; ++k) {
     239     2405641 :           inputdata[(nargs-argstart)*(1+ngder)*i + (nargs-argstart) + (j-argstart)*ngder + k] = myarg->getGridDerivative( i, k );
     240             :         }
     241             :       }
     242             :     }
     243             :   }
     244        2762 : }
     245             : 
     246             : template <class T>
     247      971687 : void FunctionOfGrid<T>::performTask( std::size_t task_index,
     248             :                                      const function::FunctionData<T>& actiondata,
     249             :                                      ParallelActionsInput& input,
     250             :                                      ParallelActionsOutput& output ) {
     251             : 
     252      971687 :   std::size_t rank=input.ranks[actiondata.argstart];
     253      971687 :   std::size_t spacing = (input.nargs-actiondata.argstart)*(1+rank);
     254      971687 :   auto funcout = function::FunctionOutput::create( input.ncomponents,
     255             :                  output.values.data(),
     256             :                  input.nargs-actiondata.argstart,
     257             :                  output.derivatives.data() );
     258      971687 :   T::calc( actiondata.f,
     259      971687 :            input.noderiv,
     260      971687 :            View<const double>( input.inputdata + task_index*spacing,
     261             :                                input.nargs-actiondata.argstart ),
     262             :            funcout );
     263             : 
     264     2638401 :   for(unsigned j=actiondata.argstart; j<input.nargs; ++j) {
     265     1666714 :     double df = output.derivatives[j-actiondata.argstart];
     266     1666714 :     View<const double> inders( input.inputdata + task_index*spacing + (input.nargs-actiondata.argstart) + (j-actiondata.argstart)*rank, rank );
     267     5228172 :     for(unsigned k=0; k<input.ranks[actiondata.argstart]; ++k) {
     268     3561458 :       output.values[1+k] += df*inders[k];
     269             :     }
     270             :   }
     271      971687 : }
     272             : 
     273             : template <class T>
     274         638 : void FunctionOfGrid<T>::applyNonZeroRankForces( std::vector<double>& outforces ) {
     275         638 :   taskmanager.applyForces( outforces );
     276         638 : }
     277             : 
     278             : template <class T>
     279        9218 : int FunctionOfGrid<T>::getNumberOfValuesPerTask( std::size_t task_index,
     280             :     const function::FunctionData<T>& actiondata ) {
     281        9218 :   return 1;
     282             : }
     283             : 
     284             : template <class T>
     285        5967 : void FunctionOfGrid<T>::getForceIndices( std::size_t task_index,
     286             :     std::size_t colno,
     287             :     std::size_t ntotal_force,
     288             :     const function::FunctionData<T>& actiondata,
     289             :     const ParallelActionsInput& input,
     290             :     ForceIndexHolder force_indices ) {
     291             :   // The force indices are found in the same way as in FunctionOfVector so we reuse that function here
     292        5967 :   function::FunctionOfVector<T>::getForceIndices( task_index, colno, ntotal_force, actiondata, input, force_indices );
     293        5967 : }
     294             : 
     295             : }
     296             : }
     297             : #endif