100 lines
2.9 KiB
Cython
100 lines
2.9 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)
|
|
|
|
|
|
|
|
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
|