Select Git revision
frame.py 8.68 KiB
"""
This module contains classes for storing atomic data.
The Frame class may contain multiple Residues which may each contain multiple Atoms.
Both Frame and Residue are iterable. Residue is indexable with either atom numbers or names.
"""
import logging
import numpy as np
from .util import backup_file
from .parsers.cfg import CFG
logger = logging.getLogger(__name__)
np.seterr(all="raise")
# Create FileNotFoundError if using older version of Python
try:
try:
raise FileNotFoundError
except FileNotFoundError:
pass
except NameError:
class FileNotFoundError(OSError):
pass
class Atom:
"""
Hold data for a single atom
"""
__slots__ = ["name", "num", "type", "mass", "charge", "coords"]
def __init__(self, name, num, type=None, mass=None, charge=None, coords=None):
"""
Create an atom.
:param str name: The name of the atom
:param int num: The atom number
:param str type: The atom type
:param float mass: The mass of the atom
:param float charge: The charge of the atom
:param coords: The coordinates of the atom
"""
self.name = name
self.num = num
self.type = type
self.mass = mass
self.charge = charge
self.coords = coords
def __repr__(self):
return "Atom #{0} {1} type: {2} mass: {3} charge: {4}".format(
self.num, self.name, self.type, self.mass, self.charge
)
def add_missing_data(self, other):
assert self.name == other.name
assert self.num == other.num
for attr in ("type", "mass", "charge", "coords"):
if getattr(self, attr) is None:
setattr(self, attr, getattr(other, attr))
class Residue:
"""
Hold data for a residue - list of atoms
"""
__slots__ = ["name", "num", "atoms", "name_to_num"]
def __init__(self, name=None, num=None):
self.atoms = []
self.name = name
self.num = num
self.name_to_num = {}
def __iter__(self):
return iter(self.atoms)
def __getitem__(self, item):
try:
return self.atoms[self.name_to_num[item]]
except KeyError:
pass
try:
return self.atoms[item]
except TypeError as e:
e.args = ("Atom {0} does not exist in residue {1}".format(item, self.name),)
raise
def __len__(self):
return len(self.atoms)
def add_atom(self, atom):
"""
Add an Atom to this Residue and store location in index
:param atom: Atom to add to Residue
:return: None
"""
self.atoms.append(atom)
self.name_to_num[atom.name] = len(self.atoms) - 1
class Frame:
"""
Hold Atom data separated into Residues
"""
def __init__(self, gro=None, xtc=None, itp=None, frame_start=0, xtc_reader=None):
"""
Return Frame instance having read Residues and Atoms from GRO if provided
:param gro: GROMACS GRO file to read initial frame and extract residues
:param xtc: GROMACS XTC file to read subsequent frames
:param itp: GROMACS ITP file to read masses and charges
:return: Frame instance
"""
self.name = ""
self.residues = []
self.number = frame_start - 1
self.time = 0
self.numframes = 0
self.natoms = 0
self.box = np.zeros(3, dtype=np.float32)
self._xtc_buffer = None
if gro is not None:
from .framereader import get_frame_reader
self._trajreader = get_frame_reader(gro, traj=xtc, frame_start=frame_start, name=xtc_reader)
self._trajreader.initialise_frame(self)
if self._trajreader.num_atoms != self.natoms:
raise AssertionError("Number of atoms does not match between gro and xtc files.")
self.numframes += self._trajreader.num_frames
if itp is not None:
self._parse_itp(itp)
@classmethod
def instance_from_reader(cls, reader):
"""
Return Frame instance initialised from existing FrameReader object
:param FrameReader reader: FrameReader object
:return: Frame instance
"""
obj = cls()
obj._trajreader = reader
obj._trajreader.initialise_frame(obj)
return obj
def __len__(self):
return len(self.residues)
def __iter__(self):
return iter(self.residues)
def __getitem__(self, item):
return self.residues[item]
def __repr__(self):
rep = self.name + "\n"
atoms = []
for res in self.residues:
for atom in res:
atoms.append(repr(atom))
rep += "\n".join(atoms)
return rep
def yield_resname_in(self, container):
for res in self:
if res.name in container:
yield res
def next_frame(self):
"""
Read next frame from input XTC.
:return: True if successful else False
"""
result = self._trajreader.read_next(self)
if result:
self.number += 1
return result
def write_xtc(self, filename):
"""
Write frame to output XTC file.
:param filename: XTC filename to write to
"""
if self._xtc_buffer is None:
try:
import mdtraj
except ImportError as e:
if "scipy" in repr(e):
e.msg = "XTC output with MDTraj also requires Scipy"
else:
e.msg = "XTC output requires the module MDTraj (and probably Scipy)"
raise
backup_file(filename)
self._xtc_buffer = mdtraj.formats.XTCTrajectoryFile(filename, mode="w")
xyz = np.ndarray((1, self.natoms, 3), dtype=np.float32)
i = 0
for residue in self.residues:
for atom in residue.atoms:
xyz[0][i] = atom.coords
i += 1
time = np.array([self.time], dtype=np.float32)
step = np.array([self.number], dtype=np.int32)
box = np.zeros((1, 3, 3), dtype=np.float32)
for i in range(3):
box[0][i][i] = self.box[i]
self._xtc_buffer.write(xyz, time=time, step=step, box=box)
def _parse_itp(self, filename):
"""
Parse a GROMACS ITP file to extract atom charges/masses.
Optional but requires that ITP contains only a single residue.
:param filename: Filename of GROMACS ITP to read
"""
with CFG(filename) as itp:
itpres = Residue(itp["moleculetype"][0][0])
for line in itp["atoms"]:
atom = Atom(num=int(line[0]) - 1, type=line[1], name=line[4], charge=float(line[6]), mass=float(line[7]))
itpres.add_atom(atom)
for res in self.residues:
if res.name == itpres.name:
for atom, itpatom in zip(res, itpres):
atom.add_missing_data(itpatom)
def output(self, filename, format="gro"):
"""
Write coordinates from Frame to file.
:param filename: Name of file to write to
:param format: Format to write e.g. 'gro', 'lammps'
"""
outputs = {"gro": self._output_gro,
"lammps": self._output_lammps_data}
try:
outputs[format](filename)
except KeyError:
print("ERROR: Invalid output format {0}, coordinates will not be output.".format(format))
def _output_lammps_data(self, filename):
"""
Output Frame coordinates in LAMMPS data format.
:param filename: Name of DATA file to create
"""
raise NotImplementedError("LAMMPS Data output has not yet been implemented.")
def _output_gro(self, filename):
"""
Create a GROMACS GRO file from the data in this Frame
:param filename: Name of GRO file to create
"""
backup_file(filename)
with open(filename, "w") as gro:
print(self.name, file=gro)
print("{0:5d}".format(self.natoms), file=gro)
i = 1
format_string = "{0:5d}{1:5s}{2:>5s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}"
for res in self.residues:
for atom in res:
print(format_string.format(res.num, res.name,
atom.name, i,
*atom.coords), file=gro)
i += 1
print("{0:10.5f}{1:10.5f}{2:10.5f}".format(*self.box), file=gro)
def add_residue(self, residue):
"""
Add a Residue to this Frame
:param residue: Residue to add
"""
self.residues.append(residue)