src/gromacs/gmxana/gmx_msd.cpp

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