diff --git a/.gitmodules b/.gitmodules index b4d06cb..bdc4ed2 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ [submodule "gromacs"] path = gromacs - url = https://github.com/gromacs/gromacs.git + url = git://git.gromacs.org/gromacs.git diff --git a/pygmx/__init__.py b/pygmx/__init__.py index 1240731..1d522ce 100644 --- a/pygmx/__init__.py +++ b/pygmx/__init__.py @@ -7,3 +7,4 @@ from . import xtcio from . import tngio from .tpxio import TPXReader +from .enxio import EDRFile diff --git a/pygmx/enxio.pyx b/pygmx/enxio.pyx new file mode 100644 index 0000000..1fc6fd9 --- /dev/null +++ b/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 diff --git a/pygmx/fileio.pxd b/pygmx/fileio.pxd index 7188a15..afb0bc0 100644 --- a/pygmx/fileio.pxd +++ b/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) diff --git a/pygmx/mdtypes.pxd b/pygmx/mdtypes.pxd index e8b47dc..f78969d 100644 --- a/pygmx/mdtypes.pxd +++ b/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 diff --git a/pygmx/utility.pxd b/pygmx/utility.pxd index 243ff0b..d1f7fa3 100644 --- a/pygmx/utility.pxd +++ b/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 diff --git a/pygmx/xtcio.pyx b/pygmx/xtcio.pyx index ce681d9..d0d8a5e 100644 --- a/pygmx/xtcio.pyx +++ b/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 - #rvec *coords + bint has_cache + rvec *coords + array _cache - def nearest_cache(self, index): - nearest = 0 - cur_dist = 99999999 - for ind in self._cache: - dist = abs(ind - index) - if dist < cur_dist: - nearest = ind - cur_dist = dist - return nearest + @property + def cache(self): + return self._cache + + def seek(self, int frame): + if self.has_cache: + gmx_fio_seek(self.fio, self.cache[frame]) - def seek(self, index): - if index in self._cache: - gmx_fio_seek(self.fio, self._cache[index]) - 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 make_cache(self): + self._cache = get_xtc_index(self.fio) + self.has_cache = True + 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 - cdef rvec *coords # = np.empty((1,3), dtype=np.float32) + real time, prec + 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, - &coords, &self.prec, &_bOK) - self._cache[self.cur_step] = gmx_fio_ftell(self.fio) + read_first_xtc(self.fio, &self.natoms, &step, &time, box, &x, &prec, &_bOK) - read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box, - coords, &self.prec, &_bOK) + if make_cache: + self.make_cache() - self._cache[self.cur_step] = gmx_fio_ftell(self.fio) - self.timestep = time - self.start_time + def __len__(self): + return len(self.cache) - - def __getitem__(self, item): + def __getitem__(self, frame): 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) - - self.seek(item) - - read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box, - coords.data, &self.prec, &_bOK) - - return coords + if frame < len(self): + self.seek(frame) + read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box, + coords.data, &self.prec, &_bOK) + if _bOK: + return coords + else: + raise + else: + raise IndexError('Frame {} is out of range for trajectory of length {}.'.format(frame, len(self))) diff --git a/setup.py b/setup.py index 3b051a9..4f34b32 100644 --- a/setup.py +++ b/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', - sources=['pygmx/tngio.pyx'], + 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.tngio', +# sources=['pygmx/tngio.pyx'], +# include_dirs=include_dirs, +# libraries=['gromacs'], +# library_dirs=library_dirs, +# runtime_library_dirs=library_dirs, #language='c++' - ), +# ), ]