diff --git a/pygmx/__init__.py b/pygmx/__init__.py index 58572b0..63a204c 100644 --- a/pygmx/__init__.py +++ b/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, } diff --git a/pygmx/enxio.pyx b/pygmx/enxio.pyx index 1fc6fd9..9514b5f 100644 --- a/pygmx/enxio.pyx +++ b/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( - - ) - def read(self): - cdef t_enxframe *frame - cdef np.ndarray[float, ndim=2] energies = np.empty((1,self.n_etypes), np.float32) + types.append((self.etypes[i].name.decode(), self.etypes[i].unit.decode())) + return types + 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): - 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 + 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)) + times = np.resize(times, (len(times)+1,)) + else: + break + + free_enxframe(frame) + sfree(frame) + + return times, energies def __init__(self, filename): diff --git a/pygmx/tpxio.pyx b/pygmx/tpxio.pyx index e1e8038..4427ee8 100644 --- a/pygmx/tpxio.pyx +++ b/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) diff --git a/pygmx/utility.pxd b/pygmx/utility.pxd index 05a45d6..be0b8ce 100644 --- a/pygmx/utility.pxd +++ b/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): diff --git a/setup.py b/setup.py index 26370b9..195cdf8 100644 --- a/setup.py +++ b/setup.py @@ -55,13 +55,12 @@ extensions = [ library_dirs=library_dirs, language='c++' ), -# Extension('pygmx.enxio', -# sources=['pygmx/enxio.pyx'], -# include_dirs=include_dirs, -# libraries=['gromacs'], -# library_dirs=library_dirs, -# runtime_library_dirs=library_dirs, -# language='c++'), + Extension('pygmx.enxio', + sources=['pygmx/enxio.pyx'], + include_dirs=include_dirs, + libraries=['gromacs'], + library_dirs=library_dirs, + language='c++'), # Extension('pygmx.tngio', # sources=['pygmx/tngio.pyx'],