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

1from abc import ABC, abstractmethod 

2from typing import Optional 

3 

4from ase import Atoms 

5from ase.units import Bohr, Hartree, kcal, kJ, mol 

6from numpy.typing import NDArray 

7 

8from dynasor.logging_tools import logger 

9 

10 

11def get_forces_from_atoms(atoms: Atoms) -> Optional[NDArray[float]]: 

12 """Return the forces carried by :attr:`atoms`, or ``None`` if it carries none. 

13 

14 ASE stores forces read from a file in the calculator attached to the 

15 :class:`Atoms <ase.Atoms>` object rather than in ``atoms.arrays``, so both places are 

16 searched. The forces are returned unscaled, in the unit of the file. 

17 """ 

18 if atoms.calc is not None and 'forces' in atoms.calc.results: 

19 return atoms.calc.results['forces'] 

20 return atoms.arrays.get('forces') 

21 

22 

23class AbstractTrajectoryReader(ABC): 

24 """Provides a way to iterate through a molecular dynamics trajectory 

25 file. 

26 

27 Each frame/time-step is returned as a trajectory_frame. 

28 """ 

29 

30 # unit conversion tables 

31 lengthunits_to_Angstrom_table = { 

32 'Angstrom': 1.0, 

33 'nm': 10.0, 

34 'pm': 1e-2, 

35 'fm': 1e-5, 

36 } 

37 

38 timeunits_to_fs_table = { 

39 'fs': 1.0, 

40 'ps': 1000, 

41 'ns': 1000000, 

42 } 

43 

44 forceunits_to_eV_per_Angstrom_table = { 

45 'eV/Angstrom': 1.0, 

46 'eV/nm': 0.1, 

47 'kJ/mol/Angstrom': kJ / mol, 

48 'kJ/mol/nm': 0.1 * kJ / mol, 

49 'kcal/mol/Angstrom': kcal / mol, 

50 'kcal/mol/nm': 0.1 * kcal / mol, 

51 'Hartree/Bohr': Hartree / Bohr, 

52 } 

53 

54 # Total number of frames in the trajectory, if known ahead of time; readers that can 

55 # determine this up front (e.g. OvitoTrajectoryReader) should set this in their __init__. 

56 _number_of_total_frames = None 

57 

58 @abstractmethod 

59 def __iter__(self): 

60 """ Iterates through the trajectory file, frame by frame. """ 

61 pass 

62 

63 @abstractmethod 

64 def __next__(self): 

65 """ Gets next trajectory frame. """ 

66 pass 

67 

68 @abstractmethod 

69 def close(self): 

70 """Release resources held by the trajectory reader.""" 

71 pass 

72 

73 @property 

74 def number_of_total_frames(self) -> Optional[int]: 

75 """ Number of frames in the file if the reader knows it when the file is opened; 

76 ``None`` otherwise. """ 

77 return self._number_of_total_frames 

78 

79 def set_unit_scaling_factors(self, length_unit: str, time_unit: str, 

80 force_unit: Optional[str] = None): 

81 # setup units 

82 if length_unit is None: 

83 logger.info('Assuming the trajectory has the default length unit (Ångström), since no ' 

84 'unit was specified.') 

85 self.x_factor = self.lengthunits_to_Angstrom_table['Angstrom'] 

86 elif length_unit not in self.lengthunits_to_Angstrom_table: 

87 raise ValueError(f'Specified length unit {length_unit} is not an available option.') 

88 else: 

89 self.x_factor = self.lengthunits_to_Angstrom_table[length_unit] 

90 if time_unit is None: 

91 logger.info('Assuming the trajectory has the default time unit (fs), since no ' 

92 'unit was specified.') 

93 self.t_factor = self.timeunits_to_fs_table['fs'] 

94 elif time_unit not in self.timeunits_to_fs_table: 

95 raise ValueError(f'Specified time unit {time_unit} is not an available option.') 

96 else: 

97 self.t_factor = self.timeunits_to_fs_table[time_unit] 

98 # Forces are optional, so the unit is resolved without logging here and reported by 

99 # :class:`Trajectory <dynasor.Trajectory>` only for a trajectory that provides them. 

100 if force_unit is None: 

101 force_unit = 'eV/Angstrom' 

102 elif force_unit not in self.forceunits_to_eV_per_Angstrom_table: 

103 raise ValueError(f'Specified force unit {force_unit} is not an available option.') 

104 self.force_unit = force_unit 

105 self.f_factor = self.forceunits_to_eV_per_Angstrom_table[force_unit] 

106 self.v_factor = self.x_factor / self.t_factor