Coverage for dynasor/trajectory/ase_trajectory_reader.py: 100%

1from itertools import count 

2from typing import Optional 

3import numpy as np 

4from ase import io 

5from ase import units 

6from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

7from dynasor.trajectory.trajectory_frame import ReaderFrame 

8 

9 

10class ASETrajectoryReader(AbstractTrajectoryReader): 

11 """Read ASE trajectory file 

12 

13 Parameters 

14 ---------- 

15 filename 

16 Name of input file. 

17 length_unit 

18 Unit of length for the input trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). 

19 time_unit 

20 Unit of time for the input trajectory (``'fs'``, ``'ps'``, ``'ns'``). 

21 """ 

22 

23 def __init__( 

24 self, 

25 filename: str, 

26 length_unit: Optional[str] = 'Angstrom', 

27 time_unit: Optional[str] = 'fs', 

28 ): 

29 self._frame_index = count(0) 

30 self._atoms = io.iread(filename, index=':') 

31 self._open = True 

32 

33 # setup units 

34 self.set_unit_scaling_factors(length_unit, time_unit) 

35 

36 def __iter__(self): 

37 return self 

38 

39 def close(self): 

40 self._open = False 

41 self._atoms.close() 

42 

43 def __next__(self): 

44 if not self._open: 

45 raise StopIteration 

46 ind = next(self._frame_index) 

47 a = next(self._atoms) 

48 if 'momenta' in a.arrays: 

49 # ASE returns velocities per ASE internal time unit, not per fs. 

50 vel = units.fs * self.v_factor * a.get_velocities() 

51 else: 

52 vel = None 

53 return ReaderFrame( 

54 frame_index=ind, 

55 n_atoms=len(a), 

56 cell=self.x_factor * a.cell.array, 

57 positions=self.x_factor * a.get_positions(), 

58 velocities=vel, 

59 atom_types=np.array(list(a.symbols)), 

60 )