Move read_tps_conf() to confio.h
[alexxy/gromacs.git] / src / gromacs / gmxana / gmx_msd.c
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
5  * Copyright (c) 2001-2004, The GROMACS development team.
6  * Copyright (c) 2013,2014,2015, 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 <math.h>
40 #include <string.h>
41
42 #include "gromacs/commandline/pargs.h"
43 #include "gromacs/fileio/confio.h"
44 #include "gromacs/fileio/trxio.h"
45 #include "gromacs/fileio/xvgr.h"
46 #include "gromacs/gmxana/gmx_ana.h"
47 #include "gromacs/gmxana/gstat.h"
48 #include "gromacs/legacyheaders/macros.h"
49 #include "gromacs/legacyheaders/typedefs.h"
50 #include "gromacs/legacyheaders/viewit.h"
51 #include "gromacs/math/utilities.h"
52 #include "gromacs/math/vec.h"
53 #include "gromacs/pbcutil/rmpbc.h"
54 #include "gromacs/statistics/statistics.h"
55 #include "gromacs/topology/index.h"
56 #include "gromacs/utility/fatalerror.h"
57 #include "gromacs/utility/futil.h"
58 #include "gromacs/utility/smalloc.h"
59
60 #define FACTOR  1000.0  /* Convert nm^2/ps to 10e-5 cm^2/s */
61 /* NORMAL = total diffusion coefficient (default). X,Y,Z is diffusion
62    coefficient in X,Y,Z direction. LATERAL is diffusion coefficient in
63    plane perpendicular to axis
64  */
65 typedef enum {
66     NOT_USED, NORMAL, X, Y, Z, LATERAL
67 } 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 *curr, int nx, atom_id index[], int nx0, rvec xc[],
98                           rvec dcom, gmx_bool bTen, matrix mat);
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     {
136         snew(curr->datam, nrgrp);
137     }
138     for (i = 0; (i < nrgrp); i++)
139     {
140         curr->ndata[i] = NULL;
141         curr->data[i]  = NULL;
142         if (bTen)
143         {
144             curr->datam[i] = NULL;
145         }
146     }
147     curr->time = NULL;
148     curr->lsq  = NULL;
149     curr->nmol = nmol;
150     if (curr->nmol > 0)
151     {
152         snew(curr->mass, curr->nmol);
153         for (i = 0; i < curr->nmol; i++)
154         {
155             curr->mass[i] = 1;
156         }
157     }
158     else
159     {
160         if (bMass)
161         {
162             atoms = &top->atoms;
163             snew(curr->mass, atoms->nr);
164             for (i = 0; (i < atoms->nr); i++)
165             {
166                 curr->mass[i] = atoms->atom[i].m;
167             }
168         }
169     }
170
171     return curr;
172 }
173
174 static void corr_print(t_corr *curr, gmx_bool bTen, const char *fn, const char *title,
175                        const char *yaxis,
176                        real msdtime, real beginfit, real endfit,
177                        real *DD, real *SigmaD, char *grpname[],
178                        const output_env_t oenv)
179 {
180     FILE *out;
181     int   i, j;
182
183     out = xvgropen(fn, title, output_env_get_xvgr_tlabel(oenv), yaxis, oenv);
184     if (DD)
185     {
186         fprintf(out, "# MSD gathered over %g %s with %d restarts\n",
187                 msdtime, output_env_get_time_unit(oenv), curr->nrestart);
188         fprintf(out, "# Diffusion constants fitted from time %g to %g %s\n",
189                 beginfit, endfit, output_env_get_time_unit(oenv));
190         for (i = 0; i < curr->ngrp; i++)
191         {
192             fprintf(out, "# D[%10s] = %.4f (+/- %.4f) (1e-5 cm^2/s)\n",
193                     grpname[i], DD[i], SigmaD[i]);
194         }
195     }
196     for (i = 0; i < curr->nframes; i++)
197     {
198         fprintf(out, "%10g", output_env_conv_time(oenv, curr->time[i]));
199         for (j = 0; j < curr->ngrp; j++)
200         {
201             fprintf(out, "  %10g", curr->data[j][i]);
202             if (bTen)
203             {
204                 fprintf(out, " %10g %10g %10g %10g %10g %10g",
205                         curr->datam[j][i][XX][XX],
206                         curr->datam[j][i][YY][YY],
207                         curr->datam[j][i][ZZ][ZZ],
208                         curr->datam[j][i][YY][XX],
209                         curr->datam[j][i][ZZ][XX],
210                         curr->datam[j][i][ZZ][YY]);
211             }
212         }
213         fprintf(out, "\n");
214     }
215     xvgrclose(out);
216 }
217
218 /* called from corr_loop, to do the main calculations */
219 static void calc_corr(t_corr *curr, int nr, int nx, atom_id index[], rvec xc[],
220                       gmx_bool bRmCOMM, rvec com, t_calc_func *calc1, gmx_bool bTen)
221 {
222     int    nx0;
223     real   g;
224     matrix mat;
225     rvec   dcom;
226
227     /* Check for new starting point */
228     if (curr->nlast < curr->nrestart)
229     {
230         if ((thistime(curr) >= (curr->nlast*curr->delta_t)) && (nr == 0))
231         {
232             memcpy(curr->x0[curr->nlast], xc, curr->ncoords*sizeof(xc[0]));
233             curr->n_offs[curr->nlast] = curr->nframes;
234             copy_rvec(com, curr->com[curr->nlast]);
235             curr->nlast++;
236         }
237     }
238
239     /* nx0 appears to be the number of new starting points,
240      * so for all starting points, call calc1.
241      */
242     for (nx0 = 0; (nx0 < curr->nlast); nx0++)
243     {
244         if (bRmCOMM)
245         {
246             rvec_sub(com, curr->com[nx0], dcom);
247         }
248         else
249         {
250             clear_rvec(dcom);
251         }
252         g = calc1(curr, nx, index, nx0, xc, dcom, bTen, mat);
253 #ifdef DEBUG2
254         printf("g[%d]=%g\n", nx0, g);
255 #endif
256         curr->data[nr][in_data(curr, nx0)] += g;
257         if (bTen)
258         {
259             m_add(curr->datam[nr][in_data(curr, nx0)], mat,
260                   curr->datam[nr][in_data(curr, nx0)]);
261         }
262         curr->ndata[nr][in_data(curr, nx0)]++;
263     }
264 }
265
266 /* the non-mass-weighted mean-squared displacement calcuation */
267 static real calc1_norm(t_corr *curr, int nx, atom_id index[], int nx0, rvec xc[],
268                        rvec dcom, gmx_bool bTen, matrix mat)
269 {
270     int  i, ix, m, m2;
271     real g, r, r2;
272     rvec rv;
273
274     g = 0.0;
275     clear_mat(mat);
276
277     for (i = 0; (i < nx); i++)
278     {
279         ix = index[i];
280         r2 = 0.0;
281         switch (curr->type)
282         {
283             case NORMAL:
284                 for (m = 0; (m < DIM); m++)
285                 {
286                     rv[m] = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
287                     r2   += rv[m]*rv[m];
288                     if (bTen)
289                     {
290                         for (m2 = 0; m2 <= m; m2++)
291                         {
292                             mat[m][m2] += rv[m]*rv[m2];
293                         }
294                     }
295                 }
296                 break;
297             case X:
298             case Y:
299             case Z:
300                 r = xc[ix][curr->type-X] - curr->x0[nx0][ix][curr->type-X] -
301                     dcom[curr->type-X];
302                 r2 += r*r;
303                 break;
304             case LATERAL:
305                 for (m = 0; (m < DIM); m++)
306                 {
307                     if (m != curr->axis)
308                     {
309                         r   = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
310                         r2 += r*r;
311                     }
312                 }
313                 break;
314             default:
315                 gmx_fatal(FARGS, "Error: did not expect option value %d", curr->type);
316         }
317         g += r2;
318     }
319     g /= nx;
320     msmul(mat, 1.0/nx, mat);
321
322     return g;
323 }
324
325 /* calculate the com of molecules in x and put it into xa */
326 static void calc_mol_com(int nmol, int *molindex, t_block *mols, t_atoms *atoms,
327                          rvec *x, rvec *xa)
328 {
329     int  m, mol, i, d;
330     rvec xm;
331     real mass, mtot;
332
333     for (m = 0; m < nmol; m++)
334     {
335         mol = molindex[m];
336         clear_rvec(xm);
337         mtot = 0;
338         for (i = mols->index[mol]; i < mols->index[mol+1]; i++)
339         {
340             mass = atoms->atom[i].m;
341             for (d = 0; d < DIM; d++)
342             {
343                 xm[d] += mass*x[i][d];
344             }
345             mtot += mass;
346         }
347         svmul(1/mtot, xm, xa[m]);
348     }
349 }
350
351 static real calc_one_mw(t_corr *curr, int ix, int nx0, rvec xc[], real *tm,
352                         rvec dcom, gmx_bool bTen, matrix mat)
353 {
354     real r2, r, mm;
355     rvec rv;
356     int  m, m2;
357
358     mm = curr->mass[ix];
359     if (mm == 0)
360     {
361         return 0;
362     }
363     (*tm) += mm;
364     r2     = 0.0;
365     switch (curr->type)
366     {
367         case NORMAL:
368             for (m = 0; (m < DIM); m++)
369             {
370                 rv[m] = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
371                 r2   += mm*rv[m]*rv[m];
372                 if (bTen)
373                 {
374                     for (m2 = 0; m2 <= m; m2++)
375                     {
376                         mat[m][m2] += mm*rv[m]*rv[m2];
377                     }
378                 }
379             }
380             break;
381         case X:
382         case Y:
383         case Z:
384             r  = xc[ix][curr->type-X] - curr->x0[nx0][ix][curr->type-X] -
385                 dcom[curr->type-X];
386             r2 = mm*r*r;
387             break;
388         case LATERAL:
389             for (m = 0; (m < DIM); m++)
390             {
391                 if (m != curr->axis)
392                 {
393                     r   = xc[ix][m] - curr->x0[nx0][ix][m] - dcom[m];
394                     r2 += mm*r*r;
395                 }
396             }
397             break;
398         default:
399             gmx_fatal(FARGS, "Options got screwed. Did not expect value %d\n", curr->type);
400     } /* end switch */
401     return r2;
402 }
403
404 /* the normal, mass-weighted mean-squared displacement calcuation */
405 static real calc1_mw(t_corr *curr, int nx, atom_id index[], int nx0, rvec xc[],
406                      rvec dcom, gmx_bool bTen, matrix mat)
407 {
408     int  i;
409     real g, tm;
410
411     g = tm = 0.0;
412     clear_mat(mat);
413     for (i = 0; (i < nx); i++)
414     {
415         g += calc_one_mw(curr, index[i], nx0, xc, &tm, dcom, bTen, mat);
416     }
417
418     g /= tm;
419     if (bTen)
420     {
421         msmul(mat, 1/tm, mat);
422     }
423
424     return g;
425 }
426
427 /* prepare the coordinates by removing periodic boundary crossings.
428    gnx = the number of atoms/molecules
429    index = the indices
430    xcur = the current coordinates
431    xprev = the previous coordinates
432    box = the box matrix */
433 static void prep_data(gmx_bool bMol, int gnx, atom_id index[],
434                       rvec xcur[], rvec xprev[], matrix box)
435 {
436     int  i, m, ind;
437     rvec hbox;
438
439     /* Remove periodicity */
440     for (m = 0; (m < DIM); m++)
441     {
442         hbox[m] = 0.5*box[m][m];
443     }
444
445     for (i = 0; (i < gnx); i++)
446     {
447         if (bMol)
448         {
449             ind = i;
450         }
451         else
452         {
453             ind = index[i];
454         }
455
456         for (m = DIM-1; m >= 0; m--)
457         {
458             if (hbox[m] == 0)
459             {
460                 continue;
461             }
462             while (xcur[ind][m]-xprev[ind][m] <= -hbox[m])
463             {
464                 rvec_inc(xcur[ind], box[m]);
465             }
466             while (xcur[ind][m]-xprev[ind][m] >  hbox[m])
467             {
468                 rvec_dec(xcur[ind], box[m]);
469             }
470         }
471     }
472 }
473
474 /* calculate the center of mass for a group
475    gnx = the number of atoms/molecules
476    index = the indices
477    xcur = the current coordinates
478    xprev = the previous coordinates
479    box = the box matrix
480    atoms = atom data (for mass)
481    com(output) = center of mass  */
482 static void calc_com(gmx_bool bMol, int gnx, atom_id index[],
483                      rvec xcur[], rvec xprev[], matrix box, t_atoms *atoms,
484                      rvec com)
485 {
486     int    i, m, ind;
487     real   mass;
488     double tmass;
489     dvec   sx;
490
491     clear_dvec(sx);
492     tmass = 0;
493     mass  = 1;
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, atom_id 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 void printmol(t_corr *curr, const char *fn,
548               const char *fn_pdb, int *molindex, t_topology *top,
549               rvec *x, int ePBC, matrix box, const 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     int        *mol2a   = NULL;
558
559     out = xvgropen(fn, "Diffusion Coefficients / Molecule", "Molecule", "D", oenv);
560
561     if (fn_pdb)
562     {
563         if (top->atoms.pdbinfo == NULL)
564         {
565             snew(top->atoms.pdbinfo, top->atoms.nr);
566         }
567         pdbinfo = top->atoms.pdbinfo;
568         mol2a   = top->mols.index;
569     }
570
571     Dav       = D2av = 0;
572     sqrtD_max = 0;
573     for (i = 0; (i < curr->nmol); i++)
574     {
575         lsq1 = gmx_stats_init();
576         for (j = 0; (j < curr->nrestart); j++)
577         {
578             real xx, yy, dx, dy;
579
580             while (gmx_stats_get_point(curr->lsq[j][i], &xx, &yy, &dx, &dy, 0) == estatsOK)
581             {
582                 gmx_stats_add_point(lsq1, xx, yy, dx, dy);
583             }
584         }
585         gmx_stats_get_ab(lsq1, elsqWEIGHT_NONE, &a, &b, NULL, NULL, NULL, NULL);
586         gmx_stats_done(lsq1);
587         sfree(lsq1);
588         D     = a*FACTOR/curr->dim_factor;
589         if (D < 0)
590         {
591             D   = 0;
592         }
593         Dav  += D;
594         D2av += sqr(D);
595         fprintf(out, "%10d  %10g\n", i, D);
596         if (pdbinfo)
597         {
598             sqrtD = sqrt(D);
599             if (sqrtD > sqrtD_max)
600             {
601                 sqrtD_max = sqrtD;
602             }
603             for (j = mol2a[molindex[i]]; j < mol2a[molindex[i]+1]; j++)
604             {
605                 pdbinfo[j].bfac = sqrtD;
606             }
607         }
608     }
609     xvgrclose(out);
610     do_view(oenv, fn, "-graphtype bar");
611
612     /* Compute variance, stddev and error */
613     Dav  /= curr->nmol;
614     D2av /= curr->nmol;
615     VarD  = D2av - sqr(Dav);
616     printf("<D> = %.4f Std. Dev. = %.4f Error = %.4f\n",
617            Dav, sqrt(VarD), sqrt(VarD/curr->nmol));
618
619     if (fn_pdb && x)
620     {
621         scale = 1;
622         while (scale*sqrtD_max > 10)
623         {
624             scale *= 0.1;
625         }
626         while (scale*sqrtD_max < 0.1)
627         {
628             scale *= 10;
629         }
630         for (i = 0; i < top->atoms.nr; i++)
631         {
632             pdbinfo[i].bfac *= scale;
633         }
634         write_sto_conf(fn_pdb, "molecular MSD", &top->atoms, x, NULL, ePBC, box);
635     }
636 }
637
638 /* this is the main loop for the correlation type functions
639  * fx and nx are file pointers to things like read_first_x and
640  * read_next_x
641  */
642 int corr_loop(t_corr *curr, const char *fn, t_topology *top, int ePBC,
643               gmx_bool bMol, int gnx[], atom_id *index[],
644               t_calc_func *calc1, gmx_bool bTen, int *gnx_com, atom_id *index_com[],
645               real dt, real t_pdb, rvec **x_pdb, matrix box_pdb,
646               const output_env_t oenv)
647 {
648     rvec            *x[2];  /* the coordinates to read */
649     rvec            *xa[2]; /* the coordinates to calculate displacements for */
650     rvec             com = {0};
651     real             t, t_prev = 0;
652     int              natoms, i, j, cur = 0, maxframes = 0;
653     t_trxstatus     *status;
654 #define        prev (1-cur)
655     matrix           box;
656     gmx_bool         bFirst;
657     gmx_rmpbc_t      gpbc = NULL;
658
659     natoms = read_first_x(oenv, &status, fn, &curr->t0, &(x[cur]), box);
660 #ifdef DEBUG
661     fprintf(stderr, "Read %d atoms for first frame\n", natoms);
662 #endif
663     if ((gnx_com != NULL) && natoms < top->atoms.nr)
664     {
665         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);
666     }
667
668     snew(x[prev], natoms);
669
670     if (bMol)
671     {
672         curr->ncoords = curr->nmol;
673         snew(xa[0], curr->ncoords);
674         snew(xa[1], curr->ncoords);
675     }
676     else
677     {
678         curr->ncoords = natoms;
679         xa[0]         = x[0];
680         xa[1]         = x[1];
681     }
682
683     bFirst = TRUE;
684     t      = curr->t0;
685     if (x_pdb)
686     {
687         *x_pdb = NULL;
688     }
689
690     if (bMol)
691     {
692         gpbc = gmx_rmpbc_init(&top->idef, ePBC, natoms);
693     }
694
695     /* the loop over all frames */
696     do
697     {
698         if (x_pdb && ((bFirst && t_pdb < t) ||
699                       (!bFirst &&
700                        t_pdb > t - 0.5*(t - t_prev) &&
701                        t_pdb < t + 0.5*(t - t_prev))))
702         {
703             if (*x_pdb == NULL)
704             {
705                 snew(*x_pdb, natoms);
706             }
707             for (i = 0; i < natoms; i++)
708             {
709                 copy_rvec(x[cur][i], (*x_pdb)[i]);
710             }
711             copy_mat(box, box_pdb);
712         }
713
714
715         /* check whether we've reached a restart point */
716         if (bRmod(t, curr->t0, dt))
717         {
718             curr->nrestart++;
719
720             srenew(curr->x0, curr->nrestart);
721             snew(curr->x0[curr->nrestart-1], curr->ncoords);
722             srenew(curr->com, curr->nrestart);
723             srenew(curr->n_offs, curr->nrestart);
724             srenew(curr->lsq, curr->nrestart);
725             snew(curr->lsq[curr->nrestart-1], curr->nmol);
726             for (i = 0; i < curr->nmol; i++)
727             {
728                 curr->lsq[curr->nrestart-1][i]  = gmx_stats_init();
729             }
730
731             if (debug)
732             {
733                 fprintf(debug, "Extended data structures because of new restart %d\n",
734                         curr->nrestart);
735             }
736         }
737         /* create or extend the frame-based arrays */
738         if (curr->nframes >= maxframes-1)
739         {
740             if (maxframes == 0)
741             {
742                 for (i = 0; (i < curr->ngrp); i++)
743                 {
744                     curr->ndata[i] = NULL;
745                     curr->data[i]  = NULL;
746                     if (bTen)
747                     {
748                         curr->datam[i] = NULL;
749                     }
750                 }
751                 curr->time = NULL;
752             }
753             maxframes += 10;
754             for (i = 0; (i < curr->ngrp); i++)
755             {
756                 srenew(curr->ndata[i], maxframes);
757                 srenew(curr->data[i], maxframes);
758                 if (bTen)
759                 {
760                     srenew(curr->datam[i], maxframes);
761                 }
762                 for (j = maxframes-10; j < maxframes; j++)
763                 {
764                     curr->ndata[i][j] = 0;
765                     curr->data[i][j]  = 0;
766                     if (bTen)
767                     {
768                         clear_mat(curr->datam[i][j]);
769                     }
770                 }
771             }
772             srenew(curr->time, maxframes);
773         }
774
775         /* set the time */
776         curr->time[curr->nframes] = t - curr->t0;
777
778         /* for the first frame, the previous frame is a copy of the first frame */
779         if (bFirst)
780         {
781             memcpy(xa[prev], xa[cur], curr->ncoords*sizeof(xa[prev][0]));
782             bFirst = FALSE;
783         }
784
785         /* make the molecules whole */
786         if (bMol)
787         {
788             gmx_rmpbc(gpbc, natoms, box, x[cur]);
789         }
790
791         /* calculate the molecules' centers of masses and put them into xa */
792         if (bMol)
793         {
794             calc_mol_com(gnx[0], index[0], &top->mols, &top->atoms, x[cur], xa[cur]);
795         }
796
797         /* first remove the periodic boundary condition crossings */
798         for (i = 0; i < curr->ngrp; i++)
799         {
800             prep_data(bMol, gnx[i], index[i], xa[cur], xa[prev], box);
801         }
802
803         /* calculate the center of mass */
804         if (gnx_com)
805         {
806             prep_data(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box);
807             calc_com(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box,
808                      &top->atoms, com);
809         }
810
811         /* loop over all groups in index file */
812         for (i = 0; (i < curr->ngrp); i++)
813         {
814             /* calculate something useful, like mean square displacements */
815             calc_corr(curr, i, gnx[i], index[i], xa[cur], (gnx_com != NULL), com,
816                       calc1, bTen);
817         }
818         cur    = prev;
819         t_prev = t;
820
821         curr->nframes++;
822     }
823     while (read_next_x(oenv, status, &t, x[cur], box));
824     fprintf(stderr, "\nUsed %d restart points spaced %g %s over %g %s\n\n",
825             curr->nrestart,
826             output_env_conv_time(oenv, dt), output_env_get_time_unit(oenv),
827             output_env_conv_time(oenv, curr->time[curr->nframes-1]),
828             output_env_get_time_unit(oenv) );
829
830     if (bMol)
831     {
832         gmx_rmpbc_done(gpbc);
833     }
834
835     close_trj(status);
836
837     return natoms;
838 }
839
840 static void index_atom2mol(int *n, int *index, t_block *mols)
841 {
842     int nat, i, nmol, mol, j;
843
844     nat  = *n;
845     i    = 0;
846     nmol = 0;
847     mol  = 0;
848     while (i < nat)
849     {
850         while (index[i] > mols->index[mol])
851         {
852             mol++;
853             if (mol >= mols->nr)
854             {
855                 gmx_fatal(FARGS, "Atom index out of range: %d", index[i]+1);
856             }
857         }
858         for (j = mols->index[mol]; j < mols->index[mol+1]; j++)
859         {
860             if (i >= nat || index[i] != j)
861             {
862                 gmx_fatal(FARGS, "The index group does not consist of whole molecules");
863             }
864             i++;
865         }
866         index[nmol++] = mol;
867     }
868
869     fprintf(stderr, "Split group of %d atoms into %d molecules\n", nat, nmol);
870
871     *n = nmol;
872 }
873
874 void do_corr(const char *trx_file, const char *ndx_file, const char *msd_file,
875              const char *mol_file, const char *pdb_file, real t_pdb,
876              int nrgrp, t_topology *top, int ePBC,
877              gmx_bool bTen, gmx_bool bMW, gmx_bool bRmCOMM,
878              int type, real dim_factor, int axis,
879              real dt, real beginfit, real endfit, const output_env_t oenv)
880 {
881     t_corr        *msd;
882     int           *gnx;   /* the selected groups' sizes */
883     atom_id      **index; /* selected groups' indices */
884     char         **grpname;
885     int            i, i0, i1, j, N, nat_trx;
886     real          *DD, *SigmaD, a, a2, b, r, chi2;
887     rvec          *x;
888     matrix         box;
889     int           *gnx_com     = NULL; /* the COM removal group size  */
890     atom_id      **index_com   = NULL; /* the COM removal group atom indices */
891     char         **grpname_com = NULL; /* the COM removal group name */
892
893     snew(gnx, nrgrp);
894     snew(index, nrgrp);
895     snew(grpname, nrgrp);
896
897     fprintf(stderr, "\nSelect a group to calculate mean squared displacement for:\n");
898     get_index(&top->atoms, ndx_file, nrgrp, gnx, index, grpname);
899
900     if (bRmCOMM)
901     {
902         snew(gnx_com, 1);
903         snew(index_com, 1);
904         snew(grpname_com, 1);
905
906         fprintf(stderr, "\nNow select a group for center of mass removal:\n");
907         get_index(&top->atoms, ndx_file, 1, gnx_com, index_com, grpname_com);
908     }
909
910     if (mol_file)
911     {
912         index_atom2mol(&gnx[0], index[0], &top->mols);
913     }
914
915     msd = init_corr(nrgrp, type, axis, dim_factor,
916                     mol_file == NULL ? 0 : gnx[0], bTen, bMW, dt, top,
917                     beginfit, endfit);
918
919     nat_trx =
920         corr_loop(msd, trx_file, top, ePBC, mol_file ? gnx[0] : 0, gnx, index,
921                   (mol_file != NULL) ? calc1_mol : (bMW ? calc1_mw : calc1_norm),
922                   bTen, gnx_com, index_com, dt, t_pdb,
923                   pdb_file ? &x : NULL, box, oenv);
924
925     /* Correct for the number of points */
926     for (j = 0; (j < msd->ngrp); j++)
927     {
928         for (i = 0; (i < msd->nframes); i++)
929         {
930             msd->data[j][i] /= msd->ndata[j][i];
931             if (bTen)
932             {
933                 msmul(msd->datam[j][i], 1.0/msd->ndata[j][i], msd->datam[j][i]);
934             }
935         }
936     }
937
938     if (mol_file)
939     {
940         if (pdb_file && x == NULL)
941         {
942             fprintf(stderr, "\nNo frame found need time tpdb = %g ps\n"
943                     "Can not write %s\n\n", t_pdb, pdb_file);
944         }
945         i             = top->atoms.nr;
946         top->atoms.nr = nat_trx;
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       = (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     = (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] = fabs(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     char            title[256];
1119     const char     *trx_file, *tps_file, *ndx_file, *msd_file, *mol_file, *pdb_file;
1120     rvec           *xdum;
1121     gmx_bool        bTop;
1122     int             axis, type;
1123     real            dim_factor;
1124     output_env_t    oenv;
1125
1126     if (!parse_common_args(&argc, argv,
1127                            PCA_CAN_VIEW | PCA_CAN_BEGIN | PCA_CAN_END | PCA_TIME_UNIT,
1128                            NFILE, fnm, asize(pa), pa, asize(desc), desc, 0, NULL, &oenv))
1129     {
1130         return 0;
1131     }
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     return 0;
1204 }