First draft of a python wrapper for gromacs library

pygmx/xtcio.pyx

@@ -0,0 +1,95 @@
+# 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