292 lines
8.5 KiB
Cython
Executable File
292 lines
8.5 KiB
Cython
Executable File
# Wrapper for xtcio.h: I/O for xtc files.
|
|
|
|
from cpython.array cimport array
|
|
import cython
|
|
from array import array
|
|
from xdrlib import Unpacker
|
|
|
|
import numpy as np
|
|
cimport numpy as np
|
|
import os
|
|
|
|
from .utility cimport *
|
|
from .math cimport *
|
|
from .fileio cimport *
|
|
from .gromacs.reader import INDEX_MAGIC, SubscriptableReader, XTCFrame
|
|
from .errors import InvalidIndexException, InvalidMagicException, XTCError
|
|
|
|
|
|
cdef extern from "gromacs/fileio/xtcio.h":
|
|
t_fileio *open_xtc(const char *filename, const char *mode) nogil
|
|
|
|
void close_xtc(t_fileio *fio) nogil
|
|
|
|
int read_first_xtc(t_fileio *fio,
|
|
int *natoms, gmx_int64_t *step, real *time,
|
|
matrix box, rvec **x, real *prec, gmx_bool *_bOK)
|
|
|
|
int read_next_xtc(t_fileio *fio,
|
|
int natoms, gmx_int64_t *step, real *time,
|
|
matrix box, rvec *x, real *prec, gmx_bool *_bOK) nogil
|
|
|
|
int write_xtc(t_fileio *fio, int natoms, gmx_int64_t step, real time, rvec *box, rvec *x, real prec)
|
|
|
|
if sizeof(real) == 4:
|
|
np_real = np.float32
|
|
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, frame, state = 1
|
|
gmx_int64_t step
|
|
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]
|
|
|
|
|
|
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
|
|
int natoms
|
|
gmx_int64_t cur_step
|
|
real start_time, timestep, prec, cur_time
|
|
bint has_cache, has_times
|
|
array _cache, _times
|
|
public str filename
|
|
|
|
@property
|
|
def cache(self):
|
|
if self.has_cache:
|
|
return self._cache
|
|
|
|
@cache.setter
|
|
def cache(self, indices):
|
|
self._cache = array('L')
|
|
for i in indices:
|
|
self._cache.append(i)
|
|
self.has_cache = True
|
|
|
|
@property
|
|
def times(self):
|
|
if self.has_times:
|
|
return self._times
|
|
|
|
@times.setter
|
|
def times(self, times):
|
|
self._times = array('f')
|
|
for t in times:
|
|
self._times.append(t)
|
|
self.has_times = True
|
|
|
|
# @property
|
|
# def filename(self):
|
|
# return self._filename.decode()
|
|
|
|
def seek(self, int frame):
|
|
if self.has_cache:
|
|
gmx_fio_seek(self.fio, self.cache[frame])
|
|
|
|
def make_cache(self):
|
|
self._cache = get_xtc_index(self.fio)
|
|
self.has_cache = True
|
|
|
|
def load_cache(self, indexfile, ignore_time=False):
|
|
|
|
xtc_stat = os.stat(self.filename)
|
|
c_time = int(xtc_stat.st_ctime)
|
|
m_time = int(xtc_stat.st_mtime)
|
|
size = xtc_stat.st_size
|
|
|
|
with open(indexfile, 'rb') as fd:
|
|
unpacker = Unpacker(fd.read())
|
|
# first 4 * 8 bytes are used for checks, followed by N * (8 + 4) bytes of data
|
|
length = int((len(unpacker.get_buffer()) - 32) / 12)
|
|
|
|
if length < 0:
|
|
raise InvalidIndexException
|
|
|
|
if unpacker.unpack_hyper() != INDEX_MAGIC:
|
|
raise InvalidMagicException
|
|
if unpacker.unpack_hyper() != c_time and not ignore_time:
|
|
raise InvalidIndexException
|
|
if unpacker.unpack_hyper() != m_time and not ignore_time:
|
|
raise InvalidIndexException
|
|
if unpacker.unpack_hyper() != size:
|
|
raise InvalidIndexException
|
|
|
|
self._cache = array('L')
|
|
self._times = array('f')
|
|
|
|
try:
|
|
while True:
|
|
self._cache.append(unpacker.unpack_hyper())
|
|
self._times.append(unpacker.unpack_float())
|
|
except EOFError:
|
|
pass
|
|
self.has_cache = True
|
|
self.has_times = True
|
|
|
|
def __cinit__(self):
|
|
self.has_cache = False
|
|
self.has_times = False
|
|
|
|
def __init__(self, filename, indexfile=None, make_cache=False, ignore_timestamps=False):
|
|
if isinstance(filename, str):
|
|
self.filename = filename
|
|
filename = filename.encode()
|
|
else:
|
|
self.filename = filename.decode()
|
|
|
|
cdef:
|
|
gmx_int64_t step
|
|
matrix box
|
|
gmx_bool _bOK
|
|
real time, prec
|
|
rvec *x
|
|
|
|
if not os.path.exists(filename):
|
|
raise OSError('File not found: {}'.format(filename))
|
|
if filename.decode().split('.')[-1] != 'xtc':
|
|
raise XTCError('File is not of xtc type: {}'.format(filename))
|
|
|
|
self.fio = open_xtc(filename, b'r')
|
|
read_first_xtc(self.fio, &self.natoms, &step, &time, box, &x, &prec, &_bOK)
|
|
|
|
if indexfile is not None:
|
|
try:
|
|
self.load_cache(indexfile, ignore_time=ignore_timestamps)
|
|
except InvalidIndexException:
|
|
if make_cache:
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
if make_cache:
|
|
self.make_cache()
|
|
|
|
def __len__(self):
|
|
if self.has_cache:
|
|
return len(self.cache)
|
|
|
|
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)
|
|
if frame < len(self):
|
|
self.seek(frame)
|
|
read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
|
|
<rvec *>coords.data, &self.prec, &_bOK)
|
|
if _bOK:
|
|
frame = XTCFrame()
|
|
frame._coordinates = coords
|
|
frame.index = self.cur_step
|
|
frame.time = time
|
|
frame.box = box
|
|
return frame
|
|
else:
|
|
raise
|
|
else:
|
|
raise IndexError('Frame {} is out of range for trajectory of length {}.'.format(frame, len(self)))
|
|
|
|
def __getstate__(self):
|
|
# state = self.__dict__.copy()
|
|
state = {}
|
|
state['natoms'] = self.natoms
|
|
state['cur_step'] = self.cur_step
|
|
state['start_time'] = self.start_time
|
|
state['timestep'] = self.timestep
|
|
state['prec'] = self.prec
|
|
state['cur_time'] = self.cur_time
|
|
state['has_cache'] = self.has_cache
|
|
state['has_times'] = self.has_times
|
|
state['_cache'] = self._cache
|
|
state['filename'] = self.filename
|
|
return state
|
|
|
|
def __setstate__(self, state):
|
|
self.natoms = state.pop('natoms')
|
|
self.cur_step = state.pop('cur_step')
|
|
self.start_time = state.pop('start_time')
|
|
self.timestep = state.pop('timestep')
|
|
self.prec = state.pop('prec')
|
|
self.cur_time = state.pop('cur_time')
|
|
self.has_cache = state.pop('has_cache')
|
|
self.has_times = state.pop('has_times')
|
|
self._cache = state.pop('_cache')
|
|
self.filename = state.pop('filename')
|
|
self.fio = open_xtc(self.filename.encode(), b'r')
|
|
|
|
#self.__dict__.update(state)
|
|
|
|
|
|
@cython.binding(True)
|
|
def append_xtcfile(filename, step, time, box, coords, prec):
|
|
if isinstance(filename, str):
|
|
filename = filename.encode()
|
|
|
|
cdef np.ndarray[real, ndim=2] b = np.asarray(box, dtype=np.float32)
|
|
cdef np.ndarray[real, ndim=2] x = np.asarray(coords, dtype=np.float32)
|
|
|
|
cdef t_fileio *fio = open_xtc(filename, b'a')
|
|
write_xtc(fio, len(coords), step, time, <rvec *>b.data, <rvec *>x.data, <real> prec)
|
|
close_xtc(fio)
|
|
|
|
|
|
|
|
|