96 lines
2.7 KiB
Cython
96 lines
2.7 KiB
Cython
# Wrapper for xtcio.h: I/O for xtc files.
|
|
|
|
import numpy as np
|
|
cimport numpy as np
|
|
|
|
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)
|
|
|
|
void close_xtc(t_fileio *fio)
|
|
|
|
int read_first_xtc(t_fileio *fio,
|
|
int *natoms, int *step, real *time,
|
|
matrix box, rvec **x, real *prec, gmx_bool *_bOK)
|
|
|
|
int read_next_xtc(t_fileio *fio,
|
|
int natoms, int *step, real *time,
|
|
matrix box, rvec *x, real *prec, gmx_bool *_bOK)
|
|
|
|
|
|
if sizeof(real) == 4:
|
|
np_real = np.float32
|
|
else:
|
|
np_real = np.float
|
|
|
|
|
|
cdef class XTCReader:
|
|
cdef:
|
|
t_fileio *fio
|
|
int natoms, cur_step
|
|
real start_time, timestep, prec, cur_time
|
|
dict _cache
|
|
#rvec *coords
|
|
|
|
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
|
|
|
|
|
|
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 __init__(self, filename):
|
|
if isinstance(filename, str):
|
|
filename = filename.encode()
|
|
|
|
cdef:
|
|
matrix box
|
|
gmx_bool _bOK
|
|
real time
|
|
cdef rvec *coords # = np.empty((1,3), dtype=np.float32)
|
|
|
|
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_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
|
|
coords, &self.prec, &_bOK)
|
|
|
|
self._cache[<int>self.cur_step] = gmx_fio_ftell(self.fio)
|
|
self.timestep = time - self.start_time
|
|
|
|
|
|
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)
|
|
|
|
self.seek(item)
|
|
|
|
read_next_xtc(self.fio, self.natoms, &self.cur_step, &time, box,
|
|
<rvec *>coords.data, &self.prec, &_bOK)
|
|
|
|
return coords
|