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