6491e164e25d019eac722ed06eb9b92af58c1ef7
[alexxy/gromacs.git] / src / gromacs / gmxana / gmx_enemat.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
5  * Copyright (c) 2001-2004, The GROMACS development team.
6  * Copyright (c) 2013,2014,2015,2016,2017,2018,2019, by the GROMACS development team, led by
7  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
8  * and including many others, as listed in the AUTHORS file in the
9  * top-level source directory and at http://www.gromacs.org.
10  *
11  * GROMACS is free software; you can redistribute it and/or
12  * modify it under the terms of the GNU Lesser General Public License
13  * as published by the Free Software Foundation; either version 2.1
14  * of the License, or (at your option) any later version.
15  *
16  * GROMACS is distributed in the hope that it will be useful,
17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
19  * Lesser General Public License for more details.
20  *
21  * You should have received a copy of the GNU Lesser General Public
22  * License along with GROMACS; if not, see
23  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
24  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
25  *
26  * If you want to redistribute modifications to GROMACS, please
27  * consider that scientific software is very special. Version
28  * control is crucial - bugs must be traceable. We will be happy to
29  * consider code for inclusion in the official distribution, but
30  * derived work must not be called official GROMACS. Details are found
31  * in the README & COPYING files - if they are missing, get the
32  * official version at http://www.gromacs.org.
33  *
34  * To help us fund GROMACS development, we humbly ask that you cite
35  * the research papers on the package. Check out http://www.gromacs.org.
36  */
37 #include "gmxpre.h"
38
39 #include <cmath>
40 #include <cstring>
41
42 #include "gromacs/commandline/pargs.h"
43 #include "gromacs/fileio/enxio.h"
44 #include "gromacs/fileio/matio.h"
45 #include "gromacs/fileio/trxio.h"
46 #include "gromacs/fileio/xvgr.h"
47 #include "gromacs/gmxana/gmx_ana.h"
48 #include "gromacs/gmxana/gstat.h"
49 #include "gromacs/math/functions.h"
50 #include "gromacs/math/units.h"
51 #include "gromacs/math/vec.h"
52 #include "gromacs/mdlib/energyoutput.h"
53 #include "gromacs/mdtypes/enerdata.h"
54 #include "gromacs/mdtypes/md_enums.h"
55 #include "gromacs/trajectory/energyframe.h"
56 #include "gromacs/utility/arraysize.h"
57 #include "gromacs/utility/cstringutil.h"
58 #include "gromacs/utility/fatalerror.h"
59 #include "gromacs/utility/futil.h"
60 #include "gromacs/utility/smalloc.h"
61 #include "gromacs/utility/strdb.h"
62
63
64 static int search_str2(int nstr, char** str, char* key)
65 {
66     int i, n;
67     int keylen = std::strlen(key);
68     /* Linear search */
69     n = 0;
70     while ((n < keylen) && ((key[n] < '0') || (key[n] > '9')))
71     {
72         n++;
73     }
74     for (i = 0; (i < nstr); i++)
75     {
76         if (gmx_strncasecmp(str[i], key, n) == 0)
77         {
78             return i;
79         }
80     }
81
82     return -1;
83 }
84
85 int gmx_enemat(int argc, char* argv[])
86 {
87     const char* desc[] = {
88         "[THISMODULE] extracts an energy matrix from the energy file ([TT]-f[tt]).",
89         "With [TT]-groups[tt] a file must be supplied with on each",
90         "line a group of atoms to be used. For these groups matrix of",
91         "interaction energies will be extracted from the energy file",
92         "by looking for energy groups with names corresponding to pairs",
93         "of groups of atoms, e.g. if your [TT]-groups[tt] file contains::",
94         "",
95         "    2",
96         "    Protein",
97         "    SOL",
98         "",
99         "then energy groups with names like 'Coul-SR:Protein-SOL' and ",
100         "'LJ:Protein-SOL' are expected in the energy file (although",
101         "[THISMODULE] is most useful if many groups are analyzed",
102         "simultaneously). Matrices for different energy types are written",
103         "out separately, as controlled by the",
104         "[TT]-[no]coul[tt], [TT]-[no]coulr[tt], [TT]-[no]coul14[tt], ",
105         "[TT]-[no]lj[tt], [TT]-[no]lj14[tt], ",
106         "[TT]-[no]bham[tt] and [TT]-[no]free[tt] options.",
107         "Finally, the total interaction energy energy per group can be ",
108         "calculated ([TT]-etot[tt]).[PAR]",
109
110         "An approximation of the free energy can be calculated using:",
111         "[MATH]E[SUB]free[sub] = E[SUB]0[sub] + kT ",
112         "[LOG][CHEVRON][EXP](E-E[SUB]0[sub])/kT[exp][chevron][log][math], where ",
113         "'[MATH][CHEVRON][chevron][math]'",
114         "stands for time-average. A file with reference free energies",
115         "can be supplied to calculate the free energy difference",
116         "with some reference state. Group names (e.g. residue names)",
117         "in the reference file should correspond to the group names",
118         "as used in the [TT]-groups[tt] file, but a appended number",
119         "(e.g. residue number) in the [TT]-groups[tt] will be ignored",
120         "in the comparison."
121     };
122     static gmx_bool bSum      = FALSE;
123     static gmx_bool bMeanEmtx = TRUE;
124     static int      skip = 0, nlevels = 20;
125     static real     cutmax = 1e20, cutmin = -1e20, reftemp = 300.0;
126     static gmx_bool bCoulSR = TRUE, bCoul14 = FALSE;
127     static gmx_bool bLJSR = TRUE, bLJ14 = FALSE, bBhamSR = FALSE, bFree = TRUE;
128     t_pargs         pa[] = {
129         { "-sum",
130           FALSE,
131           etBOOL,
132           { &bSum },
133           "Sum the energy terms selected rather than display them all" },
134         { "-skip", FALSE, etINT, { &skip }, "Skip number of frames between data points" },
135         { "-mean",
136           FALSE,
137           etBOOL,
138           { &bMeanEmtx },
139           "with [TT]-groups[tt] extracts matrix of mean energies instead of "
140           "matrix for each timestep" },
141         { "-nlevels", FALSE, etINT, { &nlevels }, "number of levels for matrix colors" },
142         { "-max", FALSE, etREAL, { &cutmax }, "max value for energies" },
143         { "-min", FALSE, etREAL, { &cutmin }, "min value for energies" },
144         { "-coulsr", FALSE, etBOOL, { &bCoulSR }, "extract Coulomb SR energies" },
145         { "-coul14", FALSE, etBOOL, { &bCoul14 }, "extract Coulomb 1-4 energies" },
146         { "-ljsr", FALSE, etBOOL, { &bLJSR }, "extract Lennard-Jones SR energies" },
147         { "-lj14", FALSE, etBOOL, { &bLJ14 }, "extract Lennard-Jones 1-4 energies" },
148         { "-bhamsr", FALSE, etBOOL, { &bBhamSR }, "extract Buckingham SR energies" },
149         { "-free", FALSE, etBOOL, { &bFree }, "calculate free energy" },
150         { "-temp",
151           FALSE,
152           etREAL,
153           { &reftemp },
154           "reference temperature for free energy calculation" }
155     };
156     /* We will define egSP more energy-groups:
157        egTotal (total energy) */
158 #define egTotal egNR
159 #define egSP 1
160     gmx_bool          egrp_use[egNR + egSP];
161     ener_file_t       in;
162     FILE*             out;
163     int               timecheck = 0;
164     gmx_enxnm_t*      enm       = nullptr;
165     t_enxframe*       fr;
166     int               teller = 0;
167     real              sum;
168     gmx_bool          bCont, bRef;
169     gmx_bool          bCutmax, bCutmin;
170     real **           eneset, *time = nullptr;
171     int *             set, i, j, prevk, k, m = 0, n, nre, nset, nenergy;
172     char**            groups = nullptr;
173     char              groupname[255], fn[255];
174     int               ngroups;
175     t_rgb             rlo, rhi, rmid;
176     real              emax, emid, emin;
177     real ***          emat, **etot, *groupnr;
178     double            beta, expE, **e, *eaver, *efree = nullptr, edum;
179     char              label[234];
180     char **           ereflines, **erefres = nullptr;
181     real *            eref = nullptr, *edif = nullptr;
182     int               neref = 0;
183     gmx_output_env_t* oenv;
184
185     t_filenm fnm[] = { { efEDR, "-f", nullptr, ffOPTRD },
186                        { efDAT, "-groups", "groups", ffREAD },
187                        { efDAT, "-eref", "eref", ffOPTRD },
188                        { efXPM, "-emat", "emat", ffWRITE },
189                        { efXVG, "-etot", "energy", ffWRITE } };
190 #define NFILE asize(fnm)
191
192     if (!parse_common_args(&argc, argv, PCA_CAN_VIEW | PCA_CAN_TIME, NFILE, fnm, asize(pa), pa,
193                            asize(desc), desc, 0, nullptr, &oenv))
194     {
195         return 0;
196     }
197
198     for (i = 0; (i < egNR + egSP); i++)
199     {
200         egrp_use[i] = FALSE;
201     }
202     egrp_use[egCOULSR] = bCoulSR;
203     egrp_use[egLJSR]   = bLJSR;
204     egrp_use[egBHAMSR] = bBhamSR;
205     egrp_use[egCOUL14] = bCoul14;
206     egrp_use[egLJ14]   = bLJ14;
207     egrp_use[egTotal]  = TRUE;
208
209     bRef = opt2bSet("-eref", NFILE, fnm);
210     in   = open_enx(ftp2fn(efEDR, NFILE, fnm), "r");
211     do_enxnms(in, &nre, &enm);
212
213     if (nre == 0)
214     {
215         gmx_fatal(FARGS, "No energies!\n");
216     }
217
218     bCutmax = opt2parg_bSet("-max", asize(pa), pa);
219     bCutmin = opt2parg_bSet("-min", asize(pa), pa);
220
221     nenergy = 0;
222
223     /* Read groupnames from input file and construct selection of
224        energy groups from it*/
225
226     fprintf(stderr, "Will read groupnames from inputfile\n");
227     ngroups = get_lines(opt2fn("-groups", NFILE, fnm), &groups);
228     fprintf(stderr, "Read %d groups\n", ngroups);
229     snew(set, static_cast<size_t>(gmx::square(ngroups) * egNR / 2));
230     n     = 0;
231     prevk = 0;
232     for (i = 0; (i < ngroups); i++)
233     {
234         for (j = i; (j < ngroups); j++)
235         {
236             for (m = 0; (m < egNR); m++)
237             {
238                 if (egrp_use[m])
239                 {
240                     sprintf(groupname, "%s:%s-%s", egrp_nm[m], groups[i], groups[j]);
241                     bool foundMatch = false;
242                     for (k = prevk; (k < prevk + nre); k++)
243                     {
244                         if (std::strcmp(enm[k % nre].name, groupname) == 0)
245                         {
246                             set[n++]   = k;
247                             foundMatch = true;
248                             break;
249                         }
250                     }
251                     if (!foundMatch)
252                     {
253                         fprintf(stderr,
254                                 "WARNING! could not find group %s (%d,%d) "
255                                 "in energy file\n",
256                                 groupname, i, j);
257                     }
258                     else
259                     {
260                         prevk = k;
261                     }
262                 }
263             }
264         }
265     }
266     fprintf(stderr, "\n");
267     if (n == 0)
268     {
269         // Return an error, can't do what the user asked for
270         fprintf(stderr,
271                 "None of the specified energy groups were found in this .edr file.\n"
272                 "Perhaps you used the wrong groups, the wrong files, or didn't use a .tpr\n"
273                 "that was made from an .mdp file that specified these energy groups.\n");
274         return 1;
275     }
276     nset = n;
277     snew(eneset, nset + 1);
278     fprintf(stderr, "Will select half-matrix of energies with %d elements\n", n);
279
280     /* Start reading energy frames */
281     snew(fr, 1);
282     do
283     {
284         do
285         {
286             bCont = do_enx(in, fr);
287             if (bCont)
288             {
289                 timecheck = check_times(fr->t);
290             }
291         } while (bCont && (timecheck < 0));
292
293         if (timecheck == 0)
294         {
295 #define DONTSKIP(cnt) (skip) ? (((cnt) % skip) == 0) : TRUE
296
297             if (bCont)
298             {
299                 fprintf(stderr, "\rRead frame: %d, Time: %.3f", teller, fr->t);
300                 fflush(stderr);
301
302                 if ((nenergy % 1000) == 0)
303                 {
304                     srenew(time, nenergy + 1000);
305                     for (i = 0; (i <= nset); i++)
306                     {
307                         srenew(eneset[i], nenergy + 1000);
308                     }
309                 }
310                 time[nenergy] = fr->t;
311                 sum           = 0;
312                 for (i = 0; (i < nset); i++)
313                 {
314                     eneset[i][nenergy] = fr->ener[set[i]].e;
315                     sum += fr->ener[set[i]].e;
316                 }
317                 if (bSum)
318                 {
319                     eneset[nset][nenergy] = sum;
320                 }
321                 nenergy++;
322             }
323             teller++;
324         }
325     } while (bCont && (timecheck == 0));
326
327     fprintf(stderr, "\n");
328
329     fprintf(stderr,
330             "Will build energy half-matrix of %d groups, %d elements, "
331             "over %d frames\n",
332             ngroups, nset, nenergy);
333
334     snew(emat, egNR + egSP);
335     for (j = 0; (j < egNR + egSP); j++)
336     {
337         if (egrp_use[m])
338         {
339             snew(emat[j], ngroups);
340             for (i = 0; (i < ngroups); i++)
341             {
342                 snew(emat[j][i], ngroups);
343             }
344         }
345     }
346     snew(groupnr, ngroups);
347     for (i = 0; (i < ngroups); i++)
348     {
349         groupnr[i] = i + 1;
350     }
351     rlo.r  = 1.0;
352     rlo.g  = 0.0;
353     rlo.b  = 0.0;
354     rmid.r = 1.0;
355     rmid.g = 1.0;
356     rmid.b = 1.0;
357     rhi.r  = 0.0;
358     rhi.g  = 0.0;
359     rhi.b  = 1.0;
360     if (bMeanEmtx)
361     {
362         snew(e, ngroups);
363         for (i = 0; (i < ngroups); i++)
364         {
365             snew(e[i], nenergy);
366         }
367         n = 0;
368         for (i = 0; (i < ngroups); i++)
369         {
370             for (j = i; (j < ngroups); j++)
371             {
372                 for (m = 0; (m < egNR); m++)
373                 {
374                     if (egrp_use[m])
375                     {
376                         for (k = 0; (k < nenergy); k++)
377                         {
378                             emat[m][i][j] += eneset[n][k];
379                             e[i][k] += eneset[n][k]; /* *0.5; */
380                             e[j][k] += eneset[n][k]; /* *0.5; */
381                         }
382                         n++;
383                         emat[egTotal][i][j] += emat[m][i][j];
384                         emat[m][i][j] /= nenergy;
385                         emat[m][j][i] = emat[m][i][j];
386                     }
387                 }
388                 emat[egTotal][i][j] /= nenergy;
389                 emat[egTotal][j][i] = emat[egTotal][i][j];
390             }
391         }
392         if (bFree)
393         {
394             if (bRef)
395             {
396                 fprintf(stderr, "Will read reference energies from inputfile\n");
397                 neref = get_lines(opt2fn("-eref", NFILE, fnm), &ereflines);
398                 fprintf(stderr, "Read %d reference energies\n", neref);
399                 snew(eref, neref);
400                 snew(erefres, neref);
401                 for (i = 0; (i < neref); i++)
402                 {
403                     snew(erefres[i], 5);
404                     sscanf(ereflines[i], "%s %lf", erefres[i], &edum);
405                     eref[i] = edum;
406                 }
407             }
408             snew(eaver, ngroups);
409             for (i = 0; (i < ngroups); i++)
410             {
411                 for (k = 0; (k < nenergy); k++)
412                 {
413                     eaver[i] += e[i][k];
414                 }
415                 eaver[i] /= nenergy;
416             }
417             beta = 1.0 / (BOLTZ * reftemp);
418             snew(efree, ngroups);
419             snew(edif, ngroups);
420             for (i = 0; (i < ngroups); i++)
421             {
422                 expE = 0;
423                 for (k = 0; (k < nenergy); k++)
424                 {
425                     expE += std::exp(beta * (e[i][k] - eaver[i]));
426                 }
427                 efree[i] = std::log(expE / nenergy) / beta + eaver[i];
428                 if (bRef)
429                 {
430                     n = search_str2(neref, erefres, groups[i]);
431                     if (n != -1)
432                     {
433                         edif[i] = efree[i] - eref[n];
434                     }
435                     else
436                     {
437                         edif[i] = efree[i];
438                         fprintf(stderr,
439                                 "WARNING: group %s not found "
440                                 "in reference energies.\n",
441                                 groups[i]);
442                     }
443                 }
444                 else
445                 {
446                     edif[i] = 0;
447                 }
448             }
449         }
450
451         emid             = 0.0; /*(emin+emax)*0.5;*/
452         egrp_nm[egTotal] = "total";
453         for (m = 0; (m < egNR + egSP); m++)
454         {
455             if (egrp_use[m])
456             {
457                 emin = 1e10;
458                 emax = -1e10;
459                 for (i = 0; (i < ngroups); i++)
460                 {
461                     for (j = i; (j < ngroups); j++)
462                     {
463                         if (emat[m][i][j] > emax)
464                         {
465                             emax = emat[m][i][j];
466                         }
467                         else if (emat[m][i][j] < emin)
468                         {
469                             emin = emat[m][i][j];
470                         }
471                     }
472                 }
473                 if (emax == emin)
474                 {
475                     fprintf(stderr,
476                             "Matrix of %s energy is uniform at %f "
477                             "(will not produce output).\n",
478                             egrp_nm[m], emax);
479                 }
480                 else
481                 {
482                     fprintf(stderr, "Matrix of %s energy ranges from %f to %f\n", egrp_nm[m], emin, emax);
483                     if ((bCutmax) || (emax > cutmax))
484                     {
485                         emax = cutmax;
486                     }
487                     if ((bCutmin) || (emin < cutmin))
488                     {
489                         emin = cutmin;
490                     }
491                     if ((emax == cutmax) || (emin == cutmin))
492                     {
493                         fprintf(stderr, "Energy range adjusted: %f to %f\n", emin, emax);
494                     }
495
496                     sprintf(fn, "%s%s", egrp_nm[m], ftp2fn(efXPM, NFILE, fnm));
497                     sprintf(label, "%s Interaction Energies", egrp_nm[m]);
498                     out = gmx_ffopen(fn, "w");
499                     if (emin >= emid)
500                     {
501                         write_xpm(out, 0, label, "Energy (kJ/mol)", "Residue Index",
502                                   "Residue Index", ngroups, ngroups, groupnr, groupnr, emat[m],
503                                   emid, emax, rmid, rhi, &nlevels);
504                     }
505                     else if (emax <= emid)
506                     {
507                         write_xpm(out, 0, label, "Energy (kJ/mol)", "Residue Index",
508                                   "Residue Index", ngroups, ngroups, groupnr, groupnr, emat[m],
509                                   emin, emid, rlo, rmid, &nlevels);
510                     }
511                     else
512                     {
513                         write_xpm3(out, 0, label, "Energy (kJ/mol)", "Residue Index",
514                                    "Residue Index", ngroups, ngroups, groupnr, groupnr, emat[m],
515                                    emin, emid, emax, rlo, rmid, rhi, &nlevels);
516                     }
517                     gmx_ffclose(out);
518                 }
519             }
520         }
521         snew(etot, egNR + egSP);
522         for (m = 0; (m < egNR + egSP); m++)
523         {
524             snew(etot[m], ngroups);
525             for (i = 0; (i < ngroups); i++)
526             {
527                 for (j = 0; (j < ngroups); j++)
528                 {
529                     etot[m][i] += emat[m][i][j];
530                 }
531             }
532         }
533
534         out = xvgropen(ftp2fn(efXVG, NFILE, fnm), "Mean Energy", "Residue", "kJ/mol", oenv);
535         xvgr_legend(out, 0, nullptr, oenv);
536         j = 0;
537         if (output_env_get_print_xvgr_codes(oenv))
538         {
539             char str1[STRLEN], str2[STRLEN];
540             if (output_env_get_xvg_format(oenv) == exvgXMGR)
541             {
542                 sprintf(str1, "@ legend string ");
543                 sprintf(str2, " ");
544             }
545             else
546             {
547                 sprintf(str1, "@ s");
548                 sprintf(str2, " legend ");
549             }
550
551             for (m = 0; (m < egNR + egSP); m++)
552             {
553                 if (egrp_use[m])
554                 {
555                     fprintf(out, "%s%d%s \"%s\"\n", str1, j++, str2, egrp_nm[m]);
556                 }
557             }
558             if (bFree)
559             {
560                 fprintf(out, "%s%d%s \"%s\"\n", str1, j++, str2, "Free");
561             }
562             if (bFree)
563             {
564                 fprintf(out, "%s%d%s \"%s\"\n", str1, j++, str2, "Diff");
565             }
566             fprintf(out, "@TYPE xy\n");
567             fprintf(out, "#%3s", "grp");
568
569             for (m = 0; (m < egNR + egSP); m++)
570             {
571                 if (egrp_use[m])
572                 {
573                     fprintf(out, " %9s", egrp_nm[m]);
574                 }
575             }
576             if (bFree)
577             {
578                 fprintf(out, " %9s", "Free");
579             }
580             if (bFree)
581             {
582                 fprintf(out, " %9s", "Diff");
583             }
584             fprintf(out, "\n");
585         }
586         for (i = 0; (i < ngroups); i++)
587         {
588             fprintf(out, "%3.0f", groupnr[i]);
589             for (m = 0; (m < egNR + egSP); m++)
590             {
591                 if (egrp_use[m])
592                 {
593                     fprintf(out, " %9.5g", etot[m][i]);
594                 }
595             }
596             if (bFree)
597             {
598                 fprintf(out, " %9.5g", efree[i]);
599             }
600             if (bRef)
601             {
602                 fprintf(out, " %9.5g", edif[i]);
603             }
604             fprintf(out, "\n");
605         }
606         xvgrclose(out);
607     }
608     else
609     {
610         fprintf(stderr,
611                 "While typing at your keyboard, suddenly...\n"
612                 "...nothing happens.\nWARNING: Not Implemented Yet\n");
613         /*
614             out=ftp2FILE(efMAT,NFILE,fnm,"w");
615             n=0;
616             emin=emax=0.0;
617             for (k=0; (k<nenergy); k++) {
618               for (i=0; (i<ngroups); i++)
619             for (j=i+1; (j<ngroups); j++)
620               emat[i][j]=eneset[n][k];
621               sprintf(label,"t=%.0f ps",time[k]);
622               write_matrix(out,ngroups,1,ngroups,groupnr,emat,label,emin,emax,nlevels);
623               n++;
624             }
625             gmx_ffclose(out);
626          */
627     }
628     close_enx(in);
629
630     return 0;
631 }