ReferencePME.cpp 28.6 KB
Newer Older
1
/*
2
 * Reference implementation of PME reciprocal space interactions.
3
 *
4
 * Copyright (c) 2009-2025, Erik Lindahl, Rossen Apostolov, Szilard Pall, Peter Eastman, Evan Pretti
5
6
 * All rights reserved.
 * Contact: lindahl@cbr.su.se Stockholm University, Sweden.
7
8
 *
 * Redistribution and use in source and binary forms, with or without
9
10
 * modification, are permitted provided that the following conditions are met:
 *
11
12
13
14
 * Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer. Redistributions in binary
 * form must reproduce the above copyright notice, this list of conditions and
 * the following disclaimer in the documentation and/or other materials provided
15
 * with the distribution.
16
17
 * Neither the name of the author/university nor the names of its contributors may
 * be used to endorse or promote products derived from this software without
18
19
 * specific prior written permission.
 *
20
21
22
23
24
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
25
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
26
 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
27
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29
30
31
 * POSSIBILITY OF SUCH DAMAGE.
 */

Peter Eastman's avatar
Peter Eastman committed
32
33
#include <cmath>
#include <vector>
34

35
#include "ReferencePME.h"
36
#include "ReferenceForce.h"
peastman's avatar
peastman committed
37
#include "SimTKOpenMMRealType.h"
38

Peter Eastman's avatar
Peter Eastman committed
39
40
41
42
43
44
#ifdef _MSC_VER
  #define POCKETFFT_NO_VECTORS
#endif
#include "pocketfft_hdronly.h"

using namespace std;
45

46
namespace OpenMM {
47

Peter Eastman's avatar
Peter Eastman committed
48
49
50
51
52
53
/* Only called once from constructor, performance does not matter! */
void ReferencePME::calculate_bsplines_moduli() {
    int nmax = 0;
    for (int d = 0; d < 3; d++) {
        nmax = (ngrid[d] > nmax) ? ngrid[d] : nmax;
        bsplines_moduli[d].resize(ngrid[d]);
54
    }
55

Peter Eastman's avatar
Peter Eastman committed
56
    /* temp storage in this routine */
Peter Eastman's avatar
Peter Eastman committed
57
58
59
    vector<double> data(order);
    vector<double> ddata(order);
    vector<double> bsplines_data(nmax);
Peter Eastman's avatar
Peter Eastman committed
60
61
62
63
64

    data[order-1]=0;
    data[1]=0;
    data[0]=1;

Peter Eastman's avatar
Peter Eastman committed
65
66
    for (int k = 3; k < order; k++) {
        double div=1.0/(k-1.0);
Peter Eastman's avatar
Peter Eastman committed
67
        data[k-1]=0;
Peter Eastman's avatar
Peter Eastman committed
68
        for (int l = 1; l < (k-1); l++)
Peter Eastman's avatar
Peter Eastman committed
69
70
71
72
73
74
            data[k-l-1]=div*(l*data[k-l-2]+(k-l)*data[k-l-1]);
        data[0]=div*data[0];
    }

    /* differentiate */
    ddata[0]=-data[0];
Peter Eastman's avatar
Peter Eastman committed
75
    for (int k = 1 ; k < order; k++)
Peter Eastman's avatar
Peter Eastman committed
76
77
        ddata[k]=data[k-1]-data[k];

Peter Eastman's avatar
Peter Eastman committed
78
    double div=1.0/(order-1);
Peter Eastman's avatar
Peter Eastman committed
79
80
    data[order-1]=0;

Peter Eastman's avatar
Peter Eastman committed
81
    for (int l = 1; l < (order-1); l++)
Peter Eastman's avatar
Peter Eastman committed
82
83
84
        data[order-l-1]=div*(l*data[order-l-2]+(order-l)*data[order-l-1]);
    data[0]=div*data[0];

Peter Eastman's avatar
Peter Eastman committed
85
    for (int i = 0; i < nmax; i++)
Peter Eastman's avatar
Peter Eastman committed
86
        bsplines_data[i]=0;
Peter Eastman's avatar
Peter Eastman committed
87
    for (int i = 1; i <= order; i++)
Peter Eastman's avatar
Peter Eastman committed
88
89
90
        bsplines_data[i]=data[i-1];

    /* Evaluate the actual bspline moduli for X/Y/Z */
Peter Eastman's avatar
Peter Eastman committed
91
92
93
94
95
96
97
    for (int d = 0; d < 3; d++) {
        int ndata = ngrid[d];
        for (int i = 0; i < ndata; i++) {
            double sc = 0;
            double ss = 0;
            for (int j = 0; j < ndata; j++) {
                double arg=(2.0*M_PI*i*j)/ndata;
Peter Eastman's avatar
Peter Eastman committed
98
99
100
                sc+=bsplines_data[j]*cos(arg);
                ss+=bsplines_data[j]*sin(arg);
            }
Peter Eastman's avatar
Peter Eastman committed
101
            bsplines_moduli[d][i]=sc*sc+ss*ss;
Peter Eastman's avatar
Peter Eastman committed
102
        }
Peter Eastman's avatar
Peter Eastman committed
103
104
105
        for (int i = 0; i < ndata; i++) {
            if (bsplines_moduli[d][i]<1.0e-7)
                bsplines_moduli[d][i]=(bsplines_moduli[d][(i-1+ndata)%ndata]+bsplines_moduli[d][(i+1)%ndata])/2;
Peter Eastman's avatar
Peter Eastman committed
106
107
        }
    }
108
109
}

Peter Eastman's avatar
Peter Eastman committed
110
111
void ReferencePME::update_grid_index_and_fraction(const vector<Vec3>& atomCoordinates, const Vec3 recipBoxVectors[3]) {
    for (int i = 0; i < natoms; i++) {
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        /* Index calculation (Look mom, no conditionals!):
         *
         * Both for Cuda and modern CPUs it is nice to avoid conditionals, but we still need to apply periodic boundary conditions.
         * Instead of having loops to add/subtract the box dimension, we do it this way:
         *
         * 1. First add the box size, to make sure this atom coordinate isnt -0.1 or something.
         *    After this we assume all fractional box positions are *positive*.
         *    The reason for this is that we always want to round coordinates _down_ to get
         *    their grid index, and when taking the integer part of -3.4 we would get -3, not -4 as we want.
         *    Since we anyway need the grid indices to fall in the central box, it is more convenient
         *    to first manipulate the coordinates to be positive.
         * 2. Convert to integer grid index
         *    Since we have added a whole box unit in step 1, this index might actually be larger than
         *    the grid dimension. Examples, assuming 10*10*10nm box and grid dimension 100*100*100 (spacing 0.1 nm):
         *
         *    coordinate is { 0.543 , 6.235 , -0.73 }
         *
         *    x[i][d]/box[d]                      becomes   { 0.0543 , 0.6235 , -0.073 }
         *    (x[i][d]/box[d] + 1.0)              becomes   { 1.0543 , 1.6235 , 0.927 }
         *    (x[i][d]/box[d] + 1.0)*ngrid[d]     becomes   { 105.43 , 162.35 , 92.7 }
         *
         *    integer part is now { 105 , 162 , 92 }
         *
         *    The fraction is calculates as t-ti, which becomes { 0.43 , 0.35 , 0.7 }
         *
         * 3. Take the first integer index part (which can be larger than the grid) modulo the grid dimension
         *
         *    Now we get { 5 , 62 , 92 }
         *
         *    Voila, both index and fraction, entirely without conditionals. The one limitation here is that
         *    we only add one box length, so if the particle had a coordinate <=-10.0, we would be screwed.
         *    In principle we can of course add 100.0, but that just moves the problem, it doesnt solve it.
         *    In practice, MD programs will apply PBC to reset particles inside the central box to avoid
         *    numerical problems, so this shouldnt cause any problems.
         *    (And, by adding 100.0 box lengths, we would lose a bit of numerical accuracy here!)
         */
peastman's avatar
peastman committed
148
        Vec3 coord = atomCoordinates[i];
Peter Eastman's avatar
Peter Eastman committed
149
150
151
152
153
154
155
        for (int d = 0; d < 3; d++) {
            double t = coord[0]*recipBoxVectors[0][d]+coord[1]*recipBoxVectors[1][d]+coord[2]*recipBoxVectors[2][d];
            t = (t-floor(t))*ngrid[d];
            int ti = (int) t;

            particlefraction[i][d] = t - ti;
            particleindex[i][d]    = ti % ngrid[d];
Peter Eastman's avatar
Peter Eastman committed
156
157
        }
    }
158
159
160
}


161
/* Ugly bspline calculation taken from Tom Dardens reference equations.
162
163
164
165
 * This probably very sub-optimal in Cuda? Separate kernel?
 *
 * In practice, it might help to require order=4 for the cuda port.
 */
Peter Eastman's avatar
Peter Eastman committed
166
167
168
void ReferencePME::update_bsplines() {
    for (int i = 0; i < natoms; i++) {
        for (int j = 0; j < 3; j++) {
Peter Eastman's avatar
Peter Eastman committed
169
            /* dr is relative offset from lower cell limit */
Peter Eastman's avatar
Peter Eastman committed
170
            double dr = particlefraction[i][j];
Peter Eastman's avatar
Peter Eastman committed
171

Peter Eastman's avatar
Peter Eastman committed
172
173
            double* data  = &(bsplines_theta[j][i*order]);
            double* ddata = &(bsplines_dtheta[j][i*order]);
Peter Eastman's avatar
Peter Eastman committed
174
175
176
177
            data[order-1] = 0;
            data[1]       = dr;
            data[0]       = 1-dr;

Peter Eastman's avatar
Peter Eastman committed
178
179
            for (int k = 3; k < order; k++) {
                double div = 1.0/(k-1.0);
Peter Eastman's avatar
Peter Eastman committed
180
                data[k-1] = div*dr*data[k-2];
Peter Eastman's avatar
Peter Eastman committed
181
                for (int l = 1; l < (k-1); l++)
Peter Eastman's avatar
Peter Eastman committed
182
183
184
185
186
187
188
                    data[k-l-1] = div*((dr+l)*data[k-l-2]+(k-l-dr)*data[k-l-1]);
                data[0] = div*(1-dr)*data[0];
            }

            /* differentiate */
            ddata[0] = -data[0];

Peter Eastman's avatar
Peter Eastman committed
189
            for (int k = 1; k < order; k++)
Peter Eastman's avatar
Peter Eastman committed
190
191
                ddata[k] = data[k-1]-data[k];

Peter Eastman's avatar
Peter Eastman committed
192
            double div = 1.0/(order-1);
Peter Eastman's avatar
Peter Eastman committed
193
194
            data[order-1] = div*dr*data[order-2];

Peter Eastman's avatar
Peter Eastman committed
195
            for (int l = 1; l < (order-1); l++)
Peter Eastman's avatar
Peter Eastman committed
196
197
198
                data[order-l-1] = div*((dr+l)*data[order-l-2]+(order-l-dr)*data[order-l-1]);
            data[0] = div*(1-dr)*data[0];
        }
199
200
201
202
    }
}


Peter Eastman's avatar
Peter Eastman committed
203
void ReferencePME::grid_spread_charge(const vector<double>& charges) {
204
    /* Reset the grid */
Peter Eastman's avatar
Peter Eastman committed
205
206
    for (int i = 0; i < ngrid[0]*ngrid[1]*ngrid[2]; i++)
        grid[i] = complex<double>(0, 0);
207

Peter Eastman's avatar
Peter Eastman committed
208
209
    for (int i = 0; i < natoms; i++) {
        double q = charges[i];
210

211
        /* Grid index for the actual atom position */
Peter Eastman's avatar
Peter Eastman committed
212
213
214
        int x0index = particleindex[i][0];
        int y0index = particleindex[i][1];
        int z0index = particleindex[i][2];
215

216
        /* Bspline factors for this atom in each dimension , calculated from fractional coordinates */
Peter Eastman's avatar
Peter Eastman committed
217
218
219
        double* thetax  = &(bsplines_theta[0][i*order]);
        double* thetay  = &(bsplines_theta[1][i*order]);
        double* thetaz  = &(bsplines_theta[2][i*order]);
220

221
        /* Loop over norder*norder*norder (typically 5*5*5) neighbor cells.
222
223
         *
         * As a neat optimization, we only spread in the forward direction, but apply PBC!
224
         *
225
         * Since we are going to do an FFT on the grid, it doesn't matter where the data is,
226
         * in frequency space the result will be the same.
227
228
         *
         * So, the influence function (bsplines) will probably be something like (0.15,0.35,0.35,0.15),
229
         * with largest weight 2-3 steps forward (you don't need to understand that for the implementation :-)
230
231
         * Effectively, you can look at this as translating the entire grid.
         *
232
         * Why do we do this stupid thing?
233
         *
234
         * 1) The loops get much simpler
235
236
237
         * 2) Just looking forward will hopefully get us more cache hits
         * 3) When we parallelize things, we only need to communicate in one direction instead of two!
         */
238

Peter Eastman's avatar
Peter Eastman committed
239
        for (int ix = 0; ix < order; ix++) {
240
            /* Calculate index, apply PBC so we spread to index 0/1/2 when a particle is close to the upper limit of the grid */
Peter Eastman's avatar
Peter Eastman committed
241
            int xindex = (x0index + ix) % ngrid[0];
242

Peter Eastman's avatar
Peter Eastman committed
243
244
            for (int iy = 0; iy < order; iy++) {
                int yindex = (y0index + iy) % ngrid[1];
245

Peter Eastman's avatar
Peter Eastman committed
246
                for (int iz = 0; iz < order; iz++) {
Peter Eastman's avatar
Peter Eastman committed
247
                    /* Can be optimized, but we keep it simple here */
Peter Eastman's avatar
Peter Eastman committed
248
                    int zindex = (z0index + iz) % ngrid[2];
249
                    /* Calculate index in the charge grid */
Peter Eastman's avatar
Peter Eastman committed
250
                    int index = xindex*ngrid[1]*ngrid[2] + yindex*ngrid[2] + zindex;
251
                    /* Add the charge times the bspline spread/interpolation factors to this grid position */
Peter Eastman's avatar
Peter Eastman committed
252
                    grid[index] += q*thetax[ix]*thetay[iy]*thetaz[iz];
Peter Eastman's avatar
Peter Eastman committed
253
254
255
256
                }
            }
        }
    }
257
258
259
260
}



Peter Eastman's avatar
Peter Eastman committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
void ReferencePME::pme_reciprocal_convolution(const Vec3 periodicBoxVectors[3], const Vec3 recipBoxVectors[3], double& energy) {
    int nx = ngrid[0];
    int ny = ngrid[1];
    int nz = ngrid[2];

    double one_4pi_eps = ONE_4PI_EPS0/epsilon_r;
    double factor = M_PI*M_PI/(ewaldcoeff*ewaldcoeff);
    double boxfactor = M_PI*periodicBoxVectors[0][0]*periodicBoxVectors[1][1]*periodicBoxVectors[2][2];

    double esum = 0;

    double maxkx = (nx+1)/2;
    double maxky = (ny+1)/2;
    double maxkz = (nz+1)/2;

    for (int kx = 0; kx < nx; kx++) {
277
        /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
278
279
280
        double mx  = (kx<maxkx) ? kx : (kx-nx);
        double mhx = mx*recipBoxVectors[0][0];
        double bx  = boxfactor*bsplines_moduli[0][kx];
281

Peter Eastman's avatar
Peter Eastman committed
282
        for (int ky = 0; ky < ny; ky++) {
283
            /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
284
285
286
            double my  = (ky<maxky) ? ky : (ky-ny);
            double mhy = mx*recipBoxVectors[1][0]+my*recipBoxVectors[1][1];
            double by  = bsplines_moduli[1][ky];
287

Peter Eastman's avatar
Peter Eastman committed
288
            for (int kz = 0; kz < nz; kz++) {
289
                /* Pointer to the grid cell in question */
Peter Eastman's avatar
Peter Eastman committed
290
                complex<double>& ptr = grid[kx*ny*nz + ky*nz + kz];
291
292
293
294
295
296

                /* The zero frequency term is undefined due to division by the frequency below.  Set this term to zero;
                 * in the case that the net charge of the system is non-zero, this is equivalent to applying a uniform
                 * neutralizing background charge density.  The contribution to the energy and charge derivatives of
                 * this neutralizing plasma is applied elsewhere.  If this term is not zeroed, however, energies and
                 * forces will be unaffected but charge derivatives for non-neutral systems will be incorrect!
297
                 */
Peter Eastman's avatar
Peter Eastman committed
298
299
                if (kx==0 && ky==0 && kz==0) {
                    ptr = 0;
Peter Eastman's avatar
Peter Eastman committed
300
301
                    continue;
                }
302
303

                /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
304
305
                double mz        = (kz<maxkz) ? kz : (kz-nz);
                double mhz       = mx*recipBoxVectors[2][0]+my*recipBoxVectors[2][1]+mz*recipBoxVectors[2][2];
306
307

                /* Get grid data for this frequency */
Peter Eastman's avatar
Peter Eastman committed
308
309
                double d1        = ptr.real();
                double d2        = ptr.imag();
310

311
                /* Calculate the convolution - see the Essman/Darden paper for the equation! */
Peter Eastman's avatar
Peter Eastman committed
312
313
314
                double m2        = mhx*mhx+mhy*mhy+mhz*mhz;
                double bz        = bsplines_moduli[2][kz];
                double denom     = m2*bx*by*bz;
315

Peter Eastman's avatar
Peter Eastman committed
316
                double eterm     = one_4pi_eps*exp(-factor*m2)/denom;
317
318

                /* write back convolution data to grid */
Peter Eastman's avatar
Peter Eastman committed
319
320
                ptr.real(d1*eterm);
                ptr.imag(d2*eterm);
321

Peter Eastman's avatar
Peter Eastman committed
322
                double struct2   = (d1*d1+d2*d2);
323

Peter Eastman's avatar
Peter Eastman committed
324
                /* Long-range PME contribution to the energy for this frequency */
Peter Eastman's avatar
Peter Eastman committed
325
                double ets2      = eterm*struct2;
Peter Eastman's avatar
Peter Eastman committed
326
327
328
329
                esum     += ets2;
            }
        }
    }
330

331
    /* The factor 0.5 is nothing special, but it is better to have it here than inside the loop :-) */
Peter Eastman's avatar
Peter Eastman committed
332
    energy = 0.5*esum;
333
334
335
}


Peter Eastman's avatar
Peter Eastman committed
336
337
338
339
void ReferencePME::dpme_reciprocal_convolution(const Vec3 periodicBoxVectors[3], const Vec3 recipBoxVectors[3], double& energy) {
    int nx = ngrid[0];
    int ny = ngrid[1];
    int nz = ngrid[2];
340

Peter Eastman's avatar
Peter Eastman committed
341
    double boxfactor = -2*M_PI*sqrt(M_PI) / (6.0*periodicBoxVectors[0][0]*periodicBoxVectors[1][1]*periodicBoxVectors[2][2]);
342

Peter Eastman's avatar
Peter Eastman committed
343
    double esum = 0;
344

Peter Eastman's avatar
Peter Eastman committed
345
346
347
    double maxkx = (nx+1)/2;
    double maxky = (ny+1)/2;
    double maxkz = (nz+1)/2;
348

Peter Eastman's avatar
Peter Eastman committed
349
    double bfac = M_PI / ewaldcoeff;
peastman's avatar
peastman committed
350
    double fac1 = 2.0*M_PI*M_PI*M_PI*sqrt(M_PI);
Peter Eastman's avatar
Peter Eastman committed
351
352
    double fac2 = ewaldcoeff*ewaldcoeff*ewaldcoeff;
    double fac3 = -2.0*ewaldcoeff*M_PI*M_PI;
353

Peter Eastman's avatar
Peter Eastman committed
354
    for (int kx = 0; kx < nx; kx++) {
355
        /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
356
357
358
        double mx  = ((kx<maxkx) ? kx : (kx-nx));
        double mhx = mx*recipBoxVectors[0][0];
        double bx  = bsplines_moduli[0][kx];
359

Peter Eastman's avatar
Peter Eastman committed
360
        for (int ky = 0; ky < ny; ky++) {
361
            /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
362
363
364
            double my  = ((ky<maxky) ? ky : (ky-ny));
            double mhy = mx*recipBoxVectors[1][0]+my*recipBoxVectors[1][1];
            double by  = bsplines_moduli[1][ky];
365

Peter Eastman's avatar
Peter Eastman committed
366
            for (int kz = 0; kz < nz; kz++) {
367
368
369
370
371
                /*
                 * Unlike the Coulombic case, there's an m=0 term so all terms are considered here.
                 */

                /* Calculate frequency. Grid indices in the upper half correspond to negative frequencies! */
Peter Eastman's avatar
Peter Eastman committed
372
373
                double mz  = ((kz<maxkz) ? kz : (kz-nz));
                double mhz = mx*recipBoxVectors[2][0]+my*recipBoxVectors[2][1]+mz*recipBoxVectors[2][2];
374
375

                /* Pointer to the grid cell in question */
Peter Eastman's avatar
Peter Eastman committed
376
                complex<double>& ptr = grid[kx*ny*nz + ky*nz + kz];
377
378

                /* Get grid data for this frequency */
Peter Eastman's avatar
Peter Eastman committed
379
380
                double d1 = ptr.real();
                double d2 = ptr.imag();
381
382

                /* Calculate the convolution - see the Essman/Darden paper for the equation! */
Peter Eastman's avatar
Peter Eastman committed
383
384
385
                double m2    = mhx*mhx+mhy*mhy+mhz*mhz;
                double bz    = bsplines_moduli[2][kz];
                double denom = boxfactor / (bx*by*bz);
386

Peter Eastman's avatar
Peter Eastman committed
387
388
389
390
391
392
                double m = sqrt(m2);
                double m3 = m*m2;
                double b = bfac*m;
                double expfac = -b*b;
                double erfcterm = erfc(b);
                double expterm = exp(expfac);
393

Peter Eastman's avatar
Peter Eastman committed
394
                double eterm = (fac1*erfcterm*m3 + expterm*(fac2 + fac3*m2)) * denom;
395
396

                /* write back convolution data to grid */
Peter Eastman's avatar
Peter Eastman committed
397
398
                ptr.real(d1*eterm);
                ptr.imag(d2*eterm);
399

Peter Eastman's avatar
Peter Eastman committed
400
                double struct2 = (d1*d1+d2*d2);
401
402

                /* Long-range PME contribution to the energy for this frequency */
Peter Eastman's avatar
Peter Eastman committed
403
404
                double ets2 = eterm*struct2;
                esum += ets2;
405
406
407
408
            }
        }
    }
    // Remember the C6 energy is attractive, hence the negative sign.
Peter Eastman's avatar
Peter Eastman committed
409
    energy = 0.5*esum;
410
411
412
}


Peter Eastman's avatar
Peter Eastman committed
413
void ReferencePME::grid_interpolate_force(const Vec3 recipBoxVectors[3], const vector<double>& charges, vector<Vec3>& forces) {
414
    /* This is almost identical to the charge spreading routine! */
415

Peter Eastman's avatar
Peter Eastman committed
416
417
    for (int i = 0; i < natoms; i++) {
        double fx = 0, fy = 0, fz = 0;
418

Peter Eastman's avatar
Peter Eastman committed
419
        double q = charges[i];
420

421
        /* Grid index for the actual atom position */
Peter Eastman's avatar
Peter Eastman committed
422
423
424
        int x0index = particleindex[i][0];
        int y0index = particleindex[i][1];
        int z0index = particleindex[i][2];
425

426
        /* Bspline factors for this atom in each dimension , calculated from fractional coordinates */
Peter Eastman's avatar
Peter Eastman committed
427
428
429
430
431
432
        double* thetax  = &(bsplines_theta[0][i*order]);
        double* thetay  = &(bsplines_theta[1][i*order]);
        double* thetaz  = &(bsplines_theta[2][i*order]);
        double* dthetax = &(bsplines_dtheta[0][i*order]);
        double* dthetay = &(bsplines_dtheta[1][i*order]);
        double* dthetaz = &(bsplines_dtheta[2][i*order]);
433

Peter Eastman's avatar
Peter Eastman committed
434
        /* See grid_spread_charge() for comments about the order here, and only interpolation in one direction */
435

436
        /* Since we will add order^3 (typically 5*5*5=125) terms to the force on each particle, we use temporary fx/fy/fz
437
438
         * variables, and only add it to memory forces[] at the end.
         */
Peter Eastman's avatar
Peter Eastman committed
439
440
        for (int ix = 0; ix < order; ix++) {
            int xindex = (x0index + ix) % ngrid[0];
441
            /* Get both the bspline factor and its derivative with respect to the x coordinate! */
Peter Eastman's avatar
Peter Eastman committed
442
443
            double tx  = thetax[ix];
            double dtx = dthetax[ix];
444

Peter Eastman's avatar
Peter Eastman committed
445
446
            for (int iy = 0; iy < order; iy++) {
                int yindex = (y0index + iy) % ngrid[1];
447
                /* bspline + derivative wrt y */
Peter Eastman's avatar
Peter Eastman committed
448
449
                double ty  = thetay[iy];
                double dty = dthetay[iy];
450

Peter Eastman's avatar
Peter Eastman committed
451
                for (int iz = 0; iz < order; iz++) {
Peter Eastman's avatar
Peter Eastman committed
452
                    /* Can be optimized, but we keep it simple here */
Peter Eastman's avatar
Peter Eastman committed
453
                    int zindex = (z0index + iz) % ngrid[2];
454
                    /* bspline + derivative wrt z */
Peter Eastman's avatar
Peter Eastman committed
455
456
457
                    double tz  = thetaz[iz];
                    double dtz = dthetaz[iz];
                    int index  = xindex*ngrid[1]*ngrid[2] + yindex*ngrid[2] + zindex;
458

459
460
                    /* Get the fft+convoluted+ifft:d data from the grid, which must be real by definition */
                    /* Checking that the imaginary part is indeed zero might be a good check :-) */
Peter Eastman's avatar
Peter Eastman committed
461
                    double gridvalue = grid[index].real();
Peter Eastman's avatar
Peter Eastman committed
462
463

                    /* The d component of the force is calculated by taking the derived bspline in dimension d, normal bsplines in the other two */
Peter Eastman's avatar
Peter Eastman committed
464
465
466
                    fx += dtx*ty*tz*gridvalue;
                    fy += tx*dty*tz*gridvalue;
                    fz += tx*ty*dtz*gridvalue;
Peter Eastman's avatar
Peter Eastman committed
467
468
469
                }
            }
        }
470
        /* Update memory force, note that we multiply by charge and some box stuff */
Peter Eastman's avatar
Peter Eastman committed
471
472
473
        forces[i][0] -= q*(fx*ngrid[0]*recipBoxVectors[0][0]);
        forces[i][1] -= q*(fx*ngrid[0]*recipBoxVectors[1][0]+fy*ngrid[1]*recipBoxVectors[1][1]);
        forces[i][2] -= q*(fx*ngrid[0]*recipBoxVectors[2][0]+fy*ngrid[1]*recipBoxVectors[2][1]+fz*ngrid[2]*recipBoxVectors[2][2]);
Peter Eastman's avatar
Peter Eastman committed
474
    }
475
476
477
}


Peter Eastman's avatar
Peter Eastman committed
478
479
480
void ReferencePME::grid_interpolate_charge_derivatives(const Vec3 recipBoxVectors[3], const vector<double>& charges,
            vector<double>& chargeDerivatives, const vector<int>& chargeIndices) {
    /* This is similar to grid_interpolate_force() */
481
    
Peter Eastman's avatar
Peter Eastman committed
482
483
484
485
    int nderiv = chargeIndices.size();
    for (int ideriv = 0; ideriv < nderiv; ideriv++) {
        int i = chargeIndices[ideriv];
        double dq = 0;
486
487

        /* Grid index for the actual atom position */
Peter Eastman's avatar
Peter Eastman committed
488
489
490
        int x0index = particleindex[i][0];
        int y0index = particleindex[i][1];
        int z0index = particleindex[i][2];
491
492

        /* Bspline factors for this atom in each dimension , calculated from fractional coordinates */
Peter Eastman's avatar
Peter Eastman committed
493
494
495
        double* thetax  = &(bsplines_theta[0][i*order]);
        double* thetay  = &(bsplines_theta[1][i*order]);
        double* thetaz  = &(bsplines_theta[2][i*order]);
496

Peter Eastman's avatar
Peter Eastman committed
497
        /* See grid_spread_charge() for comments about the order here, and only interpolation in one direction */
498
499
500
501
502

        /* Since we will add order^3 (typically 5*5*5=125) terms to the charge
         * derivative on each particle, we use a temporary dq variable, and only
         * add it to memory forces[] at the end.
         */
Peter Eastman's avatar
Peter Eastman committed
503
504
        for (int ix = 0; ix < order; ix++) {
            int xindex = (x0index + ix) % ngrid[0];
505
            /* Get the bspline factor with respect to the x coordinate */
Peter Eastman's avatar
Peter Eastman committed
506
            double tx  = thetax[ix];
507

Peter Eastman's avatar
Peter Eastman committed
508
509
            for (int iy = 0; iy < order; iy++) {
                int yindex = (y0index + iy) % ngrid[1];
510
                /* bspline wrt y */
Peter Eastman's avatar
Peter Eastman committed
511
                double ty  = thetay[iy];
512

Peter Eastman's avatar
Peter Eastman committed
513
                for (int iz = 0; iz < order; iz++) {
514
                    /* Can be optimized, but we keep it simple here */
Peter Eastman's avatar
Peter Eastman committed
515
                    int zindex = (z0index + iz) % ngrid[2];
516
                    /* bspline wrt z */
Peter Eastman's avatar
Peter Eastman committed
517
518
                    double tz  = thetaz[iz];
                    int index  = xindex*ngrid[1]*ngrid[2] + yindex*ngrid[2] + zindex;
519
520
521

                    /* Get the fft+convoluted+ifft:d data from the grid, which must be real by definition */
                    /* Checking that the imaginary part is indeed zero might be a good check :-) */
Peter Eastman's avatar
Peter Eastman committed
522
                    double gridvalue = grid[index].real();
523
524

                    /* The d component of the force is calculated by taking the derived bspline in dimension d, normal bsplines in the other two */
Peter Eastman's avatar
Peter Eastman committed
525
                    dq += tx*ty*tz*gridvalue;
526
527
528
529
530
531
532
533
                }
            }
        }

        chargeDerivatives[ideriv] += dq;
    }
}

Peter Eastman's avatar
Peter Eastman committed
534
535
536
537
538
539
ReferencePME::ReferencePME(double ewaldcoeff, int natoms, const int ngrid[3], int pme_order, double epsilon_r) :
                    order(pme_order), epsilon_r(epsilon_r), ewaldcoeff(ewaldcoeff), natoms(natoms) {
    for (int d = 0; d < 3; d++) {
        this->ngrid[d] = ngrid[d];
        bsplines_theta[d].resize(pme_order*natoms);
        bsplines_dtheta[d].resize(pme_order*natoms);
Peter Eastman's avatar
Peter Eastman committed
540
    }
541

Peter Eastman's avatar
Peter Eastman committed
542
543
    particlefraction.resize(natoms);
    particleindex.resize(natoms);
544

Peter Eastman's avatar
Peter Eastman committed
545
    /* Allocate charge grid storage */
Peter Eastman's avatar
Peter Eastman committed
546
    grid.resize(ngrid[0]*ngrid[1]*ngrid[2]);
547

Peter Eastman's avatar
Peter Eastman committed
548
    /* Setup bspline moduli (see Essman paper) */
Peter Eastman's avatar
Peter Eastman committed
549
    calculate_bsplines_moduli();
550
551
}

Peter Eastman's avatar
Peter Eastman committed
552
553
void ReferencePME::exec(const vector<Vec3>& atomCoordinates, vector<Vec3>& forces, const vector<double>& charges,
            const Vec3 periodicBoxVectors[3], double& energy) {
554
    /* Routine is called with coordinates in x, a box, and charges in q */
555

peastman's avatar
peastman committed
556
    Vec3 recipBoxVectors[3];
557
    ReferenceForce::invertBoxVectors(periodicBoxVectors, recipBoxVectors);
558
    
559
    /* Before we can do the actual interpolation, we need to recalculate and update
Peter Eastman's avatar
Peter Eastman committed
560
     * the indices for each particle in the charge grid (initialized in constructor),
561
562
     * and what its fractional offset in this grid cell is.
     */
563

Peter Eastman's avatar
Peter Eastman committed
564
    /* Update charge grid indices and fractional offsets for each atom.
565
     * The indices/fractions are stored internally in the pme datatype
566
     */
Peter Eastman's avatar
Peter Eastman committed
567
    update_grid_index_and_fraction(atomCoordinates, recipBoxVectors);
568

Peter Eastman's avatar
Peter Eastman committed
569
    /* Calculate bsplines (and their differentials) from current fractional coordinates, store in pme structure */
Peter Eastman's avatar
Peter Eastman committed
570
    update_bsplines();
571

Peter Eastman's avatar
Peter Eastman committed
572
    /* Spread the charges on grid (using newly calculated bsplines in the pme structure) */
Peter Eastman's avatar
Peter Eastman committed
573
    grid_spread_charge(charges);
574

Peter Eastman's avatar
Peter Eastman committed
575
    /* do 3d-fft */
Peter Eastman's avatar
Peter Eastman committed
576
    vector<size_t> shape = {(size_t) ngrid[0], (size_t) ngrid[1], (size_t) ngrid[2]};
Peter Eastman's avatar
Peter Eastman committed
577
    vector<size_t> axes = {0, 1, 2};
Peter Eastman's avatar
Peter Eastman committed
578
579
    vector<ptrdiff_t> stride = {(ptrdiff_t) (ngrid[1]*ngrid[2]*sizeof(complex<double>)),
                                (ptrdiff_t) (ngrid[2]*sizeof(complex<double>)),
Peter Eastman's avatar
Peter Eastman committed
580
                                (ptrdiff_t) sizeof(complex<double>)};
Peter Eastman's avatar
Peter Eastman committed
581
    pocketfft::c2c(shape, stride, stride, axes, true, grid.data(), grid.data(), 1.0, 0);
582

Peter Eastman's avatar
Peter Eastman committed
583
    /* solve in k-space */
Peter Eastman's avatar
Peter Eastman committed
584
    pme_reciprocal_convolution(periodicBoxVectors, recipBoxVectors, energy);
585

Peter Eastman's avatar
Peter Eastman committed
586
    /* do 3d-invfft */
Peter Eastman's avatar
Peter Eastman committed
587
    pocketfft::c2c(shape, stride, stride, axes, false, grid.data(), grid.data(), 1.0, 0);
588

Peter Eastman's avatar
Peter Eastman committed
589
    /* Get the particle forces from the grid and bsplines in the pme structure */
Peter Eastman's avatar
Peter Eastman committed
590
    grid_interpolate_force(recipBoxVectors, charges, forces);
591
592
}

Peter Eastman's avatar
Peter Eastman committed
593
594
void ReferencePME::exec_charge_derivatives(const vector<Vec3>& atomCoordinates, vector<double>& chargeDerivatives,
            const vector<int>& chargeIndices, const vector<double>& charges, const Vec3 periodicBoxVectors[3]) {
595
596
597
598
599
600
    /* Routine is called with coordinates in x, a box, and charges in q */

    Vec3 recipBoxVectors[3];
    ReferenceForce::invertBoxVectors(periodicBoxVectors, recipBoxVectors);
    
    /* Before we can do the actual interpolation, we need to recalculate and update
Peter Eastman's avatar
Peter Eastman committed
601
     * the indices for each particle in the charge grid (initialized in constructor),
602
603
604
605
606
607
     * and what its fractional offset in this grid cell is.
     */

    /* Update charge grid indices and fractional offsets for each atom.
     * The indices/fractions are stored internally in the pme datatype
     */
Peter Eastman's avatar
Peter Eastman committed
608
    update_grid_index_and_fraction(atomCoordinates, recipBoxVectors);
609
610

    /* Calculate bsplines (and their differentials) from current fractional coordinates, store in pme structure */
Peter Eastman's avatar
Peter Eastman committed
611
    update_bsplines();
612
613

    /* Spread the charges on grid (using newly calculated bsplines in the pme structure) */
Peter Eastman's avatar
Peter Eastman committed
614
    grid_spread_charge(charges);
615
616

    /* do 3d-fft */
Peter Eastman's avatar
Peter Eastman committed
617
    vector<size_t> shape = {(size_t) ngrid[0], (size_t) ngrid[1], (size_t) ngrid[2]};
618
    vector<size_t> axes = {0, 1, 2};
Peter Eastman's avatar
Peter Eastman committed
619
620
    vector<ptrdiff_t> stride = {(ptrdiff_t) (ngrid[1]*ngrid[2]*sizeof(complex<double>)),
                                (ptrdiff_t) (ngrid[2]*sizeof(complex<double>)),
621
                                (ptrdiff_t) sizeof(complex<double>)};
Peter Eastman's avatar
Peter Eastman committed
622
    pocketfft::c2c(shape, stride, stride, axes, true, grid.data(), grid.data(), 1.0, 0);
623
624
625

    /* solve in k-space */
    double energy;
Peter Eastman's avatar
Peter Eastman committed
626
    pme_reciprocal_convolution(periodicBoxVectors, recipBoxVectors, energy);
627
628

    /* do 3d-invfft */
Peter Eastman's avatar
Peter Eastman committed
629
    pocketfft::c2c(shape, stride, stride, axes, false, grid.data(), grid.data(), 1.0, 0);
630
631

    /* Get the charge derivatives from the grid and bsplines in the pme structure */
Peter Eastman's avatar
Peter Eastman committed
632
    grid_interpolate_charge_derivatives(recipBoxVectors, charges, chargeDerivatives, chargeIndices);
633
634
}

Peter Eastman's avatar
Peter Eastman committed
635
636
void ReferencePME::exec_dpme(const vector<Vec3>& atomCoordinates, vector<Vec3>& forces, const vector<double>& c6s,
            const Vec3 periodicBoxVectors[3], double& energy) {
637
638
    /* Routine is called with coordinates in x, a box, and charges in q */

peastman's avatar
peastman committed
639
    Vec3 recipBoxVectors[3];
640
    ReferenceForce::invertBoxVectors(periodicBoxVectors, recipBoxVectors);
641
642

    /* Before we can do the actual interpolation, we need to recalculate and update
Peter Eastman's avatar
Peter Eastman committed
643
     * the indices for each particle in the charge grid (initialized in constructor),
644
645
646
647
648
649
     * and what its fractional offset in this grid cell is.
     */

    /* Update charge grid indices and fractional offsets for each atom.
     * The indices/fractions are stored internally in the pme datatype
     */
Peter Eastman's avatar
Peter Eastman committed
650
    update_grid_index_and_fraction(atomCoordinates, recipBoxVectors);
651
652

    /* Calculate bsplines (and their differentials) from current fractional coordinates, store in pme structure */
Peter Eastman's avatar
Peter Eastman committed
653
    update_bsplines();
654
655

    /* Spread the charges on grid (using newly calculated bsplines in the pme structure) */
Peter Eastman's avatar
Peter Eastman committed
656
    grid_spread_charge(c6s);
657
658

    /* do 3d-fft */
Peter Eastman's avatar
Peter Eastman committed
659
    vector<size_t> shape = {(size_t) ngrid[0], (size_t) ngrid[1], (size_t) ngrid[2]};
Peter Eastman's avatar
Peter Eastman committed
660
    vector<size_t> axes = {0, 1, 2};
Peter Eastman's avatar
Peter Eastman committed
661
662
    vector<ptrdiff_t> stride = {(ptrdiff_t) (ngrid[1]*ngrid[2]*sizeof(complex<double>)),
                                (ptrdiff_t) (ngrid[2]*sizeof(complex<double>)),
Peter Eastman's avatar
Peter Eastman committed
663
                                (ptrdiff_t) sizeof(complex<double>)};
Peter Eastman's avatar
Peter Eastman committed
664
    pocketfft::c2c(shape, stride, stride, axes, true, grid.data(), grid.data(), 1.0, 0);
665
666

    /* solve in k-space */
Peter Eastman's avatar
Peter Eastman committed
667
    dpme_reciprocal_convolution(periodicBoxVectors, recipBoxVectors, energy);
668
669

    /* do 3d-invfft */
Peter Eastman's avatar
Peter Eastman committed
670
    pocketfft::c2c(shape, stride, stride, axes, false, grid.data(), grid.data(), 1.0, 0);
671
672

    /* Get the particle forces from the grid and bsplines in the pme structure */
Peter Eastman's avatar
Peter Eastman committed
673
    grid_interpolate_force(recipBoxVectors, c6s, forces);
674
}
675
676

} // namespace OpenMM