IPKM/python-pygmx
3
0
Fork 0
Code Issues Pull Requests Activity

First draft of a python wrapper for gromacs library

Browse Source
This commit is contained in:
Niels Müller
2016-03-16 13:07:02 +01:00
parent b49941595b
commit 58d5d3a6c4
12 changed files with 18621 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
pygmx/__init__.py Normal file
View File

@ -0,0 +1,9 @@
"""
Python wrapper for gromacs tools.
"""
from . import tpxio
from . import xtcio
from . import tngio
from .tpxio import TPXReader

12
pygmx/fileio.pxd Normal file
View File

@ -0,0 +1,12 @@
from utility cimport *
cdef extern from "gromacs/fileio/gmxfio.h":
ctypedef struct t_fileio:
pass
int gmx_fio_seek(t_fileio *fio, gmx_off_t fpos)
gmx_off_t gmx_fio_ftell(t_fileio *fio)
int xtc_seek_time(t_fileio *fio, real time, int natoms, gmx_bool bSeekForwardOnly)

7
pygmx/math.pxd Normal file
View File

@ -0,0 +1,7 @@
from utility cimport real
cdef extern from "gromacs/math/vectypes.h":
ctypedef real rvec[3]
ctypedef real matrix[3][3]
ctypedef real tensor[3][3]

175
pygmx/mdtypes.pxd Normal file
View File

@ -0,0 +1,175 @@
from utility cimport *
from math cimport *
cdef extern from "gromacs/mdtypes/inputrec.h":
ctypedef struct t_simtemp:
pass
ctypedef struct t_lambda:
pass
ctypedef struct t_expanded:
pass
ctypedef struct t_rot:
pass
ctypedef struct t_IMD:
pass
ctypedef struct t_grpopts:
int ngtc
real *ref_t # Coupling temperature per group */
ctypedef struct t_cosines:
pass
ctypedef struct t_swapcoords:
pass
ctypedef struct t_inputrec:
int eI # Integration method */
int nsteps # number of steps to be taken */
int simulation_part # Used in checkpointing to separate chunks */
int init_step # start at a stepcount >0 (used w. convert-tpr) */
int nstcalcenergy # frequency of energy calc. and T/P coupl. upd. */
int cutoff_scheme # group or verlet cutoffs */
int ns_type # which ns method should we use? */
int nstlist # number of steps before pairlist is generated */
int ndelta # number of cells per rlong */
int nstcomm # number of steps after which center of mass */
# motion is removed */
int comm_mode # Center of mass motion removal algorithm */
int nstlog # number of steps after which print to logfile */
int nstxout # number of steps after which X is output */
int nstvout # id. for V */
int nstfout # id. for F */
int nstenergy # number of steps after which energies printed */
int nstxout_compressed # id. for compressed trj (.xtc,.tng) */
double init_t # initial time (ps) */
double delta_t # time step (ps) */
double x_compression_precision # precision of x in compressed trajectory file */
double fourier_spacing # requested fourier_spacing, when nk? not set */
int nkx, nky, nkz # number of k vectors in each spatial dimension*/
# for fourier methods for long range electrost.*/
int pme_order # interpolation order for PME */
double ewald_rtol # double space tolerance for Ewald, determines */
# the double/reciprocal space relative weight */
double ewald_rtol_lj # double space tolerance for LJ-Ewald */
int ewald_geometry # normal/3d ewald, or pseudo-2d LR corrections */
double epsilon_surface # Epsilon for PME dipole correction */
int ljpme_combination_rule # Type of combination rule in LJ-PME */
int ePBC # Type of periodic boundary conditions */
int bPeriodicMols # Periodic molecules */
bint bContinuation # Continuation run: starting state is correct */
int etc # temperature coupling */
int nsttcouple # interval in steps for temperature coupling */
bint bPrintNHChains # whether to print nose-hoover chains */
int epc # pressure coupling */
int epct # pressure coupling type */
int nstpcouple # interval in steps for pressure coupling */
double tau_p # pressure coupling time (ps) */
tensor ref_p # reference pressure (kJ/(mol nm^3)) */
tensor compress # compressability ((mol nm^3)/kJ) */
int refcoord_scaling # How to scale absolute reference coordinates */
rvec posres_com # The COM of the posres atoms */
rvec posres_comB # The B-state COM of the posres atoms */
int andersen_seed # Random seed for Andersen thermostat (obsolete) */
double verletbuf_tol # Per atom pair energy drift tolerance (kJ/mol/ps/atom) for list buffer */
double rlist # short range pairlist cut-off (nm) */
double rtpi # Radius for test particle insertion */
int coulombtype # Type of electrostatics treatment */
int coulomb_modifier # Modify the Coulomb interaction */
double rcoulomb_switch # Coulomb switch range start (nm) */
double rcoulomb # Coulomb cutoff (nm) */
double epsilon_r # relative dielectric constant */
double epsilon_rf # relative dielectric constant of the RF */
int implicit_solvent # No (=explicit water), or GBSA solvent models */
int gb_algorithm # Algorithm to use for calculation Born radii */
int nstgbradii # Frequency of updating Generalized Born radii */
double rgbradii # Cutoff for GB radii calculation */
double gb_saltconc # Salt concentration (M) for GBSA models */
double gb_epsilon_solvent # dielectric coeff. of implicit solvent */
double gb_obc_alpha # 1st scaling factor for Bashford-Case GB */
double gb_obc_beta # 2nd scaling factor for Bashford-Case GB */
double gb_obc_gamma # 3rd scaling factor for Bashford-Case GB */
double gb_dielectric_offset # Dielectric offset for Still/HCT/OBC */
int sa_algorithm # Algorithm for SA part of GBSA */
double sa_surface_tension # Energy factor for SA part of GBSA */
int vdwtype # Type of Van der Waals treatment */
int vdw_modifier # Modify the VdW interaction */
double rvdw_switch # Van der Waals switch range start (nm) */
double rvdw # Van der Waals cutoff (nm) */
int eDispCorr # Perform Long range dispersion corrections */
double tabext # Extension of the table beyond the cut-off, *
# * as well as the table length for 1-4 interac. */
double shake_tol # tolerance for shake */
int efep # free energy calculations */
t_lambda *fepvals # Data for the FEP state */
bint bSimTemp # Whether to do simulated tempering */
t_simtemp *simtempvals # Variables for simulated tempering */
bint bExpanded # Whether expanded ensembles are used */
t_expanded *expandedvals # Expanded ensemble parameters */
int eDisre # Type of distance restraining */
double dr_fc # force constant for ta_disre */
int eDisreWeighting # type of weighting of pairs in one restraints */
bint bDisreMixed # Use comb of time averaged and instan. viol's */
int nstdisreout # frequency of writing pair distances to enx */
double dr_tau # time constant for memory function in disres */
double orires_fc # force constant for orientational restraints */
double orires_tau # time constant for memory function in orires */
int nstorireout # frequency of writing tr(SD) to enx */
double em_stepsize # The stepsize for updating */
double em_tol # The tolerance */
int niter # Number of iterations for convergence of */
# steepest descent in relax_shells */
double fc_stepsize # Stepsize for directional minimization */
# in relax_shells */
int nstcgsteep # number of steps after which a steepest */
# descents step is done while doing cg */
int nbfgscorr # Number of corrections to the hessian to keep */
int eConstrAlg # Type of constraint algorithm */
int nProjOrder # Order of the LINCS Projection Algorithm */
double LincsWarnAngle # If bond rotates more than %g degrees, warn */
int nLincsIter # Number of iterations in the final Lincs step */
bint bShakeSOR # Use successive overrelaxation for shake */
double bd_fric # Friction coefficient for BD (amu/ps) */
int ld_seed # Random seed for SD and BD */
int nwall # The number of walls */
int wall_type # The type of walls */
double wall_r_linpot # The potentail is linear for r<=wall_r_linpot */
int wall_atomtype[2] # The atom type for walls */
double wall_density[2] # Number density for walls */
double wall_ewald_zfac # Scaling factor for the box for Ewald */
# COM pulling data */
bint bPull # Do we do COM pulling? */
#struct pull_params_t *pull # The data for center of mass pulling */
#struct pull_t *pull_work # The COM pull force calculation data structure TODO this pointer should live somewhere else */
# Enforced rotation data */
bint bRot # Calculate enforced rotation potential(s)? */
t_rot *rot # The data for enforced rotation potentials */
int eSwapCoords # Do ion/water position exchanges (CompEL)? */
t_swapcoords *swap
bint bIMD # Allow interactive MD sessions for this .tpr? */
t_IMD *imd # Interactive molecular dynamics */
double cos_accel # Acceleration for viscosity calculation */
tensor deform # Triclinic deformation velocities (nm/ps) */
int userint1 # User determined parameters */
int userint2
int userint3
int userint4
double userdouble1
double userdouble2
double userdouble3
double userdouble4
t_grpopts opts # Group options */
t_cosines ex[0] # Electric field stuff (spatial part) */
t_cosines et[0] # Electric field stuff (time part) */
bint bQMMM # QM/MM calculation */
int QMconstraints # constraints on QM bonds */
int QMMMscheme # Scheme: ONIOM or normal */
double scalefactor # factor for scaling the MM charges in QM calc.*/
# Fields for removed features go here (better caching) */
bint bAdress # Whether AdResS is enabled - always false if a valid .tpr was read
bint useTwinRange # Whether twin-range scheme is active - always false if a valid .tpr was read

495
pygmx/reader.py Normal file
View File

@ -0,0 +1,495 @@
import xdrlib
import io
import os
# from .coordinates import decompress
from ..utils import hash_anything, merge_hashes
from functools import partial
import numpy as np
TRR_MAGIC = 1993
XTC_MAGIC = 1995
INDEX_MAGIC = 2015
def index_filename_for_xtc(xtcfile):
"""
Get the filename for the index file of a xtc file.
In general, the index file is located in the same directory as the xtc file.
If the xtc file is located in a read-only directory (for the current user)
and does not exist, the index file will be instead located in a subdirectory
of ~/.xtcindex, of the current user.
The directory for user index files can be changed by setting the environment
variable 'XTCINDEXDIR'.
Example:
xtcfile = '/data/niels/test.xtc'
# case 1: index file exists, or current user is niels
index_filename_for_xtc(xtcfile) == '/data/niels/.test.xtc.xtcindex'
# case 2: index file doesn't exist, and nor write permission in /data/niels
index_filename_for_xtc(xtcfile) == '~/.xtcindex/data/niels/.test.xtc.xtcindex'
Warning:
At this point, the index file is not checked for validity. If an invalid
index file is present in the xtc files directory, this will be used and an
error will be risen by XTCReader.
Warning:
On most systems, the home directory is on the local drive so that the indexing
has probably to do be done on every system if it can not be saved in the directory
of the xtc file.
"""
base = os.path.basename(xtcfile)
filename = "." + base + ".xtcindex"
dir = os.path.abspath(os.path.dirname(xtcfile))
if not os.path.exists(os.path.join(dir, filename)) and not os.access(dir, os.W_OK):
if 'XTCINDEXDIR' in os.environ:
index_dir = os.environ['XTCINDEXDIR']
else:
index_dir = os.path.join(os.environ['HOME'], '.xtcindex')
dir = os.path.join(index_dir, dir.lstrip('/'))
os.makedirs(dir, exist_ok=True)
return os.path.join(dir, filename)
class NumpyUnpacker(xdrlib.Unpacker):
def unpack_float_array(self, n_items):
i = self.get_position()
j = i + 4 * n_items
self.set_position(j)
data = self.get_buffer()[i:j]
if len(data) < 4:
raise EOFError
ret = np.frombuffer(data, '>f')
if len(ret) != n_items:
raise EOFError
return ret
def unpack_double_array(self, n_items):
i = self.get_position()
j = i + 8 * n_items
self.set_position(j)
data = self.get_buffer()[i:j]
if len(data) < 8:
raise EOFError
ret = np.frombuffer(data, '>d')
if len(ret) != n_items:
raise EOFError
return ret
class InvalidMagicException(Exception):
pass
class InvalidIndexException(Exception):
pass
class UnknownLenError(Exception):
pass
class SubscriptableReader:
def __init__(self, fd):
self.fd = fd
self.len = os.stat(self.fd.fileno()).st_size
def __getitem__(self, r):
if isinstance(r, slice):
if r.start is not None:
if r.start < 0:
self.fd.seek(r.start, io.SEEK_END)
else:
self.fd.seek(r.start, io.SEEK_SET)
if r.step is not None:
raise NotImplementedError
return self.fd.read(r.stop - r.start)
else:
self.fd.seek(r, io.SEEK_SET)
return self.fd.read(1)
def __len__(self):
return self.len
class XTCFrame:
__slots__ = ['_coordinates', 'index', 'time', 'box', '_compressed_coordinates']
def __init__(self):
self._coordinates = None
@property
def coordinates(self):
if self._coordinates is None:
self._coordinates = self._compressed_coordinates()
return self._coordinates
class BaseReader:
def __init__(self, file):
self.fd = open(file, 'rb')
self.filename = file
self.reader = SubscriptableReader(self.fd)
self.xdr_reader = NumpyUnpacker(self.reader)
def __exit__(self, *_):
self.fd.close()
def __enter__(self, *_):
return self
def get_position(self):
return self.fd.tell()
def set_position(self, pos):
return self.fd.seek(pos)
def skip_frames(self, num):
for i in range(num):
self.skip_frame()
def skip_frame(self):
self._read_header()
self._read_frame()
def dump_frame(self):
raise NotImplemented
def _read_header(self):
raise NotImplemented
def _read_frame(self):
raise NotImplemented
XTC_HEADER_SIZE = 4 * 4
XTC_FRAME_HEADER_SIZE = 9 * 4 + 10 * 4
class XTCReader(BaseReader):
len_available = False
def __init__(self, xtcfile, load_cache_file=True):
super().__init__(xtcfile)
self._cache = [0]
self._times = None
self.current_frame = 0
index_file_name = index_filename_for_xtc(xtcfile)
if load_cache_file:
self.load_index(index_file_name)
def load_index(self, filename):
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(filename, 'rb') as index_file_fd:
# TODO: Is NumpyUnpacker even necessary at this point?
# Seems like xdrlib.Unpacker would be sufficient here ...
reader = NumpyUnpacker(SubscriptableReader(index_file_fd))
if reader.unpack_hyper() != INDEX_MAGIC:
raise InvalidMagicException
if reader.unpack_hyper() != c_time:
raise InvalidIndexException
if reader.unpack_hyper() != m_time:
raise InvalidIndexException
if reader.unpack_hyper() != size:
raise InvalidIndexException
self._cache = []
self._times = []
try:
while True:
self._cache.append(reader.unpack_hyper())
self._times.append(reader.unpack_float())
except EOFError:
pass
self.len_available = True
def _raw_header(self):
if len(self._cache) == self.current_frame:
self._cache.append(self.get_position())
return self.fd.read(XTC_HEADER_SIZE)
def _raw_frame(self):
frame_header = self.fd.read(XTC_FRAME_HEADER_SIZE)
blob_size = xdrlib.Unpacker(frame_header[-4:]).unpack_int()
blob_size = (blob_size + 3) // 4 * 4 # Padding to 4 bytes
frame_blob = self.fd.read(blob_size)
self.current_frame += 1
return frame_header + frame_blob
def _unpack_header(self, raw_header):
unpacker = xdrlib.Unpacker(raw_header)
magic = unpacker.unpack_int()
if magic != XTC_MAGIC:
raise InvalidMagicException
n_atoms = unpacker.unpack_int()
step = unpacker.unpack_int()
time = unpacker.unpack_float()
return n_atoms, step, time
def _read_header(self):
raw_header = self._raw_header()
return self._unpack_header(raw_header)
def _unpack_frame(self, raw_frame):
unpacker = xdrlib.Unpacker(raw_frame)
raw_box = unpacker.unpack_farray(9, unpacker.unpack_float)
box = np.array(raw_box).reshape(3, 3)
num_coords = unpacker.unpack_int()
precision = unpacker.unpack_float()
maxint = unpacker.unpack_farray(3, unpacker.unpack_int)
minint = unpacker.unpack_farray(3, unpacker.unpack_int)
smallindex = unpacker.unpack_int()
blob_len = unpacker.unpack_int()
blob = unpacker.unpack_fopaque(blob_len)
return box, precision, num_coords, minint, maxint, smallindex, blob
def _read_frame(self):
raw_frame = self._raw_frame()
return self._unpack_frame(raw_frame)
def decode_frame(self, header=None, body=None):
n_atoms, step, time = header or self._read_header()
box, precision, num_coords, minint, maxint, smallindex, blob = body or self._read_frame()
coordinates = partial(decompress, num_coords, precision, minint, maxint, smallindex, blob)
frame = XTCFrame()
frame._compressed_coordinates = coordinates
frame.index = step
frame.time = time
frame.box = box
return frame
def dump_frame(self):
"""
:return: Tuple: The binary data of the frame, the frame itself
"""
raw_header = self._raw_header()
header = self._unpack_header(raw_header)
raw_body = self._raw_frame()
body = self._unpack_frame(raw_body)
return (raw_header + raw_body, self.decode_frame(header, body))
def cached_position(self, item):
try:
return self._cache[item]
except IndexError:
return None
def __getitem__(self, item):
position = self.cached_position(item)
if position is not None:
self.set_position(position)
self.current_frame = item
return self.decode_frame()
# TODO: Use elif and one single return at the end
# Also: skip_frames is in fact not implemented at all!
# TODO: Catch EOFError and raise IndexError
if self.current_frame <= item:
self.skip_frames(item - self.current_frame)
return self.decode_frame()
else:
self.set_position(0)
self.current_frame = 0
self.skip_frames(item)
return self.decode_frame()
def __len__(self):
if self.len_available:
return len(self._cache)
raise UnknownLenError
def times_of_indices(self, indices):
return [self[i].time for i in indices]
def __hash__(self):
return merge_hashes(hash(self.filename), hash_anything(self._cache))
class TRRHeader:
__slots__ = ['ir_size', 'e_size', 'box_size', 'vir_size', 'pres_size', 'top_size', 'sym_size', 'x_size', 'v_size',
'f_size', 'n_atoms', 'step', 'nre', 't', 'λ', 'is_double']
@property
def frame_size(self):
return self.ir_size + \
self.e_size + \
self.box_size + \
self.vir_size + \
self.top_size + \
self.sym_size + \
self.x_size + \
self.v_size + \
self.f_size
class TRRFrame:
__slots__ = ['header', 'box', 'x', 'v', 'f']
class TRRReader(BaseReader):
def __iter__(self):
return TRRIterator(self.filename)
def _unpack_float(self, is_double):
if is_double:
return self.xdr_reader.unpack_double()
else:
return self.xdr_reader.unpack_float()
def _unpack_np_float(self, is_double, *dim):
total = np.product(dim)
if is_double:
box = self.xdr_reader.unpack_double_array(total)
else:
box = self.xdr_reader.unpack_float_array(total)
box = box.reshape(*dim)
return box
def _read_header(self):
data = self.fd.read(8) # 4 Magic + 4 VersionStringLen
unpacker = NumpyUnpacker(data)
magic = unpacker.unpack_int()
if magic != TRR_MAGIC:
raise InvalidMagicException
version_string_len = unpacker.unpack_int()
data = self.fd.read(version_string_len + 55) # 4 Magic + 4 VersionStringLen
unpacker.reset(data)
version = unpacker.unpack_string()
assert version == b'GMX_trn_file'
header = TRRHeader()
header.ir_size = unpacker.unpack_int()
header.e_size = unpacker.unpack_int()
header.box_size = unpacker.unpack_int()
header.vir_size = unpacker.unpack_int()
header.pres_size = unpacker.unpack_int()
header.top_size = unpacker.unpack_int()
header.sym_size = unpacker.unpack_int()
header.x_size = unpacker.unpack_int()
header.v_size = unpacker.unpack_int()
header.f_size = unpacker.unpack_int()
header.n_atoms = unpacker.unpack_int()
header.step = unpacker.unpack_int()
header.nre = unpacker.unpack_int()
if header.x_size is not None:
header.is_double = (header.x_size // header.n_atoms) == 8
elif header.v_size is not None:
header.is_double = (header.v_size // header.n_atoms) == 8
elif header.f_size is not None:
header.is_double = (header.f_size // header.n_atoms) == 8
if header.is_double:
data = self.fd.read(16)
else:
data = self.fd.read(8)
unpacker.reset(data)
self.xdr_reader = unpacker
header.t = self._unpack_float(header.is_double)
header.λ = self._unpack_float(header.is_double)
return header
def _read_frame(self, header):
frame = TRRFrame()
frame.header = header
data = self.fd.read(header.frame_size)
self.xdr_reader = NumpyUnpacker(data)
frame.box = None
if header.box_size:
frame.box = self._unpack_np_float(header.is_double, 3, 3)
if header.vir_size:
for i in range(9):
self._unpack_float(header.is_double)
if header.pres_size:
for i in range(9):
self._unpack_float(header.is_double)
frame.x = None
if header.x_size:
frame.x = self._unpack_np_float(header.is_double, header.n_atoms, 3)
frame.v = None
if header.v_size:
frame.v = self._unpack_np_float(header.is_double, header.n_atoms, 3)
frame.f = None
if header.f_size:
frame.f = self._unpack_np_float(header.is_double, header.n_atoms, 3)
return frame
def decode_frame(self):
h = self._read_header()
return self._read_frame(h)
def get_position(self):
return self.fd.tell()
def set_position(self, pos):
self.fd.seek(pos)
class TRRIterator():
def __init__(self, file):
self.__reader = TRRReader(file)
self.__reader.__enter__()
def __next__(self):
try:
return self.__reader.decode_frame()
except EOFError as err:
self.__reader.__exit__()
raise StopIteration
def __iter__(self):
return self

8838
pygmx/tngio.c Normal file
View File

File diff suppressed because it is too large Load Diff

148
pygmx/tngio.pyx Normal file
View File

@ -0,0 +1,148 @@
import numpy as np
cimport numpy as np
#from libc.stdlib cimport free
cdef extern from "tng/tng_io_fwd.h":
cdef struct tng_trajectory:
pass
ctypedef tng_trajectory *tng_trajectory_t
cdef extern from "tng/tng_io.h":
ctypedef int int64_t
cdef int64_t TNG_TRAJ_POSITIONS "TNG_TRAJ_POSITIONS"
cdef int64_t TNG_TRAJ_VELOCITIES "TNG_TRAJ_VELOCITIES"
cdef int64_t TNG_TRAJ_FORCES "TNG_TRAJ_FORCES"
ctypedef enum tng_function_status:
TNG_SUCCESS, TNG_FAILURE, TNG_CRITICAL
tng_function_status tng_util_trajectory_open(const char *filename, const char mode, tng_trajectory_t *tng_data_p)
tng_function_status tng_util_trajectory_close(tng_trajectory_t *tng_data_p)
tng_function_status tng_util_time_of_frame_get(const tng_trajectory_t tng_data, const int64_t frame_nr, double *time)
tng_function_status tng_util_pos_read_range(const tng_trajectory_t tng_data,
const int64_t first_frame,
const int64_t last_frame,
float **positions,
int64_t *stride_length)
tng_function_status tng_util_vel_read_range(const tng_trajectory_t tng_data,
const int64_t first_frame,
const int64_t last_frame,
float **velocities,
int64_t *stride_length)
tng_function_status tng_util_force_read_range(const tng_trajectory_t tng_data,
const int64_t first_frame,
const int64_t last_frame,
float **forces,
int64_t *stride_length)
tng_function_status tng_util_box_shape_read_range(const tng_trajectory_t tng_data,
const int64_t first_frame,
const int64_t last_frame,
float **box_shape,
int64_t *stride_length)
tng_function_status tng_num_frames_get(const tng_trajectory_t tng_data, int64_t *n)
tng_function_status tng_num_particles_get(const tng_trajectory_t tng_data, int64_t *n);
tng_function_status tng_medium_stride_length_get(const tng_trajectory_t tng_data, int64_t *len);
tng_function_status tng_long_stride_length_get(const tng_trajectory_t tng_data, int64_t *len);
tng_function_status tng_util_trajectory_next_frame_present_data_blocks_find(
const tng_trajectory_t tng_data,
int64_t current_frame,
const int64_t n_requested_data_block_ids,
const int64_t *requested_data_block_ids,
int64_t *next_frame,
int64_t *n_data_blocks_in_next_frame,
int64_t **data_block_ids_in_next_frame)
tng_function_status tng_util_particle_data_next_frame_read(
const tng_trajectory_t tng_data,
const int64_t block_id,
void **values,
char *data_type,
int64_t *retrieved_frame_number,
double *retrieved_time)
cdef check_status(status, message=''):
if status != TNG_SUCCESS:
print(message)
#cdef ptr_to_array(float *ptr, np.npy_intp *shape, int np):
# arr = np.PyArray_SimpleNewFromData(np, shape, np.NPY_FLOAT, ptr)
# return arr
cdef class TNGTrajectory:
cdef:
tng_trajectory_t trajectory
int nframes, natoms, stride
def get_position_interval(self):
cdef:
double time
char data_type
int frame1, frame2
cdef np.ndarray[float, ndim=2] positions = np.empty((self.natoms, 3), dtype=np.float32)
cdef void *data = <void *>positions.data
tng_util_particle_data_next_frame_read(self.trajectory, TNG_TRAJ_POSITIONS, &data, &data_type, &frame1, &time)
tng_util_particle_data_next_frame_read(self.trajectory, TNG_TRAJ_POSITIONS, &data, &data_type, &frame2, &time)
del positions
return frame2 - frame1
def __len__(self):
return self.nframes
def __getitem__(self, frame):
cdef np.ndarray[float, ndim=2] positions = np.empty((self.natoms, 3), dtype=np.float32)
cdef float *pos = <float *>positions.data
#check_status(
if (TNG_SUCCESS !=
tng_util_pos_read_range(self.trajectory, frame, frame, &pos, &self.stride)#,
#'Reading of frame not successful.'
):
raise IndexError()
return positions
def __init__(self, filename):
assert filename, 'Filename must not be null.'
if isinstance(filename, str):
filename = filename.encode()
if (tng_util_trajectory_open(filename, b'r', &self.trajectory) != TNG_SUCCESS):
raise OSError('Trajectory could not be opened.')
tng_num_frames_get(self.trajectory, &self.nframes)
tng_num_particles_get(self.trajectory, &self.natoms)
from math import ceil
class TNGReader:
def __init__(self, filename):
self.trajectory = TNGTrajectory(filename)
self.nstxout = self.trajectory.get_position_interval()
def __len__(self):
return ceil(len(self.trajectory) / self.nstxout)
def __getitem__(self, item):
try:
return self.trajectory[item * self.nstxout]
except IndexError:
raise IndexError('Trajectory frame {} out of range.'.format(item))

107
pygmx/topology.pxd Normal file
View File

@ -0,0 +1,107 @@
from utility cimport *
from math cimport *
from mdtypes cimport *
cdef extern from "gromacs/topology/atoms.h":
ctypedef struct t_atom:
real m, q; # Mass and charge */
real mB, qB; # Mass and charge for Free Energy calc */
unsigned short type; # Atom type */
unsigned short typeB; # Atom type for Free Energy calc */
int ptype; # Particle type */
int resind; # Index into resinfo (in t_atoms) */
int atomnumber; # Atomic Number or 0 */
char elem[4]; # Element name */
ctypedef struct t_resinfo:
char **name; # Pointer to the residue name */
int nr; # Residue number */
unsigned char ic; # Code for insertion of residues */
int chainnum; # Iincremented at TER or new chain id */
char chainid; # Chain identifier written/read to pdb */
char **rtp; # rtp building block name (optional) */
ctypedef struct t_pdbinfo:
pass
ctypedef struct t_atoms:
int nr # Nr of atoms */
t_atom *atom # Array of atoms (dim: nr) */
# The following entries will not */
# always be used (nres==0) */
char ***atomname # Array of pointers to atom name */
# use: (*(atomname[i])) */
char ***atomtype # Array of pointers to atom types */
# use: (*(atomtype[i])) */
char ***atomtypeB # Array of pointers to B atom types */
# use: (*(atomtypeB[i])) */
int nres # The number of resinfo entries */
t_resinfo *resinfo # Array of residue names and numbers */
t_pdbinfo *pdbinfo # PDB Information, such as aniso. Bfac */
ctypedef struct t_atomtypes:
pass
#ctypedef t_atoms *t_atoms_ptr
cdef extern from "gromacs/topology/symtab.h":
ctypedef struct t_symtab:
pass
cdef extern from "gromacs/topology/block.h":
ctypedef struct t_block:
pass
ctypedef struct t_blocka:
pass
cdef extern from "gromacs/topology/idef.h":
ctypedef struct t_ilist:
pass
ctypedef struct gmx_ffparams_t:
pass
cdef extern from "gromacs/topology/topology.h":
ctypedef struct gmx_moltype_t:
char **name; # Name of the molecule type */
t_atoms atoms; # The atoms in this molecule */
t_ilist ilist[0]; # Interaction list with local indices */
t_block cgs; # The charge groups */
t_blocka excls; # The exclusions */
ctypedef struct gmx_molblock_t:
int type; # The molcule type index in mtop.moltype */
int nmol; # The number of molecules in this block */
int natoms_mol; # The number of atoms in one molecule */
int nposres_xA; # The number of posres coords for top A */
rvec *posres_xA; # The posres coords for top A */
int nposres_xB; # The number of posres coords for top B */
rvec *posres_xB; # The posres coords for top B */
ctypedef struct gmx_groups_t:
pass
ctypedef struct gmx_mtop_t:
char **name # Name of the topology */
gmx_ffparams_t ffparams
int nmoltype
gmx_moltype_t *moltype
int nmolblock
gmx_molblock_t *molblock
bint bIntermolecularInteractions # Are there intermolecular
# interactions? */
t_ilist *intermolecular_ilist # List of intermolecular interactions
# using system wide atom indices,
# either NULL or size F_NRE */
int natoms
int maxres_renum # Parameter for residue numbering */
int maxresnr # The maximum residue number in moltype */
t_atomtypes atomtypes # Atomtype properties */
t_block mols # The molecules */
gmx_groups_t groups
t_symtab symtab # The symbol table */
# cdef extern from "gromacs/topology/topology.h":
# generate a t_atoms struct for the system from gmx_mtop_t
# t_atoms* mtop2atoms(gmx_mtop_t *mtop)

8530
pygmx/tpxio.c Normal file
View File

File diff suppressed because it is too large Load Diff

191
pygmx/tpxio.pyx Normal file
View File

@ -0,0 +1,191 @@
# cython: c_string_type=unicode, c_string_encoding=utf8
# Cython wrapper around tpxio.cpp
# from cython cimport view
# from array import array
import numpy as np
cimport numpy as np
from utility cimport *
from math cimport *
from mdtypes cimport *
from topology cimport *
cdef extern from "gromacs/fileio/tpxio.h":
void read_tpxheader(const char *fn, t_tpxheader *tpx, bint TopOnlyOK)
cdef struct t_tpxheader:
bint bIr # Non zero if input_rec is present */
bint bBox # Non zero if a box is present */
bint bTop # Non zero if a topology is present */
bint bX # Non zero if coordinates are present */
bint bV # Non zero if velocities are present */
bint bF # Non zero if forces are present (no longer
# supported, but retained so old .tpr can be read) */
int natoms # The total number of atoms */
int ngtc # The number of temperature coupling groups */
real _lambda # Current value of lambda */
int fep_state # Current value of the alchemical state --
int read_tpx(const char *fname,
t_inputrec *ir,
double box[3][3],
int *natoms,
rvec *x,
rvec *v,
gmx_mtop_t *mtop)
def cstr(instr):
if isinstance(instr, str):
instr = instr.encode()
return instr
cdef atoms_from_topology(gmx_mtop_t *topology):
cdef:
t_atoms c_atoms
int moltype, resind
atoms = []
masses = []
charges = []
residues = []
for i_molblock in range(topology.nmolblock):
moltype = topology.molblock[i_molblock].type
c_atoms = topology.moltype[moltype].atoms
mol_atoms = []
mol_q = []
mol_m = []
for i_atom in range(c_atoms.nr):
resind = c_atoms.atom[i_atom].resind
resname = c_atoms.resinfo[resind].name[0]
if resname not in residues:
residues.append(resname)
resid = residues.index(resname) + 1
mol_atoms.append((
resid,
resname,
c_atoms.atomname[0][i_atom],
))
mol_q.append(c_atoms.atom[i_atom].q)
mol_m.append(c_atoms.atom[i_atom].m)
n_mol = topology.molblock[i_molblock].nmol
atoms += mol_atoms * n_mol
charges += mol_q * n_mol
masses += mol_m * n_mol
return np.array(atoms)
ctypedef object (*atom_func)(t_atom)
cdef atom_charge(t_atom atom):
return atom.q
cdef atom_mass(t_atom atom):
return atom.m
cdef class TPXReader:
cdef:
t_tpxheader header
t_inputrec input_record
gmx_mtop_t topology
real box[3][3]
readonly int n_atoms, n_tcouple_groups, n_mol
readonly char *topology_name
cdef _map_atoms(self, object (*func)(t_atom)):
cdef t_atoms atoms
map = []
for i_molblock in range(self.topology.nmolblock):
moltype = self.topology.molblock[i_molblock].type
nmol = self.topology.molblock[i_molblock].nmol
atoms = self.topology.moltype[moltype].atoms
mol_map = []
for i_atom in range(atoms.nr):
mol_map.append(func(atoms.atom[i_atom]))
map += mol_map * nmol
return map
property atoms:
"""
Get a list of tuples describing all atoms in the system:
Returns:
List of tuples: (atom_name, residue_id, resiude_name)
"""
def __get__(self):
return atoms_from_topology(&self.topology)
property mass2:
def __get__(self):
return np.array(self._map_atoms(atom_mass))
property mass:
"Get the masses of all atoms in the system."
def __get__(self):
cdef t_atoms atoms
masses = []
for i_molblock in range(self.topology.nmolblock):
moltype = self.topology.molblock[i_molblock].type
nmol = self.topology.molblock[i_molblock].nmol
atoms = self.topology.moltype[moltype].atoms
mol_masses = []
for i_atom in range(atoms.nr):
mol_masses.append(atoms.atom[i_atom].m)
masses += mol_masses * nmol
return np.array(masses)
property charge:
"Get the partial charge of all atoms in the system."
def __get__(self):
cdef t_atoms atoms
charges = []
for i_molblock in range(self.topology.nmolblock):
moltype = self.topology.molblock[i_molblock].type
nmol = self.topology.molblock[i_molblock].nmol
atoms = self.topology.moltype[moltype].atoms
mol_charges = []
for i_atom in range(atoms.nr):
mol_charges.append(atoms.atom[i_atom].q)
charges += mol_charges * nmol
return np.array(charges)
property type:
"Get the type of all atoms in the system."
def __get__(self):
cdef t_atoms atoms
types = []
for i_molblock in range(self.topology.nmolblock):
moltype = self.topology.molblock[i_molblock].type
nmol = self.topology.molblock[i_molblock].nmol
atoms = self.topology.moltype[moltype].atoms
mol_type = []
for i_atom in range(atoms.nr):
mol_type.append(atoms.atom[i_atom].type)
types += mol_type * nmol
return np.array(types)
def __cinit__(self, filename):
filename = cstr(filename)
read_tpxheader(<char *>filename, &self.header, True)
self.n_atoms = self.header.natoms
self.n_tcouple_groups = self.header.ngtc
cdef np.ndarray[real, ndim=2] coordinates = np.empty((self.n_atoms, 3), dtype=np.float32)
cdef np.ndarray[real, ndim=2] velocites = np.empty((self.n_atoms, 3), dtype=np.float32)
return_code = read_tpx(
<char *>filename,
&self.input_record,
self.box,
&self.n_atoms,
<rvec *>coordinates.data,
<rvec *>velocites.data,
&self.topology
)
self.topology_name = self.topology.name[0]

14
pygmx/utility.pxd Normal file
View File

@ -0,0 +1,14 @@
# C-API in gromacs/utility
#cdef extern from "inttypes.h":
ctypedef int __int64
cdef extern from "gromacs/utility/basedefinitions.h":
ctypedef int gmx_bool
ctypedef __int64 gmx_int64_t
cdef extern from "gromacs/utility/real.h":
ctypedef double real
cdef extern from "gromacs/utility/futil.h":
ctypedef gmx_int64_t gmx_off_t

95
pygmx/xtcio.pyx Normal file
View File

@ -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