Add TNG writing and reading support
[alexxy/gromacs.git] / src / gromacs / fileio / trxio.c
index 7367fa92aa907b4fc24239c9082c743a8937c86a..74b4831b8bfd05b737a433c60123708f731d6e0a 100644 (file)
@@ -54,6 +54,8 @@
 #include "trxio.h"
 #include "tpxio.h"
 #include "trnio.h"
+#include "tngio.h"
+#include "tngio_for_tools.h"
 #include "names.h"
 #include "vec.h"
 #include "futil.h"
@@ -77,15 +79,16 @@ typedef enum {
 
 struct t_trxstatus
 {
-    int             __frame;
-    t_trxframe     *xframe;
-    int             nxframe;
-    t_fileio       *fio;
-    eFileFormat     eFF;
-    int             NATOMS;
-    double          DT, BOX[3];
-    gmx_bool        bReadBox;
-    char           *persistent_line; /* Persistent line for reading g96 trajectories */
+    int                     __frame;
+    t_trxframe             *xframe;
+    int                     nxframe;
+    t_fileio               *fio;
+    tng_trajectory_t        tng;
+    eFileFormat             eFF;
+    int                     NATOMS;
+    double                  DT, BOX[3];
+    gmx_bool                bReadBox;
+    char                   *persistent_line; /* Persistent line for reading g96 trajectories */
 };
 
 /* utility functions */
@@ -162,6 +165,7 @@ static void status_init(t_trxstatus *status)
     status->fio             = NULL;
     status->__frame         = -1;
     status->persistent_line = NULL;
+    status->tng             = NULL;
 }
 
 
@@ -219,46 +223,53 @@ int prec2ndec(real prec)
     return (int)(log(prec)/log(10.0)+0.5);
 }
 
+real ndec2prec(int ndec)
+{
+    return pow(10.0, ndec);
+}
 
 t_fileio *trx_get_fileio(t_trxstatus *status)
 {
     return status->fio;
 }
 
-
+tng_trajectory_t trx_get_tng(t_trxstatus *status)
+{
+    return status->tng;
+}
 
 void clear_trxframe(t_trxframe *fr, gmx_bool bFirst)
 {
-    fr->not_ok  = 0;
-    fr->bTitle  = FALSE;
-    fr->bStep   = FALSE;
-    fr->bTime   = FALSE;
-    fr->bLambda = FALSE;
+    fr->not_ok    = 0;
+    fr->bTitle    = FALSE;
+    fr->bStep     = FALSE;
+    fr->bTime     = FALSE;
+    fr->bLambda   = FALSE;
     fr->bFepState = FALSE;
-    fr->bAtoms  = FALSE;
-    fr->bPrec   = FALSE;
-    fr->bX      = FALSE;
-    fr->bV      = FALSE;
-    fr->bF      = FALSE;
-    fr->bBox    = FALSE;
+    fr->bAtoms    = FALSE;
+    fr->bPrec     = FALSE;
+    fr->bX        = FALSE;
+    fr->bV        = FALSE;
+    fr->bF        = FALSE;
+    fr->bBox      = FALSE;
     if (bFirst)
     {
-        fr->flags   = 0;
-        fr->bDouble = FALSE;
-        fr->natoms  = -1;
-        fr->t0      = 0;
-        fr->tpf     = 0;
-        fr->tppf    = 0;
-        fr->title   = NULL;
-        fr->step    = 0;
-        fr->time    = 0;
-        fr->lambda  = 0;
+        fr->flags     = 0;
+        fr->bDouble   = FALSE;
+        fr->natoms    = -1;
+        fr->t0        = 0;
+        fr->tpf       = 0;
+        fr->tppf      = 0;
+        fr->title     = NULL;
+        fr->step      = 0;
+        fr->time      = 0;
+        fr->lambda    = 0;
         fr->fep_state = 0;
-        fr->atoms   = NULL;
-        fr->prec    = 0;
-        fr->x       = NULL;
-        fr->v       = NULL;
-        fr->f       = NULL;
+        fr->atoms     = NULL;
+        fr->prec      = 0;
+        fr->x         = NULL;
+        fr->v         = NULL;
+        fr->f         = NULL;
         clear_mat(fr->box);
         fr->bPBC   = FALSE;
         fr->ePBC   = -1;
@@ -276,7 +287,7 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
 {
     char  title[STRLEN];
     rvec *xout = NULL, *vout = NULL, *fout = NULL;
-    int   i;
+    int   i, ftp = -1;
     real  prec;
 
     if (fr->bPrec)
@@ -288,24 +299,39 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
         prec = 1000.0;
     }
 
-    switch (gmx_fio_getftp(status->fio))
+    if (status->tng)
+    {
+        ftp = efTNG;
+    }
+    else if (status->fio)
+    {
+        ftp = gmx_fio_getftp(status->fio);
+    }
+    else
+    {
+        gmx_incons("No input file available");
+    }
+
+    switch (ftp)
     {
         case efTRJ:
         case efTRR:
+        case efTNG:
             break;
         default:
             if (!fr->bX)
             {
                 gmx_fatal(FARGS, "Need coordinates to write a %s trajectory",
-                          ftp2ext(gmx_fio_getftp(status->fio)));
+                          ftp2ext(ftp));
             }
             break;
     }
 
-    switch (gmx_fio_getftp(status->fio))
+    switch (ftp)
     {
         case efTRJ:
         case efTRR:
+        case efTNG:
             if (fr->bV)
             {
                 snew(vout, nind);
@@ -338,8 +364,11 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
             break;
     }
 
-    switch (gmx_fio_getftp(status->fio))
+    switch (ftp)
     {
+        case efTNG:
+            gmx_write_tng_from_trxframe(status->tng, fr, nind);
+            break;
         case efXTC:
             write_xtc(status->fio, nind, fr->step, fr->time, fr->box, xout, prec);
             break;
@@ -355,10 +384,10 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
             if (!fr->bAtoms)
             {
                 gmx_fatal(FARGS, "Can not write a %s file without atom names",
-                          ftp2ext(gmx_fio_getftp(status->fio)));
+                          ftp2ext(ftp));
             }
             sprintf(title, "frame t= %.3f", fr->time);
-            if (gmx_fio_getftp(status->fio) == efGRO)
+            if (ftp == efGRO)
             {
                 write_hconf_indexed_p(gmx_fio_getfp(status->fio), title, fr->atoms, nind, ind,
                                       prec2ndec(prec),
@@ -378,15 +407,16 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
             break;
         default:
             gmx_fatal(FARGS, "Sorry, write_trxframe_indexed can not write %s",
-                      ftp2ext(gmx_fio_getftp(status->fio)));
+                      ftp2ext(ftp));
             break;
     }
 
-    switch (gmx_fio_getftp(status->fio))
+    switch (ftp)
     {
         case efTRN:
         case efTRJ:
         case efTRR:
+        case efTNG:
             if (vout)
             {
                 sfree(vout);
@@ -407,6 +437,60 @@ int write_trxframe_indexed(t_trxstatus *status, t_trxframe *fr, int nind,
     return 0;
 }
 
+void trjtools_gmx_prepare_tng_writing(const char       *filename,
+                                      char              filemode,
+                                      t_trxstatus      *in,
+                                      t_trxstatus     **out,
+                                      const char       *infile,
+                                      const int         natoms,
+                                      const gmx_mtop_t *mtop,
+                                      const atom_id    *index,
+                                      const char       *index_group_name)
+{
+    if (filemode != 'w' && filemode != 'a')
+    {
+        gmx_incons("Sorry, can only prepare for TNG output.");
+    }
+
+    if (*out == NULL)
+    {
+        snew((*out), 1);
+    }
+    status_init(*out);
+
+    if (in != NULL)
+    {
+        gmx_prepare_tng_writing(filename,
+                                filemode,
+                                &in->tng,
+                                &(*out)->tng,
+                                natoms,
+                                mtop,
+                                index,
+                                index_group_name);
+    }
+    else if (efTNG == fn2ftp(infile))
+    {
+        tng_trajectory_t tng_in;
+        gmx_tng_open(infile, 'r', &tng_in);
+
+        gmx_prepare_tng_writing(filename,
+                                filemode,
+                                &tng_in,
+                                &(*out)->tng,
+                                natoms,
+                                mtop,
+                                index,
+                                index_group_name);
+    }
+}
+
+void write_tng_frame(t_trxstatus *status,
+                     t_trxframe  *frame)
+{
+    gmx_write_tng_from_trxframe(status->tng, frame, -1);
+}
+
 int write_trxframe(t_trxstatus *status, t_trxframe *fr, gmx_conect gc)
 {
     char title[STRLEN];
@@ -421,6 +505,14 @@ int write_trxframe(t_trxstatus *status, t_trxframe *fr, gmx_conect gc)
         prec = 1000.0;
     }
 
+    if (status->tng)
+    {
+        gmx_tng_set_compression_precision(status->tng, prec);
+        write_tng_frame(status, fr);
+
+        return 0;
+    }
+
     switch (gmx_fio_getftp(status->fio))
     {
         case efTRJ:
@@ -507,7 +599,11 @@ int write_trx(t_trxstatus *status, int nind, const atom_id *ind, t_atoms *atoms,
 
 void close_trx(t_trxstatus *status)
 {
-    gmx_fio_close(status->fio);
+    gmx_tng_close(&status->tng);
+    if (status->fio)
+    {
+        gmx_fio_close(status->fio);
+    }
     sfree(status);
 }
 
@@ -870,6 +966,7 @@ gmx_bool read_next_frame(const output_env_t oenv, t_trxstatus *status, t_trxfram
     int      ct;
     gmx_bool bOK, bRet, bMissingData = FALSE, bSkip = FALSE;
     int      dummy = 0;
+    int      ftp;
 
     bRet = FALSE;
     pt   = fr->time;
@@ -880,7 +977,16 @@ gmx_bool read_next_frame(const output_env_t oenv, t_trxstatus *status, t_trxfram
         fr->tppf = fr->tpf;
         fr->tpf  = fr->time;
 
-        switch (gmx_fio_getftp(status->fio))
+        if (status->tng)
+        {
+            /* Special treatment for TNG files */
+            ftp = efTNG;
+        }
+        else
+        {
+            ftp = gmx_fio_getftp(status->fio);
+        }
+        switch (ftp)
         {
             case efTRJ:
             case efTRR:
@@ -932,6 +1038,9 @@ gmx_bool read_next_frame(const output_env_t oenv, t_trxstatus *status, t_trxfram
                     fr->not_ok = DATA_NOT_OK;
                 }
                 break;
+            case efTNG:
+                bRet = gmx_read_next_tng_frame(status->tng, fr, NULL, 0);
+                break;
             case efPDB:
                 bRet = pdb_next_x(status, gmx_fio_getfp(status->fio), fr);
                 break;
@@ -991,9 +1100,11 @@ gmx_bool read_next_frame(const output_env_t oenv, t_trxstatus *status, t_trxfram
 int read_first_frame(const output_env_t oenv, t_trxstatus **status,
                      const char *fn, t_trxframe *fr, int flags)
 {
-    t_fileio *fio;
-    gmx_bool  bFirst, bOK;
-    int       dummy = 0;
+    t_fileio      *fio;
+    gmx_bool       bFirst, bOK;
+    int            dummy = 0;
+    int            ftp   = fn2ftp(fn);
+    gmx_int64_t   *tng_ids;
 
     clear_trxframe(fr, TRUE);
     fr->flags = flags;
@@ -1006,8 +1117,16 @@ int read_first_frame(const output_env_t oenv, t_trxstatus **status,
     (*status)->nxframe = 1;
     initcount(*status);
 
-    fio = (*status)->fio = gmx_fio_open(fn, "r");
-    switch (gmx_fio_getftp(fio))
+    if (efTNG == ftp)
+    {
+        /* Special treatment for TNG files */
+        gmx_tng_open(fn, 'r', &(*status)->tng);
+    }
+    else
+    {
+        fio = (*status)->fio = gmx_fio_open(fn, "r");
+    }
+    switch (ftp)
     {
         case efTRJ:
         case efTRR:
@@ -1071,6 +1190,20 @@ int read_first_frame(const output_env_t oenv, t_trxstatus **status,
             }
             bFirst = FALSE;
             break;
+        case efTNG:
+            fr->step = -1;
+            if (!gmx_read_next_tng_frame((*status)->tng, fr, NULL, 0))
+            {
+                fr->not_ok = DATA_NOT_OK;
+                fr->natoms = 0;
+                printincomp(*status, fr);
+            }
+            else
+            {
+                printcount(*status, oenv, fr->time, FALSE);
+            }
+            bFirst = FALSE;
+            break;
         case efPDB:
             pdb_first_x(*status, gmx_fio_getfp(fio), fr);
             if (fr->natoms)
@@ -1161,7 +1294,12 @@ gmx_bool read_next_x(const output_env_t oenv, t_trxstatus *status, real *t,
 
 void close_trj(t_trxstatus *status)
 {
-    gmx_fio_close(status->fio);
+    gmx_tng_close(&status->tng);
+    if (status->fio)
+    {
+        gmx_fio_close(status->fio);
+    }
+
     /* The memory in status->xframe is lost here,
      * but the read_first_x/read_next_x functions are deprecated anyhow.
      * read_first_frame/read_next_frame and close_trx should be used.