1from math import pi 

2 

3import numpy as np 

4 

5from ase.atoms import Atoms 

6from ase.calculators.calculator import Calculator, kpts2ndarray 

7from ase.units import Bohr, Ha 

8 

9 

10def make_test_dft_calculation(): 

11 a = b = 2.0 

12 c = 6.0 

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

14 symbols='H', 

15 pbc=(1, 1, 0), 

16 cell=(a, b, c), 

17 calculator=TestCalculator()) 

18 return atoms 

19 

20 

21class TestCalculator: 

22 def __init__(self, nk=8): 

23 assert nk % 2 == 0 

24 bzk = [] 

25 weights = [] 

26 ibzk = [] 

27 w = 1.0 / nk**2 

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

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

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

31 bzk.append(k) 

32 if i >= j > 0: 

33 ibzk.append(k) 

34 if i == j: 

35 weights.append(4 * w) 

36 else: 

37 weights.append(8 * w) 

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

39 self.bzk = np.array(bzk) 

40 self.ibzk = np.array(ibzk) 

41 self.weights = np.array(weights) 

42 

43 # Calculate eigenvalues and wave functions: 

44 self.init() 

45 

46 def init(self): 

47 nibzk = len(self.weights) 

48 nbands = 1 

49 

50 V = -1.0 

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

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

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

54 

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

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

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

58 for i in range(2): 

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

60 for j in range(2): 

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

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

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

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

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

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

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

68 

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

70 assert spin == 0 and band == 0 

71 return self.psi[kpt] 

72 

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

74 assert spin == 0 

75 return self.eps[kpt] 

76 

77 def get_number_of_bands(self): 

78 return 1 

79 

80 def get_k_point_weights(self): 

81 return self.weights 

82 

83 def get_number_of_spins(self): 

84 return 1 

85 

86 def get_fermi_level(self): 

87 return 0.0 

88 

89 def get_pseudo_density(self): 

90 n = 0.0 

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

92 if eps >= 0.0: 

93 continue 

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

95 

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

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

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

99 n /= 8 

100 return n 

101 

102 

103class TestPotential(Calculator): 

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

105 

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

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

108 E = 0.0 

109 R = atoms.positions 

110 F = np.zeros_like(R) 

111 for a, r in enumerate(R): 

112 D = R - r 

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

114 x = d - 1.0 

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

116 d[a] = 1 

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

118 energy = 0.25 * E 

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

120 

121 

122class FreeElectrons(Calculator): 

123 """Free-electron band calculator. 

124 

125 Parameters: 

126 

127 nvalence: int 

128 Number of electrons 

129 kpts: dict 

130 K-point specification. 

131 

132 Example: 

133 

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

135 """ 

136 

137 implemented_properties = ['energy'] 

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

139 'nvalence': 0.0, 

140 'nbands': 20, 

141 'gridsize': 7} 

142 

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

144 Calculator.calculate(self, atoms) 

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

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

147 n = self.parameters.gridsize 

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

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

150 eps.sort() 

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

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

153 

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

155 assert spin == 0 

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

157 

158 def get_fermi_level(self): 

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

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

161 return 0.5 * kF**2 * Ha 

162 

163 def get_ibz_k_points(self): 

164 return self.kpts.copy() 

165 

166 def get_number_of_spins(self): 

167 return 1 

168 

169 

170def numeric_force(atoms, a, i, d=0.001): 

171 """Compute numeric force on atom with index a, Cartesian component i, 

172 with finite step of size d 

173 """ 

174 p0 = atoms.get_positions() 

175 p = p0.copy() 

176 p[a, i] += d 

177 atoms.set_positions(p, apply_constraint=False) 

178 eplus = atoms.get_potential_energy() 

179 p[a, i] -= 2 * d 

180 atoms.set_positions(p, apply_constraint=False) 

181 eminus = atoms.get_potential_energy() 

182 atoms.set_positions(p0, apply_constraint=False) 

183 return (eminus - eplus) / (2 * d) 

184 

185 

186def gradient_test(atoms, indices=None): 

187 """ 

188 Use numeric_force to compare analytical and numerical forces on atoms 

189 

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

191 """ 

192 if indices is None: 

193 indices = range(len(atoms)) 

194 f = atoms.get_forces()[indices] 

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

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

197 fn = np.zeros((len(indices), 3)) 

198 for idx, i in enumerate(indices): 

199 for j in range(3): 

200 fn[idx, j] = numeric_force(atoms, i, j, eps) 

201 print('{0:16.12f} {1:20.12f}'.format(eps, abs(fn - f).max())) 

202 return f, fn