1# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*- 2# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 3# 4# MDAnalysis --- https://www.mdanalysis.org 5# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors 6# (see the file AUTHORS for the full list of names) 7# 8# Released under the GNU Public Licence, v2 or any higher version 9# 10# Please cite your use of MDAnalysis in published work: 11# 12# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler, 13# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein. 14# MDAnalysis: A Python package for the rapid analysis of molecular dynamics 15# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th 16# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy. 17# 18# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein. 19# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations. 20# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787 21# 22 23 24"""\ 25Base classes --- :mod:`MDAnalysis.coordinates.base` 26=================================================== 27 28Derive other Timestep, FrameIterator, Reader and Writer classes from the classes 29in this module. The derived classes must follow the :ref:`Trajectory API` 30in :mod:`MDAnalysis.coordinates.__init__`. 31 32Timestep 33-------- 34 35A :class:`Timestep` holds information for the current time frame in 36the trajectory. It is one of the central data structures in 37MDAnalysis. 38 39.. class:: Timestep 40 41 .. automethod:: __init__ 42 .. automethod:: from_coordinates 43 .. automethod:: from_timestep 44 .. automethod:: _init_unitcell 45 .. autoattribute:: n_atoms 46 .. attribute::`frame` 47 48 frame number (0-based) 49 50 .. versionchanged:: 0.11.0 51 Frames now 0-based; was 1-based 52 53 .. autoattribute:: time 54 .. autoattribute:: dt 55 .. autoattribute:: positions 56 .. autoattribute:: velocities 57 .. autoattribute:: forces 58 .. autoattribute:: has_positions 59 .. autoattribute:: has_velocities 60 .. autoattribute:: has_forces 61 .. attribute:: _pos 62 63 :class:`numpy.ndarray` of dtype :class:`~numpy.float32` of shape 64 (*n_atoms*, 3) and internal FORTRAN order, holding the raw 65 cartesian coordinates (in MDAnalysis units, i.e. Å). 66 67 .. Note:: 68 69 Normally one does not directly access :attr:`_pos` but uses 70 the :meth:`~MDAnalysis.core.groups.AtomGroup.coordinates` 71 method of an :class:`~MDAnalysis.core.groups.AtomGroup` but 72 sometimes it can be faster to directly use the raw 73 coordinates. Any changes to this array are immediately 74 reflected in atom positions. If the frame is written to a new 75 trajectory then the coordinates are changed. If a new 76 trajectory frame is loaded, then *all* contents of 77 :attr:`_pos` are overwritten. 78 79 .. attribute:: _velocities 80 81 :class:`numpy.ndarray` of dtype :class:`~numpy.float32`. of shape 82 (*n_atoms*, 3), holding the raw velocities (in MDAnalysis 83 units, i.e. typically Å/ps). 84 85 .. Note:: 86 87 Normally velocities are accessed through the 88 :attr:`velocities` or the 89 :meth:`~MDAnalysis.core.groups.AtomGroup.velocities` 90 method of an :class:`~MDAnalysis.core.groups.AtomGroup` 91 92 :attr:`~Timestep._velocities` only exists if the :attr:`has_velocities` 93 flag is True 94 95 .. versionadded:: 0.7.5 96 97 .. attribute:: _forces 98 99 :class:`numpy.ndarray` of dtype :class:`~numpy.float32`. of shape 100 (*n_atoms*, 3), holding the forces 101 102 :attr:`~Timestep._forces` only exists if :attr:`has_forces` 103 is True 104 105 .. versionadded:: 0.11.0 106 Added as optional to :class:`Timestep` 107 108 .. autoattribute:: dimensions 109 .. autoattribute:: triclinic_dimensions 110 .. autoattribute:: volume 111 .. automethod:: __getitem__ 112 .. automethod:: __eq__ 113 .. automethod:: __iter__ 114 .. automethod:: copy 115 .. automethod:: copy_slice 116 117 118FrameIterators 119-------------- 120 121Iterator classes used by the by the :class:`ProtoReader`. 122 123.. autoclass:: FrameIteratorBase 124 125.. autoclass:: FrameIteratorSliced 126 127.. autoclass:: FrameIteratorAll 128 129.. autoclass:: FrameIteratorIndices 130 131 132Readers 133------- 134 135Readers know how to take trajectory data in a given format and present it in a 136common API to the user in MDAnalysis. There are two types of readers: 137 1381. Readers for *multi frame trajectories*, i.e., file formats that typically 139 contain many frames. These readers are typically derived from 140 :class:`ReaderBase`. 141 1422. Readers for *single frame formats*: These file formats only contain a single 143 coordinate set. These readers are derived from 144 :class`:SingleFrameReaderBase`. 145 146The underlying low-level readers handle closing of files in different 147ways. Typically, the MDAnalysis readers try to ensure that files are always 148closed when a reader instance is garbage collected, which relies on 149implementing a :meth:`~ReaderBase.__del__` method. However, in some cases, this 150is not necessary (for instance, for the single frame formats) and then such a 151method can lead to undesirable side effects (such as memory leaks). In this 152case, :class:`ProtoReader` should be used. 153 154 155.. autoclass:: ReaderBase 156 :members: 157 :inherited-members: 158 159.. autoclass:: SingleFrameReaderBase 160 :members: 161 :inherited-members: 162 163.. autoclass:: ProtoReader 164 :members: 165 166 167 168Writers 169------- 170 171Writers know how to write information in a :class:`Timestep` to a trajectory 172file. 173 174.. autoclass:: WriterBase 175 :members: 176 :inherited-members: 177 178 179Helper classes 180-------------- 181 182The following classes contain basic functionality that all readers and 183writers share. 184 185.. autoclass:: IOBase 186 :members: 187 188""" 189from __future__ import absolute_import 190import six 191from six.moves import range 192 193import numpy as np 194import numbers 195import copy 196import warnings 197import weakref 198 199from . import core 200from .. import NoDataError 201from .. import ( 202 _READERS, 203 _SINGLEFRAME_WRITERS, 204 _MULTIFRAME_WRITERS, 205) 206from .. import units 207from ..auxiliary.base import AuxReader 208from ..auxiliary.core import auxreader 209from ..core import flags 210from ..lib.util import asiterable, Namespace 211 212 213class Timestep(object): 214 """Timestep data for one frame 215 216 :Methods: 217 218 ``ts = Timestep(n_atoms)`` 219 220 create a timestep object with space for n_atoms 221 222 .. versionchanged:: 0.11.0 223 Added :meth:`from_timestep` and :meth:`from_coordinates` constructor 224 methods. 225 :class:`Timestep` init now only accepts integer creation. 226 :attr:`n_atoms` now a read only property. 227 :attr:`frame` now 0-based instead of 1-based. 228 Attributes `status` and `step` removed. 229 """ 230 order = 'F' 231 232 def __init__(self, n_atoms, **kwargs): 233 """Create a Timestep, representing a frame of a trajectory 234 235 Parameters 236 ---------- 237 n_atoms : int 238 The total number of atoms this Timestep describes 239 positions : bool, optional 240 Whether this Timestep has position information [``True``] 241 velocities : bool (optional) 242 Whether this Timestep has velocity information [``False``] 243 forces : bool (optional) 244 Whether this Timestep has force information [``False``] 245 reader : Reader (optional) 246 A weak reference to the owning Reader. Used for 247 when attributes require trajectory manipulation (e.g. dt) 248 dt : float (optional) 249 The time difference between frames (ps). If :attr:`time` 250 is set, then `dt` will be ignored. 251 time_offset : float (optional) 252 The starting time from which to calculate time (in ps) 253 254 255 .. versionchanged:: 0.11.0 256 Added keywords for `positions`, `velocities` and `forces`. 257 Can add and remove position/velocity/force information by using 258 the ``has_*`` attribute. 259 """ 260 # readers call Reader._read_next_timestep() on init, incrementing 261 # self.frame to 0 262 self.frame = -1 263 self._n_atoms = n_atoms 264 265 self.data = {} 266 267 for att in ('dt', 'time_offset'): 268 try: 269 self.data[att] = kwargs[att] 270 except KeyError: 271 pass 272 try: 273 # do I have a hook back to the Reader? 274 self._reader = weakref.ref(kwargs['reader']) 275 except KeyError: 276 pass 277 278 # Stupid hack to make it allocate first time round 279 # ie we have to go from not having, to having positions 280 # to make the Timestep allocate 281 self._has_positions = False 282 self._has_velocities = False 283 self._has_forces = False 284 285 # These will allocate the arrays if the has flag 286 # gets set to True 287 self.has_positions = kwargs.get('positions', True) 288 self.has_velocities = kwargs.get('velocities', False) 289 self.has_forces = kwargs.get('forces', False) 290 291 self._unitcell = self._init_unitcell() 292 293 # set up aux namespace for adding auxiliary data 294 self.aux = Namespace() 295 296 297 @classmethod 298 def from_timestep(cls, other, **kwargs): 299 """Create a copy of another Timestep, in the format of this Timestep 300 301 .. versionadded:: 0.11.0 302 """ 303 ts = cls(other.n_atoms, 304 positions=other.has_positions, 305 velocities=other.has_velocities, 306 forces=other.has_forces, 307 **kwargs) 308 ts.frame = other.frame 309 ts.dimensions = other.dimensions 310 try: 311 ts.positions = other.positions.copy(order=cls.order) 312 except NoDataError: 313 pass 314 try: 315 ts.velocities = other.velocities.copy(order=cls.order) 316 except NoDataError: 317 pass 318 try: 319 ts.forces = other.forces.copy(order=cls.order) 320 except NoDataError: 321 pass 322 323 # Optional attributes that don't live in .data 324 # should probably iron out these last kinks 325 for att in ('_frame',): 326 try: 327 setattr(ts, att, getattr(other, att)) 328 except AttributeError: 329 pass 330 331 if hasattr(ts, '_reader'): 332 other._reader = weakref.ref(ts._reader()) 333 334 ts.data = copy.deepcopy(other.data) 335 336 return ts 337 338 @classmethod 339 def from_coordinates(cls, 340 positions=None, 341 velocities=None, 342 forces=None, 343 **kwargs): 344 """Create an instance of this Timestep, from coordinate data 345 346 Can pass position, velocity and force data to form a Timestep. 347 348 .. versionadded:: 0.11.0 349 """ 350 has_positions = positions is not None 351 has_velocities = velocities is not None 352 has_forces = forces is not None 353 354 lens = [len(a) for a in [positions, velocities, forces] 355 if a is not None] 356 if not lens: 357 raise ValueError("Must specify at least one set of data") 358 n_atoms = max(lens) 359 # Check arrays are matched length? 360 if not all(val == n_atoms for val in lens): 361 raise ValueError("Lengths of input data mismatched") 362 363 ts = cls(n_atoms, 364 positions=has_positions, 365 velocities=has_velocities, 366 forces=has_forces, 367 **kwargs) 368 if has_positions: 369 ts.positions = positions 370 if has_velocities: 371 ts.velocities = velocities 372 if has_forces: 373 ts.forces = forces 374 375 return ts 376 377 def _init_unitcell(self): 378 """Create custom datastructure for :attr:`_unitcell`.""" 379 # override for other Timesteps 380 return np.zeros((6), np.float32) 381 382 def __eq__(self, other): 383 """Compare with another Timestep 384 385 .. versionadded:: 0.11.0 386 """ 387 if not isinstance(other, Timestep): 388 return False 389 390 if not self.frame == other.frame: 391 return False 392 393 if not self.n_atoms == other.n_atoms: 394 return False 395 396 if not self.has_positions == other.has_positions: 397 return False 398 if self.has_positions: 399 if not (self.positions == other.positions).all(): 400 return False 401 402 if not self.has_velocities == other.has_velocities: 403 return False 404 if self.has_velocities: 405 if not (self.velocities == other.velocities).all(): 406 return False 407 408 if not self.has_forces == other.has_forces: 409 return False 410 if self.has_forces: 411 if not (self.forces == other.forces).all(): 412 return False 413 414 return True 415 416 def __ne__(self, other): 417 return not self == other 418 419 def __getitem__(self, atoms): 420 """Get a selection of coordinates 421 422 ``ts[i]`` 423 424 return coordinates for the i'th atom (0-based) 425 426 ``ts[start:stop:skip]`` 427 428 return an array of coordinates, where start, stop and skip 429 correspond to atom indices, 430 :attr:`MDAnalysis.core.groups.Atom.index` (0-based) 431 """ 432 if isinstance(atoms, numbers.Integral): 433 return self._pos[atoms] 434 elif isinstance(atoms, (slice, np.ndarray)): 435 return self._pos[atoms] 436 else: 437 raise TypeError 438 439 def __len__(self): 440 return self.n_atoms 441 442 def __iter__(self): 443 """Iterate over coordinates 444 445 ``for x in ts`` 446 447 iterate of the coordinates, atom by atom 448 """ 449 for i in range(self.n_atoms): 450 yield self[i] 451 452 def __repr__(self): 453 desc = "< Timestep {0}".format(self.frame) 454 try: 455 tail = " with unit cell dimensions {0} >".format(self.dimensions) 456 except NotImplementedError: 457 tail = " >" 458 return desc + tail 459 460 def copy(self): 461 """Make an independent ("deep") copy of the whole :class:`Timestep`.""" 462 return self.__deepcopy__() 463 464 def __deepcopy__(self): 465 return self.from_timestep(self) 466 467 def copy_slice(self, sel): 468 """Make a new `Timestep` containing a subset of the original `Timestep`. 469 470 Parameters 471 ---------- 472 sel : array_like or slice 473 The underlying position, velocity, and force arrays are sliced 474 using a :class:`list`, :class:`slice`, or any array-like. 475 476 Returns 477 ------- 478 :class:`Timestep` 479 A `Timestep` object of the same type containing all header 480 information and all atom information relevant to the selection. 481 482 Note 483 ---- 484 The selection must be a 0 based :class:`slice` or array of the atom indices 485 in this :class:`Timestep` 486 487 Example 488 ------- 489 Using a Python :class:`slice` object:: 490 491 new_ts = ts.copy_slice(slice(start, stop, step)) 492 493 Using a list of indices:: 494 495 new_ts = ts.copy_slice([0, 2, 10, 20, 23]) 496 497 498 .. versionadded:: 0.8 499 .. versionchanged:: 0.11.0 500 Reworked to follow new Timestep API. Now will strictly only 501 copy official attributes of the Timestep. 502 503 """ 504 # Detect the size of the Timestep by doing a dummy slice 505 try: 506 pos = self.positions[sel, :] 507 except NoDataError: 508 # It's cool if there's no Data, we'll live 509 pos = None 510 except: 511 raise TypeError("Selection type must be compatible with slicing" 512 " the coordinates") 513 try: 514 vel = self.velocities[sel, :] 515 except NoDataError: 516 vel = None 517 except: 518 raise TypeError("Selection type must be compatible with slicing" 519 " the coordinates") 520 try: 521 force = self.forces[sel, :] 522 except NoDataError: 523 force = None 524 except: 525 raise TypeError("Selection type must be compatible with slicing" 526 " the coordinates") 527 528 new_TS = self.__class__.from_coordinates( 529 positions=pos, 530 velocities=vel, 531 forces=force) 532 533 new_TS._unitcell = self._unitcell.copy() 534 535 new_TS.frame = self.frame 536 537 for att in ('_frame',): 538 try: 539 setattr(new_TS, att, getattr(self, att)) 540 except AttributeError: 541 pass 542 543 if hasattr(self, '_reader'): 544 new_TS._reader = weakref.ref(self._reader()) 545 546 new_TS.data = copy.deepcopy(self.data) 547 548 return new_TS 549 550 @property 551 def n_atoms(self): 552 """A read only view of the number of atoms this Timestep has 553 554 .. versionchanged:: 0.11.0 555 Changed to read only property 556 """ 557 # In future could do some magic here to make setting n_atoms 558 # resize the coordinate arrays, but 559 # - not sure if that is ever useful 560 # - not sure how to manage existing data upon extension 561 return self._n_atoms 562 563 @property 564 def has_positions(self): 565 """A boolean of whether this Timestep has position data 566 567 This can be changed to ``True`` or ``False`` to allocate space for 568 or remove the data. 569 570 .. versionadded:: 0.11.0 571 """ 572 return self._has_positions 573 574 @has_positions.setter 575 def has_positions(self, val): 576 if val and not self._has_positions: 577 # Setting this will always reallocate position data 578 # ie 579 # True -> False -> True will wipe data from first True state 580 self._pos = np.zeros((self.n_atoms, 3), dtype=np.float32, 581 order=self.order) 582 self._has_positions = True 583 elif not val: 584 # Unsetting val won't delete the numpy array 585 self._has_positions = False 586 587 @property 588 def positions(self): 589 """A record of the positions of all atoms in this Timestep 590 591 Setting this attribute will add positions to the Timestep if they 592 weren't originally present. 593 594 Returns 595 ------- 596 positions : numpy.ndarray with dtype numpy.float32 597 position data of shape ``(n_atoms, 3)`` for all atoms 598 599 Raises 600 ------ 601 :exc:`MDAnalysis.exceptions.NoDataError` 602 if the Timestep has no position data 603 604 605 .. versionchanged:: 0.11.0 606 Now can raise :exc:`NoDataError` when no position data present 607 """ 608 if self.has_positions: 609 return self._pos 610 else: 611 raise NoDataError("This Timestep has no positions") 612 613 @positions.setter 614 def positions(self, new): 615 self.has_positions = True 616 self._pos[:] = new 617 618 @property 619 def _x(self): 620 """A view onto the x dimension of position data 621 622 .. versionchanged:: 0.11.0 623 Now read only 624 """ 625 return self.positions[:, 0] 626 627 @property 628 def _y(self): 629 """A view onto the y dimension of position data 630 631 .. versionchanged:: 0.11.0 632 Now read only 633 """ 634 return self.positions[:, 1] 635 636 @property 637 def _z(self): 638 """A view onto the z dimension of position data 639 640 .. versionchanged:: 0.11.0 641 Now read only 642 """ 643 return self.positions[:, 2] 644 645 @property 646 def has_velocities(self): 647 """A boolean of whether this Timestep has velocity data 648 649 This can be changed to ``True`` or ``False`` to allocate space for 650 or remove the data. 651 652 .. versionadded:: 0.11.0 653 """ 654 return self._has_velocities 655 656 @has_velocities.setter 657 def has_velocities(self, val): 658 if val and not self._has_velocities: 659 self._velocities = np.zeros((self.n_atoms, 3), dtype=np.float32, 660 order=self.order) 661 self._has_velocities = True 662 elif not val: 663 self._has_velocities = False 664 665 @property 666 def velocities(self): 667 """A record of the velocities of all atoms in this Timestep 668 669 Setting this attribute will add velocities to the Timestep if they 670 weren't originally present. 671 672 Returns 673 ------- 674 velocities : numpy.ndarray with dtype numpy.float32 675 velocity data of shape ``(n_atoms, 3)`` for all atoms 676 677 Raises 678 ------ 679 :exc:`MDAnalysis.exceptions.NoDataError` 680 if the Timestep has no velocity data 681 682 683 .. versionadded:: 0.11.0 684 """ 685 if self.has_velocities: 686 return self._velocities 687 else: 688 raise NoDataError("This Timestep has no velocities") 689 690 @velocities.setter 691 def velocities(self, new): 692 self.has_velocities = True 693 self._velocities[:] = new 694 695 @property 696 def has_forces(self): 697 """A boolean of whether this Timestep has force data 698 699 This can be changed to ``True`` or ``False`` to allocate space for 700 or remove the data. 701 702 .. versionadded:: 0.11.0 703 """ 704 return self._has_forces 705 706 @has_forces.setter 707 def has_forces(self, val): 708 if val and not self._has_forces: 709 self._forces = np.zeros((self.n_atoms, 3), dtype=np.float32, 710 order=self.order) 711 self._has_forces = True 712 elif not val: 713 self._has_forces = False 714 715 @property 716 def forces(self): 717 """A record of the forces of all atoms in this Timestep 718 719 Setting this attribute will add forces to the Timestep if they 720 weren't originally present. 721 722 Returns 723 ------- 724 forces : numpy.ndarray with dtype numpy.float32 725 force data of shape ``(n_atoms, 3)`` for all atoms 726 727 Raises 728 ------ 729 :exc:`MDAnalysis.exceptions.NoDataError` 730 if the Timestep has no force data 731 732 733 .. versionadded:: 0.11.0 734 """ 735 if self.has_forces: 736 return self._forces 737 else: 738 raise NoDataError("This Timestep has no forces") 739 740 @forces.setter 741 def forces(self, new): 742 self.has_forces = True 743 self._forces[:] = new 744 745 @property 746 def dimensions(self): 747 """unitcell dimensions (*A*, *B*, *C*, *alpha*, *beta*, *gamma*) 748 749 lengths *a*, *b*, *c* are in the MDAnalysis length unit (Å), and 750 angles are in degrees. 751 752 Setting dimensions will populate the underlying native format 753 description of box dimensions 754 """ 755 # The actual Timestep._unitcell depends on the underlying 756 # trajectory format. It can be e.g. six floats representing 757 # the box edges and angles or the 6 unique components of the 758 # box matrix or the full box matrix. 759 return self._unitcell 760 761 @dimensions.setter 762 def dimensions(self, box): 763 self._unitcell[:] = box 764 765 @property 766 def volume(self): 767 """volume of the unitcell""" 768 return core.box_volume(self.dimensions) 769 770 @property 771 def triclinic_dimensions(self): 772 """The unitcell dimensions represented as triclinic vectors 773 774 Returns 775 ------- 776 numpy.ndarray 777 A (3, 3) numpy.ndarray of unit cell vectors 778 779 Examples 780 -------- 781 The unitcell for a given system can be queried as either three 782 vectors lengths followed by their respective angle, or as three 783 triclinic vectors. 784 785 >>> ts.dimensions 786 array([ 13., 14., 15., 90., 90., 90.], dtype=float32) 787 >>> ts.triclinic_dimensions 788 array([[ 13., 0., 0.], 789 [ 0., 14., 0.], 790 [ 0., 0., 15.]], dtype=float32) 791 792 Setting the attribute also works:: 793 794 >>> ts.triclinic_dimensions = [[15, 0, 0], [5, 15, 0], [5, 5, 15]] 795 >>> ts.dimensions 796 array([ 15. , 15.81138802, 16.58312416, 67.58049774, 797 72.45159912, 71.56504822], dtype=float32) 798 799 See Also 800 -------- 801 :func:`MDAnalysis.lib.mdamath.triclinic_vectors` 802 803 804 .. versionadded:: 0.11.0 805 """ 806 return core.triclinic_vectors(self.dimensions) 807 808 @triclinic_dimensions.setter 809 def triclinic_dimensions(self, new): 810 """Set the unitcell for this Timestep as defined by triclinic vectors 811 812 .. versionadded:: 0.11.0 813 """ 814 self.dimensions = core.triclinic_box(*new) 815 816 @property 817 def dt(self): 818 """The time difference in ps between timesteps 819 820 Note 821 ---- 822 This defaults to 1.0 ps in the absence of time data 823 824 825 .. versionadded:: 0.11.0 826 """ 827 try: 828 return self.data['dt'] 829 except KeyError: 830 pass 831 try: 832 dt = self.data['dt'] = self._reader()._get_dt() 833 return dt 834 except AttributeError: 835 pass 836 warnings.warn("Reader has no dt information, set to 1.0 ps") 837 return 1.0 838 839 @dt.setter 840 def dt(self, new): 841 self.data['dt'] = new 842 843 @dt.deleter 844 def dt(self): 845 del self.data['dt'] 846 847 @property 848 def time(self): 849 """The time in ps of this timestep 850 851 This is calculated as:: 852 853 time = ts.data['time_offset'] + ts.time 854 855 Or, if the trajectory doesn't provide time information:: 856 857 time = ts.data['time_offset'] + ts.frame * ts.dt 858 859 .. versionadded:: 0.11.0 860 """ 861 offset = self.data.get('time_offset', 0) 862 try: 863 return self.data['time'] + offset 864 except KeyError: 865 return self.dt * self.frame + offset 866 867 @time.setter 868 def time(self, new): 869 self.data['time'] = new 870 871 @time.deleter 872 def time(self): 873 del self.data['time'] 874 875 876class FrameIteratorBase(object): 877 """ 878 Base iterable over the frames of a trajectory. 879 880 A frame iterable has a length that can be accessed with the :func:`len` 881 function, and can be indexed similarly to a full trajectory. When indexed, 882 indices are resolved relative to the iterable and not relative to the 883 trajectory. 884 885 .. versionadded:: 0.19.0 886 887 """ 888 def __init__(self, trajectory): 889 self._trajectory = trajectory 890 891 def __len__(self): 892 raise NotImplementedError() 893 894 @staticmethod 895 def _avoid_bool_list(frames): 896 if isinstance(frames, list) and frames and isinstance(frames[0], bool): 897 return np.array(frames, dtype=bool) 898 return frames 899 900 @property 901 def trajectory(self): 902 return self._trajectory 903 904 905class FrameIteratorSliced(FrameIteratorBase): 906 """ 907 Iterable over the frames of a trajectory on the basis of a slice. 908 909 Parameters 910 ---------- 911 trajectory: ProtoReader 912 The trajectory over which to iterate. 913 frames: slice 914 A slice to select the frames of interest. 915 916 See Also 917 -------- 918 FrameIteratorBase 919 920 .. versionadded:: 0.19.0 921 922 """ 923 def __init__(self, trajectory, frames): 924 # It would be easier to store directly a range object, as it would 925 # store its parameters in a single place, calculate its length, and 926 # take care of most the indexing. Though, doing so is not compatible 927 # with python 2 where xrange (or range with six) is only an iterator. 928 super(FrameIteratorSliced, self).__init__(trajectory) 929 self._start, self._stop, self._step = trajectory.check_slice_indices( 930 frames.start, frames.stop, frames.step, 931 ) 932 933 def __len__(self): 934 start, stop, step = self.start, self.stop, self.step 935 if (step > 0 and start < stop): 936 # We go from a lesser number to a larger one. 937 return int(1 + (stop - 1 - start) // step) 938 elif (step < 0 and start > stop): 939 # We count backward from a larger number to a lesser one. 940 return int(1 + (start - 1 - stop) // (-step)) 941 else: 942 # The range is empty. 943 return 0 944 945 def __iter__(self): 946 for i in range(self.start, self.stop, self.step): 947 yield self.trajectory[i] 948 self.trajectory.rewind() 949 950 def __getitem__(self, frame): 951 if isinstance(frame, numbers.Integral): 952 length = len(self) 953 if not -length < frame < length: 954 raise IndexError('Index {} is out of range of the range of length {}.' 955 .format(frame, length)) 956 if frame < 0: 957 frame = len(self) + frame 958 frame = self.start + frame * self.step 959 return self.trajectory._read_frame_with_aux(frame) 960 elif isinstance(frame, slice): 961 start = self.start + (frame.start or 0) * self.step 962 if frame.stop is None: 963 stop = self.stop 964 else: 965 stop = self.start + (frame.stop or 0) * self.step 966 step = (frame.step or 1) * self.step 967 968 if step > 0: 969 start = max(0, start) 970 else: 971 stop = max(0, stop) 972 973 new_slice = slice(start, stop, step) 974 return FrameIteratorSliced(self.trajectory, new_slice) 975 else: 976 # Indexing with a lists of bools does not behave the same in all 977 # version of numpy. 978 frame = self._avoid_bool_list(frame) 979 frames = np.array(list(range(self.start, self.stop, self.step)))[frame] 980 return FrameIteratorIndices(self.trajectory, frames) 981 982 @property 983 def start(self): 984 return self._start 985 986 @property 987 def stop(self): 988 return self._stop 989 990 @property 991 def step(self): 992 return self._step 993 994 995class FrameIteratorAll(FrameIteratorBase): 996 """ 997 Iterable over all the frames of a trajectory. 998 999 Parameters 1000 ---------- 1001 trajectory: ProtoReader 1002 The trajectory over which to iterate. 1003 1004 See Also 1005 -------- 1006 FrameIteratorBase 1007 1008 .. versionadded:: 0.19.0 1009 1010 """ 1011 def __init__(self, trajectory): 1012 super(FrameIteratorAll, self).__init__(trajectory) 1013 1014 def __len__(self): 1015 return self.trajectory.n_frames 1016 1017 def __iter__(self): 1018 return iter(self.trajectory) 1019 1020 def __getitem__(self, frame): 1021 return self.trajectory[frame] 1022 1023 1024class FrameIteratorIndices(FrameIteratorBase): 1025 """ 1026 Iterable over the frames of a trajectory listed in a sequence of indices. 1027 1028 Parameters 1029 ---------- 1030 trajectory: ProtoReader 1031 The trajectory over which to iterate. 1032 frames: sequence 1033 A sequence of indices. 1034 1035 See Also 1036 -------- 1037 FrameIteratorBase 1038 """ 1039 def __init__(self, trajectory, frames): 1040 super(FrameIteratorIndices, self).__init__(trajectory) 1041 self._frames = [] 1042 for frame in frames: 1043 if not isinstance(frame, numbers.Integral): 1044 raise TypeError("Frames indices must be integers.") 1045 frame = trajectory._apply_limits(frame) 1046 self._frames.append(frame) 1047 self._frames = tuple(self._frames) 1048 1049 def __len__(self): 1050 return len(self.frames) 1051 1052 def __iter__(self): 1053 for frame in self.frames: 1054 yield self.trajectory._read_frame_with_aux(frame) 1055 1056 def __getitem__(self, frame): 1057 if isinstance(frame, numbers.Integral): 1058 frame = self.frames[frame] 1059 return self.trajectory._read_frame_with_aux(frame) 1060 else: 1061 frame = self._avoid_bool_list(frame) 1062 frames = np.array(self.frames)[frame] 1063 return FrameIteratorIndices(self.trajectory, frames) 1064 1065 @property 1066 def frames(self): 1067 return self._frames 1068 1069 1070class IOBase(object): 1071 """Base class bundling common functionality for trajectory I/O. 1072 1073 .. versionchanged:: 0.8 1074 Added context manager protocol. 1075 """ 1076 1077 #: dict with units of of *time* and *length* (and *velocity*, *force*, 1078 #: ... for formats that support it) 1079 units = {'time': None, 'length': None, 'velocity': None} 1080 1081 def convert_pos_from_native(self, x, inplace=True): 1082 """Conversion of coordinate array x from native units to base units. 1083 1084 Parameters 1085 ---------- 1086 x : array_like 1087 Positions to transform 1088 inplace : bool (optional) 1089 Whether to modify the array inplace, overwriting previous data 1090 1091 Note 1092 ---- 1093 By default, the input `x` is modified in place and also returned. 1094 In-place operations improve performance because allocating new arrays 1095 is avoided. 1096 1097 1098 .. versionchanged:: 0.7.5 1099 Keyword `inplace` can be set to ``False`` so that a 1100 modified copy is returned *unless* no conversion takes 1101 place, in which case the reference to the unmodified `x` is 1102 returned. 1103 1104 """ 1105 f = units.get_conversion_factor( 1106 'length', self.units['length'], flags['length_unit']) 1107 if f == 1.: 1108 return x 1109 if not inplace: 1110 return f * x 1111 x *= f 1112 return x 1113 1114 def convert_velocities_from_native(self, v, inplace=True): 1115 """Conversion of velocities array *v* from native to base units 1116 1117 Parameters 1118 ---------- 1119 v : array_like 1120 Velocities to transform 1121 inplace : bool (optional) 1122 Whether to modify the array inplace, overwriting previous data 1123 1124 Note 1125 ---- 1126 By default, the input *v* is modified in place and also returned. 1127 In-place operations improve performance because allocating new arrays 1128 is avoided. 1129 1130 1131 .. versionadded:: 0.7.5 1132 """ 1133 f = units.get_conversion_factor( 1134 'speed', self.units['velocity'], flags['speed_unit']) 1135 if f == 1.: 1136 return v 1137 if not inplace: 1138 return f * v 1139 v *= f 1140 return v 1141 1142 def convert_forces_from_native(self, force, inplace=True): 1143 """Conversion of forces array *force* from native to base units 1144 1145 Parameters 1146 ---------- 1147 force : array_like 1148 Forces to transform 1149 inplace : bool (optional) 1150 Whether to modify the array inplace, overwriting previous data 1151 1152 Note 1153 ---- 1154 By default, the input *force* is modified in place and also returned. 1155 In-place operations improve performance because allocating new arrays 1156 is avoided. 1157 1158 .. versionadded:: 0.7.7 1159 """ 1160 f = units.get_conversion_factor( 1161 'force', self.units['force'], flags['force_unit']) 1162 if f == 1.: 1163 return force 1164 if not inplace: 1165 return f * force 1166 force *= f 1167 return force 1168 1169 def convert_time_from_native(self, t, inplace=True): 1170 """Convert time *t* from native units to base units. 1171 1172 Parameters 1173 ---------- 1174 t : array_like 1175 Time values to transform 1176 inplace : bool (optional) 1177 Whether to modify the array inplace, overwriting previous data 1178 1179 Note 1180 ---- 1181 By default, the input `t` is modified in place and also returned 1182 (although note that scalar values `t` are passed by value in Python and 1183 hence an in-place modification has no effect on the caller.) In-place 1184 operations improve performance because allocating new arrays is 1185 avoided. 1186 1187 1188 .. versionchanged:: 0.7.5 1189 Keyword `inplace` can be set to ``False`` so that a 1190 modified copy is returned *unless* no conversion takes 1191 place, in which case the reference to the unmodified `x` is 1192 returned. 1193 1194 """ 1195 f = units.get_conversion_factor( 1196 'time', self.units['time'], flags['time_unit']) 1197 if f == 1.: 1198 return t 1199 if not inplace: 1200 return f * t 1201 t *= f 1202 return t 1203 1204 def convert_pos_to_native(self, x, inplace=True): 1205 """Conversion of coordinate array `x` from base units to native units. 1206 1207 Parameters 1208 ---------- 1209 x : array_like 1210 Positions to transform 1211 inplace : bool (optional) 1212 Whether to modify the array inplace, overwriting previous data 1213 1214 Note 1215 ---- 1216 By default, the input `x` is modified in place and also returned. 1217 In-place operations improve performance because allocating new arrays 1218 is avoided. 1219 1220 1221 .. versionchanged:: 0.7.5 1222 Keyword `inplace` can be set to ``False`` so that a 1223 modified copy is returned *unless* no conversion takes 1224 place, in which case the reference to the unmodified `x` is 1225 returned. 1226 1227 """ 1228 f = units.get_conversion_factor( 1229 'length', flags['length_unit'], self.units['length']) 1230 if f == 1.: 1231 return x 1232 if not inplace: 1233 return f * x 1234 x *= f 1235 return x 1236 1237 def convert_velocities_to_native(self, v, inplace=True): 1238 """Conversion of coordinate array `v` from base to native units 1239 1240 Parameters 1241 ---------- 1242 v : array_like 1243 Velocities to transform 1244 inplace : bool (optional) 1245 Whether to modify the array inplace, overwriting previous data 1246 1247 Note 1248 ---- 1249 By default, the input `v` is modified in place and also returned. 1250 In-place operations improve performance because allocating new arrays 1251 is avoided. 1252 1253 1254 .. versionadded:: 0.7.5 1255 """ 1256 f = units.get_conversion_factor( 1257 'speed', flags['speed_unit'], self.units['velocity']) 1258 if f == 1.: 1259 return v 1260 if not inplace: 1261 return f * v 1262 v *= f 1263 return v 1264 1265 def convert_forces_to_native(self, force, inplace=True): 1266 """Conversion of force array *force* from base to native units. 1267 1268 Parameters 1269 ---------- 1270 force : array_like 1271 Forces to transform 1272 inplace : bool (optional) 1273 Whether to modify the array inplace, overwriting previous data 1274 1275 Note 1276 ---- 1277 By default, the input `force` is modified in place and also returned. 1278 In-place operations improve performance because allocating new arrays 1279 is avoided. 1280 1281 1282 .. versionadded:: 0.7.7 1283 """ 1284 f = units.get_conversion_factor( 1285 'force', flags['force_unit'], self.units['force']) 1286 if f == 1.: 1287 return force 1288 if not inplace: 1289 return f * force 1290 force *= f 1291 return force 1292 1293 def convert_time_to_native(self, t, inplace=True): 1294 """Convert time *t* from base units to native units. 1295 1296 Parameters 1297 ---------- 1298 t : array_like 1299 Time values to transform 1300 inplace : bool, optional 1301 Whether to modify the array inplace, overwriting previous data 1302 1303 Note 1304 ---- 1305 By default, the input *t* is modified in place and also 1306 returned. (Also note that scalar values *t* are passed by 1307 value in Python and hence an in-place modification has no 1308 effect on the caller.) 1309 1310 .. versionchanged:: 0.7.5 1311 Keyword *inplace* can be set to ``False`` so that a 1312 modified copy is returned *unless* no conversion takes 1313 place, in which case the reference to the unmodified *x* is 1314 returned. 1315 1316 """ 1317 f = units.get_conversion_factor( 1318 'time', flags['time_unit'], self.units['time']) 1319 if f == 1.: 1320 return t 1321 if not inplace: 1322 return f * t 1323 t *= f 1324 return t 1325 1326 def close(self): 1327 """Close the trajectory file.""" 1328 pass 1329 1330 def __enter__(self): 1331 return self 1332 1333 def __exit__(self, exc_type, exc_val, exc_tb): 1334 # see http://docs.python.org/2/library/stdtypes.html#typecontextmanager 1335 self.close() 1336 return False # do not suppress exceptions 1337 1338 1339class _Readermeta(type): 1340 # Auto register upon class creation 1341 def __init__(cls, name, bases, classdict): 1342 type.__init__(type, name, bases, classdict) 1343 try: 1344 fmt = asiterable(classdict['format']) 1345 except KeyError: 1346 pass 1347 else: 1348 for f in fmt: 1349 f = f.upper() 1350 _READERS[f] = cls 1351 1352 1353class ProtoReader(six.with_metaclass(_Readermeta, IOBase)): 1354 """Base class for Readers, without a :meth:`__del__` method. 1355 1356 Extends :class:`IOBase` with most attributes and methods of a generic 1357 Reader, with the exception of a :meth:`__del__` method. It should be used 1358 as base for Readers that do not need :meth:`__del__`, especially since 1359 having even an empty :meth:`__del__` might lead to memory leaks. 1360 1361 See the :ref:`Trajectory API` definition in 1362 :mod:`MDAnalysis.coordinates.__init__` for the required attributes and 1363 methods. 1364 1365 See Also 1366 -------- 1367 :class:`ReaderBase` 1368 1369 1370 .. versionchanged:: 0.11.0 1371 Frames now 0-based instead of 1-based 1372 """ 1373 1374 #: The appropriate Timestep class, e.g. 1375 #: :class:`MDAnalysis.coordinates.xdrfile.XTC.Timestep` for XTC. 1376 _Timestep = Timestep 1377 1378 def __init__(self): 1379 # initialise list to store added auxiliary readers in 1380 # subclasses should now call super 1381 self._auxs = {} 1382 self._transformations=[] 1383 1384 def __len__(self): 1385 return self.n_frames 1386 1387 @classmethod 1388 def parse_n_atoms(cls, filename, **kwargs): 1389 """Read the coordinate file and deduce the number of atoms 1390 1391 Returns 1392 ------- 1393 n_atoms : int 1394 the number of atoms in the coordinate file 1395 1396 Raises 1397 ------ 1398 NotImplementedError 1399 when the number of atoms can't be deduced 1400 """ 1401 raise NotImplementedError("{} cannot deduce the number of atoms" 1402 "".format(cls.__name__)) 1403 1404 def next(self): 1405 """Forward one step to next frame.""" 1406 try: 1407 ts = self._read_next_timestep() 1408 except (EOFError, IOError): 1409 self.rewind() 1410 raise StopIteration 1411 else: 1412 for auxname in self.aux_list: 1413 ts = self._auxs[auxname].update_ts(ts) 1414 1415 ts = self._apply_transformations(ts) 1416 1417 return ts 1418 1419 def __next__(self): 1420 """Forward one step to next frame when using the `next` builtin.""" 1421 return self.next() 1422 1423 def rewind(self): 1424 """Position at beginning of trajectory""" 1425 self._reopen() 1426 self.next() 1427 1428 @property 1429 def dt(self): 1430 """Time between two trajectory frames in picoseconds.""" 1431 return self.ts.dt 1432 1433 @property 1434 def totaltime(self): 1435 """Total length of the trajectory 1436 1437 The time is calculated as ``(n_frames - 1) * dt``, i.e., we assume that 1438 the first frame no time as elapsed. Thus, a trajectory with two frames will 1439 be considered to have a length of a single time step `dt` and a 1440 "trajectory" with a single frame will be reported as length 0. 1441 1442 """ 1443 return (self.n_frames - 1) * self.dt 1444 1445 @property 1446 def frame(self): 1447 """Frame number of the current time step. 1448 1449 This is a simple short cut to :attr:`Timestep.frame`. 1450 """ 1451 return self.ts.frame 1452 1453 @property 1454 def time(self): 1455 """Time of the current frame in MDAnalysis time units (typically ps). 1456 1457 This is either read straight from the Timestep, or calculated as 1458 time = :attr:`Timestep.frame` * :attr:`Timestep.dt` 1459 """ 1460 return self.ts.time 1461 1462 @property 1463 def trajectory(self): 1464 # Makes a reader effectively commpatible with a FrameIteratorBase 1465 return self 1466 1467 def Writer(self, filename, **kwargs): 1468 """A trajectory writer with the same properties as this trajectory.""" 1469 raise NotImplementedError( 1470 "Sorry, there is no Writer for this format in MDAnalysis. " 1471 "Please file an enhancement request at " 1472 "https://github.com/MDAnalysis/mdanalysis/issues") 1473 1474 def OtherWriter(self, filename, **kwargs): 1475 """Returns a writer appropriate for *filename*. 1476 1477 Sets the default keywords *start*, *step* and *dt* (if 1478 available). *n_atoms* is always set from :attr:`Reader.n_atoms`. 1479 1480 1481 See Also 1482 -------- 1483 :meth:`Reader.Writer` and :func:`MDAnalysis.Writer` 1484 1485 """ 1486 kwargs['n_atoms'] = self.n_atoms # essential 1487 kwargs.setdefault('start', self.frame) 1488 try: 1489 kwargs.setdefault('dt', self.dt) 1490 except KeyError: 1491 pass 1492 return core.writer(filename, **kwargs) 1493 1494 def _read_next_timestep(self, ts=None): # pragma: no cover 1495 # Example from DCDReader: 1496 # if ts is None: 1497 # ts = self.ts 1498 # ts.frame = self._read_next_frame(etc) 1499 # return ts 1500 raise NotImplementedError( 1501 "BUG: Override _read_next_timestep() in the trajectory reader!") 1502 1503 def __iter__(self): 1504 """ Iterate over trajectory frames. """ 1505 self._reopen() 1506 return self 1507 1508 def _reopen(self): 1509 """Should position Reader to just before first frame 1510 1511 Calling next after this should return the first frame 1512 """ 1513 pass 1514 1515 def _apply_limits(self, frame): 1516 if frame < 0: 1517 frame += len(self) 1518 if frame < 0 or frame >= len(self): 1519 raise IndexError("Index {} exceeds length of trajectory ({})." 1520 "".format(frame, len(self))) 1521 return frame 1522 1523 def __getitem__(self, frame): 1524 """Return the Timestep corresponding to *frame*. 1525 1526 If `frame` is a integer then the corresponding frame is 1527 returned. Negative numbers are counted from the end. 1528 1529 If frame is a :class:`slice` then an iterator is returned that 1530 allows iteration over that part of the trajectory. 1531 1532 Note 1533 ---- 1534 *frame* is a 0-based frame index. 1535 """ 1536 if isinstance(frame, numbers.Integral): 1537 frame = self._apply_limits(frame) 1538 return self._read_frame_with_aux(frame) 1539 elif isinstance(frame, (list, np.ndarray)): 1540 if len(frame) != 0 and isinstance(frame[0], (bool, np.bool_)): 1541 # Avoid having list of bools 1542 frame = np.asarray(frame, dtype=np.bool) 1543 # Convert bool array to int array 1544 frame = np.arange(len(self))[frame] 1545 return FrameIteratorIndices(self, frame) 1546 elif isinstance(frame, slice): 1547 start, stop, step = self.check_slice_indices( 1548 frame.start, frame.stop, frame.step) 1549 if start == 0 and stop == len(self) and step == 1: 1550 return FrameIteratorAll(self) 1551 else: 1552 return FrameIteratorSliced(self, frame) 1553 else: 1554 raise TypeError("Trajectories must be an indexed using an integer," 1555 " slice or list of indices") 1556 1557 def _read_frame(self, frame): 1558 """Move to *frame* and fill timestep with data.""" 1559 raise TypeError("{0} does not support direct frame indexing." 1560 "".format(self.__class__.__name__)) 1561 # Example implementation in the DCDReader: 1562 # self._jump_to_frame(frame) 1563 # ts = self.ts 1564 # ts.frame = self._read_next_frame(ts._x, ts._y, ts._z, 1565 # ts._unitcell, 1) 1566 # return ts 1567 1568 def _read_frame_with_aux(self, frame): 1569 """Move to *frame*, updating ts with trajectory, transformations and auxiliary data.""" 1570 ts = self._read_frame(frame) # pylint: disable=assignment-from-no-return 1571 for aux in self.aux_list: 1572 ts = self._auxs[aux].update_ts(ts) 1573 1574 ts = self._apply_transformations(ts) 1575 1576 return ts 1577 1578 def _sliced_iter(self, start, stop, step): 1579 """Generator for slicing a trajectory. 1580 1581 *start* *stop* and *step* are 3 integers describing the slice. 1582 Error checking is not done past this point. 1583 1584 A :exc:`NotImplementedError` is raised if random access to 1585 frames is not implemented. 1586 """ 1587 # override with an appropriate implementation e.g. using self[i] might 1588 # be much slower than skipping steps in a next() loop 1589 try: 1590 for i in range(start, stop, step): 1591 yield self._read_frame_with_aux(i) 1592 self.rewind() 1593 except TypeError: # if _read_frame not implemented 1594 raise TypeError("{0} does not support slicing." 1595 "".format(self.__class__.__name__)) 1596 1597 def check_slice_indices(self, start, stop, step): 1598 """Check frame indices are valid and clip to fit trajectory. 1599 1600 The usage follows standard Python conventions for :func:`range` but see 1601 the warning below. 1602 1603 Parameters 1604 ---------- 1605 start : int or None 1606 Starting frame index (inclusive). ``None`` corresponds to the default 1607 of 0, i.e., the initial frame. 1608 stop : int or None 1609 Last frame index (exclusive). ``None`` corresponds to the default 1610 of n_frames, i.e., it includes the last frame of the trajectory. 1611 step : int or None 1612 step size of the slice, ``None`` corresponds to the default of 1, i.e, 1613 include every frame in the range `start`, `stop`. 1614 1615 Returns 1616 ------- 1617 start, stop, step : tuple (int, int, int) 1618 Integers representing the slice 1619 1620 Warning 1621 ------- 1622 The returned values `start`, `stop` and `step` give the expected result 1623 when passed in :func:`range` but gives unexpected behavior when passed 1624 in a :class:`slice` when ``stop=None`` and ``step=-1`` 1625 1626 This can be a problem for downstream processing of the output from this 1627 method. For example, slicing of trajectories is implemented by passing 1628 the values returned by :meth:`check_slice_indices` to :func:`range` :: 1629 1630 range(start, stop, step) 1631 1632 and using them as the indices to randomly seek to. On the other hand, 1633 in :class:`MDAnalysis.analysis.base.AnalysisBase` the values returned 1634 by :meth:`check_slice_indices` are used to splice the trajectory by 1635 creating a :class:`slice` instance :: 1636 1637 slice(start, stop, step) 1638 1639 This creates a discrepancy because these two lines are not equivalent:: 1640 1641 range(10, -1, -1) # [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0] 1642 range(10)[slice(10, -1, -1)] # [] 1643 1644 """ 1645 1646 slice_dict = {'start': start, 'stop': stop, 'step': step} 1647 for varname, var in slice_dict.items(): 1648 if isinstance(var, numbers.Integral): 1649 slice_dict[varname] = int(var) 1650 elif (var is None): 1651 pass 1652 else: 1653 raise TypeError("{0} is not an integer".format(varname)) 1654 1655 start = slice_dict['start'] 1656 stop = slice_dict['stop'] 1657 step = slice_dict['step'] 1658 1659 if step == 0: 1660 raise ValueError("Step size is zero") 1661 1662 nframes = len(self) 1663 step = step or 1 1664 1665 if start is None: 1666 start = 0 if step > 0 else nframes - 1 1667 elif start < 0: 1668 start += nframes 1669 if start < 0: 1670 start = 0 1671 1672 if step < 0 and start >= nframes: 1673 start = nframes - 1 1674 1675 if stop is None: 1676 stop = nframes if step > 0 else -1 1677 elif stop < 0: 1678 stop += nframes 1679 1680 if step > 0 and stop > nframes: 1681 stop = nframes 1682 1683 return start, stop, step 1684 1685 def __repr__(self): 1686 return ("<{cls} {fname} with {nframes} frames of {natoms} atoms>" 1687 "".format( 1688 cls=self.__class__.__name__, 1689 fname=self.filename, 1690 nframes=self.n_frames, 1691 natoms=self.n_atoms 1692 )) 1693 1694 def add_auxiliary(self, auxname, auxdata, format=None, **kwargs): 1695 """Add auxiliary data to be read alongside trajectory. 1696 1697 Auxiliary data may be any data timeseries from the trajectory additional 1698 to that read in by the trajectory reader. *auxdata* can be an 1699 :class:`~MDAnalysis.auxiliary.base.AuxReader` instance, or the data 1700 itself as e.g. a filename; in the latter case an appropriate 1701 :class:`~MDAnalysis.auxiliary.base.AuxReader` is guessed from the 1702 data/file format. An appropriate `format` may also be directly provided 1703 as a key word argument. 1704 1705 On adding, the AuxReader is initially matched to the current timestep 1706 of the trajectory, and will be updated when the trajectory timestep 1707 changes (through a call to :meth:`next()` or jumping timesteps with 1708 ``trajectory[i]``). 1709 1710 The representative value(s) of the auxiliary data for each timestep (as 1711 calculated by the :class:`~MDAnalysis.auxiliary.base.AuxReader`) are 1712 stored in the current timestep in the ``ts.aux`` namespace under *auxname*; 1713 e.g. to add additional pull force data stored in pull-force.xvg:: 1714 1715 u = MDAnalysis.Universe(PDB, XTC) 1716 u.trajectory.add_auxiliary('pull', 'pull-force.xvg') 1717 1718 The representative value for the current timestep may then be accessed 1719 as ``u.trajectory.ts.aux.pull`` or ``u.trajectory.ts.aux['pull']``. 1720 1721 See Also 1722 -------- 1723 :meth:`remove_auxiliary` 1724 1725 Note 1726 ---- 1727 Auxiliary data is assumed to be time-ordered, with no duplicates. See 1728 the :ref:`Auxiliary API`. 1729 """ 1730 if auxname in self.aux_list: 1731 raise ValueError("Auxiliary data with name {name} already " 1732 "exists".format(name=auxname)) 1733 if isinstance(auxdata, AuxReader): 1734 aux = auxdata 1735 aux.auxname = auxname 1736 else: 1737 aux = auxreader(auxdata, format=format, auxname=auxname, **kwargs) 1738 self._auxs[auxname] = aux 1739 self.ts = aux.update_ts(self.ts) 1740 1741 def remove_auxiliary(self, auxname): 1742 """Clear data and close the :class:`~MDAnalysis.auxiliary.base.AuxReader` 1743 for the auxiliary *auxname*. 1744 1745 See Also 1746 -------- 1747 :meth:`add_auxiliary` 1748 """ 1749 aux = self._check_for_aux(auxname) 1750 aux.close() 1751 del aux 1752 delattr(self.ts.aux, auxname) 1753 1754 @property 1755 def aux_list(self): 1756 """ Lists the names of added auxiliary data. """ 1757 return self._auxs.keys() 1758 1759 def _check_for_aux(self, auxname): 1760 """ Check for the existance of an auxiliary *auxname*. If present, 1761 return the AuxReader; if not, raise ValueError 1762 """ 1763 if auxname in self.aux_list: 1764 return self._auxs[auxname] 1765 else: 1766 raise ValueError("No auxiliary named {name}".format(name=auxname)) 1767 1768 def next_as_aux(self, auxname): 1769 """ Move to the next timestep for which there is at least one step from 1770 the auxiliary *auxname* within the cutoff specified in *auxname*. 1771 1772 This allows progression through the trajectory without encountering 1773 ``NaN`` representative values (unless these are specifically part of the 1774 auxiliary data). 1775 1776 If the auxiliary cutoff is not set, where auxiliary steps are less frequent 1777 (``auxiliary.dt > trajectory.dt``), this allows progression at the 1778 auxiliary pace (rounded to nearest timestep); while if the auxiliary 1779 steps are more frequent, this will work the same as calling 1780 :meth:`next()`. 1781 1782 See the :ref:`Auxiliary API`. 1783 1784 See Also 1785 -------- 1786 :meth:`iter_as_aux` 1787 """ 1788 1789 aux = self._check_for_aux(auxname) 1790 ts = self.ts 1791 # catch up auxiliary if it starts earlier than trajectory 1792 while aux.step_to_frame(aux.step + 1, ts) is None: 1793 next(aux) 1794 # find the next frame that'll have a representative value 1795 next_frame = aux.next_nonempty_frame(ts) 1796 if next_frame is None: 1797 # no more frames with corresponding auxiliary values; stop iteration 1798 raise StopIteration 1799 # some readers set self._frame to -1, rather than self.frame, on 1800 # _reopen; catch here or doesn't read first frame 1801 while self.frame != next_frame or getattr(self, '_frame', 0) == -1: 1802 # iterate trajectory until frame is reached 1803 ts = self.next() 1804 return ts 1805 1806 def iter_as_aux(self, auxname): 1807 """Iterate through timesteps for which there is at least one assigned 1808 step from the auxiliary *auxname* within the cutoff specified in *auxname*. 1809 1810 See Also 1811 -------- 1812 :meth:`next_as_aux` 1813 :meth:`iter_auxiliary` 1814 """ 1815 aux = self._check_for_aux(auxname) 1816 self._reopen() 1817 aux._restart() 1818 while True: 1819 try: 1820 yield self.next_as_aux(auxname) 1821 except StopIteration: 1822 return 1823 1824 def iter_auxiliary(self, auxname, start=None, stop=None, step=None, 1825 selected=None): 1826 """ Iterate through the auxiliary *auxname* independently of the trajectory. 1827 1828 Will iterate over the specified steps of the auxiliary (defaults to all 1829 steps). Allows to access all values in an auxiliary, including those out 1830 of the time range of the trajectory, without having to also iterate 1831 through the trajectory. 1832 1833 After interation, the auxiliary will be repositioned at the current step. 1834 1835 Parameters 1836 ---------- 1837 auxname : str 1838 Name of the auxiliary to iterate over. 1839 (start, stop, step) : optional 1840 Options for iterating over a slice of the auxiliary. 1841 selected : lst | ndarray, optional 1842 List of steps to iterate over. 1843 1844 Yields 1845 ------ 1846 :class:`~MDAnalysis.auxiliary.base.AuxStep` object 1847 1848 See Also 1849 -------- 1850 :meth:`iter_as_aux` 1851 """ 1852 aux = self._check_for_aux(auxname) 1853 if selected is not None: 1854 selection = selected 1855 else: 1856 selection = slice(start, stop, step) 1857 for i in aux[selection]: 1858 yield i 1859 aux.read_ts(self.ts) 1860 1861 def get_aux_attribute(self, auxname, attrname): 1862 """Get the value of *attrname* from the auxiliary *auxname* 1863 1864 Parameters 1865 ---------- 1866 auxname : str 1867 Name of the auxiliary to get value for 1868 attrname : str 1869 Name of gettable attribute in the auxiliary reader 1870 1871 See Also 1872 -------- 1873 :meth:`set_aux_attribute` 1874 """ 1875 aux = self._check_for_aux(auxname) 1876 return getattr(aux, attrname) 1877 1878 def set_aux_attribute(self, auxname, attrname, new): 1879 """ Set the value of *attrname* in the auxiliary *auxname*. 1880 1881 Parameters 1882 ---------- 1883 auxname : str 1884 Name of the auxiliary to alter 1885 attrname : str 1886 Name of settable attribute in the auxiliary reader 1887 new 1888 New value to try set *attrname* to 1889 1890 See Also 1891 -------- 1892 :meth:`get_aux_attribute` 1893 :meth:`rename_aux` - to change the *auxname* attribute 1894 """ 1895 aux = self._check_for_aux(auxname) 1896 if attrname == 'auxname': 1897 self.rename_aux(auxname, new) 1898 else: 1899 setattr(aux, attrname, new) 1900 1901 def rename_aux(self, auxname, new): 1902 """ Change the name of the auxiliary *auxname* to *new*. 1903 1904 Provided there is not already an auxiliary named *new*, the auxiliary 1905 name will be changed in ts.aux namespace, the trajectory's 1906 list of added auxiliaries, and in the auxiliary reader itself. 1907 1908 Parameters 1909 ---------- 1910 auxname : str 1911 Name of the auxiliary to rename 1912 new : str 1913 New name to try set 1914 1915 Raises 1916 ------ 1917 ValueError 1918 If the name *new* is already in use by an existing auxiliary. 1919 """ 1920 aux = self._check_for_aux(auxname) 1921 if new in self.aux_list: 1922 raise ValueError("Auxiliary data with name {name} already " 1923 "exists".format(name=new)) 1924 aux.auxname = new 1925 self._auxs[new] = self._auxs.pop(auxname) 1926 setattr(self.ts.aux, new, self.ts.aux[auxname]) 1927 delattr(self.ts.aux, auxname) 1928 1929 def get_aux_descriptions(self, auxnames=None): 1930 """Get descriptions to allow reloading the specified auxiliaries. 1931 1932 If no auxnames are provided, defaults to the full list of added 1933 auxiliaries. 1934 1935 Passing the resultant description to ``add_auxiliary()`` will allow 1936 recreation of the auxiliary. e.g., to duplicate all auxiliaries into a 1937 second trajectory:: 1938 1939 descriptions = trajectory_1.get_aux_descriptions() 1940 for aux in descriptions: 1941 trajectory_2.add_auxiliary(**aux) 1942 1943 1944 Returns 1945 ------- 1946 list 1947 List of dictionaries of the args/kwargs describing each auxiliary. 1948 1949 See Also 1950 -------- 1951 :meth:`MDAnalysis.auxiliary.base.AuxReader.get_description` 1952 """ 1953 if not auxnames: 1954 auxnames = self.aux_list 1955 descriptions = [self._auxs[aux].get_description() for aux in auxnames] 1956 return descriptions 1957 1958 @property 1959 def transformations(self): 1960 """ Returns the list of transformations""" 1961 return self._transformations 1962 1963 @transformations.setter 1964 def transformations(self, transformations): 1965 if not self._transformations: 1966 self._transformations = transformations 1967 else: 1968 raise ValueError("Transformations are already set") 1969 1970 def add_transformations(self, *transformations): 1971 """Add all transformations to be applied to the trajectory. 1972 1973 This function take as list of transformations as an argument. These 1974 transformations are functions that will be called by the Reader and given 1975 a :class:`Timestep` object as argument, which will be transformed and returned 1976 to the Reader. 1977 The transformations can be part of the :mod:`~MDAnalysis.transformations` 1978 module, or created by the user, and are stored as a list `transformations`. 1979 This list can only be modified once, and further calls of this function will 1980 raise an exception. 1981 1982 .. code-block:: python 1983 1984 u = MDAnalysis.Universe(topology, coordinates) 1985 workflow = [some_transform, another_transform, this_transform] 1986 u.trajectory.add_transformations(*workflow) 1987 1988 The transformations are applied in the order given in the list 1989 `transformations`, i.e., the first transformation is the first 1990 or innermost one to be applied to the :class:`Timestep`. The 1991 example above would be equivalent to 1992 1993 .. code-block:: python 1994 1995 for ts in u.trajectory: 1996 ts = this_transform(another_transform(some_transform(ts))) 1997 1998 1999 Parameters 2000 ---------- 2001 transform_list : list 2002 list of all the transformations that will be applied to the coordinates 2003 in the order given in the list 2004 2005 See Also 2006 -------- 2007 :mod:`MDAnalysis.transformations` 2008 2009 """ 2010 2011 try: 2012 self.transformations = transformations 2013 except ValueError: 2014 raise ValueError("Can't add transformations again. Please create new Universe object") 2015 else: 2016 self.ts = self._apply_transformations(self.ts) 2017 2018 2019 # call reader here to apply the newly added transformation on the 2020 # current loaded frame? 2021 2022 def _apply_transformations(self, ts): 2023 """Applies all the transformations given by the user """ 2024 2025 for transform in self.transformations: 2026 ts = transform(ts) 2027 2028 return ts 2029 2030 2031 2032class ReaderBase(ProtoReader): 2033 """Base class for trajectory readers that extends :class:`ProtoReader` with a 2034 :meth:`__del__` method. 2035 2036 New Readers should subclass :class:`ReaderBase` and properly implement a 2037 :meth:`close` method, to ensure proper release of resources (mainly file 2038 handles). Readers that are inherently safe in this regard should subclass 2039 :class:`ProtoReader` instead. 2040 2041 See the :ref:`Trajectory API` definition in 2042 :mod:`MDAnalysis.coordinates.__init__` for the required attributes and 2043 methods. 2044 2045 See Also 2046 -------- 2047 :class:`ProtoReader` 2048 2049 2050 .. versionchanged:: 0.11.0 2051 Most of the base Reader class definitions were offloaded to 2052 :class:`ProtoReader` so as to allow the subclassing of ReaderBases without a 2053 :meth:`__del__` method. Created init method to create common 2054 functionality, all ReaderBase subclasses must now :func:`super` through this 2055 class. Added attribute :attr:`_ts_kwargs`, which is created in init. 2056 Provides kwargs to be passed to :class:`Timestep` 2057 2058 """ 2059 2060 def __init__(self, filename, convert_units=None, **kwargs): 2061 super(ReaderBase, self).__init__() 2062 2063 self.filename = filename 2064 2065 if convert_units is None: 2066 convert_units = flags['convert_lengths'] 2067 self.convert_units = convert_units 2068 2069 ts_kwargs = {} 2070 for att in ('dt', 'time_offset'): 2071 try: 2072 val = kwargs[att] 2073 except KeyError: 2074 pass 2075 else: 2076 ts_kwargs[att] = val 2077 2078 self._ts_kwargs = ts_kwargs 2079 2080 def copy(self): 2081 """Return independent copy of this Reader. 2082 2083 New Reader will have its own file handle and can seek/iterate 2084 independently of the original. 2085 2086 Will also copy the current state of the Timestep held in 2087 the original Reader 2088 """ 2089 new = self.__class__(self.filename, 2090 n_atoms=self.n_atoms) 2091 if self.transformations: 2092 new.add_transformations(*self.transformations) 2093 # seek the new reader to the same frame we started with 2094 new[self.ts.frame] 2095 # then copy over the current Timestep in case it has 2096 # been modified since initial load 2097 new.ts = self.ts.copy() 2098 for auxname, auxread in self._auxs.items(): 2099 new.add_auxiliary(auxname, auxread.copy()) 2100 return new 2101 2102 def __del__(self): 2103 for aux in self.aux_list: 2104 self._auxs[aux].close() 2105 self.close() 2106 2107 2108class _Writermeta(type): 2109 # Auto register this format upon class creation 2110 def __init__(cls, name, bases, classdict): 2111 type.__init__(type, name, bases, classdict) 2112 try: 2113 # grab the string which describes this format 2114 # could be either 'PDB' or ['PDB', 'ENT'] for multiple formats 2115 fmt = asiterable(classdict['format']) 2116 except KeyError: 2117 # not required however 2118 pass 2119 else: 2120 # does the Writer support single and multiframe writing? 2121 single = classdict.get('singleframe', True) 2122 multi = classdict.get('multiframe', False) 2123 2124 if single: 2125 for f in fmt: 2126 f = f.upper() 2127 _SINGLEFRAME_WRITERS[f] = cls 2128 if multi: 2129 for f in fmt: 2130 f = f.upper() 2131 _MULTIFRAME_WRITERS[f] = cls 2132 2133 2134class WriterBase(six.with_metaclass(_Writermeta, IOBase)): 2135 """Base class for trajectory writers. 2136 2137 See Trajectory API definition in :mod:`MDAnalysis.coordinates.__init__` for 2138 the required attributes and methods. 2139 """ 2140 2141 def convert_dimensions_to_unitcell(self, ts, inplace=True): 2142 """Read dimensions from timestep *ts* and return appropriate unitcell. 2143 2144 The default is to return ``[A,B,C,alpha,beta,gamma]``; if this 2145 is not appropriate then this method has to be overriden. 2146 """ 2147 # override if the native trajectory format does NOT use 2148 # [A,B,C,alpha,beta,gamma] 2149 lengths, angles = ts.dimensions[:3], ts.dimensions[3:] 2150 if not inplace: 2151 lengths = lengths.copy() 2152 lengths = self.convert_pos_to_native(lengths) 2153 return np.concatenate([lengths, angles]) 2154 2155 def write(self, obj): 2156 """Write current timestep, using the supplied `obj`. 2157 2158 Parameters 2159 ---------- 2160 obj : :class:`~MDAnalysis.core.groups.AtomGroup` or :class:`~MDAnalysis.core.universe.Universe` or a :class:`Timestep` 2161 write coordinate information associate with `obj` 2162 2163 Note 2164 ---- 2165 The size of the `obj` must be the same as the number of atoms provided 2166 when setting up the trajectory. 2167 """ 2168 if isinstance(obj, Timestep): 2169 ts = obj 2170 else: 2171 try: 2172 ts = obj.ts 2173 except AttributeError: 2174 try: 2175 # special case: can supply a Universe, too... 2176 ts = obj.trajectory.ts 2177 except AttributeError: 2178 raise TypeError("No Timestep found in obj argument") 2179 return self.write_next_timestep(ts) 2180 2181 def __del__(self): 2182 self.close() 2183 2184 def __repr__(self): 2185 try: 2186 return "< {0!s} {1!r} for {2:d} atoms >".format(self.__class__.__name__, self.filename, self.n_atoms) 2187 except (TypeError, AttributeError): 2188 # no trajectory loaded yet or a Writer that does not need e.g. 2189 # self.n_atoms 2190 return "< {0!s} {1!r} >".format(self.__class__.__name__, self.filename) 2191 2192 def has_valid_coordinates(self, criteria, x): 2193 """Returns ``True`` if all values are within limit values of their formats. 2194 2195 Due to rounding, the test is asymmetric (and *min* is supposed to be negative): 2196 2197 min < x <= max 2198 2199 Parameters 2200 ---------- 2201 criteria : dict 2202 dictionary containing the *max* and *min* values in native units 2203 x : numpy.ndarray 2204 ``(x, y, z)`` coordinates of atoms selected to be written out 2205 2206 Returns 2207 ------- 2208 bool 2209 """ 2210 x = np.ravel(x) 2211 return np.all(criteria["min"] < x) and np.all(x <= criteria["max"]) 2212 2213 # def write_next_timestep(self, ts=None) 2214 2215 2216class SingleFrameReaderBase(ProtoReader): 2217 """Base class for Readers that only have one frame. 2218 2219 To use this base class, define the method :meth:`_read_first_frame` to 2220 read from file `self.filename`. This should populate the attribute 2221 `self.ts` with a :class:`Timestep` object. 2222 2223 .. versionadded:: 0.10.0 2224 .. versionchanged:: 0.11.0 2225 Added attribute "_ts_kwargs" for subclasses 2226 Keywords "dt" and "time_offset" read into _ts_kwargs 2227 """ 2228 _err = "{0} only contains a single frame" 2229 2230 def __init__(self, filename, convert_units=None, n_atoms=None, **kwargs): 2231 super(SingleFrameReaderBase, self).__init__() 2232 2233 self.filename = filename 2234 if convert_units is None: 2235 convert_units = flags['convert_lengths'] 2236 self.convert_units = convert_units 2237 2238 self.n_frames = 1 2239 self.n_atom = n_atoms 2240 2241 ts_kwargs = {} 2242 for att in ('dt', 'time_offset'): 2243 try: 2244 val = kwargs[att] 2245 except KeyError: 2246 pass 2247 else: 2248 ts_kwargs[att] = val 2249 2250 self._ts_kwargs = ts_kwargs 2251 self._read_first_frame() 2252 2253 def copy(self): 2254 """Return independent copy of this Reader. 2255 2256 New Reader will have its own file handle and can seek/iterate 2257 independently of the original. 2258 2259 Will also copy the current state of the Timestep held in 2260 the original Reader 2261 """ 2262 new = self.__class__(self.filename, 2263 n_atoms=self.n_atoms) 2264 new.ts = self.ts.copy() 2265 for auxname, auxread in self._auxs.items(): 2266 new.add_auxiliary(auxname, auxread.copy()) 2267 # since the transformations have already been applied to the frame 2268 # simply copy the property 2269 new.transformations = self.transformations 2270 2271 return new 2272 2273 def _read_first_frame(self): # pragma: no cover 2274 # Override this in subclasses to create and fill a Timestep 2275 pass 2276 2277 def rewind(self): 2278 self._read_first_frame() 2279 for auxname, auxread in self._auxs.items(): 2280 self.ts = auxread.update_ts(self.ts) 2281 super(SingleFrameReaderBase, self)._apply_transformations(self.ts) 2282 2283 def _reopen(self): 2284 pass 2285 2286 def next(self): 2287 raise StopIteration(self._err.format(self.__class__.__name__)) 2288 2289 def __iter__(self): 2290 yield self.ts 2291 return 2292 2293 def _read_frame(self, frame): 2294 if frame != 0: 2295 raise IndexError(self._err.format(self.__class__.__name__)) 2296 2297 return self.ts 2298 2299 def close(self): 2300 # all single frame readers should use context managers to access 2301 # self.filename. Explicitly setting it to the null action in case 2302 # the IOBase.close method is ever changed from that. 2303 pass 2304 2305 def add_transformations(self, *transformations): 2306 """ Add all transformations to be applied to the trajectory. 2307 2308 This function take as list of transformations as an argument. These 2309 transformations are functions that will be called by the Reader and given 2310 a :class:`Timestep` object as argument, which will be transformed and returned 2311 to the Reader. 2312 The transformations can be part of the :mod:`~MDAnalysis.transformations` 2313 module, or created by the user, and are stored as a list `transformations`. 2314 This list can only be modified once, and further calls of this function will 2315 raise an exception. 2316 2317 .. code-block:: python 2318 2319 u = MDAnalysis.Universe(topology, coordinates) 2320 workflow = [some_transform, another_transform, this_transform] 2321 u.trajectory.add_transformations(*workflow) 2322 2323 Parameters 2324 ---------- 2325 transform_list : list 2326 list of all the transformations that will be applied to the coordinates 2327 2328 See Also 2329 -------- 2330 :mod:`MDAnalysis.transformations` 2331 """ 2332 #Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations` 2333 #to avoid unintended behaviour where the coordinates of each frame are transformed 2334 #multiple times when iterating over the trajectory. 2335 #In this method, the trajectory is modified all at once and once only. 2336 2337 super(SingleFrameReaderBase, self).add_transformations(*transformations) 2338 for transform in self.transformations: 2339 self.ts = transform(self.ts) 2340 2341 def _apply_transformations(self, ts): 2342 """ Applies the transformations to the timestep.""" 2343 # Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations` 2344 # to avoid applying the same transformations multiple times on each frame 2345 2346 return ts 2347