Updates to nblib listed forces kernels
[alexxy/gromacs.git] / api / nblib / listed_forces / dataflow.hpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2019, by the GROMACS development team, led by
5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6  * and including many others, as listed in the AUTHORS file in the
7  * top-level source directory and at http://www.gromacs.org.
8  *
9  * GROMACS is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public License
11  * as published by the Free Software Foundation; either version 2.1
12  * of the License, or (at your option) any later version.
13  *
14  * GROMACS 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 GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with GROMACS; if not, see
21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
23  *
24  * If you want to redistribute modifications to GROMACS, please
25  * consider that scientific software is very special. Version
26  * control is crucial - bugs must be traceable. We will be happy to
27  * consider code for inclusion in the official distribution, but
28  * derived work must not be called official GROMACS. Details are found
29  * in the README & COPYING files - if they are missing, get the
30  * official version at http://www.gromacs.org.
31  *
32  * To help us fund GROMACS development, we humbly ask that you cite
33  * the research papers on the package. Check out http://www.gromacs.org.
34  */
35 /*! \internal \file
36  * \brief
37  * Implements the data flow from ListedInteractionData and coordinates
38  * down to the individual interaction type kernels
39  *
40  * Intended for internal use inside ListedCalculator only.
41  *
42  * \author Victor Holanda <victor.holanda@cscs.ch>
43  * \author Joe Jordan <ejjordan@kth.se>
44  * \author Prashanth Kanduri <kanduri@cscs.ch>
45  * \author Sebastian Keller <keller@cscs.ch>
46  * \author Artem Zhmurov <zhmurov@gmail.com>
47  */
48
49 #ifndef NBLIB_LISTEDFORCES_DATAFLOW_HPP
50 #define NBLIB_LISTEDFORCES_DATAFLOW_HPP
51
52 #include <tuple>
53 #include <vector>
54
55 #include "nblib/listed_forces/traits.h"
56 #include "nblib/listed_forces/kernels.hpp"
57 #include "nblib/util/util.hpp"
58 #include "nblib/pbc.hpp"
59 #include "nblib/vector.h"
60 #include "gromacs/utility/arrayref.h"
61
62 #define NBLIB_ALWAYS_INLINE __attribute((always_inline))
63
64 namespace nblib
65 {
66
67 //! \brief returns the address of an element in the shiftForces buffer
68 template<class ShiftForce>
69 inline ShiftForce* accessShiftForces(int shiftIndex, gmx::ArrayRef<ShiftForce> shiftForces)
70 {
71     return shiftForces.data() + shiftIndex;
72 }
73
74 //! \brief dummy op in case shift forces are not computed (will be removed by the compiler)
75 inline std::nullptr_t accessShiftForces([[maybe_unused]] int shiftIndex,
76                                         [[maybe_unused]] gmx::ArrayRef<std::nullptr_t> shiftForces)
77 {
78     return nullptr;
79 }
80
81 template<class TwoCenterType, class BasicVector, class ShiftForce>
82 inline NBLIB_ALWAYS_INLINE
83 auto computeTwoCenter(const TwoCenterType& parameters, const BasicVector& dx, BasicVector* fi, BasicVector* fj,
84                       ShiftForce* sh_f, ShiftForce* sh_fc)
85 {
86     using ValueType = BasicVectorValueType_t<BasicVector>;
87
88     ValueType dr2 = dot(dx, dx);
89     ValueType dr  = std::sqrt(dr2);
90     auto [force, energy] = bondKernel(dr, parameters);
91
92     // avoid division by 0
93     if (dr2 != 0.0)
94     {
95         force /= dr;
96         spreadTwoCenterForces(force, dx, fi, fj, sh_f, sh_fc);
97     }
98
99     return energy;
100 }
101
102 /*! \brief calculate two-center interactions
103  *
104  * \tparam Force buffer type
105  * \tparam TwoCenterType The bond type to compute; used for type deduction
106  * \tparam Cartesian vector type
107  * \tparam PBC type
108  * \param[in] index The atom and parameter indices used for computing the interaction
109  * \param[in] bondInstances The full type-specific interaction list
110  * \param[in] x The coordinates
111  * \param[in/out] forces The forces
112  * \param[in] pbc Object used for computing distances accounting for PBC's
113  * \return Computed kernel energies
114  */
115 template <class Buffer, class TwoCenterType, class BasicVector, class ShiftForce, class Pbc,
116           std::enable_if_t<Contains<TwoCenterType, SupportedTwoCenterTypes>{}>* = nullptr>
117 inline NBLIB_ALWAYS_INLINE
118 auto dispatchInteraction(InteractionIndex<TwoCenterType> index,
119                          gmx::ArrayRef<const TwoCenterType> bondInstances,
120                          gmx::ArrayRef<const BasicVector> x,
121                          Buffer* forces,
122                          gmx::ArrayRef<ShiftForce> shiftForces,
123                          const Pbc& pbc)
124 {
125     KernelEnergy<BasicVectorValueType_t<BasicVector>> energy;
126
127     int i = std::get<0>(index);
128     int j = std::get<1>(index);
129     BasicVector   xi   = x[i];
130     BasicVector   xj   = x[j];
131     TwoCenterType bond = bondInstances[std::get<2>(index)];
132
133     BasicVector dx;
134     // calculate xi - xj modulo pbc
135     int sIdx = pbc.dxAiuc(xi, xj, dx);
136
137     ShiftForce* sh_f  = accessShiftForces(sIdx, shiftForces);
138     ShiftForce* sh_fc = accessShiftForces(gmx::c_centralShiftIndex, shiftForces);
139
140     energy.carrier() = computeTwoCenter(bond, dx, &(*forces)[i], &(*forces)[j], sh_f, sh_fc);
141     return energy;
142 }
143
144
145 template<class ThreeCenterType, class BasicVector, class ShiftForce>
146 inline NBLIB_ALWAYS_INLINE
147 std::enable_if_t<HasTwoCenterAggregate<ThreeCenterType>::value, BasicVectorValueType_t<BasicVector>>
148 addTwoCenterAggregate(const ThreeCenterType& parameters, const BasicVector& rij, const BasicVector& rkj,
149                       BasicVector* fi, BasicVector* fj, BasicVector* fk,
150                       ShiftForce* shf_ij, ShiftForce* shf_kj, ShiftForce* shf_c)
151 {
152     if (parameters.manifest == ThreeCenterType::Cargo::ij)
153     {
154         // i-j bond
155         return computeTwoCenter(parameters.twoCenter(), rij, fi, fj, shf_ij, shf_c);
156     }
157     if (parameters.manifest == ThreeCenterType::Cargo::jk)
158     {
159         // j-k bond
160         return computeTwoCenter(parameters.twoCenter(), rkj, fk, fj, shf_kj, shf_c);
161     }
162
163     // aggregate is empty
164     return 0.0;
165 };
166
167 template<class ThreeCenterType, class BasicVector, class ShiftForce>
168 inline NBLIB_ALWAYS_INLINE
169 std::enable_if_t<!HasTwoCenterAggregate<ThreeCenterType>::value, BasicVectorValueType_t<BasicVector>>
170 addTwoCenterAggregate([[maybe_unused]] const ThreeCenterType& parameters,
171                       [[maybe_unused]] const BasicVector& rij,
172                       [[maybe_unused]] const BasicVector& rkj,
173                       [[maybe_unused]] BasicVector* fi,
174                       [[maybe_unused]] BasicVector* fj,
175                       [[maybe_unused]] BasicVector* fk,
176                       [[maybe_unused]] ShiftForce* shf_ij,
177                       [[maybe_unused]] ShiftForce* shf_kj,
178                       [[maybe_unused]] ShiftForce* shf_c)
179 {
180     return 0.0;
181 };
182
183 template<class ThreeCenterType, class BasicVector, class ShiftForce>
184 inline NBLIB_ALWAYS_INLINE
185 auto computeThreeCenter(const ThreeCenterType& parameters, const BasicVector& rij, const BasicVector& rkj,
186                         const BasicVector& /* rik */, BasicVector* fi, BasicVector* fj, BasicVector* fk,
187                         ShiftForce* shf_ij, ShiftForce* shf_kj, ShiftForce* shf_c)
188 {
189     using ValueType = BasicVectorValueType_t<BasicVector>;
190     // calculate 3-center common quantities: angle between x1-x2 and x2-x3
191     // Todo: after sufficient evaluation, switch over to atan2 based algorithm
192     ValueType costh = basicVectorCosAngle(rij, rkj); /* 25 */
193     ValueType theta = std::acos(costh);    /* 10 */
194
195     // call type-specific angle kernel, e.g. harmonic, restricted, quartic, etc.
196     auto [force, energy] = threeCenterKernel(theta, parameters);
197
198     spreadThreeCenterForces(costh, force, rij, rkj, fi, fj, fk, shf_ij, shf_kj, shf_c);
199
200     return energy;
201 }
202
203 /*! \brief Calculate three-center interactions
204  *
205  * \tparam Force buffer type
206  * \tparam Three centre interaction parameters
207  * \tparam Cartesian vector type
208  * \tparam PBC type
209  * \param[in] index
210  * \param[in] Bond parameters
211  * \param[in] x coordinate array
212  * \param[in/out] Force buffer
213  * \param[in] PBC
214  * \return Computed kernel energies
215  */
216 template <class Buffer, class ThreeCenterType, class BasicVector, class ShiftForce, class Pbc,
217           std::enable_if_t<Contains<ThreeCenterType, SupportedThreeCenterTypes>{}>* = nullptr>
218 inline NBLIB_ALWAYS_INLINE
219 auto dispatchInteraction(InteractionIndex<ThreeCenterType> index,
220                          gmx::ArrayRef<const ThreeCenterType> parameters,
221                          gmx::ArrayRef<const BasicVector> x,
222                          Buffer* forces,
223                          gmx::ArrayRef<ShiftForce> shiftForces,
224                          const Pbc& pbc)
225 {
226     KernelEnergy<BasicVectorValueType_t<BasicVector>> energy;
227
228     //! fetch input data: position vectors x1-x3 and interaction parameters
229     int i = std::get<0>(index);
230     int j = std::get<1>(index);
231     int k = std::get<2>(index);
232     BasicVector xi = x[i];
233     BasicVector xj = x[j];
234     BasicVector xk = x[k];
235     ThreeCenterType threeCenterParameters = parameters[std::get<3>(index)];
236
237     BasicVector fi{0,0,0}, fj{0,0,0}, fk{0,0,0};
238
239     BasicVector rij, rkj, rik;
240     int sIdx_ij = pbc.dxAiuc(xi, xj, rij); /* 3 */
241     int sIdx_kj = pbc.dxAiuc(xk, xj, rkj); /* 3 */
242     pbc.dxAiuc(xi, xk, rik); /* 3 */
243
244     ShiftForce* shf_ij = accessShiftForces(sIdx_ij, shiftForces);
245     ShiftForce* shf_kj = accessShiftForces(sIdx_kj, shiftForces);
246     ShiftForce* shf_c  = accessShiftForces(gmx::c_centralShiftIndex, shiftForces);
247
248     energy.carrier()            = computeThreeCenter(threeCenterParameters, rij, rkj, rik, &fi, &fj, &fk, shf_ij, shf_kj, shf_c);
249     energy.twoCenterAggregate() = addTwoCenterAggregate(threeCenterParameters, rij, rkj, &fi, &fj, &fk, shf_ij, shf_kj, shf_c);
250
251     (*forces)[i] += fi;
252     (*forces)[j] += fj;
253     (*forces)[k] += fk;
254
255     return energy;
256 }
257
258 template<class FourCenterType, class BasicVector>
259 inline NBLIB_ALWAYS_INLINE
260 std::enable_if_t<HasThreeCenterAggregate<FourCenterType>::value, BasicVectorValueType_t<BasicVector>>
261 addThreeCenterAggregate(const FourCenterType& parameters,
262                         const BasicVector& rij, const BasicVector& rkj, const BasicVector& rkl,
263                         BasicVector* fi, BasicVector* fj, BasicVector* fk, BasicVector* fl)
264 {
265     using ValueType = BasicVectorValueType_t<BasicVector>;
266     if (parameters.manifest == FourCenterType::Cargo::j)
267     {
268         return computeThreeCenter(parameters.threeCenter(), rij, rkj, fi, fj, fk);
269     }
270     if (parameters.manifest == FourCenterType::Cargo::k)
271     {
272         return computeThreeCenter(parameters.threeCenter(), ValueType(-1.0)*rkj, ValueType(-1.0)*rkl, fj, fk, fl);
273         //return computeThreeCenter(parameters.threeCenter(), rkj, rkl, fj, fk, fl);
274     }
275
276     // aggregate is empty
277     return 0.0;
278 };
279
280 template<class FourCenterType, class BasicVector>
281 inline NBLIB_ALWAYS_INLINE
282 std::enable_if_t<!HasThreeCenterAggregate<FourCenterType>::value, BasicVectorValueType_t<BasicVector>>
283 addThreeCenterAggregate([[maybe_unused]] const FourCenterType& parameters,
284                         [[maybe_unused]] const BasicVector& rij,
285                         [[maybe_unused]] const BasicVector& rkj,
286                         [[maybe_unused]] const BasicVector& rkl,
287                         [[maybe_unused]] BasicVector* fi,
288                         [[maybe_unused]] BasicVector* fj,
289                         [[maybe_unused]] BasicVector* fk,
290                         [[maybe_unused]] BasicVector* fl)
291 {
292     return 0.0;
293 };
294
295 /*! \brief Calculate four-center interactions
296  *
297  * \tparam Force buffer type
298  * \tparam FourCenterType The bond type to compute; used for type deduction
299  * \tparam Cartesian vector type
300  * \tparam PBC type
301  * \param[in] index The atom and parameter indices used for computing the interaction
302  * \param[in] parameters The full type-specific interaction list
303  * \param[in] x The coordinates
304  * \param[in/out] forces The forces
305  * \param[in] pbc Object used for computing distances accounting for PBC's
306  * \return Computed kernel energies
307  */
308 template <class Buffer, class FourCenterType, class BasicVector, class ShiftForce, class Pbc,
309           std::enable_if_t<Contains<FourCenterType, SupportedFourCenterTypes>{}>* = nullptr>
310 inline NBLIB_ALWAYS_INLINE
311 auto dispatchInteraction(InteractionIndex<FourCenterType> index,
312                          gmx::ArrayRef<const FourCenterType> parameters,
313                          gmx::ArrayRef<const BasicVector> x,
314                          Buffer* forces,
315                          gmx::ArrayRef<ShiftForce> shiftForces,
316                          const Pbc& pbc)
317 {
318     using RealScalar = BasicVectorValueType_t<BasicVector>;
319     KernelEnergy<RealScalar> energy;
320
321     int i = std::get<0>(index);
322     int j = std::get<1>(index);
323     int k = std::get<2>(index);
324     int l = std::get<3>(index);
325
326     BasicVector xi = x[i];
327     BasicVector xj = x[j];
328     BasicVector xk = x[k];
329     BasicVector xl = x[l];
330
331     BasicVector fi{0,0,0}, fj{0,0,0}, fk{0,0,0}, fl{0,0,0};
332
333     BasicVector dxIJ, dxKJ, dxKL, dxLJ;
334     int sIdx_ij = pbc.dxAiuc(xi, xj, dxIJ);
335     int sIdx_kj = pbc.dxAiuc(xk, xj, dxKJ);
336     pbc.dxAiuc(xk, xl, dxKL);
337
338     int sIdx_lj = pbc.dxAiuc(xl, xj, dxLJ);
339
340     ShiftForce* shf_ij = accessShiftForces(sIdx_ij, shiftForces);
341     ShiftForce* shf_kj = accessShiftForces(sIdx_kj, shiftForces);
342     ShiftForce* shf_lj = accessShiftForces(sIdx_lj, shiftForces);
343     ShiftForce* shf_c  = accessShiftForces(gmx::c_centralShiftIndex, shiftForces);
344
345     FourCenterType fourCenterTypeParams = parameters[std::get<4>(index)];
346
347     BasicVector m, n;
348     RealScalar phi = dihedralPhi(dxIJ, dxKJ, dxKL, &m, &n);
349
350     auto [force, kernelEnergy] = fourCenterKernel(phi, fourCenterTypeParams);
351
352     energy.carrier()              = kernelEnergy;
353     energy.threeCenterAggregate() = addThreeCenterAggregate(fourCenterTypeParams, dxIJ, dxKJ, dxKL, &fi, &fj, &fk, &fl);
354
355     spreadFourCenterForces(force, dxIJ, dxKJ, dxKL, m, n, &fi, &fj, &fk, &fl, shf_ij, shf_kj, shf_lj, shf_c);
356
357     (*forces)[i] += fi;
358     (*forces)[j] += fj;
359     (*forces)[k] += fk;
360     (*forces)[l] += fl;
361
362     return energy;
363 }
364
365 /*! \brief Calculate five-center interactions
366  *
367  * \tparam Force buffer type
368  * \tparam FiveCenterType The bond type to compute; used for type deduction
369  * \tparam Cartesian vector type
370  * \tparam PBC type
371  * \param[in] index The atom and parameter indices used for computing the interaction
372  * \param[in] parameters The full type-specific interaction list
373  * \param[in] x The coordinates
374  * \param[in/out] forces The forces
375  * \param[in] pbc Object used for computing distances accounting for PBC's
376  * \return Computed kernel energies
377  */
378 template <class Buffer, class FiveCenterType, class BasicVector, class ShiftForce, class Pbc,
379           std::enable_if_t<Contains<FiveCenterType, SupportedFiveCenterTypes>{}>* = nullptr>
380 inline NBLIB_ALWAYS_INLINE
381 auto dispatchInteraction(InteractionIndex<FiveCenterType> index,
382                          gmx::ArrayRef<const FiveCenterType> parameters,
383                          gmx::ArrayRef<const BasicVector> x,
384                          Buffer* forces,
385                          [[maybe_unused]] gmx::ArrayRef<ShiftForce> shiftForces,
386                          const Pbc& pbc)
387 {
388     KernelEnergy<BasicVectorValueType_t<BasicVector>> energy;
389
390     int i = std::get<0>(index);
391     int j = std::get<1>(index);
392     int k = std::get<2>(index);
393     int l = std::get<3>(index);
394     int m = std::get<4>(index);
395
396     BasicVector xi = x[i];
397     BasicVector xj = x[j];
398     BasicVector xk = x[k];
399     BasicVector xl = x[l];
400     BasicVector xm = x[m];
401
402     BasicVector dxIJ, dxJK, dxKL, dxLM;
403     pbc.dxAiuc(xi, xj, dxIJ);
404     pbc.dxAiuc(xj, xk, dxJK);
405     pbc.dxAiuc(xk, xl, dxKL);
406     pbc.dxAiuc(xl, xm, dxLM);
407
408     FiveCenterType fiveCenterTypeParams = parameters[std::get<5>(index)];
409
410     // this dispatch function is not in use yet, because CMap is not yet implemented
411     // we don't want to add [[maybe_unused]] in the signature
412     // and we also don't want compiler warnings, so we cast to void
413     (void)fiveCenterTypeParams;
414     (void)forces;
415
416     return energy;
417 }
418
419
420 /*! \brief implement a loop over bonds for a given BondType and Kernel
421  *  corresponds to e.g. the "bonds" function at Gromacs:bonded.cpp@450
422  *
423  * \param[in] indices interaction atom pair indices + bond parameter index
424  * \param[in] interactionParameters bond/interaction parameters
425  * \param[in] x coordinate input
426  * \param[in/out] forces The forces
427  * \param[in] pbc Object used for computing distances accounting for PBC's
428  * \return Computed kernel energies
429  */
430 template <class Index, class InteractionType, class Buffer, class ShiftForce, class Pbc>
431 auto computeForces(gmx::ArrayRef<const Index> indices,
432                    gmx::ArrayRef<const InteractionType> parameters,
433                    gmx::ArrayRef<const Vec3> x,
434                    Buffer* forces,
435                    gmx::ArrayRef<ShiftForce> shiftForces,
436                    const Pbc& pbc)
437 {
438     KernelEnergy<BasicVectorValueType_t<Vec3>> energy;
439
440     for (const auto& index : indices)
441     {
442         energy += dispatchInteraction(index, parameters, x, forces, shiftForces, pbc);
443     }
444
445     return energy;
446 }
447
448 //! \brief convenience overload without shift forces
449 template <class Index, class InteractionType, class Buffer, class Pbc>
450 auto computeForces(gmx::ArrayRef<const Index> indices,
451                    gmx::ArrayRef<const InteractionType> parameters,
452                    gmx::ArrayRef<const Vec3> x,
453                    Buffer* forces,
454                    const Pbc& pbc)
455 {
456     return computeForces(indices, parameters, x, forces, gmx::ArrayRef<std::nullptr_t>{}, pbc);
457 }
458
459 /*! \brief implement a loop over bond types and accumulate their force contributions
460  *
461  * \param[in] interactions interaction pairs and bond parameters
462  * \param[in] x coordinate input
463  * \param[in/out] forces output force buffer
464  * \param[in] pbc Object used for computing distances accounting for PBC's
465  * \return Computed kernel energies
466  */
467 template<class Buffer, class ShiftForce, class Pbc>
468 auto reduceListedForces(const ListedInteractionData& interactions,
469                         gmx::ArrayRef<const Vec3> x,
470                         Buffer* forces,
471                         gmx::ArrayRef<ShiftForce> shiftForces,
472                         const Pbc& pbc)
473 {
474     using ValueType = BasicVectorValueType_t<Vec3>;
475     std::array<ValueType, std::tuple_size<ListedInteractionData>::value> energies{0};
476     energies.fill(0);
477
478     // calculate one bond type
479     auto computeForceType = [forces, x, shiftForces, &energies, &pbc](const auto& interactionElement) {
480         using InteractionType = typename std::decay_t<decltype(interactionElement)>::type;
481
482         gmx::ArrayRef<const InteractionIndex<InteractionType>> indices(interactionElement.indices);
483         gmx::ArrayRef<const InteractionType>                   parameters(interactionElement.parameters);
484
485         KernelEnergy<ValueType> energy = computeForces(indices, parameters, x, forces, shiftForces, pbc);
486
487         energies[CarrierIndex<InteractionType>{}] += energy.carrier();
488         energies[TwoCenterAggregateIndex<InteractionType>{}] += energy.twoCenterAggregate();
489         energies[ThreeCenterAggregateIndex<InteractionType>{}] += energy.threeCenterAggregate();
490         // energies[FepIndex{}] += energy.freeEnergyDerivative();
491     };
492
493     // calculate all bond types, returns a tuple with the energies for each type
494     for_each_tuple(computeForceType, interactions);
495
496     return energies;
497 }
498
499 //! \brief convenience overload without shift forces
500 template<class Buffer, class Pbc>
501 auto reduceListedForces(const ListedInteractionData& interactions,
502                         gmx::ArrayRef<const Vec3> x,
503                         Buffer* forces,
504                         const Pbc& pbc)
505 {
506     return reduceListedForces(interactions, x, forces, gmx::ArrayRef<std::nullptr_t>{}, pbc);
507 }
508
509 } // namespace nblib
510
511 #undef NBLIB_ALWAYS_INLINE
512
513 #endif // NBLIB_LISTEDFORCES_DATAFLOW_HPP