LjPmeMgr.C

Go to the documentation of this file.

/* modified from NAMD */

#include "LjPmeMgr.h"
#include "Lattice.h"
#include "LjPmeBase.h"

void LjPmeMgr::initialize(const SimParameters *simParams, const int nAtoms) {
#ifndef USE_LJPME
  NAMD_die("In order to use LJ-PME, you must build with -DUSE_LJPME");
#else
#if defined(NAMD_FFTW_3) || ! defined(NAMD_FFTW)
  NAMD_die("LJ-PME requires FFTW 2");
#endif
#endif
  if (initialized) {
                NAMD_die("LjPmeMgr has already been initialized!");
  }
  int numRecipPes = 1;
  selfEnergy = recipEnergy = 0.0;
  setSelf = false;
  numAtoms = nAtoms;
  // Store grid informations
  myGrid.K1 = simParams->LJPMEGridSizeX;
  myGrid.K2 = simParams->LJPMEGridSizeY;
  myGrid.K3 = simParams->LJPMEGridSizeZ;
  myGrid.order = simParams->LJPMEInterpOrder;
  myGrid.dim2 = myGrid.K2;
  myGrid.dim3 = 2 * (myGrid.K3 / 2 + 1);
  myGrid.block1 = (myGrid.K1 + numRecipPes - 1) / numRecipPes;
  myGrid.block2 = (myGrid.K2 + numRecipPes - 1) / numRecipPes;

  // Allocate memory
  myKSpace = new LjPmeKSpace(myGrid);
  myRealSpace = new LjPmeRealSpace(myGrid, numAtoms);
  dataArr = new double[4*numAtoms];
  if (dataArr == 0) NAMD_die("can't allocate LJ-PME Manager dataArr buffer");

  qsize = myGrid.K1 * myGrid.dim2 * myGrid.dim3;
  fsize = myGrid.K1 * myGrid.dim2;
  q_arr = new float *[fsize];
  if (q_arr == 0) NAMD_die("can't allocate LJ-PME Manager q_arr buffer");
  memset((void *)q_arr, 0, fsize * sizeof(float *));

  // kludge so we won't segfault
  for (int i = 0; i < fsize; i++)
  {
    q_arr[i] = new float[myGrid.dim3];
    if (q_arr[i] == 0) NAMD_die("can't allocate LJ-PME Manager q_arr[i] buffer");
    memset((void *)q_arr[i], 0, myGrid.dim3 * sizeof(float));
  }
  // end kludge

  f_arr = new char[fsize];
  if (f_arr == 0) NAMD_die("can't allocate LJ-PME Manager f_arr buffer");
  fz_arr = new char[myGrid.dim3];
  if (fz_arr == 0) NAMD_die("can't allocate LJ-PME Manager fz_arr buffer");

  qGrid = new float[qsize];
  if (qGrid == 0) NAMD_die("can't allocate LJ-PME Manager qGrid buffer");

  this->optimizeFFT();
  memset((void *)qGrid, 0, qsize * sizeof(float));
  initialized = true;
}

LjPmeMgr::~LjPmeMgr() {
  setSelf = false;
  initialized = false;
  if (myRealSpace) delete myRealSpace;
  if (myKSpace) delete myKSpace;
  if (dataArr) delete [] dataArr;
  if (qGrid) delete [] qGrid;
  if (f_arr) delete [] f_arr;
  if (fz_arr) delete [] fz_arr;
  if (work) delete [] work;

  for (int i = 0; i < fsize; ++i) {
    if (q_arr[i]) {
      delete [] q_arr[i];
    }
  }
}

void LjPmeMgr::optimizeFFT() {
  int n[3];
  n[0] = myGrid.K1;
  n[1] = myGrid.K2;
  n[2] = myGrid.K3;
#if defined(USE_LJPME) && ! defined(NAMD_FFTW_3)
  /*
  *  see if using FFTW_ESTIMATE makes start up faster
  */
  #ifdef NAMD_FFTW
  #ifdef NAMD_FFTW_3
  work = new fftwf_complex[n[0]];

  iout << iINFO << "Optimizing 4 LJ-PME FFT steps.  1..." << endi;
  forward_plan_yz = fftwf_plan_many_dft_r2c(2, n+1, 
                myGrid.K1, qGrid, NULL, 1, myGrid.dim2 * myGrid.dim3,
                (fftwf_complex *)qGrid, NULL, 1, 
                myGrid.dim2 * (myGrid.dim3 / 2), FFTW_ESTIMATE);

  iout << " 2..." << endi;
  int zdim = myGrid.dim3;
  int xStride=myGrid.dim2 *( myGrid.dim3 / 2);
  forward_plan_x = fftwf_plan_many_dft(1, n, xStride,
                                                (fftwf_complex *)qGrid, NULL, xStride, 1,
                                                (fftwf_complex *)qGrid, NULL, xStride, 1,
                                                FFTW_FORWARD, FFTW_ESTIMATE);

  iout << " 3..." << endi;          
  backward_plan_x = fftwf_plan_many_dft(1, n, xStride,
                                                (fftwf_complex *)qGrid, NULL, xStride, 1,
                                                (fftwf_complex *)qGrid, NULL, xStride, 1,
                                                FFTW_BACKWARD, FFTW_ESTIMATE);

  iout << " 4..." << endi;
        backward_plan_yz = fftwf_plan_many_dft_c2r(2, n+1, 
                                                      myGrid.K1,(fftwf_complex *)qGrid,
                                                      NULL, 1, myGrid.dim2 * (myGrid.dim3 / 2),
                                                      qGrid, NULL, 1, myGrid.dim2 * myGrid.dim3,
                                                      FFTW_ESTIMATE);
  iout << "   Done.\n" << endi;
  #else
  work = new fftw_complex[n[0]];

  iout << iINFO << "Optimizing 4 LJ-PME FFT steps.  1..." << endi;
  forward_plan_yz = rfftwnd_create_plan_specific(2, n+1, FFTW_REAL_TO_COMPLEX,
                  FFTW_ESTIMATE | FFTW_IN_PLACE | FFTW_USE_WISDOM,
                  qGrid, 1, 0, 0);
  iout << " 2..." << endi;
  forward_plan_x = fftw_create_plan_specific(n[0], FFTW_FORWARD,
                 FFTW_ESTIMATE | FFTW_IN_PLACE | FFTW_USE_WISDOM, (fftw_complex *)qGrid,
                 myGrid.dim2 * myGrid.dim3 / 2, work, 1);
  iout << " 3..." << endi;
  backward_plan_x = fftw_create_plan_specific(n[0], FFTW_BACKWARD,
                  FFTW_ESTIMATE | FFTW_IN_PLACE | FFTW_USE_WISDOM, (fftw_complex *)qGrid,
                  myGrid.dim2 * myGrid.dim3 / 2, work, 1);
  iout << " 4..." << endi;
  backward_plan_yz = rfftwnd_create_plan_specific(2, n + 1, FFTW_COMPLEX_TO_REAL,
                   FFTW_ESTIMATE | FFTW_IN_PLACE | FFTW_USE_WISDOM,
                   qGrid, 1, 0, 0);
  iout << "   Done.\n" << endi;
  #endif
  #else
  NAMD_die("Sorry, FFTW must be compiled in to use LJ-PME.");
  #endif
#endif // USE_LJPME
}

void LjPmeMgr::computeLongRange(const double *ljPmeCoord,
        const Lattice &lattice, const double &alphaLJ,
        double *force, double &energy, double virial[][3],
        bool doEnergy) {
  for (int i = 0; i < fsize; ++i) {
    if (q_arr[i]) {
      memset((void *)(q_arr[i]), 0, myGrid.dim3 * sizeof(float));
    }
  }
  //memset((void *)qGrid, 0, qsize * sizeof(float)); // Do I need this?
  memset((void *)f_arr, 0, fsize * sizeof(char));
  memset((void *)fz_arr, 0, myGrid.dim3 * sizeof(char));
  recipEnergy = 0.0;
  for(int i = 0; i < 6; i++) {
    recipVirial[i] = 0.0;
  }

  // Store the charge and scaled coordinates in dataArr buffer
  this->setScaledCoordinates(ljPmeCoord, lattice);
  // Compute and store the self term
  if(!setSelf) {
    selfEnergy = this->selfCompute(alphaLJ);
    setSelf = true;
  }
  // Split charges into the grid
  myRealSpace->fill_charges(q_arr, f_arr, fz_arr, dataArr);
  this->gridCalculation(alphaLJ, lattice);
  myRealSpace->compute_scaledForces(q_arr, dataArr, force, lattice);

  // Add energy and virial
  if(doEnergy) {
    energy += recipEnergy + selfEnergy;
    virial[0][0] += recipVirial[0];
    virial[0][1] += recipVirial[1];
    virial[0][2] += recipVirial[2];
    virial[1][0] += recipVirial[1];
    virial[1][1] += recipVirial[3];
    virial[1][2] += recipVirial[4];
    virial[2][0] += recipVirial[2];
    virial[2][1] += recipVirial[4];
    virial[2][2] += recipVirial[5];
  }
}

void LjPmeMgr::setScaledCoordinates(const double *refPos, const Lattice &lattice) {
  Vector origin = lattice.origin();
  Vector recip1 = lattice.a_r();
  Vector recip2 = lattice.b_r();
  Vector recip3 = lattice.c_r();
  double ox = origin.x;
  double oy = origin.y;
  double oz = origin.z;
  double r1x = recip1.x;
  double r1y = recip1.y;
  double r1z = recip1.z;
  double r2x = recip2.x;
  double r2y = recip2.y;
  double r2z = recip2.z;
  double r3x = recip3.x;
  double r3y = recip3.y;
  double r3z = recip3.z;
  int K1 = myGrid.K1;
  int K2 = myGrid.K2;
  int K3 = myGrid.K3;
  double shift1 = ((K1 + myGrid.order - 1)/2)/(double)K1;
  double shift2 = ((K2 + myGrid.order - 1)/2)/(double)K2;
  double shift3 = ((K3 + myGrid.order - 1)/2)/(double)K3;

  for (int i=0; i<numAtoms; i++) {
    int index = 4*i;
    double px = refPos[index]   - ox;
    double py = refPos[index+1] - oy;
    double pz = refPos[index+2] - oz;
    double c3Term = refPos[index+3];
    double sx = shift1 + px*r1x + py*r1y + pz*r1z;
    double sy = shift2 + px*r2x + py*r2y + pz*r2z;
    double sz = shift3 + px*r3x + py*r3y + pz*r3z;
    px = K1 * ( sx - floor(sx) );
    py = K2 * ( sy - floor(sy) );
    pz = K3 * ( sz - floor(sz) );
    //  Check for rare rounding condition where K * ( 1 - epsilon ) == K
    //  which was observed with g++ on Intel x86 architecture.
    if ( px == K1 ) px = 0;
    if ( py == K2 ) py = 0;
    if ( pz == K3 ) pz = 0;

    dataArr[index] = px;
    dataArr[index+1] = py;
    dataArr[index+2] = pz;
    dataArr[index+3] = c3Term;
  }   
}

void LjPmeMgr::gridCalculation(const double &alpha, const Lattice &lattice)
{
#if defined(USE_LJPME) && ! defined(NAMD_FFTW_3)
  // Part 1:
  int i, j, k = 0;
  for (i = 0; i < myGrid.K1 * myGrid.K2; i++) {
    for (j = 0; j < myGrid.dim3; j++) {
      qGrid[k++] = q_arr[i][j];
    }
  }
  #ifdef NAMD_FFTW
  #ifdef NAMD_FFTW_3
  fftwf_execute(forward_plan_yz);
  #else
  rfftwnd_real_to_complex(forward_plan_yz, myGrid.K1,
                          (fftw_real *)qGrid, 1, myGrid.dim2 * myGrid.dim3,
                          0, 0, 0);
  #endif
  #endif

  // Part2:
  int zdim = myGrid.dim3;
  int ny = myGrid.dim2;

  // finish forward FFT (x dimension)
  #ifdef NAMD_FFTW
  #ifdef NAMD_FFTW_3
  fftwf_execute(forward_plan_x);
  #else
  fftw(forward_plan_x, ny * zdim / 2, (fftw_complex *)qGrid,
       ny * zdim / 2, 1, work, 1, 0);
  #endif
  #endif

  // Calculate energy and virial
  double tempVir[6];
  recipEnergy += myKSpace->compute_energy(qGrid, lattice, alpha, tempVir);
  for(i = 0; i < 6; i++) {
    recipVirial[i] += tempVir[i];
  }

  #ifdef NAMD_FFTW
  #ifdef NAMD_FFTW_3
  fftwf_execute(backward_plan_x);
  #else
  // start backward FFT (x dimension)
  fftw(backward_plan_x, ny * zdim / 2, (fftw_complex *)qGrid,
       ny * zdim / 2, 1, work, 1, 0);
  #endif
  #endif

  // Part 3:
  #ifdef NAMD_FFTW
  #ifdef NAMD_FFTW_3
  fftwf_execute(backward_plan_yz);
  #else
  rfftwnd_complex_to_real(backward_plan_yz, myGrid.K1,
                          (fftw_complex *)qGrid, 1, myGrid.dim2 * myGrid.dim3 / 2,
                          0, 0, 0);
  #endif
  #endif
  k = 0;
  for (i = 0; i < myGrid.K1 * myGrid.K2; i++) {
    for (j = 0; j < myGrid.dim3; j++) {
      q_arr[i][j] = qGrid[k++];
    }
  }
#endif // USE_LJPME
}

double LjPmeMgr::selfCompute(const double &alphaLJ){
        double energy = 0.0;
        double c3Term;
  double alpha6 = pow(alphaLJ, 6.0);
        for(int i=0; i < numAtoms; i++) {
                c3Term = dataArr[4*i+3];
                energy += c3Term*c3Term;
        }
  return (energy * alpha6 / 12.0);
}