Add initial support for loading trajectories as coordinates & box vec…

…tors providers
insilichem · Jun 12, 2018 · 38e95e1 · 38e95e1
1 parent 82a9e08
commit 38e95e1
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 24 deletions.
diff --git a/docs/index.rst b/docs/index.rst
@@ -19,9 +19,9 @@ A command line application to launch molecular dynamics simulations with OpenMM
 + No coding required - just a YAML input file!
 + Smart support for different input file formats:
     + **Topology**: PDB/PDBx, Mol2, Amber's PRMTOP, Charmm's PSF, Gromacs' TOP, Desmond's DMS
-    + **Positions**: PDB, COOR, INPCRD, CRD, GRO
+    + **Positions**: PDB, COOR, INPCRD, CRD, GRO, DCD
     + **Velocities**: PDB, VEL
-    + **Box vectors**: PDB, XSC, CSV, INPCRD, GRO
+    + **Box vectors**: XSC, CSV, PDB, GRO, INPCRD, DCD
     + A fallback method is implemented and will attempt to load everything else that might be supported by `ParmEd <http://parmed.github.io/ParmEd/html/index.html>`_.
 + Choose your preferred **trajectory** format (PDB, PDBx, DCD, HDF5, NETCDF, MDCRD) and **checkpoints** (Amber's, CHARMM's, OpenMM XML states).
 + Live report of simulation progress, with estimated ETA and speed.

diff --git a/docs/input.rst b/docs/input.rst
@@ -94,16 +94,16 @@ This is possible due to the chosen default values for almost every key. Specific
 Top-level parameters
 --------------------
 
-*All the parametes are optional except stated otherwise.*
+*All the parameters are optional except stated otherwise.*
 
 Input options
 .............
 
-- ``topology``: Main input file. Should contain, at least, the topology, but it can also contain positions, velocities, PBC vectors, forcefields... Required.
-- ``positions``: File with the initial coordinates of the system. Overrides those in topology, if needed. Required if the topology does not provide positions.
-- ``velocities``: File containing the initial velocities of this stage. If not set, they will be set to the requested temperature.
-- ``box_vectors``: File with replacement periodic box vectors, instead of those in the topology or positions file. Supported file types: XSC.
-- ``checkpoint``: Restart simulation from this file. It can provide one or more of the options above.
+- ``topology``: Main input file. Should contain, at least, the topology, but it can also contain positions, velocities, PBC vectors, forcefields... Required. Supports PDB/PDBx, Mol2, Amber's PRMTOP, Charmm's PSF, Gromacs' TOP, Desmond's DMS.
+- ``positions``: File with the initial coordinates of the system. Overrides those in topology, if needed. Required if the topology does not provide positions. If the file is a trajectory, a frame must be specified with a list: ``[ path_to_trajectory.dcd, 1044]``. Supports PDB, COOR, INPCRD, CRD, GRO, DCD.
+- ``velocities``: File containing the initial velocities of this stage. If not set, they will be set to the requested temperature. Supports PDB, VEL.
+- ``box_vectors``: File with replacement periodic box vectors, instead of those in the topology or positions file. If the file is a trajectory, a frame must be specified with a list: ``[ path_to_trajectory.dcd, 1044]``. Supports XSC, CSV, PDB, GRO, INPCRD, DCD.
+- ``checkpoint``: Restart simulation from this file. It can provide one or more of the options above. Supports STATE.XML, RS.
 - ``forcefield``: Which forcefields should be used, if not provided in topology. Required for PDB topologies. More details on :ref:`forcefields`.
 - ``charmm_parameters``: CHARMM forcefield. Required for PSF topologies.
 

diff --git a/ommprotocol/io.py b/ommprotocol/io.py
@@ -49,7 +49,7 @@
 from parmed.openmm import RestartReporter, NetCDFReporter, MdcrdReporter
 from openmoltools.utils import create_ffxml_file
 # Own
-from .utils import sanitize_path_for_file, ignored_exceptions, warned_getattr, extant_file
+from .utils import sanitize_args_for_file, ignored_exceptions, warned_getattr, extant_file
 from ._version import get_versions
 __version__ = get_versions()['version']
 del get_versions
@@ -460,7 +460,8 @@ def _loaders(cls, ext):
                 'coor': cls.from_namd,
                 'gro': cls.from_gromacs,
                 'inpcrd': cls.from_amber,
-                'crd': cls.from_charmm}[ext]
+                'crd': cls.from_charmm,
+                'dcd': cls.from_mdtraj}[ext]
 
     @classmethod
     def from_pdb(cls, path):
@@ -484,6 +485,10 @@ def from_gromacs(cls, path):
     def from_charmm(cls, path):
         return CharmmCrdFile(path).positions
 
+    @classmethod
+    def from_mdtraj(cls, path, frame):
+        return mdtraj.load_frame(path, frame).xyz[0]
+
     @classmethod
     def from_parmed(cls, path):
         return parmed.load_file(path, structure=True).positions
@@ -540,7 +545,8 @@ def _loaders(cls, ext):
                 'csv': cls.from_csv,
                 'pdb': cls.from_pdb,
                 'gro': cls.from_gromacs,
-                'inpcrd': cls.from_amber}[ext]
+                'inpcrd': cls.from_amber,
+                'dcd': cls.from_mdtraj}[ext]
 
     @classmethod
     def from_xsc(cls, path):
@@ -616,6 +622,10 @@ def from_amber(cls, path):
     def from_gromacs(cls, path):
         return GromacsGroFile(path).getPeriodicBoxVectors()
 
+    @classmethod
+    def from_mdtraj(cls, path, frame):
+        return mdtraj.load_frame(path, frame).unitcell_vectors[0]
+
     @classmethod
     def from_parmed(cls, path):
         return parmed.load_file(path, structure=True, hasbox=True).box
@@ -906,37 +916,37 @@ def prepare_handler(cfg):
     Load all files into single object.
     """
     positions, velocities, box = None, None, None
-    _path = cfg['_path']
+    _path = cfg.get('_path', './')
     forcefield = cfg.pop('forcefield', None)
-    topology = sanitize_path_for_file(cfg.pop('topology'), _path)
+    topology_args = sanitize_args_for_file(cfg.pop('topology'), _path)
 
     if 'checkpoint' in cfg:
-        restart_path = sanitize_path_for_file(cfg.pop('checkpoint'), _path)
-        restart = Restart.load(restart_path)
+        restart_args = sanitize_args_for_file(cfg.pop('checkpoint'), _path)
+        restart = Restart.load(*restart_args)
         positions = restart.positions
         velocities = restart.velocities
         box = restart.box
 
     if 'positions' in cfg:
-        positions_path = sanitize_path_for_file(cfg.pop('positions'), _path)
-        positions = Positions.load(positions_path)
-        box = BoxVectors.load(positions_path)
+        positions_args = sanitize_args_for_file(cfg.pop('positions'), _path)
+        positions = Positions.load(*positions_args)
+        box = BoxVectors.load(*positions_args)
 
     if 'velocities' in cfg:
-        velocities_path = sanitize_path_for_file(cfg.pop('velocities'), _path)
-        velocities = Velocities.load(velocities_path)
+        velocities_args = sanitize_args_for_file(cfg.pop('velocities'), _path)
+        velocities = Velocities.load(*velocities_args)
 
     if 'box' in cfg:
-        box_path = sanitize_path_for_file(cfg.pop('box'), _path)
-        box = BoxVectors.load(box_path)
+        box_args = sanitize_args_for_file(cfg.pop('box'), _path)
+        box = BoxVectors.load(*box_args)
 
     options = {}
     for key in 'positions velocities box forcefield'.split():
         value = locals()[key]
         if value is not None:
             options[key] = value
 
-    return SystemHandler.load(topology, **options)
+    return SystemHandler.load(*topology_args, **options)
 
 
 def prepare_system_options(cfg, defaults=None):

diff --git a/ommprotocol/utils.py b/ommprotocol/utils.py
@@ -27,6 +27,7 @@
 
 
 if sys.version_info.major == 3:
+    basestring = str
     raw_input = input
 
 
@@ -106,6 +107,13 @@ def random_string(length=5):
     return ''.join(random.choice(string.ascii_letters) for _ in range(length))
 
 
+def sanitize_args_for_file(args, config_file):
+    if isinstance(args, basestring):
+        args = (args,)
+    path = sanitize_path_for_file(args[0], config_file)
+    return (path,) + args[1:]
+
+
 def sanitize_path_for_file(path, config_file):
     basepath = os.path.dirname(config_file)
     path = os.path.expanduser(path)

diff --git a/tests/test_io.py b/tests/test_io.py
@@ -7,7 +7,7 @@
 
 import pytest
 from distutils.spawn import find_executable
-from ommprotocol.io import SystemHandler, Positions
+from ommprotocol.io import SystemHandler, Positions, prepare_handler
 from ommprotocol.md import Stage
 from conftest import get_file
 
@@ -32,3 +32,21 @@ def test_formats(top, pos):
     handler = SystemHandler.load(top, **kw)
     stage = Stage(handler, **STAGE_OPTIONS)
     stage.run()
+
+
+def test_prepare_handler():
+    cfg = {'topology': get_file('input.prmtop'),
+           'positions': get_file('input.inpcrd')}
+    handler = prepare_handler(cfg)
+    assert handler.topology.getNumAtoms() == len(handler.positions) == 2269
+    stage = Stage(handler, **STAGE_OPTIONS)
+    stage.run()
+
+
+def test_prepare_handler_tuples():
+    cfg = {'topology': get_file('input.prmtop'),
+           'positions': (get_file('input.inpcrd'),) }
+    handler = prepare_handler(cfg)
+    assert handler.topology.getNumAtoms() == len(handler.positions) == 2269
+    stage = Stage(handler, **STAGE_OPTIONS)
+    stage.run()