Write forces and velocities to compressed TNG.
[alexxy/gromacs.git] / src / gromacs / mdlib / mdoutf.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2013,2014,2015, by the GROMACS development team, led by
5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6  * and including many others, as listed in the AUTHORS file in the
7  * top-level source directory and at http://www.gromacs.org.
8  *
9  * GROMACS is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public License
11  * as published by the Free Software Foundation; either version 2.1
12  * of the License, or (at your option) any later version.
13  *
14  * GROMACS is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with GROMACS; if not, see
21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
23  *
24  * If you want to redistribute modifications to GROMACS, please
25  * consider that scientific software is very special. Version
26  * control is crucial - bugs must be traceable. We will be happy to
27  * consider code for inclusion in the official distribution, but
28  * derived work must not be called official GROMACS. Details are found
29  * in the README & COPYING files - if they are missing, get the
30  * official version at http://www.gromacs.org.
31  *
32  * To help us fund GROMACS development, we humbly ask that you cite
33  * the research papers on the package. Check out http://www.gromacs.org.
34  */
35 #include "gmxpre.h"
36
37 #include "mdoutf.h"
38
39 #include "gromacs/commandline/filenm.h"
40 #include "gromacs/domdec/domdec.h"
41 #include "gromacs/domdec/domdec_struct.h"
42 #include "gromacs/fileio/checkpoint.h"
43 #include "gromacs/fileio/copyrite.h"
44 #include "gromacs/fileio/gmxfio.h"
45 #include "gromacs/fileio/tngio.h"
46 #include "gromacs/fileio/trrio.h"
47 #include "gromacs/fileio/xtcio.h"
48 #include "gromacs/fileio/xvgr.h"
49 #include "gromacs/math/vec.h"
50 #include "gromacs/mdlib/mdrun.h"
51 #include "gromacs/mdlib/trajectory_writing.h"
52 #include "gromacs/mdtypes/commrec.h"
53 #include "gromacs/mdtypes/md_enums.h"
54 #include "gromacs/timing/wallcycle.h"
55 #include "gromacs/utility/fatalerror.h"
56 #include "gromacs/utility/smalloc.h"
57
58 struct gmx_mdoutf {
59     t_fileio         *fp_trn;
60     t_fileio         *fp_xtc;
61     tng_trajectory_t  tng;
62     tng_trajectory_t  tng_low_prec;
63     int               x_compression_precision; /* only used by XTC output */
64     ener_file_t       fp_ene;
65     const char       *fn_cpt;
66     gmx_bool          bKeepAndNumCPT;
67     int               eIntegrator;
68     gmx_bool          bExpanded;
69     int               elamstats;
70     int               simulation_part;
71     FILE             *fp_dhdl;
72     FILE             *fp_field;
73     int               natoms_global;
74     int               natoms_x_compressed;
75     gmx_groups_t     *groups; /* for compressed position writing */
76     gmx_wallcycle_t   wcycle;
77 };
78
79
80 gmx_mdoutf_t init_mdoutf(FILE *fplog, int nfile, const t_filenm fnm[],
81                          int mdrun_flags, const t_commrec *cr,
82                          const t_inputrec *ir, gmx_mtop_t *top_global,
83                          const gmx_output_env_t *oenv, gmx_wallcycle_t wcycle)
84 {
85     gmx_mdoutf_t  of;
86     char          filemode[3];
87     gmx_bool      bAppendFiles, bCiteTng = FALSE;
88     int           i;
89
90     snew(of, 1);
91
92     of->fp_trn       = NULL;
93     of->fp_ene       = NULL;
94     of->fp_xtc       = NULL;
95     of->tng          = NULL;
96     of->tng_low_prec = NULL;
97     of->fp_dhdl      = NULL;
98     of->fp_field     = NULL;
99
100     of->eIntegrator             = ir->eI;
101     of->bExpanded               = ir->bExpanded;
102     of->elamstats               = ir->expandedvals->elamstats;
103     of->simulation_part         = ir->simulation_part;
104     of->x_compression_precision = static_cast<int>(ir->x_compression_precision);
105     of->wcycle                  = wcycle;
106
107     if (MASTER(cr))
108     {
109         bAppendFiles = (mdrun_flags & MD_APPENDFILES);
110
111         of->bKeepAndNumCPT = (mdrun_flags & MD_KEEPANDNUMCPT);
112
113         sprintf(filemode, bAppendFiles ? "a+" : "w+");
114
115         if (EI_DYNAMICS(ir->eI) &&
116             ir->nstxout_compressed > 0)
117         {
118             const char *filename;
119             filename = ftp2fn(efCOMPRESSED, nfile, fnm);
120             switch (fn2ftp(filename))
121             {
122                 case efXTC:
123                     of->fp_xtc                  = open_xtc(filename, filemode);
124                     break;
125                 case efTNG:
126                     gmx_tng_open(filename, filemode[0], &of->tng_low_prec);
127                     if (filemode[0] == 'w')
128                     {
129                         gmx_tng_prepare_low_prec_writing(of->tng_low_prec, top_global, ir);
130                     }
131                     bCiteTng = TRUE;
132                     break;
133                 default:
134                     gmx_incons("Invalid reduced precision file format");
135             }
136         }
137         if ((EI_DYNAMICS(ir->eI) || EI_ENERGY_MINIMIZATION(ir->eI))
138 #ifndef GMX_FAHCORE
139             &&
140             !(EI_DYNAMICS(ir->eI) &&
141               ir->nstxout == 0 &&
142               ir->nstvout == 0 &&
143               ir->nstfout == 0)
144 #endif
145             )
146         {
147             const char *filename;
148             filename = ftp2fn(efTRN, nfile, fnm);
149             switch (fn2ftp(filename))
150             {
151                 case efTRR:
152                 case efTRN:
153                     /* If there is no uncompressed coordinate output and
154                        there is compressed TNG output write forces
155                        and/or velocities to the TNG file instead. */
156                     if (ir->nstxout != 0 || ir->nstxout_compressed == 0 ||
157                         !of->tng_low_prec)
158                     {
159                         of->fp_trn = gmx_trr_open(filename, filemode);
160                     }
161                     break;
162                 case efTNG:
163                     gmx_tng_open(filename, filemode[0], &of->tng);
164                     if (filemode[0] == 'w')
165                     {
166                         gmx_tng_prepare_md_writing(of->tng, top_global, ir);
167                     }
168                     bCiteTng = TRUE;
169                     break;
170                 default:
171                     gmx_incons("Invalid full precision file format");
172             }
173         }
174         if (EI_DYNAMICS(ir->eI) || EI_ENERGY_MINIMIZATION(ir->eI))
175         {
176             of->fp_ene = open_enx(ftp2fn(efEDR, nfile, fnm), filemode);
177         }
178         of->fn_cpt = opt2fn("-cpo", nfile, fnm);
179
180         if ((ir->efep != efepNO || ir->bSimTemp) && ir->fepvals->nstdhdl > 0 &&
181             (ir->fepvals->separate_dhdl_file == esepdhdlfileYES ) &&
182             EI_DYNAMICS(ir->eI))
183         {
184             if (bAppendFiles)
185             {
186                 of->fp_dhdl = gmx_fio_fopen(opt2fn("-dhdl", nfile, fnm), filemode);
187             }
188             else
189             {
190                 of->fp_dhdl = open_dhdl(opt2fn("-dhdl", nfile, fnm), ir, oenv);
191             }
192         }
193
194         if (opt2bSet("-field", nfile, fnm) &&
195             (ir->ex[XX].n || ir->ex[YY].n || ir->ex[ZZ].n))
196         {
197             if (bAppendFiles)
198             {
199                 of->fp_field = gmx_fio_fopen(opt2fn("-field", nfile, fnm),
200                                              filemode);
201             }
202             else
203             {
204                 of->fp_field = xvgropen(opt2fn("-field", nfile, fnm),
205                                         "Applied electric field", "Time (ps)",
206                                         "E (V/nm)", oenv);
207             }
208         }
209
210         /* Set up atom counts so they can be passed to actual
211            trajectory-writing routines later. Also, XTC writing needs
212            to know what (and how many) atoms might be in the XTC
213            groups, and how to look up later which ones they are. */
214         of->natoms_global       = top_global->natoms;
215         of->groups              = &top_global->groups;
216         of->natoms_x_compressed = 0;
217         for (i = 0; (i < top_global->natoms); i++)
218         {
219             if (ggrpnr(of->groups, egcCompressedX, i) == 0)
220             {
221                 of->natoms_x_compressed++;
222             }
223         }
224     }
225
226     if (bCiteTng)
227     {
228         please_cite(fplog, "Lundborg2014");
229     }
230
231     return of;
232 }
233
234 FILE *mdoutf_get_fp_field(gmx_mdoutf_t of)
235 {
236     return of->fp_field;
237 }
238
239 ener_file_t mdoutf_get_fp_ene(gmx_mdoutf_t of)
240 {
241     return of->fp_ene;
242 }
243
244 FILE *mdoutf_get_fp_dhdl(gmx_mdoutf_t of)
245 {
246     return of->fp_dhdl;
247 }
248
249 gmx_wallcycle_t mdoutf_get_wcycle(gmx_mdoutf_t of)
250 {
251     return of->wcycle;
252 }
253
254 void mdoutf_write_to_trajectory_files(FILE *fplog, t_commrec *cr,
255                                       gmx_mdoutf_t of,
256                                       int mdof_flags,
257                                       gmx_mtop_t *top_global,
258                                       gmx_int64_t step, double t,
259                                       t_state *state_local, t_state *state_global,
260                                       rvec *f_local, rvec *f_global)
261 {
262     rvec *local_v;
263     rvec *global_v;
264
265     /* MRS -- defining these variables is to manage the difference
266      * between half step and full step velocities, but there must be a better way . . . */
267
268     local_v  = state_local->v;
269     global_v = state_global->v;
270
271     if (DOMAINDECOMP(cr))
272     {
273         if (mdof_flags & MDOF_CPT)
274         {
275             dd_collect_state(cr->dd, state_local, state_global);
276         }
277         else
278         {
279             if (mdof_flags & (MDOF_X | MDOF_X_COMPRESSED))
280             {
281                 dd_collect_vec(cr->dd, state_local, state_local->x,
282                                state_global->x);
283             }
284             if (mdof_flags & MDOF_V)
285             {
286                 dd_collect_vec(cr->dd, state_local, local_v,
287                                global_v);
288             }
289         }
290         if (mdof_flags & MDOF_F)
291         {
292             dd_collect_vec(cr->dd, state_local, f_local, f_global);
293         }
294     }
295     else
296     {
297         if (mdof_flags & MDOF_CPT)
298         {
299             /* All pointers in state_local are equal to state_global,
300              * but we need to copy the non-pointer entries.
301              */
302             state_global->lambda = state_local->lambda;
303             state_global->veta   = state_local->veta;
304             state_global->vol0   = state_local->vol0;
305             copy_mat(state_local->box, state_global->box);
306             copy_mat(state_local->boxv, state_global->boxv);
307             copy_mat(state_local->svir_prev, state_global->svir_prev);
308             copy_mat(state_local->fvir_prev, state_global->fvir_prev);
309             copy_mat(state_local->pres_prev, state_global->pres_prev);
310         }
311     }
312
313     if (MASTER(cr))
314     {
315         if (mdof_flags & MDOF_CPT)
316         {
317             fflush_tng(of->tng);
318             fflush_tng(of->tng_low_prec);
319             ivec one_ivec = { 1, 1, 1 };
320             write_checkpoint(of->fn_cpt, of->bKeepAndNumCPT,
321                              fplog, cr,
322                              DOMAINDECOMP(cr) ? cr->dd->nc : one_ivec,
323                              DOMAINDECOMP(cr) ? cr->dd->nnodes : cr->nnodes,
324                              of->eIntegrator, of->simulation_part,
325                              of->bExpanded, of->elamstats, step, t, state_global);
326         }
327
328         if (mdof_flags & (MDOF_X | MDOF_V | MDOF_F))
329         {
330             if (of->fp_trn)
331             {
332                 gmx_trr_write_frame(of->fp_trn, step, t, state_local->lambda[efptFEP],
333                                     state_local->box, top_global->natoms,
334                                     (mdof_flags & MDOF_X) ? state_global->x : NULL,
335                                     (mdof_flags & MDOF_V) ? global_v : NULL,
336                                     (mdof_flags & MDOF_F) ? f_global : NULL);
337                 if (gmx_fio_flush(of->fp_trn) != 0)
338                 {
339                     gmx_file("Cannot write trajectory; maybe you are out of disk space?");
340                 }
341             }
342
343             /* If a TNG file is open for uncompressed coordinate output also write
344                velocities and forces to it. */
345             else if (of->tng)
346             {
347                 gmx_fwrite_tng(of->tng, FALSE, step, t, state_local->lambda[efptFEP],
348                                state_local->box,
349                                top_global->natoms,
350                                (mdof_flags & MDOF_X) ? state_global->x : NULL,
351                                (mdof_flags & MDOF_V) ? global_v : NULL,
352                                (mdof_flags & MDOF_F) ? f_global : NULL);
353             }
354             /* If only a TNG file is open for compressed coordinate output (no uncompressed
355                coordinate output) also write forces and velocities to it. */
356             else if (of->tng_low_prec)
357             {
358                 gmx_fwrite_tng(of->tng_low_prec, FALSE, step, t, state_local->lambda[efptFEP],
359                                state_local->box,
360                                top_global->natoms,
361                                (mdof_flags & MDOF_X) ? state_global->x : NULL,
362                                (mdof_flags & MDOF_V) ? global_v : NULL,
363                                (mdof_flags & MDOF_F) ? f_global : NULL);
364             }
365         }
366         if (mdof_flags & MDOF_X_COMPRESSED)
367         {
368             rvec *xxtc = NULL;
369
370             if (of->natoms_x_compressed == of->natoms_global)
371             {
372                 /* We are writing the positions of all of the atoms to
373                    the compressed output */
374                 xxtc = state_global->x;
375             }
376             else
377             {
378                 /* We are writing the positions of only a subset of
379                    the atoms to the compressed output, so we have to
380                    make a copy of the subset of coordinates. */
381                 int i, j;
382
383                 snew(xxtc, of->natoms_x_compressed);
384                 for (i = 0, j = 0; (i < of->natoms_global); i++)
385                 {
386                     if (ggrpnr(of->groups, egcCompressedX, i) == 0)
387                     {
388                         copy_rvec(state_global->x[i], xxtc[j++]);
389                     }
390                 }
391             }
392             if (write_xtc(of->fp_xtc, of->natoms_x_compressed, step, t,
393                           state_local->box, xxtc, of->x_compression_precision) == 0)
394             {
395                 gmx_fatal(FARGS, "XTC error - maybe you are out of disk space?");
396             }
397             gmx_fwrite_tng(of->tng_low_prec,
398                            TRUE,
399                            step,
400                            t,
401                            state_local->lambda[efptFEP],
402                            state_local->box,
403                            of->natoms_x_compressed,
404                            xxtc,
405                            NULL,
406                            NULL);
407             if (of->natoms_x_compressed != of->natoms_global)
408             {
409                 sfree(xxtc);
410             }
411         }
412     }
413 }
414
415 void mdoutf_tng_close(gmx_mdoutf_t of)
416 {
417     if (of->tng || of->tng_low_prec)
418     {
419         wallcycle_start(of->wcycle, ewcTRAJ);
420         gmx_tng_close(&of->tng);
421         gmx_tng_close(&of->tng_low_prec);
422         wallcycle_stop(of->wcycle, ewcTRAJ);
423     }
424 }
425
426 void done_mdoutf(gmx_mdoutf_t of)
427 {
428     if (of->fp_ene != NULL)
429     {
430         close_enx(of->fp_ene);
431     }
432     if (of->fp_xtc)
433     {
434         close_xtc(of->fp_xtc);
435     }
436     if (of->fp_trn)
437     {
438         gmx_trr_close(of->fp_trn);
439     }
440     if (of->fp_dhdl != NULL)
441     {
442         gmx_fio_fclose(of->fp_dhdl);
443     }
444     if (of->fp_field != NULL)
445     {
446         /* This is opened sometimes with xvgropen, sometimes with
447          * gmx_fio_fopen, so we use the least common denominator for closing.
448          */
449         gmx_fio_fclose(of->fp_field);
450     }
451
452     gmx_tng_close(&of->tng);
453     gmx_tng_close(&of->tng_low_prec);
454
455     sfree(of);
456 }