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