8cbd28a23ee918de0913931d91a131e79b96a0be
[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,2017,2018,2019, 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/compat/make_unique.h"
45 #include "gromacs/fileio/confio.h"
46 #include "gromacs/fileio/trxio.h"
47 #include "gromacs/fileio/xvgr.h"
48 #include "gromacs/gmxana/gmx_ana.h"
49 #include "gromacs/gmxana/gstat.h"
50 #include "gromacs/math/functions.h"
51 #include "gromacs/math/utilities.h"
52 #include "gromacs/math/vec.h"
53 #include "gromacs/math/vectypes.h"
54 #include "gromacs/pbcutil/rmpbc.h"
55 #include "gromacs/statistics/statistics.h"
56 #include "gromacs/topology/index.h"
57 #include "gromacs/topology/topology.h"
58 #include "gromacs/utility/arraysize.h"
59 #include "gromacs/utility/fatalerror.h"
60 #include "gromacs/utility/futil.h"
61 #include "gromacs/utility/gmxassert.h"
62 #include "gromacs/utility/smalloc.h"
63
64 static constexpr double diffusionConversionFactor =  1000.0;  /* Convert nm^2/ps to 10e-5 cm^2/s */
65 /* NORMAL = total diffusion coefficient (default). X,Y,Z is diffusion
66    coefficient in X,Y,Z direction. LATERAL is diffusion coefficient in
67    plane perpendicular to axis
68  */
69 typedef enum {
70     NOT_USED, NORMAL, X, Y, Z, LATERAL
71 } msd_type;
72
73 // TODO : Group related fields into a struct
74 struct t_corr {
75     real                                         t0;         /* start time and time increment between  */
76     real                                         delta_t;    /* time between restart points */
77     real                                         beginfit,   /* the begin/end time for fits as reals between */
78                                                  endfit;     /* 0 and 1 */
79     real                                         dim_factor; /* the dimensionality factor for the diffusion
80                                                                 constant */
81     std::vector< std::vector<real> >             data;       /* the displacement data. First index is the group
82                                                                 number, second is frame number */
83     std::vector<real>                            time;       /* frame time */
84     std::vector<real>                            mass;       /* masses for mass-weighted msd */
85     matrix                                     **datam;
86     std::vector< std::vector<gmx::RVec> >        x0;         /* original positions */
87     std::vector<gmx::RVec>                       com;        /* center of mass correction for each frame */
88     gmx_stats_t                                **lsq;        /* fitting stats for individual molecule msds */
89     msd_type                                     type;       /* the type of msd to calculate (lateral, etc.)*/
90     int                                          axis;       /* the axis along which to calculate */
91     int                                          ncoords;
92     int                                          nrestart;   /* number of restart points */
93     int                                          nmol;       /* number of molecules (for bMol) */
94     int                                          nframes;    /* number of frames */
95     int                                          nlast;
96     int                                          ngrp;       /* number of groups to use for msd calculation */
97     std::vector<int>                             n_offs;
98     std::vector< std::vector<int> >              ndata;      /* the number of msds (particles/mols) per data
99                                                                 point. */
100     t_corr(int nrgrp, int type, int axis, real dim_factor, int nrmol,
101            gmx_bool bTen, gmx_bool bMass, real dt, const t_topology *top,
102            real beginfit, real endfit) :
103         t0(0),
104         delta_t(dt),
105         beginfit((1 - 2*GMX_REAL_EPS)*beginfit),
106         endfit((1 + 2*GMX_REAL_EPS)*endfit),
107         dim_factor(dim_factor),
108         data(nrgrp, std::vector<real>()),
109         datam(nullptr),
110         lsq(nullptr),
111         type(static_cast<msd_type>(type)),
112         axis(axis),
113         ncoords(0),
114         nrestart(0),
115         nmol(nrmol),
116         nframes(0),
117         nlast(0),
118         ngrp(nrgrp),
119         ndata(nrgrp, std::vector<int>())
120     {
121
122         if (bTen)
123         {
124             snew(datam, nrgrp);
125             for (int i = 0; i < nrgrp; i++)
126             {
127                 datam[i] = nullptr;
128             }
129         }
130
131         if (nmol > 0)
132         {
133             mass.resize(nmol, 1);
134         }
135         else
136         {
137             if (bMass)
138             {
139                 const t_atoms *atoms = &top->atoms;
140                 mass.resize(atoms->nr);
141                 for (int i = 0; (i < atoms->nr); i++)
142                 {
143                     mass[i] = atoms->atom[i].m;
144                 }
145             }
146         }
147     }
148     ~t_corr()
149     {
150         for (int i = 0; i < nrestart; i++)
151         {
152             for (int j = 0; j < nmol; j++)
153             {
154                 gmx_stats_free(lsq[i][j]);
155             }
156         }
157         sfree(lsq);
158     }
159 };
160
161 typedef real t_calc_func (t_corr *curr, int nx, const int index[], int nx0, rvec xc[],
162                           const rvec dcom, gmx_bool bTen, matrix mat);
163
164 static real thistime(t_corr *curr)
165 {
166     return curr->time[curr->nframes];
167 }
168
169 static int in_data(t_corr *curr, int nx00)
170 {
171     return curr->nframes-curr->n_offs[nx00];
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 gmx_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).c_str(), curr->nrestart);
188         fprintf(out, "# Diffusion constants fitted from time %g to %g %s\n",
189                 beginfit, endfit, output_env_get_time_unit(oenv).c_str());
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, int 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             std::memcpy(curr->x0[curr->nlast].data()->as_vec(), 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 calculation */
267 static real calc1_norm(t_corr *curr, int nx, const int index[], int nx0, rvec xc[],
268                        const 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, const int *molindex, const t_block *mols, const 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                         const 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, const int index[], int nx0, rvec xc[],
406                      const 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, const int 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, int index[],
483                      rvec xcur[], rvec xprev[], matrix box, const 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
494     prep_data(bMol, gnx, index, xcur, xprev, box);
495     for (i = 0; (i < gnx); i++)
496     {
497         if (bMol)
498         {
499             ind = i;
500         }
501         else
502         {
503             ind = index[i];
504         }
505
506
507         mass = atoms->atom[ind].m;
508         for (m = 0; m < DIM; m++)
509         {
510             sx[m] += mass*xcur[ind][m];
511         }
512         tmass += mass;
513     }
514     for (m = 0; m < DIM; m++)
515     {
516         com[m] = sx[m]/tmass;
517     }
518 }
519
520
521 static real calc1_mol(t_corr *curr, int nx, const int gmx_unused index[], int nx0, rvec xc[],
522                       const rvec dcom, gmx_bool bTen, matrix mat)
523 {
524     int  i;
525     real g, tm, gtot, tt;
526
527     tt   = curr->time[in_data(curr, nx0)];
528     gtot = 0;
529     tm   = 0;
530     clear_mat(mat);
531     for (i = 0; (i < nx); i++)
532     {
533         g = calc_one_mw(curr, i, nx0, xc, &tm, dcom, bTen, mat);
534         /* We don't need to normalize as the mass was set to 1 */
535         gtot += g;
536         if (tt >= curr->beginfit && (curr->endfit < 0 || tt <= curr->endfit))
537         {
538             gmx_stats_add_point(curr->lsq[nx0][i], tt, g, 0, 0);
539         }
540     }
541     msmul(mat, 1.0/nx, mat);
542
543     return gtot/nx;
544 }
545
546 static void printmol(t_corr *curr, const char *fn,
547                      const char *fn_pdb, const int *molindex, const t_topology *top,
548                      rvec *x, int ePBC, matrix box, const gmx_output_env_t *oenv)
549 {
550     FILE       *out;
551     gmx_stats_t lsq1;
552     int         i, j;
553     real        a, b, D, Dav, D2av, VarD, sqrtD, sqrtD_max, scale;
554     t_pdbinfo  *pdbinfo = nullptr;
555     const int  *mol2a   = nullptr;
556
557     out = xvgropen(fn, "Diffusion Coefficients / Molecule", "Molecule", "D (1e-5 cm^2/s)", oenv);
558
559     if (fn_pdb)
560     {
561         pdbinfo = top->atoms.pdbinfo;
562         mol2a   = top->mols.index;
563     }
564
565     Dav       = D2av = 0;
566     sqrtD_max = 0;
567     for (i = 0; (i < curr->nmol); i++)
568     {
569         lsq1 = gmx_stats_init();
570         for (j = 0; (j < curr->nrestart); j++)
571         {
572             real xx, yy, dx, dy;
573
574             while (gmx_stats_get_point(curr->lsq[j][i], &xx, &yy, &dx, &dy, 0) == estatsOK)
575             {
576                 gmx_stats_add_point(lsq1, xx, yy, dx, dy);
577             }
578         }
579         gmx_stats_get_ab(lsq1, elsqWEIGHT_NONE, &a, &b, nullptr, nullptr, nullptr, nullptr);
580         gmx_stats_free(lsq1);
581         D     = a*diffusionConversionFactor/curr->dim_factor;
582         if (D < 0)
583         {
584             D   = 0;
585         }
586         Dav  += D;
587         D2av += gmx::square(D);
588         fprintf(out, "%10d  %10g\n", i, D);
589         if (pdbinfo)
590         {
591             sqrtD = std::sqrt(D);
592             if (sqrtD > sqrtD_max)
593             {
594                 sqrtD_max = sqrtD;
595             }
596             for (j = mol2a[molindex[i]]; j < mol2a[molindex[i]+1]; j++)
597             {
598                 pdbinfo[j].bfac = sqrtD;
599             }
600         }
601     }
602     xvgrclose(out);
603     do_view(oenv, fn, "-graphtype bar");
604
605     /* Compute variance, stddev and error */
606     Dav  /= curr->nmol;
607     D2av /= curr->nmol;
608     VarD  = D2av - gmx::square(Dav);
609     printf("<D> = %.4f Std. Dev. = %.4f Error = %.4f\n",
610            Dav, std::sqrt(VarD), std::sqrt(VarD/curr->nmol));
611
612     if (fn_pdb && x)
613     {
614         scale = 1;
615         while (scale*sqrtD_max > 10)
616         {
617             scale *= 0.1;
618         }
619         while (scale*sqrtD_max < 0.1)
620         {
621             scale *= 10;
622         }
623         GMX_RELEASE_ASSERT(pdbinfo != nullptr, "Internal error - pdbinfo not set for PDB input");
624         for (i = 0; i < top->atoms.nr; i++)
625         {
626             pdbinfo[i].bfac *= scale;
627         }
628         write_sto_conf(fn_pdb, "molecular MSD", &top->atoms, x, nullptr, ePBC, box);
629     }
630 }
631
632 /* this is the main loop for the correlation type functions
633  * fx and nx are file pointers to things like read_first_x and
634  * read_next_x
635  */
636 static int corr_loop(t_corr *curr, const char *fn, const t_topology *top, int ePBC,
637                      gmx_bool bMol, int gnx[], int *index[],
638                      t_calc_func *calc1, gmx_bool bTen, gmx::ArrayRef<const int> gnx_com, int *index_com[],
639                      real dt, real t_pdb, rvec **x_pdb, matrix box_pdb,
640                      const gmx_output_env_t *oenv)
641 {
642     rvec            *x[2];  /* the coordinates to read */
643     rvec            *xa[2]; /* the coordinates to calculate displacements for */
644     rvec             com = {0};
645     real             t, t_prev = 0;
646     int              natoms, i, j, cur = 0, maxframes = 0;
647     t_trxstatus     *status;
648 #define        prev (1-cur)
649     matrix           box;
650     gmx_bool         bFirst;
651     gmx_rmpbc_t      gpbc = nullptr;
652
653     natoms = read_first_x(oenv, &status, fn, &curr->t0, &(x[cur]), box);
654 #ifdef DEBUG
655     fprintf(stderr, "Read %d atoms for first frame\n", natoms);
656 #endif
657     if ((!gnx_com.empty()) && natoms < top->atoms.nr)
658     {
659         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);
660     }
661
662     snew(x[prev], natoms);
663
664     // if com is requested, the data structure needs to be large enough to do this
665     // to prevent overflow
666     if (bMol && gnx_com.empty())
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 = nullptr;
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 == nullptr)
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             curr->x0.resize(curr->nrestart);
717             curr->x0[curr->nrestart-1].resize(curr->ncoords);
718             curr->com.resize(curr->nrestart);
719             curr->n_offs.resize(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                     if (bTen)
741                     {
742                         curr->datam[i] = nullptr;
743                     }
744                 }
745             }
746             maxframes += 10;
747             for (i = 0; (i < curr->ngrp); i++)
748             {
749                 curr->ndata[i].resize(maxframes);
750                 curr->data[i].resize(maxframes);
751                 if (bTen)
752                 {
753                     srenew(curr->datam[i], maxframes);
754                 }
755                 for (j = maxframes-10; j < maxframes; j++)
756                 {
757                     curr->ndata[i][j] = 0;
758                     curr->data[i][j]  = 0;
759                     if (bTen)
760                     {
761                         clear_mat(curr->datam[i][j]);
762                     }
763                 }
764             }
765             curr->time.resize(maxframes);
766         }
767
768         /* set the time */
769         curr->time[curr->nframes] = t - curr->t0;
770
771         /* make the molecules whole */
772         if (bMol)
773         {
774             gmx_rmpbc(gpbc, natoms, box, x[cur]);
775         }
776
777         /* calculate the molecules' centers of masses and put them into xa */
778         // NOTE and WARNING! If above both COM removal and individual molecules have been
779         // requested, x and xa point to the same memory, and the coordinate
780         // data becomes overwritten by the molecule data.
781         if (bMol)
782         {
783             calc_mol_com(gnx[0], index[0], &top->mols, &top->atoms, x[cur], xa[cur]);
784         }
785
786         /* for the first frame, the previous frame is a copy of the first frame */
787         if (bFirst)
788         {
789             std::memcpy(xa[prev], xa[cur], curr->ncoords*sizeof(xa[prev][0]));
790             bFirst = FALSE;
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.empty())
801         {
802             calc_com(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box,
803                      &top->atoms, com);
804         }
805
806         /* loop over all groups in index file */
807         for (i = 0; (i < curr->ngrp); i++)
808         {
809             /* calculate something useful, like mean square displacements */
810             calc_corr(curr, i, gnx[i], index[i], xa[cur], (!gnx_com.empty()), com,
811                       calc1, bTen);
812         }
813         cur    = prev;
814         t_prev = t;
815
816         curr->nframes++;
817     }
818     while (read_next_x(oenv, status, &t, x[cur], box));
819     fprintf(stderr, "\nUsed %d restart points spaced %g %s over %g %s\n\n",
820             curr->nrestart,
821             output_env_conv_time(oenv, dt), output_env_get_time_unit(oenv).c_str(),
822             output_env_conv_time(oenv, curr->time[curr->nframes-1]),
823             output_env_get_time_unit(oenv).c_str() );
824
825     if (bMol)
826     {
827         gmx_rmpbc_done(gpbc);
828     }
829
830     close_trx(status);
831
832     return natoms;
833 }
834
835 static void index_atom2mol(int *n, int *index, const t_block *mols)
836 {
837     int nat, i, nmol, mol, j;
838
839     nat  = *n;
840     i    = 0;
841     nmol = 0;
842     mol  = 0;
843     while (i < nat)
844     {
845         while (index[i] > mols->index[mol])
846         {
847             mol++;
848             if (mol >= mols->nr)
849             {
850                 gmx_fatal(FARGS, "Atom index out of range: %d", index[i]+1);
851             }
852         }
853         for (j = mols->index[mol]; j < mols->index[mol+1]; j++)
854         {
855             if (i >= nat || index[i] != j)
856             {
857                 gmx_fatal(FARGS, "The index group does not consist of whole molecules");
858             }
859             i++;
860         }
861         index[nmol++] = mol;
862     }
863
864     fprintf(stderr, "Split group of %d atoms into %d molecules\n", nat, nmol);
865
866     *n = nmol;
867 }
868
869 static void do_corr(const char *trx_file, const char *ndx_file, const char *msd_file,
870                     const char *mol_file, const char *pdb_file, real t_pdb,
871                     int nrgrp, t_topology *top, int ePBC,
872                     gmx_bool bTen, gmx_bool bMW, gmx_bool bRmCOMM,
873                     int type, real dim_factor, int axis,
874                     real dt, real beginfit, real endfit, const gmx_output_env_t *oenv)
875 {
876     std::unique_ptr<t_corr> msd;
877     std::vector<int>        gnx, gnx_com; /* the selected groups' sizes */
878     int                   **index;        /* selected groups' indices */
879     char                  **grpname;
880     int                     i, i0, i1, j, N, nat_trx;
881     std::vector<real>       SigmaD, DD;
882     real                    a, a2, b, r, chi2;
883     rvec                   *x = nullptr;
884     matrix                  box;
885     int                   **index_com   = nullptr; /* the COM removal group atom indices */
886     char                  **grpname_com = nullptr; /* the COM removal group name */
887
888     gnx.resize(nrgrp);
889     snew(index, nrgrp);
890     snew(grpname, nrgrp);
891
892     fprintf(stderr, "\nSelect a group to calculate mean squared displacement for:\n");
893     get_index(&top->atoms, ndx_file, nrgrp, gnx.data(), index, grpname);
894
895     if (bRmCOMM)
896     {
897         gnx_com.resize(1);
898         snew(index_com, 1);
899         snew(grpname_com, 1);
900
901         fprintf(stderr, "\nNow select a group for center of mass removal:\n");
902         get_index(&top->atoms, ndx_file, 1, gnx_com.data(), index_com, grpname_com);
903     }
904
905     if (mol_file)
906     {
907         index_atom2mol(&gnx[0], index[0], &top->mols);
908     }
909
910     msd = gmx::compat::make_unique<t_corr>(nrgrp, type, axis, dim_factor,
911                                            mol_file == nullptr ? 0 : gnx[0],
912                                            bTen, bMW, dt, top, beginfit, endfit);
913
914     nat_trx =
915         corr_loop(msd.get(), trx_file, top, ePBC, mol_file ? gnx[0] != 0 : false, gnx.data(), index,
916                   (mol_file != nullptr) ? calc1_mol : (bMW ? calc1_mw : calc1_norm),
917                   bTen, gnx_com, index_com, dt, t_pdb,
918                   pdb_file ? &x : nullptr, box, oenv);
919
920     /* Correct for the number of points */
921     for (j = 0; (j < msd->ngrp); j++)
922     {
923         for (i = 0; (i < msd->nframes); i++)
924         {
925             msd->data[j][i] /= msd->ndata[j][i];
926             if (bTen)
927             {
928                 msmul(msd->datam[j][i], 1.0/msd->ndata[j][i], msd->datam[j][i]);
929             }
930         }
931     }
932
933     if (mol_file)
934     {
935         if (pdb_file && x == nullptr)
936         {
937             fprintf(stderr, "\nNo frame found need time tpdb = %g ps\n"
938                     "Can not write %s\n\n", t_pdb, pdb_file);
939         }
940         i             = top->atoms.nr;
941         top->atoms.nr = nat_trx;
942         if (pdb_file && top->atoms.pdbinfo == nullptr)
943         {
944             snew(top->atoms.pdbinfo, top->atoms.nr);
945         }
946         printmol(msd.get(), mol_file, pdb_file, index[0], top, x, ePBC, box, oenv);
947         top->atoms.nr = i;
948     }
949
950     if (beginfit == -1)
951     {
952         i0       = gmx::roundToInt(0.1*(msd->nframes - 1));
953         beginfit = msd->time[i0];
954     }
955     else
956     {
957         for (i0 = 0; i0 < msd->nframes && msd->time[i0] < beginfit; i0++)
958         {
959             ;
960         }
961     }
962
963     if (endfit == -1)
964     {
965         i1     = gmx::roundToInt(0.9*(msd->nframes - 1)) + 1;
966         endfit = msd->time[i1-1];
967     }
968     else
969     {
970         for (i1 = i0; i1 < msd->nframes && msd->time[i1] <= endfit; i1++)
971         {
972             ;
973         }
974     }
975     fprintf(stdout, "Fitting from %g to %g %s\n\n", beginfit, endfit,
976             output_env_get_time_unit(oenv).c_str());
977
978     N = i1-i0;
979     if (N <= 2)
980     {
981         fprintf(stdout, "Not enough points for fitting (%d).\n"
982                 "Can not determine the diffusion constant.\n", N);
983     }
984     else
985     {
986         DD.resize(msd->ngrp);
987         SigmaD.resize(msd->ngrp);
988         for (j = 0; j < msd->ngrp; j++)
989         {
990             if (N >= 4)
991             {
992                 lsq_y_ax_b(N/2, &(msd->time[i0]), &(msd->data[j][i0]), &a, &b, &r, &chi2);
993                 lsq_y_ax_b(N/2, &(msd->time[i0+N/2]), &(msd->data[j][i0+N/2]), &a2, &b, &r, &chi2);
994                 SigmaD[j] = std::abs(a-a2);
995             }
996             else
997             {
998                 SigmaD[j] = 0;
999             }
1000             lsq_y_ax_b(N, &(msd->time[i0]), &(msd->data[j][i0]), &(DD[j]), &b, &r, &chi2);
1001             DD[j]     *= diffusionConversionFactor/msd->dim_factor;
1002             SigmaD[j] *= diffusionConversionFactor/msd->dim_factor;
1003             if (DD[j] > 0.01 && DD[j] < 1e4)
1004             {
1005                 fprintf(stdout, "D[%10s] %.4f (+/- %.4f) 1e-5 cm^2/s\n",
1006                         grpname[j], DD[j], SigmaD[j]);
1007             }
1008             else
1009             {
1010                 fprintf(stdout, "D[%10s] %.4g (+/- %.4g) 1e-5 cm^2/s\n",
1011                         grpname[j], DD[j], SigmaD[j]);
1012             }
1013         }
1014     }
1015     /* Print mean square displacement */
1016     corr_print(msd.get(), bTen, msd_file,
1017                "Mean Square Displacement",
1018                "MSD (nm\\S2\\N)",
1019                msd->time[msd->nframes-1], beginfit, endfit, DD.data(), SigmaD.data(), grpname, oenv);
1020 }
1021
1022 int gmx_msd(int argc, char *argv[])
1023 {
1024     const char        *desc[] = {
1025         "[THISMODULE] computes the mean square displacement (MSD) of atoms from",
1026         "a set of initial positions. This provides an easy way to compute",
1027         "the diffusion constant using the Einstein relation.",
1028         "The time between the reference points for the MSD calculation",
1029         "is set with [TT]-trestart[tt].",
1030         "The diffusion constant is calculated by least squares fitting a",
1031         "straight line (D*t + c) through the MSD(t) from [TT]-beginfit[tt] to",
1032         "[TT]-endfit[tt] (note that t is time from the reference positions,",
1033         "not simulation time). An error estimate given, which is the difference",
1034         "of the diffusion coefficients obtained from fits over the two halves",
1035         "of the fit interval.[PAR]",
1036         "There are three, mutually exclusive, options to determine different",
1037         "types of mean square displacement: [TT]-type[tt], [TT]-lateral[tt]",
1038         "and [TT]-ten[tt]. Option [TT]-ten[tt] writes the full MSD tensor for",
1039         "each group, the order in the output is: trace xx yy zz yx zx zy.[PAR]",
1040         "If [TT]-mol[tt] is set, [THISMODULE] plots the MSD for individual molecules",
1041         "(including making molecules whole across periodic boundaries): ",
1042         "for each individual molecule a diffusion constant is computed for ",
1043         "its center of mass. The chosen index group will be split into ",
1044         "molecules.[PAR]",
1045         "The default way to calculate a MSD is by using mass-weighted averages.",
1046         "This can be turned off with [TT]-nomw[tt].[PAR]",
1047         "With the option [TT]-rmcomm[tt], the center of mass motion of a ",
1048         "specific group can be removed. For trajectories produced with ",
1049         "GROMACS this is usually not necessary, ",
1050         "as [gmx-mdrun] usually already removes the center of mass motion.",
1051         "When you use this option be sure that the whole system is stored",
1052         "in the trajectory file.[PAR]",
1053         "The diffusion coefficient is determined by linear regression of the MSD,",
1054         "where, unlike for the normal output of D, the times are weighted",
1055         "according to the number of reference points, i.e. short times have",
1056         "a higher weight. Also when [TT]-beginfit[tt] is -1, fitting starts at 10%",
1057         "and when [TT]-endfit[tt] is -1, fitting goes to 90%.",
1058         "Using this option one also gets an accurate error estimate",
1059         "based on the statistics between individual molecules.",
1060         "Note that this diffusion coefficient and error estimate are only",
1061         "accurate when the MSD is completely linear between",
1062         "[TT]-beginfit[tt] and [TT]-endfit[tt].[PAR]",
1063         "Option [TT]-pdb[tt] writes a [REF].pdb[ref] file with the coordinates of the frame",
1064         "at time [TT]-tpdb[tt] with in the B-factor field the square root of",
1065         "the diffusion coefficient of the molecule.",
1066         "This option implies option [TT]-mol[tt]."
1067     };
1068     static const char *normtype[] = { nullptr, "no", "x", "y", "z", nullptr };
1069     static const char *axtitle[]  = { nullptr, "no", "x", "y", "z", nullptr };
1070     static int         ngroup     = 1;
1071     static real        dt         = 10;
1072     static real        t_pdb      = 0;
1073     static real        beginfit   = -1;
1074     static real        endfit     = -1;
1075     static gmx_bool    bTen       = FALSE;
1076     static gmx_bool    bMW        = TRUE;
1077     static gmx_bool    bRmCOMM    = FALSE;
1078     t_pargs            pa[]       = {
1079         { "-type",    FALSE, etENUM, {normtype},
1080           "Compute diffusion coefficient in one direction" },
1081         { "-lateral", FALSE, etENUM, {axtitle},
1082           "Calculate the lateral diffusion in a plane perpendicular to" },
1083         { "-ten",      FALSE, etBOOL, {&bTen},
1084           "Calculate the full tensor" },
1085         { "-ngroup",  FALSE, etINT,  {&ngroup},
1086           "Number of groups to calculate MSD for" },
1087         { "-mw",      FALSE, etBOOL, {&bMW},
1088           "Mass weighted MSD" },
1089         { "-rmcomm",      FALSE, etBOOL, {&bRmCOMM},
1090           "Remove center of mass motion" },
1091         { "-tpdb", FALSE, etTIME, {&t_pdb},
1092           "The frame to use for option [TT]-pdb[tt] (%t)" },
1093         { "-trestart", FALSE, etTIME, {&dt},
1094           "Time between restarting points in trajectory (%t)" },
1095         { "-beginfit", FALSE, etTIME, {&beginfit},
1096           "Start time for fitting the MSD (%t), -1 is 10%" },
1097         { "-endfit", FALSE, etTIME, {&endfit},
1098           "End time for fitting the MSD (%t), -1 is 90%" }
1099     };
1100
1101     t_filenm           fnm[] = {
1102         { efTRX, nullptr, nullptr,  ffREAD },
1103         { efTPS, nullptr, nullptr,  ffREAD },
1104         { efNDX, nullptr, nullptr,  ffOPTRD },
1105         { efXVG, nullptr, "msd", ffWRITE },
1106         { efXVG, "-mol", "diff_mol", ffOPTWR },
1107         { efPDB, "-pdb", "diff_mol", ffOPTWR }
1108     };
1109 #define NFILE asize(fnm)
1110
1111     t_topology        top;
1112     int               ePBC;
1113     matrix            box;
1114     const char       *trx_file, *tps_file, *ndx_file, *msd_file, *mol_file, *pdb_file;
1115     rvec             *xdum;
1116     gmx_bool          bTop;
1117     int               axis, type;
1118     real              dim_factor;
1119     gmx_output_env_t *oenv;
1120
1121     if (!parse_common_args(&argc, argv,
1122                            PCA_CAN_VIEW | PCA_CAN_BEGIN | PCA_CAN_END | PCA_TIME_UNIT,
1123                            NFILE, fnm, asize(pa), pa, asize(desc), desc, 0, nullptr, &oenv))
1124     {
1125         return 0;
1126     }
1127     trx_file = ftp2fn_null(efTRX, NFILE, fnm);
1128     tps_file = ftp2fn_null(efTPS, NFILE, fnm);
1129     ndx_file = ftp2fn_null(efNDX, NFILE, fnm);
1130     msd_file = ftp2fn_null(efXVG, NFILE, fnm);
1131     pdb_file = opt2fn_null("-pdb", NFILE, fnm);
1132     if (pdb_file)
1133     {
1134         mol_file = opt2fn("-mol", NFILE, fnm);
1135     }
1136     else
1137     {
1138         mol_file = opt2fn_null("-mol", NFILE, fnm);
1139     }
1140
1141     if (ngroup < 1)
1142     {
1143         gmx_fatal(FARGS, "Must have at least 1 group (now %d)", ngroup);
1144     }
1145     if (mol_file && ngroup > 1)
1146     {
1147         gmx_fatal(FARGS, "With molecular msd can only have 1 group (now %d)",
1148                   ngroup);
1149     }
1150
1151
1152     if (mol_file)
1153     {
1154         bMW  = TRUE;
1155         fprintf(stderr, "Calculating diffusion coefficients for molecules.\n");
1156     }
1157
1158     GMX_RELEASE_ASSERT(normtype[0] != nullptr, "Options inconsistency; normtype[0] is NULL");
1159     GMX_RELEASE_ASSERT(axtitle[0] != nullptr, "Options inconsistency; axtitle[0] is NULL");
1160
1161     if (normtype[0][0] != 'n')
1162     {
1163         type       = normtype[0][0] - 'x' + X; /* See defines above */
1164         dim_factor = 2.0;
1165     }
1166     else
1167     {
1168         type       = NORMAL;
1169         dim_factor = 6.0;
1170     }
1171     if ((type == NORMAL) && (axtitle[0][0] != 'n'))
1172     {
1173         type       = LATERAL;
1174         dim_factor = 4.0;
1175         axis       = (axtitle[0][0] - 'x'); /* See defines above */
1176     }
1177     else
1178     {
1179         axis = 0;
1180     }
1181
1182     if (bTen && type != NORMAL)
1183     {
1184         gmx_fatal(FARGS, "Can only calculate the full tensor for 3D msd");
1185     }
1186
1187     bTop = read_tps_conf(tps_file, &top, &ePBC, &xdum, nullptr, box, bMW || bRmCOMM);
1188     if (mol_file && !bTop)
1189     {
1190         gmx_fatal(FARGS,
1191                   "Could not read a topology from %s. Try a tpr file instead.",
1192                   tps_file);
1193     }
1194
1195     do_corr(trx_file, ndx_file, msd_file, mol_file, pdb_file, t_pdb, ngroup,
1196             &top, ePBC, bTen, bMW, bRmCOMM, type, dim_factor, axis, dt, beginfit, endfit,
1197             oenv);
1198
1199     done_top(&top);
1200     view_all(oenv, NFILE, fnm);
1201
1202     return 0;
1203 }