Molecular dynamics of a Lennard-Jones fluid (OpenGL)

ljmdgl.c

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#if defined(Macintosh) || defined(__APPLE__)
/* to compile:  gcc -framework OpenGL -framework GLUT lj3.c */
#include <GLUT/glut.h>
#else
#include <GL/glut.h>
#endif

#define INLINE __inline
#define D        3
#define N        256
#define DOF      (D*N-D*(D+1)/2)  /* degrees of freedom */
double rho = 0.8;
double T = 1.0;
double rc = 2.5;
double dt = 0.002; /* time step for Newton's equation */
double L;
double x[N][3]; /* from 0 to 1 */
double v[N][3], f[N][3]; /* velocity and force */
double U, Us, K, Vir, p;
double ushift; /* potential energy shift due to truncation, used to check energy conservation */
double Utail; /* tail correction for potential energy */
double ptail; /* tail correction for pressure */
int thermostat = 1; /* use a thermostat */
double thermdt = 0.01; /* time step for thermostat */
int step = 0;

typedef struct { double cnt, xsm, x2sm; } av_t;
INLINE void av_clear(av_t *av) { av->cnt = av->xsm = av->x2sm = 0; }
INLINE void av_add(av_t *av, double x) { av->cnt += 1; av->xsm += x; av->x2sm += x*x; }
INLINE double av_getav(av_t *av) { return av->cnt > 0 ? av->xsm/av->cnt : 0; }

av_t avU = {0, 0, 0}, avK = {0, 0, 0}, avp = {0, 0, 0};

int delay = 100; /* time delay between frames */
int speed =  50; /* number of steps per frame */
int verbose = 0;
int pause = 0; /* currently in pause */

/* Tausworthe random number generator */
INLINE double rnd0(void)
{
  static unsigned long s1 = 1239, s2 = 123981, s3 = 19861826;
  #define TAUS(s,a,b,c,d) (( ((s&c)<<d) ^ (((s<<a)^s)>>b) ) & 0xfffffffful)
  s1 = TAUS(s1, 13, 19, 4294967294UL, 12);
  s2 = TAUS(s2, 2,  25, 4294967288UL, 4);
  s3 = TAUS(s3, 3,  11, 4294967280UL, 17);
  return (s1^s2^s3)/4294967296.0;
}

/* Gaussian distribution with zero mean and unit variance */
INLINE double grand0(void)
{
  double x, y, u, v, q;
  do {
    u = 1 - rnd0();
    v = 1.7156*(rnd0() - .5);  /* >= 2*sqrt(2/e) */
    x = u - 0.449871;
    y = fabs(v) + 0.386595;
    q = x*x  + y*(0.196*y - 0.25472*x);
    if (q < 0.27597) break;
  } while (q > 0.27846 || v*v > -4*u*u*log(u));
  return v/u;
}

/*  to the nearest image */
INLINE double pbc(double x)
  { return L * (x - ((int)((x)+1000.5) - 1000)); }

INLINE double *vpbc(double v[])
  { v[0] = pbc(v[0]); v[1] = pbc(v[1]); v[2] = pbc(v[2]); return v; }

INLINE double vdot(const double a[], const double b[])
  { return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]; }

INLINE double vsqr(const double x[]) { return vdot(x, x); }

INLINE double *vdif(double c[], const double a[], const double b[])
 { c[0] = a[0]-b[0]; c[1] = a[1]-b[1]; c[2] = a[2]-b[2]; return c; }

INLINE double pbcdist2(double dx[], const double a[], const double b[])
  { return vsqr(vpbc(vdif(dx, a, b))); }

INLINE void vsinc(double x[], const double dx[], double s)
  { x[0] += dx[0]*s; x[1] += dx[1]*s; x[2] += dx[2]*s; }

static void resettime(void)
{
  step = 0;
  av_clear(&avU);
  av_clear(&avK);
  av_clear(&avp);
}

static void setrho(double r)
{
  double irc3, irc6;
  rho = r;
  L = pow(N/rho, 1.0/D);
  if ((rc = 2.5) > L*.5) rc = L*.5;
  irc3 = 1.0/(rc*rc*rc); irc6 = irc3*irc3;
  ushift = 4*irc6*(irc6-1);
  Utail = 8*M_PI*rho*N/9*(irc6 - 3)*irc3;
  ptail = 32*M_PI*rho*rho/9*(irc6 - 1.5)*irc3;
  printf("T %.2f, rho %.2f, N %d, L %6.3f, ushift %6.3f, utail %6.3f, ptail %6.3f\n", T, rho, N, L, ushift, Utail/N, ptail);
  resettime();
}

/* initialize an md system */
static void initmd(void)
{
  int i, j, k, n1, id;
  double a;

  setrho(rho);

  /* init. a fcc lattic */
  n1 = (int) (pow(2*N, 1.0/D) + .999999); /* # of particles per side */
  a = 1./n1;
  for (id = 0, i = 0; i < n1 && id < N; i++)
    for (j = 0; j < n1 && id < N; j++)
      for (k = 0; k < n1 && id < N; k++) {
        if ((i+j+k) % 2 != 0) continue;
        x[id][0] = (i+.5)*a;
        x[id][1] = (j+.5)*a;
        x[id][2] = (k+.5)*a;
        id++;
      }

  for (id = 0; id < N; id++)
    for (j = 0; j < D; j++) v[id][j] = 0.;
}

/* compute force and virial, return energy */
static void force(void)
{
  double dx[3], dr2, dr6, fs, tmp;
  int i, j, d, prcnt = 0;

  for (i = 0; i < N; i++) f[i][0] = f[i][1] = f[i][2] = 0;
  for (U = Vir = 0, i = 0; i < N - 1; i++) {
    for (j = i+1; j < N; j++) {
      dr2 = pbcdist2(dx, x[i], x[j]);
      if (dr2 > rc*rc) continue;
      dr2 = 1/dr2;
      dr6 = dr2*dr2*dr2;
      fs = dr6*(48*dr6-24); /* f.r */
      Vir += fs;
      fs *= dr2;
      for (d = 0; d < D; d++) {
        tmp = dx[d] * fs;
        f[i][d] += tmp;
        f[j][d] -= tmp;
      }
      U += 4*dr6*(dr6-1);
      prcnt++;
    }
  }
  Us = U - prcnt*ushift; /* shifted energy */
  U += Utail; /* unshifted energy */
  p = rho*T+Vir/(3*L*L*L)+ptail;
}

/* remove the center of mass motion */
static void rmcom(void)
{
  int i;
  double vc[D] = {0, 0, 0};

  for (i = 0; i < N; i++) vsinc(vc, v[i], 1);
  for (i = 0; i < D; i++) vc[i] /= N;
  for (i = 0; i < N; i++) vsinc(v[i], vc, -1);
}

/* velocity rescaling thermostat */
static void vrescale(void)
{
  int i, j;
  double ekav = .5f*T*DOF, ek2, s;
  double amp;

  amp = 2*sqrt(K*ekav*thermdt/DOF);
  ek2 = K + (ekav - K)*thermdt + amp*grand0();
  if (ek2 < 0) ek2 = 0;
  s = sqrt(ek2/K);
  for (i = 0; i < N; i++)
    for (j = 0; j < D; j++)
      v[i][j] *= s;
  K = ek2;
}

/* velocity verlet */
static void vv(void)
{
  int i;
  double dth = dt*.5, dtl = dt/L;

  for (i = 0; i < N; i++) { /* VV part 1 */
    vsinc(v[i], f[i], dth);
    vsinc(x[i], v[i], dtl);
  }

  force(); /* calculate the new force */

  for (i = 0; i < N; i++) /* VV part 2 */
    vsinc(v[i], f[i], dth);

  rmcom();
  for (K = 0, i = 0; i < N; i++) K += .5 * vsqr(v[i]);
  if (thermostat) vrescale();
  step++;
  av_add(&avU, U);
  av_add(&avK, K);
  av_add(&avp, p);
}

static void timer(int ival)
{
  int i;

  if (ival != 0 || pause) return;
  for (i = 0; i < speed; i++) vv();
  if (step % 10000 < speed)
    printf("t = %.3f, U/N %6.3f(%6.3f), K/N %5.3f(%5.3f), (Us + K)/N %6.3f, p %5.3f(%5.3f)\n",
      step*dt, U/N, av_getav(&avU)/N, K/N, av_getav(&avK)/N, (Us+K)/N, p, av_getav(&avp));
  glutPostRedisplay();
  glutTimerFunc(delay, timer, 0);
}

double zoomscale = 0.85;
enum {MS_NONE, MS_ZOOM, MS_ROTATE};
int msaction = MS_NONE, msdown, msx, msy;

INLINE double wrap(double x) {
  x += 1000;  x -= (int)x; /* put back into the box */
  x = .5*(1-zoomscale) + zoomscale*x; /* apply zoom scale */
  return 2*x - 1; /* to (-1, 1) */
}

static void display(void)
{
  int i;
  double r = .5*1.122/L * zoomscale * 2; /* last 2 for (0, 1) --> (-1, 1) */

  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  for (i = 0; i < N; i++) {
    glPushMatrix();
    glTranslated(wrap(x[i][0]), wrap(x[i][1]), wrap(x[i][2]));
    glutSolidSphere(r, 20, 20);
    glPopMatrix();
  }
  glutSwapBuffers();
}

enum {T_MINUS, T_PLUS, RHO_MINUS, RHO_PLUS, SPEED_MINUS, SPEED_PLUS, TSTAT,
  SCALE_MINUS, SCALE_PLUS, X_MINUS, X_PLUS, Y_MINUS, Y_PLUS, Z_MINUS, Z_PLUS,
  FULLSCREEN, PAUSE, QUIT, MENULAST};

struct { int id; unsigned char key; const char *desc; } menukey[MENULAST] = {
  {T_MINUS,     't', "- temperature"},
  {T_PLUS,      'T', "+ temperature"},
  {RHO_MINUS,   'r', "- density"},
  {RHO_PLUS,    'R', "+ density"},
  {SPEED_MINUS, 'p', "- speed"},
  {SPEED_PLUS,  'P', "+ speed"},
  {TSTAT,       'm', "Turn on/off thermostat"},
  {SCALE_MINUS, 's', "Zoom out"},
  {SCALE_PLUS,  'S', "Zoom in"},
  {X_MINUS,     'x', "Rotate around -x"},
  {X_PLUS,      'X', "Rotate around +x"},
  {Y_MINUS,     'y', "Rotate around -y"},
  {Y_PLUS,      'Y', "Rotate around +y"},
  {Z_MINUS,     'z', "Rotate around -z"},
  {Z_PLUS,      'Z', "Rotate around +z"},
  {FULLSCREEN,  'f', "Toggle full screen"},
  {PAUSE,       ' ', "Pause/restart"},
  {QUIT,        'q', "Quit"}
};

/* change color of balls according to temperature and density */
static void setcolor(void)
{
  GLfloat diffuse[] = {1.f, 1.f, 0.f, 1.f}, f, g, fs, gs;
  f = (GLfloat) (T/3); if (f > 1) f = 1; else if (f < 0) f = 0;
  fs = (GLfloat) sqrt(1-f*f);
  g = (GLfloat) rho; if (g > 1) g = 1;
  gs = (GLfloat) sqrt(1-g*g);
  diffuse[0] = gs * f;
  diffuse[1] = gs * fs;
  diffuse[2] = g;
  glLightfv(GL_LIGHT0, GL_DIFFUSE, diffuse);
}

static void menu(int id)
{
  GLfloat mat[4][4];

  if (id == T_MINUS || id == T_PLUS) {
    double T1 = T + (id == T_PLUS ? 0.05 : -0.05);
    if (T1 > 0.049) T = T1;
    setcolor();
    resettime();
  } else if (id == RHO_MINUS || id == RHO_PLUS) {
    double r1 = rho + (id == RHO_PLUS ? 0.05 : -0.05);
    if (r1 > 0 && r1 < 1.21) setrho(r1);
    setcolor();
    resettime();
  } else if (id == SPEED_MINUS || id == SPEED_PLUS) {
    int sp1 = speed + (id == SPEED_PLUS ? 10 : -10);
    if (sp1 >= 10 && sp1 < 1000) speed = sp1;
  } else if (id == TSTAT) {
    thermostat = !thermostat;
    resettime();
  } else if (id == SCALE_MINUS || id == SCALE_PLUS) {
    double s1 = zoomscale + (id == SCALE_PLUS ? 0.1 : -0.1);
    if (s1 >= 0.1) zoomscale = s1;
    glutPostRedisplay();
  } else if (id == X_MINUS || id == X_PLUS) {
    glGetFloatv(GL_MODELVIEW_MATRIX, (GLfloat *) mat);
    glRotatef((id == X_PLUS) ? 5.f : -5.f, mat[0][0], mat[1][0], mat[2][0]);
    glutPostRedisplay();
  } else if (id == Y_MINUS || id == Y_PLUS) {
    glGetFloatv(GL_MODELVIEW_MATRIX, (GLfloat *) mat);
    glRotatef((id == Y_PLUS) ? 5.f : -5.f, mat[0][1], mat[1][1], mat[2][1]);
    glutPostRedisplay();
  } else if (id == Z_MINUS || id == Z_PLUS) {
    glGetFloatv(GL_MODELVIEW_MATRIX, (GLfloat *) mat);
    glRotatef((id == Z_PLUS) ? 5.f : -5.f, mat[0][2], mat[1][2], mat[2][2]);
    glutPostRedisplay();
  } else if (id == FULLSCREEN) {
    static int full = 0, x, y, w, h;
    full = !full;
    if (full) {
      x = glutGet(GLUT_WINDOW_X);
      y = glutGet(GLUT_WINDOW_Y);
      w = glutGet(GLUT_WINDOW_WIDTH);
      h = glutGet(GLUT_WINDOW_HEIGHT);
      glutFullScreen();
    } else {
      glutPositionWindow(x, y);
      glutReshapeWindow(w, h);
    }
  } else if (id == PAUSE) {
    pause = !pause;
    if (!pause) timer(0);
  } else if (id == QUIT) {
    exit(0);
  }
  printf("%s: Tstat %s, rho %.2f, T %.2f, speed %d\n",
      menukey[id].desc, (thermostat ? "on" : "off"), rho, T, speed);
}

static void initmenu(void)
{
  int i;
  char s[64];

  glutCreateMenu(menu);
  for (i = 0; i < MENULAST; i++) {
    sprintf(s, "%s, key: %c\n", menukey[i].desc, menukey[i].key);
    glutAddMenuEntry(s, menukey[i].id);
  }
  glutAttachMenu(GLUT_RIGHT_BUTTON);
}

static void keypress(unsigned char c, int x, int y)
{
  int i;

  (void) x; (void) y;
  if (c == 27) exit(0);
  for (i = 0; i < MENULAST; i++)
    if (c == menukey[i].key)
      menu(menukey[i].id);
}

static void usage(void)
{
  int i;

  printf("Hotkeys:\n");
  for (i = 0; i < MENULAST; i++)
    printf(" %c: %s\n", menukey[i].key, menukey[i].desc);
}

static void mouse(int button, int state, int x, int y)
{
  if (state == GLUT_DOWN) {
    msdown++;
    if (verbose >= 1) printf("mouse %d down at %d,%d; %d\n", button, msx, msy, msdown);
    if (button == 3) {
      menu(SCALE_PLUS);
    } else if (button == 4) {
      menu(SCALE_MINUS);
    }
  } else if (--msdown <= 0) {
    msdown = 0;
    msaction = MS_NONE;
  }
  msx = x;
  msy = y;
}

static void motion(int x, int y)
{
  if (x == msx && y == msy) return;
  if (verbose >= 2) printf("mouse at %d,%d; %d\n", x, y, msdown);
  if (msdown) {
    float angx = (float)( (y - msy)*360.f/glutGet(GLUT_WINDOW_HEIGHT) );
    float angy = (float)( (x - msx)*360.f/glutGet(GLUT_WINDOW_WIDTH) );
    float mat[4][4];

    glGetFloatv(GL_MODELVIEW_MATRIX, (GLfloat *) mat);
    glRotated(angx, mat[0][0], mat[1][0], mat[2][0]);
    glRotated(angy, mat[0][1], mat[1][1], mat[2][1]);
    glutPostRedisplay();
  }
  msx = x;
  msy = y;
}

static void reshape(int w, int h)
{
  double xs = 1, ys = 1;
  if (w > h) xs = 1.*w/h;
  else ys = 1.*h/w;
  glViewport(0, 0, w, h);
  glMatrixMode(GL_PROJECTION);
  glLoadIdentity();
  glOrtho(-xs, xs, -ys, ys, -20, 20);
  glMatrixMode(GL_MODELVIEW);
}

static void initgui(void)
{
  GLfloat light_position0[] = {1.0f, 2.0f, 3.0f, 0.0f};  /* Infinite light location. */
  GLfloat bg[] = {0.5f, 0.5f, 0.5f, 1.0f};

  glMaterialfv(GL_FRONT_AND_BACK, GL_AMBIENT, bg);
  glLightfv(GL_LIGHT0, GL_POSITION, light_position0);
  glEnable(GL_LIGHT0);
  glEnable(GL_LIGHTING);
  glEnable(GL_DEPTH_TEST);

  setcolor();
  reshape(glutGet(GLUT_WINDOW_WIDTH), glutGet(GLUT_WINDOW_HEIGHT));
}

int main(int argc, char **argv)
{
  usage();
  initmd();

  glutInit(&argc, argv);
  glutInitWindowSize(800, 800);
  glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH);
  glutCreateWindow("Lennard-Jones simulation 3D");
  initgui();
  initmenu();
  glutDisplayFunc(display);
  glutReshapeFunc(reshape);
  glutKeyboardFunc(keypress);
  glutMouseFunc(mouse);
  glutMotionFunc(motion);
  glutTimerFunc(delay, timer, 0);
  glutMainLoop();
  return 0;
}