107 lines
3.0 KiB
Cython
107 lines
3.0 KiB
Cython
|
|
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)
|
|
|
|
void free_enxframe(t_enxframe *ef)
|
|
|
|
|
|
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((self.etypes[i].name.decode(), self.etypes[i].unit.decode()))
|
|
return types
|
|
|
|
def read(self):
|
|
cdef:
|
|
t_enxframe *frame
|
|
snew(frame, 1)
|
|
energies = array('f')
|
|
times = array('f')
|
|
|
|
i = 0
|
|
while do_enx(self.efile, frame):
|
|
times.append(frame.t)
|
|
for j in range(frame.nre):
|
|
energies.append(frame.ener[j].e)
|
|
|
|
free_enxframe(frame)
|
|
sfree(frame)
|
|
|
|
return np.array(times), np.array(energies).reshape(len(times), self.n_etypes)
|
|
|
|
|
|
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)
|
|
|
|
#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
|