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