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.abstract_trajectory_reader import get_forces_from_atoms 

8from dynasor.trajectory.trajectory_frame import ReaderFrame 

9 

10 

11class ASETrajectoryReader(AbstractTrajectoryReader): 

12 """Read ASE trajectory file 

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 force_unit 

23 Unit of force for the input trajectory (``'eV/Angstrom'``, ``'eV/nm'``, 

24 ``'kJ/mol/Angstrom'``, ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, 

25 ``'Hartree/Bohr'``). 

26 """ 

27 

28 def __init__( 

29 self, 

30 filename: str, 

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

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

33 force_unit: Optional[str] = None, 

34 ): 

35 self._frame_index = count(0) 

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

37 self._open = True 

38 

39 # setup units 

40 self.set_unit_scaling_factors(length_unit, time_unit, force_unit) 

41 

42 def __iter__(self): 

43 return self 

44 

45 def close(self): 

46 self._open = False 

47 self._atoms.close() 

48 

49 def __next__(self): 

50 if not self._open: 

51 raise StopIteration 

52 ind = next(self._frame_index) 

53 a = next(self._atoms) 

54 if 'momenta' in a.arrays: 

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

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

57 else: 

58 vel = None 

59 forces = get_forces_from_atoms(a) 

60 return ReaderFrame( 

61 frame_index=ind, 

62 n_atoms=len(a), 

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

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

65 velocities=vel, 

66 forces=None if forces is None else self.f_factor * forces, 

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

68 )