Skip to content

Commit

Permalink
Update to chemfiles 0.10
Browse files Browse the repository at this point in the history
  • Loading branch information
@Luthaf
Luthaf committed Apr 8, 2021
1 parent fe26097 commit eacff4e
Showing 1 changed file with 59 additions and 37 deletions.
96 changes: 59 additions & 37 deletions package/MDAnalysis/coordinates/chemfiles.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@
Reading trajectories with *chemfiles* --- :mod:`MDAnalysis.coordinates.chemfiles`
=================================================================================
MDAnalysis interoperates with the `chemfiles`_ library. The *chemfiles* C++ library
MDAnalysis interoperates with the `chemfiles`_ library. The *chemfiles* C++ library
supports an expanding set of file formats, some of which are not natively supported by
MDAnalysis. Using the *CHEMFILES* reader you can use `chemfiles`_ for the low-level
MDAnalysis. Using the *CHEMFILES* reader you can use `chemfiles`_ for the low-level
file reading. Check the list of `chemfile-supported file formats <formats>`_.
.. _chemfiles: https://chemfiles.org
Expand All @@ -37,7 +37,7 @@
Using the CHEMFILES reader
--------------------------
When reading, set the ``format="CHEMFILES"`` keyword argument and I/O is delegated to
When reading, set the ``format="CHEMFILES"`` keyword argument and I/O is delegated to
`chemfiles`_. For example::
>>> import MDAnalysis as mda
Expand Down Expand Up @@ -82,7 +82,7 @@
from distutils.version import LooseVersion
import warnings

from . import base, core
from . import base

try:
import chemfiles
Expand All @@ -94,6 +94,7 @@

class MockTrajectory:
pass

chemfiles = types.ModuleType("chemfiles")
chemfiles.Trajectory = MockTrajectory
else:
Expand All @@ -105,14 +106,14 @@ class MockTrajectory:
MIN_CHEMFILES_VERSION = LooseVersion("0.9")
#: Lowest version of chemfiles that is *not supported*
#: by MDAnalysis.
MAX_CHEMFILES_VERSION = LooseVersion("0.10")
MAX_CHEMFILES_VERSION = LooseVersion("0.11")


def check_chemfiles_version():
"""Check if an appropriate *chemfiles* is available
Returns ``True`` if a usable chemfiles version is available,
with :data:`MIN_CHEMFILES_VERSION` <= version <
with :data:`MIN_CHEMFILES_VERSION` <= version <
:data:`MAX_CHEMFILES_VERSION`
.. versionadded:: 1.0.0
Expand All @@ -127,8 +128,9 @@ def check_chemfiles_version():
wrong = version < MIN_CHEMFILES_VERSION or version >= MAX_CHEMFILES_VERSION
if wrong:
warnings.warn(
"unsupported Chemfiles version {}, we need a version >{} and <{}"
.format(version, MIN_CHEMFILES_VERSION, MAX_CHEMFILES_VERSION)
"unsupported Chemfiles version {}, we need a version >{} and <{}".format(
version, MIN_CHEMFILES_VERSION, MAX_CHEMFILES_VERSION
)
)
return not wrong

Expand All @@ -142,8 +144,9 @@ class ChemfilesReader(base.ReaderBase):
.. versionadded:: 1.0.0
"""
format = 'chemfiles'
units = {'time': 'fs', 'length': 'Angstrom'}

format = "chemfiles"
units = {"time": "fs", "length": "Angstrom"}

def __init__(self, filename, chemfiles_format="", **kwargs):
"""
Expand All @@ -161,8 +164,9 @@ def __init__(self, filename, chemfiles_format="", **kwargs):
"""
if not check_chemfiles_version():
raise RuntimeError("Please install Chemfiles > {}"
"".format(MIN_CHEMFILES_VERSION))
raise RuntimeError(
"Please install Chemfiles > {}" "".format(MIN_CHEMFILES_VERSION)
)
super(ChemfilesReader, self).__init__(filename, **kwargs)
self._format = chemfiles_format
self._cached_n_atoms = None
Expand All @@ -184,7 +188,7 @@ def _open(self):
if isinstance(self.filename, chemfiles.Trajectory):
self._file = self.filename
else:
self._file = ChemfilesPicklable(self.filename, 'r', self._format)
self._file = ChemfilesPicklable(self.filename, "r", self._format)

def close(self):
"""close reader"""
Expand Down Expand Up @@ -217,7 +221,7 @@ def _read_frame(self, i):
def _read_next_timestep(self, ts=None):
"""copy next frame into timestep"""
if self._step >= self.n_frames:
raise IOError('trying to go over trajectory limit')
raise IOError("trying to go over trajectory limit")
if ts is None:
ts = self.ts
self.ts = ts
Expand Down Expand Up @@ -262,17 +266,26 @@ class ChemfilesWriter(base.WriterBase):
.. versionadded:: 1.0.0
"""
format = 'chemfiles'

format = "chemfiles"
multiframe = True

# chemfiles mostly[1] uses these units for the in-memory representation,
# and convert into the file format units when writing.
#
# [1] mostly since some format don't have a specified unit
# (XYZ for example), so then chemfiles just assume they are in A and fs.
units = {'time': 'fs', 'length': 'Angstrom'}

def __init__(self, filename, n_atoms=0, mode="w", chemfiles_format="", topology=None, **kwargs):
units = {"time": "fs", "length": "Angstrom"}

def __init__(
self,
filename,
n_atoms=0,
mode="w",
chemfiles_format="",
topology=None,
**kwargs
):
"""
Parameters
----------
Expand All @@ -298,8 +311,9 @@ def __init__(self, filename, n_atoms=0, mode="w", chemfiles_format="", topology=
"""
if not check_chemfiles_version():
raise RuntimeError("Please install Chemfiles > {}"
"".format(MIN_CHEMFILES_VERSION))
raise RuntimeError(
"Please install Chemfiles > {}" "".format(MIN_CHEMFILES_VERSION)
)
self.filename = filename
self.n_atoms = n_atoms
if mode != "a" and mode != "w":
Expand Down Expand Up @@ -343,15 +357,15 @@ def _write_next_frame(self, obj):
Use AtomGroup or Universe as an input instead.
"""
if hasattr(obj, "atoms"):
if hasattr(obj, 'universe'):
if hasattr(obj, "universe"):
# For AtomGroup and children (Residue, ResidueGroup, Segment)
ts_full = obj.universe.trajectory.ts
if ts_full.n_atoms == obj.atoms.n_atoms:
ts = ts_full
else:
# Only populate a time step with the selected atoms.
ts = ts_full.copy_slice(atoms.indices)
elif hasattr(obj, 'trajectory'):
ts = ts_full.copy_slice(obj.atoms.indices)
elif hasattr(obj, "trajectory"):
# For Universe only --- get everything
ts = obj.trajectory.ts
else:
Expand All @@ -375,7 +389,17 @@ def _timestep_to_chemfiles(self, ts):
if ts.has_velocities:
frame.add_velocities()
frame.velocities[:] = ts.velocities[:]
frame.cell = chemfiles.UnitCell(*ts.dimensions)

lengths = ts.dimensions[:3]
angles = ts.dimensions[3:]
if angles[0] == 0.0 and angles[1] == 0.0 and angles[2] == 0.0:
angles = [90, 90, 90]

if chemfiles.__version__.startswith("0.9"):
frame.cell = chemfiles.UnitCell(*lengths, *angles)
else:
frame.cell = chemfiles.UnitCell(lengths, angles)

return frame

def _topology_to_chemfiles(self, obj, n_atoms):
Expand All @@ -391,21 +415,21 @@ def _topology_to_chemfiles(self, obj, n_atoms):
# (1) add all atoms to the topology
residues = {}
for atom in obj.atoms:
name = getattr(atom, 'name', "")
type = getattr(atom, 'type', name)
name = getattr(atom, "name", "")
type = getattr(atom, "type", name)
chemfiles_atom = chemfiles.Atom(name, type)

if hasattr(atom, 'altLoc'):
if hasattr(atom, "altLoc"):
chemfiles_atom["altloc"] = str(atom.altLoc)

if hasattr(atom, 'segid'):
if hasattr(atom, "segid"):
chemfiles_atom["segid"] = str(atom.segid)

if hasattr(atom, 'segindex'):
if hasattr(atom, "segindex"):
chemfiles_atom["segindex"] = int(atom.segindex)

if hasattr(atom, 'resid'):
resname = getattr(atom, 'resname', "")
if hasattr(atom, "resid"):
resname = getattr(atom, "resname", "")
if atom.resid not in residues.keys():
residues[atom.resid] = chemfiles.Residue(resname, atom.resid)
residue = residues[atom.resid]
Expand Down Expand Up @@ -486,13 +510,11 @@ class ChemfilesPicklable(chemfiles.Trajectory):
.. versionadded:: 2.0.0
"""

def __init__(self, path, mode="r", format=""):
if mode != 'r':
raise ValueError("Only read mode ('r') "
"files can be pickled.")
super().__init__(path=path,
mode=mode,
format=format)
if mode != "r":
raise ValueError("Only read mode ('r') " "files can be pickled.")
super().__init__(path=path, mode=mode, format=format)

def __getstate__(self):
return self.path, self._Trajectory__mode, self._Trajectory__format
Expand Down

0 comments on commit eacff4e

Please sign in to comment.