Apply re-formatting to C++ in src/ tree.
[alexxy/gromacs.git] / src / gromacs / gmxana / gmx_trjorder.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,2017,2018 by the GROMACS development team.
7  * Copyright (c) 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 <cstdlib>
42 #include <cstring>
43
44 #include "gromacs/commandline/pargs.h"
45 #include "gromacs/fileio/confio.h"
46 #include "gromacs/fileio/trxio.h"
47 #include "gromacs/fileio/xvgr.h"
48 #include "gromacs/gmxana/gmx_ana.h"
49 #include "gromacs/math/vec.h"
50 #include "gromacs/pbcutil/pbc.h"
51 #include "gromacs/pbcutil/rmpbc.h"
52 #include "gromacs/topology/index.h"
53 #include "gromacs/topology/topology.h"
54 #include "gromacs/trajectory/trajectoryframe.h"
55 #include "gromacs/utility/arraysize.h"
56 #include "gromacs/utility/fatalerror.h"
57 #include "gromacs/utility/futil.h"
58 #include "gromacs/utility/smalloc.h"
59
60 typedef struct
61 {
62     int  i;
63     real d2;
64 } t_order;
65
66 static t_order* order;
67
68 static int ocomp(const void* a, const void* b)
69 {
70     const t_order *oa, *ob;
71
72     oa = reinterpret_cast<const t_order*>(a);
73     ob = reinterpret_cast<const t_order*>(b);
74
75     if (oa->d2 < ob->d2)
76     {
77         return -1;
78     }
79     else
80     {
81         return 1;
82     }
83 }
84
85 int gmx_trjorder(int argc, char* argv[])
86 {
87     const char* desc[] = {
88         "[THISMODULE] orders molecules according to the smallest distance",
89         "to atoms in a reference group",
90         "or on z-coordinate (with option [TT]-z[tt]).",
91         "With distance ordering, it will ask for a group of reference",
92         "atoms and a group of molecules. For each frame of the trajectory",
93         "the selected molecules will be reordered according to the shortest",
94         "distance between atom number [TT]-da[tt] in the molecule and all the",
95         "atoms in the reference group. The center of mass of the molecules can",
96         "be used instead of a reference atom by setting [TT]-da[tt] to 0.",
97         "All atoms in the trajectory are written",
98         "to the output trajectory.[PAR]",
99         "[THISMODULE] can be useful for e.g. analyzing the n waters closest to a",
100         "protein.",
101         "In that case the reference group would be the protein and the group",
102         "of molecules would consist of all the water atoms. When an index group",
103         "of the first n waters is made, the ordered trajectory can be used",
104         "with any GROMACS program to analyze the n closest waters.",
105         "[PAR]",
106         "If the output file is a [REF].pdb[ref] file, the distance to the reference target",
107         "will be stored in the B-factor field in order to color with e.g. Rasmol.",
108         "[PAR]",
109         "With option [TT]-nshell[tt] the number of molecules within a shell",
110         "of radius [TT]-r[tt] around the reference group are printed."
111     };
112     static int      na = 3, ref_a = 1;
113     static real     rcut = 0;
114     static gmx_bool bCOM = FALSE, bZ = FALSE;
115     t_pargs         pa[] = {
116         { "-na", FALSE, etINT, { &na }, "Number of atoms in a molecule" },
117         { "-da", FALSE, etINT, { &ref_a }, "Atom used for the distance calculation, 0 is COM" },
118         { "-com",
119           FALSE,
120           etBOOL,
121           { &bCOM },
122           "Use the distance to the center of mass of the reference group" },
123         { "-r",
124           FALSE,
125           etREAL,
126           { &rcut },
127           "Cutoff used for the distance calculation when computing the number of molecules in a "
128           "shell around e.g. a protein" },
129         { "-z", FALSE, etBOOL, { &bZ }, "Order molecules on z-coordinate" }
130     };
131     FILE*             fp;
132     t_trxstatus*      out;
133     t_trxstatus*      status;
134     gmx_bool          bNShell, bPDBout;
135     t_topology        top;
136     PbcType           pbcType;
137     rvec *            x, *xsol, xcom, dx;
138     matrix            box;
139     t_pbc             pbc;
140     gmx_rmpbc_t       gpbc;
141     real              t, totmass, mass, rcut2 = 0, n2;
142     int               natoms, nwat, ncut;
143     char**            grpname;
144     int               i, j, d, *isize, isize_ref      = 0, isize_sol;
145     int               sa, sr, *swi, **index, *ind_ref = nullptr, *ind_sol;
146     gmx_output_env_t* oenv;
147     t_filenm          fnm[] = { { efTRX, "-f", nullptr, ffREAD },
148                        { efTPS, nullptr, nullptr, ffREAD },
149                        { efNDX, nullptr, nullptr, ffOPTRD },
150                        { efTRO, "-o", "ordered", ffOPTWR },
151                        { efXVG, "-nshell", "nshell", ffOPTWR } };
152 #define NFILE asize(fnm)
153
154     if (!parse_common_args(
155                 &argc, argv, PCA_CAN_TIME, NFILE, fnm, asize(pa), pa, asize(desc), desc, 0, nullptr, &oenv))
156     {
157         return 0;
158     }
159
160     read_tps_conf(ftp2fn(efTPS, NFILE, fnm), &top, &pbcType, &x, nullptr, box, TRUE);
161     sfree(x);
162
163     /* get index groups */
164     printf("Select %sa group of molecules to be ordered:\n", bZ ? "" : "a group of reference atoms and ");
165     snew(grpname, 2);
166     snew(index, 2);
167     snew(isize, 2);
168     get_index(&top.atoms, ftp2fn_null(efNDX, NFILE, fnm), bZ ? 1 : 2, isize, index, grpname);
169
170     if (!bZ)
171     {
172         isize_ref = isize[0];
173         isize_sol = isize[1];
174         ind_ref   = index[0];
175         ind_sol   = index[1];
176     }
177     else
178     {
179         isize_sol = isize[0];
180         ind_sol   = index[0];
181     }
182
183     natoms = read_first_x(oenv, &status, ftp2fn(efTRX, NFILE, fnm), &t, &x, box);
184     if (natoms > top.atoms.nr)
185     {
186         gmx_fatal(FARGS, "Number of atoms in the run input file is larger than in the trjactory");
187     }
188     for (i = 0; (i < 2); i++)
189     {
190         for (j = 0; (j < isize[i]); j++)
191         {
192             if (index[i][j] > natoms)
193             {
194                 gmx_fatal(FARGS,
195                           "An atom number in group %s is larger than the number of atoms in the "
196                           "trajectory",
197                           grpname[i]);
198             }
199         }
200     }
201
202     if ((isize_sol % na) != 0)
203     {
204         gmx_fatal(FARGS,
205                   "Number of atoms in the molecule group (%d) is not a multiple of na (%d)",
206                   isize[1],
207                   na);
208     }
209
210     nwat = isize_sol / na;
211     if (ref_a > na)
212     {
213         gmx_fatal(FARGS,
214                   "The reference atom can not be larger than the number of atoms in a molecule");
215     }
216     ref_a--;
217     snew(xsol, nwat);
218     snew(order, nwat);
219     snew(swi, natoms);
220     for (i = 0; (i < natoms); i++)
221     {
222         swi[i] = i;
223     }
224
225     out     = nullptr;
226     fp      = nullptr;
227     bNShell = ((opt2bSet("-nshell", NFILE, fnm)) || (opt2parg_bSet("-r", asize(pa), pa)));
228     bPDBout = FALSE;
229     if (bNShell)
230     {
231         rcut2 = rcut * rcut;
232         fp = xvgropen(opt2fn("-nshell", NFILE, fnm), "Number of molecules", "Time (ps)", "N", oenv);
233         printf("Will compute the number of molecules within a radius of %g\n", rcut);
234     }
235     if (!bNShell || opt2bSet("-o", NFILE, fnm))
236     {
237         bPDBout = (fn2ftp(opt2fn("-o", NFILE, fnm)) == efPDB);
238         if (bPDBout && !top.atoms.pdbinfo)
239         {
240             fprintf(stderr, "Creating pdbfino records\n");
241             snew(top.atoms.pdbinfo, top.atoms.nr);
242         }
243         out = open_trx(opt2fn("-o", NFILE, fnm), "w");
244     }
245     gpbc = gmx_rmpbc_init(&top.idef, pbcType, natoms);
246     do
247     {
248         gmx_rmpbc(gpbc, natoms, box, x);
249         set_pbc(&pbc, pbcType, box);
250
251         if (ref_a == -1)
252         {
253             /* Calculate the COM of all solvent molecules */
254             for (i = 0; i < nwat; i++)
255             {
256                 totmass = 0;
257                 clear_rvec(xsol[i]);
258                 for (j = 0; j < na; j++)
259                 {
260                     sa   = ind_sol[i * na + j];
261                     mass = top.atoms.atom[sa].m;
262                     totmass += mass;
263                     for (d = 0; d < DIM; d++)
264                     {
265                         xsol[i][d] += mass * x[sa][d];
266                     }
267                 }
268                 svmul(1.0 / totmass, xsol[i], xsol[i]);
269             }
270         }
271         else
272         {
273             /* Copy the reference atom of all solvent molecules */
274             for (i = 0; i < nwat; i++)
275             {
276                 copy_rvec(x[ind_sol[i * na + ref_a]], xsol[i]);
277             }
278         }
279
280         if (bZ)
281         {
282             for (i = 0; (i < nwat); i++)
283             {
284                 sa          = ind_sol[na * i];
285                 order[i].i  = sa;
286                 order[i].d2 = xsol[i][ZZ];
287             }
288         }
289         else if (bCOM)
290         {
291             totmass = 0;
292             clear_rvec(xcom);
293             for (i = 0; i < isize_ref; i++)
294             {
295                 mass = top.atoms.atom[ind_ref[i]].m;
296                 totmass += mass;
297                 for (j = 0; j < DIM; j++)
298                 {
299                     xcom[j] += mass * x[ind_ref[i]][j];
300                 }
301             }
302             svmul(1 / totmass, xcom, xcom);
303             for (i = 0; (i < nwat); i++)
304             {
305                 sa = ind_sol[na * i];
306                 pbc_dx(&pbc, xcom, xsol[i], dx);
307                 order[i].i  = sa;
308                 order[i].d2 = norm2(dx);
309             }
310         }
311         else
312         {
313             /* Set distance to first atom */
314             for (i = 0; (i < nwat); i++)
315             {
316                 sa = ind_sol[na * i];
317                 pbc_dx(&pbc, x[ind_ref[0]], xsol[i], dx);
318                 order[i].i  = sa;
319                 order[i].d2 = norm2(dx);
320             }
321             for (j = 1; (j < isize_ref); j++)
322             {
323                 sr = ind_ref[j];
324                 for (i = 0; (i < nwat); i++)
325                 {
326                     pbc_dx(&pbc, x[sr], xsol[i], dx);
327                     n2 = norm2(dx);
328                     if (n2 < order[i].d2)
329                     {
330                         order[i].d2 = n2;
331                     }
332                 }
333             }
334         }
335
336         if (bNShell)
337         {
338             ncut = 0;
339             for (i = 0; (i < nwat); i++)
340             {
341                 if (order[i].d2 <= rcut2)
342                 {
343                     ncut++;
344                 }
345             }
346             fprintf(fp, "%10.3f  %8d\n", t, ncut);
347         }
348         if (out)
349         {
350             qsort(order, nwat, sizeof(*order), ocomp);
351             for (i = 0; (i < nwat); i++)
352             {
353                 for (j = 0; (j < na); j++)
354                 {
355                     swi[ind_sol[na * i] + j] = order[i].i + j;
356                 }
357             }
358
359             /* Store the distance as the B-factor */
360             if (bPDBout)
361             {
362                 for (i = 0; (i < nwat); i++)
363                 {
364                     for (j = 0; (j < na); j++)
365                     {
366                         top.atoms.pdbinfo[order[i].i + j].bfac = std::sqrt(order[i].d2);
367                     }
368                 }
369             }
370             write_trx(out, natoms, swi, &top.atoms, 0, t, box, x, nullptr, nullptr);
371         }
372     } while (read_next_x(oenv, status, &t, x, box));
373     close_trx(status);
374     if (out)
375     {
376         close_trx(out);
377     }
378     if (fp)
379     {
380         xvgrclose(fp);
381     }
382     gmx_rmpbc_done(gpbc);
383
384     return 0;
385 }