Coverage for /builds/ase/ase/ase/calculators/test.py: 89.43%

123 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3from math import pi 

4 

5import numpy as np 

6 

7from ase.atoms import Atoms 

8from ase.calculators.calculator import Calculator, kpts2ndarray 

9from ase.calculators.fd import calculate_numerical_forces 

10from ase.units import Bohr, Ha 

11 

12 

13def make_test_dft_calculation(): 

14 a = b = 2.0 

15 c = 6.0 

16 atoms = Atoms(positions=[(0, 0, c / 2)], 

17 symbols='H', 

18 pbc=(1, 1, 0), 

19 cell=(a, b, c), 

20 calculator=TestCalculator()) 

21 return atoms 

22 

23 

24class TestCalculator: 

25 def __init__(self, nk=8): 

26 assert nk % 2 == 0 

27 bzk = [] 

28 weights = [] 

29 ibzk = [] 

30 w = 1.0 / nk**2 

31 for i in range(-nk + 1, nk, 2): 

32 for j in range(-nk + 1, nk, 2): 

33 k = (0.5 * i / nk, 0.5 * j / nk, 0) 

34 bzk.append(k) 

35 if i >= j > 0: 

36 ibzk.append(k) 

37 if i == j: 

38 weights.append(4 * w) 

39 else: 

40 weights.append(8 * w) 

41 assert abs(sum(weights) - 1.0) < 1e-12 

42 self.bzk = np.array(bzk) 

43 self.ibzk = np.array(ibzk) 

44 self.weights = np.array(weights) 

45 

46 # Calculate eigenvalues and wave functions: 

47 self.init() 

48 

49 def init(self): 

50 nibzk = len(self.weights) 

51 nbands = 1 

52 

53 V = -1.0 

54 self.eps = 2 * V * (np.cos(2 * pi * self.ibzk[:, 0]) + 

55 np.cos(2 * pi * self.ibzk[:, 1])) 

56 self.eps.shape = (nibzk, nbands) 

57 

58 self.psi = np.zeros((nibzk, 20, 20, 60), complex) 

59 phi = np.empty((2, 2, 20, 20, 60)) 

60 z = np.linspace(-1.5, 1.5, 60, endpoint=False) 

61 for i in range(2): 

62 x = np.linspace(0, 1, 20, endpoint=False) - i 

63 for j in range(2): 

64 y = np.linspace(0, 1, 20, endpoint=False) - j 

65 r = (((x[:, None]**2 + 

66 y**2)[:, :, None] + 

67 z**2)**0.5).clip(0, 1) 

68 phi = 1.0 - r**2 * (3.0 - 2.0 * r) 

69 phase = np.exp(pi * 2j * np.dot(self.ibzk, (i, j, 0))) 

70 self.psi += phase[:, None, None, None] * phi 

71 

72 def get_pseudo_wave_function(self, band=0, kpt=0, spin=0): 

73 assert spin == 0 and band == 0 

74 return self.psi[kpt] 

75 

76 def get_eigenvalues(self, kpt=0, spin=0): 

77 assert spin == 0 

78 return self.eps[kpt] 

79 

80 def get_number_of_bands(self): 

81 return 1 

82 

83 def get_k_point_weights(self): 

84 return self.weights 

85 

86 def get_number_of_spins(self): 

87 return 1 

88 

89 def get_fermi_level(self): 

90 return 0.0 

91 

92 def get_pseudo_density(self): 

93 n = 0.0 

94 for w, eps, psi in zip(self.weights, self.eps[:, 0], self.psi): 

95 if eps >= 0.0: 

96 continue 

97 n += w * (psi * psi.conj()).real 

98 

99 n[1:] += n[:0:-1].copy() 

100 n[:, 1:] += n[:, :0:-1].copy() 

101 n += n.transpose((1, 0, 2)).copy() 

102 n /= 8 

103 return n 

104 

105 

106class TestPotential(Calculator): 

107 implemented_properties = ['energy', 'forces'] 

108 

109 def calculate(self, atoms, properties, system_changes): 

110 Calculator.calculate(self, atoms, properties, system_changes) 

111 E = 0.0 

112 R = atoms.positions 

113 F = np.zeros_like(R) 

114 for a, r in enumerate(R): 

115 D = R - r 

116 d = (D**2).sum(1)**0.5 

117 x = d - 1.0 

118 E += np.vdot(x, x) 

119 d[a] = 1 

120 F -= (x / d)[:, None] * D 

121 energy = 0.25 * E 

122 self.results = {'energy': energy, 'forces': F} 

123 

124 

125class FreeElectrons(Calculator): 

126 """Free-electron band calculator. 

127 

128 Parameters: 

129 

130 nvalence: int 

131 Number of electrons 

132 kpts: dict 

133 K-point specification. 

134 

135 Example: 

136 >>> from ase.calculators.test import FreeElectrons 

137 >>> calc = FreeElectrons(nvalence=1, kpts={'path': 'GXL'}) 

138 """ 

139 

140 implemented_properties = ['energy'] 

141 default_parameters = {'kpts': np.zeros((1, 3)), 

142 'nvalence': 0.0, 

143 'nbands': 20, 

144 'gridsize': 7} 

145 

146 def calculate(self, atoms, properties, system_changes): 

147 Calculator.calculate(self, atoms) 

148 self.kpts = kpts2ndarray(self.parameters.kpts, atoms) 

149 icell = atoms.cell.reciprocal() * 2 * np.pi * Bohr 

150 n = self.parameters.gridsize 

151 offsets = np.indices((n, n, n)).T.reshape((n**3, 1, 3)) - n // 2 

152 eps = 0.5 * (np.dot(self.kpts + offsets, icell)**2).sum(2).T 

153 eps.sort() 

154 self.eigenvalues = eps[:, :self.parameters.nbands] * Ha 

155 self.results = {'energy': 0.0} 

156 

157 def get_eigenvalues(self, kpt, spin=0): 

158 assert spin == 0 

159 return self.eigenvalues[kpt].copy() 

160 

161 def get_fermi_level(self): 

162 v = self.atoms.get_volume() / Bohr**3 

163 kF = (self.parameters.nvalence / v * 3 * np.pi**2)**(1 / 3) 

164 return 0.5 * kF**2 * Ha 

165 

166 def get_ibz_k_points(self): 

167 return self.kpts.copy() 

168 

169 def get_number_of_spins(self): 

170 return 1 

171 

172 

173def gradient_test(atoms, indices=None): 

174 """ 

175 Use numeric_force to compare analytical and numerical forces on atoms 

176 

177 If indices is None, test is done on all atoms. 

178 """ 

179 if indices is None: 

180 indices = range(len(atoms)) 

181 f = atoms.get_forces()[indices] 

182 print('{:>16} {:>20}'.format('eps', 'max(abs(df))')) 

183 for eps in np.logspace(-1, -8, 8): 

184 fn = calculate_numerical_forces(atoms, eps, indices) 

185 print(f'{eps:16.12f} {abs(fn - f).max():20.12f}') 

186 return f, fn