From dece1c63b0aa7281560e1746288425cabc3b5e3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guillermo=20P=C3=A9rez-Hern=C3=A1ndez?= Date: Tue, 28 Jun 2016 00:06:38 +0200 Subject: [PATCH] [coor/api/save_traj] new optarg save_molecules, minor bugfix (#841) --- pyemma/coordinates/api.py | 13 +++++++++++-- pyemma/coordinates/tests/test_save_traj.py | 4 ++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/pyemma/coordinates/api.py b/pyemma/coordinates/api.py index b4ff9c216..70f0b4eb4 100644 --- a/pyemma/coordinates/api.py +++ b/pyemma/coordinates/api.py @@ -545,7 +545,7 @@ def discretizer(reader, return disc -def save_traj(traj_inp, indexes, outfile, top=None, stride = 1, chunksize=1000, verbose=False): +def save_traj(traj_inp, indexes, outfile, top=None, stride = 1, chunksize=1000, image_molecules=False, verbose=False): r""" Saves a sequence of frames as a single trajectory. Extracts the specified sequence of time/trajectory indexes from traj_inp @@ -597,6 +597,11 @@ def save_traj(traj_inp, indexes, outfile, top=None, stride = 1, chunksize=1000, is a :py:func:`pyemma.coordinates.data.feature_reader.FeatureReader` object, this input variable will be ignored and :py:obj:`traj_inp.chunksize` will be used instead. + image_molecules: boolean, default is False + If set to true, :py:obj:`save_traj` will call the method traj.image_molecules and try to correct for broken + molecules accross periodic boundary conditions. + (http://mdtraj.org/1.7.2/api/generated/mdtraj.Trajectory.html#mdtraj.Trajectory.image_molecules) + verbose : boolean, default is False Verbose output while looking for :py:obj`indexes` in the :py:obj:`traj_inp.trajfiles` @@ -648,10 +653,14 @@ def save_traj(traj_inp, indexes, outfile, top=None, stride = 1, chunksize=1000, if len(trajfiles) < indexes[:, 0].max(): raise ValueError("traj_inp contains %u trajfiles, " "but indexes will ask for file nr. %u" - % (len(trajfiles), indexes[0].max())) + % (len(trajfiles), indexes[:,0].max())) traj = frames_from_files(trajfiles, top, indexes, chunksize, stride, reader=reader) + # Avoid broken molecules + if image_molecules: + traj.image_molecules(inplace=True) + # Return to memory as an mdtraj trajectory object if outfile is None: return traj diff --git a/pyemma/coordinates/tests/test_save_traj.py b/pyemma/coordinates/tests/test_save_traj.py index aee085686..66bab27d7 100644 --- a/pyemma/coordinates/tests/test_save_traj.py +++ b/pyemma/coordinates/tests/test_save_traj.py @@ -94,6 +94,10 @@ def test_reader_input_returns_trajectory(self): self.assertTrue(isinstance(save_traj(self.reader, self.sets, None), md.Trajectory)) + def test_reader_input_returns_trajectory_w_image_molecules(self): + self.assertTrue(isinstance(save_traj(self.reader, self.sets, None, image_molecules=True), + md.Trajectory)) + def test_list_input_save_IO(self): # Test that we're saving to disk alright save_traj(self.trajfiles, self.sets, self.outfile, top=self.pdbfile)