Skip to content

Commit c9503b5

Browse files
Fixed NewReader to work with auxs
1 parent 5535e3b commit c9503b5

File tree

1 file changed

+15
-10
lines changed
  • package/MDAnalysis/coordinates

1 file changed

+15
-10
lines changed

package/MDAnalysis/coordinates/base.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1642,6 +1642,7 @@ def __init__(self, filename, convert_units=None, **kwargs):
16421642
convert_units = flags['convert_lengths']
16431643
self.convert_units = convert_units
16441644

1645+
self._auxs = {}
16451646
ts_kwargs = {}
16461647
for att in ('dt', 'time_offset'):
16471648
try:
@@ -1657,30 +1658,29 @@ def _update_last_fh_position(self, pos):
16571658
"""Update the last known position of the file handle"""
16581659
self._last_fh_pos = pos
16591660

1660-
def rewind(self):
1661-
self._reopen()
1662-
self.next()
1663-
16641661
def _full_iter(self):
16651662
self._reopen()
16661663
with util.openany(self.filename, 'r') as self._file:
16671664
while True:
16681665
try:
1669-
yield self._read_next_timestep()
1666+
ts = self._read_next_timestep()
1667+
for auxname in self.aux_list:
1668+
ts = self._auxs[auxname].update_ts(ts)
1669+
yield ts
16701670
except (EOFError, IOError):
16711671
self.rewind()
16721672
raise StopIteration
16731673

1674-
def _sliced_iter(self, frames):
1674+
def _sliced_iter(self, start, stop, step):
16751675
with util.openany(self.filename, 'r') as self._file:
1676-
for f in frames:
1677-
yield self._read_frame(f)
1676+
for f in range(start, stop, step):
1677+
yield self._read_frame_with_aux(f)
16781678
self.rewind()
16791679
raise StopIteration
16801680

16811681
def _goto_frame(self, i):
16821682
with util.openany(self.filename, 'r') as self._file:
1683-
ts = self._read_frame(i)
1683+
ts = self._read_frame_with_aux(i)
16841684
return ts
16851685

16861686
def __iter__(self):
@@ -1700,7 +1700,10 @@ def apply_limits(frame):
17001700
elif isinstance(item, (list, np.ndarray)):
17011701
return self._sliced_iter(item)
17021702
elif isinstance(item, slice): # TODO Fix me!
1703-
return self._sliced_iter(item)
1703+
start, stop, step = self.check_slice_indices(
1704+
item.start, item.stop, item.step)
1705+
1706+
return self._sliced_iter(start, stop, step)
17041707

17051708
def __next__(self):
17061709
with util.openany(self.filename, 'r') as self._file:
@@ -1712,6 +1715,8 @@ def __next__(self):
17121715
raise StopIteration
17131716
else:
17141717
self._last_fh_pos = self._file.tell()
1718+
for auxname in self.aux_list:
1719+
ts = self._auxs[auxname].update_ts(ts)
17151720
return ts
17161721

17171722
next = __next__

0 commit comments

Comments
 (0)