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,2021, 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     fprintf(stdout, "Wrote per-molecule output to %s\n", fn);
 637     do_view(oenv, fn, "-graphtype bar");
 638
 639     /* Compute variance, stddev and error */
 640     Dav /= curr->nmol;
 641     D2av /= curr->nmol;
 642     VarD = D2av - gmx::square(Dav);
 643     printf("<D> = %.4f Std. Dev. = %.4f Error = %.4f\n", Dav, std::sqrt(VarD), std::sqrt(VarD / curr->nmol));
 644
 645     if (fn_pdb && x)
 646     {
 647         scale = 1;
 648         while (scale * sqrtD_max > 10)
 649         {
 650             scale *= 0.1;
 651         }
 652         while (scale * sqrtD_max < 0.1)
 653         {
 654             scale *= 10;
 655         }
 656         GMX_RELEASE_ASSERT(pdbinfo != nullptr, "Internal error - pdbinfo not set for PDB input");
 657         for (i = 0; i < top->atoms.nr; i++)
 658         {
 659             pdbinfo[i].bfac *= scale;
 660         }
 661         write_sto_conf(fn_pdb, "molecular MSD", &top->atoms, x, nullptr, pbcType, box);
 662         fprintf(stdout, "Wrote frame for -tpdb to %s\n", fn_pdb);
 663     }
 664 }
 665
 666 /* this is the main loop for the correlation type functions
 667  * fx and nx are file pointers to things like read_first_x and
 668  * read_next_x
 669  */
 670 static int corr_loop(t_corr*                  curr,
 671                      const char*              fn,
 672                      const t_topology*        top,
 673                      PbcType                  pbcType,
 674                      gmx_bool                 bMol,
 675                      int                      gnx[],
 676                      int*                     index[],
 677                      t_calc_func*             calc1,
 678                      gmx_bool                 bTen,
 679                      gmx::ArrayRef<const int> gnx_com,
 680                      int*                     index_com[],
 681                      real                     dt,
 682                      real                     t_pdb,
 683                      rvec**                   x_pdb,
 684                      matrix                   box_pdb,
 685                      const gmx_output_env_t*  oenv)
 686 {
 687     rvec*        x[2];  /* the coordinates to read */
 688     rvec*        xa[2]; /* the coordinates to calculate displacements for */
 689     rvec         com = { 0 };
 690     real         t, t_prev = 0;
 691     int          natoms, i, j, cur = 0, maxframes = 0;
 692     t_trxstatus* status;
 693 #define prev (1 - cur)
 694     matrix      box;
 695     gmx_bool    bFirst;
 696     gmx_rmpbc_t gpbc = nullptr;
 697
 698     natoms = read_first_x(oenv, &status, fn, &curr->t0, &(x[cur]), box);
 699 #ifdef DEBUG
 700     fprintf(stderr, "Read %d atoms for first frame\n", natoms);
 701 #endif
 702     if ((!gnx_com.empty()) && natoms < top->atoms.nr)
 703     {
 704         fprintf(stderr,
 705                 "WARNING: The trajectory only contains part of the system (%d of %d atoms) and "
 706                 "therefore the COM motion of only this part of the system will be removed\n",
 707                 natoms,
 708                 top->atoms.nr);
 709     }
 710
 711     snew(x[prev], natoms);
 712
 713     // if com is requested, the data structure needs to be large enough to do this
 714     // to prevent overflow
 715     if (bMol && gnx_com.empty())
 716     {
 717         curr->ncoords = curr->nmol;
 718         snew(xa[0], curr->ncoords);
 719         snew(xa[1], curr->ncoords);
 720     }
 721     else
 722     {
 723         curr->ncoords = natoms;
 724         xa[0]         = x[0];
 725         xa[1]         = x[1];
 726     }
 727
 728     bFirst = TRUE;
 729     t      = curr->t0;
 730     if (x_pdb)
 731     {
 732         *x_pdb = nullptr;
 733     }
 734
 735     if (bMol)
 736     {
 737         gpbc = gmx_rmpbc_init(&top->idef, pbcType, natoms);
 738     }
 739
 740     /* the loop over all frames */
 741     do
 742     {
 743         if (x_pdb
 744             && ((bFirst && t_pdb < t)
 745                 || (!bFirst && t_pdb > t - 0.5 * (t - t_prev) && t_pdb < t + 0.5 * (t - t_prev))))
 746         {
 747             if (*x_pdb == nullptr)
 748             {
 749                 snew(*x_pdb, natoms);
 750             }
 751             for (i = 0; i < natoms; i++)
 752             {
 753                 copy_rvec(x[cur][i], (*x_pdb)[i]);
 754             }
 755             copy_mat(box, box_pdb);
 756         }
 757
 758
 759         /* check whether we've reached a restart point */
 760         if (bRmod(t, curr->t0, dt))
 761         {
 762             curr->nrestart++;
 763
 764             curr->x0.resize(curr->nrestart);
 765             curr->x0[curr->nrestart - 1].resize(curr->ncoords);
 766             curr->com.resize(curr->nrestart);
 767             curr->n_offs.resize(curr->nrestart);
 768             srenew(curr->lsq, curr->nrestart);
 769             snew(curr->lsq[curr->nrestart - 1], curr->nmol);
 770             for (i = 0; i < curr->nmol; i++)
 771             {
 772                 curr->lsq[curr->nrestart - 1][i] = gmx_stats_init();
 773             }
 774
 775             if (debug)
 776             {
 777                 fprintf(debug, "Extended data structures because of new restart %d\n", curr->nrestart);
 778             }
 779         }
 780         /* create or extend the frame-based arrays */
 781         if (curr->nframes >= maxframes - 1)
 782         {
 783             if (maxframes == 0)
 784             {
 785                 for (i = 0; (i < curr->ngrp); i++)
 786                 {
 787                     if (bTen)
 788                     {
 789                         curr->datam[i] = nullptr;
 790                     }
 791                 }
 792             }
 793             maxframes += 10;
 794             for (i = 0; (i < curr->ngrp); i++)
 795             {
 796                 curr->ndata[i].resize(maxframes);
 797                 curr->data[i].resize(maxframes);
 798                 if (bTen)
 799                 {
 800                     srenew(curr->datam[i], maxframes);
 801                 }
 802                 for (j = maxframes - 10; j < maxframes; j++)
 803                 {
 804                     curr->ndata[i][j] = 0;
 805                     curr->data[i][j]  = 0;
 806                     if (bTen)
 807                     {
 808                         clear_mat(curr->datam[i][j]);
 809                     }
 810                 }
 811             }
 812             curr->time.resize(maxframes);
 813         }
 814
 815         /* set the time */
 816         curr->time[curr->nframes] = t - curr->t0;
 817
 818         /* make the molecules whole */
 819         if (bMol)
 820         {
 821             gmx_rmpbc(gpbc, natoms, box, x[cur]);
 822         }
 823
 824         /* calculate the molecules' centers of masses and put them into xa */
 825         // NOTE and WARNING! If above both COM removal and individual molecules have been
 826         // requested, x and xa point to the same memory, and the coordinate
 827         // data becomes overwritten by the molecule data.
 828         if (bMol)
 829         {
 830             calc_mol_com(gnx[0], index[0], &top->mols, &top->atoms, x[cur], xa[cur]);
 831         }
 832
 833         /* for the first frame, the previous frame is a copy of the first frame */
 834         if (bFirst)
 835         {
 836             std::memcpy(xa[prev], xa[cur], curr->ncoords * sizeof(xa[prev][0]));
 837             bFirst = FALSE;
 838         }
 839
 840         /* first remove the periodic boundary condition crossings */
 841         for (i = 0; i < curr->ngrp; i++)
 842         {
 843             prep_data(bMol, gnx[i], index[i], xa[cur], xa[prev], box);
 844         }
 845
 846         /* calculate the center of mass */
 847         if (!gnx_com.empty())
 848         {
 849             GMX_RELEASE_ASSERT(index_com != nullptr,
 850                                "Center-of-mass removal must have valid index group");
 851             calc_com(bMol, gnx_com[0], index_com[0], xa[cur], xa[prev], box, &top->atoms, com);
 852         }
 853
 854         /* loop over all groups in index file */
 855         for (i = 0; (i < curr->ngrp); i++)
 856         {
 857             /* calculate something useful, like mean square displacements */
 858             calc_corr(curr, i, gnx[i], index[i], xa[cur], (!gnx_com.empty()), com, calc1, bTen);
 859         }
 860         cur    = prev;
 861         t_prev = t;
 862
 863         curr->nframes++;
 864     } while (read_next_x(oenv, status, &t, x[cur], box));
 865     fprintf(stderr,
 866             "\nUsed %d restart points spaced %g %s over %g %s\n\n",
 867             curr->nrestart,
 868             output_env_conv_time(oenv, dt),
 869             output_env_get_time_unit(oenv).c_str(),
 870             output_env_conv_time(oenv, curr->time[curr->nframes - 1]),
 871             output_env_get_time_unit(oenv).c_str());
 872
 873     if (bMol)
 874     {
 875         gmx_rmpbc_done(gpbc);
 876     }
 877
 878     close_trx(status);
 879
 880     return natoms;
 881 }
 882
 883 static void index_atom2mol(int* n, int* index, const t_block* mols)
 884 {
 885     int nat, i, nmol, mol, j;
 886
 887     nat  = *n;
 888     i    = 0;
 889     nmol = 0;
 890     mol  = 0;
 891     while (i < nat)
 892     {
 893         while (index[i] > mols->index[mol])
 894         {
 895             mol++;
 896             if (mol >= mols->nr)
 897             {
 898                 gmx_fatal(FARGS, "Atom index out of range: %d", index[i] + 1);
 899             }
 900         }
 901         for (j = mols->index[mol]; j < mols->index[mol + 1]; j++)
 902         {
 903             if (i >= nat || index[i] != j)
 904             {
 905                 gmx_fatal(FARGS, "The index group does not consist of whole molecules");
 906             }
 907             i++;
 908         }
 909         index[nmol++] = mol;
 910     }
 911
 912     fprintf(stderr, "Split group of %d atoms into %d molecules\n", nat, nmol);
 913
 914     *n = nmol;
 915 }
 916
 917 static void do_corr(const char*             trx_file,
 918                     const char*             ndx_file,
 919                     const char*             msd_file,
 920                     const char*             mol_file,
 921                     const char*             pdb_file,
 922                     real                    t_pdb,
 923                     int                     nrgrp,
 924                     t_topology*             top,
 925                     PbcType                 pbcType,
 926                     gmx_bool                bTen,
 927                     gmx_bool                bMW,
 928                     gmx_bool                bRmCOMM,
 929                     int                     type,
 930                     real                    dim_factor,
 931                     int                     axis,
 932                     real                    dt,
 933                     real                    beginfit,
 934                     real                    endfit,
 935                     const gmx_output_env_t* oenv)
 936 {
 937     std::unique_ptr<t_corr> msd;
 938     std::vector<int>        gnx, gnx_com; /* the selected groups' sizes */
 939     int**                   index;        /* selected groups' indices */
 940     char**                  grpname;
 941     int                     i, i0, i1, j, N, nat_trx;
 942     std::vector<real>       SigmaD, DD;
 943     real                    a, a2, b, r, chi2;
 944     rvec*                   x = nullptr;
 945     matrix                  box;
 946     int**                   index_com   = nullptr; /* the COM removal group atom indices */
 947     char**                  grpname_com = nullptr; /* the COM removal group name */
 948
 949     gnx.resize(nrgrp);
 950     snew(index, nrgrp);
 951     snew(grpname, nrgrp);
 952
 953     fprintf(stderr, "\nSelect a group to calculate mean squared displacement for:\n");
 954     get_index(&top->atoms, ndx_file, nrgrp, gnx.data(), index, grpname);
 955
 956     if (bRmCOMM)
 957     {
 958         gnx_com.resize(1);
 959         snew(index_com, 1);
 960         snew(grpname_com, 1);
 961
 962         fprintf(stderr, "\nNow select a group for center of mass removal:\n");
 963         get_index(&top->atoms, ndx_file, 1, gnx_com.data(), index_com, grpname_com);
 964     }
 965
 966     if (mol_file)
 967     {
 968         index_atom2mol(&gnx[0], index[0], &top->mols);
 969     }
 970
 971     msd = std::make_unique<t_corr>(
 972             nrgrp, type, axis, dim_factor, mol_file == nullptr ? 0 : gnx[0], bTen, bMW, dt, top, beginfit, endfit);
 973
 974     nat_trx = corr_loop(msd.get(),
 975                         trx_file,
 976                         top,
 977                         pbcType,
 978                         mol_file ? gnx[0] != 0 : false,
 979                         gnx.data(),
 980                         index,
 981                         (mol_file != nullptr) ? calc1_mol : (bMW ? calc1_mw : calc1_norm),
 982                         bTen,
 983                         gnx_com,
 984                         index_com,
 985                         dt,
 986                         t_pdb,
 987                         pdb_file ? &x : nullptr,
 988                         box,
 989                         oenv);
 990
 991     /* Correct for the number of points */
 992     for (j = 0; (j < msd->ngrp); j++)
 993     {
 994         for (i = 0; (i < msd->nframes); i++)
 995         {
 996             msd->data[j][i] /= msd->ndata[j][i];
 997             if (bTen)
 998             {
 999                 msmul(msd->datam[j][i], 1.0 / msd->ndata[j][i], msd->datam[j][i]);
1000             }
1001         }
1002     }
1003
1004     if (mol_file)
1005     {
1006         if (pdb_file && x == nullptr)
1007         {
1008             fprintf(stderr,
1009                     "\nNo frame found need time tpdb = %g ps\n"
1010                     "Can not write %s\n\n",
1011                     t_pdb,
1012                     pdb_file);
1013         }
1014         i             = top->atoms.nr;
1015         top->atoms.nr = nat_trx;
1016         if (pdb_file && top->atoms.pdbinfo == nullptr)
1017         {
1018             snew(top->atoms.pdbinfo, top->atoms.nr);
1019         }
1020         printmol(msd.get(), mol_file, pdb_file, index[0], top, x, pbcType, box, oenv);
1021         top->atoms.nr = i;
1022     }
1023
1024     if (beginfit == -1)
1025     {
1026         i0       = gmx::roundToInt(0.1 * (msd->nframes - 1));
1027         beginfit = msd->time[i0];
1028     }
1029     else
1030     {
1031         for (i0 = 0; i0 < msd->nframes && msd->time[i0] < beginfit; i0++) {}
1032     }
1033
1034     if (endfit == -1)
1035     {
1036         i1     = gmx::roundToInt(0.9 * (msd->nframes - 1)) + 1;
1037         endfit = msd->time[i1 - 1];
1038     }
1039     else
1040     {
1041         for (i1 = i0; i1 < msd->nframes && msd->time[i1] <= endfit; i1++) {}
1042     }
1043     fprintf(stdout, "Fitting from %g to %g %s\n\n", beginfit, endfit, output_env_get_time_unit(oenv).c_str());
1044
1045     N = i1 - i0;
1046     if (N <= 2)
1047     {
1048         fprintf(stdout,
1049                 "Not enough points for fitting (%d).\n"
1050                 "Can not determine the diffusion constant.\n",
1051                 N);
1052     }
1053     else
1054     {
1055         DD.resize(msd->ngrp);
1056         SigmaD.resize(msd->ngrp);
1057         for (j = 0; j < msd->ngrp; j++)
1058         {
1059             if (N >= 4)
1060             {
1061                 lsq_y_ax_b(N / 2, &(msd->time[i0]), &(msd->data[j][i0]), &a, &b, &r, &chi2);
1062                 lsq_y_ax_b(N / 2, &(msd->time[i0 + N / 2]), &(msd->data[j][i0 + N / 2]), &a2, &b, &r, &chi2);
1063                 SigmaD[j] = std::abs(a - a2);
1064             }
1065             else
1066             {
1067                 SigmaD[j] = 0;
1068             }
1069             lsq_y_ax_b(N, &(msd->time[i0]), &(msd->data[j][i0]), &(DD[j]), &b, &r, &chi2);
1070             DD[j] *= diffusionConversionFactor / msd->dim_factor;
1071             SigmaD[j] *= diffusionConversionFactor / msd->dim_factor;
1072             if (DD[j] > 0.01 && DD[j] < 1e4)
1073             {
1074                 fprintf(stdout, "D[%10s] %.4f (+/- %.4f) 1e-5 cm^2/s\n", grpname[j], DD[j], SigmaD[j]);
1075             }
1076             else
1077             {
1078                 fprintf(stdout, "D[%10s] %.4g (+/- %.4g) 1e-5 cm^2/s\n", grpname[j], DD[j], SigmaD[j]);
1079             }
1080         }
1081     }
1082     /* Print mean square displacement */
1083     corr_print(msd.get(),
1084                bTen,
1085                msd_file,
1086                "Mean Square Displacement",
1087                "MSD (nm\\S2\\N)",
1088                msd->time[msd->nframes - 1],
1089                beginfit,
1090                endfit,
1091                DD.data(),
1092                SigmaD.data(),
1093                grpname,
1094                oenv);
1095 }
1096
1097 int gmx_msd(int argc, char* argv[])
1098 {
1099     const char* desc[] = {
1100         "[THISMODULE] computes the mean square displacement (MSD) of atoms from",
1101         "a set of initial positions. This provides an easy way to compute",
1102         "the diffusion constant using the Einstein relation.",
1103         "The time between the reference points for the MSD calculation",
1104         "is set with [TT]-trestart[tt].",
1105         "The diffusion constant is calculated by least squares fitting a",
1106         "straight line (D*t + c) through the MSD(t) from [TT]-beginfit[tt] to",
1107         "[TT]-endfit[tt] (note that t is time from the reference positions,",
1108         "not simulation time). An error estimate given, which is the difference",
1109         "of the diffusion coefficients obtained from fits over the two halves",
1110         "of the fit interval.[PAR]",
1111         "There are three, mutually exclusive, options to determine different",
1112         "types of mean square displacement: [TT]-type[tt], [TT]-lateral[tt]",
1113         "and [TT]-ten[tt]. Option [TT]-ten[tt] writes the full MSD tensor for",
1114         "each group, the order in the output is: trace xx yy zz yx zx zy.[PAR]",
1115         "If [TT]-mol[tt] is set, [THISMODULE] plots the MSD for individual molecules",
1116         "(including making molecules whole across periodic boundaries): ",
1117         "for each individual molecule a diffusion constant is computed for ",
1118         "its center of mass. The chosen index group will be split into ",
1119         "molecules.[PAR]",
1120         "The default way to calculate a MSD is by using mass-weighted averages.",
1121         "This can be turned off with [TT]-nomw[tt].[PAR]",
1122         "With the option [TT]-rmcomm[tt], the center of mass motion of a ",
1123         "specific group can be removed. For trajectories produced with ",
1124         "GROMACS this is usually not necessary, ",
1125         "as [gmx-mdrun] usually already removes the center of mass motion.",
1126         "When you use this option be sure that the whole system is stored",
1127         "in the trajectory file.[PAR]",
1128         "The diffusion coefficient is determined by linear regression of the MSD.",
1129         "When [TT]-beginfit[tt] is -1, fitting starts at 10%",
1130         "and when [TT]-endfit[tt] is -1, fitting goes to 90%.",
1131         "Using this option one also gets an accurate error estimate",
1132         "based on the statistics between individual molecules.",
1133         "Note that this diffusion coefficient and error estimate are only",
1134         "accurate when the MSD is completely linear between",
1135         "[TT]-beginfit[tt] and [TT]-endfit[tt].[PAR]",
1136         "Option [TT]-pdb[tt] writes a [REF].pdb[ref] file with the coordinates of the frame",
1137         "at time [TT]-tpdb[tt] with in the B-factor field the square root of",
1138         "the diffusion coefficient of the molecule.",
1139         "This option implies option [TT]-mol[tt]."
1140     };
1141     static const char* normtype[] = { nullptr, "no", "x", "y", "z", nullptr };
1142     static const char* axtitle[]  = { nullptr, "no", "x", "y", "z", nullptr };
1143     static int         ngroup     = 1;
1144     static real        dt         = 10;
1145     static real        t_pdb      = 0;
1146     static real        beginfit   = -1;
1147     static real        endfit     = -1;
1148     static gmx_bool    bTen       = FALSE;
1149     static gmx_bool    bMW        = TRUE;
1150     static gmx_bool    bRmCOMM    = FALSE;
1151     t_pargs            pa[]       = {
1152         { "-type", FALSE, etENUM, { normtype }, "Compute diffusion coefficient in one direction" },
1153         { "-lateral",
1154           FALSE,
1155           etENUM,
1156           { axtitle },
1157           "Calculate the lateral diffusion in a plane perpendicular to" },
1158         { "-ten", FALSE, etBOOL, { &bTen }, "Calculate the full tensor" },
1159         { "-ngroup", FALSE, etINT, { &ngroup }, "Number of groups to calculate MSD for" },
1160         { "-mw", FALSE, etBOOL, { &bMW }, "Mass weighted MSD" },
1161         { "-rmcomm", FALSE, etBOOL, { &bRmCOMM }, "Remove center of mass motion" },
1162         { "-tpdb", FALSE, etTIME, { &t_pdb }, "The frame to use for option [TT]-pdb[tt] (%t)" },
1163         { "-trestart", FALSE, etTIME, { &dt }, "Time between restarting points in trajectory (%t)" },
1164         { "-beginfit",
1165           FALSE,
1166           etTIME,
1167           { &beginfit },
1168           "Start time for fitting the MSD (%t), -1 is 10%" },
1169         { "-endfit", FALSE, etTIME, { &endfit }, "End time for fitting the MSD (%t), -1 is 90%" }
1170     };
1171
1172     t_filenm fnm[] = {
1173         { efTRX, nullptr, nullptr, ffREAD },    { efTPS, nullptr, nullptr, ffREAD },
1174         { efNDX, nullptr, nullptr, ffOPTRD },   { efXVG, nullptr, "msd", ffWRITE },
1175         { efXVG, "-mol", "diff_mol", ffOPTWR }, { efPDB, "-pdb", "diff_mol", ffOPTWR }
1176     };
1177 #define NFILE asize(fnm)
1178
1179     t_topology        top;
1180     PbcType           pbcType;
1181     matrix            box;
1182     const char *      trx_file, *tps_file, *ndx_file, *msd_file, *mol_file, *pdb_file;
1183     rvec*             xdum;
1184     gmx_bool          bTop;
1185     int               axis, type;
1186     real              dim_factor;
1187     gmx_output_env_t* oenv;
1188
1189     if (!parse_common_args(&argc,
1190                            argv,
1191                            PCA_CAN_VIEW | PCA_CAN_BEGIN | PCA_CAN_END | PCA_TIME_UNIT,
1192                            NFILE,
1193                            fnm,
1194                            asize(pa),
1195                            pa,
1196                            asize(desc),
1197                            desc,
1198                            0,
1199                            nullptr,
1200                            &oenv))
1201     {
1202         return 0;
1203     }
1204     trx_file = ftp2fn_null(efTRX, NFILE, fnm);
1205     tps_file = ftp2fn_null(efTPS, NFILE, fnm);
1206     ndx_file = ftp2fn_null(efNDX, NFILE, fnm);
1207     msd_file = ftp2fn_null(efXVG, NFILE, fnm);
1208     pdb_file = opt2fn_null("-pdb", NFILE, fnm);
1209     if (pdb_file)
1210     {
1211         mol_file = opt2fn("-mol", NFILE, fnm);
1212     }
1213     else
1214     {
1215         mol_file = opt2fn_null("-mol", NFILE, fnm);
1216     }
1217
1218     if (ngroup < 1)
1219     {
1220         gmx_fatal(FARGS, "Must have at least 1 group (now %d)", ngroup);
1221     }
1222     if (mol_file && ngroup > 1)
1223     {
1224         gmx_fatal(FARGS, "With molecular msd can only have 1 group (now %d)", ngroup);
1225     }
1226
1227
1228     if (mol_file)
1229     {
1230         bMW = TRUE;
1231         fprintf(stderr, "Calculating diffusion coefficients for molecules.\n");
1232     }
1233
1234     GMX_RELEASE_ASSERT(normtype[0] != nullptr, "Options inconsistency; normtype[0] is NULL");
1235     GMX_RELEASE_ASSERT(axtitle[0] != nullptr, "Options inconsistency; axtitle[0] is NULL");
1236
1237     if (normtype[0][0] != 'n')
1238     {
1239         type       = normtype[0][0] - 'x' + X; /* See defines above */
1240         dim_factor = 2.0;
1241     }
1242     else
1243     {
1244         type       = NORMAL;
1245         dim_factor = 6.0;
1246     }
1247     if ((type == NORMAL) && (axtitle[0][0] != 'n'))
1248     {
1249         type       = LATERAL;
1250         dim_factor = 4.0;
1251         axis       = (axtitle[0][0] - 'x'); /* See defines above */
1252     }
1253     else
1254     {
1255         axis = 0;
1256     }
1257
1258     if (bTen && type != NORMAL)
1259     {
1260         gmx_fatal(FARGS, "Can only calculate the full tensor for 3D msd");
1261     }
1262
1263     bTop = read_tps_conf(tps_file, &top, &pbcType, &xdum, nullptr, box, bMW || bRmCOMM);
1264     if (mol_file && !bTop)
1265     {
1266         gmx_fatal(FARGS, "Could not read a topology from %s. Try a tpr file instead.", tps_file);
1267     }
1268
1269     do_corr(trx_file,
1270             ndx_file,
1271             msd_file,
1272             mol_file,
1273             pdb_file,
1274             t_pdb,
1275             ngroup,
1276             &top,
1277             pbcType,
1278             bTen,
1279             bMW,
1280             bRmCOMM,
1281             type,
1282             dim_factor,
1283             axis,
1284             dt,
1285             beginfit,
1286             endfit,
1287             oenv);
1288
1289     done_top(&top);
1290     view_all(oenv, NFILE, fnm);
1291
1292     return 0;
1293 }