Coverage for local_installation/dynasor/trajectory/ase_trajectory_reader.py: 39%

1from ase import io 

2from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

3from dynasor.trajectory.trajectory_frame import ReaderFrame 

4from itertools import count 

5 

6 

7class ASETrajectoryReader(AbstractTrajectoryReader): 

8 """Read ASE trajectory file 

9 

10 ... 

11 

12 Parameters 

13 ---------- 

14 filename 

15 Name of input file. 

16 x_factor 

17 Conversion factor between the length unit used in the trajectory and the internal 

18 dynasor length unit. 

19 t_factor 

20 Conversion factor between the time unit used in the trajectory and the internal 

21 dynasor time unit. 

22 """ 

23 

24 def __init__(self, 

25 filename: str, 

26 x_factor: float = 0.1, 

27 t_factor: float = 1.0): 

28 self.x_factor = x_factor 

29 self.t_factor = t_factor 

30 self.v_factor = x_factor / t_factor 

31 self._frame_index = count(0) 

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

33 

34 def __iter__(self): 

35 return self 

36 

37 def close(self): 

38 if not self._fh.closed: 

39 self._fh.close() 

40 

41 def __next__(self): 

42 ind = self._frame_index 

43 next(self._frame_index) 

44 a = next(self._atoms) 

45 return ReaderFrame( 

46 frame_index=ind, 

47 n_atoms=len(a), 

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

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

50 velocities=self.v_factor * a.get_velocities(), 

51 )