5d4dec51b1469954da48ae8a0bd226252fad0d57
[alexxy/gromacs.git] / src / tools / gmx_msd.c
1 /*
2  * 
3  *                This source code is part of
4  * 
5  *                 G   R   O   M   A   C   S
6  * 
7  *          GROningen MAchine for Chemical Simulations
8  * 
9  *                        VERSION 3.2.0
10  * Written by David van der Spoel, Erik Lindahl, Berk Hess, and others.
11  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
12  * Copyright (c) 2001-2004, The GROMACS development team,
13  * check out http://www.gromacs.org for more information.
14
15  * This program is free software; you can redistribute it and/or
16  * modify it under the terms of the GNU General Public License
17  * as published by the Free Software Foundation; either version 2
18  * of the License, or (at your option) any later version.
19  * 
20  * If you want to redistribute modifications, please consider that
21  * scientific software is very special. Version control is crucial -
22  * bugs must be traceable. We will be happy to consider code for
23  * inclusion in the official distribution, but derived work must not
24  * be called official GROMACS. Details are found in the README & COPYING
25  * files - if they are missing, get the official version at www.gromacs.org.
26  * 
27  * To help us fund GROMACS development, we humbly ask that you cite
28  * the papers on the package - you can find them in the top README file.
29  * 
30  * For more info, check our website at http://www.gromacs.org
31  * 
32  * And Hey:
33  * Green Red Orange Magenta Azure Cyan Skyblue
34  */
35 #ifdef HAVE_CONFIG_H
36 #include <config.h>
37 #endif
38
39 #include <string.h>
40 #include <ctype.h>
41 #include <math.h>
42
43 #include "sysstuff.h"
44 #include "smalloc.h"
45 #include "macros.h"
46 #include "statutil.h"
47 #include "maths.h"
48 #include "futil.h"
49 #include "index.h"
50 #include "copyrite.h"
51 #include "typedefs.h"
52 #include "xvgr.h"
53 #include "gstat.h"
54 #include "gmx_statistics.h"
55 #include "tpxio.h"
56 #include "pbc.h"
57 #include "vec.h"
58 #include "confio.h"
59 #include "gmx_ana.h"
60
61
62 #define FACTOR  1000.0  /* Convert nm^2/ps to 10e-5 cm^2/s */
63 /* NORMAL = total diffusion coefficient (default). X,Y,Z is diffusion 
64    coefficient in X,Y,Z direction. LATERAL is diffusion coefficient in
65    plane perpendicular to axis
66 */
67 typedef enum { NOT_USED, NORMAL, X, Y, Z, LATERAL } msd_type;
68
69 typedef struct {
70   real    t0;           /* start time and time increment between  */
71   real    delta_t;      /* time between restart points */
72   real    beginfit,     /* the begin/end time for fits as reals between */
73           endfit;       /* 0 and 1 */
74   real    dim_factor;   /* the dimensionality factor for the diffusion 
75                            constant */
76   real    **data;       /* the displacement data. First index is the group 
77                            number, second is frame number */
78   real    *time;        /* frame time */
79   real    *mass;        /* masses for mass-weighted msd */
80   matrix  **datam;
81   rvec    **x0;         /* original positions */
82   rvec    *com;         /* center of mass correction for each frame */
83   gmx_stats_t **lsq;    /* fitting stats for individual molecule msds */
84   msd_type type;        /* the type of msd to calculate (lateral, etc.)*/
85   int       axis;       /* the axis along which to calculate */
86   int       ncoords;
87   int       nrestart;   /* number of restart points */
88   int       nmol;       /* number of molecules (for bMol) */
89   int       nframes;    /* number of frames */ 
90   int       nlast;
91   int       ngrp;       /* number of groups to use for msd calculation */
92   int       *n_offs;
93   int       **ndata;    /* the number of msds (particles/mols) per data 
94                            point. */
95 } t_corr;
96
97 typedef real t_calc_func(t_corr *,int,atom_id[],int,rvec[],rvec,bool,matrix,
98                          const output_env_t oenv);
99                               
100 static real thistime(t_corr *curr) 
101 {
102   return curr->time[curr->nframes]; 
103 }
104
105 static bool in_data(t_corr *curr,int nx00) 
106
107   return curr->nframes-curr->n_offs[nx00]; 
108 }
109
110 t_corr *init_corr(int nrgrp,int type,int axis,real dim_factor,
111                   int nmol,bool bTen,bool bMass,real dt,t_topology *top,
112                   real beginfit,real endfit)
113 {
114   t_corr  *curr;
115   t_atoms *atoms;
116   int     i;
117
118   snew(curr,1);
119   curr->type      = (msd_type)type;
120   curr->axis      = axis;
121   curr->ngrp      = nrgrp;
122   curr->nrestart  = 0;
123   curr->delta_t   = dt;
124   curr->beginfit  = (1 - 2*GMX_REAL_EPS)*beginfit;
125   curr->endfit    = (1 + 2*GMX_REAL_EPS)*endfit;
126   curr->x0        = NULL;
127   curr->n_offs    = NULL;
128   curr->nframes   = 0;
129   curr->nlast     = 0;
130   curr->dim_factor = dim_factor;
131   
132   snew(curr->ndata,nrgrp);
133   snew(curr->data,nrgrp);
134   if (bTen)
135     snew(curr->datam,nrgrp);
136   for(i=0; (i<nrgrp); i++) {
137     curr->ndata[i] = NULL;
138     curr->data[i]  = NULL;
139     if (bTen)
140       curr->datam[i] = NULL;
141   }
142   curr->time = NULL;
143   curr->lsq  = NULL;
144   curr->nmol = nmol;
145   if (curr->nmol > 0) {
146     snew(curr->mass,curr->nmol);
147     for(i=0; i<curr->nmol; i++)
148       curr->mass[i] = 1;
149   }
150   else {
151     if (bMass) {
152       atoms = &top->atoms;
153       snew(curr->mass,atoms->nr);
154       for(i=0; (i<atoms->nr); i++) {
155         curr->mass[i] = atoms->atom[i].m;
156       }
157     }
158   }
159
160   return curr;
161 }
162
163 static void done_corr(t_corr *curr)
164 {
165   int i;
166   
167   sfree(curr->n_offs);
168   for(i=0; (i<curr->nrestart); i++)
169     sfree(curr->x0[i]);
170   sfree(curr->x0);
171 }
172
173 static void corr_print(t_corr *curr,bool bTen,const char *fn,const char *title,
174                        const char *yaxis,
175                        real msdtime,real beginfit,real endfit,
176                        real *DD,real *SigmaD,char *grpname[],
177                        const output_env_t oenv)
178 {
179   FILE *out;
180   int  i,j;
181   
182   out=xvgropen(fn,title,output_env_get_xvgr_tlabel(oenv),yaxis,oenv);
183   if (DD) {
184     fprintf(out,"# MSD gathered over %g %s with %d restarts\n",
185             msdtime,output_env_get_time_unit(oenv),curr->nrestart);
186     fprintf(out,"# Diffusion constants fitted from time %g to %g %s\n",
187             beginfit,endfit,output_env_get_time_unit(oenv));
188     for(i=0; i<curr->ngrp; i++) 
189       fprintf(out,"# D[%10s] = %.4f (+/- %.4f) (1e-5 cm^2/s)\n",
190               grpname[i],DD[i],SigmaD[i]);
191   }
192   for(i=0; i<curr->nframes; i++) {
193     fprintf(out,"%10g",output_env_conv_time(oenv,curr->time[i]));
194     for(j=0; j<curr->ngrp; j++) {
195       fprintf(out,"  %10g",curr->data[j][i]);
196       if (bTen) {
197         fprintf(out," %10g %10g %10g %10g %10g %10g",
198                 curr->datam[j][i][XX][XX],
199                 curr->datam[j][i][YY][YY],
200                 curr->datam[j][i][ZZ][ZZ],
201                 curr->datam[j][i][YY][XX],
202                 curr->datam[j][i][ZZ][XX],
203                 curr->datam[j][i][ZZ][YY]);
204       }
205     }
206     fprintf(out,"\n");
207   }
208   ffclose(out);
209 }
210
211 /* called from corr_loop, to do the main calculations */
212 static void calc_corr(t_corr *curr,int nr,int nx,atom_id index[],rvec xc[],
213                       bool bRmCOMM,rvec com,t_calc_func *calc1,bool bTen,
214                       const output_env_t oenv)
215 {
216   int  nx0;
217   real g;
218   matrix mat;
219   rvec dcom;
220   
221   /* Check for new starting point */
222   if (curr->nlast < curr->nrestart) {
223     if ((thistime(curr) >= (curr->nlast*curr->delta_t)) && (nr==0)) {
224       memcpy(curr->x0[curr->nlast],xc,curr->ncoords*sizeof(xc[0]));
225       curr->n_offs[curr->nlast]=curr->nframes;
226       copy_rvec(com,curr->com[curr->nlast]);
227       curr->nlast++;
228     }
229   }
230
231   /* nx0 appears to be the number of new starting points,
232    * so for all starting points, call calc1. 
233    */
234   for(nx0=0; (nx0<curr->nlast); nx0++) {
235     if (bRmCOMM) {
236       rvec_sub(com,curr->com[nx0],dcom);
237     } else {
238       clear_rvec(dcom);
239     }
240     g = calc1(curr,nx,index,nx0,xc,dcom,bTen,mat,oenv);
241 #ifdef DEBUG2
242     printf("g[%d]=%g\n",nx0,g);
243 #endif
244     curr->data[nr][in_data(curr,nx0)] += g;
245     if (bTen) {
246       m_add(curr->datam[nr][in_data(curr,nx0)],mat,
247             curr->datam[nr][in_data(curr,nx0)]);
248     }
249     curr->ndata[nr][in_data(curr,nx0)]++;
250   }
251 }
252
253 /* the non-mass-weighted mean-squared displacement calcuation */
254 static real calc1_norm(t_corr *curr,int nx,atom_id index[],int nx0,rvec xc[],
255                       rvec dcom,bool bTen,matrix mat, const output_env_t oenv)
256 {
257   int  i,ix,m,m2;
258   real g,r,r2;
259   rvec rv;
260   
261   g=0.0;
262   clear_mat(mat);
263
264   for(i=0; (i<nx); i++) {
265     ix=index[i];
266     r2=0.0;
267     switch (curr->type) {
268     case NORMAL:
269       for(m=0; (m<DIM); m++) {
270         rv[m] = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
271         r2   += rv[m]*rv[m];
272         if (bTen) {
273           for(m2=0; m2<=m; m2++)
274             mat[m][m2] += rv[m]*rv[m2];
275         }
276       }
277       break;
278     case X:
279     case Y:
280     case Z:
281       r = xc[ix][curr->type-X] - curr->x0[nx0][ix][curr->type-X] - 
282                 dcom[curr->type-X];
283       r2 += r*r;
284       break;
285     case LATERAL:
286       for(m=0; (m<DIM); m++) {
287         if (m != curr->axis) {
288           r   = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
289           r2 += r*r;
290         }
291       }
292       break;
293     default:
294       gmx_fatal(FARGS,"Error: did not expect option value %d",curr->type);
295     }
296     g+=r2;
297   }
298   g/=nx;
299   msmul(mat,1.0/nx,mat);
300
301   return g;
302 }
303
304 /* calculate the com of molecules in x and put it into xa */
305 static void calc_mol_com(int nmol,int *molindex,t_block *mols,t_atoms *atoms,
306                          rvec *x,rvec *xa)
307 {
308   int  m,mol,i,d;
309   rvec xm;
310   real mass,mtot;
311
312   for(m=0; m<nmol; m++) {
313     mol = molindex[m];
314     clear_rvec(xm);
315     mtot = 0;
316     for(i=mols->index[mol]; i<mols->index[mol+1]; i++) {
317       mass = atoms->atom[i].m;
318       for(d=0; d<DIM; d++)
319         xm[d] += mass*x[i][d];
320       mtot += mass;
321     }
322     svmul(1/mtot,xm,xa[m]);
323   }
324 }
325
326 static real calc_one_mw(t_corr *curr,int ix,int nx0,rvec xc[],real *tm,
327                         rvec dcom,bool bTen,matrix mat)
328 {
329   real r2,r,mm;
330   rvec rv;
331   int  m,m2;
332   
333   mm=curr->mass[ix];
334   if (mm == 0)
335     return 0;
336   (*tm) += mm;
337   r2     = 0.0;
338   switch (curr->type) {
339   case NORMAL:
340     for(m=0; (m<DIM); m++) {
341       rv[m] = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
342       r2   += mm*rv[m]*rv[m];
343       if (bTen) {
344         for(m2=0; m2<=m; m2++)
345           mat[m][m2] += mm*rv[m]*rv[m2];
346       }
347     }
348     break;
349   case X:
350   case Y:
351   case Z:
352     r  = xc[ix][curr->type-X] - curr->x0[nx0][ix][curr->type-X] - 
353               dcom[curr->type-X];
354     r2 = mm*r*r;
355       break;
356   case LATERAL:
357     for(m=0; (m<DIM); m++) {
358       if (m != curr->axis) {
359         r   = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
360         r2 += mm*r*r;
361       }
362     }
363     break;
364   default:
365     gmx_fatal(FARGS,"Options got screwed. Did not expect value %d\n",curr->type);
366   } /* end switch */
367   return r2;
368 }
369
370 /* the normal, mass-weighted mean-squared displacement calcuation */
371 static real calc1_mw(t_corr *curr,int nx,atom_id index[],int nx0,rvec xc[],
372                      rvec dcom,bool bTen,matrix mat,const output_env_t oenv)
373 {
374   int  i;
375   real g,tm;
376
377   g=tm=0.0;
378   clear_mat(mat);
379   for(i=0; (i<nx); i++) 
380     g += calc_one_mw(curr,index[i],nx0,xc,&tm,dcom,bTen,mat);
381   
382   g/=tm;
383   if (bTen)
384     msmul(mat,1/tm,mat);
385
386   return g;
387 }
388
389 /* prepare the coordinates by removing periodic boundary crossings.
390    gnx = the number of atoms/molecules
391    index = the indices
392    xcur = the current coordinates
393    xprev = the previous coordinates
394    box = the box matrix */
395 static void prep_data(bool bMol,int gnx,atom_id index[],
396                       rvec xcur[],rvec xprev[],matrix box)
397 {
398     int  i,m,ind;
399     rvec hbox;
400
401     /* Remove periodicity */
402     for(m=0; (m<DIM); m++)
403         hbox[m]=0.5*box[m][m];
404     
405     for(i=0; (i<gnx); i++) {
406         if (bMol)
407             ind = i;
408         else
409             ind = index[i];
410
411         for(m=DIM-1; m>=0; m--) 
412         {
413             while(xcur[ind][m]-xprev[ind][m] <= -hbox[m])
414                 rvec_inc(xcur[ind],box[m]);
415             while(xcur[ind][m]-xprev[ind][m] >  hbox[m])
416                 rvec_dec(xcur[ind],box[m]);
417         }      
418     }
419 }
420
421 /* calculate the center of mass for a group 
422    gnx = the number of atoms/molecules
423    index = the indices
424    xcur = the current coordinates
425    xprev = the previous coordinates
426    box = the box matrix
427    atoms = atom data (for mass)
428    com(output) = center of mass  */
429 static void calc_com(bool bMol, int gnx, atom_id index[], 
430                      rvec xcur[],rvec xprev[],matrix box, t_atoms *atoms,
431                      rvec com)
432 {
433     int  i,m,ind;
434     real mass;
435     double tmass;
436     dvec sx;
437
438     clear_dvec(sx);
439     tmass = 0;
440     mass = 1;
441
442     prep_data(bMol, gnx, index, xcur, xprev, box);
443     for(i=0; (i<gnx); i++) 
444     {
445         if (bMol)
446             ind = i;
447         else
448             ind = index[i];
449
450
451         mass = atoms->atom[ind].m;
452         for(m=0; m<DIM; m++)
453             sx[m] += mass*xcur[ind][m];
454         tmass += mass;
455     }
456     for(m=0; m<DIM; m++)
457         com[m] = sx[m]/tmass;
458 }
459
460
461 static real calc1_mol(t_corr *curr,int nx,atom_id index[],int nx0,rvec xc[],
462                       rvec dcom,bool bTen,matrix mat, const output_env_t oenv)
463 {
464   int  i;
465   real g,tm,gtot,tt;
466
467   tt = curr->time[in_data(curr,nx0)];
468   gtot = 0;
469   tm = 0;
470   clear_mat(mat);
471   for(i=0; (i<nx); i++) {
472     g = calc_one_mw(curr,i,nx0,xc,&tm,dcom,bTen,mat);
473     /* We don't need to normalize as the mass was set to 1 */
474     gtot += g;
475     if (tt >= curr->beginfit && (curr->endfit < 0 || tt <= curr->endfit))
476       gmx_stats_add_point(curr->lsq[nx0][i],tt,g,0,0);
477   }
478   msmul(mat,1.0/nx,mat);
479
480   return gtot/nx;
481 }
482
483 void printmol(t_corr *curr,const char *fn,
484               const char *fn_pdb,int *molindex,t_topology *top,
485               rvec *x,int ePBC,matrix box, const output_env_t oenv)
486 {
487 #define NDIST 100
488   FILE  *out;
489   gmx_stats_t lsq1;
490   int   i,j;
491   real  a,b,D,Dav,D2av,VarD,sqrtD,sqrtD_max,scale;
492   t_pdbinfo *pdbinfo=NULL;
493   int   *mol2a=NULL;
494
495   out=xvgropen(fn,"Diffusion Coefficients / Molecule","Molecule","D",oenv);
496   
497   if (fn_pdb) {
498     if (top->atoms.pdbinfo == NULL)
499       snew(top->atoms.pdbinfo,top->atoms.nr);
500     pdbinfo = top->atoms.pdbinfo;
501     mol2a = top->mols.index;
502   }
503
504   Dav = D2av = 0;
505   sqrtD_max = 0;
506   for(i=0; (i<curr->nmol); i++) {
507     lsq1 = gmx_stats_init();
508     for(j=0; (j<curr->nrestart); j++) {
509       real xx,yy,dx,dy;
510       
511       while(gmx_stats_get_point(curr->lsq[j][i],&xx,&yy,&dx,&dy) == estatsOK)
512           gmx_stats_add_point(lsq1,xx,yy,dx,dy);
513     }
514     gmx_stats_get_ab(lsq1,elsqWEIGHT_NONE,&a,&b,NULL,NULL,NULL,NULL);
515     gmx_stats_done(lsq1);
516     sfree(lsq1);
517     D     = a*FACTOR/curr->dim_factor;
518     if (D < 0)
519       D   = 0;
520     Dav  += D;
521     D2av += sqr(D);
522     fprintf(out,"%10d  %10g\n",i,D);
523     if (pdbinfo) {
524       sqrtD = sqrt(D);
525       if (sqrtD > sqrtD_max)
526         sqrtD_max = sqrtD;
527       for(j=mol2a[molindex[i]]; j<mol2a[molindex[i]+1]; j++)
528         pdbinfo[j].bfac = sqrtD;
529     }
530   }
531   ffclose(out);
532   do_view(oenv,fn,"-graphtype bar");
533   
534   /* Compute variance, stddev and error */
535   Dav  /= curr->nmol;
536   D2av /= curr->nmol;
537   VarD  = D2av - sqr(Dav);
538   printf("<D> = %.4f Std. Dev. = %.4f Error = %.4f\n",
539          Dav,sqrt(VarD),sqrt(VarD/curr->nmol));
540
541   if (fn_pdb && x) {
542     scale = 1;
543     while (scale*sqrtD_max > 10)
544       scale *= 0.1;
545     while (scale*sqrtD_max < 0.1)
546       scale *= 10;
547     for(i=0; i<top->atoms.nr; i++)
548       pdbinfo[i].bfac *= scale;
549     write_sto_conf(fn_pdb,"molecular MSD",&top->atoms,x,NULL,ePBC,box);
550   }
551 }
552
553 /* this is the main loop for the correlation type functions 
554  * fx and nx are file pointers to things like read_first_x and
555  * read_next_x
556  */
557 int corr_loop(t_corr *curr,const char *fn,t_topology *top,int ePBC,
558               bool bMol,int gnx[],atom_id *index[],
559               t_calc_func *calc1,bool bTen, int *gnx_com, atom_id *index_com[],
560               real dt, real t_pdb,rvec **x_pdb,matrix box_pdb, 
561               const output_env_t oenv)
562 {
563   rvec         *x[2]; /* the coordinates to read */
564   rvec         *xa[2]; /* the coordinates to calculate displacements for */
565   rvec         com={0};
566   real         t,t_prev=0;
567   int          natoms,i,j,cur=0,maxframes=0;
568   t_trxstatus *status;
569 #define        prev (1-cur)
570   matrix       box;
571   bool         bFirst;
572   gmx_rmpbc_t  gpbc=NULL;
573
574   natoms = read_first_x(oenv,&status,fn,&curr->t0,&(x[cur]),box);
575 #ifdef DEBUG
576   fprintf(stderr,"Read %d atoms for first frame\n",natoms);
577 #endif
578   if ((gnx_com!=NULL) && natoms < top->atoms.nr) {
579     fprintf(stderr,"WARNING: The trajectory only contains part of the system (%d of %d atoms) and therefore the COM motion of only this part of the system will be removed\n",natoms,top->atoms.nr);
580   }
581
582   snew(x[prev],natoms);
583
584   if (bMol) {
585     curr->ncoords = curr->nmol;
586     snew(xa[0],curr->ncoords);
587     snew(xa[1],curr->ncoords);
588   } else {
589     curr->ncoords = natoms;
590     xa[0] = x[0];
591     xa[1] = x[1];
592   }
593
594   bFirst = TRUE;
595   t=curr->t0;
596   if (x_pdb)
597     *x_pdb = NULL;
598
599   if (bMol)
600     gpbc = gmx_rmpbc_init(&top->idef,ePBC,natoms,box);
601
602   /* the loop over all frames */
603   do 
604   {
605     if (x_pdb && ((bFirst && t_pdb < t) || 
606                   (!bFirst && 
607                    t_pdb > t - 0.5*(t - t_prev) && 
608                    t_pdb < t + 0.5*(t - t_prev)))) 
609     {
610         if (*x_pdb == NULL)
611             snew(*x_pdb,natoms);
612         for(i=0; i<natoms; i++)
613             copy_rvec(x[cur][i],(*x_pdb)[i]);
614         copy_mat(box,box_pdb);
615     }
616       
617
618     /* check whether we've reached a restart point */
619     if (bRmod(t,curr->t0,dt)) {
620       curr->nrestart++;
621   
622       srenew(curr->x0,curr->nrestart);
623       snew(curr->x0[curr->nrestart-1],curr->ncoords);
624       srenew(curr->com,curr->nrestart);
625       srenew(curr->n_offs,curr->nrestart);
626       srenew(curr->lsq,curr->nrestart);
627       snew(curr->lsq[curr->nrestart-1],curr->nmol);
628       for(i=0;i<curr->nmol;i++)
629           curr->lsq[curr->nrestart-1][i]  = gmx_stats_init();
630
631       if (debug)
632         fprintf(debug,"Extended data structures because of new restart %d\n",
633                 curr->nrestart);
634     }
635     /* create or extend the frame-based arrays */
636     if (curr->nframes >= maxframes-1) {
637       if (maxframes==0) {
638         for(i=0; (i<curr->ngrp); i++) {
639           curr->ndata[i] = NULL;
640           curr->data[i]  = NULL;
641           if (bTen)
642             curr->datam[i] = NULL;
643         }
644         curr->time = NULL;
645       }
646       maxframes+=10;
647       for(i=0; (i<curr->ngrp); i++) {
648         srenew(curr->ndata[i],maxframes);
649         srenew(curr->data[i],maxframes);
650         if (bTen)
651           srenew(curr->datam[i],maxframes);
652         for(j=maxframes-10; j<maxframes; j++) {
653           curr->ndata[i][j]=0;
654           curr->data[i][j]=0;
655           if (bTen)
656             clear_mat(curr->datam[i][j]);
657         }
658       }
659       srenew(curr->time,maxframes);
660     }
661
662     /* set the time */
663     curr->time[curr->nframes] = t - curr->t0;
664
665     /* for the first frame, the previous frame is a copy of the first frame */
666     if (bFirst) {
667       memcpy(xa[prev],xa[cur],curr->ncoords*sizeof(xa[prev][0]));
668       bFirst = FALSE;
669     }
670
671     /* make the molecules whole */
672     if (bMol)
673       gmx_rmpbc(gpbc,natoms,box,x[cur]);
674
675     /* first remove the periodic boundary condition crossings */
676     for(i=0;i<curr->ngrp;i++)
677     {
678         prep_data(bMol, gnx[i], index[i], xa[cur], xa[prev], box);
679     }
680
681     /* calculate the center of mass */
682     if (gnx_com)
683     {
684         prep_data(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box);
685         calc_com(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box,
686                  &top->atoms, com);
687     }
688
689     /* calculate the molecules' centers of masses and put them into xa */
690     if (bMol)
691         calc_mol_com(gnx[0],index[0],&top->mols,&top->atoms, x[cur],xa[cur]);
692     
693     /* loop over all groups in index file */
694     for(i=0; (i<curr->ngrp); i++) 
695     {
696       /* calculate something useful, like mean square displacements */
697       calc_corr(curr,i,gnx[i],index[i],xa[cur], (gnx_com!=NULL),com,
698                 calc1,bTen,oenv);
699     }
700     cur=prev;
701     t_prev = t;
702     
703     curr->nframes++;
704   } while (read_next_x(oenv,status,&t,natoms,x[cur],box));
705   fprintf(stderr,"\nUsed %d restart points spaced %g %s over %g %s\n\n", 
706           curr->nrestart, 
707           output_env_conv_time(oenv,dt), output_env_get_time_unit(oenv),
708           output_env_conv_time(oenv,curr->time[curr->nframes-1]), 
709           output_env_get_time_unit(oenv) );
710   
711   if (bMol)
712     gmx_rmpbc_done(gpbc);
713
714   close_trj(status);
715
716   return natoms;
717 }
718  
719 static void index_atom2mol(int *n,int *index,t_block *mols)
720 {
721   int nat,i,nmol,mol,j;
722
723   nat = *n;
724   i = 0;
725   nmol = 0;
726   mol = 0;
727   while (i < nat) {
728     while (index[i] > mols->index[mol]) {
729       mol++;
730       if (mol >= mols->nr)
731         gmx_fatal(FARGS,"Atom index out of range: %d",index[i]+1);
732     }
733     for(j=mols->index[mol]; j<mols->index[mol+1]; j++) {
734       if (i >= nat || index[i] != j)
735         gmx_fatal(FARGS,"The index group does not consist of whole molecules");
736       i++;
737     }
738     index[nmol++] = mol;
739   }
740
741   fprintf(stderr,"Split group of %d atoms into %d molecules\n",nat,nmol);
742   
743   *n = nmol;
744 }
745                             
746 void do_corr(const char *trx_file, const char *ndx_file, const char *msd_file, 
747              const char *mol_file, const char *pdb_file,real t_pdb,
748              int nrgrp, t_topology *top,int ePBC,
749              bool bTen,bool bMW,bool bRmCOMM,
750              int type,real dim_factor,int axis,
751              real dt,real beginfit,real endfit,const output_env_t oenv)
752 {
753   t_corr       *msd;
754   int          *gnx; /* the selected groups' sizes */
755   atom_id      **index; /* selected groups' indices */
756   char         **grpname; 
757   int          i,i0,i1,j,N,nat_trx;
758   real         *DD,*SigmaD,a,a2,b,r,chi2;
759   rvec         *x;
760   matrix       box;
761   int          *gnx_com=NULL; /* the COM removal group size  */
762   atom_id      **index_com=NULL; /* the COM removal group atom indices */
763   char         **grpname_com=NULL; /* the COM removal group name */
764
765   snew(gnx,nrgrp);
766   snew(index,nrgrp);
767   snew(grpname,nrgrp);
768    
769   fprintf(stderr, "\nSelect a group to calculate mean squared displacement for:\n");
770   get_index(&top->atoms,ndx_file,nrgrp,gnx,index,grpname);
771
772   if (bRmCOMM)
773   {
774       snew(gnx_com,1);
775       snew(index_com,1);
776       snew(grpname_com,1);
777
778       fprintf(stderr, "\nNow select a group for center of mass removal:\n");
779       get_index(&top->atoms, ndx_file, 1, gnx_com, index_com, grpname_com);
780   }
781   
782   if (mol_file)
783     index_atom2mol(&gnx[0],index[0],&top->mols);
784
785   msd = init_corr(nrgrp,type,axis,dim_factor,
786                   mol_file==NULL ? 0 : gnx[0],bTen,bMW,dt,top,
787                   beginfit,endfit);
788   
789   nat_trx =
790     corr_loop(msd,trx_file,top,ePBC,mol_file ? gnx[0] : 0,gnx,index,
791               (mol_file!=NULL) ? calc1_mol : (bMW ? calc1_mw : calc1_norm),
792               bTen, gnx_com, index_com, dt,t_pdb,
793               pdb_file ? &x : NULL,box,oenv);
794
795   /* Correct for the number of points */
796   for(j=0; (j<msd->ngrp); j++) {
797     for(i=0; (i<msd->nframes); i++) {
798       msd->data[j][i] /= msd->ndata[j][i];
799       if (bTen)
800         msmul(msd->datam[j][i],1.0/msd->ndata[j][i],msd->datam[j][i]);
801     }
802   }
803
804   if (mol_file) {
805     if (pdb_file && x == NULL) {
806       fprintf(stderr,"\nNo frame found need time tpdb = %g ps\n"
807               "Can not write %s\n\n",t_pdb,pdb_file);
808     }
809     i = top->atoms.nr;
810     top->atoms.nr = nat_trx;
811     printmol(msd,mol_file,pdb_file,index[0],top,x,ePBC,box,oenv);
812     top->atoms.nr = i;
813   }
814
815   DD     = NULL;
816   SigmaD = NULL;
817
818   if (beginfit == -1) {
819     i0 = (int)(0.1*(msd->nframes - 1) + 0.5);
820     beginfit = msd->time[i0];
821   } else
822   for(i0=0; i0<msd->nframes && msd->time[i0]<beginfit; i0++) 
823     ;
824
825   if (endfit == -1) {
826     i1 = (int)(0.9*(msd->nframes - 1) + 0.5) + 1;
827     endfit = msd->time[i1-1];
828   } else
829     for(i1=i0; i1<msd->nframes && msd->time[i1]<=endfit; i1++)
830                       ;
831   fprintf(stdout,"Fitting from %g to %g %s\n\n",beginfit,endfit,
832                  output_env_get_time_unit(oenv));
833
834   N = i1-i0;
835   if (N <= 2) {
836     fprintf(stdout,"Not enough points for fitting (%d).\n"
837             "Can not determine the diffusion constant.\n",N);
838   } else {
839     snew(DD,msd->ngrp);
840     snew(SigmaD,msd->ngrp);
841     for(j=0; j<msd->ngrp; j++) {
842       if (N >= 4) {
843         lsq_y_ax_b(N/2,&(msd->time[i0]),&(msd->data[j][i0]),&a,&b,&r,&chi2);
844         lsq_y_ax_b(N/2,&(msd->time[i0+N/2]),&(msd->data[j][i0+N/2]),&a2,&b,&r,&chi2);
845         SigmaD[j] = fabs(a-a2);
846       } else
847         SigmaD[j] = 0;
848       lsq_y_ax_b(N,&(msd->time[i0]),&(msd->data[j][i0]),&(DD[j]),&b,&r,&chi2);
849       DD[j]     *= FACTOR/msd->dim_factor;
850       SigmaD[j] *= FACTOR/msd->dim_factor;
851       if (DD[j] > 0.01 && DD[j] < 1e4)
852           fprintf(stdout,"D[%10s] %.4f (+/- %.4f) 1e-5 cm^2/s\n",
853                   grpname[j],DD[j],SigmaD[j]);
854       else
855           fprintf(stdout,"D[%10s] %.4g (+/- %.4g) 1e-5 cm^2/s\n",
856                   grpname[j],DD[j], SigmaD[j]);
857     }
858   }
859   /* Print mean square displacement */
860   corr_print(msd,bTen,msd_file,
861              "Mean Square Displacement",
862              "MSD (nm\\S2\\N)",
863              msd->time[msd->nframes-1],beginfit,endfit,DD,SigmaD,grpname,oenv);
864 }
865
866 int gmx_msd(int argc,char *argv[])
867 {
868   const char *desc[] = {
869     "g_msd computes the mean square displacement (MSD) of atoms from",
870     "a set of initial positions. This provides an easy way to compute",
871     "the diffusion constant using the Einstein relation.",
872     "The time between the reference points for the MSD calculation",
873     "is set with [TT]-trestart[tt].",
874     "The diffusion constant is calculated by least squares fitting a",
875     "straight line (D*t + c) through the MSD(t) from [TT]-beginfit[tt] to",
876     "[TT]-endfit[tt] (note that t is time from the reference positions,",
877     "not simulation time). An error estimate given, which is the difference",
878     "of the diffusion coefficients obtained from fits over the two halves",
879     "of the fit interval.[PAR]",
880     "There are three, mutually exclusive, options to determine different",
881     "types of mean square displacement: [TT]-type[tt], [TT]-lateral[tt]",
882     "and [TT]-ten[tt]. Option [TT]-ten[tt] writes the full MSD tensor for",
883     "each group, the order in the output is: trace xx yy zz yx zx zy.[PAR]",
884     "If [TT]-mol[tt] is set, g_msd plots the MSD for individual molecules: ",
885     "for each individual molecule a diffusion constant is computed for ",
886     "its center of mass. The chosen index group will be split into ",
887     "molecules.[PAR]",
888     "The default way to calculate a MSD is by using mass-weighted averages.",
889     "This can be turned off with [TT]-nomw[tt].[PAR]",
890     "With the option [TT]-rmcomm[tt], the center of mass motion of a ",
891     "specific group can be removed. For trajectories produced with ",
892     "GROMACS this is usually not necessary, ",
893     "as mdrun usually already removes the center of mass motion.",
894     "When you use this option be sure that the whole system is stored",
895     "in the trajectory file.[PAR]",
896     "The diffusion coefficient is determined by linear regression of the MSD,",
897     "where, unlike for the normal output of D, the times are weighted",
898     "according to the number of reference points, i.e. short times have",
899     "a higher weight. Also when [TT]-beginfit[tt]=-1,fitting starts at 10%",
900     "and when [TT]-endfit[tt]=-1, fitting goes to 90%.",
901     "Using this option one also gets an accurate error estimate",
902     "based on the statistics between individual molecules.",
903     "Note that this diffusion coefficient and error estimate are only",
904     "accurate when the MSD is completely linear between",
905     "[TT]-beginfit[tt] and [TT]-endfit[tt].[PAR]",
906     "Option [TT]-pdb[tt] writes a pdb file with the coordinates of the frame",
907     "at time [TT]-tpdb[tt] with in the B-factor field the square root of",
908     "the diffusion coefficient of the molecule.",
909     "This option implies option [TT]-mol[tt]."
910   };
911   static const char *normtype[]= { NULL,"no","x","y","z",NULL };
912   static const char *axtitle[] = { NULL,"no","x","y","z",NULL };
913   static int  ngroup     = 1;
914   static real dt         = 10; 
915   static real t_pdb      = 0; 
916   static real beginfit   = -1; 
917   static real endfit     = -1; 
918   static bool bTen       = FALSE;
919   static bool bMW        = TRUE;
920   static bool bRmCOMM    = FALSE;
921   t_pargs pa[] = {
922     { "-type",    FALSE, etENUM, {normtype},
923       "Compute diffusion coefficient in one direction" },
924     { "-lateral", FALSE, etENUM, {axtitle}, 
925       "Calculate the lateral diffusion in a plane perpendicular to" },
926     { "-ten",      FALSE, etBOOL, {&bTen},
927       "Calculate the full tensor" },
928     { "-ngroup",  FALSE, etINT,  {&ngroup},
929       "Number of groups to calculate MSD for" },
930     { "-mw",      FALSE, etBOOL, {&bMW},
931       "Mass weighted MSD" },
932     { "-rmcomm",      FALSE, etBOOL, {&bRmCOMM},
933       "Remove center of mass motion" },
934     { "-tpdb",FALSE, etTIME, {&t_pdb},
935       "The frame to use for option -pdb (%t)" },
936     { "-trestart",FALSE, etTIME, {&dt},
937       "Time between restarting points in trajectory (%t)" },
938     { "-beginfit",FALSE, etTIME, {&beginfit},
939       "Start time for fitting the MSD (%t), -1 is 10%" },
940     { "-endfit",FALSE, etTIME, {&endfit},
941       "End time for fitting the MSD (%t), -1 is 90%" }
942   };
943
944   t_filenm fnm[] = { 
945     { efTRX, NULL, NULL,  ffREAD },
946     { efTPS, NULL, NULL,  ffREAD }, 
947     { efNDX, NULL, NULL,  ffOPTRD },
948     { efXVG, NULL, "msd", ffWRITE },
949     { efXVG, "-mol", "diff_mol",ffOPTWR },
950     { efPDB, "-pdb", "diff_mol", ffOPTWR }
951   };
952 #define NFILE asize(fnm)
953
954   t_topology  top;
955   int         ePBC;
956   matrix      box;
957   char        title[256];
958   const char  *trx_file, *tps_file, *ndx_file, *msd_file, *mol_file, *pdb_file;
959   rvec        *xdum;
960   bool        bTop;
961   int         axis,type;
962   real        dim_factor;
963   output_env_t oenv;
964
965   CopyRight(stderr,argv[0]);
966
967   parse_common_args(&argc,argv,
968                     PCA_CAN_VIEW | PCA_CAN_BEGIN | PCA_CAN_END | PCA_TIME_UNIT | PCA_BE_NICE,
969                     NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL,&oenv);
970   trx_file = ftp2fn_null(efTRX,NFILE,fnm);
971   tps_file = ftp2fn_null(efTPS,NFILE,fnm);
972   ndx_file = ftp2fn_null(efNDX,NFILE,fnm);
973   msd_file = ftp2fn_null(efXVG,NFILE,fnm);
974   pdb_file = opt2fn_null("-pdb",NFILE,fnm);
975   if (pdb_file)
976     mol_file = opt2fn("-mol",NFILE,fnm);
977   else
978     mol_file = opt2fn_null("-mol",NFILE,fnm);
979   
980   if (ngroup < 1)
981     gmx_fatal(FARGS,"Must have at least 1 group (now %d)",ngroup);
982   if (mol_file && ngroup > 1)
983     gmx_fatal(FARGS,"With molecular msd can only have 1 group (now %d)",
984               ngroup);
985
986
987   if (mol_file) {
988     bMW  = TRUE;
989     fprintf(stderr,"Calculating diffusion coefficients for molecules.\n");
990   }
991
992   if (normtype[0][0]!='n') {
993     type = normtype[0][0] - 'x' + X; /* See defines above */
994     dim_factor = 2.0;
995   }
996   else {
997     type       = NORMAL;
998     dim_factor = 6.0;
999   }
1000   if ((type==NORMAL) && (axtitle[0][0]!='n')) {
1001     type=LATERAL;
1002     dim_factor = 4.0;
1003     axis = (axtitle[0][0] - 'x'); /* See defines above */
1004   }
1005   else
1006     axis = 0;
1007
1008   if (bTen && type != NORMAL)
1009     gmx_fatal(FARGS,"Can only calculate the full tensor for 3D msd");
1010
1011   bTop = read_tps_conf(tps_file,title,&top,&ePBC,&xdum,NULL,box,bMW||bRmCOMM); 
1012   if (mol_file && !bTop)
1013     gmx_fatal(FARGS,
1014               "Could not read a topology from %s. Try a tpr file instead.",
1015               tps_file);
1016     
1017   do_corr(trx_file,ndx_file,msd_file,mol_file,pdb_file,t_pdb,ngroup,
1018           &top,ePBC,bTen,bMW,bRmCOMM,type,dim_factor,axis,dt,beginfit,endfit,
1019           oenv);
1020   
1021   view_all(oenv,NFILE, fnm);
1022   
1023   thanx(stderr);
1024   
1025   return 0;
1026 }