src/gromacs/gmxana/gmx_msd.c

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