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