Coverage for gpaw/new/tb/builder.py: 94%

1from __future__ import annotations 

2 

3from math import pi 

4from types import SimpleNamespace 

5 

6import numpy as np 

7from ase.data import atomic_numbers, covalent_radii 

8from ase.neighborlist import neighbor_list 

9from ase.units import Bohr, Ha 

10 

11from gpaw.core.arrays import DistributedArrays 

12from gpaw.core.atom_arrays import AtomArraysLayout 

13from gpaw.core.domain import Domain 

14from gpaw.core.matrix import Matrix 

15from gpaw.lcao.tci import TCIExpansions 

16from gpaw.lfc import BasisFunctions 

17from gpaw.mpi import MPIComm, serial_comm 

18from gpaw.new import zips 

19from gpaw.new.lcao.builder import LCAODFTComponentsBuilder, create_lcao_ibzwfs 

20from gpaw.new.lcao.hamiltonian import CollinearHamiltonianMatrixCalculator 

21from gpaw.new.lcao.wave_functions import LCAOWaveFunctions 

22from gpaw.new.pot_calc import PotentialCalculator 

23from gpaw.setup import Setup 

24from gpaw.spline import Spline 

25from gpaw.utilities.timing import NullTimer 

26from gpaw.typing import Array3D 

27from gpaw.new.scf import SCFContext 

28 

29 

30class TBHamiltonianMatrixCalculator(CollinearHamiltonianMatrixCalculator): 

31 def _calculate_potential_matrix(self, 

32 wfs: LCAOWaveFunctions, 

33 V_xMM: Array3D = None) -> Matrix: 

34 return wfs.V_MM 

35 

36 

37class TBHamiltonian: 

38 def __init__(self, 

39 basis: BasisFunctions): 

40 self.basis = basis 

41 

42 def apply(self): 

43 raise NotImplementedError 

44 

45 def create_hamiltonian_matrix_calculator( 

46 self, 

47 potential) -> TBHamiltonianMatrixCalculator: 

48 ncomponents = potential.dH_asii.dims[0] 

49 dH_saii = [{a: dH_sii[s] 

50 for a, dH_sii in potential.dH_asii.items()} 

51 for s in range(ncomponents)] 

52 

53 V_sxMM = [np.zeros(0) for _ in range(ncomponents)] 

54 

55 return TBHamiltonianMatrixCalculator(V_sxMM, dH_saii, self.basis) 

56 

57 

58class NoGrid(Domain): 

59 def __init__(self, *args, **kwargs): 

60 super().__init__(*args, **kwargs) 

61 self._gd = SimpleNamespace( 

62 get_grid_spacings=lambda: [0, 0, 0], 

63 cell_cv=self.cell_cv, 

64 pbc_c=self.pbc_c, 

65 N_c=[0, 0, 0], 

66 dv=0.0) 

67 self.size = (0, 0, 0) 

68 self.zerobc_c = np.zeros(3, bool) 

69 

70 def empty(self, shape=(), comm=serial_comm, xp=None): 

71 return DummyFunctions(self, shape, comm) 

72 

73 def ranks_from_fractional_positions(self, relpos_ac): 

74 return np.zeros(len(relpos_ac), int) 

75 

76 

77class DummyFunctions(DistributedArrays[NoGrid]): 

78 def __init__(self, 

79 grid: NoGrid, 

80 dims: int | tuple[int, ...] = (), 

81 comm: MPIComm = serial_comm): 

82 DistributedArrays. __init__(self, dims, (), 

83 comm, grid.comm, None, np.nan, 

84 grid.dtype) 

85 self.desc = grid 

86 

87 def integrate(self, other=None): 

88 if other is None: 

89 return np.ones(self.dims) 

90 return np.zeros(self.dims + other.dims) 

91 

92 def new(self, zeroed=False): 

93 return self 

94 

95 def __getitem__(self, index): 

96 if isinstance(index, int): 

97 dims = self.dims[1:] 

98 else: 

99 dims = self.dims 

100 return DummyFunctions(self.desc, dims, comm=self.comm) 

101 

102 def moment(self): 

103 return np.zeros(3) 

104 

105 def to_xp(self, xp): 

106 return self 

107 

108 

109class PSCoreDensities: 

110 xp = np 

111 

112 def __init__(self, grid, relpos_ac): 

113 self.layout = AtomArraysLayout([1] * len(relpos_ac), 

114 grid.comm) 

115 

116 def to_uniform_grid(self, out, scale): 

117 pass 

118 

119 

120class TBPotentialCalculator(PotentialCalculator): 

121 def __init__(self, 

122 xc, 

123 setups, 

124 atoms, 

125 domain_comm): 

126 super().__init__(xc, None, setups, 

127 relpos_ac=atoms.get_scaled_positions(), 

128 environment=None) 

129 self.atoms = atoms.copy() 

130 self.domain_comm = domain_comm 

131 self.force_av = None 

132 self.stress_vv = None 

133 

134 def calculate_pseudo_potential(self, density, ibzwfs, vHt_r): 

135 vt_sR = density.nt_sR 

136 

137 atoms = self.atoms 

138 energy, force_av, stress_vv = pairpot(atoms) 

139 energy /= Ha 

140 self.force_av = force_av * Bohr / Ha 

141 

142 vol = abs(np.linalg.det(atoms.cell[atoms.pbc][:, atoms.pbc])) 

143 self.stress_vv = stress_vv / vol * Bohr**atoms.pbc.sum() / Ha 

144 

145 V_aL = AtomArraysLayout([9] * len(self.atoms), 

146 self.domain_comm).zeros() 

147 return ({'kinetic': 0.0, 

148 'coulomb': 0.0, 

149 'zero': 0.0, 

150 'xc': energy, 

151 'external': 0.0}, 

152 vt_sR, 

153 None, 

154 DummyFunctions(density.nt_sR.desc), 

155 V_aL, 

156 np.nan) 

157 

158 def _move(self, relpos_ac, ndensities): 

159 self.atoms.set_scaled_positions(relpos_ac) 

160 self.force_av = None 

161 self.stress_vv = None 

162 

163 def force_contributions(self, density, potential): 

164 return {}, {}, {a: self.force_av[a:a + 1] 

165 for a in density.D_asii.keys()} 

166 

167 def stress_contribution(self, ibzwfs, density, potential): 

168 return self.stress_vv 

169 

170 

171class DummyXC: 

172 no_forces = False 

173 xc = None 

174 exx_fraction = 0.0 

175 

176 def calculate_paw_correction(self, setup, D_sp, dH_sp): 

177 return 0.0 

178 

179 

180class TBSCFLoop: 

181 def __init__(self, hamiltonian, occ_calc, eigensolver, comm): 

182 self.hamiltonian = hamiltonian 

183 self.occ_calc = occ_calc 

184 self.eigensolver = eigensolver 

185 self.comm = comm 

186 

187 def iterate(self, 

188 ibzwfs, 

189 density, 

190 potential, 

191 energies, 

192 pot_calc, 

193 convergence=None, 

194 maxiter=None, 

195 calculate_forces=None, 

196 log=None): 

197 self.eigensolver.iterate(ibzwfs, density, potential, self.hamiltonian) 

198 e_band, e_entropy, e_extrapolation = ibzwfs.calculate_occs( 

199 self.occ_calc, 

200 nelectrons=density.nvalence - density.charge) 

201 

202 energies.set(band=e_band, 

203 entropy=e_entropy, 

204 extrapolation=e_extrapolation) 

205 

206 yield SCFContext( 

207 log, 

208 1, 

209 energies, 

210 ibzwfs, density, potential, 

211 0.0, 0.0, 0.0, 

212 self.comm, calculate_forces, 

213 pot_calc, False) 

214 

215 # potential, _, _ = pot_calc.calculate( 

216 # density, None, potential.vHt_x) 

217 

218 

219class DummyBasis: 

220 def __init__(self, setups): 

221 self.my_atom_indices = np.arange(len(setups)) 

222 self.Mstart = 0 

223 self.Mstop = setups.nao 

224 

225 def add_to_density(self, nt_sR, f_asi): 

226 pass 

227 

228 def construct_density(self, rho_MM, nt_G, q): 

229 pass 

230 

231 

232class TBDFTComponentsBuilder(LCAODFTComponentsBuilder): 

233 def check_cell(self, cell): 

234 pass 

235 

236 def create_uniform_grids(self): 

237 grid = NoGrid( 

238 self.atoms.cell.complete() / Bohr, 

239 self.atoms.pbc, 

240 dtype=self.dtype, 

241 comm=self.communicators['d']) 

242 return grid, grid 

243 

244 def get_pseudo_core_densities(self): 

245 return PSCoreDensities(self.grid, self.relpos_ac) 

246 

247 def get_pseudo_core_ked(self): 

248 return PSCoreDensities(self.grid, self.relpos_ac) 

249 

250 def create_basis_set(self): 

251 self.basis = DummyBasis(self.setups) 

252 return self.basis 

253 

254 def create_hamiltonian_operator(self): 

255 return TBHamiltonian(self.basis) 

256 

257 def create_potential_calculator(self): 

258 xc = DummyXC() 

259 return TBPotentialCalculator(xc, self.setups, self.atoms, 

260 self.communicators['d']) 

261 

262 def create_scf_loop(self): 

263 occ_calc = self.create_occupation_number_calculator() 

264 hamiltonian = self.create_hamiltonian_operator() 

265 eigensolver = self.create_eigensolver(hamiltonian) 

266 return TBSCFLoop(hamiltonian, occ_calc, eigensolver, 

267 self.communicators['w']) 

268 

269 def create_ibz_wave_functions(self, 

270 basis: BasisFunctions, 

271 potential, 

272 *, 

273 coefficients=None): 

274 assert self.communicators['w'].size == 1 

275 

276 ibzwfs, tciexpansions = create_lcao_ibzwfs( 

277 basis, 

278 self.ibz, self.communicators, self.setups, 

279 self.relpos_ac, self.grid, self.dtype, 

280 self.nbands, self.ncomponents, self.atomdist, self.nelectrons) 

281 

282 vtphit: dict[Setup, list[Spline]] = {} 

283 

284 for setup in self.setups.setups.values(): 

285 try: 

286 vt_r = setup.vt_g 

287 except AttributeError: 

288 vt_r = calculate_pseudo_potential(setup, self.xc.xc)[0] 

289 

290 vt_r[-1] = 0.0 # ??? 

291 vt = setup.rgd.spline(vt_r, points=300) 

292 vtphit_j = [] 

293 for phit in setup.basis_functions_J: 

294 rc = phit.get_cutoff() 

295 r_g = np.linspace(0, rc, 150) 

296 vt_g = vt.map(r_g) / (4 * pi)**0.5 

297 phit_g = phit.map(r_g) 

298 vtphit_j.append(Spline.from_data(phit.l, rc, vt_g * phit_g)) 

299 vtphit[setup] = vtphit_j 

300 

301 vtciexpansions = TCIExpansions([s.basis_functions_J 

302 for s in self.setups], 

303 [vtphit[s] for s in self.setups], 

304 tciexpansions.I_a) 

305 

306 kpt_qc = np.array([wfs.kpt_c for wfs in ibzwfs]) 

307 manytci = vtciexpansions.get_manytci_calculator( 

308 self.setups, self.grid._gd, self.relpos_ac, 

309 kpt_qc, self.dtype, NullTimer()) 

310 

311 manytci.Pindices = manytci.Mindices 

312 my_atom_indices = basis.my_atom_indices 

313 

314 for wfs, V_MM in zips(ibzwfs, manytci.P_qIM(my_atom_indices)): 

315 V_MM = V_MM.toarray() 

316 V_MM += V_MM.T.conj().copy() 

317 M1 = 0 

318 for m in manytci.Mindices.nm_a: 

319 M2 = M1 + m 

320 V_MM[M1:M2, M1:M2] *= 0.5 

321 M1 = M2 

322 wfs.V_MM = Matrix(M2, M2, data=V_MM) 

323 

324 return ibzwfs 

325 

326 

327def pairpot(atoms): 

328 """Simple pair-potential for testing. 

329 

330 >>> from ase import Atoms 

331 >>> r = covalent_radii[1] 

332 >>> atoms = Atoms('H2', [(0, 0, 0), (0, 0, 2 * r)]) 

333 >>> e, f, s = pairpot(atoms) 

334 >>> print(f'{e:.6f} eV') 

335 -9.677419 eV 

336 >>> f 

337 array([[0., 0., 0.], 

338 [0., 0., 0.]]) 

339 

340 """ 

341 radii = {} 

342 symbol_a = atoms.symbols 

343 for symbol in symbol_a: 

344 radii[symbol] = covalent_radii[atomic_numbers[symbol]] 

345 

346 r0 = {} 

347 for s1, r1 in radii.items(): 

348 for s2, r2 in radii.items(): 

349 r0[(s1, s2)] = r1 + r2 

350 rcutmax = 2 * max(r0.values(), default=1.0) 

351 

352 energy = 0.0 

353 force_av = np.zeros((len(atoms), 3)) 

354 stress_vv = np.zeros((3, 3)) 

355 

356 for i, j, d, D_v in zips(*neighbor_list('ijdD', atoms, rcutmax)): 

357 d0 = r0[(symbol_a[i], symbol_a[j])] 

358 e0 = 6.0 / d0 

359 x = d0 / d 

360 if x > 0.5: 

361 energy += 0.5 * e0 * (-5 + x * (24 + x * (-36 + 16 * x))) 

362 f = -0.5 * e0 * (24 + x * (-72 + 48 * x)) * d0 / d**2 

363 F_v = D_v * f / d 

364 force_av[i] += F_v 

365 force_av[j] -= F_v 

366 # print(i, j, d, D_v, F_v) 

367 stress_vv += np.outer(F_v, D_v) 

368 

369 return energy, force_av, stress_vv 

370 

371 

372def calculate_pseudo_potential(setup: Setup, xc): 

373 phit_jg = np.array(setup.data.phit_jg) 

374 rgd = setup.rgd 

375 

376 # Density: 

377 nt_g = np.einsum('jg, j, jg -> g', 

378 phit_jg, setup.f_j, phit_jg) / (4 * pi) 

379 nt_g += setup.data.nct_g * (1 / (4 * pi)**0.5) 

380 

381 # XC: 

382 vt_g = rgd.zeros() 

383 xc.calculate_spherical(rgd, nt_g[np.newaxis], vt_g[np.newaxis]) 

384 

385 # Zero-potential: 

386 vt_g += setup.data.vbar_g / (4 * pi)**0.5 

387 

388 # Coulomb: 

389 g_g = setup.ghat_l[0].map(rgd.r_g) 

390 Q = -rgd.integrate(nt_g) / rgd.integrate(g_g) 

391 rhot_g = nt_g + Q * g_g 

392 vHtr_g = rgd.poisson(rhot_g) 

393 

394 W = rgd.integrate(g_g * vHtr_g, n=-1) / (4 * pi)**0.5 

395 

396 vtr_g = vt_g * rgd.r_g + vHtr_g 

397 

398 vtr_g[1:] /= rgd.r_g[1:] 

399 vtr_g[0] = vtr_g[1] 

400 

401 return vtr_g * (4 * pi)**0.5, W 

402 

403 

404def poly(): 

405 """Polynomium used for pair potential.""" 

406 import matplotlib.pyplot as plt 

407 c = np.linalg.solve([[1, 0.5, 0.25, 0.125], 

408 [1, 1, 1, 1], 

409 [0, 1, 1, 0.75], 

410 [0, 1, 2, 3]], 

411 [0, -1, 0, 0]) 

412 print(c) 

413 d = np.linspace(0.5, 2, 101) 

414 plt.plot(d, c[0] + c[1] / d + c[2] / d**2 + c[3] / d**3) 

415 plt.show() 

416 

417 

418if __name__ == '__main__': 

419 poly()