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