192 lines
6.3 KiB
Cython
192 lines
6.3 KiB
Cython
# 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]
|