Coverage for /builds/debichem-team/python-ase/ase/calculators/fd.py: 100.00%

1from collections.abc import Iterable 

2from typing import Optional 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.calculators.calculator import BaseCalculator, all_properties 

8 

9 

10class FiniteDifferenceCalculator(BaseCalculator): 

11 """Wrapper calculator using the finite-difference method. 

12 

13 The forces and the stress are computed using the finite-difference method. 

14 

15 .. versionadded:: 3.24.0 

16 """ 

17 

18 implemented_properties = all_properties 

19 

20 def __init__( 

21 self, 

22 calc: BaseCalculator, 

23 eps_disp: float = 1e-6, 

24 eps_strain: float = 1e-6, 

25 ) -> None: 

26 """ 

27 

28 Parameters 

29 ---------- 

30 calc : :class:`~ase.calculators.calculator.BaseCalculator` 

31 ASE Calculator object to be wrapped. 

32 eps_disp : float, default 1e-6 

33 Displacement used for computing forces. 

34 eps_strain : float, default 1e-6 

35 Strain used for computing stress. 

36 

37 """ 

38 super().__init__() 

39 self.calc = calc 

40 self.eps_disp = eps_disp 

41 self.eps_strain = eps_strain 

42 

43 def calculate(self, atoms: Atoms, properties, system_changes) -> None: 

44 atoms = atoms.copy() # copy to not mess up original `atoms` 

45 atoms.calc = self.calc 

46 self.results = { 

47 'energy': atoms.get_potential_energy(), 

48 'forces': calculate_numerical_forces(atoms, eps=self.eps_disp), 

49 'stress': calculate_numerical_stress(atoms, eps=self.eps_strain), 

50 } 

51 self.results['free_energy'] = self.results['energy'] 

52 

53 

54def _numeric_force( 

55 atoms: Atoms, 

56 iatom: int, 

57 icart: int, 

58 eps: float = 1e-6, 

59) -> float: 

60 """Calculate numerical force on a specific atom along a specific direction. 

61 

62 Parameters 

63 ---------- 

64 atoms : :class:`~ase.Atoms` 

65 ASE :class:`~ase.Atoms` object. 

66 iatom : int 

67 Index of atoms. 

68 icart : {0, 1, 2} 

69 Index of Cartesian component. 

70 eps : float, default 1e-6 

71 Displacement. 

72 

73 """ 

74 p0 = atoms.get_positions() 

75 p = p0.copy() 

76 p[iatom, icart] = p0[iatom, icart] + eps 

77 atoms.set_positions(p, apply_constraint=False) 

78 eplus = atoms.get_potential_energy() 

79 p[iatom, icart] = p0[iatom, icart] - eps 

80 atoms.set_positions(p, apply_constraint=False) 

81 eminus = atoms.get_potential_energy() 

82 atoms.set_positions(p0, apply_constraint=False) 

83 return (eminus - eplus) / (2 * eps) 

84 

85 

86def calculate_numerical_forces( 

87 atoms: Atoms, 

88 eps: float = 1e-6, 

89 iatoms: Optional[Iterable[int]] = None, 

90 icarts: Optional[Iterable[int]] = None, 

91) -> np.ndarray: 

92 """Calculate forces numerically based on the finite-difference method. 

93 

94 Parameters 

95 ---------- 

96 atoms : :class:`~ase.Atoms` 

97 ASE :class:`~ase.Atoms` object. 

98 eps : float, default 1e-6 

99 Displacement. 

100 iatoms : Optional[Iterable[int]] 

101 Indices of atoms for which forces are computed. 

102 By default, all atoms are considered. 

103 icarts : Optional[Iterable[int]] 

104 Indices of Cartesian coordinates for which forces are computed. 

105 By default, all three coordinates are considered. 

106 

107 Returns 

108 ------- 

109 forces : np.ndarray 

110 Forces computed numerically based on the finite-difference method. 

111 

112 """ 

113 if iatoms is None: 

114 iatoms = range(len(atoms)) 

115 if icarts is None: 

116 icarts = [0, 1, 2] 

117 return np.array( 

118 [ 

119 [_numeric_force(atoms, iatom, icart, eps) for icart in icarts] 

120 for iatom in iatoms 

121 ] 

122 ) 

123 

124 

125def calculate_numerical_stress( 

126 atoms: Atoms, 

127 eps: float = 1e-6, 

128 voigt: bool = True, 

129) -> np.ndarray: 

130 """Calculate stress numerically based on the finite-difference method. 

131 

132 Parameters 

133 ---------- 

134 atoms : :class:`~ase.Atoms` 

135 ASE :class:`~ase.Atoms` object. 

136 eps : float, default 1e-6 

137 Strain in the Voigt notation. 

138 voigt : bool, default True 

139 If True, the stress is returned in the Voigt notation. 

140 

141 Returns 

142 ------- 

143 stress : np.ndarray 

144 Stress computed numerically based on the finite-difference method. 

145 

146 """ 

147 stress = np.zeros((3, 3), dtype=float) 

148 

149 cell = atoms.cell.copy() 

150 volume = atoms.get_volume() 

151 for i in range(3): 

152 x = np.eye(3) 

153 x[i, i] = 1.0 + eps 

154 atoms.set_cell(cell @ x, scale_atoms=True) 

155 eplus = atoms.get_potential_energy(force_consistent=True) 

156 

157 x[i, i] = 1.0 - eps 

158 atoms.set_cell(cell @ x, scale_atoms=True) 

159 eminus = atoms.get_potential_energy(force_consistent=True) 

160 

161 stress[i, i] = (eplus - eminus) / (2 * eps * volume) 

162 x[i, i] = 1.0 

163 

164 j = i - 2 

165 x[i, j] = x[j, i] = +0.5 * eps 

166 atoms.set_cell(cell @ x, scale_atoms=True) 

167 eplus = atoms.get_potential_energy(force_consistent=True) 

168 

169 x[i, j] = x[j, i] = -0.5 * eps 

170 atoms.set_cell(cell @ x, scale_atoms=True) 

171 eminus = atoms.get_potential_energy(force_consistent=True) 

172 

173 stress[i, j] = stress[j, i] = (eplus - eminus) / (2 * eps * volume) 

174 

175 atoms.set_cell(cell, scale_atoms=True) 

176 

177 return stress.flat[[0, 4, 8, 5, 2, 1]] if voigt else stress