277 lines
9.0 KiB
Cython
Executable File
277 lines
9.0 KiB
Cython
Executable File
# cython: c_string_type=unicode, c_string_encoding=utf8
|
|
# Cython wrapper around tpxio.cpp
|
|
|
|
from libc cimport stdio
|
|
from libc.string cimport memcpy
|
|
|
|
import numpy as np
|
|
import cython
|
|
cimport numpy as np
|
|
|
|
from .utility cimport *
|
|
from .math cimport *
|
|
from .mdtypes cimport *
|
|
from .topology cimport *
|
|
|
|
|
|
cdef extern from "gromacs/fileio/tpxio.h":
|
|
ctypedef struct t_tpxheader:
|
|
int bIr # Non zero if input_rec is present */
|
|
int bBox # Non zero if a box is present */
|
|
int bTop # Non zero if a topology is present */
|
|
int bX # Non zero if coordinates are present */
|
|
int bV # Non zero if velocities are present */
|
|
int 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 --
|
|
|
|
|
|
void read_tpxheader(const char *fn, t_tpxheader *tpx, bint TopOnlyOK, int *version, int *generation)
|
|
int read_tpx(const char *fname,
|
|
t_inputrec *ir,
|
|
matrix box,
|
|
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 = []
|
|
residues = []
|
|
res_id = 0
|
|
for i_molblock in range(topology.nmolblock):
|
|
moltype = topology.molblock[i_molblock].type
|
|
c_atoms = topology.moltype[moltype].atoms
|
|
for n in range(topology.molblock[i_molblock].nmol):
|
|
mol_atoms = []
|
|
res_id += 1
|
|
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((
|
|
res_id,
|
|
resname,
|
|
c_atoms.atomname[i_atom][0],
|
|
))
|
|
atoms += mol_atoms
|
|
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 index_groups_from_topology(gmx_mtop_t *topology):
|
|
# retrieve the index groups from the topology->groups ?
|
|
pass
|
|
|
|
|
|
cdef open_tpx(const char* filename, t_inputrec *ir, matrix box, int *natoms, gmx_mtop_t *top):
|
|
|
|
# unterdrücke GMX output durch internen stderr buffer...
|
|
cdef char buffer[stdio.BUFSIZ]
|
|
stdio.setbuf(stdio.stderr, buffer)
|
|
cdef rvec *x, *v
|
|
return_code = read_tpx(filename, ir, box, natoms, x, v, top)
|
|
|
|
for i in range(stdio.BUFSIZ):
|
|
buffer[i] = 0
|
|
stdio.fflush(stdio.stderr)
|
|
stdio.fseek(stdio.stderr, 0, stdio.SEEK_END)
|
|
stdio.setbuf(stdio.stderr, NULL)
|
|
return return_code
|
|
|
|
|
|
cdef read_ffparams(gmx_ffparams_t *ffparams, gmx_bool bShowNumbers):
|
|
cdef char buffer[stdio.BUFSIZ]
|
|
stdio.setbuf(stdio.stdout, buffer)
|
|
pr_ffparams(stdio.stdout, 0, '', ffparams, bShowNumbers)
|
|
stdio.fflush(stdio.stderr)
|
|
stdio.fseek(stdio.stderr, 0, stdio.SEEK_END)
|
|
stdio.setbuf(stdio.stderr, NULL)
|
|
return buffer
|
|
|
|
|
|
cdef class TPXReader:
|
|
cdef:
|
|
t_tpxheader header
|
|
t_inputrec input_record
|
|
gmx_mtop_t topology
|
|
gmx_localtop_t *local_top
|
|
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)
|
|
|
|
property molecules:
|
|
"Get molecule indices from topology."
|
|
def __get__(self):
|
|
mols = [0] * self.n_atoms
|
|
molid = 0
|
|
for i in range(self.topology.mols.nr):
|
|
molid += 1
|
|
for j in range(self.topology.mols.index[i], self.topology.mols.index[i+1]):
|
|
mols[j] = molid
|
|
return mols
|
|
# @property
|
|
# def nsteps(self):
|
|
# return self.input_record.nsteps
|
|
|
|
# @property
|
|
# def nstxout(self):
|
|
# return self.input_record.nstxout
|
|
|
|
# @property
|
|
# def nstvout(self):
|
|
# return self.input_record.nstvout
|
|
|
|
# @property
|
|
# def nstfout(self):
|
|
# return self.input_record.nstfout
|
|
|
|
# @property
|
|
# def nstxout_compressed(self):
|
|
# return self.input_record.nstxout_compressed
|
|
|
|
# @property
|
|
# def nstenergy(self):
|
|
# return self.input_record.nstenergy
|
|
|
|
|
|
def read_ff(self):
|
|
return read_ffparams(&self.topology.ffparams, True)
|
|
|
|
def __cinit__(self, filename):
|
|
filename = cstr(filename)
|
|
|
|
# Arrays for coordinates etc can be NULL, if they shouldn't be read
|
|
# cdef int version, generation
|
|
# read_tpxheader(<char *>filename, &self.header, True, &version, &generation)
|
|
# 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)
|
|
# cdef np.ndarray[real, ndim=2] forces = np.empty((self.n_atoms, 3), dtype=np.float32)
|
|
open_tpx(
|
|
<char *>filename,
|
|
NULL,
|
|
self.box,
|
|
&self.n_atoms,
|
|
&self.topology
|
|
)
|
|
self.topology_name = self.topology.name[0]
|
|
self.local_top = gmx_mtop_generate_local_top(&self.topology, True)
|
|
|
|
|
|
@cython.binding(True)
|
|
def make_xtcframe_whole(coords, box, TPXReader reader):
|
|
|
|
cdef int natoms = reader.topology.natoms
|
|
cdef gmx_rmpbc_t gpbc = gmx_rmpbc_init(&reader.local_top.idef, -1, natoms)
|
|
|
|
cdef np.ndarray[real, ndim=2] b = np.asarray(box, dtype=np.float32)
|
|
cdef np.ndarray[real, ndim=2] x = np.array(coords, dtype=np.float32).copy()
|
|
|
|
gmx_rmpbc(gpbc, natoms, <rvec *>b.data, <rvec *>x.data)
|
|
|
|
# free up memory - in fact, we should free memory of 'top' too
|
|
# however, it shares memory with 'reader.topology', so we can not free it correctly
|
|
gmx_rmpbc_done(gpbc)
|
|
return x
|
|
|
|
|