From d178df5c7f60ea57d3cb6775449f29395fe75d20 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niels=20M=C3=BCller?= <niels.mueller@physik.tu-darmstadt.de>
Date: Fri, 18 Aug 2017 13:22:58 +0200
Subject: [PATCH] Add function to read single xtc frmames.

---
 pygmx/__init__.py |  2 +-
 pygmx/fileio.pxd  | 27 ++++++++++++++++++++++++++-
 pygmx/xtcio.pyx   | 33 ++++++++++++++++++++++++++++++---
 3 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/pygmx/__init__.py b/pygmx/__init__.py
index d598446..d0cbfe5 100644
--- a/pygmx/__init__.py
+++ b/pygmx/__init__.py
@@ -8,7 +8,7 @@ Trajectories in trr format may be read with `.gromacs.reader.TRRReader`, which i
 import os
 
 from .tpxio import TPXReader
-from .xtcio import XTCReader
+from .xtcio import XTCReader, read_xtcframe
 from .enxio import EDRFile
 from .errors import FileTypeError
 from .gromacs.reader import index_filename_for_xtc
diff --git a/pygmx/fileio.pxd b/pygmx/fileio.pxd
index afb0bc0..8bde95c 100644
--- a/pygmx/fileio.pxd
+++ b/pygmx/fileio.pxd
@@ -1,4 +1,4 @@
-# from libc.stdio cimport FILE
+from libc.stdio cimport FILE
 
 from utility cimport *
 
@@ -6,11 +6,36 @@ from utility cimport *
 #    ctypedef struct XDR:
 #        pass
 
+cdef extern from "gromacs/fileio/filetypes.h":
+    int fn2ftp(const char *fn)
+
+#cdef extern from "gromacs/fileio/gmx_internal_xdr.h":
+#    ctypedef struct XDR:
+#        pass
+#        
+#    ctypedef enum xdr_op:
+#        XDR_ENCODE = 0
+#        XDR_DECODE = 1
+#        XDR_FREE   = 2
+#
+#    void xdrstdio_create (XDR *__xdrs, FILE *__file, xdr_op __xop)
 
 cdef extern from "gromacs/fileio/gmxfio.h":
 
     ctypedef struct t_fileio:
         pass
+        # FILE           *fp               # the file pointer */
+        # gmx_bool        bRead             # the file is open for reading */
+        # gmx_bool        bDouble           # write doubles instead of floats */
+        # gmx_bool        bReadWrite        # the file is open for reading and writing */
+        # char        *fn                   # the file name */
+        # # XDR         *xdr                  # the xdr data pointer */
+        # int  xdrmode              # the xdr mode */
+        # int          iFTP                 # the file type identifier */
+
+        # t_fileio    *next, *prev          # next and previous file pointers in the linked list */
+        # # tMPI_Lock_t  mtx;                  # content locking mutex. This is a fast lock
+    
 
     void gmx_fio_rewind(t_fileio *fio)
 
diff --git a/pygmx/xtcio.pyx b/pygmx/xtcio.pyx
index 4b93a32..1694ac1 100644
--- a/pygmx/xtcio.pyx
+++ b/pygmx/xtcio.pyx
@@ -16,9 +16,9 @@ from .errors import InvalidIndexException, InvalidMagicException, XTCError
 
 
 cdef extern from "gromacs/fileio/xtcio.h":
-    t_fileio *open_xtc(const char *filename, const char *mode)
+    t_fileio *open_xtc(const char *filename, const char *mode) nogil
 
-    void close_xtc(t_fileio *fio)
+    void close_xtc(t_fileio *fio) nogil
 
     int read_first_xtc(t_fileio *fio,
                        int *natoms, gmx_int64_t *step, real *time,
@@ -26,7 +26,7 @@ cdef extern from "gromacs/fileio/xtcio.h":
 
     int read_next_xtc(t_fileio *fio,
                       int natoms, gmx_int64_t *step, real *time,
-                      matrix box, rvec *x, real *prec, gmx_bool *_bOK)
+                      matrix box, rvec *x, real *prec, gmx_bool *_bOK) nogil
 
 
 if sizeof(real) == 4:
@@ -69,6 +69,33 @@ cdef array get_xtc_index(t_fileio *fio):
     return cache[:-1]
 
 
+def read_xtcframe(fname, pos, natoms):
+    cdef t_fileio *fio = open_xtc(fname.encode(), 'r')
+    gmx_fio_seek(fio, pos)
+
+    cdef:
+        matrix box
+        gmx_bool _bOK
+        real time, prec
+        gmx_int64_t cur_step
+        np.ndarray[real, ndim=2] coords = np.empty((natoms, 3), dtype=np_real)
+        int nratms = natoms
+
+    with nogil:
+        read_next_xtc(fio, nratms, &cur_step, &time, box, <rvec *>coords.data, &prec, &_bOK)
+        close_xtc(fio)
+
+    if _bOK:
+        frame = XTCFrame()
+        frame._coordinates = coords
+        frame.index = cur_step
+        frame.time = time
+        frame.box = box
+        return frame
+    else:
+        raise
+
+
 cdef class XTCReader:
     cdef:
         t_fileio *fio