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 @abstractmethod 

55 def __iter__(self): 

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

57 pass 

58 

59 @abstractmethod 

60 def __next__(self): 

61 """ Gets next trajectory frame. """ 

62 pass 

63 

64 @abstractmethod 

65 def close(self): 

66 """ Closes down, release resources etc. """ 

67 pass 

68 

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

70 force_unit: Optional[str] = None): 

71 # setup units 

72 if length_unit is None: 

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

74 'unit was specified.') 

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

76 elif length_unit not in self.lengthunits_to_Angstrom_table: 

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

78 else: 

79 self.x_factor = self.lengthunits_to_Angstrom_table[length_unit] 

80 if time_unit is None: 

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

82 'unit was specified.') 

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

84 elif time_unit not in self.timeunits_to_fs_table: 

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

86 else: 

87 self.t_factor = self.timeunits_to_fs_table[time_unit] 

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

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

90 if force_unit is None: 

91 force_unit = 'eV/Angstrom' 

92 elif force_unit not in self.forceunits_to_eV_per_Angstrom_table: 

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

94 self.force_unit = force_unit 

95 self.f_factor = self.forceunits_to_eV_per_Angstrom_table[force_unit] 

96 self.v_factor = self.x_factor / self.t_factor