src/gromacs/gmxana/gmx_msd.cpp

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