Skip to content

Commit

Permalink
Second batch of @orbeckst comments
Browse files Browse the repository at this point in the history
  • Loading branch information
jbarnoud committed Jul 12, 2018
1 parent 167d285 commit 44faa4c
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 8 deletions.
2 changes: 1 addition & 1 deletion package/MDAnalysis/coordinates/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1528,7 +1528,7 @@ def __getitem__(self, frame):
frame = self._apply_limits(frame)
return self._read_frame_with_aux(frame)
elif isinstance(frame, (list, np.ndarray)):
if isinstance(frame[0], (bool, np.bool_)):
if len(frame) != 0 and isinstance(frame[0], (bool, np.bool_)):
# Avoid having list of bools
frame = np.asarray(frame, dtype=np.bool)
# Convert bool array to int array
Expand Down
19 changes: 13 additions & 6 deletions package/MDAnalysis/core/groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -2546,7 +2546,7 @@ def write(self, filename=None, file_format="PDB",
>>> # Write the trajectory in XTC format
>>> ag.write('trajectory.xtc', frames='all')
>>> # Write every other frame of the trajectory in PBD format
>>> ag.abswrite('trajectory.pdb', frames=u.trajectory[::2]
>>> ag.write('trajectory.pdb', frames=u.trajectory[::2])
Parameters
----------
Expand All @@ -2565,19 +2565,22 @@ def write(self, filename=None, file_format="PDB",
file. ``"all"``: write out all bonds, both the original defined and
those guessed by MDAnalysis. ``None``: do not write out bonds.
Default is ``"conect"``.
frames:
frames: array-like or slice or FrameIteratorBase or str, optional
An ensemble of frames to write. The ensemble can be an list or
array of frame indices, a mask of booleans, an instance of
:class:`slice`, or the value returned when a trajectory is indexed.
By default, 'frames' is set to ``None`` and only the current frame
is written.
By default, `frames` is set to ``None`` and only the current frame
is written. If `frames` is set to "all", then all the frame from
trajectory are written.
.. versionchanged:: 0.9.0 Merged with write_selection. This method can
now write both selections out.
.. versionchanged:: 0.19.0
Can write multiframe trajectories with the 'frames' argument.
"""
# TODO: Add a 'verbose' option alongside 'frames'.

# check that AtomGroup actually has any atoms (Issue #434)
if len(self.atoms) == 0:
raise IndexError("Cannot write an AtomGroup with 0 atoms")
Expand Down Expand Up @@ -2647,8 +2650,12 @@ def write(self, filename=None, file_format="PDB",
if frames is None:
w.write(self.atoms)
else:
for _ in trj_frames:
w.write(self.atoms)
current_frame = trj.ts.frame
try:
for _ in trj_frames:
w.write(self.atoms)
finally:
trj[current_frame]
return

try:
Expand Down
12 changes: 11 additions & 1 deletion testsuite/MDAnalysisTests/coordinates/test_reader_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,9 @@ class TestMultiFrameReader(_Multi):
(1, 5, -1), # Stop less than start
(-100, None, None),
(100, None, None), # Outside of range of trajectory
(-2, 10, -2)
(-2, 10, -2),
(0, 0, 1), # empty
(10, 1, 2), # empty
])
def test_slice(self, start, stop, step, reader):
"""Compare the slice applied to trajectory, to slice of list"""
Expand Down Expand Up @@ -226,6 +228,8 @@ def test_getitem(self, slice_cls, sl, reader):
slice(10, 0, -1),
slice(2, 7, 2),
slice(7, 2, -2),
slice(7, 2, 1), # empty
slice(0, 0, 1), # empty
])
def test_getitem_len(self, sl, reader):
traj_iterable = reader[sl]
Expand All @@ -234,6 +238,12 @@ def test_getitem_len(self, sl, reader):
ref = self.reference[sl]
assert len(traj_iterable) == len(ref)

@pytest.mark.parametrize('iter_type', (list, np.array))
def test_getitem_len_empty(self, reader, iter_type):
# Indexing a numpy array with an empty array tends to break.
traj_iterable = reader[iter_type([])]
assert len(traj_iterable) == 0

# All the sl1 slice must be 5 frames long so that the sl2 can be a mask
@pytest.mark.parametrize('sl1', [
[0, 1, 2, 3, 4],
Expand Down
7 changes: 7 additions & 0 deletions testsuite/MDAnalysisTests/core/test_atomgroup.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,13 @@ def test_incompatible_trajectories(self, tmpdir):
with pytest.raises(ValueError):
u1.atoms.write(destination, frames=u2.trajectory)

def test_write_no_traj_move(self, u, tmpdir):
destination = str(tmpdir / 'test.dcd')
u.trajectory[10]
u.atoms.write(destination, frames=[1, 2, 3])
assert u.trajectory.ts.frame == 10


def test_write_selection(self, u, tmpdir):
with tmpdir.as_cwd():
u.atoms.write("test.vmd")
Expand Down

0 comments on commit 44faa4c

Please sign in to comment.