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 dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

6from dynasor.trajectory.trajectory_frame import ReaderFrame 

7 

8 

9class ASETrajectoryReader(AbstractTrajectoryReader): 

10 """Read ASE trajectory file 

11 

12 ... 

13 

14 Parameters 

15 ---------- 

16 filename 

17 Name of input file. 

18 length_unit 

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

20 time_unit 

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

22 """ 

23 

24 def __init__( 

25 self, 

26 filename: str, 

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

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

29 ): 

30 self._frame_index = count(0) 

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

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 pass 

41 

42 def __next__(self): 

43 ind = next(self._frame_index) 

44 a = next(self._atoms) 

45 if 'momenta' in a.arrays: 

46 vel = self.v_factor * a.get_velocities() 

47 else: 

48 vel = None 

49 return ReaderFrame( 

50 frame_index=ind, 

51 n_atoms=len(a), 

52 cell=self.x_factor * a.cell.array.copy('F'), 

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

54 velocities=vel, 

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

56 )