Bundle commit of all recent changes on pygmx.

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "gromacs"]
 	path = gromacs
-	url = https://github.com/gromacs/gromacs.git 
+	url = git://git.gromacs.org/gromacs.git 

pygmx/__init__.py

@@ -7,3 +7,4 @@ from . import xtcio
 from . import tngio
 
 from .tpxio import TPXReader
+from .enxio import EDRFile

pygmx/enxio.pyx

@@ -0,0 +1,99 @@
+
+from cpython.array cimport array
+from array import array
+
+cimport numpy as np
+import numpy as np
+
+from utility cimport *
+from mdtypes cimport *
+
+cdef extern from "gromacs/fileio/enxio.h":
+    ctypedef struct gmx_enxnm_t:
+        char *name
+        char *unit
+
+    ctypedef struct ener_file_t:
+        pass
+
+    ctypedef struct t_enxblock:
+        pass
+
+    ctypedef struct t_enxframe:
+        double          t            # Timestamp of this frame	                     */
+        gmx_int64_t     step         # MD step	                             */
+        gmx_int64_t     nsteps       # The number of steps between frames            */
+        double          dt           # The MD time step                              */
+        int             nsum         # The number of terms for the sums in ener      */
+        int             nre          # Number of energies			     */
+        int             e_size       # Size (in bytes) of energies		     */
+        int             e_alloc      # Allocated size (in elements) of ener          */
+        t_energy       *ener         # The energies                                  */
+        int             nblock       # Number of following energy blocks             */
+        t_enxblock     *block        # The blocks                                    */
+        int             nblock_alloc # The number of blocks allocated                */
+
+    ener_file_t open_enx(const char *fn, const char *mode)
+    void close_enx(ener_file_t ef)
+
+    void do_enxnms(ener_file_t ef, int *nre, gmx_enxnm_t **enms)
+
+    gmx_bool do_enx(ener_file_t ef, t_enxframe *fr)
+
+
+
+cdef class EDRFile:
+
+    cdef:
+        ener_file_t efile
+        gmx_enxnm_t *etypes
+        int n_etypes
+        t_enxframe *frames
+
+    @property
+    def types(self):
+        types = []
+        for i in range(self.n_etypes):
+            types.append(
+
+            )
+    def read(self):
+        cdef t_enxframe *frame
+        cdef np.ndarray[float, ndim=2] energies = np.empty((1,self.n_etypes), np.float32)
+
+        snew(frame, 1)
+
+        while do_enx(self.efile, frame):
+            energies = np.resize(energies, (len(energies)+1, self.n_etypes))
+            for i in range(frame.nre):
+                energies[-1, i] = frame.ener[i].e
+        return energies
+
+
+    def __init__(self, filename):
+        filename = cstr(filename)
+
+        self.efile = open_enx(filename, b'r')
+
+        self.etypes = NULL
+        do_enxnms(self.efile, &self.n_etypes, &self.etypes)
+        print(self.n_etypes)
+
+        #do_enx(self.efile, self.frames)
+
+
+    def __getitem__(self, item):
+        cdef:
+            int f = 0
+            t_enxframe frame = self.frames[f]
+
+        cdef array values = array('d')
+
+        while frame.nblock > 0:
+            if item < frame.nre:
+                values.append(frame.ener[item].e)
+            i += 1
+            frame = self.frames[i]
+            print(i, frame.t, frame.nblock)
+
+        return values

pygmx/fileio.pxd

@@ -1,12 +1,24 @@
+# from libc.stdio cimport FILE
 
 from utility cimport *
 
+#cdef extern from "gromacs/fileio/gmx_system_xdr.h":
+#    ctypedef struct XDR:
+#        pass
+
+
 cdef extern from "gromacs/fileio/gmxfio.h":
+
     ctypedef struct t_fileio:
         pass
 
+    void gmx_fio_rewind(t_fileio *fio)
+
     int gmx_fio_seek(t_fileio *fio, gmx_off_t fpos)
 
     gmx_off_t gmx_fio_ftell(t_fileio *fio)
 
+#    int xtc_seek_frame(t_fileio *fio, int frame, int natoms)
     int xtc_seek_time(t_fileio *fio, real time, int natoms, gmx_bool bSeekForwardOnly)
+
+    void gmx_fio_setdebug(t_fileio *fio, gmx_bool bDebug)

pygmx/mdtypes.pxd

@@ -2,6 +2,13 @@
 from utility cimport *
 from math cimport *
 
+cdef extern from "gromacs/legacyheaders/types/energy.h":
+    ctypedef struct t_energy:
+        real    e
+        double  eav
+        double  esum
+        
+
 cdef extern from "gromacs/legacyheaders/types/inputrec.h":
     ctypedef struct t_simtemp:
         pass

pygmx/utility.pxd

@@ -1,7 +1,7 @@
 # C-API in gromacs/utility
 
 #cdef extern from "inttypes.h":
-ctypedef int __int64
+ctypedef unsigned long __int64
 
 cdef extern from "gromacs/utility/basedefinitions.h":
     ctypedef int gmx_bool
@@ -12,3 +12,12 @@ cdef extern from "gromacs/utility/real.h":
 
 cdef extern from "gromacs/utility/futil.h":
     ctypedef gmx_int64_t    gmx_off_t
+
+cdef extern from "gromacs/utility/smalloc.h":
+    void snew(void *ptr, int nelem)
+
+
+cdef inline cstr(instr):
+    if isinstance(instr, str):
+        instr = instr.encode()
+    return instr

pygmx/xtcio.pyx

@@ -1,5 +1,8 @@
 # Wrapper for xtcio.h: I/O for xtc files.
 
+from cpython.array cimport array
+from array import array
+
 import numpy as np
 cimport numpy as np
 
@@ -7,6 +10,7 @@ from utility cimport *
 from math cimport *
 from fileio cimport *
 
+
 cdef extern from "gromacs/fileio/xtcio.h":
     t_fileio *open_xtc(const char *filename, const char *mode)
 
@@ -27,69 +31,96 @@ else:
     np_real = np.float
 
 
+#cdef array get_xtc_index_by_frames(t_fileio *fio, int length, int natoms):
+#    cdef:
+        #gmx_bool _bOK
+##        int frame = 1
+#    cdef array cache = array('L')
+#    gmx_fio_rewind(fio)
+#    cache.append(gmx_fio_ftell(fio))
+#    while xtc_seek_frame(fio, frame*500, natoms) == 0:
+#        cache.append(gmx_fio_ftell(fio))
+#        #print(frame, cache[-1])
+#        frame += 1
+#        if frame == length:
+#            break
+#    return cache
+
+cdef array get_xtc_index(t_fileio *fio):
+    cdef:
+        gmx_bool _bOK
+        int natoms, step, frame, state = 1
+        real time, prec
+        matrix box
+        rvec *x
+    cdef array cache = array('L')
+    gmx_fio_rewind(fio)
+    cache.append(gmx_fio_ftell(fio))
+    read_first_xtc(fio, &natoms, &step, &time, box, &x, &prec, &_bOK)
+    cache.append(gmx_fio_ftell(fio))
+    while read_next_xtc(fio, natoms, &step, &time, box, x, &prec, &_bOK):
+        cache.append(gmx_fio_ftell(fio))
+    # the last index is invalid
+    return cache[:-1]
+
+
 cdef class XTCReader:
     cdef:
         t_fileio *fio
         int natoms, cur_step
         real start_time, timestep, prec, cur_time
-        dict _cache
+        bint has_cache
-        #rvec *coords
+        rvec *coords
+        array _cache
 
-    def nearest_cache(self, index):
+    @property
-        nearest = 0
+    def cache(self):
-        cur_dist = 99999999
+        return self._cache
-        for ind in self._cache:
+
-            dist = abs(ind - index)
+    def seek(self, int frame):
-            if dist < cur_dist:
+        if self.has_cache:
-                nearest = ind
+            gmx_fio_seek(self.fio, self.cache[frame])
-                cur_dist = dist
-        return nearest
 
 
-    def seek(self, index):
+    def make_cache(self):
-        if index in self._cache:
+        self._cache = get_xtc_index(self.fio)
-            gmx_fio_seek(self.fio, self._cache[index])
+        self.has_cache = True
-        else:
-            nearest = self.nearest_cache(index)
-            gmx_fio_seek(self.fio, self._cache[nearest])
-            self.cur_step = index
-            self.cur_time = index * self.timestep + self.start_time
-            xtc_seek_time(self.fio, self.cur_time, self.natoms, index > nearest)
-            self._cache[index] = gmx_fio_ftell(self.fio)
 
+    def __cinit__(self):
+        self.has_cache = False
 
-    def __init__(self, filename):
+    def __init__(self, filename, make_cache=True):
         if isinstance(filename, str):
             filename = filename.encode()
 
         cdef:
+            int step
             matrix box
             gmx_bool _bOK
-            real time
+            real time, prec
-        cdef rvec *coords # = np.empty((1,3), dtype=np.float32)
+            rvec *x
 
-        self._cache = {}
         self.fio = open_xtc(filename, b'r')
-        read_first_xtc(self.fio, &self.natoms, &self.cur_step, &self.start_time, box,
+        read_first_xtc(self.fio, &self.natoms, &step, &time, box, &x, &prec, &_bOK)
-            &coords, &self.prec, &_bOK)
-        self._cache[self.cur_step] = gmx_fio_ftell(self.fio)
 
-        read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
+        if make_cache:
-            coords, &self.prec, &_bOK)
+            self.make_cache()
 
-        self._cache[<int>self.cur_step] = gmx_fio_ftell(self.fio)
+    def __len__(self):
-        self.timestep = time - self.start_time
+        return len(self.cache)
 
-
+    def __getitem__(self, frame):
-    def __getitem__(self, item):
         cdef matrix box
         cdef gmx_bool _bOK
         cdef real time
         cdef np.ndarray[real, ndim=2] coords = np.empty((self.natoms, 3), dtype=np_real)
-
+        if frame < len(self):
-        self.seek(item)
+            self.seek(frame)
-
+            read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
-        read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
+                <rvec *>coords.data, &self.prec, &_bOK)
-            <rvec *>coords.data, &self.prec, &_bOK)
+            if _bOK:
-
+                return coords
-        return coords
+            else:
+                raise
+        else:
+            raise IndexError('Frame {} is out of range for trajectory of length {}.'.format(frame, len(self)))

setup.py

@@ -26,7 +26,7 @@ elif os.path.exists('gromacs/build/lib'):
 if not library_dirs:
     raise OSError("""
         Gromacs library not found.
-        Activate a gromacs module or specify environment variable LD_LIBRARY_PATH.
+        Activate a gromacs module or set environment variable LD_LIBRARY_PATH.
         """)
 
 library_dirs.append('gromacs/src/external/tng_io/build/lib')
@@ -53,14 +53,22 @@ extensions = [
               library_dirs=library_dirs,
               runtime_library_dirs=library_dirs,
               language='c++'),
-    Extension('pygmx.tngio',
+    Extension('pygmx.enxio',
-              sources=['pygmx/tngio.pyx'],
+              sources=['pygmx/enxio.pyx'],
               include_dirs=include_dirs,
               libraries=['gromacs'],
               library_dirs=library_dirs,
               runtime_library_dirs=library_dirs,
+              language='c++'),
+
+#    Extension('pygmx.tngio',
+#              sources=['pygmx/tngio.pyx'],
+#              include_dirs=include_dirs,
+#              libraries=['gromacs'],
+#              library_dirs=library_dirs,
+#              runtime_library_dirs=library_dirs,
               #language='c++'
-              ),
+#              ),
 
 ]