Enabled & functioning edrfile reader.

pygmx/__init__.py

@@ -9,14 +9,15 @@ import os
 
 from .tpxio import TPXReader
 from .xtcio import XTCReader
-# from .enxio import EDRFile
+from .enxio import EDRFile
 from .errors import FileTypeError
 from .gromacs.reader import index_filename_for_xtc
 from .gromacs.xtcindex import index_xtcfile
 
 FILE_EXTENSIONS = {
     'xtc': XTCReader,
-    'tpr': TPXReader
+    'tpr': TPXReader,
+    'edr': EDRFile,
 }
 
 

pygmx/enxio.pyx

@@ -40,6 +40,7 @@ cdef extern from "gromacs/fileio/enxio.h":
 
     gmx_bool do_enx(ener_file_t ef, t_enxframe *fr)
 
+    void free_enxframe(t_enxframe *ef)
 
 
 cdef class EDRFile:
@@ -52,22 +53,45 @@ cdef class EDRFile:
 
     @property
     def types(self):
+
         types = []
         for i in range(self.n_etypes):
-            types.append(
+            types.append((self.etypes[i].name.decode(), self.etypes[i].unit.decode()))
-
+        return types
-            )
-    def read(self):
-        cdef t_enxframe *frame
-        cdef np.ndarray[float, ndim=2] energies = np.empty((1,self.n_etypes), np.float32)
 
+    def read(self, steps=None):
+        cdef:
+            t_enxframe *frame
+            np.ndarray[float, ndim=2] energies
+            np.ndarray[float, ndim=1] times
+        if steps is None:
+            energies = np.empty((1,self.n_etypes), np.float32)
+            times = np.empty((1,), np.float32)
+        else:
+            energies = np.empty((steps,self.n_etypes), np.float32)
+            times = np.empty((steps,), np.float32)
         snew(frame, 1)
 
-        while do_enx(self.efile, frame):
+        i = 0
+        while True:
+            cont = do_enx(self.efile, frame)
+
+            for j in range(frame.nre):
+                energies[i, j] = frame.ener[j].e
+            times[i] = frame.t
+
+            if cont:
+                i += 1
+                if len(energies) <= i:
                     energies = np.resize(energies, (len(energies)+1, self.n_etypes))
-            for i in range(frame.nre):
+                    times = np.resize(times, (len(times)+1,))
-                energies[-1, i] = frame.ener[i].e
+            else:
-        return energies
+                break
+
+        free_enxframe(frame)
+        sfree(frame)
+
+        return times, energies
 
 
     def __init__(self, filename):

pygmx/tpxio.pyx

@@ -183,6 +183,30 @@ cdef class TPXReader:
                 types += mol_type * nmol
             return np.array(types)
 
+    @property
+    def nsteps(self):
+        return self.input_record.nsteps
+
+    @property
+    def nstxout(self):
+        return self.input_record.nstxout
+
+    @property
+    def nstvout(self):
+        return self.input_record.nstvout
+
+    @property
+    def nstfout(self):
+        return self.input_record.nstfout
+
+    @property
+    def nstxout_compressed(self):
+        return self.input_record.nstxout_compressed
+        
+    @property
+    def nstenergy(self):
+        return self.input_record.nstenergy
+
 
     def __cinit__(self, filename):
         filename = cstr(filename)

pygmx/utility.pxd

@@ -15,6 +15,7 @@ cdef extern from "gromacs/utility/futil.h":
 
 cdef extern from "gromacs/utility/smalloc.h":
     void snew(void *ptr, int nelem)
+    void sfree(void *ptr)
 
 
 cdef inline cstr(instr):

setup.py

@@ -55,13 +55,12 @@ extensions = [
         library_dirs=library_dirs,
         language='c++'
     ),
-#    Extension('pygmx.enxio',
+    Extension('pygmx.enxio',
-#              sources=['pygmx/enxio.pyx'],
+              sources=['pygmx/enxio.pyx'],
-#              include_dirs=include_dirs,
+              include_dirs=include_dirs,
-#              libraries=['gromacs'],
+              libraries=['gromacs'],
-#              library_dirs=library_dirs,
+              library_dirs=library_dirs,
-#              runtime_library_dirs=library_dirs,
+              language='c++'),
-#              language='c++'),
 
 #    Extension('pygmx.tngio',
 #              sources=['pygmx/tngio.pyx'],