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