96966c9a00ab73368617ba7a440d798eda65cfbf
[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(&argc, argv, PCA_CAN_TIME, NFILE, fnm, asize(pa), pa, asize(desc), desc,
155                            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, "Number of atoms in the molecule group (%d) is not a multiple of na (%d)",
205                   isize[1], na);
206     }
207
208     nwat = isize_sol / na;
209     if (ref_a > na)
210     {
211         gmx_fatal(FARGS,
212                   "The reference atom can not be larger than the number of atoms in a molecule");
213     }
214     ref_a--;
215     snew(xsol, nwat);
216     snew(order, nwat);
217     snew(swi, natoms);
218     for (i = 0; (i < natoms); i++)
219     {
220         swi[i] = i;
221     }
222
223     out     = nullptr;
224     fp      = nullptr;
225     bNShell = ((opt2bSet("-nshell", NFILE, fnm)) || (opt2parg_bSet("-r", asize(pa), pa)));
226     bPDBout = FALSE;
227     if (bNShell)
228     {
229         rcut2 = rcut * rcut;
230         fp = xvgropen(opt2fn("-nshell", NFILE, fnm), "Number of molecules", "Time (ps)", "N", oenv);
231         printf("Will compute the number of molecules within a radius of %g\n", rcut);
232     }
233     if (!bNShell || opt2bSet("-o", NFILE, fnm))
234     {
235         bPDBout = (fn2ftp(opt2fn("-o", NFILE, fnm)) == efPDB);
236         if (bPDBout && !top.atoms.pdbinfo)
237         {
238             fprintf(stderr, "Creating pdbfino records\n");
239             snew(top.atoms.pdbinfo, top.atoms.nr);
240         }
241         out = open_trx(opt2fn("-o", NFILE, fnm), "w");
242     }
243     gpbc = gmx_rmpbc_init(&top.idef, pbcType, natoms);
244     do
245     {
246         gmx_rmpbc(gpbc, natoms, box, x);
247         set_pbc(&pbc, pbcType, box);
248
249         if (ref_a == -1)
250         {
251             /* Calculate the COM of all solvent molecules */
252             for (i = 0; i < nwat; i++)
253             {
254                 totmass = 0;
255                 clear_rvec(xsol[i]);
256                 for (j = 0; j < na; j++)
257                 {
258                     sa   = ind_sol[i * na + j];
259                     mass = top.atoms.atom[sa].m;
260                     totmass += mass;
261                     for (d = 0; d < DIM; d++)
262                     {
263                         xsol[i][d] += mass * x[sa][d];
264                     }
265                 }
266                 svmul(1.0 / totmass, xsol[i], xsol[i]);
267             }
268         }
269         else
270         {
271             /* Copy the reference atom of all solvent molecules */
272             for (i = 0; i < nwat; i++)
273             {
274                 copy_rvec(x[ind_sol[i * na + ref_a]], xsol[i]);
275             }
276         }
277
278         if (bZ)
279         {
280             for (i = 0; (i < nwat); i++)
281             {
282                 sa          = ind_sol[na * i];
283                 order[i].i  = sa;
284                 order[i].d2 = xsol[i][ZZ];
285             }
286         }
287         else if (bCOM)
288         {
289             totmass = 0;
290             clear_rvec(xcom);
291             for (i = 0; i < isize_ref; i++)
292             {
293                 mass = top.atoms.atom[ind_ref[i]].m;
294                 totmass += mass;
295                 for (j = 0; j < DIM; j++)
296                 {
297                     xcom[j] += mass * x[ind_ref[i]][j];
298                 }
299             }
300             svmul(1 / totmass, xcom, xcom);
301             for (i = 0; (i < nwat); i++)
302             {
303                 sa = ind_sol[na * i];
304                 pbc_dx(&pbc, xcom, xsol[i], dx);
305                 order[i].i  = sa;
306                 order[i].d2 = norm2(dx);
307             }
308         }
309         else
310         {
311             /* Set distance to first atom */
312             for (i = 0; (i < nwat); i++)
313             {
314                 sa = ind_sol[na * i];
315                 pbc_dx(&pbc, x[ind_ref[0]], xsol[i], dx);
316                 order[i].i  = sa;
317                 order[i].d2 = norm2(dx);
318             }
319             for (j = 1; (j < isize_ref); j++)
320             {
321                 sr = ind_ref[j];
322                 for (i = 0; (i < nwat); i++)
323                 {
324                     pbc_dx(&pbc, x[sr], xsol[i], dx);
325                     n2 = norm2(dx);
326                     if (n2 < order[i].d2)
327                     {
328                         order[i].d2 = n2;
329                     }
330                 }
331             }
332         }
333
334         if (bNShell)
335         {
336             ncut = 0;
337             for (i = 0; (i < nwat); i++)
338             {
339                 if (order[i].d2 <= rcut2)
340                 {
341                     ncut++;
342                 }
343             }
344             fprintf(fp, "%10.3f  %8d\n", t, ncut);
345         }
346         if (out)
347         {
348             qsort(order, nwat, sizeof(*order), ocomp);
349             for (i = 0; (i < nwat); i++)
350             {
351                 for (j = 0; (j < na); j++)
352                 {
353                     swi[ind_sol[na * i] + j] = order[i].i + j;
354                 }
355             }
356
357             /* Store the distance as the B-factor */
358             if (bPDBout)
359             {
360                 for (i = 0; (i < nwat); i++)
361                 {
362                     for (j = 0; (j < na); j++)
363                     {
364                         top.atoms.pdbinfo[order[i].i + j].bfac = std::sqrt(order[i].d2);
365                     }
366                 }
367             }
368             write_trx(out, natoms, swi, &top.atoms, 0, t, box, x, nullptr, nullptr);
369         }
370     } while (read_next_x(oenv, status, &t, x, box));
371     close_trx(status);
372     if (out)
373     {
374         close_trx(out);
375     }
376     if (fp)
377     {
378         xvgrclose(fp);
379     }
380     gmx_rmpbc_done(gpbc);
381
382     return 0;
383 }