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